Skip to content

Commit 41c2fcd

Browse files
Add 0079 OpenAI-compatible embeddings wire mapping
Add the retrieval-provider 0079 OpenAI-compatible embeddings wire mapping (section 8.3) over the existing OpenAIEmbeddingProvider. input_type flows into EmbeddingEvent.request_params but is a wire no-op on OpenAI's symmetric /v1/embeddings; for an asymmetric model behind a compatible endpoint, a bound client-side query_prefix / document_prefix prepends to the input per input_type. dimensions maps to the wire dimensions field, and base_url now defaults to https://api.openai.com per the spec (it was previously required). Single-source the client-side prefix rule and its application in retrieval/_events.py (client_side_prefix + apply_client_side_prefix), adopted by both OpenAI and TEI. Un-defer conformance fixtures 023-027, read an openai_embedding_provider block in the harness, and assert the request origin against the bound base_url so the override fixture verifies routing. Env-gated live tests cover embed, dimensions, and Matryoshka truncation.
1 parent 4d16551 commit 41c2fcd

7 files changed

Lines changed: 474 additions & 61 deletions

File tree

conformance.toml

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -786,10 +786,12 @@ since = "0.16.0"
786786
note = "Retrieval-provider Jina hosted wire mapping (§8 / §8.2) -- a §8 mapping in the family established by 0077. Two bundled providers ship, distinct instances (one model each) sharing the hosted endpoint: base_url defaults to https://api.jina.ai (origin only; the provider appends /v1/embeddings / /v1/rerank), and the required api_key is sent as Authorization: Bearer <key> (§8.2 *Construction*). JinaEmbeddingProvider posts POST {base_url}/v1/embeddings with {model, input, task?, dimensions?, truncate: false} -- input_type is realized as Jina's native `task` per the closed set (query -> retrieval.query, document -> retrieval.passage); absent input_type omits `task` (the symmetric default), and an unrecognized input_type is a pre-send provider_invalid_request (§7, no request issued). dimensions -> Jina's dimensions (Matryoshka) when set; truncate: false is sent explicitly (fail-loud). It parses the object envelope {model, usage: {total_tokens}, data: [{index, embedding}]}, orders the vectors by data[i].index into input order, and maps usage.total_tokens -> EmbeddingUsage.input_tokens (Jina reports usage, unlike TEI's null -- a record when the body surfaces it, never fabricated). JinaRerankProvider posts POST {base_url}/v1/rerank with {model, query, documents, top_n?, return_documents, truncation: false} -- documents maps directly onto the string array (no per-document wrapping), top_n <- top_k (omitted when None), and return_documents is sent EXPLICITLY (Jina's wire default is true but OA's is False, §2, so the mapping sends the resolved OA value). truncation: false is sent explicitly (fail-loud). It parses {model, usage: {total_tokens}, results: [{index, relevance_score, document?}]} (Cohere-shaped relevance_score), reads the document echo only when present (never auto-filled), sorts by relevance_score descending, enforces the §6 invariants (valid index, no duplicate, len(results) <= top_k), and maps usage.total_tokens -> RerankUsage.input_tokens with search_units null (Jina meters by tokens). raw is the verbatim response dict on both surfaces. Errors map per §8.2: 401/403 -> provider_authentication, 429 -> provider_rate_limit (NOT provider_unavailable), 404 -> provider_invalid_model, 422 (over-length / malformed) -> provider_invalid_request, else -> provider_unavailable. Jina enforces no per-call embed cap (server-side batching), so there is no client-side chunk-and-stitch. Both providers dispatch the typed EmbeddingEvent / RerankEvent (graph-engine §6) per call, mutually exclusive with the failure variant; input_type flows into EmbeddingEvent.request_params with absence-is-meaningful semantics, and RerankEvent.output_results is populated unconditionally on success (proposal 0089). genai_system = 'jina'. Conformance fixtures 018-022 pass, driven through the general wire-capture harness (0077) extended with an expected_wire_headers subset assertion locking the Authorization: Bearer <api_key> header on the outbound request."
787787

