feat(rag): add VoyageAI voyage-context-4 contextualized embedding support#6368
feat(rag): add VoyageAI voyage-context-4 contextualized embedding support#6368fzowl wants to merge 1 commit into
Conversation
…port 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.
📝 WalkthroughWalkthroughAdds support for VoyageAI "contextualized" models in VoyageAI contextualized embedding
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
lib/crewai/tests/rag/embeddings/test_voyageai_embedding_callable.py (1)
30-49: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winCover the single-document multi-chunk shape directly.
This test only proves concatenation across
result.results. The new behavior that matters here is flatteningresult.results[*].embeddingswhen one input auto-chunks into multiple vectors, so a fixture with a singleresultcontaining two embeddings would make the regression target much clearer.Suggested test tweak
mock_client.contextualized_embed.return_value = MagicMock( results=[ - MagicMock(embeddings=[[0.1, 0.2]]), - MagicMock(embeddings=[[0.3, 0.4]]), + MagicMock(embeddings=[[0.1, 0.2], [0.3, 0.4]]), ] ) @@ - result = fn(["aa", "bb"]) + result = fn(["aa"]) @@ 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]])🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@lib/crewai/tests/rag/embeddings/test_voyageai_embedding_callable.py` around lines 30 - 49, Update the contextualized embedding test in VoyageAIEmbeddingFunction so it directly covers the single-document multi-chunk case: keep the existing VoyageAIEmbeddingFunction and contextualized_embed setup, but make the mocked contextualized_embed response return one result whose embeddings contains multiple vectors, then assert the callable flattens result.results[*].embeddings into the final output. This should replace the current multi-result concatenation check while still verifying embed is not called and contextualized_embed is used.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In
`@lib/crewai/src/crewai/rag/embeddings/providers/voyageai/embedding_callable.py`:
- Around line 58-66: The voyage-context path in embedding_callable’s callable
currently hardcodes input_type to document, ignoring any configured value.
Update the contextualized_embed call to use the configured input_type from
self._config when appropriate, or add explicit validation in the embedding
callable to reject non-document configurations for voyage-context models. Make
the fix in the branch that handles model.startswith("voyage-context") so the
behavior matches the selected config.
---
Nitpick comments:
In `@lib/crewai/tests/rag/embeddings/test_voyageai_embedding_callable.py`:
- Around line 30-49: Update the contextualized embedding test in
VoyageAIEmbeddingFunction so it directly covers the single-document multi-chunk
case: keep the existing VoyageAIEmbeddingFunction and contextualized_embed
setup, but make the mocked contextualized_embed response return one result whose
embeddings contains multiple vectors, then assert the callable flattens
result.results[*].embeddings into the final output. This should replace the
current multi-result concatenation check while still verifying embed is not
called and contextualized_embed is used.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: 5a57c8b2-22c8-4709-bd45-3bc212afb30b
📒 Files selected for processing (2)
lib/crewai/src/crewai/rag/embeddings/providers/voyageai/embedding_callable.pylib/crewai/tests/rag/embeddings/test_voyageai_embedding_callable.py
| 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, |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🌐 Web query:
For voyageai Python client version 0.3.5, what is the signature of Client.contextualized_embed, and which input_typevalues does it support? In particular, isquery supported or is the endpoint document-only?
💡 Result:
The contextualized_embed method in the Voyage AI Python client supports both query and document types [1][2]. It is not restricted to documents only [1][2]. The method signature for voyageai.Client.contextualized_embed is as follows [1][2]: def contextualized_embed( self, inputs: Union[List[List[str]], List[str]], model: str, input_type: Optional[str] = None, output_dimension: Optional[int] = None, output_dtype: Optional[str] = "float", enable_auto_chunking: Optional[bool] = False, chunk_size: Optional[int] = 512, chunk_overlap: Optional[int] = 0, chunk_fn: Optional[Callable[[str], List[str]]] = None,) The input_type parameter accepts the values None, "query", and "document" [1][2]. When using "query", each inner list should typically contain a single query string [3][4]. When using "document", each inner list typically contains chunks from a document [3][4]. The input_type argument is optional, but setting it to "query" or "document" is recommended for retrieval tasks as it prompts the model appropriately [1][2].
Citations:
- 1: https://www.mongodb.com/docs/voyageai/models/contextualized-chunk-embeddings/
- 2: https://docs.voyageai.com/docs/contextualized-chunk-embeddings
- 3: https://www.mongodb.com/docs/api/doc/atlas-embedding-and-reranking-api/operation/operation-createcontextualizedembedding
- 4: https://docs.rs/mongodb-voyageai/latest/mongodb_voyageai/client/struct.ContextualizedEmbedBuilder.html
🏁 Script executed:
sed -n '1,160p' lib/crewai/src/crewai/rag/embeddings/providers/voyageai/embedding_callable.py
printf '\n---\n'
rg -n "voyage-context|contextualized_embed|input_type" lib/crewai/src/crewai/rag -SRepository: crewAIInc/crewAI
Length of output: 3654
Respect input_type for voyage-context models contextualized_embed() accepts both query and document, so hardcoding "document" here ignores the configured value. Either forward self._config.get("input_type") or reject non-document configs explicitly.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In
`@lib/crewai/src/crewai/rag/embeddings/providers/voyageai/embedding_callable.py`
around lines 58 - 66, The voyage-context path in embedding_callable’s callable
currently hardcodes input_type to document, ignoring any configured value.
Update the contextualized_embed call to use the configured input_type from
self._config when appropriate, or add explicit validation in the embedding
callable to reject non-document configurations for voyage-context models. Make
the fix in the branch that handles model.startswith("voyage-context") so the
behavior matches the selected config.
1 participant
Add support for VoyageAI models: voyage-context-4
Summary by CodeRabbit
New Features
Tests