Add 0092 embed chunk-and-stitch

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.

Author: chris-colinsky

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_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,79 @@ 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`.
+    """
+    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 +186,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(

src/openarmature/retrieval/providers/tei.py

@@ -55,10 +55,9 @@
     build_rerank_event,
     build_rerank_failed_event,
 )
-from .._wire import apply_client_side_prefix, document_echo
+from .._wire import apply_client_side_prefix, chunk_and_stitch_embed, document_echo
 from ..provider import (
     validate_embedding_input,
-    validate_embedding_response,
     validate_rerank_input,
     validate_rerank_response,
 )
@@ -178,9 +177,9 @@ def __init__(
         if chunk_size <= 0:
             raise ValueError(f"chunk_size must be positive (got {chunk_size})")
         # chunk_size is TEI's max-client-batch-size (default 32); it bounds the
-        # /embed batch per the §8 batch-chunking rule. Client-side embed chunk-
-        # and-stitch is proposal 0092, out of scope here -- chunk_size is bound
-        # for construction parity so an operator can already pin it.
+        # /embed batch per the §8 batch-chunking rule (proposal 0092) -- an
+        # over-cap embed call chunk-and-stitches over consecutive <=chunk_size
+        # slices.
         self._chunk_size = chunk_size
         # input_type -> TEI prompt_name map (server-side prompts) with an
         # optional client-side prefix fallback (§8.1 input_type realization).
@@ -241,14 +240,7 @@ async def embed(
         adapter_start = time.perf_counter()
         try:
             validate_embedding_input(input_strings)
-            body = self._build_request_body(input_strings, input_type, dimensions, request_extras)
-            try:
-                resp = await self._client.post("/embed", json=body)
-            except httpx.HTTPError as exc:
-                raise ProviderUnavailable(f"embedding request failed: {exc}") from exc
-            if resp.status_code != 200:
-                raise _classify_tei_http_error(resp)
-            response = self._parse_response(resp, input_strings)
+            response = await self._embed_chunked(input_strings, input_type, dimensions, request_extras)
         except LlmProviderError as exc:
             latency_ms_failed = (time.perf_counter() - adapter_start) * 1000.0
             if dispatch is not None:
@@ -331,18 +323,64 @@ def _build_request_body(
             body["dimensions"] = dimensions
         return body
 
-    def _parse_response(
+    async def _embed_chunked(
         self,
-        resp: httpx.Response,
         input_strings: list[str],
+        input_type: str | None,
+        dimensions: int | None,
+        request_extras: dict[str, Any],
     ) -> EmbeddingResponse:
-        """Parse TEI's bare vector-array response into an EmbeddingResponse."""
+        """Issue one /embed per <=chunk_size chunk and stitch the vectors."""
+
+        # THE chunk-and-stitch (§8.1 /embed max-client-batch-size cap / the §8
+        # general embed rule, proposal 0092) via the shared chunk_and_stitch_embed
+        # helper. When len(input) <= chunk_size this issues a single request (the
+        # helper's single-iteration path -- preserving the pre-0092 one-request
+        # behavior); otherwise it splits the inputs into consecutive <=chunk_size
+        # slices with IDENTICAL per-call params (the input_type realization --
+        # prompt_name or client-side prefix -- is identical per chunk; only inputs
+        # differ), concatenates the vectors IN INPUT ORDER, and stitches usage /
+        # response_id. TEI /embed reports NO usage and NO id, so the per-chunk
+        # tokens + id slots are None: the stitched usage is null (no chunk reports
+        # tokens) and response_id is null (§8 step 4 / §4). 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]], None, None, list[Any]]:
+            body = self._build_request_body(chunk, input_type, dimensions, request_extras)
+            try:
+                resp = await self._client.post("/embed", json=body)
+            except httpx.HTTPError as exc:
+                raise ProviderUnavailable(f"embedding request failed: {exc}") from exc
+            if resp.status_code != 200:
+                raise _classify_tei_http_error(resp)
+            return self._parse_chunk(resp)
+
+        return await chunk_and_stitch_embed(
+            input_strings,
+            model=self.model,
+            cap=self._chunk_size,
+            embed_chunk=_embed_one,
+        )
+
+    def _parse_chunk(
+        self,
+        resp: httpx.Response,
+    ) -> tuple[list[list[float]], None, None, list[Any]]:
+        """Parse one TEI /embed chunk into (vectors, None, None, raw)."""
         # TEI /embed returns a bare JSON array of vector arrays in input order
         # ([[float, ...], ...]) -- no envelope, no id, no usage object. The
         # vectors are already in input order (position == input index), so no
         # index-keyed reordering is needed (unlike the OpenAI data[] shape).
-        # Validates the §4 invariants (count + consistent dimensionality) via
-        # validate_embedding_response, and rejects non-numeric vector values.
+        # Rejects non-numeric vector values. TEI reports no usage and no
+        # response id, so the tokens + id slots are None; the shared
+        # chunk_and_stitch_embed helper enforces the §4 count / dimensionality
+        # invariants against the stitched result and (since no chunk reports
+        # tokens / an id) produces usage = null + response_id = null (MUST NOT
+        # fabricate). raw is the verbatim deserialized chunk response: TEI's bare
+        # vector array carried as a list, not wrapped (§4 / §8.1 per proposal
+        # 0096 -- raw is dict | list, the top-level shape the provider returned;
+        # a chunked call stitches to the LIST of these per-chunk arrays).
         try:
             body_raw = resp.json()
         except ValueError as exc:
@@ -362,19 +400,7 @@ def _parse_response(
             if not all(isinstance(x, (int, float)) and not isinstance(x, bool) for x in values):
                 raise ProviderInvalidResponse("TEI /embed response has a non-numeric vector value")
             vectors.append([float(x) for x in values])
-        dimensions = validate_embedding_response(vectors, len(input_strings))
-        # TEI returns no usage object, so usage = null (§4 -- MUST NOT
-        # fabricate). raw is the verbatim deserialized response: TEI's bare
-        # vector array carried as a list, not wrapped (§4 / §8.1 per proposal
-        # 0096 -- raw is dict | list, the top-level shape the provider returned).
-        return EmbeddingResponse(
-            vectors=vectors,
-            model=self.model,
-            usage=None,
-            response_id=None,
-            dimensions=dimensions,
-            raw=rows,
-        )
+        return vectors, None, None, rows
 
 
 class TeiRerankProvider: