Skip to content

feat(voyageai): refresh models, add contextualized embeddings + token-aware batching - #2

Open
fzowl wants to merge 3 commits into
mainfrom
feat/voyageai-refresh
Open

feat(voyageai): refresh models, add contextualized embeddings + token-aware batching#2
fzowl wants to merge 3 commits into
mainfrom
feat/voyageai-refresh

Conversation

@fzowl

@fzowl fzowl commented Aug 16, 2026

Copy link
Copy Markdown
Owner

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)

  • Contextualized models (voyage-context-4, voyage-context-3) 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. This yields deterministic per-input vectors and trivial result collection. Cross-input contextualization (inputs=[batch]) is intentionally not used, because generic embed_many callers pass unrelated texts. This is documented on the class.
  • Query path: the API rejects auto-chunking when input_type="query", so the retrieval path disables it — enable_auto_chunking = input_type != "query", and chunk_size is dropped when disabled.
  • Token-aware batching for the plain-text path (embed/embed_many and 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.
  • Fixes the async path, which previously awaited an attribute instead of the coroutine.

Rerankings

  • Default updated to the current rerank-2.5 model; help text lists the current models.

Config

  • VOYAGE_RERANK_MODEL env var wired into flowsettings.py; model options documented in .env.example.

Tests

  • Contextualized document path: flat list[str], enable_auto_chunking=True, chunk_size=32000.
  • Contextualized query path: auto-chunking disabled, no chunk_size.
  • Batching logic: split at the token boundary, a single oversized text goes through alone, and the item-count cap is respected.

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, and isort are clean on the changed files.

…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).
@fzowl

fzowl commented Aug 16, 2026

Copy link
Copy Markdown
Owner Author

VERDICT:APPROVED

Reviewed the full diff. Correct, well-tested, clean. Approving with a few non-blocking notes.

Strengths

  • Real bug fixed: the old async path did await self._aclient.embed(...).embeddings — awaiting the attribute of a coroutine, not the coroutine. New _aembed_* helpers await first, then read .embeddings. Good catch.
  • Contextualized support is coherent: flat list[str] + enable_auto_chunking=True + chunk_size for documents, auto-chunking disabled for queries (API rejects it there). Design tradeoff (one input = one document, no cross-input contextualization) is documented on the class and in the commit.
  • Token-aware batching correctly bounds by both MAX_BATCH_SIZE and per-model token budget, and the empty-batch guard sends an oversized text on its own instead of dropping it.
  • Tests cover document path, query path, token boundary, oversized-single, and item-count cap. Assertions check the actual request kwargs, not just output shape.
  • Commit hygiene clean: single commit, real author (zoltan@voyageai.com), no Co-Authored-By, no "Generated with", no .claude/ files. Conventional-commit style matches repo history.
  • Conforms to repo conventions: Param(...) fields, 3.10+ builtin generics, flake8/black/isort clean, line length under 88.

Notes (non-blocking)

  1. Default input_type for regular models changed from None to "document". Previously embed(texts, model=...) sent no input_type; now every call sends "document". Voyage produces different vectors with a document prompt prefix, so existing deployments would need a re-index to stay consistent, and query-side callers get document-typed vectors (no caller in the repo passes input_type="query"grep finds zero). The new query branch is therefore correct but currently unreachable from the app itself. Worth a note in docs/changelog about the re-index implication.
  2. Long-document silent truncation. With chunk_size=32000, an input longer than 32000 tokens is auto-chunked into multiple chunks server-side, but result.embeddings[0] keeps only the first chunk and silently discards the rest. Fine for pre-chunked RAG inputs (the common case), but a latent assumption worth a comment or guard.
  3. Minor: _iter_batches uses the sync self._client.tokenize even on the async path (blocking call in async context) and tokenizes one text per call in a loop. tokenize is local so impact is small, but a single batched tokenize(texts) call would be cleaner.

None of these are correctness bugs. Ship it; consider addressing (1) in release notes.

@fzowl

fzowl commented Aug 17, 2026

Copy link
Copy Markdown
Owner Author

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.
@fzowl

fzowl commented Aug 17, 2026

Copy link
Copy Markdown
Owner Author

Housekeeping: fixed the red upstream CI. The lock pinned voyageai==0.3.2 (no contextualized_embed), so the contextualized tests errored at patch time. Bumped to voyageai>=0.4.0 (relocked to 0.5.0), added a Python 3.13 marker fallback for the numpy<->unstructured conflict, and a uv override-dependencies pin keeping langchain-text-splitters on 0.2.x so the langchain stack is untouched. Unit tests on 3.10 and 3.11 are green.

@fzowl

fzowl commented Aug 17, 2026

Copy link
Copy Markdown
Owner Author

VERDICT:CHANGES_NEEDED

Solid, well-crafted PR overall — clean structure, good test coverage (document/query/batching edge cases), clear docstrings and .env.example docs, a genuine async bug fix (the old code awaited .embeddings on the un-awaited coroutine), and thoughtful dependency handling (the py3.13 numpy marker split and the langchain-text-splitters override are both well-justified in comments and commit messages). Commit hygiene is clean. VOYAGE_EMBEDDINGS_MODEL in .env.example correctly matches the pre-existing wiring in flowsettings.py:181.

One substantive issue blocks merge:

