Skip to content

feat(python): refresh VoyageAI models and contextualized embedding path - #1

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

feat(python): refresh VoyageAI models and contextualized embedding path#1
fzowl wants to merge 2 commits into
mainfrom
feat/voyageai-refresh

Conversation

@fzowl

@fzowl fzowl commented Aug 17, 2026

Copy link
Copy Markdown
Owner

What

Refresh the VoyageAI embedding integration against the current Voyage model catalog and correct the contextualized embedding path.

Models

  • Add text models: voyage-code-4, voyage-4-nano, voyage-3-large, voyage-code-3.
  • Add contextualized model voyage-context-4 alongside voyage-context-3.
  • Update ndims, the per-request token-limit table, docstrings, and docs to match.

Contextualized embeddings (voyage-context-*)

Each input string is embedded as its own independent document. On the document side the batch is sent as a flat list[str] with enable_auto_chunking=True and chunk_size=32000, so every input resolves to exactly one chunk and one embedding — deterministic per-input vectors and trivial result collection.

Inputs are intentionally not contextualized against one another: generic embed_many callers pass unrelated texts, so cross-input contextualization would be wrong. Previously the batch was sent as a single document's chunks (inputs=[batch]), which mixed unrelated inputs; that is fixed here.

The query/retrieval path omits enable_auto_chunking/chunk_size, because the API rejects auto-chunking when input_type is query.

Token-aware batching

The plain-text path builds batches by estimated token count against the per-model token limit (not only a fixed item count).

Why

Keep the exposed models current, make contextualized embeddings correct and deterministic for generic callers, and respect per-request token budgets.

Tests

  • Unit tests (no network) for model dimensions, contextual classification, the contextualized document path (asserts flat list[str] + auto-chunking + chunk_size=32000, one embedding per input) and the query path (asserts auto-chunking is not sent).
  • Unit tests for batching: split at the token boundary, a single oversized text still goes through alone, item-count cap respected.
  • Slow integration tests extended to voyage-context-4, voyage-code-4, voyage-3-large, and to exercise the contextual query/search path.

All non-slow unit tests pass locally (pytest python/tests/test_voyageai_embeddings.py). Ruff check/format clean.

Refresh the VoyageAI integration against the current model catalog and fix
the contextualized embedding design.

Models:
- add voyage-code-4, voyage-4-nano, voyage-3-large, voyage-code-3 text models
- add voyage-context-4 alongside voyage-context-3 (contextualized)
- update ndims, token-limit table, docstrings and docs accordingly

Contextualized embeddings (voyage-context-*):
- embed each input string as its own independent document by sending the
  batch as a flat list[str] with enable_auto_chunking=True and
  chunk_size=32000 on the document side, so every input resolves to exactly
  one chunk and one embedding. Inputs are not contextualized against each
  other, since generic embed_many callers pass unrelated texts.
- the query path omits auto-chunking (the API rejects it for query inputs).

Token-aware batching:
- retained/covered by unit tests: split at the token boundary, a single
  oversized text still goes through alone, item-count cap respected.

Tests cover both the document and query contextualized paths and the
batching logic.
@github-actions github-actions Bot added the enhancement New feature or request label Aug 17, 2026
@fzowl

fzowl commented Aug 17, 2026

Copy link
Copy Markdown
Owner Author

VERDICT:CHANGES_NEEDED

Reviewing the VoyageAI refresh (commit eee57f9c). Note: the task referenced PR lancedb#2959, which does not exist in this fork; lancedb#2959 is the merged upstream PR. This review targets the equivalent fork PR (feat/voyageai-refresh).

Core change — good

The actual source change (5 files) is correct, well-tested, and cleanly documented:

  • Contextualized document path correctly switches inputs=[batch]inputs=batch (flat list[str]) with enable_auto_chunking=True + chunk_size=32000, so each input is embedded as its own document → one vector per input. This is a genuine bug fix: the old form contextualized unrelated embed_many inputs against each other. Extracting res.embeddings[0] per result is consistent with the one-chunk-per-doc guarantee.
  • Query path (compute_query_embeddings) correctly keeps inputs=[[query]], input_type="query", and omits auto-chunking (API rejects it for query inputs). Both branches are covered by no-network unit tests asserting the exact call kwargs.
  • Model catalog, ndims, token-limit table, docstrings and docs updated consistently. _build_batches token-aware batching covered by boundary tests (split at limit, single oversize alone, item-count cap). Meets the project's "no merge without tests" bar.
  • Reranker docstring fix (OPENAI_API_KEYVOYAGE_API_KEY, drop stale rerank-english-v2.0 default) matches the actual signature.

Commit hygiene: clean — single focused commit, no Co-Authored-By, no "Generated with", no .claude/ files.

Blocking / must-fix

  1. Stale base branch. As opened, this PR's diff is ~201,522 additions across 783 files, because the branch sits on recent upstream while the PR base (fork main) is far behind. It is not reviewable or safely mergeable in this state. Sync fork main to upstream (or retarget the base) so the diff shows only the 5 VoyageAI files.
  2. Unrelated formatting churn in python/python/tests/test_embeddings_slow.py (the test_colpali / test_colpali_models assert reflows) is unrelated to VoyageAI and diverges from upstream's ruff formatting — likely to fail ruff format --check. Drop these hunks.

Minor

  • Document inputs exceeding chunk_size (32000 tokens) will auto-chunk into multiple chunks; only the first chunk embedding is kept (res.embeddings[0]), silently dropping the rest. Acceptable given the one-vector-per-row intent, but worth a docstring note or explicit truncation.
  • Token-limit table changes (e.g. voyage-context-3 32K → 120K) alter batching for existing users; fine if these match the current published API, which the author (voyageai.com) is positioned to confirm.

