mirror of
https://github.com/crewAIInc/crewAI.git
synced 2026-01-11 00:58:30 +00:00
Fix #2451: Add support for OpenRouter as an embedding provider
Co-Authored-By: Joe Moura <joao@crewai.com>
This commit is contained in:
@@ -19,6 +19,7 @@ class EmbeddingConfigurator:
|
|||||||
"huggingface": self._configure_huggingface,
|
"huggingface": self._configure_huggingface,
|
||||||
"watson": self._configure_watson,
|
"watson": self._configure_watson,
|
||||||
"custom": self._configure_custom,
|
"custom": self._configure_custom,
|
||||||
|
"openrouter": self._configure_openrouter,
|
||||||
}
|
}
|
||||||
|
|
||||||
def configure_embedder(
|
def configure_embedder(
|
||||||
@@ -210,6 +211,18 @@ class EmbeddingConfigurator:
|
|||||||
|
|
||||||
return WatsonEmbeddingFunction()
|
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
|
@staticmethod
|
||||||
def _configure_custom(config):
|
def _configure_custom(config):
|
||||||
custom_embedder = config.get("embedder")
|
custom_embedder = config.get("embedder")
|
||||||
|
|||||||
35
tests/utilities/test_embedding_configurator.py
Normal file
35
tests/utilities/test_embedding_configurator.py
Normal 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",
|
||||||
|
)
|
||||||
Reference in New Issue
Block a user