vectorstores: add WithEndpoint option for Azure AI Search (tmc#1308)

Author: thecoons

vectorstores/azureaisearch/azureaisearch_httprr_test.go

@@ -46,14 +46,11 @@ func TestStoreHTTPRR_CreateIndex(t *testing.T) {
 		apiKey = envKey
 	}
 
-	// Set endpoint via environment variable
-	os.Setenv("AZURE_AI_SEARCH_ENDPOINT", endpoint)
-	defer os.Unsetenv("AZURE_AI_SEARCH_ENDPOINT")
-
 	store, err := New(
 		WithAPIKey(apiKey),
 		WithEmbedder(&mockEmbedder{}),
 		WithHTTPClient(rr.Client()),
+		WithEndpoint(endpoint),
 	)
 	require.NoError(t, err)
 
@@ -81,14 +78,11 @@ func TestStoreHTTPRR_AddDocuments(t *testing.T) {
 		apiKey = envKey
 	}
 
-	// Set endpoint via environment variable
-	os.Setenv("AZURE_AI_SEARCH_ENDPOINT", endpoint)
-	defer os.Unsetenv("AZURE_AI_SEARCH_ENDPOINT")
-
 	store, err := New(
 		WithAPIKey(apiKey),
 		WithEmbedder(&mockEmbedder{}),
 		WithHTTPClient(rr.Client()),
+		WithEndpoint(endpoint),
 	)
 	require.NoError(t, err)
 
@@ -133,14 +127,11 @@ func TestStoreHTTPRR_SimilaritySearch(t *testing.T) {
 		apiKey = envKey
 	}
 
-	// Set endpoint via environment variable
-	os.Setenv("AZURE_AI_SEARCH_ENDPOINT", endpoint)
-	defer os.Unsetenv("AZURE_AI_SEARCH_ENDPOINT")
-
 	store, err := New(
 		WithAPIKey(apiKey),
 		WithEmbedder(&mockEmbedder{}),
 		WithHTTPClient(rr.Client()),
+		WithEndpoint(endpoint),
 	)
 	require.NoError(t, err)
 
@@ -169,14 +160,11 @@ func TestStoreHTTPRR_DeleteIndex(t *testing.T) {
 		apiKey = envKey
 	}
 
-	// Set endpoint via environment variable
-	os.Setenv("AZURE_AI_SEARCH_ENDPOINT", endpoint)
-	defer os.Unsetenv("AZURE_AI_SEARCH_ENDPOINT")
-
 	store, err := New(
 		WithAPIKey(apiKey),
 		WithEmbedder(&mockEmbedder{}),
 		WithHTTPClient(rr.Client()),
+		WithEndpoint(endpoint),
 	)
 	require.NoError(t, err)
 
@@ -203,14 +191,11 @@ func TestStoreHTTPRR_ListIndexes(t *testing.T) {
 		apiKey = envKey
 	}
 
-	// Set endpoint via environment variable
-	os.Setenv("AZURE_AI_SEARCH_ENDPOINT", endpoint)
-	defer os.Unsetenv("AZURE_AI_SEARCH_ENDPOINT")
-
 	store, err := New(
 		WithAPIKey(apiKey),
 		WithEmbedder(&mockEmbedder{}),
 		WithHTTPClient(rr.Client()),
+		WithEndpoint(endpoint),
 	)
 	require.NoError(t, err)
 

vectorstores/azureaisearch/options.go

@@ -68,6 +68,13 @@ func WithAPIKey(azureAISearchAPIKey string) Option {
 	}
 }
 
+// WithEndpoint is an option for setting the azure AI search endpoint.
+func WithEndpoint(endpoint string) Option {
+	return func(s *Store) {
+		s.azureAISearchEndpoint = strings.TrimSuffix(endpoint, "/")
+	}
+}
+
 func applyClientOptions(s *Store, opts ...Option) error {
 	for _, opt := range opts {
 		opt(s)