feat(voyageai): refresh models, add contextualized embeddings + token-aware batching - #855
Open
fzowl wants to merge 3 commits into
Open
feat(voyageai): refresh models, add contextualized embeddings + token-aware batching#855fzowl wants to merge 3 commits into
fzowl wants to merge 3 commits into
Conversation
…and token-aware batching Update the Voyage AI integration to the current catalog and add support for the contextualized-chunk embedding models. Embeddings (kotaemon/embeddings/voyageai.py): - Support voyage-context-* models via the contextualized_embed API. Each input string is embedded as its own independent document: the batch is sent as a flat list[str] with enable_auto_chunking=True and chunk_size=32000 so every string resolves to exactly one chunk and one embedding. Cross-input contextualization is intentionally not used because generic embed_many callers pass unrelated texts. - The query path drops auto-chunking (the API rejects it for input_type "query"): enable_auto_chunking = input_type != "query", and chunk_size is omitted when disabled. - Token-aware batching for the plain-text path (sync + async): batches are bounded by both the item count and the per-model total-token budget; a single oversized text is still sent on its own. - Fix the async path (previously awaited an attribute, not the coroutine). Rerankings: default to the current rerank-2.5 model. Config: VOYAGE_RERANK_MODEL env var; document model options in .env.example. Tests: cover the document and query contextualized paths plus the batching logic (token boundary, oversized single text, item-count cap).
…ng API The contextualized-chunk path calls Client.contextualized_embed with the flat-list inputs plus enable_auto_chunking/chunk_size parameters, which only exist in voyageai>=0.4.0. The lockfile pinned voyageai 0.3.2, whose Client has no contextualized_embed attribute, so the contextualized unit tests failed at patch time and the feature could not work at runtime. - Bump the specifier to voyageai>=0.4.0 and relock (voyage 0.5.0). - voyageai>=0.4.0 pulls numpy>=2.1 on Python 3.13, which conflicts with unstructured<0.16 (numpy<2); fall back to an older voyageai there via a marker split (the test matrix is 3.10/3.11). - voyageai declares langchain-text-splitters>=0.3.8, which would cascade a langchain-core 0.3 bump across the app. It only needs RecursiveCharacterTextSplitter (present since 0.2.x) for the optional local chunk_fn, never used on the server-side auto-chunking path, so override the splitter to the 0.2 line and leave the langchain stack as is.
Contributor
Author
|
CI was failing because the lockfile pinned
Only |
The retriever now tags query embeddings with input_type="query" so asymmetric models (Voyage AI) apply the query prompt instead of the document prompt, and the contextualized query path (auto-chunking disabled for queries) is reachable from the application. input_type is forwarded from the two VectorRetrieval query call sites (vector and hybrid modes). It is an optional embedding kwarg: backends that don't distinguish query from document accept and ignore it (langchain-based, endpoint, OpenAI drops it before the API call), so every other embedding backend is unchanged. Also tokenize each batch in a single call in the token-aware batcher, document the asymmetric/re-index implication and the long-input truncation contract, and cover the query path end to end.
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
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.
What
Refreshes the Voyage AI integration to the current model catalog and adds first-class support for the contextualized-chunk embedding models.
Embeddings (
kotaemon/embeddings/voyageai.py)voyage-context-4,voyage-context-3) via thecontextualized_embedAPI. Each input string is embedded as its own independent document: the batch is sent as a flatlist[str]withenable_auto_chunking=Trueandchunk_size=32000, so every string resolves to exactly one chunk and one embedding. This yields deterministic per-input vectors and trivial result collection. Cross-input contextualization (inputs=[batch]) is intentionally not used, because genericembed_manycallers pass unrelated texts. This is documented on the class.input_type="query", so the retrieval path disables it —enable_auto_chunking = input_type != "query", andchunk_sizeis dropped when disabled.embed/embed_manyand async equivalents): batches are bounded by both the item count (MAX_BATCH_SIZE) and the per-model total-token budget, not only a fixed item count. A single oversized text is still sent on its own.Rerankings
rerank-2.5model; help text lists the current models.Config
VOYAGE_RERANK_MODELenv var wired intoflowsettings.py; model options documented in.env.example.Tests
list[str],enable_auto_chunking=True,chunk_size=32000.chunk_size.Validation
The repo's Python environment has a pre-existing langchain/langchain-core version mismatch that blocks collecting the kotaemon test suite (unrelated to this change). Ran the affected tests in an isolated environment with compatible langchain pins: all 6 Voyage AI tests pass; the only unrelated failure is the pre-existing Cohere embedding test.
flake8(max-line-length 88),black, andisortare clean on the changed files.