Adopt 0093 nullable provider usage records (#218)

* Adopt 0093 nullable provider usage records

Proposal 0093 gives OA one model for "the provider reported no usage"
across both retrieval responses and both typed events: no usage yields
usage = null, never a fabricated record, zero, or estimate.

The response types and both observers were already correct. The gap was
OpenAIEmbeddingProvider, which recorded EmbeddingUsage(input_tokens=0)
when the usage block was absent or malformed. A zero asserts "zero
tokens billed", a claim the provider never made. It now surfaces
usage = null; a reported zero still yields a record carrying 0.

Reviewing that diff surfaced a second, unrelated defect. Proposal 0092
requires every capped embedding mapping to chunk-and-stitch and names
OpenAI's 2048-input cap, but the 0092 implementation adopted the shared
helper for TEI and Cohere only. An over-2048 embed call sent one
over-cap request the wire rejects, breaking the arbitrary-length input
contract. OpenAI now chunks at 2048. Conformance stayed green because
the spec ships no over-cap fixture for that mapping; tracked. The shared
helper now takes the response identity (id and model) from the first
chunk, so OpenAI keeps reporting the provider-returned model while the
wires that carry none fall back to the bound identifier.

A missing usage figure and a corrupt one were both swallowed silently,
so a gateway sending garbage looked identical to a healthy no-usage
call. The shared nonneg_int reader now logs a warning when a figure is
present but not a non-negative int, while an absent figure stays quiet.
nonneg_int lifts into _wire.py and is adopted by OpenAI, Cohere, and
Jina.

Un-defers observability fixtures 139 and 140. Fixture 143 stays
deferred behind the unimplemented section 11 embedding-metrics path,
not behind anything in 0093.

* 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.

* Correct 0092 note to all-or-nothing usage rule

The conformance note described the pre-hardening per-chunk usage
semantics (sum when reported, else null). chunk_and_stitch_embed now
combines usage all-or-nothing: it sums only when every chunk reports a
figure, and a mixed report yields usage = null rather than a partial
sum. Bring the note in line with the code.

* Close provider in try/finally in usage tests

The nullable-usage unit tests closed the provider after the assertions,
so a failing assert would skip aclose() and leak the httpx client.
Wrap embed() in try/finally to match the sibling tests.

Author: chris-colinsky

conformance.toml

@@ -887,15 +887,15 @@ note = "Retrieval-provider Cohere /v2/embed wire mapping (§8.4) -- the embed ha
 [proposals."0092"]
 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."
+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 all-or-nothing (sum EmbeddingUsage.input_tokens only when EVERY chunk reports a figure, else usage = null -- a mixed report is an unaccountable total, never a partial sum or a fabricated figure), 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. OpenAIEmbeddingProvider adopts it at §8.3's fixed 2048-input cap (closure POSTs /v1/embeddings): an over-2048 embed call now chunk-and-stitches instead of sending one over-cap request the wire rejects. OpenAI's summed-token ceiling is NOT addressed by the count-based rule -- an over-token chunk still fails loud as provider_invalid_request. Jina is the cap-free mapping the rule exempts (server-side batching), so it does not chunk. The per-chunk vector count is validated before stitching so a positional misalignment cannot pass on a compensating total, and the response identity (response_id + response_model) is taken from the FIRST chunk -- the model falls back to the bound identifier for the wires that carry none (TEI's bare array, Cohere /v2/embed), while OpenAI reports the body's model per §4. Conformance fixtures 037 (Cohere, unchanged) + 038 (TEI) pass via the general wire-capture harness; no retrieval fixtures remain deferred. The spec ships NO §8.3 over-cap fixture, so the OpenAI chunking is covered by unit tests only (2049 inputs -> 2 requests sized 2048/1) -- flagged to spec as a fixture-set gap."
 
 # Spec v0.88.0 (proposal 0093).  Nullable provider usage records:
 # EmbeddingResponse.usage / RerankResponse.usage become record|null and
 # the embedding usage emission (OTel §5.5.8 gen_ai.usage.input_tokens,
-# Langfuse §8.4.5 usageDetails.input) is conditional.  Partial: the
-# response-type widening ships; the no-usage observability fixtures
-# 139-143 stay deferred as the 0093-last PR.
+# Langfuse §8.4.5 usageDetails.input) is conditional.  One uniform model
+# for "the provider reported no usage" across both responses and both
+# typed events: no usage => usage = null, never fabricated.
 [proposals."0093"]
-status = "partial"
+status = "implemented"
 since = "0.16.0"
-note = "The response-type nullable widening ships: RerankResponse.usage is record | null from the start (0060a), and EmbeddingResponse.usage is widened to EmbeddingUsage | null here with 0077 (the OpenAIEmbeddingProvider still builds an EmbeddingUsage record; the TeiEmbeddingProvider surfaces usage = null since TEI /embed returns no usage object, never fabricated -- §4). Both bundled observers already emit the embedding usage conditionally (OTel §5.5.8 gen_ai.usage.input_tokens guarded by `event.usage is not None`; Langfuse §8.4.5 usageDetails.input built only when event.usage is not None), landed with the embedding observability in 0059b, so a null-usage EmbeddingEvent renders without a usage attribute rather than crashing. partial because the no-usage OBSERVABILITY fixtures (139 / 140 / 143, asserting the conditional-emission contract end-to-end) stay deferred as the 0093-last PR."
+note = "One uniform model across all four surfaces: a provider that reports no usage yields usage = null, never a fabricated record / zero / client-side estimate (§4 / §6). The response types were already widened (RerankResponse.usage is record | null from 0060a; EmbeddingResponse.usage from 0077), and both bundled observers already emitted the usage conditionally (OTel §5.5.8 gen_ai.usage.input_tokens guarded by `event.usage is not None`; Langfuse §8.4.5 usageDetails.input built only when a record is present), landed with the embedding observability in 0059b. Completed here: the OpenAIEmbeddingProvider no longer FABRICATES an EmbeddingUsage(input_tokens=0) when the OpenAI usage block is absent or malformed -- it surfaces usage = null (a *reported* zero still yields a record carrying 0; a malformed figure reads as not-reported rather than failing the call, since the vectors are sound and usage is secondary accounting -- the posture the Cohere and Jina mappings already took, and deliberately unlike the §6 document echo of 0097, where the corrupt field IS the payload). The other three embedding mappings were already conformant: TEI /embed and /rerank surface usage = null (§8.1 -- neither endpoint returns a usage object), and Cohere / Jina return null when the body reports no figure. The shared nonneg_int helper (one figure out of a vendor usage block: a non-negative int is reported, anything else is not) lifts into _wire.py and is adopted by all three of Cohere, Jina, and OpenAI. Conformance: the no-usage rendering fixtures 139 (OTel embedding span omits gen_ai.usage.input_tokens) + 140 (Langfuse embedding usageDetails carries no `input` key) are un-deferred and pass, joining the rerank pair 141 / 142 (wired with 0060b); retrieval fixtures 017 / 038 already assert usage: null for TEI /embed. Fixture 143 (embedding no-usage token METRIC not observed) stays deferred, and NOT for a 0093 reason: the OTelObserver records §11 GenAI metrics from the LLM per-attempt event only, so the embedding metric path does not exist yet -- 143 defers behind the same gate as 089, whose implementation (proposal 0067) realizes it. 0093 itself changes nothing in §11 or graph-engine §6 (both were already null-usage-aware), so no part of THIS proposal is left unimplemented by that carve-out; cf. 0060, marked implemented with the sibling metric fixture 109 deferred on the same grounds."

src/openarmature/retrieval/_wire.py

@@ -2,14 +2,15 @@
 
 Pure, stateless primitives reused across the vendor wire mappings:
 request-body shaping (the client-side ``input_type`` prefix), endpoint
-normalization (``base_url``), and response parsing (the rerank document echo).
-Types live in ``response.py``, the observability event builders in
-``_events.py``, and the provider protocols in ``provider.py``; only the
-mapping helpers belong here.
+normalization (``base_url``), batch chunking, and response parsing (the rerank
+document echo, usage figures). Types live in ``response.py``, the observability
+event builders in ``_events.py``, and the provider protocols in ``provider.py``;
+only the mapping helpers belong here.
 """
 
 from __future__ import annotations
 
+import logging
 from collections.abc import Awaitable, Callable
 from typing import Any, cast
 
@@ -18,6 +19,50 @@
 from .provider import validate_embedding_input, validate_embedding_response
 from .response import EmbeddingResponse, EmbeddingUsage
 
+_log = logging.getLogger(__name__)
+
+
+# retrieval-provider §4 / §6 (proposal 0093): a usage figure is reported or it
+# is not -- a mapping MUST NOT fabricate one (an empty record, a zero, or a
+# client-side estimate). This reads ONE figure out of a vendor's usage block:
+# a non-negative int is the reported count, anything else is "not reported"
+# -> None. A malformed figure does NOT fail the call: the embed / rerank itself
+# succeeded and its payload is sound, so a corrupt secondary accounting field
+# surfaces as unknown rather than sinking the whole response (unlike the §6
+# document echo, where the corrupt field IS the payload -- proposal 0097).
+# Callers fold the Nones into a record-or-null per their vendor's shape.
+#
+# Two "not reported" cases with different signals: an ABSENT figure (None -- the
+# key was missing or the value was JSON null) is the ordinary no-usage case
+# (e.g. every TEI call) and returns quietly; a PRESENT-but-unusable figure (a
+# string, a negative, a bool, a float) means the provider / gateway is
+# 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``.
+
+    ``None`` (the figure was not reported) returns quietly; any other
+    non-conforming value is logged at WARNING as a misbehaving provider before
+    returning ``None``. ``field`` names the figure in that log line.
+    """
+    # bool is an int subclass, so exclude it explicitly.
+    if isinstance(value, int) and not isinstance(value, bool) and value >= 0:
+        return value
+    if value is not None:
+        _log.warning(
+            "retrieval provider reported a %s that is not a non-negative int "
+            "(got %s %.80r); recording usage as unknown",
+            field,
+            type(value).__name__,
+            value,
+        )
+    return None
+
 
 def normalize_base_url(base_url: str, *, guard_prefix: str) -> str:
     """Strip a trailing slash from ``base_url`` and reject a doubled version prefix.
@@ -94,28 +139,44 @@ 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
-# (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.
+# (vectors, input_tokens, response_id, response_model, raw_body) for its slice --
+# so the loop / stitch / validation is mapping-agnostic. The RESPONSE IDENTITY
+# (response_id + response_model) is the FIRST chunk's, since every chunk is the
+# same call with the same params; a mapping whose wire carries no id / model
+# (TEI's bare array, Cohere /v2/embed) reports None and the bound model is the
+# fallback (§4: the response model is the provider-returned identifier, which may
+# be more specific than the bound one, when the body carries one). 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]]],
+    embed_chunk: Callable[
+        [list[str]],
+        Awaitable[tuple[list[list[float]], int | None, str | 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`.
+    ``(vectors, input_tokens, response_id, response_model, raw_body)`` for that
+    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
@@ -130,11 +191,13 @@ 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):
         chunk = input_strings[offset : offset + cap]
-        chunk_vectors, chunk_tokens, chunk_id, chunk_body = await embed_chunk(chunk)
+        chunk_vectors, chunk_tokens, chunk_id, chunk_model, 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
@@ -145,24 +208,36 @@ async def chunk_and_stitch_embed(
             )
         stitched_vectors.extend(chunk_vectors)
         chunk_bodies.append(chunk_body)
+        # The response identity is the FIRST chunk's (every chunk is the same
+        # call with the same params -- §8 batch chunking step 4 for the id, and
+        # the same reasoning for the model).
         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
+            response_model = chunk_model
+        # 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.
     raw: dict[str, Any] | list[Any] = chunk_bodies[0] if len(chunk_bodies) == 1 else chunk_bodies
     return EmbeddingResponse(
         vectors=stitched_vectors,
-        model=model,
+        model=response_model if response_model is not None else model,
         usage=usage,
         response_id=response_id,
         dimensions=dimensions,
@@ -200,5 +275,6 @@ def document_echo(value: Any) -> str | None:
     "chunk_and_stitch_embed",
     "client_side_prefix",
     "document_echo",
+    "nonneg_int",
     "normalize_base_url",
 ]