mirror of
https://github.com/crewAIInc/crewAI.git
synced 2026-10-12 02:12:31 +00:00
feat(embeddings): add openrouter as a supported embedding provider (#7127)
* feat(embeddings): add openrouter provider type definitions * feat(embeddings): implement OpenRouterProvider * feat(embeddings): export openrouter provider package symbols * feat(embeddings): register openrouter in allowed embedding providers * feat(embeddings): register openrouter provider and overloads in factory * feat(tools): add openrouter embedding service support * test(embeddings): add comprehensive openrouter provider and factory tests * test(embeddings): add openrouter build test in embedding factory * test(embeddings): add openrouter model key alias and env tests * test(tools): add openrouter tests for embedding service * docs: add openrouter embedder configuration example * docs(ar): sync openrouter embedder translation * docs(ko): sync openrouter embedder translation * docs(pt-BR): sync openrouter embedder translation * feat(embeddings): allow model alias and None fields in OpenRouterProviderConfig * fix(tools): resolve EMBEDDINGS_OPENROUTER_API_KEY before OPENROUTER_API_KEY * test(tools): add regression tests for openrouter env var precedence and fallback * feat(embeddings): drop organization_id and resolve api_key via OPENROUTER_API_KEY only * fix(tools): use OPENROUTER_API_KEY in embedding service and update tests * docs: switch openrouter knowledge example to model_name and document OPENROUTER_API_KEY * fix(tools): default openrouter model to namespaced openai/text-embedding-3-small --------- Co-authored-by: Vidit Ostwal <[email protected]>
This commit is contained in:
co-authored by
Vidit Ostwal
parent
4bfdb0df67
commit
c860613c7a
@@ -636,6 +636,21 @@ agent = Agent(
|
||||
}
|
||||
}
|
||||
)
|
||||
|
||||
# Option 4: Use OpenRouter embeddings (supports models across providers)
|
||||
# Set OPENROUTER_API_KEY environment variable if api_key is omitted in config
|
||||
crew = Crew(
|
||||
agents=[agent],
|
||||
tasks=[...],
|
||||
knowledge_sources=[knowledge_source],
|
||||
embedder={
|
||||
"provider": "openrouter",
|
||||
"config": {
|
||||
"model_name": "openai/text-embedding-3-small",
|
||||
"api_key": "your-openrouter-api-key" # Optional if OPENROUTER_API_KEY is set
|
||||
}
|
||||
}
|
||||
)
|
||||
```
|
||||
|
||||
#### إعداد تضمينات Azure OpenAI
|
||||
|
||||
@@ -636,6 +636,21 @@ agent = Agent(
|
||||
}
|
||||
}
|
||||
)
|
||||
|
||||
# Option 4: Use OpenRouter embeddings (supports models across providers)
|
||||
# Set OPENROUTER_API_KEY environment variable if api_key is omitted in config
|
||||
crew = Crew(
|
||||
agents=[agent],
|
||||
tasks=[...],
|
||||
knowledge_sources=[knowledge_source],
|
||||
embedder={
|
||||
"provider": "openrouter",
|
||||
"config": {
|
||||
"model_name": "openai/text-embedding-3-small",
|
||||
"api_key": "your-openrouter-api-key" # Optional if OPENROUTER_API_KEY is set
|
||||
}
|
||||
}
|
||||
)
|
||||
```
|
||||
|
||||
#### Configuring Azure OpenAI Embeddings
|
||||
|
||||
@@ -603,6 +603,21 @@ agent = Agent(
|
||||
}
|
||||
}
|
||||
)
|
||||
|
||||
# Option 4: Use OpenRouter embeddings (supports models across providers)
|
||||
# Set OPENROUTER_API_KEY environment variable if api_key is omitted in config
|
||||
crew = Crew(
|
||||
agents=[agent],
|
||||
tasks=[...],
|
||||
knowledge_sources=[knowledge_source],
|
||||
embedder={
|
||||
"provider": "openrouter",
|
||||
"config": {
|
||||
"model_name": "openai/text-embedding-3-small",
|
||||
"api_key": "your-openrouter-api-key" # Optional if OPENROUTER_API_KEY is set
|
||||
}
|
||||
}
|
||||
)
|
||||
```
|
||||
|
||||
#### Azure OpenAI 임베딩 구성
|
||||
|
||||
@@ -601,6 +601,21 @@ agent = Agent(
|
||||
}
|
||||
}
|
||||
)
|
||||
|
||||
# Option 4: Use OpenRouter embeddings (supports models across providers)
|
||||
# Set OPENROUTER_API_KEY environment variable if api_key is omitted in config
|
||||
crew = Crew(
|
||||
agents=[agent],
|
||||
tasks=[...],
|
||||
knowledge_sources=[knowledge_source],
|
||||
embedder={
|
||||
"provider": "openrouter",
|
||||
"config": {
|
||||
"model_name": "openai/text-embedding-3-small",
|
||||
"api_key": "your-openrouter-api-key" # Optional if OPENROUTER_API_KEY is set
|
||||
}
|
||||
}
|
||||
)
|
||||
```
|
||||
|
||||
#### Configurando Embeddings do Azure OpenAI
|
||||
|
||||
@@ -62,7 +62,7 @@ class EmbeddingService:
|
||||
def __init__(
|
||||
self,
|
||||
provider: str = "openai",
|
||||
model: str = "text-embedding-3-small",
|
||||
model: str | None = None,
|
||||
api_key: str | None = None,
|
||||
**kwargs: Any,
|
||||
):
|
||||
@@ -71,10 +71,16 @@ class EmbeddingService:
|
||||
|
||||
Args:
|
||||
provider: The embedding provider to use
|
||||
model: The model name
|
||||
model: The model name (if not provided, uses provider default)
|
||||
api_key: API key (if not provided, will look for environment variables)
|
||||
**kwargs: Additional configuration options
|
||||
"""
|
||||
if model is None:
|
||||
if provider == "openrouter":
|
||||
model = "openai/text-embedding-3-small"
|
||||
else:
|
||||
model = "text-embedding-3-small"
|
||||
|
||||
self.config = EmbeddingConfig(
|
||||
provider=provider,
|
||||
model=model,
|
||||
@@ -98,6 +104,7 @@ class EmbeddingService:
|
||||
"jina": "JINA_API_KEY",
|
||||
"ollama": None, # Ollama typically runs locally without API key
|
||||
"openai": "OPENAI_API_KEY",
|
||||
"openrouter": "OPENROUTER_API_KEY",
|
||||
"roboflow": "ROBOFLOW_API_KEY",
|
||||
"voyageai": "VOYAGE_API_KEY",
|
||||
"watsonx": "WATSONX_API_KEY",
|
||||
@@ -144,6 +151,12 @@ class EmbeddingService:
|
||||
"model_name": self.config.model,
|
||||
**self.config.extra_config,
|
||||
}
|
||||
elif self.config.provider == "openrouter":
|
||||
base_config["config"] = {
|
||||
"api_key": self.config.api_key,
|
||||
"model_name": self.config.model,
|
||||
**self.config.extra_config,
|
||||
}
|
||||
elif self.config.provider == "azure":
|
||||
base_config["config"] = {
|
||||
"api_key": self.config.api_key,
|
||||
@@ -371,6 +384,7 @@ class EmbeddingService:
|
||||
"onnx",
|
||||
"openai",
|
||||
"openclip",
|
||||
"openrouter",
|
||||
"roboflow",
|
||||
"sentence-transformer",
|
||||
"text2vec",
|
||||
@@ -388,6 +402,16 @@ class EmbeddingService:
|
||||
"""Create an OpenAI embedding service."""
|
||||
return cls(provider="openai", model=model, api_key=api_key, **kwargs)
|
||||
|
||||
@classmethod
|
||||
def create_openrouter_service(
|
||||
cls,
|
||||
model: str = "openai/text-embedding-3-small",
|
||||
api_key: str | None = None,
|
||||
**kwargs: Any,
|
||||
) -> EmbeddingService:
|
||||
"""Create an OpenRouter embedding service."""
|
||||
return cls(provider="openrouter", model=model, api_key=api_key, **kwargs)
|
||||
|
||||
@classmethod
|
||||
def create_voyage_service(
|
||||
cls, model: str = "voyage-2", api_key: str | None = None, **kwargs: Any
|
||||
|
||||
@@ -52,7 +52,7 @@ class TestEmbeddingService:
|
||||
"""Test listing supported providers."""
|
||||
providers = EmbeddingService.list_supported_providers()
|
||||
expected_providers = [
|
||||
"openai", "azure", "voyageai", "cohere", "google-generativeai",
|
||||
"openai", "openrouter", "azure", "voyageai", "cohere", "google-generativeai",
|
||||
"amazon-bedrock", "huggingface", "jina", "ollama", "sentence-transformer",
|
||||
"instructor", "watsonx", "custom"
|
||||
]
|
||||
@@ -69,6 +69,10 @@ class TestEmbeddingService:
|
||||
api_key = service._get_default_api_key("openai")
|
||||
assert api_key == "test-openai-key"
|
||||
|
||||
with patch.dict(os.environ, {"OPENROUTER_API_KEY": "test-openrouter-key"}, clear=True):
|
||||
api_key = service._get_default_api_key("openrouter")
|
||||
assert api_key == "test-openrouter-key"
|
||||
|
||||
with patch.dict(os.environ, {}, clear=True):
|
||||
api_key = service._get_default_api_key("openai")
|
||||
assert api_key is None
|
||||
@@ -326,3 +330,75 @@ class TestProviderConfigurations:
|
||||
assert call_args["provider"] == "google-generativeai"
|
||||
assert call_args["config"]["api_key"] == "test-key"
|
||||
assert call_args["config"]["model_name"] == "models/embedding-001"
|
||||
|
||||
@patch("crewai.rag.embeddings.factory.build_embedder")
|
||||
def test_openrouter_config(self, mock_build_embedder):
|
||||
"""Test OpenRouter configuration mapping."""
|
||||
mock_build_embedder.return_value = Mock()
|
||||
|
||||
service = EmbeddingService(
|
||||
provider="openrouter",
|
||||
model="openai/text-embedding-3-large",
|
||||
api_key="test-key",
|
||||
extra_config={"dimensions": 1536},
|
||||
)
|
||||
|
||||
call_args = mock_build_embedder.call_args[0][0]
|
||||
assert call_args["provider"] == "openrouter"
|
||||
assert call_args["config"]["api_key"] == "test-key"
|
||||
assert call_args["config"]["model_name"] == "openai/text-embedding-3-large"
|
||||
assert call_args["config"]["dimensions"] == 1536
|
||||
|
||||
@patch("crewai.rag.embeddings.factory.build_embedder")
|
||||
def test_create_openrouter_service(self, mock_build_embedder):
|
||||
"""Test create_openrouter_service helper method."""
|
||||
mock_build_embedder.return_value = Mock()
|
||||
|
||||
service = EmbeddingService.create_openrouter_service(
|
||||
model="openai/text-embedding-3-small",
|
||||
api_key="test-key",
|
||||
)
|
||||
|
||||
assert service.config.provider == "openrouter"
|
||||
assert service.config.model == "openai/text-embedding-3-small"
|
||||
assert service.config.api_key == "test-key"
|
||||
|
||||
@patch("crewai.rag.embeddings.factory.build_embedder")
|
||||
def test_openrouter_initialization_with_env_key_only(
|
||||
self, mock_build_embedder
|
||||
):
|
||||
"""Test OpenRouter initialization when OPENROUTER_API_KEY is set in env."""
|
||||
mock_build_embedder.return_value = Mock()
|
||||
|
||||
with patch.dict(
|
||||
os.environ,
|
||||
{"OPENROUTER_API_KEY": "test-openrouter-key"},
|
||||
clear=True,
|
||||
):
|
||||
service = EmbeddingService(
|
||||
provider="openrouter",
|
||||
model="openai/text-embedding-3-small",
|
||||
)
|
||||
|
||||
assert service.config.api_key == "test-openrouter-key"
|
||||
mock_build_embedder.assert_called_once()
|
||||
call_args = mock_build_embedder.call_args[0][0]
|
||||
assert call_args["provider"] == "openrouter"
|
||||
assert call_args["config"]["api_key"] == "test-openrouter-key"
|
||||
|
||||
@patch("crewai.rag.embeddings.factory.build_embedder")
|
||||
def test_openrouter_default_model_without_model_arg(self, mock_build_embedder):
|
||||
"""Test constructing EmbeddingService(provider='openrouter') without a model uses namespaced default."""
|
||||
mock_build_embedder.return_value = Mock()
|
||||
|
||||
service = EmbeddingService(
|
||||
provider="openrouter",
|
||||
api_key="test-key",
|
||||
)
|
||||
|
||||
assert service.config.provider == "openrouter"
|
||||
assert service.config.model == "openai/text-embedding-3-small"
|
||||
mock_build_embedder.assert_called_once()
|
||||
call_args = mock_build_embedder.call_args[0][0]
|
||||
assert call_args["provider"] == "openrouter"
|
||||
assert call_args["config"]["model_name"] == "openai/text-embedding-3-small"
|
||||
|
||||
@@ -74,6 +74,7 @@ if TYPE_CHECKING:
|
||||
from crewai.rag.embeddings.providers.onnx.types import ONNXProviderSpec
|
||||
from crewai.rag.embeddings.providers.openai.types import OpenAIProviderSpec
|
||||
from crewai.rag.embeddings.providers.openclip.types import OpenCLIPProviderSpec
|
||||
from crewai.rag.embeddings.providers.openrouter.types import OpenRouterProviderSpec
|
||||
from crewai.rag.embeddings.providers.roboflow.types import RoboflowProviderSpec
|
||||
from crewai.rag.embeddings.providers.sentence_transformer.types import (
|
||||
SentenceTransformerProviderSpec,
|
||||
@@ -102,6 +103,7 @@ PROVIDER_PATHS = {
|
||||
"onnx": "crewai.rag.embeddings.providers.onnx.onnx_provider.ONNXProvider",
|
||||
"openai": "crewai.rag.embeddings.providers.openai.openai_provider.OpenAIProvider",
|
||||
"openclip": "crewai.rag.embeddings.providers.openclip.openclip_provider.OpenCLIPProvider",
|
||||
"openrouter": "crewai.rag.embeddings.providers.openrouter.openrouter_provider.OpenRouterProvider",
|
||||
"roboflow": "crewai.rag.embeddings.providers.roboflow.roboflow_provider.RoboflowProvider",
|
||||
"sentence-transformer": "crewai.rag.embeddings.providers.sentence_transformer.sentence_transformer_provider.SentenceTransformerProvider",
|
||||
"text2vec": "crewai.rag.embeddings.providers.text2vec.text2vec_provider.Text2VecProvider",
|
||||
@@ -206,6 +208,12 @@ def build_embedder_from_dict(
|
||||
) -> OpenCLIPEmbeddingFunction: ...
|
||||
|
||||
|
||||
@overload
|
||||
def build_embedder_from_dict(
|
||||
spec: OpenRouterProviderSpec,
|
||||
) -> OpenAIEmbeddingFunction: ...
|
||||
|
||||
|
||||
@overload
|
||||
def build_embedder_from_dict(
|
||||
spec: Text2VecProviderSpec,
|
||||
@@ -337,6 +345,10 @@ def build_embedder(spec: RoboflowProviderSpec) -> RoboflowEmbeddingFunction: ...
|
||||
def build_embedder(spec: OpenCLIPProviderSpec) -> OpenCLIPEmbeddingFunction: ...
|
||||
|
||||
|
||||
@overload
|
||||
def build_embedder(spec: OpenRouterProviderSpec) -> OpenAIEmbeddingFunction: ...
|
||||
|
||||
|
||||
@overload
|
||||
def build_embedder(spec: Text2VecProviderSpec) -> Text2VecEmbeddingFunction: ...
|
||||
|
||||
|
||||
@@ -0,0 +1,16 @@
|
||||
"""OpenRouter embedding providers."""
|
||||
|
||||
from crewai.rag.embeddings.providers.openrouter.openrouter_provider import (
|
||||
OpenRouterProvider,
|
||||
)
|
||||
from crewai.rag.embeddings.providers.openrouter.types import (
|
||||
OpenRouterProviderConfig,
|
||||
OpenRouterProviderSpec,
|
||||
)
|
||||
|
||||
|
||||
__all__ = [
|
||||
"OpenRouterProvider",
|
||||
"OpenRouterProviderConfig",
|
||||
"OpenRouterProviderSpec",
|
||||
]
|
||||
@@ -0,0 +1,59 @@
|
||||
"""OpenRouter embeddings provider."""
|
||||
|
||||
from typing import Any
|
||||
|
||||
from chromadb.utils.embedding_functions.openai_embedding_function import (
|
||||
OpenAIEmbeddingFunction,
|
||||
)
|
||||
from pydantic import AliasChoices, Field, model_validator
|
||||
|
||||
from crewai.rag.core.base_embeddings_provider import BaseEmbeddingsProvider
|
||||
|
||||
|
||||
class OpenRouterProvider(BaseEmbeddingsProvider[OpenAIEmbeddingFunction]):
|
||||
"""OpenRouter embeddings provider."""
|
||||
|
||||
@model_validator(mode="before")
|
||||
@classmethod
|
||||
def _normalize_model_alias(cls, data: Any) -> Any:
|
||||
if isinstance(data, dict) and "model" in data and "model_name" not in data:
|
||||
data = data.copy()
|
||||
data["model_name"] = data["model"]
|
||||
return data
|
||||
|
||||
embedding_callable: type[OpenAIEmbeddingFunction] = Field(
|
||||
default=OpenAIEmbeddingFunction,
|
||||
description="OpenAI-compatible embedding function class",
|
||||
)
|
||||
api_key: str = Field(
|
||||
description="OpenRouter API key",
|
||||
validation_alias=AliasChoices(
|
||||
"OPENROUTER_API_KEY",
|
||||
),
|
||||
)
|
||||
model_name: str = Field(
|
||||
default="openai/text-embedding-3-small",
|
||||
description="Model name to use for embeddings",
|
||||
validation_alias=AliasChoices(
|
||||
"model_name",
|
||||
),
|
||||
)
|
||||
api_base: str = Field(
|
||||
default="https://openrouter.ai/api/v1",
|
||||
description="Base URL for OpenRouter API requests",
|
||||
validation_alias=AliasChoices(
|
||||
"OPENROUTER_API_BASE",
|
||||
"api_base",
|
||||
),
|
||||
)
|
||||
default_headers: dict[str, Any] | None = Field(
|
||||
default=None, description="Default headers for API requests"
|
||||
)
|
||||
dimensions: int | None = Field(
|
||||
default=None,
|
||||
description="Embedding dimensions",
|
||||
validation_alias=AliasChoices(
|
||||
"OPENROUTER_DIMENSIONS",
|
||||
"dimensions",
|
||||
),
|
||||
)
|
||||
@@ -0,0 +1,23 @@
|
||||
"""Type definitions for OpenRouter embedding providers."""
|
||||
|
||||
from typing import Annotated, Any, Literal
|
||||
|
||||
from typing_extensions import Required, TypedDict
|
||||
|
||||
|
||||
class OpenRouterProviderConfig(TypedDict, total=False):
|
||||
"""Configuration for OpenRouter provider."""
|
||||
|
||||
api_key: str
|
||||
model: str
|
||||
model_name: Annotated[str, "openai/text-embedding-3-small"]
|
||||
api_base: str
|
||||
default_headers: dict[str, Any] | None
|
||||
dimensions: int | None
|
||||
|
||||
|
||||
class OpenRouterProviderSpec(TypedDict, total=False):
|
||||
"""OpenRouter provider specification."""
|
||||
|
||||
provider: Required[Literal["openrouter"]]
|
||||
config: OpenRouterProviderConfig
|
||||
@@ -21,6 +21,7 @@ from crewai.rag.embeddings.providers.ollama.types import OllamaProviderSpec
|
||||
from crewai.rag.embeddings.providers.onnx.types import ONNXProviderSpec
|
||||
from crewai.rag.embeddings.providers.openai.types import OpenAIProviderSpec
|
||||
from crewai.rag.embeddings.providers.openclip.types import OpenCLIPProviderSpec
|
||||
from crewai.rag.embeddings.providers.openrouter.types import OpenRouterProviderSpec
|
||||
from crewai.rag.embeddings.providers.roboflow.types import RoboflowProviderSpec
|
||||
from crewai.rag.embeddings.providers.sentence_transformer.types import (
|
||||
SentenceTransformerProviderSpec,
|
||||
@@ -42,6 +43,7 @@ ProviderSpec: TypeAlias = (
|
||||
| ONNXProviderSpec
|
||||
| OpenAIProviderSpec
|
||||
| OpenCLIPProviderSpec
|
||||
| OpenRouterProviderSpec
|
||||
| RoboflowProviderSpec
|
||||
| SentenceTransformerProviderSpec
|
||||
| Text2VecProviderSpec
|
||||
@@ -64,6 +66,7 @@ AllowedEmbeddingProviders = Literal[
|
||||
"onnx",
|
||||
"openai",
|
||||
"openclip",
|
||||
"openrouter",
|
||||
"roboflow",
|
||||
"sentence-transformer",
|
||||
"text2vec",
|
||||
|
||||
@@ -15,6 +15,7 @@ from crewai.rag.embeddings.providers.sentence_transformer.sentence_transformer_p
|
||||
)
|
||||
from crewai.rag.embeddings.providers.instructor.instructor_provider import InstructorProvider
|
||||
from crewai.rag.embeddings.providers.openclip.openclip_provider import OpenCLIPProvider
|
||||
from crewai.rag.embeddings.providers.openrouter.openrouter_provider import OpenRouterProvider
|
||||
|
||||
|
||||
class TestGoogleProviderAlias:
|
||||
@@ -56,6 +57,25 @@ class TestModelKeyBackwardCompatibility:
|
||||
|
||||
assert provider.model_name == "text-embedding-3-large"
|
||||
|
||||
def test_openrouter_provider_accepts_model_key(self):
|
||||
"""Test OpenRouter provider accepts 'model' as alias for 'model_name'."""
|
||||
provider = OpenRouterProvider(
|
||||
api_key="test-key",
|
||||
model="openai/text-embedding-3-large",
|
||||
)
|
||||
assert provider.model_name == "openai/text-embedding-3-large"
|
||||
|
||||
def test_openrouter_provider_ignores_chat_model_env(self, monkeypatch):
|
||||
"""Test OpenRouter embeddings don't inherit unrelated chat model env vars."""
|
||||
monkeypatch.setenv("OPENAI_MODEL_NAME", "gpt-5.5")
|
||||
monkeypatch.setenv("MODEL", "gpt-5.5")
|
||||
monkeypatch.delenv("EMBEDDINGS_OPENROUTER_MODEL_NAME", raising=False)
|
||||
monkeypatch.delenv("OPENROUTER_MODEL_NAME", raising=False)
|
||||
|
||||
provider = OpenRouterProvider(api_key="test-key")
|
||||
|
||||
assert provider.model_name == "openai/text-embedding-3-small"
|
||||
|
||||
def test_azure_provider_ignores_openai_chat_model_env(self, monkeypatch):
|
||||
"""Test Azure embeddings don't inherit the OpenAI chat model env var."""
|
||||
monkeypatch.setenv("OPENAI_MODEL_NAME", "gpt-5.5")
|
||||
|
||||
@@ -40,6 +40,36 @@ class TestEmbeddingFactory:
|
||||
assert call_kwargs["api_key"] == "test-key"
|
||||
assert call_kwargs["model_name"] == "text-embedding-3-small"
|
||||
|
||||
@patch("crewai.rag.embeddings.factory.import_and_validate_definition")
|
||||
def test_build_embedder_openrouter(self, mock_import):
|
||||
"""Test building OpenRouter embedder."""
|
||||
mock_provider_class = MagicMock()
|
||||
mock_provider_instance = MagicMock()
|
||||
mock_embedding_function = MagicMock()
|
||||
|
||||
mock_import.return_value = mock_provider_class
|
||||
mock_provider_class.return_value = mock_provider_instance
|
||||
mock_provider_instance.embedding_callable.return_value = mock_embedding_function
|
||||
|
||||
config = {
|
||||
"provider": "openrouter",
|
||||
"config": {
|
||||
"api_key": "test-key",
|
||||
"model_name": "openai/text-embedding-3-small",
|
||||
},
|
||||
}
|
||||
|
||||
build_embedder(config)
|
||||
|
||||
mock_import.assert_called_once_with(
|
||||
"crewai.rag.embeddings.providers.openrouter.openrouter_provider.OpenRouterProvider"
|
||||
)
|
||||
mock_provider_class.assert_called_once()
|
||||
|
||||
call_kwargs = mock_provider_class.call_args.kwargs
|
||||
assert call_kwargs["api_key"] == "test-key"
|
||||
assert call_kwargs["model_name"] == "openai/text-embedding-3-small"
|
||||
|
||||
@patch("crewai.rag.embeddings.factory.import_and_validate_definition")
|
||||
def test_build_embedder_azure(self, mock_import):
|
||||
"""Test building Azure embedder."""
|
||||
|
||||
@@ -0,0 +1,157 @@
|
||||
"""Test OpenRouter embedder configuration with factory."""
|
||||
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
from pydantic import ValidationError
|
||||
|
||||
from crewai.rag.embeddings.factory import build_embedder
|
||||
from crewai.rag.embeddings.providers.openrouter.openrouter_provider import (
|
||||
OpenRouterProvider,
|
||||
)
|
||||
|
||||
|
||||
class TestOpenRouterEmbedderFactory:
|
||||
"""Test OpenRouter embedder configuration with factory function."""
|
||||
|
||||
@patch("crewai.rag.embeddings.factory.import_and_validate_definition")
|
||||
def test_openrouter_with_nested_config(self, mock_import):
|
||||
"""Test OpenRouter configuration with nested config key."""
|
||||
mock_provider_class = MagicMock()
|
||||
mock_provider_instance = MagicMock()
|
||||
mock_embedding_function = MagicMock()
|
||||
|
||||
mock_import.return_value = mock_provider_class
|
||||
mock_provider_class.return_value = mock_provider_instance
|
||||
mock_provider_instance.embedding_callable.return_value = mock_embedding_function
|
||||
|
||||
embedder_config = {
|
||||
"provider": "openrouter",
|
||||
"config": {
|
||||
"api_key": "test-openrouter-key",
|
||||
"model_name": "openai/text-embedding-3-large",
|
||||
"api_base": "https://openrouter.ai/api/v1",
|
||||
"dimensions": 1536,
|
||||
},
|
||||
}
|
||||
|
||||
result = build_embedder(embedder_config)
|
||||
|
||||
mock_import.assert_called_once_with(
|
||||
"crewai.rag.embeddings.providers.openrouter.openrouter_provider.OpenRouterProvider"
|
||||
)
|
||||
|
||||
call_kwargs = mock_provider_class.call_args.kwargs
|
||||
assert call_kwargs["api_key"] == "test-openrouter-key"
|
||||
assert call_kwargs["model_name"] == "openai/text-embedding-3-large"
|
||||
assert call_kwargs["api_base"] == "https://openrouter.ai/api/v1"
|
||||
assert call_kwargs["dimensions"] == 1536
|
||||
|
||||
assert result == mock_embedding_function
|
||||
|
||||
@patch("crewai.rag.embeddings.factory.import_and_validate_definition")
|
||||
def test_openrouter_with_model_alias(self, mock_import):
|
||||
"""Test OpenRouter configuration with 'model' alias instead of 'model_name'."""
|
||||
mock_provider_class = MagicMock()
|
||||
mock_provider_instance = MagicMock()
|
||||
mock_embedding_function = MagicMock()
|
||||
|
||||
mock_import.return_value = mock_provider_class
|
||||
mock_provider_class.return_value = mock_provider_instance
|
||||
mock_provider_instance.embedding_callable.return_value = mock_embedding_function
|
||||
|
||||
embedder_config = {
|
||||
"provider": "openrouter",
|
||||
"config": {
|
||||
"api_key": "test-openrouter-key",
|
||||
"model": "cohere/embed-multilingual-v3.0",
|
||||
},
|
||||
}
|
||||
|
||||
result = build_embedder(embedder_config)
|
||||
|
||||
mock_import.assert_called_once_with(
|
||||
"crewai.rag.embeddings.providers.openrouter.openrouter_provider.OpenRouterProvider"
|
||||
)
|
||||
|
||||
call_kwargs = mock_provider_class.call_args.kwargs
|
||||
assert call_kwargs["api_key"] == "test-openrouter-key"
|
||||
assert call_kwargs["model"] == "cohere/embed-multilingual-v3.0"
|
||||
|
||||
assert result == mock_embedding_function
|
||||
|
||||
@patch("crewai.rag.embeddings.factory.import_and_validate_definition")
|
||||
def test_openrouter_import_error(self, mock_import):
|
||||
"""Test handling of import errors for OpenRouter provider."""
|
||||
mock_import.side_effect = ImportError("Failed to import OpenRouter provider")
|
||||
|
||||
embedder_config = {
|
||||
"provider": "openrouter",
|
||||
"config": {"api_key": "test-key"},
|
||||
}
|
||||
|
||||
with pytest.raises(ImportError) as exc_info:
|
||||
build_embedder(embedder_config)
|
||||
|
||||
assert "Failed to import provider openrouter" in str(exc_info.value)
|
||||
|
||||
|
||||
class TestOpenRouterProviderDirect:
|
||||
"""Test OpenRouterProvider Pydantic settings model directly."""
|
||||
|
||||
def test_default_values(self):
|
||||
"""Test default values for OpenRouterProvider."""
|
||||
provider = OpenRouterProvider(api_key="test-key")
|
||||
|
||||
assert provider.api_key == "test-key"
|
||||
assert provider.model_name == "openai/text-embedding-3-small"
|
||||
assert provider.api_base == "https://openrouter.ai/api/v1"
|
||||
assert provider.dimensions is None
|
||||
assert provider.default_headers is None
|
||||
|
||||
def test_custom_values(self):
|
||||
"""Test custom configuration values."""
|
||||
provider = OpenRouterProvider(
|
||||
api_key="test-custom-key",
|
||||
model="openai/text-embedding-3-large",
|
||||
api_base="https://custom.openrouter.ai/api/v1",
|
||||
dimensions=3072,
|
||||
default_headers={"HTTP-Referer": "https://crewai.com"},
|
||||
)
|
||||
|
||||
assert provider.api_key == "test-custom-key"
|
||||
assert provider.model_name == "openai/text-embedding-3-large"
|
||||
assert provider.api_base == "https://custom.openrouter.ai/api/v1"
|
||||
assert provider.dimensions == 3072
|
||||
assert provider.default_headers == {"HTTP-Referer": "https://crewai.com"}
|
||||
|
||||
def test_missing_api_key_raises_validation_error(self, monkeypatch):
|
||||
"""Test that missing API key raises ValidationError when no env vars set."""
|
||||
monkeypatch.delenv("OPENROUTER_API_KEY", raising=False)
|
||||
|
||||
with pytest.raises(ValidationError):
|
||||
OpenRouterProvider()
|
||||
|
||||
def test_env_var_openrouter_api_key(self, monkeypatch):
|
||||
"""Test resolving API key from OPENROUTER_API_KEY env var."""
|
||||
monkeypatch.setenv("OPENROUTER_API_KEY", "sk-or-env-key")
|
||||
|
||||
provider = OpenRouterProvider()
|
||||
assert provider.api_key == "sk-or-env-key"
|
||||
|
||||
def test_model_alias_normalization(self):
|
||||
"""Test 'model' parameter maps to 'model_name'."""
|
||||
provider = OpenRouterProvider(
|
||||
api_key="test-key",
|
||||
model="cohere/embed-multilingual-v3.0",
|
||||
)
|
||||
assert provider.model_name == "cohere/embed-multilingual-v3.0"
|
||||
|
||||
def test_model_name_takes_precedence(self):
|
||||
"""Test that model_name takes precedence over model if both are given."""
|
||||
provider = OpenRouterProvider(
|
||||
api_key="test-key",
|
||||
model="openai/text-embedding-3-small",
|
||||
model_name="openai/text-embedding-3-large",
|
||||
)
|
||||
assert provider.model_name == "openai/text-embedding-3-large"
|
||||
Reference in New Issue
Block a user