788788
# Spec v0.74.0 (proposal 0079). Retrieval-provider OpenAI-compatible
789-
# embeddings wire mapping (§8.3). Not-yet; retrieval-provider fixtures
790-
# 023-027 defer with it.
789+
# embeddings wire mapping (§8.3), in the §8 mapping family established by
790+
# 0077.
791791
[proposals."0079"]
792-
status = "not-yet"
792+
status = "implemented"
793+
since = "0.16.0"
794+
note = "Retrieval-provider OpenAI-compatible embeddings wire mapping (§8 / §8.3) -- a §8 mapping in the family established by 0077, formalized over the OpenAIEmbeddingProvider that shipped with the embedding capability (0059). base_url defaults to https://api.openai.com (origin only; the provider appends /v1/embeddings) and is overridable for any OpenAI-compatible backend (vLLM / LocalAI / TEI's OpenAI surface) -- gen_ai.system 'openai' names the WIRE SURFACE, not the backing deployment (fixture 025); the api_key is sent as Authorization: Bearer <key> (§8.3 *Construction*). OpenAIEmbeddingProvider posts POST {base_url}/v1/embeddings with {model, input, dimensions?} -- input is always the array form (§3), and dimensions -> the wire dimensions field (Matryoshka) when set. input_type is a NO-OP on the base wire: the OpenAI /v1/embeddings wire has no query/document field, so an absent input_type is the correct symmetric default and the mapping does not error on a supplied one -- the wire body is byte-identical to the no-input_type request, input sent verbatim (fixture 026). input_type STILL flows into EmbeddingEvent.request_params with absence-is-meaningful semantics (empty when absent), decoupled from the wire body so it never leaks onto the wire. For an asymmetric model behind a compatible endpoint the construction MAY bind the §8.1 client-side query_prefix / document_prefix -- input_type then selects which prefix to prepend to each input client-side, the only way to express the query/document distinction on a wire with no input_type field (fixture 027); the un-prefixed caller intent is what reaches the event. encoding_format ('base64') rides the extras-pass-through bag (the mapping does not send it by default; OpenAI's wire default is 'float'). It parses the object envelope {object, data: [{object, index, embedding}], model, usage: {prompt_tokens, total_tokens}}, orders the vectors by data[i].index into input order, and maps usage.prompt_tokens -> EmbeddingUsage.input_tokens (total_tokens == prompt_tokens; embedding has no output tokens); response_id is null (OpenAI embeddings carry no id). Errors map per §8.3: 401/403 -> provider_authentication, 429 -> provider_rate_limit, 404 -> provider_invalid_model, 400/422 (malformed / oversized) -> provider_invalid_request, else (5xx) -> provider_unavailable. The client-side embed chunk-and-stitch over OpenAI's 2048-input cap (0092) stays out of scope here. genai_system = 'openai'. Conformance fixtures 023-027 pass, driven through the general wire-capture harness (0077) with the expected_wire_headers Authorization: Bearer assertion (0078). The §8.1 client-side prefix lookup is now a shared retrieval helper (client_side_prefix), adopted by both TEI (0077) and OpenAI."
793795

794796
# Spec v0.75.0 (proposal 0080). PromptGroup arity enforcement
795797
# (prompt-management §10 / §11 -- construct-time raise plus the new

src/openarmature/retrieval/_events.py

Lines changed: 48 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -62,6 +62,52 @@ def normalize_base_url(base_url: str, *, guard_prefix: str) -> str:
6262
return normalized
6363

6464