The new input_type parameter is never plumbed through from the retriever. invoke/ainvoke now default input_type="document", but the retrieval query path never overrides it — both call sites embed the query with no input_type:

  • libs/kotaemon/kotaemon/indices/vectorindex.py:169 (vector mode): emb = self.embedding(text)[0].embedding
  • libs/kotaemon/kotaemon/indices/vectorindex.py:188 (hybrid mode): same

Consequences:

  1. Query-side regression for regular models. Previously no input_type was sent (symmetric). Now every query gets the document prompt prepended, which is a mismatch that measurably hurts retrieval quality for asymmetric Voyage models. This affects all existing Voyage embedding users, silently.
  2. The contextualized query branch is dead in-app. The nicely-built query path (auto-chunking disabled for input_type="query") is only reachable via direct API calls — the tests exercise it, but the app never does. A voyage-context-* query would go through the document branch with auto-chunking on.

Fix: forward input_type="query" from the retriever query call sites (or otherwise distinguish query vs. document at those two self.embedding(text) calls). Without it, the query handling that is the headline of this PR doesn't actually run in the application.

Minor / non-blocking:

  • _iter_batches calls self._client.tokenize([text], ...) once per text; batching a single tokenize(texts, ...) call per window would avoid the per-item overhead.
  • The contextualized tests mock tokenize, so real tokenization of voyage-context-* model names is unverified — worth a quick real-SDK sanity check that tokenize(model="voyage-context-4") resolves.
  • _embed_contextualized keeps only result.embeddings[0], correct under the one-input-one-chunk assumption; fine as long as inputs stay under chunk_size, which is the documented contract.

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.
@fzowl

fzowl commented Aug 17, 2026

Copy link
Copy Markdown
Owner Author

Housekeeping: addressed the blocking query-plumbing review.

The retriever now forwards input_type="query" from both VectorRetrieval query call sites (vector + hybrid, vectorindex.py), via a small _embed_query helper. So asymmetric Voyage models get the query prompt on searches, and the contextualized query path (auto-chunking disabled for queries) is now reachable from the app rather than only from tests.

input_type is treated as an optional embedding kwarg. A theflow node resolves to a proxy during pipeline execution, so a runtime capability check on the embedding object isn't reliable; instead every backend tolerates the kwarg — langchain-based and endpoint embeddings accept-and-ignore it, and the OpenAI backend drops it before the API call. All other backends are behaviorally unchanged (the existing AzureOpenAI retrieval test stays green).

Also: the token-aware batcher now tokenizes each window in a single tokenize call (was one call per text); the class docstring documents the asymmetric/re-index implication for indexes built under the old symmetric behavior, plus the long-input truncation contract.

Tests: added an end-to-end retrieval test asserting the query is embedded with input_type="query" (driven through the real pipeline so the node proxy is exercised) and a regular-model query-path test. Ran the Voyage embedding suite and the indexing/retrieval suite in an isolated env (py3.13 here blocks the full install via llama-hub, and langchain pins were relaxed to import the package): 10 passed. flake8 (max-line-length 88, extend-ignore E203), black and isort clean on the changed files. Remaining CI validation runs on the branch.

@fzowl

fzowl commented Aug 17, 2026

Copy link
Copy Markdown
Owner Author

VERDICT:APPROVED

Reviewed diff, commits, tests, docs. Solid work — ship it.

Correctness

  • Async bug genuinely fixed: old await self._aclient.embed(...).embeddings awaited an attribute, not the coroutine. New _aembed_* helpers await the call. Good catch.
  • Contextualized doc/query split is right: enable_auto_chunking = input_type != "query", chunk_size dropped for queries (API rejects it). Matches Voyage semantics.
  • Batch ordering preserved: batches processed sequentially, embeddings.extend(...) in order, _format_output zips input to embedding 1:1.
  • _iter_batches token boundary + oversized-text-alone + item-count cap logic correct and directly covered by tests.

Fit / breaking-change risk — low

  • input_type forwarding verified across ALL embedding backends: base run passes **kwargs to invoke; endpoint/langchain/fastembed/tei accept-and-ignore, OpenAI explicitly pops it. No backend breaks.
  • Retriever change (_embed_query) scoped to vector + hybrid query call sites. Indexing path untouched.
  • Conventions matched: Param usage, docstrings, black/isort/flake8 clean per description.

Tests / docs — good

  • 6 Voyage tests cover doc path, query path, and all three batching branches. End-to-end retrieval test drives the real theflow pipeline, not a direct method call.
  • .env.example, class docstring (incl. re-index warning for asymmetric switch), and reranker help text all updated.

Commit hygiene — clean

  • Zero AI trace (no Co-Authored-By / Generated-with), no .claude/ files. Commit messages explain the why (dependency cascade, numpy/unstructured conflict, splitter override).

Non-blocking notes

  1. Silent truncation: contextualized doc path keeps only result.embeddings[0] — a doc exceeding chunk_size (32k) loses trailing chunks silently. Documented as the pre-chunked-RAG contract, so acceptable, but it is a real behavior.
  2. Py3.13 fallback: voyageai; python_version >= '3.13' has no lower bound, so it can resolve below 0.4.0 where contextualized_embed doesn't exist — contextualized models would fail at runtime on 3.13. Forced by the numpy<2 / unstructured conflict and 3.13 isn't in the test matrix, so fine for now; worth a follow-up when unstructured allows numpy>=2.1.
  3. Minor: _iter_batches calls the sync self._client.tokenize even on the async path — blocks the event loop briefly. Negligible in practice.

None block merge.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant