Skip to content

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

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

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

Conversation

@fzowl

@fzowl fzowl commented Aug 17, 2026

Copy link
Copy Markdown
Contributor

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).
…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
Contributor Author

CI was failing because the lockfile pinned voyageai==0.3.2, whose Client predates the contextualized_embed API — the two contextualized-embedding tests errored at patch time and the feature could not work at runtime. I've pushed a fix:

  • Require voyageai>=0.4.0 (the version that adds the flat-list contextualized_embed with enable_auto_chunking/chunk_size) and relocked to 0.5.0.
  • On Python 3.13 voyageai>=0.4.0 wants numpy>=2.1, which conflicts with unstructured<0.16 (numpy<2), so a marker keeps the older voyageai there; the test matrix (3.10/3.11) gets 0.5.0.
  • voyageai declares langchain-text-splitters>=0.3.8, which would otherwise 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 invoked on the server-side auto-chunking path, so a uv override-dependencies keeps the splitter on the 0.2 line and leaves the langchain stack untouched.

Only voyageai changes version in uv.lock; the rest of that diff is marker re-serialization from the 3.13 split. Unit tests on 3.10 and 3.11 are green.

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