feat(voyageai): refresh models, add contextualized embeddings + token-aware batching - #2
feat(voyageai): refresh models, add contextualized embeddings + token-aware batching#2fzowl 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).
|
VERDICT:APPROVED Reviewed the full diff. Correct, well-tested, clean. Approving with a few non-blocking notes. Strengths
Notes (non-blocking)
None of these are correctness bugs. Ship it; consider addressing (1) in release notes. |
|
Upstream PR opened: Cinnamon#855 This PR stays open as the working channel until the upstream PR is resolved. |
…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.
|
Housekeeping: fixed the red upstream CI. The lock pinned |
|
VERDICT:CHANGES_NEEDED Solid, well-crafted PR overall — clean structure, good test coverage (document/query/batching edge cases), clear docstrings and One substantive issue blocks merge: The new
Consequences:
Fix: forward Minor / non-blocking:
|
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.
|
Housekeeping: addressed the blocking query-plumbing review. The retriever now forwards
Also: the token-aware batcher now tokenizes each window in a single Tests: added an end-to-end retrieval test asserting the query is embedded with |
|
VERDICT:APPROVED Reviewed diff, commits, tests, docs. Solid work — ship it. Correctness
Fit / breaking-change risk — low
Tests / docs — good
Commit hygiene — clean
Non-blocking notes
None block merge. |
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.