mirror of
https://github.com/crewAIInc/crewAI.git
synced 2026-01-11 00:58:30 +00:00
Fix issue #2832: Add VoyageAI embedding function implementation
Co-Authored-By: Joe Moura <joao@crewai.com>
This commit is contained in:
47
tests/utilities/test_embedding_configurator.py
Normal file
47
tests/utilities/test_embedding_configurator.py
Normal file
@@ -0,0 +1,47 @@
|
||||
import os
|
||||
import pytest
|
||||
from unittest.mock import patch, MagicMock
|
||||
|
||||
from crewai.utilities.embedding_configurator import EmbeddingConfigurator
|
||||
|
||||
|
||||
class TestEmbeddingConfigurator:
|
||||
def test_configure_voyageai_embedder(self):
|
||||
"""Test that the VoyageAI embedder is configured correctly."""
|
||||
with patch(
|
||||
"crewai.utilities.embedding_configurator.VoyageAIEmbeddingFunction"
|
||||
) as mock_voyageai:
|
||||
mock_instance = MagicMock()
|
||||
mock_voyageai.return_value = mock_instance
|
||||
|
||||
config = {"api_key": "test-key"}
|
||||
model_name = "voyage-3"
|
||||
|
||||
configurator = EmbeddingConfigurator()
|
||||
embedder = configurator._configure_voyageai(config, model_name)
|
||||
|
||||
mock_voyageai.assert_called_once_with(
|
||||
model_name=model_name, api_key="test-key"
|
||||
)
|
||||
assert embedder == mock_instance
|
||||
|
||||
def test_configure_embedder_with_voyageai(self):
|
||||
"""Test that the embedder configurator correctly handles VoyageAI provider."""
|
||||
with patch(
|
||||
"crewai.utilities.embedding_configurator.VoyageAIEmbeddingFunction"
|
||||
) as mock_voyageai:
|
||||
mock_instance = MagicMock()
|
||||
mock_voyageai.return_value = mock_instance
|
||||
|
||||
embedder_config = {
|
||||
"provider": "voyageai",
|
||||
"config": {"api_key": "test-key", "model": "voyage-3"},
|
||||
}
|
||||
|
||||
configurator = EmbeddingConfigurator()
|
||||
embedder = configurator.configure_embedder(embedder_config)
|
||||
|
||||
mock_voyageai.assert_called_once_with(
|
||||
model_name="voyage-3", api_key="test-key"
|
||||
)
|
||||
assert embedder == mock_instance
|
||||
Reference in New Issue
Block a user