Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 6 additions & 6 deletions conformance.toml
Original file line number Diff line number Diff line change
Expand Up @@ -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."
122 changes: 99 additions & 23 deletions src/openarmature/retrieval/_wire.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand All @@ -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.
Expand Down Expand Up @@ -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
Expand All @@ -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
Expand All @@ -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,
Expand Down Expand Up @@ -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",
]
Loading