Add 0092 embed chunk-and-stitch (TEI + shared helper)#217
Merged
Conversation
Implement the general section 8 embed batch-chunking rule and TEI /embed chunk-and-stitch (fixture 038). TeiEmbeddingProvider.embed() now chunks by its chunk_size (TEI's max-client-batch-size) over consecutive slices, concatenates vectors in input order, and produces usage null / response_id null (TEI reports neither). Lift the shared chunk_and_stitch_embed helper into retrieval/_wire.py: the loop, per-chunk count validation, input-order concat, None-aware input_tokens sum, first-chunk response_id, and section 4 validation, driven by a per-mapping embed_chunk closure. TEI /embed and Cohere /v2/embed both adopt it; Cohere's own chunk path migrates onto it unchanged. Un-defer fixture 038 and flip conformance.toml.
There was a problem hiding this comment.
Pull request overview
Implements retrieval-provider §8 embedding Batch chunking (proposal 0092) and applies it to TEI /embed (fixture 038) while refactoring Cohere /v2/embed to share the same orchestration helper.
Changes:
- Added a shared
chunk_and_stitch_embed()helper inretrieval/_wire.pyto chunk requests, stitch vectors in input order, combine nullable usage, and validate stitched invariants. - Updated
TeiEmbeddingProvider.embed()to chunk by constructionchunk_sizeand stitch responses (preserving the single-request path when within the cap). - Migrated Cohere embed chunking onto the shared helper and enabled conformance fixture 038 (no deferred retrieval fixtures remain).
Reviewed changes
Copilot reviewed 7 out of 7 changed files in this pull request and generated 1 comment.
Show a summary per file
| File | Description |
|---|---|
src/openarmature/retrieval/_wire.py |
Adds shared embedding chunk-and-stitch helper implementing the §8 rule. |
src/openarmature/retrieval/providers/tei.py |
Routes TEI /embed through shared chunking helper with cap = chunk_size. |
src/openarmature/retrieval/providers/cohere.py |
Refactors Cohere /v2/embed chunking to use the shared helper. |
tests/unit/test_retrieval_provider.py |
Adds unit tests covering TEI embed chunking/stitching behaviors and edge cases. |
tests/integration/test_tei_wire.py |
Adds live integration test exercising TEI embed chunk-and-stitch with small chunk_size. |
tests/conformance/test_retrieval_provider.py |
Removes deferral for fixture 038 (retrieval conformance fully enabled). |
conformance.toml |
Marks proposal 0092 as implemented and documents the behavior. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
The shared helper assumed cap was positive; a cap <= 0 (a caller misconfiguration) surfaced as a raw range() error (0) or a misleading empty-stitch validation failure (< 0). Raise ValueError eagerly with a clear message, mirroring the chunk_size guard TeiEmbeddingProvider already applies at construction.
An empty input reaching the shared helper fell through to
validate_embedding_response and raised provider_invalid_response
("provider returned no vectors"), misclassifying a caller-side invalid
request as an invalid response. Call validate_embedding_input up front so
an empty (or non-string) input raises provider_invalid_request. Providers
already validate before the helper; this protects a direct/future caller.
The up-front validate_embedding_input() comment claimed the guard covers empty and non-string inputs, but the validator only rejects an empty list. Drop the non-string claim so the comment matches the code.
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.
Implements the general retrieval-provider §8 Batch chunking rule and its TEI
/embedrealization (fixture 038), and folds the shared chunk-and-stitch orchestration into_wire.pyso TEI and Cohere embed share one implementation.The rule
When an embedding provider enforces a per-call input cap and the caller exceeds it, the mapping splits into consecutive
≤ capslices (preserving order), issues one request per slice with every field except the chunked input identical, concatenates the per-chunk vectors in input order, sumsEmbeddingUsage.input_tokenswhen the provider reports usage (elseusage = null), and takesresponse_idfrom the first chunk.Changes
chunk_and_stitch_embedhelper inretrieval/_wire.py— owns the loop, the per-chunk vector-count validation, the input-order concat, the None-awareinput_tokenssum, the first-chunkresponse_id, therawdict-or-list discriminator (0096), and the §4 validation on the stitched result. A per-mappingembed_chunkclosure owns the wire shaping / POST / parse and returns(vectors, input_tokens, response_id, raw_body)per slice./embednow chunks by itschunk_size(the existing construction param — TEI'smax-client-batch-size, previously bound but unused for embed). TEI reports no usage and no id, so the stitch yieldsusage = nulland a nullresponse_id. The single-request path (len(input) ≤ chunk_size) is preserved./v2/embedmigrates its own chunk-and-stitch onto the shared helper — a pure refactor (fixture 037 and its tests are unchanged).Tests
Un-defers fixture 038 (retrieval conformance now has no deferred fixtures); 4 TEI embed chunk unit tests (chunk boundaries, input-order stitch, usage/id null, per-chunk-count mismatch, single-request path); and an env-gated live TEI embed chunk test, verified against the local endpoint.