65+
# retrieval-provider §8.1 client-side prefix rule, reused by §8.3: input_type
66+
# "query" selects query_prefix, "document" selects document_prefix; any other
67+
# value (including None / absent) selects no prefix -- the symmetric default.
68+
def client_side_prefix(
69+
input_type: str | None,
70+
*,
71+
query_prefix: str | None,
72+
document_prefix: str | None,
73+
) -> str | None:
74+
"""Return the client-side prefix for ``input_type``, or ``None``.
75+
76+
``"query"`` maps to ``query_prefix``, ``"document"`` to ``document_prefix``;
77+
any other value (including ``None``) maps to ``None`` (no prefix).
78+
"""
79+
if input_type == "query":
80+
return query_prefix
81+
if input_type == "document":
82+
return document_prefix
83+
return None
84+
85+
86+
# Apply the §8.1 client-side prefix to an input list for input_type: prepend the
87+
# resolved prefix to every string, or return the list unchanged when input_type
88+
# resolves to no prefix. The single source of the "prefix the inputs" step
89+
# shared by the OpenAI (§8.3) and TEI (§8.1) embed mappings -- keeps the copy /
90+
# no-copy semantics from drifting between providers.
91+
def apply_client_side_prefix(
92+
input_strings: list[str],
93+
input_type: str | None,
94+
*,
95+
query_prefix: str | None,
96+
document_prefix: str | None,
97+
) -> list[str]:
98+
"""Return ``input_strings`` prefixed per ``input_type``, or unchanged.
99+
100+
Resolves the prefix via :func:`client_side_prefix`; when it is ``None`` the
101+
list is returned as-is (no defensive copy -- the result is only serialized
102+
onto the wire), otherwise a new list with the prefix prepended to each
103+
string.
104+
"""
105+
prefix = client_side_prefix(input_type, query_prefix=query_prefix, document_prefix=document_prefix)
106+
if prefix is None:
107+
return input_strings
108+
return [prefix + s for s in input_strings]
109+
110+
65111
def build_embedding_event(
66112
response: EmbeddingResponse,
67113
latency_ms: float,
@@ -282,9 +328,11 @@ def build_rerank_failed_event(
282328

283329

284330
__all__ = [
331+
"apply_client_side_prefix",
285332
"build_embedding_event",
286333
"build_embedding_failed_event",
287334
"build_rerank_event",
288335
"build_rerank_failed_event",
336+
"client_side_prefix",
289337
"normalize_base_url",
290338
]

src/openarmature/retrieval/providers/openai.py

Lines changed: 64 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -41,7 +41,12 @@
4141
)
4242
from openarmature.observability.correlation import current_dispatch
4343

44-
from .._events import build_embedding_event, build_embedding_failed_event, normalize_base_url
44+
from .._events import (
45+
apply_client_side_prefix,
46+
build_embedding_event,
47+
build_embedding_failed_event,
48+
normalize_base_url,
49+
)
4550
from ..provider import validate_embedding_input, validate_embedding_response
4651
from ..response import EmbeddingResponse, EmbeddingRuntimeConfig, EmbeddingUsage
4752

@@ -82,17 +87,20 @@ def _classify_embedding_http_error(resp: httpx.Response) -> LlmProviderError:
8287

8388
# Absence-is-meaningful per observability §5.5.2: only caller-supplied keys
8489
# appear in the event's request_params -- "the field was not supplied for
85-
# this call", distinct from a supplied zero. The input_type field joins this
86-
# set with proposal 0079 (OpenAI-compatible input_type is its wire mapping's
87-
# job, not 0077's). For embedding the event-param names coincide with the
88-
# wire-body keys, so the same dict feeds both the event and the body.
90+
# this call", distinct from a supplied zero. For OpenAI, input_type joins
91+
# dimensions in this set (proposal 0079). input_type does NOT feed the wire
92+
# body verbatim -- the OpenAI wire has no query/document field, so input_type
93+
# is a wire no-op (realized only as a client-side prefix when one is bound),
94+
# and the wire body is built separately.
8995
def _request_params_from_config(config: EmbeddingRuntimeConfig | None) -> dict[str, Any]:
9096
"""Extract the supplied embedding request parameters for the event."""
9197
if config is None:
9298
return {}
9399
out: dict[str, Any] = {}
94100
if config.dimensions is not None:
95101
out["dimensions"] = config.dimensions
102+
if config.input_type is not None:
103+
out["input_type"] = config.input_type
96104
return out
97105

98106

@@ -102,9 +110,16 @@ def _request_params_from_config(config: EmbeddingRuntimeConfig | None) -> dict[s
102110
class OpenAIEmbeddingProvider:
103111
"""OpenAI ``/v1/embeddings`` wire-compatible embedding provider.
104112
105-
Construct with a base URL (host root), the bound embedding model, and
106-
an optional API key + transport. ``embed()`` posts to
107-
``/v1/embeddings``.
113+
Construct with the bound embedding model and an optional API key +
114+
transport. ``base_url`` is the host root and defaults to the OpenAI
115+
origin (``https://api.openai.com``), overridable for any OpenAI-compatible
116+
backend. ``embed()`` posts to ``/v1/embeddings``.
117+
118+
The optional ``query_prefix`` / ``document_prefix`` bind the client-side
119+
asymmetric prefixes -- off by default (pure-symmetric OpenAI). When bound
120+
(for an asymmetric model served behind a compatible endpoint),
121+
``input_type`` selects which prefix to prepend to each input before
122+
sending, since the OpenAI wire carries no query/document field.
108123
109124
``ready()`` verifies the bound model per the ``readiness_probe``
110125
argument:
@@ -120,13 +135,15 @@ class OpenAIEmbeddingProvider:
120135
def __init__(
121136
self,
122137
*,
123-
base_url: str,
138+
base_url: str = "https://api.openai.com",
124139
model: str,
125140
api_key: str | None = None,
126141
transport: httpx.AsyncBaseTransport | None = None,
127142
timeout: float = 60.0,
128143
genai_system: str = "openai",
129144
readiness_probe: Literal["embed", "models", "both"] = "embed",
145+
query_prefix: str | None = None,
146+
document_prefix: str | None = None,
130147
populate_caller_metadata: bool = True,
131148
) -> None:
132149
# base_url is the host root; the provider appends the /v1 routes, so
@@ -143,6 +160,11 @@ def __init__(
143160
f"readiness_probe must be one of {sorted(_VALID_READINESS_PROBES)} (got {readiness_probe!r})"
144161
)
145162
self._readiness_probe = readiness_probe
163+
# Optional client-side asymmetric prefixes (§8.1, reused by §8.3): the
164+
# OpenAI wire has no query/document field, so input_type is realized as
165+
# a prefix prepend only when these are bound. Off by default (symmetric).
166+
self._query_prefix = query_prefix
167+
self._document_prefix = document_prefix
146168
# ``genai_system`` surfaces as gen_ai.system on the embedding span.
147169
self._genai_system = genai_system
148170
self._populate_caller_metadata = populate_caller_metadata
@@ -229,10 +251,12 @@ async def embed(
229251
input_strings = list(input)
230252
request_params = _request_params_from_config(config)
231253
request_extras = dict(config.model_extra or {}) if config is not None else {}
254+
input_type = config.input_type if config is not None else None
255+
dimensions = config.dimensions if config is not None else None
232256
adapter_start = time.perf_counter()
233257
try:
234258
validate_embedding_input(input_strings)
235-
body = self._build_request_body(input_strings, request_params, request_extras)
259+
body = self._build_request_body(input_strings, input_type, dimensions, request_extras)
236260
try:
237261
resp = await self._client.post("/v1/embeddings", json=body)
238262
except httpx.HTTPError as exc:
@@ -281,18 +305,38 @@ async def embed(
281305
def _build_request_body(
282306
self,
283307
input_strings: list[str],
284-
request_params: dict[str, Any],
308+
input_type: str | None,
309+
dimensions: int | None,
285310
request_extras: dict[str, Any],
286311
) -> dict[str, Any]:
287-
"""Build the /v1/embeddings request body."""
288-
# Extras are merged FIRST so the bound model, the input list, and the
289-
# declared request params always win: a caller's undeclared extra
290-
# named "model" or "input" must not clobber the wire identity. The
291-
# request params are the same dict the event carries (dimensions for
292-
# embedding), keeping the wire body and the observed request_params
293-
# provably identical. Proposal 0077's input_type maps to a per-vendor
294-
# wire key at the wire layer.
295-
return {**request_extras, "model": self.model, "input": input_strings, **request_params}
312+
"""Build the /v1/embeddings request body.
313+
314+
The OpenAI wire has no query/document field, so ``input_type`` is not
315+
realized on the wire -- it selects a client-side ``query_prefix`` /
316+
``document_prefix`` prepend when one is bound, and is otherwise a wire
317+
no-op (``input`` sent verbatim). ``dimensions`` maps to the wire
318+
``dimensions`` field (Matryoshka) when set.
319+
"""
320+
# input_type is realized as a client-side prefix (§8.1, reused by §8.3):
321+
# the OpenAI wire has NO query/document/input_type/task field, so an
322+
# input_type with no bound prefix leaves the input VERBATIM (the
323+
# symmetric no-op). input_type MUST NEVER land on the wire. Extras merge
324+
# FIRST so the managed keys (model, input, dimensions) always win: a
325+
# caller's undeclared extra named "model" or "input" must not clobber
326+
# the wire identity. encoding_format ("base64") rides the extras bag as
327+
# a request pass-through only -- _parse_response decodes float embeddings
328+
# (not base64 strings), so a base64 response raises
329+
# provider_invalid_response; base64 is not an end-to-end response format.
330+
inputs = apply_client_side_prefix(
331+
input_strings,
332+
input_type,
333+
query_prefix=self._query_prefix,
334+
document_prefix=self._document_prefix,
335+
)
336+
body: dict[str, Any] = {**request_extras, "model": self.model, "input": inputs}
337+
if dimensions is not None:
338+
body["dimensions"] = dimensions
339+
return body
296340

297341
def _parse_response(
298342
self,

src/openarmature/retrieval/providers/tei.py

Lines changed: 7 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -50,6 +50,7 @@
5050
from openarmature.observability.correlation import current_dispatch
5151

5252
from .._events import (
53+
apply_client_side_prefix,
5354
build_embedding_event,
5455
build_embedding_failed_event,
5556
build_rerank_event,
@@ -317,24 +318,19 @@ def _build_request_body(
317318
else:
318319
# Client-side prefix fallback: prepend for models without
319320
# configured prompts. Only "query" / "document" have prefixes.
320-
prefix = self._prefix_for(input_type)
321-
if prefix is not None:
322-
inputs = [prefix + s for s in inputs]
321+
inputs = apply_client_side_prefix(
322+
inputs,
323+
input_type,
324+
query_prefix=self._query_prefix,
325+
document_prefix=self._document_prefix,
326+
)
323327
body: dict[str, Any] = {**request_extras, "inputs": inputs}
324328
if prompt_name is not None:
325329
body["prompt_name"] = prompt_name
326330
if dimensions is not None:
327331
body["dimensions"] = dimensions
328332
return body
329333

330-
def _prefix_for(self, input_type: str) -> str | None:
331-
"""Return the client-side prefix for ``input_type``, or ``None``."""
332-
if input_type == "query":
333-
return self._query_prefix
334-
if input_type == "document":
335-
return self._document_prefix
336-
return None
337-
338334
def _parse_response(
339335
self,
340336
resp: httpx.Response,

0 commit comments

Comments
 (0)