From 342f2972fa771c59f68a0a8a981824fbe867408d Mon Sep 17 00:00:00 2001 From: chris-colinsky Date: Sat, 27 Jun 2026 20:22:46 -0700 Subject: [PATCH 1/4] Add retrieval-provider embedding capability Add the embedding surface of the new retrieval-provider capability: an EmbeddingProvider protocol, the EmbeddingResponse / EmbeddingUsage / EmbeddingRuntimeConfig types, and an OpenAIEmbeddingProvider reference impl posting to an OpenAI-compatible /v1/embeddings. Typed EmbeddingEvent / EmbeddingFailedEvent join the observer event union, dispatched on every embed(); the bundled OTel and Langfuse observers safely skip them for now (the embedding span and Embedding observation are a follow-up). The provider rejects malformed responses (missing, empty, or non-numeric vectors and non-permutation indices), guards a trailing /v1 on base_url, and offers a configurable readiness probe defaulting to a universal one-input embed so it works against backends without a /v1/models catalog (e.g. TEI's OpenAI surface). Conformance fixtures 001-005 pass; unit tests cover the guard, the readiness modes, and the observer safe-handling. conformance.toml keeps 0059 not-yet; the observability rendering (074-083), the cross-provider helper cleanups, and the docs are v0.16.0 follow-ups. --- src/openarmature/graph/events.py | 122 +++++ src/openarmature/graph/observer.py | 8 +- src/openarmature/observability/correlation.py | 121 ++--- .../observability/langfuse/observer.py | 10 + .../observability/otel/observer.py | 11 + src/openarmature/retrieval/__init__.py | 26 + src/openarmature/retrieval/provider.py | 101 ++++ .../retrieval/providers/__init__.py | 7 + .../retrieval/providers/openai.py | 462 ++++++++++++++++++ src/openarmature/retrieval/response.py | 94 ++++ tests/conformance/test_retrieval_provider.py | 205 ++++++++ tests/unit/test_retrieval_provider.py | 162 ++++++ 12 files changed, 1236 insertions(+), 93 deletions(-) create mode 100644 src/openarmature/retrieval/__init__.py create mode 100644 src/openarmature/retrieval/provider.py create mode 100644 src/openarmature/retrieval/providers/__init__.py create mode 100644 src/openarmature/retrieval/providers/openai.py create mode 100644 src/openarmature/retrieval/response.py create mode 100644 tests/conformance/test_retrieval_provider.py create mode 100644 tests/unit/test_retrieval_provider.py diff --git a/src/openarmature/graph/events.py b/src/openarmature/graph/events.py index 79bb4ec..c6eddcf 100644 --- a/src/openarmature/graph/events.py +++ b/src/openarmature/graph/events.py @@ -32,6 +32,7 @@ if TYPE_CHECKING: from openarmature.llm.messages import ToolCall from openarmature.llm.response import Usage + from openarmature.retrieval.response import EmbeddingUsage # Sentinel empty metadata mapping for events constructed without a # live caller-metadata snapshot (test helpers, synthetic events). @@ -749,6 +750,125 @@ class LlmRetryAttemptEvent: output_tool_calls: list["ToolCall"] = field(default_factory=list["ToolCall"]) +# Spec: realizes graph-engine §6 + observability §5.5.9 -- the typed +# EmbeddingEvent / EmbeddingFailedEvent pair (proposal 0059, +# retrieval-provider capability). Dispatched on the observer delivery +# queue per EmbeddingProvider.embed() call: the success variant after the +# response is parsed + validated, the failure variant alongside a raised +# §7 category exception (mutually exclusive per call). Scalar +# fan_out_index / branch_name only; the lineage chains arrive uniformly +# across the provider events with proposal 0084 (v0.81.0). input_strings / +# request_extras are payload-bearing, populated unconditionally; observer- +# side privacy gates (OTel disable_provider_payload, Langfuse equivalents) +# apply at rendering, symmetric with LlmCompletionEvent. +@dataclass(frozen=True) +class EmbeddingEvent: + """A typed embedding provider call event delivered to observers. + + Carries identity, scoping, and outcome data for a successful + ``EmbeddingProvider.embed()`` call. Observer code filters by type + discrimination (``isinstance(event, EmbeddingEvent)``). + + The identity / scoping / request-side fields mirror + ``LlmCompletionEvent``'s convention; the outcome fields are + embedding-specific: + + - ``input_strings``: the input strings the call was made with; + non-nullable, populated unconditionally (privacy gating is + observer-side at rendering). + - ``input_count``: ``len(input_strings)``; a convenience field. + - ``dimensions``: the output vector dimensionality from the response; + ``None`` when the response surfaced no determinate dimensionality. + - ``response_model`` / ``response_id``: the provider-returned model + and response identifiers; ``None`` when the provider returned none. + - ``usage``: the embedding token record; ``None`` when the call + returned no usage. + - ``request_params``: the embedding request parameters the caller + supplied (e.g. ``dimensions``). 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 embed-call time; ``None`` outside a binding. + - ``call_id``: a per-call disambiguator, always present, freshly + minted per ``embed()`` 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 + # EmbeddingUsage is a string-typed forward reference per the + # TYPE_CHECKING import -- keeps the runtime import direction + # graph -> retrieval off the module-load path. + usage: "EmbeddingUsage | None" + latency_ms: float | None + input_strings: list[str] + input_count: int + dimensions: int | None + 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 EmbeddingEvent (proposal 0059). Dispatched +# whenever EmbeddingProvider.embed() raises a §7 category exception -- +# covers both the provider-caught path and the pre-send validation raise +# (provider_invalid_request on an empty input list). Dispatched ALONGSIDE +# the exception, not in place of it; mutually exclusive with EmbeddingEvent +# on the same call. The response-side fields are absent (no response). +@dataclass(frozen=True) +class EmbeddingFailedEvent: + """A typed embedding provider call failure event delivered to observers. + + Carries identity, scoping, and failure-context data for an ``embed()`` + call that raised a retrieval-provider category exception. Observer code + filters by type discrimination + (``isinstance(event, EmbeddingFailedEvent)``). + + The identity / scoping / request-side field set mirrors + ``EmbeddingEvent``; the response-side fields are absent. Failure- + specific fields: + + - ``error_category``: the error category the call raised (one of the + embedding-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 + input_strings: list[str] + 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 @@ -905,6 +1025,8 @@ class ToolCallFailedEvent: __all__ = [ + "EmbeddingEvent", + "EmbeddingFailedEvent", "FailureIsolatedEvent", "FanOutEventConfig", "InvocationCompletedEvent", diff --git a/src/openarmature/graph/observer.py b/src/openarmature/graph/observer.py index 0f583c0..c4fea6a 100644 --- a/src/openarmature/graph/observer.py +++ b/src/openarmature/graph/observer.py @@ -35,6 +35,8 @@ from typing import Any, Literal, Protocol from .events import ( + EmbeddingEvent, + EmbeddingFailedEvent, FailureIsolatedEvent, InvocationCompletedEvent, InvocationStartedEvent, @@ -63,7 +65,9 @@ # retry to drive the per-attempt OTel span surface), # and FailureIsolatedEvent (proposal 0050 §6.3 framework-emitted event, # dispatched by FailureIsolationMiddleware when it catches an exception -# escaping the inner chain and substitutes a degraded partial update). +# 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()). ObserverEvent = ( NodeEvent | MetadataAugmentationEvent @@ -75,6 +79,8 @@ | FailureIsolatedEvent | ToolCallEvent | ToolCallFailedEvent + | EmbeddingEvent + | EmbeddingFailedEvent ) diff --git a/src/openarmature/observability/correlation.py b/src/openarmature/observability/correlation.py index 442ff4c..bc01cd0 100644 --- a/src/openarmature/observability/correlation.py +++ b/src/openarmature/observability/correlation.py @@ -37,6 +37,8 @@ if TYPE_CHECKING: from openarmature.graph.events import ( + EmbeddingEvent, + EmbeddingFailedEvent, FailureIsolatedEvent, InvocationCompletedEvent, InvocationStartedEvent, @@ -50,6 +52,25 @@ ) from openarmature.graph.observer import SubscribedObserver + # The event-record union the engine's serial dispatch worker accepts. + # Defined once and reused across the ContextVar + the get/set/reset + # accessors so a new capability event widens the contract in one place. + _DispatchEvent = ( + NodeEvent + | MetadataAugmentationEvent + | InvocationStartedEvent + | InvocationCompletedEvent + | LlmCompletionEvent + | LlmFailedEvent + | LlmRetryAttemptEvent + | FailureIsolatedEvent + | ToolCallEvent + | ToolCallFailedEvent + | EmbeddingEvent + | EmbeddingFailedEvent + ) + _DispatchFn = Callable[[_DispatchEvent], None] + # --------------------------------------------------------------------------- # Correlation ID (observability spec §3.1) @@ -221,44 +242,12 @@ def _reset_active_observers(token: Token[tuple[SubscribedObserver, ...]]) -> Non # --------------------------------------------------------------------------- -_active_dispatch_var: ContextVar[ - Callable[ - [ - NodeEvent - | MetadataAugmentationEvent - | InvocationStartedEvent - | InvocationCompletedEvent - | LlmCompletionEvent - | LlmFailedEvent - | LlmRetryAttemptEvent - | FailureIsolatedEvent - | ToolCallEvent - | ToolCallFailedEvent - ], - None, - ] - | None -] = ContextVar("openarmature.active_dispatch", default=None) - - -def current_dispatch() -> ( - Callable[ - [ - NodeEvent - | MetadataAugmentationEvent - | InvocationStartedEvent - | InvocationCompletedEvent - | LlmCompletionEvent - | LlmFailedEvent - | LlmRetryAttemptEvent - | FailureIsolatedEvent - | ToolCallEvent - | ToolCallFailedEvent - ], - None, - ] - | None -): +_active_dispatch_var: ContextVar[_DispatchFn | None] = ContextVar( + "openarmature.active_dispatch", default=None +) + + +def current_dispatch() -> _DispatchFn | None: """Return the engine's dispatch callable for the current invocation, or ``None`` outside any invocation. @@ -272,65 +261,13 @@ def current_dispatch() -> ( return _active_dispatch_var.get() -def _set_active_dispatch( - dispatch: Callable[ - [ - NodeEvent - | MetadataAugmentationEvent - | InvocationStartedEvent - | InvocationCompletedEvent - | LlmCompletionEvent - | LlmFailedEvent - | LlmRetryAttemptEvent - | FailureIsolatedEvent - | ToolCallEvent - | ToolCallFailedEvent - ], - None, - ], -) -> Token[ - Callable[ - [ - NodeEvent - | MetadataAugmentationEvent - | InvocationStartedEvent - | InvocationCompletedEvent - | LlmCompletionEvent - | LlmFailedEvent - | LlmRetryAttemptEvent - | FailureIsolatedEvent - | ToolCallEvent - | ToolCallFailedEvent - ], - None, - ] - | None -]: +def _set_active_dispatch(dispatch: _DispatchFn) -> Token[_DispatchFn | None]: """Set the engine's dispatch callable in scope. Internal — engine-only.""" return _active_dispatch_var.set(dispatch) -def _reset_active_dispatch( - token: Token[ - Callable[ - [ - NodeEvent - | MetadataAugmentationEvent - | InvocationStartedEvent - | InvocationCompletedEvent - | LlmCompletionEvent - | LlmFailedEvent - | LlmRetryAttemptEvent - | FailureIsolatedEvent - | ToolCallEvent - | ToolCallFailedEvent - ], - None, - ] - | None - ], -) -> None: +def _reset_active_dispatch(token: Token[_DispatchFn | None]) -> None: _active_dispatch_var.reset(token) diff --git a/src/openarmature/observability/langfuse/observer.py b/src/openarmature/observability/langfuse/observer.py index 7ab8139..3d15a2c 100644 --- a/src/openarmature/observability/langfuse/observer.py +++ b/src/openarmature/observability/langfuse/observer.py @@ -30,6 +30,8 @@ from typing import Any, cast from openarmature.graph.events import ( + EmbeddingEvent, + EmbeddingFailedEvent, FailureIsolatedEvent, InvocationCompletedEvent, InvocationStartedEvent, @@ -443,6 +445,8 @@ async def __call__( | FailureIsolatedEvent | ToolCallEvent | ToolCallFailedEvent + | EmbeddingEvent + | EmbeddingFailedEvent ), ) -> None: if isinstance(event, InvocationStartedEvent): @@ -484,6 +488,12 @@ async def __call__( if isinstance(event, MetadataAugmentationEvent): self._handle_metadata_augmentation(event) return + # Proposal 0059 embedding events: the bundled Langfuse Embedding + # observation is a follow-up. Until it lands the events are safely + # ignored here rather than falling through to the NodeEvent phase + # dispatch (which would AttributeError on the absent ``phase``). + if isinstance(event, EmbeddingEvent | EmbeddingFailedEvent): + 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 5a5b89b..825488e 100644 --- a/src/openarmature/observability/otel/observer.py +++ b/src/openarmature/observability/otel/observer.py @@ -100,6 +100,8 @@ from opentelemetry.trace.propagation import set_span_in_context from openarmature.graph.events import ( + EmbeddingEvent, + EmbeddingFailedEvent, FailureIsolatedEvent, InvocationCompletedEvent, InvocationStartedEvent, @@ -755,6 +757,8 @@ async def __call__( | FailureIsolatedEvent | ToolCallEvent | ToolCallFailedEvent + | EmbeddingEvent + | EmbeddingFailedEvent ), ) -> None: # Proposal 0043 invocation-boundary events: OTel has no @@ -782,6 +786,13 @@ async def __call__( # consumers, so the OTel observer ignores them here. if isinstance(event, LlmCompletionEvent | LlmFailedEvent): return + # Proposal 0059 embedding events: the bundled OTel embedding span + # (openarmature.embedding.complete) is a follow-up. Until it lands + # the events are safely ignored here rather than falling through to + # the NodeEvent phase dispatch (which would AttributeError on the + # absent ``phase`` field). + if isinstance(event, EmbeddingEvent | EmbeddingFailedEvent): + 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 new file mode 100644 index 0000000..691be52 --- /dev/null +++ b/src/openarmature/retrieval/__init__.py @@ -0,0 +1,26 @@ +"""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. +""" + +from __future__ import annotations + +from .provider import ( + EmbeddingProvider, + validate_embedding_input, + validate_embedding_response, +) +from .providers.openai import OpenAIEmbeddingProvider +from .response import EmbeddingResponse, EmbeddingRuntimeConfig, EmbeddingUsage + +__all__ = [ + "EmbeddingProvider", + "EmbeddingResponse", + "EmbeddingRuntimeConfig", + "EmbeddingUsage", + "OpenAIEmbeddingProvider", + "validate_embedding_input", + "validate_embedding_response", +] diff --git a/src/openarmature/retrieval/provider.py b/src/openarmature/retrieval/provider.py new file mode 100644 index 0000000..12c5bcb --- /dev/null +++ b/src/openarmature/retrieval/provider.py @@ -0,0 +1,101 @@ +# 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. + +"""EmbeddingProvider Protocol + 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: + +- ``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 +``provider_invalid_response`` respectively. +""" + +from __future__ import annotations + +from collections.abc import Sequence +from typing import Protocol + +from openarmature.llm.errors import ProviderInvalidRequest, ProviderInvalidResponse + +from .response import EmbeddingResponse, EmbeddingRuntimeConfig + + +class EmbeddingProvider(Protocol): + """The shape of any retrieval-provider embedding implementation. + + Implementations are bound to a single embedding model identifier; + switching models means constructing a new provider, not passing a + different argument per call. + """ + + async def ready(self) -> None: + """Verify the bound embedding model is reachable and serving.""" + ... + + async def embed( + self, + input: Sequence[str], + *, + config: EmbeddingRuntimeConfig | None = None, + ) -> EmbeddingResponse: + """Embed ``input`` into one vector per string, in input order. + + Args: + input: The strings to embed. Always a list, even for a + single-string caller (wrap as a one-element list). Not + mutated by the implementation. + config: Optional per-call request parameters. + + Returns an :class:`EmbeddingResponse` whose ``vectors[i]`` is the + embedding of ``input[i]`` -- order is preserved, never permuted. + """ + ... + + +def validate_embedding_input(input: Sequence[str]) -> None: + """Validate the input list before sending. + + Raises :class:`ProviderInvalidRequest` when the input list is empty. + """ + if not input: + raise ProviderInvalidRequest("embedding input list must not be empty") + + +def validate_embedding_response( + vectors: Sequence[Sequence[float]], + input_count: int, +) -> int: + """Validate the response invariants and return the dimensionality. + + Raises :class:`ProviderInvalidResponse` when the vector count does not + match the input count, when the response carries no vectors, or when + the vectors are not all the same length. Returns the (consistent) + dimensionality on success. + """ + if len(vectors) != input_count: + raise ProviderInvalidResponse(f"provider returned {len(vectors)} vectors for {input_count} inputs") + if not vectors: + raise ProviderInvalidResponse("provider returned no vectors") + dimensions = len(vectors[0]) + for vector in vectors: + if len(vector) != dimensions: + raise ProviderInvalidResponse("provider returned vectors with inconsistent dimensionality") + return dimensions + + +__all__ = [ + "EmbeddingProvider", + "validate_embedding_input", + "validate_embedding_response", +] diff --git a/src/openarmature/retrieval/providers/__init__.py b/src/openarmature/retrieval/providers/__init__.py new file mode 100644 index 0000000..477878f --- /dev/null +++ b/src/openarmature/retrieval/providers/__init__.py @@ -0,0 +1,7 @@ +"""Bundled retrieval-provider reference implementations.""" + +from __future__ import annotations + +from .openai import OpenAIEmbeddingProvider + +__all__ = ["OpenAIEmbeddingProvider"] diff --git a/src/openarmature/retrieval/providers/openai.py b/src/openarmature/retrieval/providers/openai.py new file mode 100644 index 0000000..3459181 --- /dev/null +++ b/src/openarmature/retrieval/providers/openai.py @@ -0,0 +1,462 @@ +# Spec: realizes the retrieval-provider §3 EmbeddingProvider protocol +# against an OpenAI-compatible POST /v1/embeddings endpoint -- the +# reference embedding provider. The wire mapping proposal 0079 formalizes +# (base_url override, encoding_format) is pre-satisfied here; 0079 pins +# the remaining wire details + its fixtures 023-027. The §7 error +# categories are shared with llm-provider; the embedding-applicable subset +# (no unsupported_content_block, no structured_output_invalid) is mapped +# from the OpenAI-shape HTTP error envelope below. Typed EmbeddingEvent / +# EmbeddingFailedEvent dispatch mirrors the llm-provider §6 path via +# current_dispatch(). FOLLOW-UP: classify_http_error / base_url +# normalization are duplicated in spirit from llm.providers.openai; +# lifting a shared OpenAI-shape HTTP helper is a multi-provider follow-on. + +"""OpenAI-compatible embedding provider. + +``OpenAIEmbeddingProvider`` issues ``POST {base_url}/v1/embeddings`` and +parses the OpenAI ``{data: [{index, embedding}], model, usage}`` envelope +into an :class:`EmbeddingResponse`. ``base_url`` is the host root (the +provider appends ``/v1/embeddings`` and ``/v1/models``), overridable for +any OpenAI-compatible backend (vLLM, LocalAI, TEI's OpenAI surface). The +``transport`` parameter is the test seam (``httpx.MockTransport``). +""" + +from __future__ import annotations + +import time +import uuid +from collections.abc import Mapping, Sequence +from typing import Any, Literal, cast + +import httpx + +from openarmature.graph.events import EmbeddingEvent, EmbeddingFailedEvent +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_embedding_input, validate_embedding_response +from ..response import EmbeddingResponse, EmbeddingRuntimeConfig, EmbeddingUsage + + +def _classify_embedding_http_error(resp: httpx.Response) -> LlmProviderError: + """Map a non-200 OpenAI-shape embeddings response to a §7 category. + + The embedding-applicable subset: 401/403 -> auth, 429 -> rate_limit, + 404 -> invalid_model, 400/422 -> invalid_request, other 5xx -> + 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 {} + 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_raw = error_block.get("message") + message = message_raw if isinstance(message_raw, str) else None + + if status in (401, 403): + return ProviderAuthentication(message or f"HTTP {status}") + if status == 429: + # The Retry-After surface is wired when a 429 embedding fixture + # lands; unfixtured at this pin, so retry_after stays unset. + 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 zero. The input_type field joins this +# set with proposal 0077. For embedding the event-param names coincide with +# the wire-body keys, so the same dict feeds both the event and the body. +def _request_params_from_config(config: EmbeddingRuntimeConfig | None) -> dict[str, Any]: + """Extract the supplied embedding request parameters for the event.""" + if config is None: + return {} + out: dict[str, Any] = {} + if config.dimensions is not None: + out["dimensions"] = config.dimensions + return out + + +_VALID_READINESS_PROBES = frozenset({"embed", "models", "both"}) + + +class OpenAIEmbeddingProvider: + """OpenAI ``/v1/embeddings`` wire-compatible embedding provider. + + Construct with a base URL (host root), the bound embedding model, and + an optional API key + transport. ``embed()`` posts to + ``/v1/embeddings``. + + ``ready()`` verifies the bound model per the ``readiness_probe`` + argument: + + - ``"embed"`` (default): a one-input ``/v1/embeddings`` probe. Works + against any OpenAI-compatible backend, including ones that do not + serve the ``/v1/models`` catalog (e.g. TEI's OpenAI surface). + - ``"models"``: a ``GET /v1/models`` catalog check. Cheaper (no embed + billed), but requires the endpoint to serve the catalog. + - ``"both"``: the catalog check, then the embed 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 = "openai", + readiness_probe: Literal["embed", "models", "both"] = "embed", + populate_caller_metadata: bool = True, + ) -> None: + # base_url is the host root; the provider appends the /v1 routes, so + # a trailing /v1 would produce a doubled /v1/v1 path that 404s (the + # sibling llm provider guards the same footgun). Trailing slashes are + # stripped. Proposal 0079 pins the full base_url-override contract + # (fixture 025); lifting a shared base_url helper is a follow-up. + normalized = base_url.rstrip("/") + if normalized.endswith("/v1"): + raise ValueError( + f"base_url should be the host root (e.g. 'https://api.openai.com'); " + f"the provider appends /v1/embeddings and /v1/models itself, so a " + f"trailing /v1 would produce a doubled /v1/v1 path. Got {base_url!r}." + ) + self.base_url = normalized + self.model = model + # ``readiness_probe`` modes are documented on the class docstring; + # the default "embed" is the universal probe. Reject an unknown mode. + if readiness_probe not in _VALID_READINESS_PROBES: + raise ValueError( + f"readiness_probe must be one of {sorted(_VALID_READINESS_PROBES)} (got {readiness_probe!r})" + ) + self._readiness_probe = readiness_probe + # ``genai_system`` surfaces as gen_ai.system on the embedding 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 embedding model is reachable and serving.""" + # Dispatches on readiness_probe (see __init__): "embed" issues a + # minimal /v1/embeddings call (universal across backends), "models" + # the /v1/models catalog check, "both" runs catalog then embed. + if self._readiness_probe in ("models", "both"): + await self._probe_models() + if self._readiness_probe in ("embed", "both"): + await self._probe_embed() + + async def _probe_embed(self) -> None: + # The universal probe: a one-input embed against /v1/embeddings. + # Surfaces provider_invalid_model (404) / provider_unavailable (5xx) + # / provider_authentication exactly as a real embed() would, against + # any OpenAI-compatible backend (OpenAI, vLLM, TEI's OpenAI surface). + body = {"model": self.model, "input": ["ready"]} + try: + resp = await self._client.post("/v1/embeddings", json=body) + except httpx.HTTPError as exc: + raise ProviderUnavailable(f"readiness probe failed: {exc}") from exc + if resp.status_code != 200: + raise _classify_embedding_http_error(resp) + + async def _probe_models(self) -> None: + # Catalog probe: GET /v1/models + bound-model presence check. + # Cheaper than _probe_embed (no embed billed) but requires the + # endpoint to serve the /v1/models catalog -- OpenAI / vLLM do, TEI + # does not. + try: + resp = await self._client.get("/v1/models") + except httpx.HTTPError as exc: + raise ProviderUnavailable(f"readiness probe failed: {exc}") from exc + if resp.status_code != 200: + raise _classify_embedding_http_error(resp) + try: + body_raw = resp.json() + except ValueError as exc: + raise ProviderInvalidResponse("models response is not valid JSON") from exc + body = cast("dict[str, Any]", body_raw) if isinstance(body_raw, dict) else {} + data_raw = body.get("data") + models = cast("list[Any]", data_raw) if isinstance(data_raw, list) else [] + ids = [cast("dict[str, Any]", m).get("id") for m in models if isinstance(m, dict)] + if self.model not in ids: + seen = sorted(i for i in ids if isinstance(i, str)) + raise ProviderInvalidModel(f"model {self.model!r} not in catalog (seen: {seen})") + + async def embed( + self, + input: Sequence[str], + *, + config: EmbeddingRuntimeConfig | None = None, + ) -> EmbeddingResponse: + """Embed ``input`` into one vector per string, in input order.""" + 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() + input_strings = list(input) + 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_embedding_input(input_strings) + body = self._build_request_body(input_strings, request_params, request_extras) + try: + resp = await self._client.post("/v1/embeddings", json=body) + except httpx.HTTPError as exc: + raise ProviderUnavailable(f"embedding request failed: {exc}") from exc + if resp.status_code != 200: + raise _classify_embedding_http_error(resp) + response = self._parse_response(resp, input_strings) + except LlmProviderError as exc: + latency_ms_failed = (time.perf_counter() - adapter_start) * 1000.0 + if dispatch is not None: + dispatch( + self._build_embedding_failed_event( + exc, + latency_ms_failed, + call_id=call_id, + input_strings=input_strings, + 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_embedding_event( + response, + latency_ms, + call_id=call_id, + input_strings=input_strings, + 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, + input_strings: list[str], + request_params: dict[str, Any], + request_extras: dict[str, Any], + ) -> dict[str, Any]: + """Build the /v1/embeddings request body.""" + # model + input always; the supplied request params (the same dict + # the event carries -- dimensions for embedding) and the + # runtime-config extras pass-through merge through. Sharing the one + # params dict keeps the wire body and the observed request_params + # provably identical. Proposal 0077's input_type maps to a per-vendor + # wire key at the wire layer. + return {"model": self.model, "input": input_strings, **request_params, **request_extras} + + def _parse_response( + self, + resp: httpx.Response, + input_strings: list[str], + ) -> EmbeddingResponse: + """Parse the OpenAI embeddings envelope into an EmbeddingResponse.""" + # Orders ``data`` by the per-entry ``index`` so output position + # matches input position, then validates the §4 invariants (count + + # consistent dimensionality, via validate_embedding_response). The + # index set MUST be a 0..n-1 permutation so each vector maps to + # exactly one input, and every entry MUST carry a non-empty numeric + # embedding -- a missing / empty / non-numeric vector is a malformed + # response, not a zero-dim result. + try: + body_raw = resp.json() + except ValueError as exc: + raise ProviderInvalidResponse("embedding response is not valid JSON") from exc + if not isinstance(body_raw, dict): + raise ProviderInvalidResponse("embedding response is not a JSON object") + body = cast("dict[str, Any]", body_raw) + data_raw = body.get("data") + if not isinstance(data_raw, list): + raise ProviderInvalidResponse("embedding response missing 'data' array") + data = cast("list[Any]", data_raw) + entries: list[dict[str, Any]] = [] + indices: list[int] = [] + for raw_entry in data: + if not isinstance(raw_entry, dict): + raise ProviderInvalidResponse("embedding response entry is not a JSON object") + entry = cast("dict[str, Any]", raw_entry) + index = entry.get("index") + if not isinstance(index, int): + raise ProviderInvalidResponse("embedding response entry missing integer 'index'") + entries.append(entry) + indices.append(index) + if sorted(indices) != list(range(len(entries))): + raise ProviderInvalidResponse("embedding response 'index' values are not a 0..n-1 permutation") + ordered = sorted(entries, key=lambda e: cast("int", e["index"])) + vectors: list[list[float]] = [] + for entry in ordered: + embedding = entry.get("embedding") + if not isinstance(embedding, list) or not embedding: + raise ProviderInvalidResponse("embedding response entry has a missing or empty 'embedding'") + try: + vectors.append([float(x) for x in cast("list[Any]", embedding)]) + except (TypeError, ValueError) as exc: + raise ProviderInvalidResponse("embedding response has a non-numeric vector value") from exc + dimensions = validate_embedding_response(vectors, len(input_strings)) + usage_block = body.get("usage") + prompt_tokens = ( + cast("dict[str, Any]", usage_block).get("prompt_tokens") + if isinstance(usage_block, dict) + else None + ) + input_tokens = prompt_tokens if isinstance(prompt_tokens, int) and prompt_tokens >= 0 else 0 + response_id = body.get("id") + model = body.get("model") + return EmbeddingResponse( + vectors=vectors, + model=model if isinstance(model, str) else self.model, + usage=EmbeddingUsage(input_tokens=input_tokens), + response_id=response_id if isinstance(response_id, str) else None, + dimensions=dimensions, + raw=body, + ) + + def _build_embedding_event( + self, + response: EmbeddingResponse, + latency_ms: float, + *, + call_id: str, + input_strings: list[str], + request_params: dict[str, Any], + request_extras: dict[str, Any], + active_prompt: Any, + active_prompt_group: Any, + ) -> EmbeddingEvent: + """Construct the typed EmbeddingEvent 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 EmbeddingEvent( + 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, + input_strings=input_strings, + input_count=len(input_strings), + dimensions=response.dimensions, + 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_embedding_failed_event( + self, + exc: LlmProviderError, + latency_ms: float, + *, + call_id: str, + input_strings: list[str], + request_params: dict[str, Any], + request_extras: dict[str, Any], + active_prompt: Any, + active_prompt_group: Any, + ) -> EmbeddingFailedEvent: + """Construct the typed EmbeddingFailedEvent 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 EmbeddingFailedEvent( + 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, + input_strings=input_strings, + 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__ = ["OpenAIEmbeddingProvider"] diff --git a/src/openarmature/retrieval/response.py b/src/openarmature/retrieval/response.py new file mode 100644 index 0000000..3dfb0ae --- /dev/null +++ b/src/openarmature/retrieval/response.py @@ -0,0 +1,94 @@ +# 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. + +The ``EmbeddingResponse`` is what ``EmbeddingProvider.embed()`` returns: +one vector per input string in input order, the model identifier, usage, +the optional provider response id, the output dimensionality, and the +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. +""" + +from __future__ import annotations + +from typing import Any + +from pydantic import BaseModel, ConfigDict, Field + + +class EmbeddingUsage(BaseModel): + """Token-accounting record for an embedding call. + + Carries ``input_tokens`` only -- an embedding call has no output + tokens (vectors are not tokens). + """ + + model_config = ConfigDict(extra="forbid") + + input_tokens: int = Field(ge=0) + + +class EmbeddingResponse(BaseModel): + """The result of an ``EmbeddingProvider.embed()`` call. + + Attributes: + vectors: One vector (a list of floats) per input string, in the + order the inputs were supplied. ``len(vectors)`` equals the + input length. + model: The model identifier the provider returned; may be more + specific than the bound identifier. + usage: The token record. + response_id: The provider-returned response id when present; + ``None`` otherwise. + dimensions: The output vector dimensionality; equals the length + of each inner vector. + raw: The parsed provider response, populated on every successful + return. Carries everything the provider returned. + """ + + model_config = ConfigDict(extra="forbid") + + vectors: list[list[float]] + model: str + usage: EmbeddingUsage + response_id: str | None = None + dimensions: int + raw: dict[str, Any] + + +# Spec §2 declared-field surface: an optional ``dimensions`` plus the +# extras pass-through bag (``extra="allow"``). Undeclared fields supplied +# by callers are forwarded to the wire body untouched by the §8 wire-format +# mapping; declared fields with value ``None`` are omitted on the wire. +class EmbeddingRuntimeConfig(BaseModel): + """Per-call embedding request parameters.""" + + model_config = ConfigDict(extra="allow") + + dimensions: int | None = None + + # Pure ergonomic, not a contract: lets callers splat a dict whose + # entries may be ``None`` without filtering at the call site, mirroring + # ``llm.response.RuntimeConfig.from_partial``. + @classmethod + def from_partial(cls, **kwargs: Any) -> EmbeddingRuntimeConfig: + """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", +] diff --git a/tests/conformance/test_retrieval_provider.py b/tests/conformance/test_retrieval_provider.py new file mode 100644 index 0000000..c780b72 --- /dev/null +++ b/tests/conformance/test_retrieval_provider.py @@ -0,0 +1,205 @@ +"""Run every spec retrieval-provider conformance fixture against OpenAIEmbeddingProvider. + +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. +""" + +from __future__ import annotations + +import json +from collections.abc import Mapping +from pathlib import Path +from typing import Any, cast + +import httpx +import pytest +import yaml + +from openarmature.llm import LlmProviderError +from openarmature.retrieval import EmbeddingRuntimeConfig, OpenAIEmbeddingProvider + +from ._deferral import skip_if_deferred + +CONFORMANCE_DIR = ( + Path(__file__).resolve().parents[2] / "openarmature-spec" / "spec" / "retrieval-provider" / "conformance" +) + +# Default bound model when a fixture's calls_embed omits one; matches the +# 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." +_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 +} + + +def _fixture_paths() -> list[Path]: + return sorted(CONFORMANCE_DIR.glob("[0-9][0-9][0-9]-*.yaml")) + + +def _fixture_id(path: Path) -> str: + return path.stem + + +def _load(path: Path) -> dict[str, Any]: + with path.open() as f: + return cast("dict[str, Any]", yaml.safe_load(f)) + + +def _build_handler( + responses: list[Mapping[str, Any]], +) -> tuple[httpx.MockTransport, list[httpx.Request]]: + """MockTransport handing back the configured responses in arrival order.""" + captured: list[httpx.Request] = [] + iterator = iter(responses) + + def handler(request: httpx.Request) -> httpx.Response: + captured.append(request) + try: + spec = next(iterator) + except StopIteration as exc: + raise AssertionError(f"mock embedding exhausted at request {len(captured)}") from exc + status = int(spec.get("status", 200)) + body = spec.get("body") + if body is None: + return httpx.Response(status) + return httpx.Response(status, content=json.dumps(body).encode("utf-8")) + + return httpx.MockTransport(handler), captured + + +def _build_provider( + responses: list[Mapping[str, Any]], + *, + model: str, +) -> tuple[OpenAIEmbeddingProvider, list[httpx.Request]]: + transport, captured = _build_handler(responses) + provider = OpenAIEmbeddingProvider( + base_url="http://mock-embed.test", + model=model, + api_key="test-key", + transport=transport, + ) + return provider, captured + + +def _build_config(config_block: Mapping[str, Any] | None) -> EmbeddingRuntimeConfig | 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 EmbeddingRuntimeConfig(**block) + + +def _assert_embedding_response( + response: Any, + expected: Mapping[str, Any], +) -> None: + """Assert the response against the expected.final_state.embedding_response + block. The block uses assertion keys (vectors_length / dimensions / ...) + and/or a literal ``vectors`` list.""" + for key, val in expected.items(): + if key == "vectors": + assert response.vectors == [[float(x) for x in v] for v in val], ( + f"vectors mismatch: {response.vectors} != {val}" + ) + elif key == "vectors_length": + assert len(response.vectors) == val + elif key == "dimensions": + assert response.dimensions == val + elif key == "inner_vector_lengths_all_equal": + assert all(len(v) == val for v in response.vectors) + elif key == "model": + assert response.model == val + elif key == "response_id": + assert response.response_id == val + elif key == "usage": + assert response.usage.input_tokens == cast("Mapping[str, Any]", val)["input_tokens"] + # Unknown keys are skipped (tolerant, matching the llm-provider + # harness): a new spec assertion key is a no-op here until wired, + # rather than breaking the whole parametrized suite. + + +def _check_success_invariants( + response: Any, + input_strings: list[str], + invariants: Mapping[str, Any], +) -> None: + for key, val in invariants.items(): + if not val: + continue + if key == "vectors_length_matches_input_length": + assert len(response.vectors) == len(input_strings) + elif key == "all_vectors_same_dimensionality": + assert len({len(v) for v in response.vectors}) <= 1 + elif key == "dimensions_field_matches_inner_vector_length": + assert not response.vectors or response.dimensions == len(response.vectors[0]) + elif key.startswith("vector_at_index_"): + # Order assertions; covered by the literal final_state.vectors check. + continue + # Unknown invariants are skipped (tolerant); the structural invariants + # above plus the baseline count assertion cover the core 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 + skip_if_deferred(fixture_id, _DEFERRED_FIXTURES) + spec = _load(fixture_path) + for case in cast("list[dict[str, Any]]", spec.get("cases", [])): + case_name = case.get("name", "") + try: + await _run_one_case(case) + except AssertionError as e: + raise AssertionError(f"case {case_name!r}: {e}") from e + + +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]) + 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)) + config = _build_config(cast("Mapping[str, Any] | None", calls_embed.get("config"))) + responses = cast("list[Mapping[str, Any]]", case.get("mock_embedding") or []) + provider, _captured = _build_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.embed(input_strings, config=config) + assert excinfo.value.category == expected_error["category"], ( + f"expected {expected_error['category']}, got {excinfo.value.category}" + ) + else: + response = await provider.embed(input_strings, config=config) + # Baseline assertion: every success case checks the core + # one-vector-per-input invariant, so a fixture without a + # final_state / invariants block still asserts something real. + assert len(response.vectors) == len(input_strings), ( + f"expected {len(input_strings)} vectors, got {len(response.vectors)}" + ) + final_state = cast("Mapping[str, Any]", expected.get("final_state") or {}) + stored = cast("str", calls_embed["stores_response_in"]) + if stored in final_state: + _assert_embedding_response(response, cast("Mapping[str, Any]", final_state[stored])) + _check_success_invariants(response, input_strings, invariants) + finally: + await provider.aclose() diff --git a/tests/unit/test_retrieval_provider.py b/tests/unit/test_retrieval_provider.py new file mode 100644 index 0000000..1753203 --- /dev/null +++ b/tests/unit/test_retrieval_provider.py @@ -0,0 +1,162 @@ +"""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). +""" + +from __future__ import annotations + +from typing import Literal + +import httpx +import pytest + +from openarmature.graph.events import EmbeddingEvent, EmbeddingFailedEvent +from openarmature.llm.errors import ProviderInvalidModel +from openarmature.retrieval import OpenAIEmbeddingProvider + + +def _ok_handler(request: httpx.Request) -> httpx.Response: + if request.url.path.endswith("/v1/embeddings"): + return httpx.Response( + 200, + json={ + "object": "list", + "model": "m", + "data": [{"object": "embedding", "index": 0, "embedding": [0.1, 0.2]}], + "usage": {"prompt_tokens": 1}, + }, + ) + return httpx.Response(200, json={"object": "list", "data": [{"id": "m"}]}) + + +def test_base_url_rejects_v1_suffix() -> None: + with pytest.raises(ValueError, match="host root"): + OpenAIEmbeddingProvider(base_url="https://api.openai.com/v1", model="m") + # The host root is accepted (no doubled /v1/v1). + OpenAIEmbeddingProvider(base_url="https://api.openai.com", model="m") + + +def test_readiness_probe_rejects_unknown_mode() -> None: + with pytest.raises(ValueError, match="readiness_probe must be one of"): + OpenAIEmbeddingProvider(base_url="http://x", model="m", readiness_probe="bogus") # type: ignore[arg-type] + + +@pytest.mark.parametrize( + ("mode", "expected_paths"), + [ + ("embed", ["/v1/embeddings"]), + ("models", ["/v1/models"]), + ("both", ["/v1/models", "/v1/embeddings"]), + ], +) +async def test_readiness_probe_modes_hit_expected_paths( + mode: Literal["embed", "models", "both"], expected_paths: list[str] +) -> None: + paths: list[str] = [] + + def handler(request: httpx.Request) -> httpx.Response: + paths.append(request.url.path) + return _ok_handler(request) + + provider = OpenAIEmbeddingProvider( + base_url="http://x", + model="m", + readiness_probe=mode, + transport=httpx.MockTransport(handler), + ) + await provider.ready() + assert paths == expected_paths + await provider.aclose() + + +async def test_readiness_models_probe_raises_when_model_absent() -> None: + def handler(request: httpx.Request) -> httpx.Response: + return httpx.Response(200, json={"object": "list", "data": [{"id": "other-model"}]}) + + provider = OpenAIEmbeddingProvider( + base_url="http://x", + model="m", + readiness_probe="models", + transport=httpx.MockTransport(handler), + ) + with pytest.raises(ProviderInvalidModel): + await provider.ready() + await provider.aclose() + + +def _embedding_event() -> EmbeddingEvent: + return EmbeddingEvent( + invocation_id="inv", + correlation_id=None, + node_name="embed_node", + namespace=("embed_node",), + attempt_index=0, + fan_out_index=None, + branch_name=None, + provider="openai", + model="m", + response_id=None, + response_model=None, + usage=None, + latency_ms=1.0, + input_strings=["x"], + input_count=1, + dimensions=2, + request_params={}, + request_extras={}, + active_prompt=None, + active_prompt_group=None, + call_id="c", + ) + + +def _embedding_failed_event() -> EmbeddingFailedEvent: + return EmbeddingFailedEvent( + invocation_id="inv", + correlation_id=None, + node_name="embed_node", + namespace=("embed_node",), + attempt_index=0, + fan_out_index=None, + branch_name=None, + provider="openai", + model="m", + latency_ms=1.0, + input_strings=["x"], + 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_safely_ignores_embedding_events() -> 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)) + # Without the skip branch these would fall through to ``event.phase`` + # and AttributeError (the events have no ``phase`` field). The bundled + # embedding span is a follow-up; here they must no-op. + await observer(_embedding_event()) + await observer(_embedding_failed_event()) + assert len(exporter.get_finished_spans()) == 0 + + +async def test_langfuse_observer_safely_ignores_embedding_events() -> None: + from openarmature.observability.langfuse import InMemoryLangfuseClient, LangfuseObserver + + observer = LangfuseObserver(client=InMemoryLangfuseClient()) + # Must not raise (same skip-branch contract as the OTel observer). + await observer(_embedding_event()) + await observer(_embedding_failed_event()) From 47ba9086f9ced3900767b4184c7940cac4f96104 Mon Sep 17 00:00:00 2001 From: chris-colinsky Date: Sun, 28 Jun 2026 08:57:44 -0700 Subject: [PATCH 2/4] Address embedding-capability review findings Fix the issues from the PR #196 re-review: - Build the embeddings request body extras-first so a caller's undeclared runtime-config extra cannot clobber the bound model or the input list. - Reject non-numeric vector values (JSON strings, bools) and bool usage counts strictly rather than coercing them, so a non-numeric vector is treated as malformed. - Dedup the observer event union: correlation and both observers now reference the single ObserverEvent alias instead of repeating the member list, so a new event widens the contract in one place. - Move a stray spec section ref out of a docstring into a comment and reword em-dash substitutes in docstrings. - Strengthen the tests: the Langfuse safe-handling test asserts no trace is created, and the conformance runner tolerates a fixture without a stores_response_in key instead of raising KeyError. - Refresh the observer dispatch docstrings to point at ObserverEvent rather than a hardcoded variant count. --- src/openarmature/graph/observer.py | 8 ++-- src/openarmature/observability/correlation.py | 41 ++++--------------- .../observability/langfuse/observer.py | 16 +------- .../observability/otel/observer.py | 16 +------- src/openarmature/retrieval/provider.py | 2 +- .../retrieval/providers/openai.py | 37 +++++++++++------ src/openarmature/retrieval/response.py | 2 +- tests/conformance/test_retrieval_provider.py | 6 +-- tests/unit/test_retrieval_provider.py | 7 +++- 9 files changed, 49 insertions(+), 86 deletions(-) diff --git a/src/openarmature/graph/observer.py b/src/openarmature/graph/observer.py index c4fea6a..fd26800 100644 --- a/src/openarmature/graph/observer.py +++ b/src/openarmature/graph/observer.py @@ -110,9 +110,9 @@ async def log_observer(event: NodeEvent | MetadataAugmentationEvent) -> None: conformance doesn't pin you to that name; any of `event`, `_event`, `e`, etc. matches. - Seven event variants reach observers. The signature is the union; - observers ``isinstance``-narrow on the first line and choose which - variants they handle. + The variants reaching observers are the :data:`ObserverEvent` members. + The signature is that union; observers ``isinstance``-narrow on the + first line and choose which variants they handle. - :class:`NodeEvent` — the started/completed/checkpoint phase events. Subject to the ``phases`` filter on @@ -784,7 +784,7 @@ def _dispatch( ) -> None: """Enqueue an event for the delivery worker. - Handles four event variants: + Handles the :data:`ObserverEvent` variants. The principal ones: - :class:`NodeEvent`: the started/completed/checkpoint pair model. For ``"started"``-phase events, also calls any subscribed diff --git a/src/openarmature/observability/correlation.py b/src/openarmature/observability/correlation.py index bc01cd0..4a336fe 100644 --- a/src/openarmature/observability/correlation.py +++ b/src/openarmature/observability/correlation.py @@ -36,40 +36,13 @@ from typing import TYPE_CHECKING if TYPE_CHECKING: - from openarmature.graph.events import ( - EmbeddingEvent, - EmbeddingFailedEvent, - FailureIsolatedEvent, - InvocationCompletedEvent, - InvocationStartedEvent, - LlmCompletionEvent, - LlmFailedEvent, - LlmRetryAttemptEvent, - MetadataAugmentationEvent, - NodeEvent, - ToolCallEvent, - ToolCallFailedEvent, - ) - from openarmature.graph.observer import SubscribedObserver - - # The event-record union the engine's serial dispatch worker accepts. - # Defined once and reused across the ContextVar + the get/set/reset - # accessors so a new capability event widens the contract in one place. - _DispatchEvent = ( - NodeEvent - | MetadataAugmentationEvent - | InvocationStartedEvent - | InvocationCompletedEvent - | LlmCompletionEvent - | LlmFailedEvent - | LlmRetryAttemptEvent - | FailureIsolatedEvent - | ToolCallEvent - | ToolCallFailedEvent - | EmbeddingEvent - | EmbeddingFailedEvent - ) - _DispatchFn = Callable[[_DispatchEvent], None] + from openarmature.graph.observer import ObserverEvent, SubscribedObserver + + # The dispatch callable accepts the full ObserverEvent union (the single + # source of truth, defined in graph.observer): a new capability event + # widens the contract in one place rather than across every dispatch + # accessor here. + _DispatchFn = Callable[[ObserverEvent], None] # --------------------------------------------------------------------------- diff --git a/src/openarmature/observability/langfuse/observer.py b/src/openarmature/observability/langfuse/observer.py index 3d15a2c..716b106 100644 --- a/src/openarmature/observability/langfuse/observer.py +++ b/src/openarmature/observability/langfuse/observer.py @@ -43,6 +43,7 @@ ToolCallEvent, ToolCallFailedEvent, ) +from openarmature.graph.observer import ObserverEvent from openarmature.observability.lineage import is_strict_prefix from .client import ( @@ -434,20 +435,7 @@ def __post_init__(self) -> None: async def __call__( self, - event: ( - NodeEvent - | MetadataAugmentationEvent - | InvocationStartedEvent - | InvocationCompletedEvent - | LlmCompletionEvent - | LlmFailedEvent - | LlmRetryAttemptEvent - | FailureIsolatedEvent - | ToolCallEvent - | ToolCallFailedEvent - | EmbeddingEvent - | EmbeddingFailedEvent - ), + event: ObserverEvent, ) -> None: if isinstance(event, InvocationStartedEvent): self._handle_invocation_started(event) diff --git a/src/openarmature/observability/otel/observer.py b/src/openarmature/observability/otel/observer.py index 825488e..a5c04df 100644 --- a/src/openarmature/observability/otel/observer.py +++ b/src/openarmature/observability/otel/observer.py @@ -113,6 +113,7 @@ ToolCallEvent, ToolCallFailedEvent, ) +from openarmature.graph.observer import ObserverEvent from openarmature.observability.lineage import is_strict_prefix from openarmature.observability.llm_event import serialize_tool_calls @@ -746,20 +747,7 @@ def _inv_state_for(self, invocation_id: str) -> _InvState: async def __call__( self, - event: ( - NodeEvent - | MetadataAugmentationEvent - | InvocationStartedEvent - | InvocationCompletedEvent - | LlmCompletionEvent - | LlmFailedEvent - | LlmRetryAttemptEvent - | FailureIsolatedEvent - | ToolCallEvent - | ToolCallFailedEvent - | EmbeddingEvent - | EmbeddingFailedEvent - ), + event: ObserverEvent, ) -> None: # Proposal 0043 invocation-boundary events: OTel has no # Trace-level input/output payload concept (a trace is a diff --git a/src/openarmature/retrieval/provider.py b/src/openarmature/retrieval/provider.py index 12c5bcb..ac7e896 100644 --- a/src/openarmature/retrieval/provider.py +++ b/src/openarmature/retrieval/provider.py @@ -58,7 +58,7 @@ async def embed( config: Optional per-call request parameters. Returns an :class:`EmbeddingResponse` whose ``vectors[i]`` is the - embedding of ``input[i]`` -- order is preserved, never permuted. + embedding of ``input[i]``; order is preserved, never permuted. """ ... diff --git a/src/openarmature/retrieval/providers/openai.py b/src/openarmature/retrieval/providers/openai.py index 3459181..c2b6326 100644 --- a/src/openarmature/retrieval/providers/openai.py +++ b/src/openarmature/retrieval/providers/openai.py @@ -56,10 +56,10 @@ def _classify_embedding_http_error(resp: httpx.Response) -> LlmProviderError: - """Map a non-200 OpenAI-shape embeddings response to a §7 category. + """Map a non-200 OpenAI-shape embeddings response to an error category. - The embedding-applicable subset: 401/403 -> auth, 429 -> rate_limit, - 404 -> invalid_model, 400/422 -> invalid_request, other 5xx -> + The applicable subset: 401/403 to auth, 429 to rate_limit, 404 to + invalid_model, 400/422 to invalid_request, and other 5xx to unavailable. Returns the exception (does not raise) so the caller raises with consistent traceback context. """ @@ -290,13 +290,14 @@ def _build_request_body( request_extras: dict[str, Any], ) -> dict[str, Any]: """Build the /v1/embeddings request body.""" - # model + input always; the supplied request params (the same dict - # the event carries -- dimensions for embedding) and the - # runtime-config extras pass-through merge through. Sharing the one - # params dict keeps the wire body and the observed request_params + # Extras are merged FIRST so the bound model, the input list, and the + # declared request params always win: a caller's undeclared extra + # named "model" or "input" must not clobber the wire identity. The + # request params are the same dict the event carries (dimensions for + # embedding), keeping the wire body and the observed request_params # provably identical. Proposal 0077's input_type maps to a per-vendor # wire key at the wire layer. - return {"model": self.model, "input": input_strings, **request_params, **request_extras} + return {**request_extras, "model": self.model, "input": input_strings, **request_params} def _parse_response( self, @@ -341,10 +342,14 @@ def _parse_response( embedding = entry.get("embedding") if not isinstance(embedding, list) or not embedding: raise ProviderInvalidResponse("embedding response entry has a missing or empty 'embedding'") - try: - vectors.append([float(x) for x in cast("list[Any]", embedding)]) - except (TypeError, ValueError) as exc: - raise ProviderInvalidResponse("embedding response has a non-numeric vector value") from exc + values = cast("list[Any]", embedding) + # Reject non-numeric vector values (JSON strings, bools) rather + # than coercing them: bool is an int subclass, and float("0.1") + # would silently accept a string, so the strict isinstance check + # is what makes "non-numeric is malformed" actually hold. + if not all(isinstance(x, (int, float)) and not isinstance(x, bool) for x in values): + raise ProviderInvalidResponse("embedding response has a non-numeric vector value") + vectors.append([float(x) for x in values]) dimensions = validate_embedding_response(vectors, len(input_strings)) usage_block = body.get("usage") prompt_tokens = ( @@ -352,7 +357,13 @@ def _parse_response( if isinstance(usage_block, dict) else None ) - input_tokens = prompt_tokens if isinstance(prompt_tokens, int) and prompt_tokens >= 0 else 0 + # bool is an int subclass, so exclude it explicitly; a malformed or + # absent usage falls back to 0 (the embed succeeded; usage is secondary). + input_tokens = ( + prompt_tokens + if isinstance(prompt_tokens, int) and not isinstance(prompt_tokens, bool) and prompt_tokens >= 0 + else 0 + ) response_id = body.get("id") model = body.get("model") return EmbeddingResponse( diff --git a/src/openarmature/retrieval/response.py b/src/openarmature/retrieval/response.py index 3dfb0ae..0ad13f1 100644 --- a/src/openarmature/retrieval/response.py +++ b/src/openarmature/retrieval/response.py @@ -30,7 +30,7 @@ class EmbeddingUsage(BaseModel): """Token-accounting record for an embedding call. - Carries ``input_tokens`` only -- an embedding call has no output + Carries ``input_tokens`` only; an embedding call has no output tokens (vectors are not tokens). """ diff --git a/tests/conformance/test_retrieval_provider.py b/tests/conformance/test_retrieval_provider.py index c780b72..d236955 100644 --- a/tests/conformance/test_retrieval_provider.py +++ b/tests/conformance/test_retrieval_provider.py @@ -4,7 +4,7 @@ ``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 +``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). @@ -197,8 +197,8 @@ async def _run_one_case(case: Mapping[str, Any]) -> None: f"expected {len(input_strings)} vectors, got {len(response.vectors)}" ) final_state = cast("Mapping[str, Any]", expected.get("final_state") or {}) - stored = cast("str", calls_embed["stores_response_in"]) - if stored in final_state: + stored = cast("str | None", calls_embed.get("stores_response_in")) + if stored is not None and stored in final_state: _assert_embedding_response(response, cast("Mapping[str, Any]", final_state[stored])) _check_success_invariants(response, input_strings, invariants) finally: diff --git a/tests/unit/test_retrieval_provider.py b/tests/unit/test_retrieval_provider.py index 1753203..fefe134 100644 --- a/tests/unit/test_retrieval_provider.py +++ b/tests/unit/test_retrieval_provider.py @@ -156,7 +156,10 @@ async def test_otel_observer_safely_ignores_embedding_events() -> None: async def test_langfuse_observer_safely_ignores_embedding_events() -> None: from openarmature.observability.langfuse import InMemoryLangfuseClient, LangfuseObserver - observer = LangfuseObserver(client=InMemoryLangfuseClient()) - # Must not raise (same skip-branch contract as the OTel observer). + client = InMemoryLangfuseClient() + observer = LangfuseObserver(client=client) + # Must not raise (same skip-branch contract as the OTel observer), and no + # trace / observation is created for the skipped events. await observer(_embedding_event()) await observer(_embedding_failed_event()) + assert client.traces == {} From a7dcfe69211bcde4bdcf93e157668bb8c2d2267f Mon Sep 17 00:00:00 2001 From: chris-colinsky Date: Sun, 28 Jun 2026 09:20:41 -0700 Subject: [PATCH 3/4] Raise invalid-response on malformed model catalog From the PR #196 review: - _probe_models now raises ProviderInvalidResponse when the /v1/models body is not a JSON object or its data field is missing or not a list, mirroring _parse_response. A malformed catalog was being degraded to an empty list and misreported as a missing model (ProviderInvalidModel); that category is now reserved for a well-formed catalog that genuinely lacks the bound model. Covered by a new unit test. - Align the _classify_embedding_http_error docstring with the code: the catch-all maps every other status (not just 5xx) to unavailable, the same as the sibling classify_http_error. --- .../retrieval/providers/openai.py | 10 +++++++--- tests/unit/test_retrieval_provider.py | 20 ++++++++++++++++++- 2 files changed, 26 insertions(+), 4 deletions(-) diff --git a/src/openarmature/retrieval/providers/openai.py b/src/openarmature/retrieval/providers/openai.py index c2b6326..55fb276 100644 --- a/src/openarmature/retrieval/providers/openai.py +++ b/src/openarmature/retrieval/providers/openai.py @@ -59,7 +59,7 @@ def _classify_embedding_http_error(resp: httpx.Response) -> LlmProviderError: """Map a non-200 OpenAI-shape embeddings response to an error category. The applicable subset: 401/403 to auth, 429 to rate_limit, 404 to - invalid_model, 400/422 to invalid_request, and other 5xx 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. """ @@ -213,9 +213,13 @@ async def _probe_models(self) -> None: body_raw = resp.json() except ValueError as exc: raise ProviderInvalidResponse("models response is not valid JSON") from exc - body = cast("dict[str, Any]", body_raw) if isinstance(body_raw, dict) else {} + if not isinstance(body_raw, dict): + raise ProviderInvalidResponse("models response is not a JSON object") + body = cast("dict[str, Any]", body_raw) data_raw = body.get("data") - models = cast("list[Any]", data_raw) if isinstance(data_raw, list) else [] + if not isinstance(data_raw, list): + raise ProviderInvalidResponse("models response missing 'data' array") + models = cast("list[Any]", data_raw) ids = [cast("dict[str, Any]", m).get("id") for m in models if isinstance(m, dict)] if self.model not in ids: seen = sorted(i for i in ids if isinstance(i, str)) diff --git a/tests/unit/test_retrieval_provider.py b/tests/unit/test_retrieval_provider.py index fefe134..3428899 100644 --- a/tests/unit/test_retrieval_provider.py +++ b/tests/unit/test_retrieval_provider.py @@ -15,7 +15,7 @@ import pytest from openarmature.graph.events import EmbeddingEvent, EmbeddingFailedEvent -from openarmature.llm.errors import ProviderInvalidModel +from openarmature.llm.errors import ProviderInvalidModel, ProviderInvalidResponse from openarmature.retrieval import OpenAIEmbeddingProvider @@ -88,6 +88,24 @@ def handler(request: httpx.Request) -> httpx.Response: await provider.aclose() +async def test_readiness_models_probe_raises_on_malformed_catalog() -> None: + # A 200 OK whose body is not a JSON object, or whose 'data' is not a + # list, is a wire-format problem (proxy/backend), not a missing model: + # it surfaces as ProviderInvalidResponse, not ProviderInvalidModel. + def handler(request: httpx.Request) -> httpx.Response: + return httpx.Response(200, json={"object": "list", "data": "not-a-list"}) + + provider = OpenAIEmbeddingProvider( + base_url="http://x", + model="m", + readiness_probe="models", + transport=httpx.MockTransport(handler), + ) + with pytest.raises(ProviderInvalidResponse): + await provider.ready() + await provider.aclose() + + def _embedding_event() -> EmbeddingEvent: return EmbeddingEvent( invocation_id="inv", From f04fcc94bcd112b64eb56627682489a49b187831 Mon Sep 17 00:00:00 2001 From: chris-colinsky Date: Sun, 28 Jun 2026 09:38:45 -0700 Subject: [PATCH 4/4] Fix Observer docstring example annotation The Observer protocol docstring example annotated its log_observer with NodeEvent | MetadataAugmentationEvent, which would be a structural conformance error now that the protocol delivers the full ObserverEvent union (embedding, tool, and provider events included). Annotate the example with ObserverEvent so copied observers type-check. --- src/openarmature/graph/observer.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/openarmature/graph/observer.py b/src/openarmature/graph/observer.py index fd26800..c35c4bb 100644 --- a/src/openarmature/graph/observer.py +++ b/src/openarmature/graph/observer.py @@ -91,7 +91,7 @@ class Observer(Protocol): signature qualifies, no subclass required. Plain functions, bound methods, and class instances with `__call__` all work:: - async def log_observer(event: NodeEvent | MetadataAugmentationEvent) -> None: + async def log_observer(event: ObserverEvent) -> None: if isinstance(event, NodeEvent): print(event.node_name, event.phase)