feat(rag): add VoyageAI voyage-context-4 contextualized embedding support#6368
Open
fzowl wants to merge 4 commits into
Open
feat(rag): add VoyageAI voyage-context-4 contextualized embedding support#6368fzowl wants to merge 4 commits into
fzowl wants to merge 4 commits 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 ChangesVoyageAI 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
- 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
fzowl
force-pushed
the
feat/embedding-model-voyage-context-4
branch
from
99e5301 to
2971943
Compare
There was a problem hiding this comment.
🧹 Nitpick comments (1)
lib/crewai/tests/rag/embeddings/test_voyageai_embedding_callable.py (1)
51-107: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd a regression test for
input_typeonvoyage-context*.These tests cover wrapping and flattening, but they never assert what gets passed when the function is configured with
input_type="query". That gap is why the current hardcoded"document"behavior can slip through.🧪 Minimal test to add
+ def test_contextualized_model_forwards_input_type(self): + 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", + input_type="query", + ) + fn(["aa"]) + + _, kwargs = mock_client.contextualized_embed.call_args + assert kwargs["input_type"] == "query"🤖 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 51 - 107, Add a regression test around VoyageAIEmbeddingFunction for voyage-context* models when configured with input_type="query"; the current contextualized embedding tests only cover document-style wrapping and miss this branch. Extend the existing test module by creating a case that instantiates VoyageAIEmbeddingFunction with input_type="query", invokes it, and asserts the mocked voyageai.Client.contextualized_embed call receives the query-specific input_type rather than the hardcoded document value. Keep the assertions aligned with the existing contextualized_embed call_args checks so the new test clearly guards the input_type wiring.
🤖 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.
Nitpick comments:
In `@lib/crewai/tests/rag/embeddings/test_voyageai_embedding_callable.py`:
- Around line 51-107: Add a regression test around VoyageAIEmbeddingFunction for
voyage-context* models when configured with input_type="query"; the current
contextualized embedding tests only cover document-style wrapping and miss this
branch. Extend the existing test module by creating a case that instantiates
VoyageAIEmbeddingFunction with input_type="query", invokes it, and asserts the
mocked voyageai.Client.contextualized_embed call receives the query-specific
input_type rather than the hardcoded document value. Keep the assertions aligned
with the existing contextualized_embed call_args checks so the new test clearly
guards the input_type wiring.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: c8757ef0-82ae-4557-aa08-97e2d74e53b2
📒 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
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Add support for VoyageAI models: voyage-context-4
Summary by CodeRabbit
New Features
Tests