diff --git a/conformance.toml b/conformance.toml index 423c117c..e6e8a999 100644 --- a/conformance.toml +++ b/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 diff --git a/src/openarmature/retrieval/_wire.py b/src/openarmature/retrieval/_wire.py index 09fc2336..0d188c9b 100644 --- a/src/openarmature/retrieval/_wire.py +++ b/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_input, 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,90 @@ 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`. + """ + # 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 + # (cap == 0) or a misleading empty-stitched validation failure (cap < 0). + if cap <= 0: + raise ValueError(f"cap must be positive (got {cap})") + # Validate the input up front so an empty input raises + # provider_invalid_request -- the caller-side contract error -- rather than + # falling through to a misclassified provider_invalid_response from the + # empty stitched-count check. Providers already call this before the helper; + # this protects a direct / future caller. + validate_embedding_input(input_strings) + 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 +197,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", diff --git a/src/openarmature/retrieval/providers/cohere.py b/src/openarmature/retrieval/providers/cohere.py index 321cf987..066a0223 100644 --- a/src/openarmature/retrieval/providers/cohere.py +++ b/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,20 +620,20 @@ 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) @@ -643,40 +641,13 @@ async def _embed_chunked( 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( diff --git a/src/openarmature/retrieval/providers/tei.py b/src/openarmature/retrieval/providers/tei.py index 9b705b60..e0433ffa 100644 --- a/src/openarmature/retrieval/providers/tei.py +++ b/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: diff --git a/tests/conformance/test_retrieval_provider.py b/tests/conformance/test_retrieval_provider.py index df38f3cd..e8e7028b 100644 --- a/tests/conformance/test_retrieval_provider.py +++ b/tests/conformance/test_retrieval_provider.py @@ -87,20 +87,11 @@ # batch-chunking rule: 028-031 (0090 Cohere rerank), 032-037 (0091 Cohere # embed), 038 (0092 TEI /embed over-cap chunk-and-stitch). 028-031 ship with # 0090 and 032-037 with 0091 (CohereEmbeddingProvider against POST /v2/embed, -# driven through the general wire-capture harness). 038 stays deferred: the TEI -# /embed over-cap chunk-and-stitch rides the general 0092 embed rule against the -# TEI embed wire mapping, unshipped at this pin. -_DEFERRED_FIXTURES: dict[str, str] = { - **{ - p.stem: ( - "embedding batch-chunking general rule (proposal 0092); the TEI " - "/embed over-cap chunk-and-stitch fixture rides the unshipped TEI " - "embed wire mapping (proposal 0077)" - ) - for p in CONFORMANCE_DIR.glob("[0-9][0-9][0-9]-*.yaml") - if int(p.stem[:3]) == 38 - }, -} +# driven through the general wire-capture harness). 038 ships with 0092: +# TeiEmbeddingProvider.embed() chunk-and-stitches by its construction chunk_size +# via the shared chunk_and_stitch_embed helper (also adopted by Cohere /v2/embed), +# so no retrieval fixtures remain deferred. +_DEFERRED_FIXTURES: dict[str, str] = {} def _fixture_paths() -> list[Path]: diff --git a/tests/integration/test_tei_wire.py b/tests/integration/test_tei_wire.py index 8263a296..222159b2 100644 --- a/tests/integration/test_tei_wire.py +++ b/tests/integration/test_tei_wire.py @@ -55,6 +55,34 @@ async def test_tei_embed_returns_vectors_with_null_usage() -> None: assert response.response_id is None +@pytest.mark.integration +@requires_embed +async def test_tei_embed_chunk_and_stitch_over_chunk_size() -> None: + """A chunk_size smaller than the input list exercises the mandatory + /embed chunk-and-stitch: one vector per input, in input order, with the + same dimensionality across the chunk boundary and no fabricated usage.""" + inputs = [ + "the moon orbits the earth", + "lunar regolith is abrasive", + "the far side never faces earth", + ] + # chunk_size 2 over 3 inputs => 2 /embed requests (sizes 2, 1). + provider = TeiEmbeddingProvider(base_url=str(_EMBED_URL), model=_EMBED_MODEL, chunk_size=2) + try: + response = await provider.embed(inputs) + finally: + await provider.aclose() + + assert len(response.vectors) == len(inputs) + dim = len(response.vectors[0]) + assert dim > 0 + assert all(len(v) == dim for v in response.vectors) + assert response.dimensions == dim + # TEI /embed reports no usage object and no id -> both null across the stitch. + assert response.usage is None + assert response.response_id is None + + @pytest.mark.integration @requires_rerank async def test_tei_rerank_returns_sorted_results() -> None: diff --git a/tests/unit/test_retrieval_provider.py b/tests/unit/test_retrieval_provider.py index 978f3e60..ff42ae56 100644 --- a/tests/unit/test_retrieval_provider.py +++ b/tests/unit/test_retrieval_provider.py @@ -1372,6 +1372,151 @@ def handler(_req: httpx.Request) -> httpx.Response: await provider.aclose() +# --- /embed chunk-and-stitch (proposal 0092) -------------------------------- + + +async def test_tei_embed_chunk_and_stitch_by_chunk_size() -> None: + # Mirror fixture 038: 5 inputs, chunk_size 2 => exactly 3 /embed requests + # with `inputs` sizes [2, 2, 1] over the consecutive slices input[0:2], + # input[2:4], input[4:5], identical per-call params on each (bare {inputs} + # form, no prompt_name / dimensions). The per-chunk bare vector arrays are + # stitched IN INPUT ORDER across the chunk boundaries; TEI reports no usage + # and no id, so usage + response_id are null. One EmbeddingEvent per embed() + # call (the whole stitched call, not per chunk). + inputs = [f"e{i}" for i in range(5)] + chunks = [ + [[0.0, 0.5], [0.1, 0.5]], + [[0.2, 0.5], [0.3, 0.5]], + [[0.4, 0.5]], + ] + responses = iter(chunks) + captured: list[dict[str, Any]] = [] + + def handler(req: httpx.Request) -> httpx.Response: + captured.append(json.loads(req.content)) + assert req.url.path == "/embed" + return httpx.Response(200, json=next(responses)) + + provider = _tei_embed_provider(handler, chunk_size=2) + events, token = _collecting_dispatch() + try: + response = await provider.embed(inputs) + finally: + await provider.aclose() + _reset_active_dispatch(token) + + # Exactly 3 requests, `inputs` sizes [2, 2, 1], consecutive slices, identical + # per-call params on each (bare {inputs} form; no prompt_name / dimensions). + assert len(captured) == 3 + assert [len(b["inputs"]) for b in captured] == [2, 2, 1] + assert captured[0]["inputs"] == inputs[0:2] + assert captured[1]["inputs"] == inputs[2:4] + assert captured[2]["inputs"] == inputs[4:5] + for b in captured: + assert set(b) == {"inputs"} + # Stitched in input order across the boundaries (vector i == input i). + assert response.vectors == [[0.0, 0.5], [0.1, 0.5], [0.2, 0.5], [0.3, 0.5], [0.4, 0.5]] + assert len(response.vectors) == 5 + assert response.dimensions == 2 + # TEI /embed reports no usage object and no id -> both stitch to null. + assert response.usage is None + assert response.response_id is None + assert response.model == "tei-embed-test" + # raw is the list of the 3 per-chunk bare vector arrays, in request order. + assert response.raw == chunks + embed_events = [e for e in events if isinstance(e, EmbeddingEvent)] + assert len(embed_events) == 1 + assert embed_events[0].input_count == 5 + assert embed_events[0].dimensions == 2 + + +async def test_tei_embed_chunk_params_identical_across_chunks() -> None: + # >chunk_size inputs with an input_type_prompt_map + input_type=query: every + # /embed chunk carries the SAME prompt_name -- identical per-call params, only + # `inputs` differs across chunks (a param that drifted across chunks fails). + captured: list[dict[str, Any]] = [] + + def handler(req: httpx.Request) -> httpx.Response: + body = json.loads(req.content) + captured.append(body) + return httpx.Response(200, json=[[0.1, 0.2] for _ in body["inputs"]]) + + provider = _tei_embed_provider( + handler, chunk_size=2, input_type_prompt_map={"query": "query", "document": "passage"} + ) + await provider.embed([f"q{i}" for i in range(5)], config=EmbeddingRuntimeConfig(input_type="query")) + await provider.aclose() + + assert [len(b["inputs"]) for b in captured] == [2, 2, 1] + # prompt_name identical on every chunk request; the wire key set is stable. + assert all(b["prompt_name"] == "query" for b in captured) + assert all(set(b) == {"inputs", "prompt_name"} for b in captured) + + +async def test_tei_embed_per_chunk_count_mismatch_raises() -> None: + # The positional stitch validates EACH chunk's vector count. chunk_size 2 + # over 5 inputs => chunks [2, 2, 1]; the first chunk returns 1 vector for its + # 2 inputs, which raises provider_invalid_response before a silent positional + # misalignment (the shared helper's per-chunk guard). + def handler(req: httpx.Request) -> httpx.Response: + n = len(json.loads(req.content)["inputs"]) + count = 1 if n == 2 else n + return httpx.Response(200, json=[[0.1, 0.2] for _ in range(count)]) + + provider = _tei_embed_provider(handler, chunk_size=2) + with pytest.raises(ProviderInvalidResponse): + await provider.embed([f"e{i}" for i in range(5)]) + await provider.aclose() + + +async def test_tei_embed_single_request_path_when_within_chunk_size() -> None: + # len(input) <= chunk_size issues exactly ONE /embed request (the pre-0092 + # single-request behavior preserved): 3 inputs, chunk_size 4 => 1 request, + # and raw is the single bare array, not a one-element list. + captured: list[dict[str, Any]] = [] + + def handler(req: httpx.Request) -> httpx.Response: + captured.append(json.loads(req.content)) + return httpx.Response(200, json=[[0.1, 0.2], [0.3, 0.4], [0.5, 0.6]]) + + provider = _tei_embed_provider(handler, chunk_size=4) + response = await provider.embed(["a", "b", "c"]) + await provider.aclose() + + assert len(captured) == 1 + assert captured[0]["inputs"] == ["a", "b", "c"] + assert response.vectors == [[0.1, 0.2], [0.3, 0.4], [0.5, 0.6]] + # A single-request call's raw is the ONE bare array, not a one-element list. + assert response.raw == [[0.1, 0.2], [0.3, 0.4], [0.5, 0.6]] + assert response.usage is None + assert response.response_id is None + + +async def test_chunk_and_stitch_embed_rejects_non_positive_cap() -> None: + # The shared helper guards a caller passing cap <= 0 (a misconfigured per-call + # limit): fail loudly rather than a raw range() error (0) or an empty stitch. + from openarmature.retrieval._wire import chunk_and_stitch_embed + + async def _noop(chunk: list[str]) -> tuple[list[list[float]], None, None, list[Any]]: + return [], None, None, [] + + for bad_cap in (0, -1): + with pytest.raises(ValueError, match="cap must be positive"): + await chunk_and_stitch_embed(["a"], model="m", cap=bad_cap, embed_chunk=_noop) + + +async def test_chunk_and_stitch_embed_rejects_empty_input() -> None: + # An empty input reaching the helper is a caller error -> provider_invalid_request + # (validate_embedding_input), NOT a misclassified provider_invalid_response. + from openarmature.retrieval._wire import chunk_and_stitch_embed + + async def _noop(chunk: list[str]) -> tuple[list[list[float]], None, None, list[Any]]: + return [], None, None, [] + + with pytest.raises(ProviderInvalidRequest): + await chunk_and_stitch_embed([], model="m", cap=4, embed_chunk=_noop) + + # --- /rerank wire + chunk-and-stitch ----------------------------------------