Add 0092 embed chunk-and-stitch (TEI + shared helper)

conformance.toml

@@ -883,11 +883,11 @@ note = "Retrieval-provider Cohere /v2/embed wire mapping (§8.4) -- the embed ha
 
 # Spec v0.87.0 (proposal 0092).  Embedding-mapping batch chunking, the
 # general §8 rule (consecutive <=cap chunks, one request each, vectors
-# concatenated in input order, usage summed).  Not-yet; the TEI /embed
-# over-cap fixture 038 defers with it (no embedding wire mapping is
-# shipped).
+# concatenated in input order, usage combined record-aware).
 [proposals."0092"]
-status = "not-yet"
+status = "implemented"
+since = "0.16.0"
+note = "Retrieval-provider general §8 *Batch chunking* rule for the embedding side + fixture 038 (TEI /embed chunk-and-stitch by the construction chunk_size). The shared chunk_and_stitch_embed helper in _wire.py realizes the rule mapping-agnostically: split the inputs into consecutive <=cap slices preserving order, issue one request per slice via a per-mapping closure (every field other than the chunked input list identical across slices), concatenate the per-chunk vectors in input order (so §4's one-vector-per-input + input-order invariants hold across the whole call, enforced against the stitched result), combine the per-chunk usage record-aware (sum EmbeddingUsage.input_tokens when the provider reports usage, else usage = null -- never fabricated), take response_id from the first chunk, and set raw to the single response for a single-request call or the list of per-chunk responses in request order (proposal 0096). TeiEmbeddingProvider.embed() adopts it with cap = its construction chunk_size (TEI's max-client-batch-size, default 32): an over-cap embed call chunk-and-stitches (fixture 038: chunk_size 2 over 5 inputs -> 3 /embed requests sized 2/2/1 with identical per-call params), and since TEI /embed reports no usage and no id the stitched usage is null and response_id is null; the len(input) <= chunk_size path issues a single /embed request (the pre-0092 behavior preserved). CohereEmbeddingProvider migrates its existing 96-input chunk path (0091, fixture 037) onto the same helper (cap = 96, closure POSTs /v2/embed) with identical behavior. The per-chunk vector count is validated before stitching so a positional misalignment cannot pass on a compensating total. Conformance fixtures 037 (Cohere, unchanged) + 038 (TEI) pass via the general wire-capture harness; no retrieval fixtures remain deferred."
 
 # Spec v0.88.0 (proposal 0093).  Nullable provider usage records:
 # EmbeddingResponse.usage / RerankResponse.usage become record|null and

src/openarmature/retrieval/_wire.py

@@ -10,10 +10,14 @@
 
 from __future__ import annotations
 
+from collections.abc import Awaitable, Callable
 from typing import Any, cast
 
 from openarmature.llm.errors import ProviderInvalidResponse
 
