mirror of
https://github.com/crewAIInc/crewAI.git
synced 2026-01-09 08:08:32 +00:00
Improve OpenRouter embedding provider with type hints, error handling, and tests
Co-Authored-By: Joe Moura <joao@crewai.com>
This commit is contained in:
@@ -212,13 +212,30 @@ class EmbeddingConfigurator:
|
|||||||
return WatsonEmbeddingFunction()
|
return WatsonEmbeddingFunction()
|
||||||
|
|
||||||
@staticmethod
|
@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 (
|
from chromadb.utils.embedding_functions.openai_embedding_function import (
|
||||||
OpenAIEmbeddingFunction,
|
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(
|
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"),
|
api_base=config.get("api_base", "https://openrouter.ai/api/v1"),
|
||||||
model_name=model_name,
|
model_name=model_name,
|
||||||
)
|
)
|
||||||
|
|||||||
@@ -34,3 +34,57 @@ def test_openrouter_embedder_configuration():
|
|||||||
api_base="https://openrouter.ai/api/v1",
|
api_base="https://openrouter.ai/api/v1",
|
||||||
model_name="test-model",
|
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)
|
||||||
|
|||||||
Reference in New Issue
Block a user