mirror of
https://github.com/crewAIInc/crewAI.git
synced 2026-01-08 15:48:29 +00:00
f28e78c5ba
Some checks failed
Notify Downstream / notify-downstream (push) Has been cancelled
Update Test Durations / update-durations (3.10) (push) Has been cancelled
Update Test Durations / update-durations (3.11) (push) Has been cancelled
Update Test Durations / update-durations (3.12) (push) Has been cancelled
Update Test Durations / update-durations (3.13) (push) Has been cancelled
Build uv cache / build-cache (3.10) (push) Has been cancelled
Build uv cache / build-cache (3.11) (push) Has been cancelled
Build uv cache / build-cache (3.12) (push) Has been cancelled
Build uv cache / build-cache (3.13) (push) Has been cancelled
- ignore line length errors globally - migrate knowledge/memory and crew query_knowledge to `SearchResult` - remove legacy chromadb utils; fix empty metadata handling - restore openai as default embedding provider; support instance-specific clients - update and fix tests for `SearchResult` migration and rag changes
251 lines
8.2 KiB
Python
251 lines
8.2 KiB
Python
"""Enhanced tests for embedding function factory."""
|
|
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
import pytest
|
|
|
|
from crewai.rag.embeddings.factory import ( # type: ignore[import-untyped]
|
|
get_embedding_function,
|
|
)
|
|
from crewai.rag.embeddings.types import EmbeddingOptions # type: ignore[import-untyped]
|
|
|
|
|
|
def test_get_embedding_function_default() -> None:
|
|
"""Test default embedding function when no config provided."""
|
|
with patch("crewai.rag.embeddings.factory.OpenAIEmbeddingFunction") as mock_openai:
|
|
mock_instance = MagicMock()
|
|
mock_openai.return_value = mock_instance
|
|
|
|
with patch(
|
|
"crewai.rag.embeddings.factory.os.getenv", return_value="test-api-key"
|
|
):
|
|
result = get_embedding_function()
|
|
|
|
mock_openai.assert_called_once_with(
|
|
api_key="test-api-key", model_name="text-embedding-3-small"
|
|
)
|
|
assert result == mock_instance
|
|
|
|
|
|
def test_get_embedding_function_with_embedding_options() -> None:
|
|
"""Test embedding function creation with EmbeddingOptions object."""
|
|
with patch("crewai.rag.embeddings.factory.OpenAIEmbeddingFunction") as mock_openai:
|
|
mock_instance = MagicMock()
|
|
mock_openai.return_value = mock_instance
|
|
|
|
options = EmbeddingOptions(
|
|
provider="openai", api_key="test-key", model="text-embedding-3-large"
|
|
)
|
|
|
|
result = get_embedding_function(options)
|
|
|
|
call_kwargs = mock_openai.call_args.kwargs
|
|
assert "api_key" in call_kwargs
|
|
assert call_kwargs["api_key"].get_secret_value() == "test-key"
|
|
# OpenAI uses model_name parameter, not model
|
|
assert result == mock_instance
|
|
|
|
|
|
def test_get_embedding_function_sentence_transformer() -> None:
|
|
"""Test sentence transformer embedding function."""
|
|
with patch(
|
|
"crewai.rag.embeddings.factory.SentenceTransformerEmbeddingFunction"
|
|
) as mock_st:
|
|
mock_instance = MagicMock()
|
|
mock_st.return_value = mock_instance
|
|
|
|
config = {"provider": "sentence-transformer", "model_name": "all-MiniLM-L6-v2"}
|
|
|
|
result = get_embedding_function(config)
|
|
|
|
mock_st.assert_called_once_with(model_name="all-MiniLM-L6-v2")
|
|
assert result == mock_instance
|
|
|
|
|
|
def test_get_embedding_function_ollama() -> None:
|
|
"""Test Ollama embedding function."""
|
|
with patch("crewai.rag.embeddings.factory.OllamaEmbeddingFunction") as mock_ollama:
|
|
mock_instance = MagicMock()
|
|
mock_ollama.return_value = mock_instance
|
|
|
|
config = {
|
|
"provider": "ollama",
|
|
"model_name": "nomic-embed-text",
|
|
"url": "http://localhost:11434",
|
|
}
|
|
|
|
result = get_embedding_function(config)
|
|
|
|
mock_ollama.assert_called_once_with(
|
|
model_name="nomic-embed-text", url="http://localhost:11434"
|
|
)
|
|
assert result == mock_instance
|
|
|
|
|
|
def test_get_embedding_function_cohere() -> None:
|
|
"""Test Cohere embedding function."""
|
|
with patch("crewai.rag.embeddings.factory.CohereEmbeddingFunction") as mock_cohere:
|
|
mock_instance = MagicMock()
|
|
mock_cohere.return_value = mock_instance
|
|
|
|
config = {
|
|
"provider": "cohere",
|
|
"api_key": "cohere-key",
|
|
"model_name": "embed-english-v3.0",
|
|
}
|
|
|
|
result = get_embedding_function(config)
|
|
|
|
mock_cohere.assert_called_once_with(
|
|
api_key="cohere-key", model_name="embed-english-v3.0"
|
|
)
|
|
assert result == mock_instance
|
|
|
|
|
|
def test_get_embedding_function_huggingface() -> None:
|
|
"""Test HuggingFace embedding function."""
|
|
with patch("crewai.rag.embeddings.factory.HuggingFaceEmbeddingFunction") as mock_hf:
|
|
mock_instance = MagicMock()
|
|
mock_hf.return_value = mock_instance
|
|
|
|
config = {
|
|
"provider": "huggingface",
|
|
"api_key": "hf-token",
|
|
"model_name": "sentence-transformers/all-MiniLM-L6-v2",
|
|
}
|
|
|
|
result = get_embedding_function(config)
|
|
|
|
mock_hf.assert_called_once_with(
|
|
api_key="hf-token", model_name="sentence-transformers/all-MiniLM-L6-v2"
|
|
)
|
|
assert result == mock_instance
|
|
|
|
|
|
def test_get_embedding_function_onnx() -> None:
|
|
"""Test ONNX embedding function."""
|
|
with patch("crewai.rag.embeddings.factory.ONNXMiniLM_L6_V2") as mock_onnx:
|
|
mock_instance = MagicMock()
|
|
mock_onnx.return_value = mock_instance
|
|
|
|
config = {"provider": "onnx"}
|
|
|
|
result = get_embedding_function(config)
|
|
|
|
mock_onnx.assert_called_once()
|
|
assert result == mock_instance
|
|
|
|
|
|
def test_get_embedding_function_google_palm() -> None:
|
|
"""Test Google PaLM embedding function."""
|
|
with patch(
|
|
"crewai.rag.embeddings.factory.GooglePalmEmbeddingFunction"
|
|
) as mock_palm:
|
|
mock_instance = MagicMock()
|
|
mock_palm.return_value = mock_instance
|
|
|
|
config = {"provider": "google-palm", "api_key": "palm-key"}
|
|
|
|
result = get_embedding_function(config)
|
|
|
|
mock_palm.assert_called_once_with(api_key="palm-key")
|
|
assert result == mock_instance
|
|
|
|
|
|
def test_get_embedding_function_amazon_bedrock() -> None:
|
|
"""Test Amazon Bedrock embedding function."""
|
|
with patch(
|
|
"crewai.rag.embeddings.factory.AmazonBedrockEmbeddingFunction"
|
|
) as mock_bedrock:
|
|
mock_instance = MagicMock()
|
|
mock_bedrock.return_value = mock_instance
|
|
|
|
config = {
|
|
"provider": "amazon-bedrock",
|
|
"region_name": "us-west-2",
|
|
"model_name": "amazon.titan-embed-text-v1",
|
|
}
|
|
|
|
result = get_embedding_function(config)
|
|
|
|
mock_bedrock.assert_called_once_with(
|
|
region_name="us-west-2", model_name="amazon.titan-embed-text-v1"
|
|
)
|
|
assert result == mock_instance
|
|
|
|
|
|
def test_get_embedding_function_jina() -> None:
|
|
"""Test Jina embedding function."""
|
|
with patch("crewai.rag.embeddings.factory.JinaEmbeddingFunction") as mock_jina:
|
|
mock_instance = MagicMock()
|
|
mock_jina.return_value = mock_instance
|
|
|
|
config = {
|
|
"provider": "jina",
|
|
"api_key": "jina-key",
|
|
"model_name": "jina-embeddings-v2-base-en",
|
|
}
|
|
|
|
result = get_embedding_function(config)
|
|
|
|
mock_jina.assert_called_once_with(
|
|
api_key="jina-key", model_name="jina-embeddings-v2-base-en"
|
|
)
|
|
assert result == mock_instance
|
|
|
|
|
|
def test_get_embedding_function_unsupported_provider() -> None:
|
|
"""Test handling of unsupported provider."""
|
|
config = {"provider": "unsupported-provider"}
|
|
|
|
with pytest.raises(ValueError, match="Unsupported provider: unsupported-provider"):
|
|
get_embedding_function(config)
|
|
|
|
|
|
def test_get_embedding_function_config_modification() -> None:
|
|
"""Test that original config dict is not modified."""
|
|
original_config = {
|
|
"provider": "openai",
|
|
"api_key": "test-key",
|
|
"model": "text-embedding-3-small",
|
|
}
|
|
config_copy = original_config.copy()
|
|
|
|
with patch("crewai.rag.embeddings.factory.OpenAIEmbeddingFunction"):
|
|
get_embedding_function(config_copy)
|
|
|
|
assert config_copy == original_config
|
|
|
|
|
|
def test_get_embedding_function_exclude_none_values() -> None:
|
|
"""Test that None values are excluded from embedding function calls."""
|
|
with patch("crewai.rag.embeddings.factory.OpenAIEmbeddingFunction") as mock_openai:
|
|
mock_instance = MagicMock()
|
|
mock_openai.return_value = mock_instance
|
|
|
|
options = EmbeddingOptions(provider="openai", api_key="test-key", model=None)
|
|
|
|
result = get_embedding_function(options)
|
|
|
|
call_kwargs = mock_openai.call_args.kwargs
|
|
assert "api_key" in call_kwargs
|
|
assert call_kwargs["api_key"].get_secret_value() == "test-key"
|
|
assert "model" not in call_kwargs
|
|
assert result == mock_instance
|
|
|
|
|
|
def test_get_embedding_function_instructor() -> None:
|
|
"""Test Instructor embedding function."""
|
|
with patch(
|
|
"crewai.rag.embeddings.factory.InstructorEmbeddingFunction"
|
|
) as mock_instructor:
|
|
mock_instance = MagicMock()
|
|
mock_instructor.return_value = mock_instance
|
|
|
|
config = {"provider": "instructor", "model_name": "hkunlp/instructor-large"}
|
|
|
|
result = get_embedding_function(config)
|
|
|
|
mock_instructor.assert_called_once_with(model_name="hkunlp/instructor-large")
|
|
assert result == mock_instance
|