Increase test coverage

Author: pamelafox

app/backend/app.py

@@ -560,7 +560,7 @@ async def setup_clients():
         # Wait until token is needed to fetch for the first time
         current_app.config[CONFIG_SPEECH_SERVICE_TOKEN] = None
 
-    openai_client = setup_openai_client(
+    openai_client, azure_openai_endpoint = setup_openai_client(
         openai_host=OPENAI_HOST,
         azure_credential=azure_credential,
         azure_openai_service=AZURE_OPENAI_SERVICE,

app/backend/prepdocs.py

@@ -169,23 +169,19 @@ def setup_embeddings_service(
         logger.info("Not setting up embeddings service")
         return None
 
-    azure_endpoint = None
-    azure_deployment = None
     if openai_host in [OpenAIHost.AZURE, OpenAIHost.AZURE_CUSTOM]:
         if azure_openai_endpoint is None:
             raise ValueError("Azure OpenAI endpoint must be provided when using Azure OpenAI embeddings")
         if azure_openai_deployment is None:
             raise ValueError("Azure OpenAI deployment must be provided when using Azure OpenAI embeddings")
-        azure_endpoint = azure_openai_endpoint
-        azure_deployment = azure_openai_deployment
 
     return OpenAIEmbeddings(
         open_ai_client=open_ai_client,
         open_ai_model_name=emb_model_name,
         open_ai_dimensions=emb_model_dimensions,
         disable_batch=disable_batch_vectors,
-        azure_deployment_name=azure_deployment,
-        azure_endpoint=azure_endpoint,
+        azure_deployment_name=azure_openai_deployment,
+        azure_endpoint=azure_openai_endpoint,
     )
 
 
@@ -197,33 +193,39 @@ def setup_openai_client(
     azure_openai_custom_url: Optional[str] = None,
     openai_api_key: Optional[str] = None,
     openai_organization: Optional[str] = None,
-):
+) -> tuple[AsyncOpenAI, Optional[str]]:
     if openai_host not in OpenAIHost:
         raise ValueError(f"Invalid OPENAI_HOST value: {openai_host}. Must be one of {[h.value for h in OpenAIHost]}.")
 
     openai_client: AsyncOpenAI
+    azure_openai_endpoint: Optional[str] = None
 
     if openai_host in [OpenAIHost.AZURE, OpenAIHost.AZURE_CUSTOM]:
+        base_url: Optional[str] = None
+        api_key_or_token: Optional[str | AsyncTokenCredential] = None
         if openai_host == OpenAIHost.AZURE_CUSTOM:
             logger.info("OPENAI_HOST is azure_custom, setting up Azure OpenAI custom client")
             if not azure_openai_custom_url:
                 raise ValueError("AZURE_OPENAI_CUSTOM_URL must be set when OPENAI_HOST is azure_custom")
-            endpoint = azure_openai_custom_url
+            base_url = azure_openai_custom_url
         else:
             logger.info("OPENAI_HOST is azure, setting up Azure OpenAI client")
             if not azure_openai_service:
                 raise ValueError("AZURE_OPENAI_SERVICE must be set when OPENAI_HOST is azure")
-            endpoint = f"https://{azure_openai_service}.openai.azure.com/openai/v1"
+            azure_openai_endpoint = "https://{azure_openai_service}.openai.azure.com/"
+            base_url = f"{azure_openai_endpoint}/openai/v1"
         if azure_openai_api_key:
             logger.info("AZURE_OPENAI_API_KEY_OVERRIDE found, using as api_key for Azure OpenAI client")
-            openai_client = AsyncOpenAI(base_url=endpoint, api_key=azure_openai_api_key)
+            api_key_or_token = azure_openai_api_key
         else:
             logger.info("Using Azure credential (passwordless authentication) for Azure OpenAI client")
-            token_provider = get_bearer_token_provider(azure_credential, "https://cognitiveservices.azure.com/.default")
-            openai_client = AsyncOpenAI(
-                base_url=endpoint,
-                api_key=token_provider,
+            api_key_or_token = get_bearer_token_provider(
+                azure_credential, "https://cognitiveservices.azure.com/.default"
             )
+        openai_client = AsyncOpenAI(
+            base_url=base_url,
+            api_key=api_key_or_token,
+        )
     elif openai_host == OpenAIHost.LOCAL:
         logger.info("OPENAI_HOST is local, setting up local OpenAI client for OPENAI_BASE_URL with no key")
         openai_client = AsyncOpenAI(
@@ -240,7 +242,7 @@ def setup_openai_client(
             api_key=openai_api_key,
             organization=openai_organization,
         )
-    return openai_client
+    return openai_client, azure_openai_endpoint
 
 
 def setup_file_processors(
@@ -349,7 +351,7 @@ async def main(strategy: Strategy, setup_index: bool = True):
     await strategy.run()
 
 
-if __name__ == "__main__":
+if __name__ == "__main__":  # pragma: no cover
     parser = argparse.ArgumentParser(
         description="Prepare documents by extracting content from PDFs, splitting content into sections, uploading to blob storage, and indexing in a search index."
     )
@@ -500,7 +502,8 @@ async def main(strategy: Strategy, setup_index: bool = True):
     emb_model_dimensions = 1536
     if os.getenv("AZURE_OPENAI_EMB_DIMENSIONS"):
         emb_model_dimensions = int(os.environ["AZURE_OPENAI_EMB_DIMENSIONS"])
-    openai_client = setup_openai_client(
+
+    openai_client, azure_openai_endpoint = setup_openai_client(
         openai_host=OPENAI_HOST,
         azure_credential=azd_credential,
         azure_openai_service=os.getenv("AZURE_OPENAI_SERVICE"),
@@ -509,17 +512,13 @@ async def main(strategy: Strategy, setup_index: bool = True):
         openai_api_key=clean_key_if_exists(os.getenv("OPENAI_API_KEY")),
         openai_organization=os.getenv("OPENAI_ORGANIZATION"),
     )
-    azure_embedding_endpoint = os.getenv("AZURE_OPENAI_ENDPOINT") or os.getenv("AZURE_OPENAI_CUSTOM_URL")
-    if not azure_embedding_endpoint and OPENAI_HOST == OpenAIHost.AZURE:
-        if service := os.getenv("AZURE_OPENAI_SERVICE"):
-            azure_embedding_endpoint = f"https://{service}.openai.azure.com"
     openai_embeddings_service = setup_embeddings_service(
         open_ai_client=openai_client,
         openai_host=OPENAI_HOST,
         emb_model_name=os.environ["AZURE_OPENAI_EMB_MODEL_NAME"],
         emb_model_dimensions=emb_model_dimensions,
         azure_openai_deployment=os.getenv("AZURE_OPENAI_EMB_DEPLOYMENT"),
-        azure_openai_endpoint=azure_embedding_endpoint,
+        azure_openai_endpoint=azure_openai_endpoint,
         disable_vectors=dont_use_vectors,
         disable_batch_vectors=args.disablebatchvectors,
     )

app/backend/prepdocslib/integratedvectorizerstrategy.py

@@ -29,7 +29,7 @@
 logger = logging.getLogger("scripts")
 
 
-class IntegratedVectorizerStrategy(Strategy):
+class IntegratedVectorizerStrategy(Strategy):  # pragma: no cover
     """
     Strategy for ingesting and vectorizing documents into a search service from files stored storage account
     """

tests/test_prepdocs.py

@@ -313,6 +313,35 @@ def test_setup_embeddings_service_requires_endpoint_for_azure() -> None:
         )
 
 
+def test_setup_embeddings_service_requires_deployment_for_azure() -> None:
+    with pytest.raises(ValueError):
+        prepdocs.setup_embeddings_service(
+            open_ai_client=MockClient(
+                MockEmbeddingsClient(
+                    openai.types.CreateEmbeddingResponse(
+                        object="list",
+                        data=[],
+                        model="text-embedding-3-large",
+                        usage=Usage(prompt_tokens=0, total_tokens=0),
+                    )
+                )
+            ),
+            openai_host=prepdocs.OpenAIHost.AZURE,
+            emb_model_name=MOCK_EMBEDDING_MODEL_NAME,
+            emb_model_dimensions=MOCK_EMBEDDING_DIMENSIONS,
+            azure_openai_deployment=None,
+            azure_openai_endpoint="https://service.openai.azure.com",
+        )
+
+
+def test_setup_openai_client_requires_valid_host() -> None:
+    with pytest.raises(ValueError, match="Invalid OPENAI_HOST value"):
+        prepdocs.setup_openai_client(
+            openai_host="invalid_host",  # type: ignore
+            azure_credential=MockAzureCredential(),
+        )
+
+
 @pytest.mark.asyncio
 async def test_openai_embeddings_use_deployment_for_azure_model():
     class RecordingEmbeddingsClient:

tests/test_searchmanager.py

@@ -12,6 +12,7 @@
     SearchIndex,
     SearchIndexPermissionFilterOption,
     SimpleField,
+    VectorSearch,
 )
 from openai.types.create_embedding_response import Usage
 
@@ -155,6 +156,73 @@ async def mock_create_or_update_index(self, index, *args, **kwargs):
     assert updated_indexes[0].fields[0].name == "storageUrl"
 
 
+@pytest.mark.asyncio
+async def test_create_index_adds_vectorizer_to_existing_index(monkeypatch, search_info):
+    """Test that a vectorizer is added to an existing index when embeddings are configured."""
+    created_indexes = []
+    updated_indexes = []
+
+    async def mock_create_index(self, index):
+        created_indexes.append(index)  # pragma: no cover
+
+    async def mock_list_index_names(self):
+        yield "test"
+
+    async def mock_get_index(self, *args, **kwargs):
+        # Return an existing index with vector_search but no vectorizers
+        # Include embedding field to avoid triggering the embedding field addition code path
+        return SearchIndex(
+            name="test",
+            fields=[
+                SimpleField(
+                    name="storageUrl",
+                    type=SearchFieldDataType.String,
+                    filterable=True,
+                ),
+                SimpleField(
+                    name="embedding",
+                    type=SearchFieldDataType.Collection(SearchFieldDataType.Single),
+                    searchable=True,
+                    vector_search_dimensions=MOCK_EMBEDDING_DIMENSIONS,
+                ),
+            ],
+            vector_search=VectorSearch(vectorizers=[]),
+        )
+
+    async def mock_create_or_update_index(self, index, *args, **kwargs):
+        updated_indexes.append(index)
+
+    monkeypatch.setattr(SearchIndexClient, "create_index", mock_create_index)
+    monkeypatch.setattr(SearchIndexClient, "list_index_names", mock_list_index_names)
+    monkeypatch.setattr(SearchIndexClient, "get_index", mock_get_index)
+    monkeypatch.setattr(SearchIndexClient, "create_or_update_index", mock_create_or_update_index)
+
+    # Create a simple mock embeddings object with just the properties we need for index creation
+    class MockEmbeddings:
+        def __init__(self):
+            self.azure_endpoint = "https://test.openai.azure.com"
+            self.azure_deployment_name = "test-deployment"
+            self.open_ai_model_name = MOCK_EMBEDDING_MODEL_NAME
+            self.open_ai_dimensions = MOCK_EMBEDDING_DIMENSIONS
+
+    embeddings = MockEmbeddings()
+
+    manager = SearchManager(search_info, embeddings=embeddings, field_name_embedding="embedding")
+    await manager.create_index()
+
+    assert len(created_indexes) == 0, "It should not have created a new index"
+    assert len(updated_indexes) == 1, "It should have updated the existing index"
+    assert updated_indexes[0].vector_search.vectorizers is not None
+    assert len(updated_indexes[0].vector_search.vectorizers) == 1, "Should have added one vectorizer"
+    # The vectorizer name for updating existing indexes uses index_name
+    assert updated_indexes[0].vector_search.vectorizers[0].vectorizer_name == "test-vectorizer"
+    # Verify the vectorizer parameters
+    vectorizer = updated_indexes[0].vector_search.vectorizers[0]
+    assert vectorizer.parameters.resource_url == "https://test.openai.azure.com"
+    assert vectorizer.parameters.deployment_name == "test-deployment"
+    assert vectorizer.parameters.model_name == MOCK_EMBEDDING_MODEL_NAME
+
+
 @pytest.mark.asyncio
 async def test_create_index_acls(monkeypatch, search_info):
     indexes = []