From cf4d5e3cdd4c5d6f3cdf19b0d3aec358cd383b01 Mon Sep 17 00:00:00 2001
From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
Date: Mon, 24 Mar 2025 07:11:11 +0000
Subject: [PATCH] Improve OpenRouter embedding provider with type hints, error
 handling, and tests

Co-Authored-By: Joe Moura <joao@crewai.com>
---
 .../utilities/embedding_configurator.py       | 23 ++++++--
 .../utilities/test_embedding_configurator.py  | 54 +++++++++++++++++++
 2 files changed, 74 insertions(+), 3 deletions(-)

diff --git a/src/crewai/utilities/embedding_configurator.py b/src/crewai/utilities/embedding_configurator.py
index 5f27211ac..81631d926 100644
--- a/src/crewai/utilities/embedding_configurator.py
+++ b/src/crewai/utilities/embedding_configurator.py
@@ -212,13 +212,30 @@ class EmbeddingConfigurator:
         return WatsonEmbeddingFunction()
 
     @staticmethod
-    def _configure_openrouter(config, model_name):
+    def _configure_openrouter(config: Dict[str, Any], model_name: str) -> EmbeddingFunction:
+        """
+        Configure OpenRouter embedding provider.
+        
+        Args:
+            config (Dict[str, Any]): Configuration dictionary containing the API key and optional settings.
+            model_name (str): Name of the embedding model to use.
+            
+        Returns:
+            OpenAIEmbeddingFunction: Configured OpenRouter embedding function.
+            
+        Raises:
+            ValueError: If the API key is not provided in the config or environment.
+        """
         from chromadb.utils.embedding_functions.openai_embedding_function import (
             OpenAIEmbeddingFunction,
         )
-
+        
+        api_key = config.get("api_key") or os.getenv("OPENROUTER_API_KEY")
+        if not api_key:
+            raise ValueError("OpenRouter API key must be provided either in config or OPENROUTER_API_KEY environment variable")
+            
         return OpenAIEmbeddingFunction(
-            api_key=config.get("api_key") or os.getenv("OPENROUTER_API_KEY"),
+            api_key=api_key,
             api_base=config.get("api_base", "https://openrouter.ai/api/v1"),
             model_name=model_name,
         )
diff --git a/tests/utilities/test_embedding_configurator.py b/tests/utilities/test_embedding_configurator.py
index fb74196d9..17f4cc7cd 100644
--- a/tests/utilities/test_embedding_configurator.py
+++ b/tests/utilities/test_embedding_configurator.py
@@ -34,3 +34,57 @@ def test_openrouter_embedder_configuration():
             api_base="https://openrouter.ai/api/v1",
             model_name="test-model",
         )
+
+
+def test_openrouter_embedder_configuration_with_env_var():
+    # Setup
+    configurator = EmbeddingConfigurator()
+    mock_openai_embedding = MagicMock()
+    
+    # Test with API key from environment variable
+    with patch.dict(os.environ, {"OPENROUTER_API_KEY": "env-key"}), \
+        patch(
+            "chromadb.utils.embedding_functions.openai_embedding_function.OpenAIEmbeddingFunction",
+            return_value=mock_openai_embedding,
+        ) as mock_embedder:
+        # Config without API key
+        embedder_config = {
+            "provider": "openrouter",
+            "config": {
+                "model": "test-model",
+            },
+        }
+        
+        # Execute
+        result = configurator.configure_embedder(embedder_config)
+        
+        # Verify
+        assert result == mock_openai_embedding
+        mock_embedder.assert_called_once_with(
+            api_key="env-key",
+            api_base="https://openrouter.ai/api/v1",
+            model_name="test-model",
+        )
+
+
+def test_openrouter_embedder_configuration_missing_api_key():
+    # Setup
+    configurator = EmbeddingConfigurator()
+    
+    # Test without API key
+    with patch.dict(os.environ, {}, clear=True), \
+        patch(
+            "chromadb.utils.embedding_functions.openai_embedding_function.OpenAIEmbeddingFunction",
+            side_effect=Exception("Should not be called"),
+        ):
+        # Config without API key
+        embedder_config = {
+            "provider": "openrouter",
+            "config": {
+                "model": "test-model",
+            },
+        }
+        
+        # Verify error is raised
+        with pytest.raises(ValueError, match="OpenRouter API key must be provided"):
+            configurator.configure_embedder(embedder_config)