Core code is approvable; fix the base-branch scope (#1) and drop the unrelated formatting hunks (lancedb#2) before merge.

- Revert unrelated assert-reflow formatting churn in test_embeddings_slow.py
  (test_colpali / test_colpali_models / voyage embedding assert) so the diff
  carries only VoyageAI changes and passes ruff format --check.
- Apply project ruff formatting to test_voyageai_embeddings.py assert messages.
- Document that a contextualized input longer than the 32000-token per-chunk
  ceiling is split server-side and only the first chunk embedding is kept.
@fzowl

fzowl commented Aug 17, 2026

Copy link
Copy Markdown
Owner Author

Addressed the review feedback in 54a63f4b:

Blocking

  1. Stale base branch — synced fork main to upstream (928c3dde), which is exactly this branch's base. The PR diff is now 5 files, +231/-23 (all VoyageAI), instead of the previous ~201k-line cross-upstream diff.
  2. Unrelated formatting churn — reverted the assert reflows in test_embeddings_slow.py (test_colpali, test_colpali_models, and the voyage embedding assert) back to upstream ruff form; kept only the legitimate content changes (added voyage-3-large/voyage-code-4 params, contextual test parametrization + query-path assertion). Also ran the project ruff formatter over test_voyageai_embeddings.py, which had the same reflow style.

Minor

  • Added a docstring note that a contextualized input exceeding the 32000-token per-chunk ceiling is split server-side and only the first chunk embedding is kept.
  • Token-limit table values were cross-checked against the langchain-voyageai reference (voyage-4/voyage-context-4/voyage-4-lite etc.) and match.

Catalog audit — current docs catalog is fully covered: text (voyage-4-large/4/4-lite/4-nano/code-4 + domain + legacy), contextualized (voyage-context-4/3), multimodal (voyage-multimodal-3.5/3), rerankers (rerank-2.5/2.5-lite). Nothing missing.

Validationruff check and ruff format --check pass on all changed files; py_compile clean. The no-network unit tests are mock-based (no logic change since the prior review); local execution is pending a full maturin build (very slow on this ARM host) and is still running.

@fzowl

fzowl commented Aug 17, 2026

Copy link
Copy Markdown
Owner Author

VERDICT:APPROVED

Reviewed the full diff, commit messages, and the surrounding voyageai.py / reranker source. This is a clean, well-scoped refresh of the VoyageAI integration and I recommend merging.

Correctness

  • Contextualized document path (_get_embed_function): the switch from inputs=[batch] (one document whose chunks were unrelated inputs) to a flat inputs=batch with enable_auto_chunking=True + chunk_size=32000 is the right fix for generic embed_many callers. Returning [res.embeddings[0] for res in result.results] correctly yields one vector per input. The old code silently contextualized unrelated texts against each other — this is a genuine behavior bug being corrected.
  • Query path: compute_query_embeddings is unchanged (inputs=[[query]], no auto-chunking), which is correct since the API rejects auto-chunking for input_type="query". The two paths are consistent.
  • Token-aware batching (_build_batches): logic verified by hand against the three unit tests — split at the 120K boundary, single oversized text emitted alone, and BATCH_SIZE item cap all behave as asserted.
  • ndims() / model lists / token-limit table: additions are consistent across all four places (token table, text_embedding_models, contextual_embedding_models, ndims). Unknown models still fall back to the 120K default.
  • Reranker doc fixes: removing the bogus default "rerank-english-v2.0" (the ctor has no default — model_name is required positional) and correcting OPENAI_API_KEY -> VOYAGE_API_KEY are both accurate against the actual code.

Fit / conventions

Matches existing module style, docstrings, and the _build_batches/_get_embed_function structure. Docs page updated to match. Conventional-commit messages (feat(python): / fix(python):).

Breaking-change risk

Low. New models are additive; doc edits are cosmetic. The one real behavior change is contextualized document embedding semantics — but the previous behavior was incorrect for the generic caller, and the new behavior + the data-loss caveat (input >32000 tokens keeps only the first chunk) is clearly documented in the docstring. Acceptable.

Tests

Good coverage: no-network unit tests for dims, contextual classification, document vs. query contextual paths (asserting exact kwargs), and all three batching edge cases; slow integration extended to voyage-context-4 / voyage-code-4 / voyage-3-large and the query/search path. Could not execute locally (bindings not built in this env), but the assertions were validated statically against the code.

Commit hygiene

Clean — no Co-Authored-By, no "Generated with" trailers, no .claude/ files. Author is zoltan@voyageai.com.

Minor (non-blocking)

  1. Confirm the per-request budget for voyage-context-* is genuinely 120K — it was previously 32K in the table (that value was really the per-chunk ceiling, now correctly separated into CONTEXT_CHUNK_SIZE). Worth a sanity check against the API reference.
  2. The if input_type != "query" guard inside the contextual embed_batch is effectively dead: the document path always passes input_type="document", and the query path bypasses batching entirely via compute_query_embeddings. Harmless defensive code, and note that test_query_path_disables_auto_chunking actually exercises compute_query_embeddings, not this guard.

Neither warrants blocking. Approving.

@fzowl

fzowl commented Aug 17, 2026

Copy link
Copy Markdown
Owner Author

Upstream PR opened: lancedb#3955

This PR stays open as the working channel until the upstream PR is resolved.

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

Labels

enhancement New feature or request Python

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant