diff --git a/conformance.toml b/conformance.toml index b71f1b1..226cb29 100644 --- a/conformance.toml +++ b/conformance.toml @@ -610,12 +610,11 @@ note = "Retrieval-provider embedding capability. The EmbeddingProvider protocol # Spec v0.70.0 (proposal 0060). Retrieval-provider rerank protocol — the # ``RerankProvider`` surface + ``RerankEvent`` / ``RerankFailedEvent`` typed # variants + OTel ``openarmature.rerank.complete`` span / Langfuse Retriever -# observation / rerank metrics. Sibling to the embedding surface (0059); -# python has not shipped retrieval-provider, so rerank is not-yet. Crossed by -# the v0.70.1 pin (adopted for fixture 110 / 0075); the 11 rerank observability -# fixtures (099-109) defer with it. +# observation / rerank metrics. Sibling to the embedding surface (0059). [proposals."0060"] -status = "not-yet" +status = "partial" +since = "0.16.0" +note = "Retrieval-provider rerank protocol. The RerankProvider protocol (retrieval-provider §5) + the RerankResponse / ScoredDocument / RerankUsage / RerankRuntimeConfig response types (§6) + the RerankRuntimeConfig runtime config (§2) ship, alongside the CohereRerankProvider reference reranker against POST /v2/rerank (§8.4, Cohere-shape): the request body is {model, query, documents (string array), top_n from top_k}, return_documents is a silent no-op (the Cohere wire has no such field), and max_tokens_per_doc rides the extras pass-through bag. The provider parses {id, model?, results: [{index, relevance_score, document?}], meta.billed_units.{search_units, input_tokens}}, sorts results by relevance_score descending, reads the document echo only when present (never auto-filled), builds a RerankUsage record only when the provider surfaces >= 1 usage figure (else usage = null, never fabricated), and enforces the §6 invariants -- valid index into the input documents, no duplicate index, len(results) <= top_k when supplied -- raising provider_invalid_response otherwise. Pre-send validation raises provider_invalid_request on an empty query, an empty documents list, or top_k <= 0 (top_k MAY exceed len(documents)). The provider dispatches the typed RerankEvent / RerankFailedEvent (graph-engine §6) per rerank() call, mutually exclusive, alongside the §7-category exception on failure; RerankEvent.output_results is populated unconditionally on success (proposal 0089). The §7 categories reuse the shared llm-provider taxonomy. Protocol fixtures 006-012 pass. partial because the observability rendering (OTel §5.5.13 openarmature.rerank.complete span + Langfuse §8.4.7 Retriever observation + the rerank-metrics surface; fixtures 099-109 / 138 / 141 / 142) lands in the 0060b follow-on -- the bundled OTel + Langfuse observers safely skip the rerank events for now. Wire-format details (0090 Cohere rerank fixtures 028-031) are a separate proposal." # Spec v0.55.0 (proposal 0065; repo pins v0.55.1). Failure-isolation # cause fidelity at non-node placements (pipeline-utilities §6.3 / diff --git a/src/openarmature/graph/events.py b/src/openarmature/graph/events.py index b092a61..a2e35a5 100644 --- a/src/openarmature/graph/events.py +++ b/src/openarmature/graph/events.py @@ -32,7 +32,7 @@ if TYPE_CHECKING: from openarmature.llm.messages import ToolCall from openarmature.llm.response import Usage - from openarmature.retrieval.response import EmbeddingUsage + from openarmature.retrieval.response import EmbeddingUsage, RerankUsage, ScoredDocument # Sentinel empty metadata mapping for events constructed without a # live caller-metadata snapshot (test helpers, synthetic events). @@ -877,6 +877,142 @@ class EmbeddingFailedEvent: caller_invocation_metadata: Mapping[str, AttributeValue] | None = None +# Spec: realizes graph-engine §6 -- the typed RerankEvent / RerankFailedEvent +# pair (proposal 0060, retrieval-provider rerank capability), the rerank +# sibling to the EmbeddingEvent / EmbeddingFailedEvent pair. Dispatched on the +# observer delivery queue per RerankProvider.rerank() call: the success variant +# after the response is parsed + validated (retrieval-provider §6), the failure +# variant alongside a raised §7 category exception (mutually exclusive per +# call). Scalar fan_out_index / branch_name only, matching the embedding pair +# (the lineage chains arrive uniformly across the provider events with proposal +# 0084). query / documents / request_extras / output_results are payload- +# bearing, populated unconditionally; observer-side privacy gates (OTel +# disable_provider_payload, Langfuse equivalents) apply at rendering, symmetric +# with EmbeddingEvent. +@dataclass(frozen=True) +class RerankEvent: + """A typed rerank provider call event delivered to observers. + + Carries identity, scoping, and outcome data for a successful + ``RerankProvider.rerank()`` call. Observer code filters by type + discrimination (``isinstance(event, RerankEvent)``). + + The identity / scoping / request-side fields mirror ``EmbeddingEvent``'s + convention; the outcome fields are rerank-specific: + + - ``query``: the query string the call was made with; non-nullable, + populated unconditionally (privacy gating is observer-side at + rendering). + - ``documents``: the input documents list; non-nullable, populated + unconditionally. Same privacy posture as ``query``. + - ``document_count``: ``len(documents)``; a convenience field. + - ``top_k``: the caller-supplied ``top_k`` value; ``None`` when the + caller passed ``None``. + - ``result_count``: ``len(output_results)``; a convenience field. + - ``output_results``: the scored results the call returned; populated + unconditionally on the success event (privacy gating is observer-side + at rendering). + - ``response_model`` / ``response_id``: the provider-returned model and + response identifiers; ``None`` when the provider returned none. + - ``usage``: the rerank usage record; ``None`` when the call returned + no usage. + - ``request_params``: the rerank request parameters the caller supplied + (e.g. ``return_documents``). Absence-is-meaningful: only supplied keys + appear; an empty mapping when none. + - ``request_extras``: the runtime-config extras pass-through bag. + - ``active_prompt`` / ``active_prompt_group``: prompt-context snapshots + at rerank-call time; ``None`` outside a binding. + - ``call_id``: a per-call disambiguator, always present, freshly minted + per ``rerank()`` call. + """ + + invocation_id: str + correlation_id: str | None + node_name: str + namespace: tuple[str, ...] + attempt_index: int + fan_out_index: int | None + branch_name: str | None + provider: str + model: str + response_id: str | None + response_model: str | None + # RerankUsage is a string-typed forward reference per the TYPE_CHECKING + # import -- keeps the runtime import direction graph -> retrieval off the + # module-load path. + usage: "RerankUsage | None" + latency_ms: float | None + query: str + documents: list[str] + document_count: int + top_k: int | None + result_count: int + # Spec graph-engine §6 (proposal 0089): the output scored results (sourced + # from RerankResponse.results at dispatch). Populated unconditionally on the + # success event -- payload-bearing like query / documents, gated observer- + # side at the rendering boundary. The §8.4.7 rerank output mapping reads + # this field. No output field on RerankFailedEvent (no response). + output_results: list["ScoredDocument"] + request_params: Mapping[str, Any] + request_extras: Mapping[str, Any] + active_prompt: Any + active_prompt_group: Any + call_id: str + caller_invocation_metadata: Mapping[str, AttributeValue] | None = None + + +# Spec: the failure sibling of RerankEvent (proposal 0060). Dispatched +# whenever RerankProvider.rerank() raises a §7 category exception -- covers +# both the provider-caught path and the pre-send validation raise +# (provider_invalid_request on an empty query / documents list / non-positive +# top_k). Dispatched ALONGSIDE the exception, not in place of it; mutually +# exclusive with RerankEvent on the same call. The response-side fields +# (usage / response_id / response_model / result_count / output_results) are +# absent (no response). +@dataclass(frozen=True) +class RerankFailedEvent: + """A typed rerank provider call failure event delivered to observers. + + Carries identity, scoping, and failure-context data for a ``rerank()`` + call that raised a retrieval-provider category exception. Observer code + filters by type discrimination (``isinstance(event, RerankFailedEvent)``). + + The identity / scoping / request-side field set mirrors ``RerankEvent``; + the response-side fields are absent. Failure-specific fields: + + - ``error_category``: the error category the call raised (one of the + rerank-applicable provider categories). Always present. + - ``error_type``: an optional impl-level / vendor-specific type or code; + ``None`` when unavailable. + - ``error_message``: a human-readable message; always present (the empty + string when the exception carried no message). + """ + + invocation_id: str + correlation_id: str | None + node_name: str + namespace: tuple[str, ...] + attempt_index: int + fan_out_index: int | None + branch_name: str | None + provider: str + model: str + latency_ms: float | None + query: str + documents: list[str] + document_count: int + top_k: int | None + request_params: Mapping[str, Any] + request_extras: Mapping[str, Any] + active_prompt: Any + active_prompt_group: Any + call_id: str + error_category: str + error_message: str + error_type: str | None = None + caller_invocation_metadata: Mapping[str, AttributeValue] | None = None + + # Spec: realizes pipeline-utilities §6.3 failure-isolation middleware # (proposal 0050). Emitted by FailureIsolationMiddleware when it # catches an exception escaping the inner chain and substitutes a @@ -1045,6 +1181,8 @@ class ToolCallFailedEvent: "MetadataAugmentationEvent", "NodeEvent", "ParallelBranchesEventConfig", + "RerankEvent", + "RerankFailedEvent", "ToolCallEvent", "ToolCallFailedEvent", ] diff --git a/src/openarmature/graph/observer.py b/src/openarmature/graph/observer.py index c35c4bb..2a53193 100644 --- a/src/openarmature/graph/observer.py +++ b/src/openarmature/graph/observer.py @@ -45,6 +45,8 @@ LlmRetryAttemptEvent, MetadataAugmentationEvent, NodeEvent, + RerankEvent, + RerankFailedEvent, ToolCallEvent, ToolCallFailedEvent, ) @@ -67,7 +69,9 @@ # dispatched by FailureIsolationMiddleware when it catches an exception # escaping the inner chain and substitutes a degraded partial update); # and EmbeddingEvent / EmbeddingFailedEvent (proposal 0059 typed embedding -# provider call events, dispatched on every EmbeddingProvider.embed()). +# provider call events, dispatched on every EmbeddingProvider.embed()); +# and RerankEvent / RerankFailedEvent (proposal 0060 typed rerank provider +# call events, dispatched on every RerankProvider.rerank()). ObserverEvent = ( NodeEvent | MetadataAugmentationEvent @@ -81,6 +85,8 @@ | ToolCallFailedEvent | EmbeddingEvent | EmbeddingFailedEvent + | RerankEvent + | RerankFailedEvent ) diff --git a/src/openarmature/observability/langfuse/observer.py b/src/openarmature/observability/langfuse/observer.py index 8f48245..06db3ab 100644 --- a/src/openarmature/observability/langfuse/observer.py +++ b/src/openarmature/observability/langfuse/observer.py @@ -40,6 +40,8 @@ LlmRetryAttemptEvent, MetadataAugmentationEvent, NodeEvent, + RerankEvent, + RerankFailedEvent, ToolCallEvent, ToolCallFailedEvent, ) @@ -485,6 +487,13 @@ async def __call__( if isinstance(event, EmbeddingEvent | EmbeddingFailedEvent): self._handle_embedding(event) return + # Proposal 0060 rerank observability (observability §8.4.7): the + # dedicated Langfuse Retriever observation renders from the typed + # rerank events. Not implemented yet (the rerank rendering lands in the + # 0060b follow-on); the observer safely skips the events for now rather + # than falling through to the NodeEvent phase dispatch. + if isinstance(event, RerankEvent | RerankFailedEvent): + return if event.phase == "started": self._open_started_observation(event) elif event.phase == "completed": diff --git a/src/openarmature/observability/otel/observer.py b/src/openarmature/observability/otel/observer.py index 9d0236f..658f1cd 100644 --- a/src/openarmature/observability/otel/observer.py +++ b/src/openarmature/observability/otel/observer.py @@ -110,6 +110,8 @@ LlmRetryAttemptEvent, MetadataAugmentationEvent, NodeEvent, + RerankEvent, + RerankFailedEvent, ToolCallEvent, ToolCallFailedEvent, ) @@ -786,6 +788,13 @@ async def __call__( if isinstance(event, EmbeddingEvent | EmbeddingFailedEvent): self._handle_embedding(event) return + # Proposal 0060 rerank observability (observability §5.5.13): the + # openarmature.rerank.complete span + rerank metrics render from the + # typed rerank events. Not implemented yet (the rerank rendering lands + # in the 0060b follow-on); the observer safely skips the events for now + # rather than falling through to the NodeEvent phase dispatch. + if isinstance(event, RerankEvent | RerankFailedEvent): + return # Proposal 0063 tool-execution observability: emit the # openarmature.tool.call span from the typed tool events. # Independent of disable_llm_spans (that flag is scoped to LLM diff --git a/src/openarmature/retrieval/__init__.py b/src/openarmature/retrieval/__init__.py index 691be52..38bc9f8 100644 --- a/src/openarmature/retrieval/__init__.py +++ b/src/openarmature/retrieval/__init__.py @@ -1,26 +1,47 @@ """The retrieval-provider capability. -The embedding provider protocol, its response types, and the bundled -reference embedding provider. The rerank protocol is a sibling surface on -the same capability. +The embedding + rerank provider protocols, their response types, and the +bundled reference providers (an OpenAI-compatible embedding provider and a +Cohere-shape reranker). Embedding and rerank are sibling surfaces on the same +capability. """ from __future__ import annotations from .provider import ( EmbeddingProvider, + RerankProvider, validate_embedding_input, validate_embedding_response, + validate_rerank_input, + validate_rerank_response, ) +from .providers.cohere import CohereRerankProvider from .providers.openai import OpenAIEmbeddingProvider -from .response import EmbeddingResponse, EmbeddingRuntimeConfig, EmbeddingUsage +from .response import ( + EmbeddingResponse, + EmbeddingRuntimeConfig, + EmbeddingUsage, + RerankResponse, + RerankRuntimeConfig, + RerankUsage, + ScoredDocument, +) __all__ = [ + "CohereRerankProvider", "EmbeddingProvider", "EmbeddingResponse", "EmbeddingRuntimeConfig", "EmbeddingUsage", "OpenAIEmbeddingProvider", + "RerankProvider", + "RerankResponse", + "RerankRuntimeConfig", + "RerankUsage", + "ScoredDocument", "validate_embedding_input", "validate_embedding_response", + "validate_rerank_input", + "validate_rerank_response", ] diff --git a/src/openarmature/retrieval/provider.py b/src/openarmature/retrieval/provider.py index ac7e896..42b1249 100644 --- a/src/openarmature/retrieval/provider.py +++ b/src/openarmature/retrieval/provider.py @@ -1,23 +1,26 @@ # Spec: realizes retrieval-provider §3 (EmbeddingProvider protocol + -# operations) and §4 (the response invariants). The §7 error categories -# are shared cross-capability with llm-provider, so the embedding path -# raises the existing llm-provider error classes (the embedding-applicable -# subset) rather than defining its own taxonomy. +# operations) + §4 (the embedding response invariants) and §5 (RerankProvider +# protocol + operations) + §6 (the rerank response invariants). The §7 error +# categories are shared cross-capability with llm-provider, so both paths +# raise the existing llm-provider error classes (the retrieval-applicable +# subset) rather than defining their own taxonomy. -"""EmbeddingProvider Protocol + input / response validation. +"""EmbeddingProvider / RerankProvider protocols + input / response validation. -An ``EmbeddingProvider`` is stateless; every call carries the full input -list. It does not cache, retry, or fall back (those compose above via -pipeline-utilities middleware). A provider exposes two operations: +An ``EmbeddingProvider`` / ``RerankProvider`` is stateless; every call +carries the full input. It does not cache, retry, or fall back (those +compose above via pipeline-utilities middleware). Each protocol exposes two +operations: - ``async ready() -> None``: verifies the bound model is reachable. -- ``async embed(input, *, config=None) -> EmbeddingResponse``: performs a - single embedding call, preserving input order. - -This module also exports two validators that bracket the per-call flow: -:func:`validate_embedding_input` (pre-send, list non-empty) and -:func:`validate_embedding_response` (post-receive, the vector-count and -dimensionality invariants). They raise ``provider_invalid_request`` and +- ``async embed(input, *, config=None) -> EmbeddingResponse`` / ``async + rerank(query, documents, *, top_k=None, config=None) -> RerankResponse``: + performs a single call, preserving order semantics. + +This module also exports the validators that bracket each per-call flow: +:func:`validate_embedding_input` / :func:`validate_rerank_input` (pre-send) +and :func:`validate_embedding_response` / :func:`validate_rerank_response` +(post-receive). They raise ``provider_invalid_request`` and ``provider_invalid_response`` respectively. """ @@ -28,7 +31,13 @@ from openarmature.llm.errors import ProviderInvalidRequest, ProviderInvalidResponse -from .response import EmbeddingResponse, EmbeddingRuntimeConfig +from .response import ( + EmbeddingResponse, + EmbeddingRuntimeConfig, + RerankResponse, + RerankRuntimeConfig, + ScoredDocument, +) class EmbeddingProvider(Protocol): @@ -94,8 +103,94 @@ def validate_embedding_response( return dimensions +class RerankProvider(Protocol): + """The shape of any retrieval-provider rerank implementation. + + Implementations are bound to a single rerank model identifier; + switching models means constructing a new provider, not passing a + different argument per call. + """ + + async def ready(self) -> None: + """Verify the bound rerank model is reachable and serving.""" + ... + + async def rerank( + self, + query: str, + documents: Sequence[str], + *, + top_k: int | None = None, + config: RerankRuntimeConfig | None = None, + ) -> RerankResponse: + """Score ``documents`` against ``query``, sorted by relevance. + + Args: + query: The query string the documents are scored against. + documents: The candidate documents. Always a list, even for a + single-document caller (wrap as a one-element list). Not + mutated by the implementation. + top_k: The maximum number of results to return. ``None`` means + "all" (up to ``len(documents)``). + config: Optional per-call request parameters. + + Returns a :class:`RerankResponse` whose ``results`` are sorted by + ``relevance_score`` descending; each result's ``index`` keys back + into the input ``documents`` list. + """ + ... + + +def validate_rerank_input( + query: str, + documents: Sequence[str], + top_k: int | None, +) -> None: + """Validate the rerank inputs before sending. + + Raises :class:`ProviderInvalidRequest` when the query is empty, the + documents list is empty, or ``top_k`` is supplied and not positive. + ``top_k`` may exceed ``len(documents)``; that is allowed. + """ + if not query: + raise ProviderInvalidRequest("rerank query must not be empty") + if not documents: + raise ProviderInvalidRequest("rerank documents list must not be empty") + if top_k is not None and top_k <= 0: + raise ProviderInvalidRequest(f"rerank top_k must be positive (got {top_k})") + + +def validate_rerank_response( + results: Sequence[ScoredDocument], + document_count: int, + top_k: int | None, +) -> None: + """Validate the rerank response invariants. + + Raises :class:`ProviderInvalidResponse` when a result's ``index`` is + out of range for the input documents, when an ``index`` appears twice, + or when ``top_k`` is supplied and the provider returned more results + than requested. + """ + seen: set[int] = set() + for result in results: + index = result.index + if index < 0 or index >= document_count: + raise ProviderInvalidResponse( + f"rerank result index {index} out of range for {document_count} documents" + ) + if index in seen: + raise ProviderInvalidResponse(f"rerank response has duplicate index {index}") + seen.add(index) + if top_k is not None and len(results) > top_k: + raise ProviderInvalidResponse(f"rerank provider returned {len(results)} results for top_k={top_k}") + + __all__ = [ "EmbeddingProvider", + "RerankProvider", "validate_embedding_input", "validate_embedding_response", + "validate_rerank_input", + "validate_rerank_response", ] diff --git a/src/openarmature/retrieval/providers/cohere.py b/src/openarmature/retrieval/providers/cohere.py new file mode 100644 index 0000000..5d631b4 --- /dev/null +++ b/src/openarmature/retrieval/providers/cohere.py @@ -0,0 +1,477 @@ +# Spec: realizes the retrieval-provider §5 RerankProvider protocol against a +# Cohere v2 POST /v2/rerank endpoint (retrieval-provider §8.4) -- the reference +# rerank provider. The §7 error categories are shared with llm-provider; the +# rerank-applicable subset (no unsupported_content_block, no +# structured_output_invalid) is mapped from the Cohere-shape HTTP error +# envelope below. Typed RerankEvent / RerankFailedEvent dispatch mirrors the +# embedding path via current_dispatch(). FOLLOW-UP: classify_http_error / +# base_url normalization are duplicated in spirit from the embedding + llm +# OpenAI providers; lifting a shared HTTP helper is a multi-provider follow-on. + +"""Cohere-shape rerank provider. + +``CohereRerankProvider`` issues ``POST {base_url}/v2/rerank`` and parses the +Cohere ``{id, model?, results: [{index, relevance_score, document?}], meta}`` +envelope into a :class:`RerankResponse`. ``base_url`` is the host root (the +provider appends ``/v2/rerank``), overridable for a proxy / private gateway. +The ``transport`` parameter is the test seam (``httpx.MockTransport``). + +The Cohere ``/v2/rerank`` wire has no ``return_documents`` parameter, so +``RerankRuntimeConfig.return_documents`` is a silent no-op on this wire (the +mapping sends no wire field for it). Cohere v2 does not echo document text, so +``ScoredDocument.document`` is null in practice; the parser still passes +through a ``document`` echo when a response carries one (a compatible gateway +or a fixture). ``max_tokens_per_doc`` rides the extras pass-through bag. +""" + +from __future__ import annotations + +import time +import uuid +from collections.abc import Mapping, Sequence +from typing import Any, cast + +import httpx + +from openarmature.graph.events import RerankEvent, RerankFailedEvent +from openarmature.llm.errors import ( + LlmProviderError, + ProviderAuthentication, + ProviderInvalidModel, + ProviderInvalidRequest, + ProviderInvalidResponse, + ProviderRateLimit, + ProviderUnavailable, +) +from openarmature.observability.correlation import ( + current_attempt_index, + current_branch_name, + current_correlation_id, + current_dispatch, + current_fan_out_index, + current_invocation_id, + current_namespace_prefix, +) +from openarmature.observability.metadata import AttributeValue, current_invocation_metadata + +from ..provider import validate_rerank_input, validate_rerank_response +from ..response import RerankResponse, RerankRuntimeConfig, RerankUsage, ScoredDocument + + +def _classify_cohere_http_error(resp: httpx.Response) -> LlmProviderError: + """Map a non-200 Cohere-shape rerank response to an error category. + + The rerank-applicable subset: 401/403 to auth, + 429 to rate_limit, 404 to invalid_model, 400/422 to invalid_request, and + every other status to unavailable. Returns the exception (does not raise) + so the caller raises with consistent traceback context. + """ + status = resp.status_code + try: + body_raw = resp.json() + except ValueError: + body_raw = {} + body: dict[str, Any] = cast("dict[str, Any]", body_raw) if isinstance(body_raw, dict) else {} + # Cohere surfaces errors as a top-level ``message`` string; the OpenAI-shape + # ``{error: {message}}`` envelope the fixtures also use is read as a + # fallback so a compatible gateway's error body still yields a message. + message_raw = body.get("message") + if not isinstance(message_raw, str): + error_block_raw = body.get("error") + error_block: dict[str, Any] = ( + cast("dict[str, Any]", error_block_raw) if isinstance(error_block_raw, dict) else {} + ) + message_candidate = error_block.get("message") + message_raw = message_candidate if isinstance(message_candidate, str) else None + message = message_raw + + if status in (401, 403): + return ProviderAuthentication(message or f"HTTP {status}") + if status == 429: + return ProviderRateLimit(message or "HTTP 429") + if status == 404: + return ProviderInvalidModel(message or "model not found") + if status in (400, 422): + return ProviderInvalidRequest(message or f"HTTP {status}") + return ProviderUnavailable(message or f"HTTP {status}") + + +# 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. return_documents is the one declared +# rerank-config field (retrieval-provider §2); it is reported on the event when +# the caller explicitly set it, even though it is a silent no-op on the Cohere +# wire (§8.4). Unlike embedding, the rerank event's request_params do NOT feed +# the wire body -- return_documents has no Cohere wire key. +def _request_params_from_config(config: RerankRuntimeConfig | None) -> dict[str, Any]: + """Extract the supplied rerank request parameters for the event.""" + if config is None: + return {} + out: dict[str, Any] = {} + if "return_documents" in config.model_fields_set: + out["return_documents"] = config.return_documents + return out + + +class CohereRerankProvider: + """Cohere ``/v2/rerank`` wire-shape rerank provider. + + Construct with a base URL (host root), the bound rerank model, and an + optional API key + transport. ``rerank()`` posts to ``/v2/rerank``. + + ``ready()`` verifies the bound model with a minimal one-document + ``/v2/rerank`` probe. The Cohere ``/v2/rerank`` wire exposes no + model-catalog probe (unlike the OpenAI-compatible embedding surface), so + there is a single universal probe. + """ + + def __init__( + self, + *, + base_url: str, + model: str, + api_key: str | None = None, + transport: httpx.AsyncBaseTransport | None = None, + timeout: float = 60.0, + genai_system: str = "cohere", + populate_caller_metadata: bool = True, + ) -> None: + # base_url is the host root; the provider appends /v2/rerank, so a + # trailing /v2 would produce a doubled /v2/v2 path that 404s (the + # sibling embedding provider guards the same footgun for /v1). Trailing + # slashes are stripped. + normalized = base_url.rstrip("/") + if normalized.endswith("/v2"): + raise ValueError( + f"base_url should be the host root (e.g. 'https://api.cohere.com'); " + f"the provider appends /v2/rerank itself, so a trailing /v2 would " + f"produce a doubled /v2/v2 path. Got {base_url!r}." + ) + self.base_url = normalized + self.model = model + # ``genai_system`` surfaces as gen_ai.system on the rerank span. + self._genai_system = genai_system + self._populate_caller_metadata = populate_caller_metadata + self._headers: dict[str, str] = {"Content-Type": "application/json"} + if api_key is not None: + self._headers["Authorization"] = f"Bearer {api_key}" + self._client: httpx.AsyncClient = httpx.AsyncClient( + base_url=self.base_url, + headers=self._headers, + transport=transport, + timeout=timeout, + ) + + async def aclose(self) -> None: + """Close the underlying HTTP client (releases the connection pool).""" + await self._client.aclose() + + async def ready(self) -> None: + """Verify the bound rerank model is reachable and serving.""" + # A minimal one-document /v2/rerank probe surfaces + # provider_invalid_model (404) / provider_unavailable (5xx) / + # provider_authentication exactly as a real rerank() would. + body = {"model": self.model, "query": "ready", "documents": ["ready"]} + try: + resp = await self._client.post("/v2/rerank", json=body) + except httpx.HTTPError as exc: + raise ProviderUnavailable(f"readiness probe failed: {exc}") from exc + if resp.status_code != 200: + raise _classify_cohere_http_error(resp) + + async def rerank( + self, + query: str, + documents: Sequence[str], + *, + top_k: int | None = None, + config: RerankRuntimeConfig | None = None, + ) -> RerankResponse: + """Score ``documents`` against ``query``, sorted by relevance.""" + dispatch = current_dispatch() + call_id = str(uuid.uuid4()) + # Snapshot prompt context at dispatch time (the node task's context); + # the delivery worker has a stale ContextVar view. Lazy import avoids + # the prompts -> graph -> ... cycle. + from openarmature.prompts.context import current_prompt_group, current_prompt_result + + active_prompt = current_prompt_result() + active_prompt_group = current_prompt_group() + documents_list = list(documents) + request_params = _request_params_from_config(config) + request_extras = dict(config.model_extra or {}) if config is not None else {} + adapter_start = time.perf_counter() + try: + validate_rerank_input(query, documents_list, top_k) + body = self._build_request_body(query, documents_list, top_k, request_extras) + try: + resp = await self._client.post("/v2/rerank", json=body) + except httpx.HTTPError as exc: + raise ProviderUnavailable(f"rerank request failed: {exc}") from exc + if resp.status_code != 200: + raise _classify_cohere_http_error(resp) + response = self._parse_response(resp, documents_list, top_k) + except LlmProviderError as exc: + latency_ms_failed = (time.perf_counter() - adapter_start) * 1000.0 + if dispatch is not None: + dispatch( + self._build_rerank_failed_event( + exc, + latency_ms_failed, + call_id=call_id, + query=query, + documents=documents_list, + top_k=top_k, + request_params=request_params, + request_extras=request_extras, + active_prompt=active_prompt, + active_prompt_group=active_prompt_group, + ), + ) + raise + latency_ms = (time.perf_counter() - adapter_start) * 1000.0 + if dispatch is not None: + dispatch( + self._build_rerank_event( + response, + latency_ms, + call_id=call_id, + query=query, + documents=documents_list, + top_k=top_k, + request_params=request_params, + request_extras=request_extras, + active_prompt=active_prompt, + active_prompt_group=active_prompt_group, + ), + ) + return response + + def _build_request_body( + self, + query: str, + documents_list: list[str], + top_k: int | None, + request_extras: dict[str, Any], + ) -> dict[str, Any]: + """Build the /v2/rerank request body.""" + # Extras are merged FIRST so the bound model, the query, the documents, + # and top_n always win: a caller's undeclared extra named "model" / + # "query" / "documents" must not clobber the wire identity. + # ``documents`` is the plain string-array form (Cohere v2 takes strings + # only). ``top_n`` maps from top_k and is omitted when the caller passed + # None. return_documents is NOT sent -- the Cohere wire has no such + # field, so it is a silent no-op (§8.4); it stays a declared config + # field (not an extra), so it never reaches request_extras. + body: dict[str, Any] = { + **request_extras, + "model": self.model, + "query": query, + "documents": documents_list, + } + # top_n is the wire mapping of the top_k parameter, so it is a managed + # key like model / query / documents: a caller extra named "top_n" must + # not set it. Leaving it would send an effective top_n the response-count + # validation (which keys off top_k) never checks. Drop any extras top_n; + # top_k is the sole source. + body.pop("top_n", None) + if top_k is not None: + body["top_n"] = top_k + return body + + def _parse_response( + self, + resp: httpx.Response, + documents_list: list[str], + top_k: int | None, + ) -> RerankResponse: + """Parse the Cohere rerank envelope into a RerankResponse.""" + # Reads the §6 result shape ({index, relevance_score, document?}), + # sorts by relevance_score descending (Cohere returns ranked results, + # but §6 mandates the sort regardless), then validates the §6 + # invariants (valid index into the input documents, no duplicate index, + # len(results) <= top_k when supplied) via validate_rerank_response. + try: + body_raw = resp.json() + except ValueError as exc: + raise ProviderInvalidResponse("rerank response is not valid JSON") from exc + if not isinstance(body_raw, dict): + raise ProviderInvalidResponse("rerank response is not a JSON object") + body = cast("dict[str, Any]", body_raw) + results_raw = body.get("results") + if not isinstance(results_raw, list): + raise ProviderInvalidResponse("rerank response missing 'results' array") + results = cast("list[Any]", results_raw) + scored: list[ScoredDocument] = [] + for raw_entry in results: + if not isinstance(raw_entry, dict): + raise ProviderInvalidResponse("rerank response entry is not a JSON object") + entry = cast("dict[str, Any]", raw_entry) + index = entry.get("index") + # bool is an int subclass, so exclude it explicitly. + if not isinstance(index, int) or isinstance(index, bool): + raise ProviderInvalidResponse("rerank response entry missing integer 'index'") + score = entry.get("relevance_score") + if not isinstance(score, (int, float)) or isinstance(score, bool): + 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 + 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) + validate_rerank_response(scored, len(documents_list), top_k) + usage = self._parse_usage(body) + response_id = body.get("id") + model = body.get("model") + return RerankResponse( + results=scored, + model=model if isinstance(model, str) else self.model, + usage=usage, + response_id=response_id if isinstance(response_id, str) else None, + raw=body, + ) + + def _parse_usage(self, body: dict[str, Any]) -> RerankUsage | None: + """Extract a RerankUsage from meta.billed_units, or None. + + A record is present only when the provider surfaces at least one usage + figure; an all-null record is never fabricated. + Cohere reports ``search_units``; ``input_tokens`` is read too so a + Cohere-compatible backend that surfaces it is honored. + """ + meta_raw = body.get("meta") + if not isinstance(meta_raw, dict): + return None + billed_raw = cast("dict[str, Any]", meta_raw).get("billed_units") + if not isinstance(billed_raw, dict): + return None + billed = cast("dict[str, Any]", billed_raw) + search_units = self._nonneg_int(billed.get("search_units")) + input_tokens = self._nonneg_int(billed.get("input_tokens")) + if search_units is None and input_tokens is None: + return None + return RerankUsage(search_units=search_units, input_tokens=input_tokens) + + @staticmethod + def _nonneg_int(value: Any) -> int | None: + """Return a non-negative int value, or None (bool excluded).""" + # bool is an int subclass, so exclude it explicitly; a malformed value + # falls back to None (the rerank succeeded; usage is secondary). + if isinstance(value, int) and not isinstance(value, bool) and value >= 0: + return value + return None + + def _build_rerank_event( + self, + response: RerankResponse, + latency_ms: float, + *, + call_id: str, + query: str, + documents: list[str], + top_k: int | None, + request_params: dict[str, Any], + request_extras: dict[str, Any], + active_prompt: Any, + active_prompt_group: Any, + ) -> RerankEvent: + """Construct the typed RerankEvent for the success path. + + Sources identity / scoping from the calling-node ContextVars and + outcome fields from the response. + """ + namespace = current_namespace_prefix() + node_name = namespace[-1] if namespace else "" + invocation_id = current_invocation_id() or "" + caller_metadata: Mapping[str, AttributeValue] | None = None + if self._populate_caller_metadata: + caller_metadata = dict(current_invocation_metadata()) + return RerankEvent( + invocation_id=invocation_id, + correlation_id=current_correlation_id(), + node_name=node_name, + namespace=namespace, + attempt_index=current_attempt_index(), + fan_out_index=current_fan_out_index(), + branch_name=current_branch_name(), + provider=self._genai_system, + model=self.model, + response_id=response.response_id, + response_model=response.model, + usage=response.usage, + latency_ms=latency_ms, + query=query, + documents=documents, + document_count=len(documents), + top_k=top_k, + result_count=len(response.results), + # Populated unconditionally on success per proposal 0089; privacy + # gating is observer-side at rendering (symmetric with query / + # documents). Sources output_results from the parsed response. + output_results=list(response.results), + request_params=request_params, + request_extras=request_extras, + active_prompt=active_prompt, + active_prompt_group=active_prompt_group, + call_id=call_id, + caller_invocation_metadata=caller_metadata, + ) + + def _build_rerank_failed_event( + self, + exc: LlmProviderError, + latency_ms: float, + *, + call_id: str, + query: str, + documents: list[str], + top_k: int | None, + request_params: dict[str, Any], + request_extras: dict[str, Any], + active_prompt: Any, + active_prompt_group: Any, + ) -> RerankFailedEvent: + """Construct the typed RerankFailedEvent for the failure path. + + ``error_type`` defaults to the exception class name (the "upstream + exception class name" style). + """ + namespace = current_namespace_prefix() + node_name = namespace[-1] if namespace else "" + invocation_id = current_invocation_id() or "" + caller_metadata: Mapping[str, AttributeValue] | None = None + if self._populate_caller_metadata: + caller_metadata = dict(current_invocation_metadata()) + return RerankFailedEvent( + invocation_id=invocation_id, + correlation_id=current_correlation_id(), + node_name=node_name, + namespace=namespace, + attempt_index=current_attempt_index(), + fan_out_index=current_fan_out_index(), + branch_name=current_branch_name(), + provider=self._genai_system, + model=self.model, + latency_ms=latency_ms, + query=query, + documents=documents, + document_count=len(documents), + top_k=top_k, + request_params=request_params, + request_extras=request_extras, + active_prompt=active_prompt, + active_prompt_group=active_prompt_group, + call_id=call_id, + error_category=exc.category, + error_type=type(exc).__name__, + error_message=str(exc), + caller_invocation_metadata=caller_metadata, + ) + + +__all__ = ["CohereRerankProvider"] diff --git a/src/openarmature/retrieval/response.py b/src/openarmature/retrieval/response.py index 0ad13f1..6bb5ff1 100644 --- a/src/openarmature/retrieval/response.py +++ b/src/openarmature/retrieval/response.py @@ -1,13 +1,15 @@ # Spec: realizes retrieval-provider §4 (EmbeddingResponse + EmbeddingUsage -# shapes and the response invariants) and §2 (the EmbeddingRuntimeConfig -# record). ``raw`` follows charter §3.1 principle 8 (Transparency over -# abstraction) -- it carries the parsed provider response verbatim -# alongside the normalized fields, so callers who need un-normalized data -# can reach through the abstraction. The ``input_type`` config field -# arrives with proposal 0077; ``dimensions`` is the one declared field -# at this pin. - -"""EmbeddingResponse, EmbeddingUsage, and EmbeddingRuntimeConfig. +# shapes and the response invariants), §6 (RerankResponse + RerankUsage + +# ScoredDocument shapes and the rerank response invariants), and §2 (the +# EmbeddingRuntimeConfig + RerankRuntimeConfig records). ``raw`` follows +# charter §3.1 principle 8 (Transparency over abstraction) -- it carries +# the parsed provider response verbatim alongside the normalized fields, so +# callers who need un-normalized data can reach through the abstraction. The +# ``input_type`` config field arrives with proposal 0077; ``dimensions`` is +# the one declared embedding-config field at this pin. ``return_documents`` +# is the one declared rerank-config field (proposal 0060). + +"""EmbeddingResponse / RerankResponse types and their runtime configs. The ``EmbeddingResponse`` is what ``EmbeddingProvider.embed()`` returns: one vector per input string in input order, the model identifier, usage, @@ -15,9 +17,15 @@ verbatim parsed provider response. ``raw`` carries everything the provider returned. -``EmbeddingRuntimeConfig`` is the optional per-call request-parameter -record. Implementations may accept additional provider-specific fields -via the extras pass-through; ``dimensions`` is the declared field. +The ``RerankResponse`` is what ``RerankProvider.rerank()`` returns: the +scored documents sorted by relevance descending, the model identifier, the +optional usage record and response id, and the verbatim parsed provider +response. + +``EmbeddingRuntimeConfig`` and ``RerankRuntimeConfig`` are the optional +per-call request-parameter records. Implementations may accept additional +provider-specific fields via the extras pass-through; ``dimensions`` is the +declared embedding field and ``return_documents`` the declared rerank field. """ from __future__ import annotations @@ -87,8 +95,97 @@ def from_partial(cls, **kwargs: Any) -> EmbeddingRuntimeConfig: return cls(**{k: v for k, v in kwargs.items() if v is not None}) +class RerankUsage(BaseModel): + """Token-accounting record for a rerank call. + + Both fields default to ``None`` and are individually nullable: a + provider may surface one figure and not the other (Cohere reports + ``search_units`` but no token count; Voyage AI reports ``input_tokens``). + """ + + # Spec §6: both fields are individually nullable; a RerankUsage record is + # present only when the provider surfaces at least one figure, and + # implementations MUST NOT fabricate an all-null record. + model_config = ConfigDict(extra="forbid") + + search_units: int | None = None + input_tokens: int | None = None + + +class ScoredDocument(BaseModel): + """A single scored result entry in a ``RerankResponse``. + + Attributes: + index: The 0-based position of this document in the original input + ``documents`` list. Load-bearing for caller-side lookup: + ``documents[result.index]`` maps a result back to its input. + relevance_score: The provider-assigned relevance score; higher = + more relevant. Provider-specific scale (not normalized here). + document: The echoed document text when the provider returns it; + ``None`` otherwise. Never fabricated from the input documents. + """ + + model_config = ConfigDict(extra="forbid") + + index: int + relevance_score: float + document: str | None = None + + +class RerankResponse(BaseModel): + """The result of a ``RerankProvider.rerank()`` call. + + Attributes: + results: The scored documents sorted by ``relevance_score`` + descending (most relevant first). Each entry's ``index`` keys + back to the original input ``documents`` list. + model: The model identifier the provider returned; may be more + specific than the bound identifier. + usage: The usage record, or ``None`` when the provider reports no + usage object. + response_id: The provider-returned response id when present; + ``None`` otherwise. + raw: The parsed provider response, populated on every successful + return. Carries everything the provider returned. + """ + + model_config = ConfigDict(extra="forbid") + + results: list[ScoredDocument] + model: str + usage: RerankUsage | None = None + response_id: str | None = None + raw: dict[str, Any] + + +# Spec §2 rerank runtime config: one declared field ``return_documents`` +# (boolean, default False) plus the extras pass-through bag (``extra="allow"``). +# Undeclared fields supplied by callers are forwarded to the wire body by the +# §8 wire-format mapping, except for the provider-reserved keys a mapping +# manages itself (e.g. the Cohere mapping owns model / query / documents / +# top_n, so a caller extra cannot clobber them). +class RerankRuntimeConfig(BaseModel): + """Per-call rerank request parameters.""" + + model_config = ConfigDict(extra="allow") + + return_documents: bool = False + + # Pure ergonomic, not a contract: lets callers splat a dict whose + # entries may be ``None`` without filtering at the call site, mirroring + # ``EmbeddingRuntimeConfig.from_partial``. + @classmethod + def from_partial(cls, **kwargs: Any) -> RerankRuntimeConfig: + """Construct a config, dropping kwargs whose value is ``None``.""" + return cls(**{k: v for k, v in kwargs.items() if v is not None}) + + __all__ = [ "EmbeddingResponse", "EmbeddingRuntimeConfig", "EmbeddingUsage", + "RerankResponse", + "RerankRuntimeConfig", + "RerankUsage", + "ScoredDocument", ] diff --git a/tests/conformance/harness/directives.py b/tests/conformance/harness/directives.py index 49180bc..0925008 100644 --- a/tests/conformance/harness/directives.py +++ b/tests/conformance/harness/directives.py @@ -412,6 +412,22 @@ class CallsEmbedSpec(_AllowExtras): stores_response_in: str +class CallsRerankSpec(_AllowExtras): + """Rerank-using node: sends ``query`` + ``documents`` to a mock rerank + provider and stores the result in ``stores_response_in``. Used by the + retrieval-provider rerank fixtures (006-012) and the rerank observability + fixtures (099-109 / 138 / 141 / 142). ``top_k``, ``config``, and ``model`` + are optional. Permissive on extras so companion knobs (e.g. + ``query_from_rendered_prompt`` on fixture 106) parse.""" + + query: str + documents: list[str] + stores_response_in: str + top_k: int | None = None + config: dict[str, Any] | None = None + model: str | None = None + + class EmitsLogSpec(_AllowExtras): """Additive companion: the node emits a log record alongside its state update. Verified by observability fixture 010 (Logs Bridge).""" @@ -447,6 +463,9 @@ class NodeSpec(_ForbidExtras): - ``calls_llm`` — see :class:`CallsLlmSpec`. - ``calls_embed`` — see :class:`CallsEmbedSpec` (embedding fixture 089, deferred at the runner; modelled so the fixture parses). + - ``calls_rerank`` — see :class:`CallsRerankSpec` (rerank fixtures 006-012 + run against the reference reranker; the rerank observability fixtures + defer at the runner until 0060b). Companion modifiers (additive, may combine with most primaries): @@ -472,6 +491,7 @@ class NodeSpec(_ForbidExtras): flaky_resume_aware: FlakyResumeAwareSpec | None = None calls_llm: CallsLlmSpec | None = None calls_embed: CallsEmbedSpec | None = None + calls_rerank: CallsRerankSpec | None = None calls_tool: CallsToolSpec | None = None # Companions — additive. @@ -506,6 +526,7 @@ class NodeSpec(_ForbidExtras): "flaky_resume_aware", "calls_llm", "calls_embed", + "calls_rerank", "calls_tool", ) @@ -708,6 +729,7 @@ class LlmCallSpec(_AllowExtras): __all__ = [ "CallsLlmSpec", + "CallsRerankSpec", "EdgeSpec", "EmitsLogSpec", "ErrorRecoveryMiddleware", diff --git a/tests/conformance/test_fixture_parsing.py b/tests/conformance/test_fixture_parsing.py index a457696..c360a2e 100644 --- a/tests/conformance/test_fixture_parsing.py +++ b/tests/conformance/test_fixture_parsing.py @@ -527,32 +527,40 @@ def _id(case: tuple[str, Path]) -> str: "observability/081-embedding-event-active-prompt-populated": ( "Proposal 0059 typed-event-collector shape; runs in test_observability" ), - # Proposal 0060 (retrieval-provider rerank, v0.70.0): the rerank - # observability fixtures (099-109) model the RerankEvent / - # RerankFailedEvent + rerank span / Langfuse Retriever / rerank-metrics - # surface, which python does not implement (0060 is not-yet; rerank lands - # with the embedding capability in v0.16.0). Sibling to embeddings (074-083). - "observability/099-rerank-event-dispatch": "Proposal 0060 rerank events; not implemented", + # Proposal 0060 (retrieval-provider rerank, v0.70.0): the rerank typed- + # collector fixtures (099-106) share the ``expected.observers`` shape of the + # embedding typed-collector fixtures (074-081) / the LLM 050-068 (which the + # harness schema models loosely, hence the same parser-deferral); they still + # RUN via test_observability. Their ``calls_rerank`` directive now parses + # (0060a); the parse-deferral is purely the typed-event-collector shape. The + # OTel-span (107), Langfuse (108), and rerank-metrics (109) fixtures parse + # fine against the existing span_tree / langfuse_trace / metrics shapes now + # that calls_rerank is modelled and are NOT deferred here -- they run in + # test_observability, where the rerank rendering stays deferred until 0060b. + "observability/099-rerank-event-dispatch": ( + "Proposal 0060 typed-event-collector shape; runs in test_observability" + ), "observability/100-rerank-failure-event-dispatch-on-provider-unavailable": ( - "Proposal 0060 rerank events; not implemented" + "Proposal 0060 typed-event-collector shape; runs in test_observability" + ), + "observability/101-rerank-event-mutual-exclusion": ( + "Proposal 0060 typed-event-collector shape; runs in test_observability" + ), + "observability/102-rerank-event-call-id-distinct": ( + "Proposal 0060 typed-event-collector shape; runs in test_observability" ), - "observability/101-rerank-event-mutual-exclusion": "Proposal 0060 rerank events; not implemented", - "observability/102-rerank-event-call-id-distinct": "Proposal 0060 rerank events; not implemented", "observability/103-rerank-event-query-and-documents-populated": ( - "Proposal 0060 rerank events; not implemented" + "Proposal 0060 typed-event-collector shape; runs in test_observability" ), "observability/104-rerank-event-request-params-populated": ( - "Proposal 0060 rerank events; not implemented" + "Proposal 0060 typed-event-collector shape; runs in test_observability" ), "observability/105-rerank-event-top-k-and-result-count-populated": ( - "Proposal 0060 rerank events; not implemented" + "Proposal 0060 typed-event-collector shape; runs in test_observability" ), "observability/106-rerank-event-active-prompt-populated": ( - "Proposal 0060 rerank events; not implemented" + "Proposal 0060 typed-event-collector shape; runs in test_observability" ), - "observability/107-otel-rerank-span-attributes": "Proposal 0060 rerank events; not implemented", - "observability/108-langfuse-rerank-observation": "Proposal 0060 rerank events; not implemented", - "observability/109-rerank-metrics-token-and-duration": "Proposal 0060 rerank events; not implemented", # Proposal 0075 (callable-branch span, fixture 110 added v0.70.1): the case # mixes a graph-style ``expected.final_state`` with the observability # ``span_tree``; the cross-capability parser's ObservabilityExpected model @@ -623,27 +631,11 @@ def _id(case: tuple[str, Path]) -> str: "observability/135-within-node-directive-execution-order": ( "Proposal 0087 within-node directive execution order; not implemented" ), - # Proposal 0089 (embedding / rerank typed-event output, v0.84.0) -- the - # rerank failure observation directive shape (rerank capability unshipped). - "observability/138-langfuse-rerank-failure-observation": ( - "Proposal 0089 rerank failure observation; rerank capability not implemented" - ), # Proposal 0086 (PromptManager default cache_ttl_seconds, v0.79.0) -- the # manager default-cache-ttl directive shape. "prompt-management/036-prompt-manager-default-cache-ttl": ( "Proposal 0086 default cache_ttl_seconds; not implemented" ), - # Proposal 0093 (nullable provider usage records, v0.88.0) -- the rerank - # no-usage fixtures use the calls_rerank directive (rerank capability - # unshipped, proposal 0060), so they defer at the parse layer like - # 099-109 / 138. The embedding no-usage fixtures (139/140/143) parse via - # the modelled calls_embed directive. - "observability/141-otel-rerank-no-usage-attributes-omitted": ( - "Proposal 0093 rerank no-usage; rerank capability not implemented (0060)" - ), - "observability/142-langfuse-rerank-no-usage-usagedetails-omitted": ( - "Proposal 0093 rerank no-usage; rerank capability not implemented (0060)" - ), } diff --git a/tests/conformance/test_observability.py b/tests/conformance/test_observability.py index 9142432..ca14172 100644 --- a/tests/conformance/test_observability.py +++ b/tests/conformance/test_observability.py @@ -248,7 +248,9 @@ def _reset_otel_global_tracer_provider(restore_to: object) -> None: ) _RERANK_DEFER = ( - "rerank capability (proposal 0060) unimplemented until v0.16.0; no rerank event/provider to record from" + "rerank observability rendering (OTel span / Langfuse Retriever observation / rerank metrics) " + "lands in 0060b; the RerankProvider + typed RerankEvent / RerankFailedEvent ship in 0060a, but " + "the bundled observers safe-skip the rerank events until the rendering is wired" ) @@ -273,8 +275,8 @@ def _reset_otel_global_tracer_provider(restore_to: object) -> None: # the LLM per-attempt event only; extending the metric instruments to the # embedding path rides the §11 embedding-metrics work, not 0059b. "089-embedding-metrics-token-and-duration": _EMBEDDING_METRICS_DEFER, - # Rerank observability (proposal 0060, v0.70.0). The rerank protocol is - # unshipped in python until v0.16.0; no rerank provider/event exists. + # Rerank observability (proposal 0060, v0.70.0). The RerankProvider + typed + # rerank events ship in 0060a; only the observer rendering defers to 0060b. **{ fixture_id: _RERANK_DEFER for fixture_id in ( @@ -290,8 +292,8 @@ def _reset_otel_global_tracer_provider(restore_to: object) -> None: "108-langfuse-rerank-observation", "109-rerank-metrics-token-and-duration", # proposal 0089 (v0.84.0) rerank failure observation -- the - # RerankFailedEvent ERROR-level Langfuse rendering; blocked on - # the same unimplemented rerank capability. + # RerankFailedEvent ERROR-level Langfuse rendering; lands with the + # 0060b rerank observability. "138-langfuse-rerank-failure-observation", ) }, @@ -399,12 +401,12 @@ def _reset_otel_global_tracer_provider(restore_to: object) -> None: "140-langfuse-embedding-no-usage-usagedetails-omitted", ) }, - # 141/142 -- rerank no-usage; also blocked on the unshipped rerank - # capability (proposal 0060). + # 141/142 -- rerank no-usage; the rendering lands with the 0060b rerank + # observability (the rerank capability ships in 0060a). **{ fixture_id: ( - "rerank no-usage rendering (proposal 0093); blocked on the unshipped " - "rerank capability (proposal 0060)" + "rerank no-usage rendering (proposal 0093); lands with the 0060b " + "rerank observability (the rerank capability ships in 0060a)" ) for fixture_id in ( "141-otel-rerank-no-usage-attributes-omitted", diff --git a/tests/conformance/test_retrieval_provider.py b/tests/conformance/test_retrieval_provider.py index 2131d64..629999b 100644 --- a/tests/conformance/test_retrieval_provider.py +++ b/tests/conformance/test_retrieval_provider.py @@ -1,16 +1,14 @@ -"""Run every spec retrieval-provider conformance fixture against OpenAIEmbeddingProvider. +"""Run every spec retrieval-provider conformance fixture against the reference providers. The fixtures (``spec/retrieval-provider/conformance/``) describe an -``EmbeddingProvider``'s behavior as OpenAI-compatible ``/v1/embeddings`` -mock responses + expected ``embed()`` outcomes. Each fixture wraps the -call in a single ``embed_node`` graph, but that node is just one -``embed()`` call, so the harness extracts the ``calls_embed`` directive -+ ``mock_embedding`` and drives the real :class:`OpenAIEmbeddingProvider` -through ``httpx.MockTransport`` directly, mirroring ``test_llm_provider`` -(no graph engine; fixtures 074-083 cover the observed-event path). - -Rerank fixtures (006-012) ride the ``RerankProvider`` and are deferred here -until it lands. +``EmbeddingProvider`` / ``RerankProvider``'s behavior as mock responses + +expected ``embed()`` / ``rerank()`` outcomes. Each fixture wraps the call in a +single-node graph, but that node is just one provider call, so the harness +extracts the ``calls_embed`` / ``calls_rerank`` directive + ``mock_embedding`` +/ ``mock_rerank`` and drives the real :class:`OpenAIEmbeddingProvider` / +:class:`CohereRerankProvider` through ``httpx.MockTransport`` directly, +mirroring ``test_llm_provider`` (no graph engine; the observed-event fixtures +cover the observer path). """ from __future__ import annotations @@ -25,7 +23,12 @@ import yaml from openarmature.llm import LlmProviderError -from openarmature.retrieval import EmbeddingRuntimeConfig, OpenAIEmbeddingProvider +from openarmature.retrieval import ( + CohereRerankProvider, + EmbeddingRuntimeConfig, + OpenAIEmbeddingProvider, + RerankRuntimeConfig, +) from ._deferral import skip_if_deferred @@ -37,8 +40,12 @@ # fixtures' mock-body model. _DEFAULT_MODEL = "text-embedding-test" -# Rerank fixtures (006-012) ride the RerankProvider, which lands with -# proposal 0060. Deferred so a green run means "what we implement passes." +# Default bound rerank model when a fixture's calls_rerank omits one; matches +# the rerank fixtures' mock-body model. +_DEFAULT_RERANK_MODEL = "rerank-test" + +# The rerank protocol fixtures (006-012) run against the reference +# CohereRerankProvider (proposal 0060) as of v0.16.0. # # The v0.84.0 pin bump also introduced the retrieval-provider WIRE-MAPPING # fixtures (013-027) for proposals 0077 (TEI) / 0078 (Jina) / 0079 @@ -53,13 +60,11 @@ # The v0.88.0 pin pulls in the Cohere wire mappings + the general embed # batch-chunking rule: 028-031 (0090 Cohere rerank), 032-037 (0091 Cohere # embed), 038 (0092 TEI /embed over-cap chunk-and-stitch). None are shipped -# -- same deferral rationale as 013-027. +# -- same deferral rationale as 013-027. (0060a ships the RerankProvider +# protocol + the Cohere-shape reference reranker for the protocol fixtures +# 006-012; the 0090 Cohere wire-mapping fixtures 028-031 assert the request- +# side wire contract this runner cannot capture, so they stay deferred.) _DEFERRED_FIXTURES: dict[str, str] = { - **{ - p.stem: "RerankProvider not implemented (proposal 0060 not-yet)" - for p in CONFORMANCE_DIR.glob("[0-9][0-9][0-9]-*.yaml") - if 6 <= int(p.stem[:3]) <= 12 - }, **{ p.stem: "TEI wire mapping (proposal 0077) not implemented" for p in CONFORMANCE_DIR.glob("[0-9][0-9][0-9]-*.yaml") @@ -127,7 +132,7 @@ def handler(request: httpx.Request) -> httpx.Response: try: spec = next(iterator) except StopIteration as exc: - raise AssertionError(f"mock embedding exhausted at request {len(captured)}") from exc + raise AssertionError(f"mock provider exhausted at request {len(captured)}") from exc status = int(spec.get("status", 200)) body = spec.get("body") if body is None: @@ -162,6 +167,31 @@ def _build_config(config_block: Mapping[str, Any] | None) -> EmbeddingRuntimeCon return EmbeddingRuntimeConfig(**block) +def _build_rerank_provider( + responses: list[Mapping[str, Any]], + *, + model: str, +) -> tuple[CohereRerankProvider, list[httpx.Request]]: + transport, captured = _build_handler(responses) + provider = CohereRerankProvider( + base_url="http://mock-rerank.test", + model=model, + api_key="test-key", + transport=transport, + ) + return provider, captured + + +def _build_rerank_config(config_block: Mapping[str, Any] | None) -> RerankRuntimeConfig | None: + if not config_block: + return None + block = dict(config_block) + extras = cast("Mapping[str, Any] | None", block.pop("extras", None)) + if extras: + block.update(extras) + return RerankRuntimeConfig(**block) + + def _assert_embedding_response( response: Any, expected: Mapping[str, Any], @@ -212,6 +242,66 @@ def _check_success_invariants( # above plus the baseline count assertion cover the core contract. +def _assert_rerank_response( + response: Any, + expected: Mapping[str, Any], +) -> None: + """Assert the rerank response against the expected.final_state. + block: the sorted results (index / relevance_score / document), model, + response_id, and usage (search_units / input_tokens, null-aware).""" + for key, val in expected.items(): + if key == "results": + results = cast("list[Mapping[str, Any]]", val) + assert len(response.results) == len(results), ( + f"result count mismatch: {len(response.results)} != {len(results)}" + ) + for got, want in zip(response.results, results, strict=True): + assert got.index == want["index"], f"index mismatch: {got.index} != {want['index']}" + assert got.relevance_score == want["relevance_score"], ( + f"relevance_score mismatch: {got.relevance_score} != {want['relevance_score']}" + ) + # ``document`` is only asserted when the expected entry carries + # the key (null included -- the omitted-echo case asserts None). + if "document" in want: + assert got.document == want["document"], ( + f"document mismatch: {got.document!r} != {want['document']!r}" + ) + elif key == "model": + assert response.model == val + elif key == "response_id": + assert response.response_id == val + elif key == "usage": + if val is None: + assert response.usage is None + else: + usage = cast("Mapping[str, Any]", val) + assert response.usage is not None, "expected a usage record, got None" + if "search_units" in usage: + assert response.usage.search_units == usage["search_units"] + if "input_tokens" in usage: + assert response.usage.input_tokens == usage["input_tokens"] + # Unknown keys are skipped (tolerant), matching the embedding harness. + + +def _check_rerank_success_invariants( + response: Any, + documents: list[str], + invariants: Mapping[str, Any], +) -> None: + for key, val in invariants.items(): + if not val: + continue + if key == "results_sorted_by_relevance_descending": + scores = [r.relevance_score for r in response.results] + assert scores == sorted(scores, reverse=True), f"results not sorted descending: {scores}" + elif key == "each_index_valid_into_input_documents": + assert all(0 <= r.index < len(documents) for r in response.results) + elif key == "result_count_at_most_document_count": + assert len(response.results) <= len(documents) + # Unknown invariants are skipped (tolerant); the structural invariants + # above cover the core §6 contract. + + @pytest.mark.parametrize("fixture_path", _fixture_paths(), ids=_fixture_id) async def test_retrieval_provider_fixture(fixture_path: Path) -> None: fixture_id = fixture_path.stem @@ -228,6 +318,9 @@ async def test_retrieval_provider_fixture(fixture_path: Path) -> None: async def _run_one_case(case: Mapping[str, Any]) -> None: entry = cast("str", case["entry"]) node = cast("Mapping[str, Any]", cast("Mapping[str, Any]", case["nodes"])[entry]) + if "calls_rerank" in node: + await _run_rerank_case(case, cast("Mapping[str, Any]", node["calls_rerank"])) + return calls_embed = cast("Mapping[str, Any]", node["calls_embed"]) input_strings = cast("list[str]", calls_embed["input"]) model = cast("str", calls_embed.get("model", _DEFAULT_MODEL)) @@ -259,3 +352,38 @@ async def _run_one_case(case: Mapping[str, Any]) -> None: _check_success_invariants(response, input_strings, invariants) finally: await provider.aclose() + + +async def _run_rerank_case(case: Mapping[str, Any], calls_rerank: Mapping[str, Any]) -> None: + query = cast("str", calls_rerank["query"]) + documents = cast("list[str]", calls_rerank["documents"]) + top_k = cast("int | None", calls_rerank.get("top_k")) + model = cast("str", calls_rerank.get("model", _DEFAULT_RERANK_MODEL)) + config = _build_rerank_config(cast("Mapping[str, Any] | None", calls_rerank.get("config"))) + responses = cast("list[Mapping[str, Any]]", case.get("mock_rerank") or []) + provider, _captured = _build_rerank_provider(responses, model=model) + expected_error = cast("Mapping[str, Any] | None", case.get("expected_error")) + expected = cast("Mapping[str, Any]", case.get("expected") or {}) + invariants = cast("Mapping[str, Any]", expected.get("invariants") or {}) + try: + if expected_error is not None: + with pytest.raises(LlmProviderError) as excinfo: + await provider.rerank(query, documents, top_k=top_k, config=config) + assert excinfo.value.category == expected_error["category"], ( + f"expected {expected_error['category']}, got {excinfo.value.category}" + ) + else: + response = await provider.rerank(query, documents, top_k=top_k, config=config) + # Baseline assertion: every success case satisfies the §6 + # at-most-len(documents) bound, so a fixture without a + # final_state / invariants block still asserts something real. + assert len(response.results) <= len(documents), ( + f"expected at most {len(documents)} results, got {len(response.results)}" + ) + final_state = cast("Mapping[str, Any]", expected.get("final_state") or {}) + stored = cast("str | None", calls_rerank.get("stores_response_in")) + if stored is not None and stored in final_state: + _assert_rerank_response(response, cast("Mapping[str, Any]", final_state[stored])) + _check_rerank_success_invariants(response, documents, invariants) + finally: + await provider.aclose() diff --git a/tests/unit/test_retrieval_provider.py b/tests/unit/test_retrieval_provider.py index d23020b..2fa732e 100644 --- a/tests/unit/test_retrieval_provider.py +++ b/tests/unit/test_retrieval_provider.py @@ -1,22 +1,43 @@ -"""Unit tests for the retrieval-provider embedding capability. - -Covers behavior the spec conformance fixtures (001-005) do not pin: the -base_url guard, the readiness-probe modes, and the bundled observers' -safe handling of the typed embedding events (the spans / observations are -a follow-up; until then the observers must skip the events rather than -fall through to the NodeEvent phase dispatch). +"""Unit tests for the retrieval-provider embedding + rerank capability. + +Covers behavior the spec conformance fixtures do not pin: the base_url +guards, the readiness-probe modes, the rerank validators + request-body +mapping + typed-event dispatch, and the bundled observers' safe handling of +the typed embedding / rerank events (the spans / observations are a follow-up; +until then the observers must skip the events rather than fall through to the +NodeEvent phase dispatch). """ from __future__ import annotations -from typing import Literal +import json +from typing import Any, Literal import httpx import pytest -from openarmature.graph.events import EmbeddingEvent, EmbeddingFailedEvent -from openarmature.llm.errors import ProviderInvalidModel, ProviderInvalidResponse -from openarmature.retrieval import OpenAIEmbeddingProvider +from openarmature.graph.events import ( + EmbeddingEvent, + EmbeddingFailedEvent, + RerankEvent, + RerankFailedEvent, +) +from openarmature.graph.observer import ObserverEvent +from openarmature.llm.errors import ( + ProviderInvalidModel, + ProviderInvalidRequest, + ProviderInvalidResponse, + ProviderUnavailable, +) +from openarmature.observability.correlation import _reset_active_dispatch, _set_active_dispatch +from openarmature.retrieval import ( + CohereRerankProvider, + OpenAIEmbeddingProvider, + RerankRuntimeConfig, + ScoredDocument, + validate_rerank_input, + validate_rerank_response, +) def _ok_handler(request: httpx.Request) -> httpx.Response: @@ -185,3 +206,416 @@ async def test_langfuse_observer_embedding_no_op_without_invocation_context() -> await observer(_embedding_event()) await observer(_embedding_failed_event()) assert client.traces == {} + + +# --------------------------------------------------------------------------- +# Rerank capability (proposal 0060) +# --------------------------------------------------------------------------- + + +def _rerank_body( + *, + id: str = "rerank-id", + model: str | None = "rerank-test", + results: list[dict[str, Any]], + search_units: int | None = 1, + input_tokens: int | None = None, +) -> dict[str, Any]: + body: dict[str, Any] = {"id": id, "results": results} + if model is not None: + body["model"] = model + billed: dict[str, Any] = {} + if search_units is not None: + billed["search_units"] = search_units + if input_tokens is not None: + billed["input_tokens"] = input_tokens + if billed: + body["meta"] = {"billed_units": billed} + return body + + +def _rerank_provider( + handler: Any, + *, + model: str = "rerank-test", +) -> CohereRerankProvider: + return CohereRerankProvider( + base_url="http://mock-rerank.test", + model=model, + api_key="test-key", + transport=httpx.MockTransport(handler), + ) + + +# --- validators ------------------------------------------------------------- + + +@pytest.mark.parametrize( + ("query", "documents", "top_k"), + [ + ("", ["a"], None), # empty query + ("q", [], None), # empty documents + ("q", ["a"], 0), # top_k == 0 + ("q", ["a"], -3), # top_k < 0 + ], +) +def test_validate_rerank_input_rejects(query: str, documents: list[str], top_k: int | None) -> None: + with pytest.raises(ProviderInvalidRequest): + validate_rerank_input(query, documents, top_k) + + +def test_validate_rerank_input_allows_top_k_exceeding_documents() -> None: + # top_k MAY exceed len(documents) -- that is allowed, not an error. + validate_rerank_input("q", ["a", "b"], 10) + + +def test_validate_rerank_response_rejects_out_of_range_index() -> None: + with pytest.raises(ProviderInvalidResponse): + validate_rerank_response([ScoredDocument(index=5, relevance_score=0.9)], 2, None) + + +def test_validate_rerank_response_rejects_duplicate_index() -> None: + results = [ + ScoredDocument(index=0, relevance_score=0.9), + ScoredDocument(index=0, relevance_score=0.4), + ] + with pytest.raises(ProviderInvalidResponse): + validate_rerank_response(results, 3, None) + + +def test_validate_rerank_response_rejects_more_results_than_top_k() -> None: + results = [ScoredDocument(index=i, relevance_score=0.5) for i in range(3)] + with pytest.raises(ProviderInvalidResponse): + validate_rerank_response(results, 3, top_k=2) + + +# --- rerank() behavior ------------------------------------------------------ + + +async def test_rerank_sorts_results_and_populates_usage() -> None: + def handler(_req: httpx.Request) -> httpx.Response: + return httpx.Response( + 200, + json=_rerank_body( + results=[ + {"index": 0, "relevance_score": 0.5}, + {"index": 1, "relevance_score": 0.9}, + {"index": 2, "relevance_score": 0.1}, + ], + search_units=1, + ), + ) + + provider = _rerank_provider(handler) + response = await provider.rerank("q", ["berlin", "paris", "madrid"]) + # Provider returned UNSORTED; the adapter sorts by relevance descending. + assert [(r.index, r.relevance_score) for r in response.results] == [(1, 0.9), (0, 0.5), (2, 0.1)] + assert response.model == "rerank-test" + assert response.response_id == "rerank-id" + assert response.usage is not None + assert response.usage.search_units == 1 + assert response.usage.input_tokens is None + assert response.raw["id"] == "rerank-id" + await provider.aclose() + + +async def test_rerank_document_echo_passthrough_and_null() -> None: + def handler(_req: httpx.Request) -> httpx.Response: + return httpx.Response( + 200, + json=_rerank_body( + results=[ + {"index": 0, "relevance_score": 0.9, "document": "alpha doc"}, + {"index": 2, "relevance_score": 0.5, "document": "gamma doc"}, + {"index": 1, "relevance_score": 0.2}, + ], + ), + ) + + provider = _rerank_provider(handler) + response = await provider.rerank("q", ["alpha doc", "beta doc", "gamma doc"]) + by_index = {r.index: r.document for r in response.results} + # Echo preserved verbatim where present; None where omitted (never + # auto-filled from the input documents list). + assert by_index == {0: "alpha doc", 2: "gamma doc", 1: None} + await provider.aclose() + + +async def test_rerank_no_usage_object_yields_null_usage() -> None: + def handler(_req: httpx.Request) -> httpx.Response: + return httpx.Response( + 200, + json=_rerank_body(results=[{"index": 0, "relevance_score": 0.9}], search_units=None), + ) + + provider = _rerank_provider(handler) + response = await provider.rerank("q", ["only"]) + # No usage object surfaced -> usage is None (never a fabricated all-null + # record). + assert response.usage is None + await provider.aclose() + + +async def test_rerank_out_of_range_index_raises() -> None: + def handler(_req: httpx.Request) -> httpx.Response: + return httpx.Response( + 200, + json=_rerank_body( + results=[ + {"index": 0, "relevance_score": 0.9}, + {"index": 5, "relevance_score": 0.4}, + ], + ), + ) + + provider = _rerank_provider(handler) + with pytest.raises(ProviderInvalidResponse): + await provider.rerank("q", ["a", "b"]) + await provider.aclose() + + +async def test_rerank_top_k_maps_to_top_n_and_return_documents_not_sent() -> None: + captured: list[dict[str, Any]] = [] + + def handler(req: httpx.Request) -> httpx.Response: + captured.append(json.loads(req.content)) + return httpx.Response( + 200, + json=_rerank_body(results=[{"index": 0, "relevance_score": 0.9}]), + ) + + provider = _rerank_provider(handler) + # return_documents=True is a silent no-op on the Cohere wire (no such field); + # max_tokens_per_doc rides the extras pass-through bag (model_validate so + # the undeclared extra is accepted, mirroring the conformance config path). + config = RerankRuntimeConfig.model_validate({"return_documents": True, "max_tokens_per_doc": 100}) + await provider.rerank("q", ["a", "b", "c"], top_k=2, config=config) + body = captured[0] + assert body["model"] == "rerank-test" + assert body["query"] == "q" + assert body["documents"] == ["a", "b", "c"] + assert body["top_n"] == 2 + assert "return_documents" not in body + assert body["max_tokens_per_doc"] == 100 + await provider.aclose() + + +# --- typed-event dispatch --------------------------------------------------- + + +def _collecting_dispatch() -> tuple[list[ObserverEvent], Any]: + events: list[ObserverEvent] = [] + + def _dispatch(event: ObserverEvent) -> None: + events.append(event) + + token = _set_active_dispatch(_dispatch) + return events, token + + +async def test_rerank_success_dispatches_rerank_event_with_fields() -> None: + def handler(_req: httpx.Request) -> httpx.Response: + return httpx.Response( + 200, + json=_rerank_body( + results=[ + {"index": 0, "relevance_score": 0.5}, + {"index": 1, "relevance_score": 0.9}, + ], + search_units=3, + ), + ) + + provider = _rerank_provider(handler) + events, token = _collecting_dispatch() + try: + await provider.rerank("q", ["a", "b"], top_k=2, config=RerankRuntimeConfig(return_documents=True)) + finally: + await provider.aclose() + _reset_active_dispatch(token) + + rerank_events = [e for e in events if isinstance(e, RerankEvent)] + failed = [e for e in events if isinstance(e, RerankFailedEvent)] + assert len(rerank_events) == 1 + assert failed == [] # mutually exclusive + event = rerank_events[0] + assert event.provider == "cohere" + assert event.model == "rerank-test" + assert event.query == "q" + assert event.documents == ["a", "b"] + assert event.document_count == 2 + assert event.top_k == 2 + assert event.result_count == 2 + # output_results populated unconditionally on success (proposal 0089). + assert [(r.index, r.relevance_score) for r in event.output_results] == [(1, 0.9), (0, 0.5)] + assert event.usage is not None + assert event.usage.search_units == 3 + # return_documents was explicitly supplied -> it appears in request_params + # even though it is a wire no-op. + assert event.request_params == {"return_documents": True} + assert event.call_id + + +async def test_rerank_request_params_empty_when_no_config() -> None: + def handler(_req: httpx.Request) -> httpx.Response: + return httpx.Response(200, json=_rerank_body(results=[{"index": 0, "relevance_score": 0.9}])) + + provider = _rerank_provider(handler) + events, token = _collecting_dispatch() + try: + await provider.rerank("q", ["a"]) + finally: + await provider.aclose() + _reset_active_dispatch(token) + + event = next(e for e in events if isinstance(e, RerankEvent)) + assert event.request_params == {} + + +async def test_rerank_failure_dispatches_rerank_failed_event_only() -> None: + def handler(_req: httpx.Request) -> httpx.Response: + return httpx.Response(503, json={"message": "down"}) + + provider = _rerank_provider(handler) + events, token = _collecting_dispatch() + try: + with pytest.raises(ProviderUnavailable): + await provider.rerank("q", ["a", "b"]) + finally: + await provider.aclose() + _reset_active_dispatch(token) + + failed = [e for e in events if isinstance(e, RerankFailedEvent)] + success = [e for e in events if isinstance(e, RerankEvent)] + assert len(failed) == 1 + assert success == [] # mutually exclusive + event = failed[0] + assert event.error_category == "provider_unavailable" + assert event.error_type == "ProviderUnavailable" + assert event.query == "q" + assert event.documents == ["a", "b"] + assert event.document_count == 2 + assert event.call_id + + +async def test_rerank_invalid_request_dispatches_failed_event_before_send() -> None: + # An empty query fails pre-send validation; the RerankFailedEvent still + # fires (dispatched alongside the exception, not in place of it). + def handler(_req: httpx.Request) -> httpx.Response: # pragma: no cover - never called + raise AssertionError("provider must not be contacted on pre-send validation failure") + + provider = _rerank_provider(handler) + events, token = _collecting_dispatch() + try: + with pytest.raises(ProviderInvalidRequest): + await provider.rerank("", ["a"]) + finally: + await provider.aclose() + _reset_active_dispatch(token) + + failed = [e for e in events if isinstance(e, RerankFailedEvent)] + assert len(failed) == 1 + assert failed[0].error_category == "provider_invalid_request" + + +# --- construction guards + readiness ---------------------------------------- + + +def test_rerank_base_url_rejects_v2_suffix() -> None: + with pytest.raises(ValueError, match="host root"): + CohereRerankProvider(base_url="https://api.cohere.com/v2", model="m") + # The host root is accepted (no doubled /v2/v2). + CohereRerankProvider(base_url="https://api.cohere.com", model="m") + + +async def test_rerank_ready_probe_surfaces_invalid_model_on_404() -> None: + def handler(_req: httpx.Request) -> httpx.Response: + return httpx.Response(404, json={"message": "model not found"}) + + provider = _rerank_provider(handler, model="nonexistent") + with pytest.raises(ProviderInvalidModel): + await provider.ready() + await provider.aclose() + + +def _rerank_event() -> RerankEvent: + return RerankEvent( + invocation_id="inv", + correlation_id=None, + node_name="rerank_node", + namespace=("rerank_node",), + attempt_index=0, + fan_out_index=None, + branch_name=None, + provider="cohere", + model="rerank-test", + response_id=None, + response_model=None, + usage=None, + latency_ms=1.0, + query="q", + documents=["a"], + document_count=1, + top_k=None, + result_count=1, + output_results=[ScoredDocument(index=0, relevance_score=0.9)], + request_params={}, + request_extras={}, + active_prompt=None, + active_prompt_group=None, + call_id="c", + ) + + +def _rerank_failed_event() -> RerankFailedEvent: + return RerankFailedEvent( + invocation_id="inv", + correlation_id=None, + node_name="rerank_node", + namespace=("rerank_node",), + attempt_index=0, + fan_out_index=None, + branch_name=None, + provider="cohere", + model="rerank-test", + latency_ms=1.0, + query="q", + documents=["a"], + document_count=1, + top_k=None, + request_params={}, + request_extras={}, + active_prompt=None, + active_prompt_group=None, + call_id="c", + error_category="provider_unavailable", + error_message="boom", + ) + + +async def test_otel_observer_rerank_no_op_without_invocation_context() -> None: + from opentelemetry.sdk.trace.export import SimpleSpanProcessor + from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter + + from openarmature.observability.otel import OTelObserver + + exporter = InMemorySpanExporter() + observer = OTelObserver(span_processor=SimpleSpanProcessor(exporter)) + # The rerank span rendering lands in 0060b; for now the observer safely + # skips the rerank events rather than falling through to the NodeEvent + # phase dispatch (which would raise on the missing ``phase`` attribute). + await observer(_rerank_event()) + await observer(_rerank_failed_event()) + assert len(exporter.get_finished_spans()) == 0 + + +async def test_langfuse_observer_rerank_no_op_without_invocation_context() -> None: + from openarmature.observability.langfuse import InMemoryLangfuseClient, LangfuseObserver + + client = InMemoryLangfuseClient() + observer = LangfuseObserver(client=client) + # The Langfuse Retriever observation lands in 0060b; for now the observer + # safely skips the rerank events. + await observer(_rerank_event()) + await observer(_rerank_failed_event()) + assert client.traces == {}