Fix #2451: Add support for OpenRouter as an embedding provider

Co-Authored-By: Joe Moura <joao@crewai.com>
This commit is contained in:
Devin AI
2025-03-24 07:04:47 +00:00
parent ed1f009c64
commit de200b9404
2 changed files with 48 additions and 0 deletions

View File

@@ -19,6 +19,7 @@ class EmbeddingConfigurator:
"huggingface": self._configure_huggingface,
"watson": self._configure_watson,
"custom": self._configure_custom,
"openrouter": self._configure_openrouter,
}
def configure_embedder(
@@ -210,6 +211,18 @@ class EmbeddingConfigurator:
return WatsonEmbeddingFunction()
@staticmethod
def _configure_openrouter(config, model_name):
from chromadb.utils.embedding_functions.openai_embedding_function import (
OpenAIEmbeddingFunction,
)
return OpenAIEmbeddingFunction(
api_key=config.get("api_key") or os.getenv("OPENROUTER_API_KEY"),
api_base=config.get("api_base", "https://openrouter.ai/api/v1"),
model_name=model_name,
)
@staticmethod
def _configure_custom(config):
custom_embedder = config.get("embedder")

View File

@@ -0,0 +1,35 @@
import os
import pytest
from unittest.mock import patch, MagicMock
from crewai.utilities.embedding_configurator import EmbeddingConfigurator
def test_openrouter_embedder_configuration():
# Setup
configurator = EmbeddingConfigurator()
mock_openai_embedding = MagicMock()
with patch(
"chromadb.utils.embedding_functions.openai_embedding_function.OpenAIEmbeddingFunction",
return_value=mock_openai_embedding,
) as mock_embedder:
# Test with provided config
embedder_config = {
"provider": "openrouter",
"config": {
"api_key": "test-key",
"model": "test-model",
},
}
# Execute
result = configurator.configure_embedder(embedder_config)
# Verify
assert result == mock_openai_embedding
mock_embedder.assert_called_once_with(
api_key="test-key",
api_base="https://openrouter.ai/api/v1",
model_name="test-model",
)