Harden 0093 chunk stitching per review

Follow-up to the adversarial and CoPilot review of the 0093 PR.

Usage accounting in chunk_and_stitch_embed is now all-or-nothing: a
total is reported only when EVERY chunk reports a figure. A chunk that
omits usage (absent or corrupt) makes the whole-call total null rather
than a partial sum, which would present a confident undercount and
diverge from the single-request path (which returns null for the same
unaccountable figure). Behaviorally invisible for the reference
providers (hosted OpenAI and Cohere report on every chunk; TEI never
reports), it only changes the unspecified mixed case an inconsistent
gateway could produce. A new unit test pins it and fails under the old
partial-sum logic.

An empty-string wire model or id from an OpenAI-compatible backend now
reads as absent (None) rather than surfacing as "", so the stitch falls
back to the bound model identifier per §4.

Documents that a chunked embed is non-atomic: a mid-call failure may
leave earlier chunks billed by the provider with no usage event, since
there is no rollback over a non-transactional wire.

Corrects the Cohere chunk-closure comment to the 5-tuple return shape
(model always null for Cohere), and notes the per-chunk warning
multiplicity is accepted.

Author: chris-colinsky

src/openarmature/retrieval/_wire.py

@@ -39,6 +39,10 @@
 # misbehaving on a field it did populate, so it logs at WARNING before returning
 # None -- otherwise a corrupt count is indistinguishable from the healthy absent
 # case and usage silently goes missing. ``field`` names the figure in the log.
+# Under batch chunking this fires once per chunk, so a persistently-corrupt
+# gateway logs N warnings for one embed() call; that per-chunk multiplicity is
+# accepted -- deduping would move the warning out of this single-request-shaped
+# reader and lose it for the non-chunked providers (Cohere / Jina rerank).
 def nonneg_int(value: Any, *, field: str = "usage figure") -> int | None:
     """Return ``value`` when it is a non-negative int, else ``None``.
 
@@ -135,8 +139,9 @@ def apply_client_side_prefix(
 # 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
+# usage contract ALL-OR-NOTHING -- sum input_tokens only when EVERY chunk reports
+# a figure, else usage = null (a mixed report is an unaccountable total, never a
+# partial sum). 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
@@ -166,6 +171,12 @@ async def chunk_and_stitch_embed(
     chunk; the per-mapping wire shaping, POST, and parse live in the closure.
     ``model`` is the bound identifier, reported when the chunk bodies carry no
     model of their own. Returns the stitched :class:`EmbeddingResponse`.
+
+    A chunked call issues its per-chunk requests sequentially and is NOT atomic:
+    if a later chunk fails, the earlier chunks have already been sent (and billed
+    by the provider) but their vectors are discarded and the call raises. There
+    is no rollback over a non-transactional wire, and the failure event carries
+    no usage, so a mid-call failure may leave earlier chunks billed unobservably.
     """
     # 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
@@ -180,7 +191,8 @@ async def chunk_and_stitch_embed(
     validate_embedding_input(input_strings)
     stitched_vectors: list[list[float]] = []
     chunk_bodies: list[Any] = []
-    input_tokens_total: int | None = None
+    input_tokens_total = 0
+    usage_fully_reported = True
     response_id: str | None = None
     response_model: str | None = None
     for offset in range(0, len(input_strings), cap):
@@ -202,15 +214,23 @@ async def chunk_and_stitch_embed(
         if offset == 0:
             response_id = chunk_id
             response_model = chunk_model
-        # 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
+        # Sum input_tokens ALL-OR-NOTHING: a total is reported only when EVERY
+        # chunk reported one. Any chunk that omits usage (a None figure, incl. an
+        # absent OR a corrupt one via nonneg_int) makes the whole-call total
+        # unaccountable -- so usage = null rather than a partial sum that would
+        # present a confident undercount and diverge from the single-request
+        # path (which returns null for the same unaccountable figure). All report
+        # (hosted OpenAI / Cohere) -> sum; none report (TEI) -> null; mixed (an
+        # inconsistent gateway) -> null (§4 / §8 batch-chunking step 4 -- never
+        # fabricate a figure you cannot fully account for).
+        if chunk_tokens is None:
+            usage_fully_reported = False
+        else:
+            input_tokens_total += 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
+    usage = EmbeddingUsage(input_tokens=input_tokens_total) if usage_fully_reported 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.

src/openarmature/retrieval/providers/cohere.py

@@ -616,9 +616,11 @@ async def _embed_chunked(
         # 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
+        # parses the chunk into (vectors, input_tokens, id, model, body) -- model
+        # is always None (the /v2/embed envelope carries no model), so the stitch
+        # falls back to the bound model; the helper loops the consecutive <=96
+        # slices, concatenates the vectors IN INPUT ORDER, sums input_tokens,
+        # takes the response identity (id + model) 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.

src/openarmature/retrieval/providers/openai.py

@@ -446,13 +446,17 @@ def _parse_chunk(
             if isinstance(usage_block, dict)
             else None
         )
+        # An empty-string id / model carries no provider identifier (§4: the
+        # reported identifier is used only when the body carries one), so treat
+        # "" as absent -> None. For the model that lets the stitch fall back to
+        # the bound identifier rather than surfacing "".
         response_id = body.get("id")
         model = body.get("model")
         return (
             vectors,
             input_tokens,
-            response_id if isinstance(response_id, str) else None,
-            model if isinstance(model, str) else None,
+            response_id if isinstance(response_id, str) and response_id else None,
+            model if isinstance(model, str) and model else None,
             body,
         )
 

tests/unit/test_retrieval_provider.py

@@ -445,6 +445,38 @@ def handler(req: httpx.Request) -> httpx.Response:
     assert isinstance(response.raw, dict)
 
 
+async def test_openai_embed_chunked_usage_is_all_or_nothing() -> None:
+    # Usage is summed ONLY when every chunk reports it. A 2049-input call
+    # (2 chunks) where chunk 1 reports prompt_tokens but chunk 2 omits the usage
+    # block yields usage = null for the WHOLE call, not a confident partial sum:
+    # the total is unaccountable, and a partial sum would both undercount and
+    # diverge from the single-request path (which returns null for the same
+    # unaccountable figure). Guards against the chunked/unchunked inconsistency.
+    inputs = [f"doc-{i:04d}" for i in range(2049)]
+    chunk_a = [[i / 10000, 0.5] for i in range(2048)]
+    body_a = _openai_embed_body(
+        data=[{"object": "embedding", "index": i, "embedding": v} for i, v in enumerate(chunk_a)],
+        prompt_tokens=4096,
+    )
+    body_b = _openai_embed_body(data=[{"object": "embedding", "index": 0, "embedding": [0.2048, 0.5]}])
+    del body_b["usage"]  # chunk 2 reports no usage block at all
+    responses = iter([body_a, body_b])
+
+    def handler(req: httpx.Request) -> httpx.Response:
+        assert req.url.path == "/v1/embeddings"
+        return httpx.Response(200, json=next(responses))
+
+    provider = _openai_embed_provider(handler)
+    try:
+        response = await provider.embed(inputs)
+    finally:
+        await provider.aclose()
+
+    assert len(response.vectors) == 2049
+    # Chunk 2 reported nothing, so the whole-call total is unaccountable -> null.
+    assert response.usage is None
+
+
 # --- nonneg_int: warn on a present-but-corrupt usage figure (proposal 0093) --
 #
 # The shared usage-figure reader distinguishes "not reported" (None -- the