+from .provider import validate_embedding_input, validate_embedding_response
+from .response import EmbeddingResponse, EmbeddingUsage
+
 
 def normalize_base_url(base_url: str, *, guard_prefix: str) -> str:
     """Strip a trailing slash from ``base_url`` and reject a doubled version prefix.
@@ -82,6 +86,90 @@ def apply_client_side_prefix(
     return [prefix + s for s in input_strings]
 
 
+# retrieval-provider §8 *Batch chunking* (the general embed rule), shared by
+# every EmbeddingProvider whose wire enforces a per-call input cap (TEI /embed's
+# max-client-batch-size, Cohere /v2/embed's 96-input cap, ...). When len(input)
+# exceeds cap the mapping MUST split the inputs into consecutive <=cap slices
+# (preserving order), issue one request per slice with EVERY field other than the
+# chunked input list identical across slices, concatenate the per-chunk vectors
+# IN INPUT ORDER (so §4's one-vector-per-input + input-order invariants hold
+# across the whole call), and combine the per-chunk usage per §4's nullable
+# usage contract -- sum input_tokens when the provider reports usage, else
+# usage = null. response_id is the FIRST chunk's id (a single-request call uses
+# that request's id). raw is that one response for a single-request call, or the
+# LIST of per-chunk responses in request order for a chunked call (proposal
+# 0096). embed_chunk owns the per-mapping wire shaping / POST / parse and returns
+# (vectors, input_tokens, response_id, raw_body) for its slice -- so the loop /
+# stitch / validation is mapping-agnostic. When len(input) <= cap this issues a
+# single request (the single-iteration path). Valid because each input's
+# embedding is independent of the others in its batch.
+async def chunk_and_stitch_embed(
+    input_strings: list[str],
+    *,
+    model: str,
+    cap: int,
+    embed_chunk: Callable[[list[str]], Awaitable[tuple[list[list[float]], int | None, str | None, Any]]],
+) -> EmbeddingResponse:
+    """Issue one embed request per ``<= cap`` chunk and stitch the vectors.
+
+    ``embed_chunk`` sends one chunk's request and returns
+    ``(vectors, input_tokens, response_id, raw_body)`` for that chunk; the
+    per-mapping wire shaping, POST, and parse live in the closure. Returns the
+    stitched :class:`EmbeddingResponse`.
+    """
+    # cap is the provider's per-call input limit; a non-positive cap is a caller
+    # misconfiguration -- fail loudly rather than surfacing a raw range() error
+    # (cap == 0) or a misleading empty-stitched validation failure (cap < 0).
+    if cap <= 0:
+        raise ValueError(f"cap must be positive (got {cap})")
+    # Validate the input up front so an empty input raises
+    # provider_invalid_request -- the caller-side contract error -- rather than
+    # falling through to a misclassified provider_invalid_response from the
+    # empty stitched-count check. Providers already call this before the helper;
+    # this protects a direct / future caller.
+    validate_embedding_input(input_strings)
+    stitched_vectors: list[list[float]] = []
+    chunk_bodies: list[Any] = []
+    input_tokens_total: int | None = None
+    response_id: str | None = None
+    for offset in range(0, len(input_strings), cap):
+        chunk = input_strings[offset : offset + cap]
+        chunk_vectors, chunk_tokens, chunk_id, chunk_body = await embed_chunk(chunk)
+        # Per-chunk count MUST match the chunk's inputs before stitching: the
+        # stitch is positional (no index re-basing), so a chunk returning the
+        # wrong vector count would silently misalign vectors that a compensating
+        # chunk lets the stitched total pass (§4 input-order).
+        if len(chunk_vectors) != len(chunk):
+            raise ProviderInvalidResponse(
+                f"embedding response returned {len(chunk_vectors)} vectors for {len(chunk)} inputs"
+            )
+        stitched_vectors.extend(chunk_vectors)
+        chunk_bodies.append(chunk_body)
+        if offset == 0:
+            response_id = chunk_id
+        # Sum input_tokens across chunks; a chunk that omits usage does not
+        # contribute, and usage stays null only when NO chunk reports it (§4 /
+        # §8 batch-chunking step 4 -- never fabricate).
+        if chunk_tokens is not None:
+            input_tokens_total = (input_tokens_total or 0) + chunk_tokens
+    # §4 cross-impl invariants (one vector per input, uniform dimensionality)
+    # are enforced against the STITCHED result.
+    dimensions = validate_embedding_response(stitched_vectors, len(input_strings))
+    usage = EmbeddingUsage(input_tokens=input_tokens_total) if input_tokens_total is not None else None
+    # raw is the verbatim deserialized response (§4 / §8 per proposal 0096, dict
+    # | list): a single call carries that response's body; a chunked call carries
+    # the LIST of per-chunk bodies in request order.
+    raw: dict[str, Any] | list[Any] = chunk_bodies[0] if len(chunk_bodies) == 1 else chunk_bodies
+    return EmbeddingResponse(
+        vectors=stitched_vectors,
+        model=model,
+        usage=usage,
+        response_id=response_id,
+        dimensions=dimensions,
+        raw=raw,
+    )
+
+
 # §6 (0097): the rerank document-echo shape rule, shared by every RerankProvider
 # (Cohere / TEI / Jina). ScoredDocument.document carries the provider's string
 # echo verbatim when present, null otherwise -- never fabricated from the input
@@ -109,6 +197,7 @@ def document_echo(value: Any) -> str | None:
 
 __all__ = [
     "apply_client_side_prefix",
+    "chunk_and_stitch_embed",
     "client_side_prefix",
     "document_echo",
     "normalize_base_url",

src/openarmature/retrieval/providers/cohere.py

@@ -64,17 +64,15 @@
     build_rerank_event,
     build_rerank_failed_event,
 )
-from .._wire import document_echo, normalize_base_url
+from .._wire import chunk_and_stitch_embed, document_echo, normalize_base_url
 from ..provider import (
     validate_embedding_input,
-    validate_embedding_response,
     validate_rerank_input,
     validate_rerank_response,
 )
 from ..response import (
     EmbeddingResponse,
     EmbeddingRuntimeConfig,
-    EmbeddingUsage,
     RerankResponse,
     RerankRuntimeConfig,
     RerankUsage,
@@ -622,61 +620,34 @@ async def _embed_chunked(
         request_extras: dict[str, Any],
     ) -> EmbeddingResponse:
         """Issue one /v2/embed per <=96 chunk and stitch the vectors."""
+
         # THE chunk-and-stitch (§8.4 *Batch chunking (96-input cap)* / the §8
-        # general embed rule). When len(input) <= 96 this issues a single
-        # request; otherwise it splits the inputs into consecutive <=96 slices,
-        # issues one /v2/embed per slice with IDENTICAL per-call params (only
-        # texts differs), concatenates the per-chunk embeddings.float IN INPUT
-        # ORDER, sums meta.billed_units.input_tokens across chunks, and takes
-        # response_id from the FIRST chunk. This is valid because each input's
-        # embedding is independent of the others in its batch.
-        stitched_vectors: list[list[float]] = []
-        chunk_bodies: list[Any] = []
-        input_tokens_total: int | None = None
-        response_id: str | None = None
-        for offset in range(0, len(input_strings), _COHERE_EMBED_MAX_INPUTS):
-            chunk = input_strings[offset : offset + _COHERE_EMBED_MAX_INPUTS]
+        # general embed rule) via the shared chunk_and_stitch_embed helper. The
+        # per-chunk closure builds the /v2/embed body with IDENTICAL per-call
+        # params (only texts differs), POSTs, classifies the HTTP error, and
+        # parses the chunk into (vectors, input_tokens, id, body); the helper
+        # loops the consecutive <=96 slices, concatenates the vectors IN INPUT
+        # ORDER, sums input_tokens, takes response_id from the FIRST chunk, and
+        # validates the §4 invariants against the stitched result. When
+        # len(input) <= 96 this issues a single request. Valid because each
+        # input's embedding is independent of the others in its batch.
+        async def _embed_one(
+            chunk: list[str],
+        ) -> tuple[list[list[float]], int | None, str | None, dict[str, Any]]:
             body = self._build_request_body(chunk, input_type, dimensions, request_extras)
             try:
                 resp = await self._client.post("/v2/embed", json=body)
             except httpx.HTTPError as exc:
                 raise ProviderUnavailable(f"embedding request failed: {exc}") from exc
             if resp.status_code != 200:
                 raise _classify_cohere_http_error(resp)
-            chunk_vectors, chunk_tokens, chunk_id, chunk_body = self._parse_chunk(resp)
-            # Per-chunk count MUST match the chunk's inputs before stitching: the
-            # stitch is positional (no index re-basing), so a chunk returning the
-            # wrong vector count would silently misalign vectors that a
-            # compensating chunk lets the stitched total pass (§4 input-order).
-            if len(chunk_vectors) != len(chunk):
-                raise ProviderInvalidResponse(
-                    f"embedding response returned {len(chunk_vectors)} vectors for {len(chunk)} inputs"
-                )
-            stitched_vectors.extend(chunk_vectors)
-            chunk_bodies.append(chunk_body)
-            if offset == 0:
-                response_id = chunk_id
-            # Sum input_tokens across chunks; a chunk that omits usage does not
-            # contribute, and usage stays null only when NO chunk reports it
-            # (§4 / §8 batch-chunking step 4 -- never fabricate).
-            if chunk_tokens is not None:
-                input_tokens_total = (input_tokens_total or 0) + chunk_tokens
-        # §4 cross-impl invariants (one vector per input, uniform
-        # dimensionality) are enforced against the STITCHED result.
-        dimensions_out = validate_embedding_response(stitched_vectors, len(input_strings))
-        usage = EmbeddingUsage(input_tokens=input_tokens_total) if input_tokens_total is not None else None
-        # raw is the verbatim deserialized response (§4 / §8 per proposal 0096,
-        # dict | list): a single /v2/embed call carries that response's object;
-        # a chunked call carries the LIST of per-chunk objects in request order,
-        # discriminated by whether the input exceeded the 96-input cap.
-        raw: dict[str, Any] | list[Any] = chunk_bodies[0] if len(chunk_bodies) == 1 else chunk_bodies
-        return EmbeddingResponse(
-            vectors=stitched_vectors,
+            return self._parse_chunk(resp)
+
+        return await chunk_and_stitch_embed(
+            input_strings,
             model=self.model,
-            usage=usage,
-            response_id=response_id,
-            dimensions=dimensions_out,
-            raw=raw,
+            cap=_COHERE_EMBED_MAX_INPUTS,
+            embed_chunk=_embed_one,
         )
 
     def _build_request_body(