Adopt 0097 rerank document-echo shape rule (#213)

Adopt proposal 0097's document-echo shape rule and lift it into a shared
document_echo helper in retrieval/_events.py, used by every RerankProvider
(Cohere, TEI, Jina). A string echo surfaces verbatim (an empty string is
present), a TextDoc object {"text": ...} unwraps to its string, a
text-less object or absent/null yields null, and a non-object scalar
(number / array / bool) raises provider_invalid_response rather than
folding to null. The verbatim echo stays on RerankResponse.raw.

This replaces three per-provider copies of the string-or-null fold;
Cohere and TEI now unwrap TextDoc objects and raise on corrupt scalars,
matching Jina. The conformance fixtures (019 cases C-G) ride the v0.92.0
pin bump.

Author: chris-colinsky

src/openarmature/retrieval/_events.py

@@ -19,15 +19,15 @@
 from __future__ import annotations
 
 from collections.abc import Mapping
-from typing import Any
+from typing import Any, cast
 
 from openarmature.graph.events import (
     EmbeddingEvent,
     EmbeddingFailedEvent,
     RerankEvent,
     RerankFailedEvent,
 )
-from openarmature.llm.errors import LlmProviderError
+from openarmature.llm.errors import LlmProviderError, ProviderInvalidResponse
 from openarmature.observability.correlation import (
     current_attempt_index,
     current_branch_name,
@@ -108,6 +108,31 @@ def apply_client_side_prefix(
     return [prefix + s for s in input_strings]
 
 
+# §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
+# documents. Vendors echo it in several shapes: a bare string, a TextDoc object
+# {"text": "..."} (Jina's return_documents form), a text-less object (an
+# ImageDoc), or absent/null. Unwrap the string from a string or a TextDoc; a
+# text-less object or absent/null yields null (an empty string is present --
+# surfaced as "", not folded to null). A NON-object scalar (number / array /
+# bool -- outside the documented anyOf[string, object, null]) is wire corruption
+# and raises provider_invalid_response (§7), NOT folded to null. The verbatim
+# echo is preserved on RerankResponse.raw regardless.
+def document_echo(value: Any) -> str | None:
+    """Extract the string document echo from a rerank result's documented shapes."""
+    if value is None:
+        return None
+    if isinstance(value, str):
+        return value
+    if isinstance(value, dict):
+        text = cast("dict[str, Any]", value).get("text")
+        return text if isinstance(text, str) else None
+    raise ProviderInvalidResponse(
+        f"rerank document echo must be a string, object, or null (got {type(value).__name__})"
+    )
+
+
 def build_embedding_event(
     response: EmbeddingResponse,
     latency_ms: float,
@@ -334,5 +359,6 @@ def build_rerank_failed_event(
     "build_rerank_event",
     "build_rerank_failed_event",
     "client_side_prefix",
+    "document_echo",
     "normalize_base_url",
 ]

src/openarmature/retrieval/providers/cohere.py

@@ -44,7 +44,7 @@
 )
 from openarmature.observability.correlation import current_dispatch
 
-from .._events import build_rerank_event, build_rerank_failed_event, normalize_base_url
+from .._events import build_rerank_event, build_rerank_failed_event, document_echo, normalize_base_url
 from ..provider import validate_rerank_input, validate_rerank_response
 from ..response import RerankResponse, RerankRuntimeConfig, RerankUsage, ScoredDocument
 
@@ -306,11 +306,10 @@ def _parse_response(
                 raise ProviderInvalidResponse(
                     "rerank response entry has a missing or non-numeric 'relevance_score'"
                 )
-            # Read ``document`` only when present; never auto-fill from the
-            # input documents list (§6 -- the provider's echo and the caller's
-            # input are two different surfaces).
-            document_raw = entry.get("document")
-            document = document_raw if isinstance(document_raw, str) else None
+            # Read the document echo per the shared §6 (0097) rule; never
+            # auto-fill from the input documents. Cohere v2 does not echo, so
+            # this is null in practice, but a compatible gateway may return one.
+            document = document_echo(entry.get("document"))
             scored.append(ScoredDocument(index=index, relevance_score=float(score), document=document))
         # §6: sort by relevance_score descending before validating / returning.
         scored.sort(key=lambda s: s.relevance_score, reverse=True)

src/openarmature/retrieval/providers/jina.py

@@ -61,6 +61,7 @@
     build_embedding_failed_event,
     build_rerank_event,
     build_rerank_failed_event,
+    document_echo,
     normalize_base_url,
 )
 from ..provider import (
@@ -153,25 +154,6 @@ def _nonneg_int(value: Any) -> int | None:
     return None
 
 
-# §6: ScoredDocument.document carries the provider's string echo verbatim when
-# present, null otherwise -- never fabricated from the input documents. Jina
-# types `document` as anyOf[string, TextDoc{text}, ImageDoc, null] (its OpenAPI
-# schema); the text reranker returns the TextDoc OBJECT form {"text": "..."}
-# when return_documents=true, while a bare string is also valid. Extract the
-# string from either shape; any other shape (an ImageDoc, a malformed entry, an
-# absent echo) yields null -- the verbatim object is preserved on
-# RerankResponse.raw regardless.
-def _document_echo(value: Any) -> str | None:
-    """Extract the string document echo from Jina's string-or-TextDoc shape."""
-    if isinstance(value, str):
-        return value
-    if isinstance(value, dict):
-        text = cast("dict[str, Any]", value).get("text")
-        if isinstance(text, str):
-            return text
-    return None
-
-
 # Absence-is-meaningful per observability §5.5.2: only caller-supplied keys
 # appear in the event's request_params -- "the field was not supplied for this
 # call", distinct from a supplied default. For Jina, input_type joins dimensions
@@ -662,12 +644,10 @@ def _parse_response(
                 raise ProviderInvalidResponse(
                     "rerank response entry has a missing or non-numeric 'relevance_score'"
                 )
-            # Read the document echo only when present; never auto-fill from the
-            # input documents list (§6 -- the provider's echo and the caller's
-            # input are two different surfaces). Jina echoes text as a TextDoc
-            # object ({"text": ...}), so _document_echo unwraps either the object
-            # or a bare-string shape onto the string field.
-            document = _document_echo(entry.get("document"))
+            # Read the document echo per the shared §6 (0097) rule; never
+            # auto-fill from the input documents. Jina echoes text as a TextDoc
+            # object ({"text": ...}), which document_echo unwraps to the string.
+            document = document_echo(entry.get("document"))
             scored.append(ScoredDocument(index=index, relevance_score=float(score), document=document))
         # §6: sort by relevance_score descending before validating / returning.
         scored.sort(key=lambda s: s.relevance_score, reverse=True)

src/openarmature/retrieval/providers/tei.py

@@ -55,6 +55,7 @@
     build_embedding_failed_event,
     build_rerank_event,
     build_rerank_failed_event,
+    document_echo,
 )
 from ..provider import (
     validate_embedding_input,
@@ -634,11 +635,10 @@ def _parse_chunk(
             score = entry.get("score")
             if not isinstance(score, (int, float)) or isinstance(score, bool):
                 raise ProviderInvalidResponse("rerank response entry has a missing or non-numeric 'score'")
-            # Read ``text`` only when present; never auto-fill from the input
-            # documents list (§6 -- the provider's echo and the caller's input
-            # are two different surfaces).
-            text_raw = entry.get("text")
-            document = text_raw if isinstance(text_raw, str) else None
+            # Read the ``text`` echo per the shared §6 (0097) rule; never
+            # auto-fill from the input documents. TEI's return_text echo is a
+            # bare string, so this is the identity path in practice.
+            document = document_echo(entry.get("text"))
             scored.append(
                 ScoredDocument(index=offset + index, relevance_score=float(score), document=document)
             )

tests/unit/test_retrieval_provider.py

@@ -637,6 +637,40 @@ def handler(req: httpx.Request) -> httpx.Response:
     await provider.aclose()
 
 
+async def test_cohere_rerank_document_echo_shared_rule() -> None:
+    # §6 (0097): Cohere is on the shared document_echo rule -- a TextDoc object
+    # echo unwraps to its text; a text-less object (or absent) yields null.
+    def handler(_req: httpx.Request) -> httpx.Response:
+        return httpx.Response(
+            200,
+            json=_rerank_body(
+                results=[
+                    {"index": 0, "relevance_score": 0.9, "document": {"text": "echoed"}},
+                    {"index": 1, "relevance_score": 0.5, "document": {"image": "..."}},
+                ]
+            ),
+        )
+
+    provider = _rerank_provider(handler)
+    response = await provider.rerank("q", ["a", "b"], config=RerankRuntimeConfig(return_documents=True))
+    await provider.aclose()
+    assert {r.index: r.document for r in response.results} == {0: "echoed", 1: None}
+
+
+async def test_cohere_rerank_non_object_document_echo_raises() -> None:
+    # §6 (0097): a non-object scalar echo is wire corruption -> provider_invalid_response.
+    def handler(_req: httpx.Request) -> httpx.Response:
+        return httpx.Response(
+            200,
+            json=_rerank_body(results=[{"index": 0, "relevance_score": 0.9, "document": 123}]),
+        )
+
+    provider = _rerank_provider(handler)
+    with pytest.raises(ProviderInvalidResponse):
+        await provider.rerank("q", ["a"], config=RerankRuntimeConfig(return_documents=True))
+    await provider.aclose()
+
+
 # --- typed-event dispatch ---------------------------------------------------
 
 
@@ -1152,6 +1186,19 @@ def handler(_req: httpx.Request) -> httpx.Response:
     await provider.aclose()
 
 
+async def test_tei_rerank_non_object_text_echo_raises() -> None:
+    # §6 (0097): TEI is on the shared document_echo rule -- its `text` echo is a
+    # bare string in practice, but a non-object scalar (wire corruption) raises
+    # provider_invalid_response rather than folding to null.
+    def handler(_req: httpx.Request) -> httpx.Response:
+        return httpx.Response(200, json=[{"index": 0, "score": 0.9, "text": 123}])
+
+    provider = _tei_rerank_provider(handler)
+    with pytest.raises(ProviderInvalidResponse):
+        await provider.rerank("q", ["a"], config=RerankRuntimeConfig(return_documents=True))
+    await provider.aclose()
+
+
 def test_tei_base_url_strips_trailing_slash() -> None:
     provider = TeiEmbeddingProvider(base_url="http://tei-embed.test/", model="m")
     assert provider.base_url == "http://tei-embed.test"
@@ -1335,9 +1382,10 @@ async def test_jina_rerank_document_echo_unwraps_textdoc_object() -> None:
     # Jina's real /v1/rerank echoes the document as a TextDoc OBJECT
     # ({"text": ...}) when return_documents=true, not a bare string -- its
     # OpenAPI types `document` as anyOf[string, TextDoc, ImageDoc, null]. The
-    # mapping unwraps either shape onto ScoredDocument.document (§6); a shape
-    # with no string text (an ImageDoc, a malformed entry) yields null, and the
-    # verbatim object is preserved on RerankResponse.raw.
+    # mapping unwraps either shape onto ScoredDocument.document (§6); a text-less
+    # object (an ImageDoc) or an absent/null echo yields null, an empty string is
+    # present (surfaced as ""), and the verbatim object is preserved on
+    # RerankResponse.raw.
     def handler(_req: httpx.Request) -> httpx.Response:
         return httpx.Response(
             200,
@@ -1347,23 +1395,47 @@ def handler(_req: httpx.Request) -> httpx.Response:
                 "results": [
                     {"index": 0, "relevance_score": 0.9, "document": {"text": "doc a"}},
                     {"index": 1, "relevance_score": 0.5, "document": "doc b"},
-                    {"index": 2, "relevance_score": 0.1, "document": {"image": "..."}},
+                    {"index": 2, "relevance_score": 0.4, "document": {"image": "..."}},
+                    {"index": 3, "relevance_score": 0.3, "document": ""},
+                    {"index": 4, "relevance_score": 0.1, "document": None},
                 ],
             },
         )
 
     provider = _jina_rerank_provider(handler)
     config = RerankRuntimeConfig(return_documents=True)
-    response = await provider.rerank("q", ["doc a", "doc b", "doc c"], config=config)
+    response = await provider.rerank("q", ["a", "b", "c", "d", "e"], config=config)
     await provider.aclose()
 
-    # TextDoc object -> its text; bare string -> itself; non-text shape -> null.
-    assert {r.index: r.document for r in response.results} == {0: "doc a", 1: "doc b", 2: None}
+    # TextDoc object -> its text; bare string -> itself; text-less object -> null;
+    # empty string -> "" (present, distinct from absent); null -> null.
+    echoes = {r.index: r.document for r in response.results}
+    assert echoes == {0: "doc a", 1: "doc b", 2: None, 3: "", 4: None}
     # The verbatim TextDoc object is still reachable on raw.
     assert isinstance(response.raw, dict)
     assert response.raw["results"][0]["document"] == {"text": "doc a"}
 
 
+async def test_jina_rerank_non_object_document_echo_raises() -> None:
+    # 0097: a non-object scalar echo (number / array / bool -- outside Jina's
+    # documented anyOf[string, TextDoc, ImageDoc, null]) is wire corruption and
+    # raises provider_invalid_response, NOT folded to null.
+    def handler(_req: httpx.Request) -> httpx.Response:
+        return httpx.Response(
+            200,
+            json={
+                "model": "jina-reranker-test",
+                "usage": {"total_tokens": 4},
+                "results": [{"index": 0, "relevance_score": 0.9, "document": 123}],
+            },
+        )
+
+    provider = _jina_rerank_provider(handler)
+    with pytest.raises(ProviderInvalidResponse):
+        await provider.rerank("q", ["doc a"], config=RerankRuntimeConfig(return_documents=True))
+    await provider.aclose()
+
+
 async def test_jina_rerank_relevance_score_parse_sorted_and_usage() -> None:
     def handler(_req: httpx.Request) -> httpx.Response:
         # UNSORTED, Cohere-shaped relevance_score; Jina reports total_tokens.