From 042bd71410bd33b3819c768227b76ec80ff5488c Mon Sep 17 00:00:00 2001 From: fzowl Date: Fri, 26 Jun 2026 15:20:20 +0200 Subject: [PATCH 1/3] feat(rag): add VoyageAI voyage-context-4 contextualized embedding support Route voyage-context-* models through the contextualized embeddings endpoint (client.contextualized_embed) with chunk_size set to the 32000 maximum. Input is passed as a flat list of strings and the per-document chunk embeddings are flattened into the returned vectors. --- .../providers/voyageai/embedding_callable.py | 22 +++- .../test_voyageai_embedding_callable.py | 104 ++++++++++++++++++ 2 files changed, 125 insertions(+), 1 deletion(-) create mode 100644 lib/crewai/tests/rag/embeddings/test_voyageai_embedding_callable.py diff --git a/lib/crewai/src/crewai/rag/embeddings/providers/voyageai/embedding_callable.py b/lib/crewai/src/crewai/rag/embeddings/providers/voyageai/embedding_callable.py index 56a13ad6bb..8835427633 100644 --- a/lib/crewai/src/crewai/rag/embeddings/providers/voyageai/embedding_callable.py +++ b/lib/crewai/src/crewai/rag/embeddings/providers/voyageai/embedding_callable.py @@ -8,6 +8,9 @@ from crewai.rag.embeddings.providers.voyageai.types import VoyageAIProviderConfig +CONTEXTUALIZED_CHUNK_SIZE = 32000 + + class VoyageAIEmbeddingFunction(EmbeddingFunction[Documents]): """Embedding function for VoyageAI models.""" @@ -50,9 +53,26 @@ def __call__(self, input: Documents) -> Embeddings: if isinstance(input, str): input = [input] + model = self._config.get("model", "voyage-2") + + if model.startswith("voyage-context"): + result = self._client.contextualized_embed( + inputs=input, + model=model, + input_type="document", + output_dtype=self._config.get("output_dtype"), + output_dimension=self._config.get("output_dimension"), + enable_auto_chunking=True, + chunk_size=CONTEXTUALIZED_CHUNK_SIZE, + ) + return cast( + Embeddings, + [embedding for r in result.results for embedding in r.embeddings], + ) + result = self._client.embed( texts=input, - model=self._config.get("model", "voyage-2"), + model=model, input_type=self._config.get("input_type"), truncation=self._config.get("truncation", True), output_dtype=self._config.get("output_dtype"), diff --git a/lib/crewai/tests/rag/embeddings/test_voyageai_embedding_callable.py b/lib/crewai/tests/rag/embeddings/test_voyageai_embedding_callable.py new file mode 100644 index 0000000000..8d3e7c6f43 --- /dev/null +++ b/lib/crewai/tests/rag/embeddings/test_voyageai_embedding_callable.py @@ -0,0 +1,104 @@ +"""Tests for the VoyageAI embedding function.""" + +from unittest.mock import MagicMock, patch + +import numpy as np + +from crewai.rag.embeddings.providers.voyageai.embedding_callable import ( + CONTEXTUALIZED_CHUNK_SIZE, + VoyageAIEmbeddingFunction, +) + + +class TestVoyageAIEmbeddingFunction: + """Test the VoyageAI embedding function call routing.""" + + def test_standard_model_uses_embed(self): + """Standard models should call the regular embed endpoint.""" + with patch("voyageai.Client") as mock_client_class: + mock_client = MagicMock() + mock_client_class.return_value = mock_client + mock_client.embed.return_value = MagicMock(embeddings=[[0.1, 0.2]]) + + fn = VoyageAIEmbeddingFunction(api_key="voyage-key", model="voyage-2") + result = fn(["aa", "bb"]) + + mock_client.embed.assert_called_once() + mock_client.contextualized_embed.assert_not_called() + assert np.allclose(result, [[0.1, 0.2]]) + + def test_contextualized_model_uses_contextualized_embed(self): + """voyage-context-4 should call the contextualized embeddings endpoint.""" + with patch("voyageai.Client") as mock_client_class: + mock_client = MagicMock() + mock_client_class.return_value = mock_client + mock_client.contextualized_embed.return_value = MagicMock( + results=[ + MagicMock(embeddings=[[0.1, 0.2]]), + MagicMock(embeddings=[[0.3, 0.4]]), + ] + ) + + fn = VoyageAIEmbeddingFunction( + api_key="voyage-key", model="voyage-context-4" + ) + result = fn(["aa", "bb"]) + + mock_client.embed.assert_not_called() + mock_client.contextualized_embed.assert_called_once() + assert np.allclose(result, [[0.1, 0.2], [0.3, 0.4]]) + + def test_contextualized_call_sets_chunk_size_to_max(self): + """chunk_size must be set to 32000 on every contextualized call.""" + with patch("voyageai.Client") as mock_client_class: + mock_client = MagicMock() + mock_client_class.return_value = mock_client + mock_client.contextualized_embed.return_value = MagicMock( + results=[MagicMock(embeddings=[[0.1, 0.2]])] + ) + + fn = VoyageAIEmbeddingFunction( + api_key="voyage-key", model="voyage-context-4" + ) + fn(["aa"]) + + _, kwargs = mock_client.contextualized_embed.call_args + assert kwargs["chunk_size"] == CONTEXTUALIZED_CHUNK_SIZE + assert CONTEXTUALIZED_CHUNK_SIZE == 32000 + + def test_contextualized_input_is_flat_list(self): + """Input must be passed as a flat List[str], not wrapped in an extra list.""" + with patch("voyageai.Client") as mock_client_class: + mock_client = MagicMock() + mock_client_class.return_value = mock_client + mock_client.contextualized_embed.return_value = MagicMock( + results=[ + MagicMock(embeddings=[[0.1, 0.2]]), + MagicMock(embeddings=[[0.3, 0.4]]), + ] + ) + + fn = VoyageAIEmbeddingFunction( + api_key="voyage-key", model="voyage-context-4" + ) + fn(["aa", "bb"]) + + _, kwargs = mock_client.contextualized_embed.call_args + assert kwargs["inputs"] == ["aa", "bb"] + + def test_contextualized_string_input_normalized_to_flat_list(self): + """A single string input is normalized to a flat list of one string.""" + with patch("voyageai.Client") as mock_client_class: + mock_client = MagicMock() + mock_client_class.return_value = mock_client + mock_client.contextualized_embed.return_value = MagicMock( + results=[MagicMock(embeddings=[[0.1, 0.2]])] + ) + + fn = VoyageAIEmbeddingFunction( + api_key="voyage-key", model="voyage-context-4" + ) + fn("aa") + + _, kwargs = mock_client.contextualized_embed.call_args + assert kwargs["inputs"] == ["aa"] From 297194325967e2272b36095a198c96834d1ef400 Mon Sep 17 00:00:00 2001 From: fzowl Date: Sun, 28 Jun 2026 15:33:10 +0200 Subject: [PATCH 2/3] fix: correct voyage-context-4 contextualized_embed API call and tests - Wrap inputs as List[List[str]] (each string = single-chunk document) - Remove invalid enable_auto_chunking and chunk_size params - Use .results[i].embeddings[0] for single-chunk extraction - Update tests to match correct API contract - mypy: 0 errors, pytest: 5/5 VoyageAI + 15/15 factory tests pass --- .../providers/voyageai/embedding_callable.py | 9 +++---- .../test_voyageai_embedding_callable.py | 25 +++++++++++-------- 2 files changed, 18 insertions(+), 16 deletions(-) diff --git a/lib/crewai/src/crewai/rag/embeddings/providers/voyageai/embedding_callable.py b/lib/crewai/src/crewai/rag/embeddings/providers/voyageai/embedding_callable.py index 8835427633..ecb8f53edf 100644 --- a/lib/crewai/src/crewai/rag/embeddings/providers/voyageai/embedding_callable.py +++ b/lib/crewai/src/crewai/rag/embeddings/providers/voyageai/embedding_callable.py @@ -56,18 +56,17 @@ def __call__(self, input: Documents) -> Embeddings: model = self._config.get("model", "voyage-2") if model.startswith("voyage-context"): - result = self._client.contextualized_embed( - inputs=input, + inputs = [[s] for s in input] + ctx_result = self._client.contextualized_embed( + inputs=inputs, model=model, input_type="document", output_dtype=self._config.get("output_dtype"), output_dimension=self._config.get("output_dimension"), - enable_auto_chunking=True, - chunk_size=CONTEXTUALIZED_CHUNK_SIZE, ) return cast( Embeddings, - [embedding for r in result.results for embedding in r.embeddings], + [r.embeddings[0] for r in ctx_result.results], ) result = self._client.embed( diff --git a/lib/crewai/tests/rag/embeddings/test_voyageai_embedding_callable.py b/lib/crewai/tests/rag/embeddings/test_voyageai_embedding_callable.py index 8d3e7c6f43..8c8803efbc 100644 --- a/lib/crewai/tests/rag/embeddings/test_voyageai_embedding_callable.py +++ b/lib/crewai/tests/rag/embeddings/test_voyageai_embedding_callable.py @@ -48,8 +48,8 @@ def test_contextualized_model_uses_contextualized_embed(self): mock_client.contextualized_embed.assert_called_once() assert np.allclose(result, [[0.1, 0.2], [0.3, 0.4]]) - def test_contextualized_call_sets_chunk_size_to_max(self): - """chunk_size must be set to 32000 on every contextualized call.""" + def test_contextualized_call_wraps_inputs_as_list_of_lists(self): + """Each input string is wrapped as its own single-chunk document (List[List[str]]).""" with patch("voyageai.Client") as mock_client_class: mock_client = MagicMock() mock_client_class.return_value = mock_client @@ -63,11 +63,14 @@ def test_contextualized_call_sets_chunk_size_to_max(self): fn(["aa"]) _, kwargs = mock_client.contextualized_embed.call_args - assert kwargs["chunk_size"] == CONTEXTUALIZED_CHUNK_SIZE - assert CONTEXTUALIZED_CHUNK_SIZE == 32000 - - def test_contextualized_input_is_flat_list(self): - """Input must be passed as a flat List[str], not wrapped in an extra list.""" + # Each string is wrapped as its own single-chunk document + assert kwargs["inputs"] == [["aa"]] + # chunk_size and enable_auto_chunking must NOT be passed + assert "chunk_size" not in kwargs + assert "enable_auto_chunking" not in kwargs + + def test_contextualized_input_is_list_of_lists(self): + """Input must be passed as List[List[str]], each inner list is one document with its chunks.""" with patch("voyageai.Client") as mock_client_class: mock_client = MagicMock() mock_client_class.return_value = mock_client @@ -84,10 +87,10 @@ def test_contextualized_input_is_flat_list(self): fn(["aa", "bb"]) _, kwargs = mock_client.contextualized_embed.call_args - assert kwargs["inputs"] == ["aa", "bb"] + assert kwargs["inputs"] == [["aa"], ["bb"]] - def test_contextualized_string_input_normalized_to_flat_list(self): - """A single string input is normalized to a flat list of one string.""" + def test_contextualized_string_input_normalized_with_wrapping(self): + """A single string input is normalized and wrapped as a single-chunk document.""" with patch("voyageai.Client") as mock_client_class: mock_client = MagicMock() mock_client_class.return_value = mock_client @@ -101,4 +104,4 @@ def test_contextualized_string_input_normalized_to_flat_list(self): fn("aa") _, kwargs = mock_client.contextualized_embed.call_args - assert kwargs["inputs"] == ["aa"] + assert kwargs["inputs"] == [["aa"]] From d595c9e32f4e8602c590df9ce35990b870e7daa0 Mon Sep 17 00:00:00 2001 From: fzowl Date: Tue, 30 Jun 2026 14:46:21 +0200 Subject: [PATCH 3/3] test(rag): mock voyageai module so callable tests run without the optional dep The VoyageAI callable tests patched `voyageai.Client` directly, which raises ModuleNotFoundError in CI where `voyageai` is an optional, uninstalled dependency. Inject a mock `voyageai` module into `sys.modules` instead so the lazy `import voyageai` inside `VoyageAIEmbeddingFunction.__init__` resolves regardless of whether the package is present. --- .../test_voyageai_embedding_callable.py | 39 +++++++++++-------- 1 file changed, 23 insertions(+), 16 deletions(-) diff --git a/lib/crewai/tests/rag/embeddings/test_voyageai_embedding_callable.py b/lib/crewai/tests/rag/embeddings/test_voyageai_embedding_callable.py index 8c8803efbc..1fc50c4c1c 100644 --- a/lib/crewai/tests/rag/embeddings/test_voyageai_embedding_callable.py +++ b/lib/crewai/tests/rag/embeddings/test_voyageai_embedding_callable.py @@ -1,23 +1,38 @@ """Tests for the VoyageAI embedding function.""" +import sys +from contextlib import contextmanager from unittest.mock import MagicMock, patch import numpy as np from crewai.rag.embeddings.providers.voyageai.embedding_callable import ( - CONTEXTUALIZED_CHUNK_SIZE, VoyageAIEmbeddingFunction, ) +@contextmanager +def mock_voyageai_client(): + """Inject a fake ``voyageai`` module so tests run without the optional dependency. + + ``VoyageAIEmbeddingFunction`` imports ``voyageai`` lazily, so patching + ``voyageai.Client`` directly fails when the package is not installed (e.g. in + CI). Registering a mock module in ``sys.modules`` makes the lazy import resolve + to the mock regardless of whether the real package is present. + """ + mock_client = MagicMock() + mock_module = MagicMock() + mock_module.Client.return_value = mock_client + with patch.dict(sys.modules, {"voyageai": mock_module}): + yield mock_client + + class TestVoyageAIEmbeddingFunction: """Test the VoyageAI embedding function call routing.""" def test_standard_model_uses_embed(self): """Standard models should call the regular embed endpoint.""" - with patch("voyageai.Client") as mock_client_class: - mock_client = MagicMock() - mock_client_class.return_value = mock_client + with mock_voyageai_client() as mock_client: mock_client.embed.return_value = MagicMock(embeddings=[[0.1, 0.2]]) fn = VoyageAIEmbeddingFunction(api_key="voyage-key", model="voyage-2") @@ -29,9 +44,7 @@ def test_standard_model_uses_embed(self): def test_contextualized_model_uses_contextualized_embed(self): """voyage-context-4 should call the contextualized embeddings endpoint.""" - with patch("voyageai.Client") as mock_client_class: - mock_client = MagicMock() - mock_client_class.return_value = mock_client + with mock_voyageai_client() as mock_client: mock_client.contextualized_embed.return_value = MagicMock( results=[ MagicMock(embeddings=[[0.1, 0.2]]), @@ -50,9 +63,7 @@ def test_contextualized_model_uses_contextualized_embed(self): def test_contextualized_call_wraps_inputs_as_list_of_lists(self): """Each input string is wrapped as its own single-chunk document (List[List[str]]).""" - with patch("voyageai.Client") as mock_client_class: - mock_client = MagicMock() - mock_client_class.return_value = mock_client + with mock_voyageai_client() as mock_client: mock_client.contextualized_embed.return_value = MagicMock( results=[MagicMock(embeddings=[[0.1, 0.2]])] ) @@ -71,9 +82,7 @@ def test_contextualized_call_wraps_inputs_as_list_of_lists(self): def test_contextualized_input_is_list_of_lists(self): """Input must be passed as List[List[str]], each inner list is one document with its chunks.""" - with patch("voyageai.Client") as mock_client_class: - mock_client = MagicMock() - mock_client_class.return_value = mock_client + with mock_voyageai_client() as mock_client: mock_client.contextualized_embed.return_value = MagicMock( results=[ MagicMock(embeddings=[[0.1, 0.2]]), @@ -91,9 +100,7 @@ def test_contextualized_input_is_list_of_lists(self): def test_contextualized_string_input_normalized_with_wrapping(self): """A single string input is normalized and wrapped as a single-chunk document.""" - with patch("voyageai.Client") as mock_client_class: - mock_client = MagicMock() - mock_client_class.return_value = mock_client + with mock_voyageai_client() as mock_client: mock_client.contextualized_embed.return_value = MagicMock( results=[MagicMock(embeddings=[[0.1, 0.2]])] )