diff --git a/src/crewai/utilities/embedding_configurator.py b/src/crewai/utilities/embedding_configurator.py index e523b60f0..71281d6d3 100644 --- a/src/crewai/utilities/embedding_configurator.py +++ b/src/crewai/utilities/embedding_configurator.py @@ -75,10 +75,36 @@ class EmbeddingConfigurator: @staticmethod def _configure_azure(config, model_name): + """ + Configure an Azure OpenAI embedding function. + + Args: + config: A dictionary containing Azure OpenAI configuration parameters. + Required parameters: + - api_key: Azure OpenAI API key + - api_base: Azure OpenAI API base URL + - api_version: Azure OpenAI API version + - deployment_id: Azure OpenAI deployment ID for the embedding model + model_name: The name of the embedding model + + Returns: + An OpenAIEmbeddingFunction configured for Azure OpenAI + + Raises: + ValueError: If required parameters are missing + """ from chromadb.utils.embedding_functions.openai_embedding_function import ( OpenAIEmbeddingFunction, ) - + + # Check if deployment_id is provided for Azure OpenAI + deployment_id = config.get("deployment_id") + if not deployment_id: + raise ValueError( + "Missing required parameter 'deployment_id' for Azure OpenAI embeddings. " + "Please provide a deployment_id in your Azure embedder configuration." + ) + return OpenAIEmbeddingFunction( api_key=config.get("api_key"), api_base=config.get("api_base"), @@ -87,7 +113,7 @@ class EmbeddingConfigurator: model_name=model_name, default_headers=config.get("default_headers"), dimensions=config.get("dimensions"), - deployment_id=config.get("deployment_id"), + deployment_id=deployment_id, organization_id=config.get("organization_id"), ) diff --git a/tests/memory/azure_memory_test.py b/tests/memory/azure_memory_test.py new file mode 100644 index 000000000..29a2941c8 --- /dev/null +++ b/tests/memory/azure_memory_test.py @@ -0,0 +1,53 @@ +import pytest +from unittest.mock import patch, MagicMock + +from crewai.utilities.embedding_configurator import EmbeddingConfigurator + + +def test_azure_embedder_missing_deployment_id(): + """Test that Azure embedder raises an error when deployment_id is missing""" + embedder_config = { + "provider": "azure", + "config": { + "model": "text-embedding-ada-002", + "api_key": "test-key", + "api_base": "https://test.openai.azure.com", + "api_version": "2023-05-15", + "api_type": "azure", + } + } + + configurator = EmbeddingConfigurator() + + with pytest.raises(ValueError) as excinfo: + configurator.configure_embedder(embedder_config) + + assert "Missing required parameter 'deployment_id'" in str(excinfo.value) + + +@patch("chromadb.utils.embedding_functions.openai_embedding_function.OpenAIEmbeddingFunction") +def test_azure_embedder_with_deployment_id(mock_openai_embedding): + """Test that Azure embedder works when deployment_id is provided""" + mock_instance = MagicMock() + mock_openai_embedding.return_value = mock_instance + + embedder_config = { + "provider": "azure", + "config": { + "model": "text-embedding-ada-002", + "api_key": "test-key", + "api_base": "https://test.openai.azure.com", + "api_version": "2023-05-15", + "api_type": "azure", + "deployment_id": "text-embedding-ada-002", + } + } + + configurator = EmbeddingConfigurator() + result = configurator.configure_embedder(embedder_config) + + assert result == mock_instance + mock_openai_embedding.assert_called_once() + # Verify deployment_id was passed correctly + call_kwargs = mock_openai_embedding.call_args.kwargs + assert call_kwargs["deployment_id"] == "text-embedding-ada-002"