Improve OpenRouter embedding provider with type hints, error handling, and tests

Co-Authored-By: Joe Moura <joao@crewai.com>
This commit is contained in:
Devin AI
2025-03-24 07:11:11 +00:00
parent 5470c56dc6
commit cf4d5e3cdd
2 changed files with 74 additions and 3 deletions

View File

@@ -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)