From 43a901a751bbe182775545132303e135ed59f482 Mon Sep 17 00:00:00 2001 From: chris-colinsky Date: Sat, 18 Jul 2026 10:27:28 -0700 Subject: [PATCH 1/5] Add retrieval docs, providers section, and example Document the retrieval-provider capability for v0.16.0, mirroring the LLM documentation layout. Concept page (concepts/retrieval.md): embedding and reranking, the input_type query/document knob, batch chunking, and the nullable-usage contract. Stays concept-focused and points to the providers section. Retrieval Providers section (retrieval-providers/), parallel to model-providers/: index.md (the bundled-providers catalog, the EmbeddingProvider / RerankProvider contract, error categories, a minimal example), tei.md (self-hosting Text Embeddings Inference for embeddings and reranking, from a first docker run to a systemd production deployment, with the providers wired to it), and authoring.md (a provider skeleton, contract checklist, and the chunking / input_type / observability concerns). API reference (reference/retrieval.md) and nav entries for both new sections; a retrieval entry in the README capabilities list. Example (examples/retrieval-rag): a lunar-corpus retrieval-augmented answering pipeline that batch-embeds a corpus, embeds the query, retrieves by cosine similarity, reranks the shortlist, and grounds an answer in the reranked passages, with an observer printing the embedding and rerank events. Exports the four retrieval event types (EmbeddingEvent, EmbeddingFailedEvent, RerankEvent, RerankFailedEvent) top-level from openarmature.graph, matching the LLM events, so observers over retrieval calls import them the same way. Regenerates the bundled AGENTS.md example index. --- README.md | 2 +- docs/concepts/retrieval.md | 176 ++++++++++++++ docs/reference/retrieval.md | 7 + docs/retrieval-providers/authoring.md | 234 +++++++++++++++++++ docs/retrieval-providers/index.md | 161 +++++++++++++ docs/retrieval-providers/tei.md | 238 +++++++++++++++++++ examples/README.md | 19 ++ examples/retrieval-rag/main.py | 320 ++++++++++++++++++++++++++ mkdocs.yml | 8 + src/openarmature/AGENTS.md | 1 + src/openarmature/graph/__init__.py | 8 + tests/test_examples_smoke.py | 1 + 12 files changed, 1174 insertions(+), 1 deletion(-) create mode 100644 docs/concepts/retrieval.md create mode 100644 docs/reference/retrieval.md create mode 100644 docs/retrieval-providers/authoring.md create mode 100644 docs/retrieval-providers/index.md create mode 100644 docs/retrieval-providers/tei.md create mode 100644 examples/retrieval-rag/main.py diff --git a/README.md b/README.md index 9f43e211..702379a9 100644 --- a/README.md +++ b/README.md @@ -177,7 +177,7 @@ A few things to notice: ## Next steps - **Quickstart**: build your first graph end-to-end. [openarmature.ai/getting-started](https://openarmature.ai/getting-started/) -- **Concepts**: typed state, reducers, graphs, composition, fan-out, parallel branches, LLMs, prompts, observability, checkpointing. [openarmature.ai/concepts](https://openarmature.ai/concepts/) +- **Concepts**: typed state, reducers, graphs, composition, fan-out, parallel branches, LLMs, retrieval, prompts, observability, checkpointing. [openarmature.ai/concepts](https://openarmature.ai/concepts/) - **Model Providers**: implement the Provider Protocol for a custom LLM backend. [openarmature.ai/model-providers/authoring](https://openarmature.ai/model-providers/authoring/) - **API reference**: auto-generated from docstrings. [openarmature.ai/reference](https://openarmature.ai/reference/) - **Examples**: ten runnable demos with walk-throughs. [openarmature.ai/examples](https://openarmature.ai/examples/) (source at [./examples/](./examples/)) diff --git a/docs/concepts/retrieval.md b/docs/concepts/retrieval.md new file mode 100644 index 00000000..3069c4ef --- /dev/null +++ b/docs/concepts/retrieval.md @@ -0,0 +1,176 @@ +# Retrieval + +The retrieval capability covers the two building blocks of a search or +RAG pipeline: turning text into vectors (**embedding**) and reordering a +candidate list by relevance to a query (**reranking**). Like the LLM +capability, both are async IO you call from inside a node body, behind a +small provider protocol so the graph does not care which vendor is on the +other end. + +Everything lives in `openarmature.retrieval`: + +```python +from openarmature.retrieval import ( + OpenAIEmbeddingProvider, + CohereRerankProvider, + EmbeddingRuntimeConfig, + RerankRuntimeConfig, +) +``` + +## Embedding text + +An `EmbeddingProvider` turns a list of strings into a list of vectors, +one vector per input, in input order: + +```python +provider = OpenAIEmbeddingProvider(model="text-embedding-3-small", api_key="sk-...") + +response = await provider.embed(["the lunar south pole", "the Sea of Tranquility"]) + +response.vectors # list[list[float]], one per input, same order +response.dimensions # the vector length (equal across all vectors) +response.model # the model the provider actually served +response.usage # an EmbeddingUsage record, or None (see below) +``` + +The input-order guarantee is the contract you build on: `vectors[i]` is +always the embedding of `input[i]`, no matter how the provider paginated +the request under the hood. `len(response.vectors) == len(input)` always +holds, and every vector has the same dimensionality. + +`embed` takes an arbitrary-length list. You do not have to pre-chunk to +fit a provider's per-request cap; the mapping does that for you (see +[Long input lists](#long-input-lists-are-chunked-for-you)). + +## Reranking candidates + +A `RerankProvider` scores a set of candidate documents against a query +and returns them sorted best-first. This is the precision step after a +cheap-and-broad first pass (vector similarity, keyword search, whatever): + +```python +provider = CohereRerankProvider(model="rerank-v3.5", api_key="...") + +response = await provider.rerank( + "where is there water ice on the Moon?", + ["Regolith is abrasive dust.", "Ice sits in shadowed polar craters.", "..."], + top_k=3, +) + +for result in response.results: # sorted by relevance_score, best first + result.index # position in the input `documents` list + result.relevance_score # higher is more relevant + result.document # the echoed text, or None (see below) +``` + +Results come back ranked. The mapping applies the sort even if a provider +returns them unsorted, so you can always trust `response.results[0]` to be +the best hit. + +`result.index` is the load-bearing field: it points back into the +`documents` list you passed. Not every provider echoes the document text +(Cohere, for one, returns scores only), so `result.document` may be +`None`. Map back to your own candidate list by index rather than relying +on the echo: + +```python +ranked = [candidates[result.index] for result in response.results] +``` + +`top_k` trims the returned results to the best K. Pass it when you only +want the top of the list; omit it to score every candidate. + +## Query vs document: `input_type` + +Some embedding models are **asymmetric**: they embed a search query and +the documents being searched with slightly different representations, and +mixing them up hurts recall. `EmbeddingRuntimeConfig.input_type` is the +portable knob for this: + +```python +# When embedding the corpus you are searching over: +await provider.embed(passages, config=EmbeddingRuntimeConfig(input_type="document")) + +# When embedding the user's query at search time: +await provider.embed([query], config=EmbeddingRuntimeConfig(input_type="query")) +``` + +Each provider realizes it the way its wire expects: on an asymmetric +model it selects the query or document representation; on a symmetric +model (OpenAI's) it is a no-op and the text is embedded verbatim. Setting +it costs nothing on the symmetric providers and keeps the same pipeline +correct if you later switch to an asymmetric one, so prefer setting it. + +## Long input lists are chunked for you + +Every hosted embedding API caps how many inputs one request may carry. +The mapping honors the contract that `embed` takes any-length input by +splitting a large list into consecutive slices under the cap, issuing one +request per slice, and stitching the vectors back together in input +order. This is invisible: you call `embed` with 10,000 strings and get +10,000 vectors, regardless of the provider's per-request limit. + +Usage is combined across the slices: the token total is reported only +when every slice reported one, otherwise `usage` is `None` for the whole +call (an honest "unknown" rather than a partial count). + +## Usage is a record or null + +`response.usage` is an `EmbeddingUsage` (or `RerankUsage`) record when the +provider reported token accounting, and `None` when it did not. The +mapping never fabricates a usage record, a zero, or a client-side +estimate: a `None` means "the provider told us nothing," which is +different from a real, reported zero. Guard before reading it: + +```python +if response.usage is not None: + print(response.usage.input_tokens) +``` + +This matters because not every provider bills the same way. A local +Text Embeddings Inference server reports no usage at all, so `usage` is +`None`. Cohere's reranker reports `search_units` but no token count, so +its `RerankUsage` carries `search_units` with `input_tokens=None`. + +## The bundled providers + +Four vendors ship as reference providers: OpenAI-compatible (embedding +only), Cohere, Jina, and TEI. All four embed; Cohere, Jina, and TEI also +rerank. The [Retrieval Providers](../retrieval-providers/index.md) section +has the full table, the protocol contract, per-provider notes, and a guide +to [self-hosting TEI](../retrieval-providers/tei.md) for embeddings and +reranking on your own hardware. Writing your own is the same exercise as a +custom LLM provider: implement the `EmbeddingProvider` or `RerankProvider` +protocol and the graph treats it like any other. + +## Observability + +When you call `embed` or `rerank` from inside a node body, the provider +dispatches a typed `EmbeddingEvent` or `RerankEvent` (and a failed +variant on error) to any attached observer, exactly like LLM completions. +The bundled OTel observer renders an `openarmature.embedding.complete` or +`openarmature.rerank.complete` span; the Langfuse observer renders a +dedicated Embedding or Retriever observation. Token usage lands on the +span or observation only when the provider reported it, matching the +nullable-usage contract above. + +Provider calls made outside a graph (for example an offline index build) +run fine but dispatch no events, since there is no observer context to +receive them. + +## Putting it together + +The [`retrieval-rag`](https://github.com/LunarCommand/openarmature-python/tree/main/examples/retrieval-rag) +example wires the whole pattern into a graph: it batch-embeds a corpus +into an index, then per query embeds the question, retrieves by cosine +similarity, reranks the shortlist, and grounds an LLM answer in the +reranked passages. + +## Where to next + +- [LLMs](llms.md) for the completion side that consumes retrieved context. +- [Observability](observability.md) for what the embedding and rerank + spans carry. +- The [`openarmature.retrieval` reference](../reference/retrieval.md) for + the full type surface. diff --git a/docs/reference/retrieval.md b/docs/reference/retrieval.md new file mode 100644 index 00000000..ef6f0be6 --- /dev/null +++ b/docs/reference/retrieval.md @@ -0,0 +1,7 @@ +# openarmature.retrieval + +::: openarmature.retrieval + options: + show_root_heading: false + show_source: false + heading_level: 2 diff --git a/docs/retrieval-providers/authoring.md b/docs/retrieval-providers/authoring.md new file mode 100644 index 00000000..a60f0ff9 --- /dev/null +++ b/docs/retrieval-providers/authoring.md @@ -0,0 +1,234 @@ +# Authoring a Provider + +When you target a backend none of the bundled providers cover (a +different vendor API, an internal embedding gateway, a hand-rolled +inference service), implement the `EmbeddingProvider` or `RerankProvider` +protocol yourself. Each is one async method plus a readiness check, so a +minimum-viable provider is short. + +If you are new to retrieval providers, read +[Retrieval Providers](index.md) for the contract guarantees first. + +## Skeleton + +A minimal embedding provider against an OpenAI-shaped `/v1/embeddings` +endpoint. Compare with `openarmature.retrieval.OpenAIEmbeddingProvider` +to see what a full implementation adds (input-list chunking, the +client-side `input_type` prefix, observability events, a readiness probe). + +```python +from collections.abc import Sequence + +import httpx + +from openarmature.llm import ( + ProviderAuthentication, + ProviderInvalidRequest, + ProviderInvalidResponse, + ProviderRateLimit, + ProviderUnavailable, +) +from openarmature.retrieval import ( + EmbeddingResponse, + EmbeddingRuntimeConfig, + EmbeddingUsage, + validate_embedding_input, + validate_embedding_response, +) + + +class MyEmbeddingProvider: + def __init__(self, *, base_url: str, model: str, api_key: str | None = None) -> None: + headers = {"Content-Type": "application/json"} + if api_key: + headers["Authorization"] = f"Bearer {api_key}" + self._client = httpx.AsyncClient(base_url=base_url, headers=headers, timeout=60.0) + self.model = model + + async def ready(self) -> None: + resp = await self._client.post( + "/v1/embeddings", json={"model": self.model, "input": ["ready"]} + ) + if resp.status_code != 200: + raise _classify(resp) + + async def embed( + self, + input: Sequence[str], + *, + config: EmbeddingRuntimeConfig | None = None, + ) -> EmbeddingResponse: + input_strings = list(input) + validate_embedding_input(input_strings) # rejects an empty list + + body = {"model": self.model, "input": input_strings} + try: + resp = await self._client.post("/v1/embeddings", json=body) + except httpx.HTTPError as exc: + raise ProviderUnavailable(str(exc)) from exc + if resp.status_code != 200: + raise _classify(resp) + + payload = resp.json() + # Order the returned data by index so vectors line up with input. + entries = sorted(payload["data"], key=lambda e: e["index"]) + vectors = [[float(x) for x in e["embedding"]] for e in entries] + # Validates the count and uniform dimensionality, returns the dim. + dimensions = validate_embedding_response(vectors, len(input_strings)) + + prompt_tokens = payload.get("usage", {}).get("prompt_tokens") + usage = ( + EmbeddingUsage(input_tokens=prompt_tokens) + if isinstance(prompt_tokens, int) and prompt_tokens >= 0 + else None # never fabricate a record the provider did not report + ) + return EmbeddingResponse( + vectors=vectors, + model=payload.get("model", self.model), + usage=usage, + response_id=payload.get("id"), + dimensions=dimensions, + raw=payload, + ) + + async def aclose(self) -> None: + await self._client.aclose() + + +def _classify(resp: httpx.Response) -> Exception: + status = resp.status_code + if status in (401, 403): + return ProviderAuthentication(f"HTTP {status}") + if status == 429: + return ProviderRateLimit("HTTP 429") + if status in (400, 413, 422): + return ProviderInvalidRequest(f"HTTP {status}") + if status >= 500: + return ProviderUnavailable(f"HTTP {status}") + return ProviderInvalidResponse(f"HTTP {status}") +``` + +A `RerankProvider` is the same shape: a `rerank(query, documents, *, +top_k=None, config=None)` method that posts the query and documents, +validates the inputs with `validate_rerank_input`, builds `ScoredDocument` +results, and validates them with `validate_rerank_response` before +returning a `RerankResponse`. + +## Contract checklist + +When you ship a provider, the following must hold: + +**Input order and shape.** + +- [ ] `embed()` MUST return one vector per input, in input order, so + `vectors[i]` is the embedding of `input[i]`. `validate_embedding_response` + enforces the count and uniform dimensionality; ordering is yours to get + right (sort by the wire's index, or trust a positional wire). +- [ ] `rerank()` results MUST be sorted by relevance, best first, and each + `ScoredDocument.index` MUST point into the input `documents` list. + `validate_rerank_response` enforces the index range and `top_k`. + +**Non-mutation and reentrancy.** + +- [ ] Inputs MUST NOT be mutated; build wire bodies from copies. +- [ ] Concurrent calls on one instance MUST be safe (httpx.AsyncClient is). + +**Boundary validation.** + +- [ ] Call `validate_embedding_input(input)` (rejects an empty list) or + `validate_rerank_input(query, documents, top_k)` (rejects an empty + query, empty documents, or non-positive `top_k`) before sending. + +**Usage is a record or null.** + +- [ ] Populate `response.usage` only when the provider reports token + accounting. MUST NOT fabricate a record, a zero, or a client-side + estimate; a provider that reports nothing yields `usage = None`. A + malformed usage figure reads as not-reported, not as an error. + +**Error mapping.** + +- [ ] Network failures and timeouts to `ProviderUnavailable`; 401/403 to + `ProviderAuthentication`; 404 (model not found) to + `ProviderInvalidModel`; 429 to `ProviderRateLimit`; 400/413/422 to + `ProviderInvalidRequest`; other 5xx to `ProviderUnavailable`; a 200 + that fails to parse to `ProviderInvalidResponse`. + +**Rerank document echo.** + +- [ ] If the wire echoes document text, put it on `ScoredDocument.document` + verbatim; otherwise leave it `None`. Never reconstruct it from the + input documents. + +## Beyond the skeleton + +The skeleton omits things real providers usually need: + +- **Chunking for a capped wire.** Most embedding APIs cap the inputs per + request. When `embed()` must accept an arbitrary-length list, split it + into consecutive slices under the cap, issue one request per slice with + every non-input field identical, concatenate the vectors in input order, + and combine usage all-or-nothing (sum only when every slice reported a + figure, else `None`). The bundled `OpenAIEmbeddingProvider` and + `TeiEmbeddingProvider` do this through a shared internal helper; a custom + provider implements the same loop. A `RerankProvider` chunks a large + candidate pool the same way. + +- **The `input_type` knob.** `EmbeddingRuntimeConfig.input_type` + (`"query"` / `"document"`) selects the asymmetric-embedding + representation. Realize it however your wire expects: a native request + field, a server-side prompt name, or a client-side prefix prepended to + each input. On a symmetric model it is a no-op. + +- **Observability events.** Dispatch a typed `EmbeddingEvent` (or + `RerankEvent`) on success and the failed variant alongside any error, so + the bundled OTel and Langfuse observers render retrieval spans and + observations. The shape mirrors the LLM path (see the + [LLM provider authoring guide](../model-providers/authoring.md#beyond-the-skeleton) + for the full dispatch sketch): capture the request-side data once, then + on success dispatch an `EmbeddingEvent` built from the response, or on an + error dispatch an `EmbeddingFailedEvent` **alongside** the raised + exception (the exception still propagates; the event goes on the observer + queue). A single call emits exactly one success or one failed event, + never both. + + ```python + import uuid + + from openarmature.graph import EmbeddingEvent + from openarmature.observability.correlation import ( + current_dispatch, + current_invocation_id, + current_namespace_prefix, + ) + + + # Inside embed(), after a successful parse: + dispatch = current_dispatch() + if dispatch is not None: + namespace = current_namespace_prefix() + dispatch( + EmbeddingEvent( + invocation_id=current_invocation_id() or "", + node_name=namespace[-1] if namespace else "", + namespace=namespace, + provider="my-provider", + model=self.model, + response_model=response.model, + response_id=response.response_id, + usage=response.usage, + input_strings=input_strings, + input_count=len(input_strings), + dimensions=response.dimensions, + output_vectors=response.vectors, + call_id=str(uuid.uuid4()), + # plus the scoping fields: correlation_id, attempt_index, + # fan_out_index, branch_name, latency_ms, request_params, + # request_extras, active_prompt, active_prompt_group. + ) + ) + ``` + +The conformance fixtures under `tests/conformance/test_retrieval_provider.py` +exercise the wire mapping end to end; a custom provider that passes them +matches the contract. diff --git a/docs/retrieval-providers/index.md b/docs/retrieval-providers/index.md new file mode 100644 index 00000000..f9f1a5d4 --- /dev/null +++ b/docs/retrieval-providers/index.md @@ -0,0 +1,161 @@ +# Retrieval Providers + +A retrieval provider is the seam between OpenArmature's graph engine and +a vector or reranking backend (a hosted API like OpenAI, Cohere, or Jina, +or a self-hosted Text Embeddings Inference server). The engine does not +know about embeddings or rerankers; nodes call providers, providers do +the wire work. For what embedding and reranking are and how they fit a +pipeline, see the [Retrieval concept page](../concepts/retrieval.md); this +section is the catalog of what ships and how to run it. + +## What ships + +Four vendors ship as reference providers. Each embedding surface and each +rerank surface is a separate provider instance bound to one model: + +| Provider | Embedding | Rerank | Default `base_url` | +|---|---|---|---| +| OpenAI-compatible | `OpenAIEmbeddingProvider` | not offered | `https://api.openai.com` | +| Cohere | `CohereEmbeddingProvider` | `CohereRerankProvider` | `https://api.cohere.com` | +| Jina | `JinaEmbeddingProvider` | `JinaRerankProvider` | `https://api.jina.ai` | +| TEI (self-hosted) | `TeiEmbeddingProvider` | `TeiRerankProvider` | none (pass your instance URL) | + +A few provider-specific notes: + +- **OpenAI-compatible** is embedding-only and symmetric. `base_url` is + overridable, so the same provider drives any OpenAI-shaped embedding + endpoint (vLLM, LocalAI, TEI's OpenAI surface). `dimensions` maps to the + wire for Matryoshka models that support truncation. +- **Cohere** embedding requires an `input_type`; the mapping sends a + sensible default when you omit it. Its reranker returns scores without + echoing documents, so map results back to your candidates by + `ScoredDocument.index`. +- **Jina** realizes `input_type` as its native task field over a fixed + set of values and reports token usage on both surfaces. +- **TEI** is the self-hosted option: you pass the instance URL, it reports + no usage, and its reranker chunks a large candidate pool the same way + embedding chunks a large input list. See [Self-hosting TEI](tei.md). + +For a backend none of these cover, write your own; see +[Authoring a Provider](authoring.md). + +## The contract + +Retrieval is two protocols, one async method each: + +```python +from collections.abc import Sequence +from typing import Protocol + +from openarmature.retrieval import ( + EmbeddingResponse, + EmbeddingRuntimeConfig, + RerankResponse, + RerankRuntimeConfig, +) + + +class EmbeddingProvider(Protocol): + async def ready(self) -> None: ... + async def embed( + self, + input: Sequence[str], + *, + config: EmbeddingRuntimeConfig | None = None, + ) -> EmbeddingResponse: ... + + +class RerankProvider(Protocol): + async def ready(self) -> None: ... + async def rerank( + self, + query: str, + documents: Sequence[str], + *, + top_k: int | None = None, + config: RerankRuntimeConfig | None = None, + ) -> RerankResponse: ... +``` + +- **`ready()`** verifies the bound model is reachable. Pre-flight check, + typically called once before invoking the graph. +- **`embed()`** returns one vector per input string, in input order. +- **`rerank()`** scores the documents against the query and returns them + sorted best-first. + +### Behaviour guarantees + +- **Input order.** `embed()` returns `len(input)` vectors, `vectors[i]` + being the embedding of `input[i]`, regardless of how the provider + paginated the request. All vectors share one dimensionality. +- **Arbitrary-length input.** `embed()` accepts any-length list; the + mapping chunks under the provider's per-request cap and stitches the + result. `rerank()` chunks a large candidate pool the same way. +- **Usage is a record or null.** `response.usage` is a token-accounting + record when the provider reports one and `None` otherwise. The mapping + never fabricates a record, a zero, or an estimate. +- **Reentrant and stateless.** Safe to call concurrently from many nodes; + no implicit state carries between calls. Inputs are never mutated. +- **No retry on transient errors.** That is middleware's job; wrap the + calling node in `RetryMiddleware` or similar. + +## Errors + +Retrieval calls raise the same canonical error categories as LLM calls +(from `openarmature.llm.errors`), mapped from the provider's HTTP +response. The retrieval-applicable subset: + +| Error | Trigger | +|---|---| +| `ProviderAuthentication` | 401 / 403 (bad or missing key) | +| `ProviderRateLimit` | 429 | +| `ProviderInvalidModel` | 404 (bound model not found) | +| `ProviderInvalidRequest` | 400 / 413 / 422 (malformed or over-length request, empty input) | +| `ProviderUnavailable` | 5xx, network failure, timeout | +| `ProviderInvalidResponse` | 200 OK that fails to parse, or a malformed response body | + +`ProviderUnavailable` and `ProviderRateLimit` are in +`TRANSIENT_CATEGORIES`, the canonical "safe to retry" set the default +retry-middleware classifier uses. A malformed usage figure does not fail +the call: the vectors or scores are sound, so the count is recorded as +unknown (`usage = None`) rather than raising. + +## A minimal example + +Direct usage of the providers, without the engine in the picture: + +```python +import asyncio + +from openarmature.retrieval import CohereRerankProvider, OpenAIEmbeddingProvider + + +async def main() -> None: + embedder = OpenAIEmbeddingProvider(model="text-embedding-3-small", api_key="sk-...") + reranker = CohereRerankProvider(model="rerank-v3.5", api_key="...") + try: + vectors = (await embedder.embed(["the lunar south pole"])).vectors + ranked = (await reranker.rerank("water on the Moon", ["...", "..."])).results + print(len(vectors), [r.index for r in ranked]) + finally: + await embedder.aclose() + await reranker.aclose() + + +asyncio.run(main()) +``` + +In a real graph you construct the providers once at startup and let nodes +call them inside their bodies, where their `EmbeddingEvent` / `RerankEvent` +reach any attached observer. + +## Where to next + +- **[Self-hosting TEI](tei.md)**: run Text Embeddings Inference for + embeddings and reranking, from a first `docker run` to a production + deployment, and wire the TEI providers to it. +- **[Authoring a Provider](authoring.md)**: implement the + `EmbeddingProvider` or `RerankProvider` protocol for a backend none of + the bundled providers cover. +- **[API reference: `openarmature.retrieval`](../reference/retrieval.md)**: + the full public surface (the providers, response types, runtime configs). diff --git a/docs/retrieval-providers/tei.md b/docs/retrieval-providers/tei.md new file mode 100644 index 00000000..6179cb3f --- /dev/null +++ b/docs/retrieval-providers/tei.md @@ -0,0 +1,238 @@ +# Self-hosting TEI + +[Text Embeddings Inference](https://huggingface.co/docs/text-embeddings-inference) +(TEI) is HuggingFace's Rust serving stack for encoder models: embedding +bi-encoders and reranking cross-encoders. It is the self-hosted backend +for `TeiEmbeddingProvider` and `TeiRerankProvider`, and the right tool for +this workload where a generation server like vLLM is not. vLLM is built +for autoregressive decoding (KV cache, token streaming), none of which an +encoder uses; TEI serves `/embed` and `/rerank` directly, with Flash +Attention and batching, at lower latency because the architecture matches. + +Self-hosting keeps your corpus and queries on your own hardware and takes +the per-token cost of a hosted embedding API to zero, at the price of +running two containers. + +## The 30-second version + +TEI serves exactly one model per container, so embedding and reranking are +two containers on two ports. A working pair: + +```bash +# Embeddings: a bi-encoder on port 8083 +docker run --rm --name tei-embed \ + --gpus '"device=0"' \ + -p 8083:80 \ + -v $HOME/.cache/huggingface:/data \ + ghcr.io/huggingface/text-embeddings-inference:86-1.9 \ + --model-id BAAI/bge-small-en-v1.5 \ + --port 80 + +# Reranking: a cross-encoder on port 8082 +docker run --rm --name tei-reranker \ + --gpus '"device=0"' \ + -p 8082:80 \ + -v $HOME/.cache/huggingface:/data \ + ghcr.io/huggingface/text-embeddings-inference:86-1.9 \ + --model-id BAAI/bge-reranker-v2-m3 \ + --port 80 \ + --auto-truncate false +``` + +TEI listens on port 80 inside the container; `-p 8083:80` maps it to a +host port. The `-v` mount is TEI's model cache, so a restart reloads from +disk instead of re-downloading. First boot pulls the model (seconds for a +small embedder, a minute or two for a larger reranker); later boots are +sub-second. + +## One model per container + +A single TEI process loads one model, and the endpoints it exposes follow +that model's architecture. An embedding bi-encoder maps one text to one +vector and answers `/embed` and `/v1/embeddings`. A reranking cross-encoder +scores a `(query, document)` pair to one number and answers `/rerank`. +They are different model families, so a full retrieval stack runs two +containers. They share the image and the cache mount and coexist on one +GPU with room to spare (a small embedder needs a few hundred MB, a +large reranker a couple of GB). + +| Endpoint | Container | Purpose | +|---|---|---| +| `POST /embed` | embedding | Dense vector(s) for the input text(s) | +| `POST /v1/embeddings` | embedding | OpenAI-compatible embeddings surface | +| `POST /rerank` | reranker | Score documents against a query | +| `GET /health` | both | Liveness (200 OK, empty body) | +| `GET /info` | both | Reports `model_id`, model type, `max_input_length` | +| `GET /metrics` | both | Prometheus metrics | + +## Choosing the image + +TEI publishes **GPU-architecture-specific images**; the wrong tag will not +start. Pick by your card's CUDA compute capability: + +| Tag prefix | Architecture | Example cards | +|---|---|---| +| `cpu-` | none | CPU-only, no GPU | +| (plain, no prefix) | Ampere 8.0 | A100, A30 | +| `86-` | Ampere 8.6 | RTX 3090, A10, A40 | +| `89-` | Ada Lovelace | RTX 4090, L4, L40S | +| `hopper-` | Hopper 9.0 | H100 | + +The suffix is the TEI version line (`1.9` in the examples, current as of +mid-2026). So an RTX 3090 uses `:86-1.9`, an RTX 4090 uses `:89-1.9`. The +tag is architecture-specific, not model-specific, so both containers use +the same image and it is cached once. HuggingFace publishes the full +matrix in the [TEI supported-models docs](https://huggingface.co/docs/text-embeddings-inference/supported_models). + +## Wiring the providers + +Point the TEI providers at the host ports. `base_url` is the instance root +(the provider appends `/embed` or `/rerank` itself), and `model` is the +model you loaded, so the observability layer reports the right identifier: + +```python +from openarmature.retrieval import TeiEmbeddingProvider, TeiRerankProvider + +embedder = TeiEmbeddingProvider( + base_url="http://localhost:8083", + model="BAAI/bge-small-en-v1.5", +) +reranker = TeiRerankProvider( + base_url="http://localhost:8082", + model="BAAI/bge-reranker-v2-m3", +) +``` + +TEI reports no token usage on either surface, so `response.usage` is +`None` for TEI calls; that is the nullable-usage contract, not a bug. + +`chunk_size` (default 32) is TEI's per-request batch cap, the server's +`--max-client-batch-size`. `embed()` splits a longer input list into +consecutive chunks under this size and stitches the vectors back in input +order, and `rerank()` chunks a large candidate pool the same way. Set it +to match the server if you raise TEI's limit: + +```python +embedder = TeiEmbeddingProvider( + base_url="http://localhost:8083", + model="BAAI/bge-small-en-v1.5", + chunk_size=64, # match TEI's --max-client-batch-size if you raise it +) +``` + +## Fail loud on over-length input + +The reranker command sets `--auto-truncate false`. By default TEI silently +clips an input that exceeds the model's token window, which quietly changes +what you scored; with auto-truncate off, an over-length `(query, document)` +pair returns a validation error instead. That surfaces through the provider +as `ProviderInvalidRequest`, so an over-length call fails loudly at the +boundary rather than returning a score computed on truncated text. Prefer +it for retrieval, where a silently-clipped document is a correctness bug, +not a warning. The embedding container defaults are usually fine, since a +single short text rarely overflows the window. + +## Readiness and smoke tests + +`ready()` on either provider issues a minimal probe against its endpoint +(TEI serves no model catalog, so it is an actual embed / rerank call). +Before wiring the providers, confirm the containers directly: + +```bash +# Liveness (both): 200 OK, empty body +curl http://localhost:8083/health +curl http://localhost:8082/health + +# Model identity and max input length +curl -s http://localhost:8083/info | jq +curl -s http://localhost:8082/info | jq + +# Embed: one vector per input; index [0] is the vector +curl -s http://localhost:8083/embed \ + -H 'Content-Type: application/json' \ + -d '{"inputs": "water ice in permanently shadowed lunar craters"}' | jq '.[0] | length' + +# Rerank: on-topic docs score high, off-topic near zero, sorted descending +curl -s http://localhost:8082/rerank \ + -H 'Content-Type: application/json' \ + -d '{ + "query": "where is there water on the Moon?", + "texts": [ + "Ice sits in permanently shadowed craters at the lunar poles.", + "The Sea of Tranquility was the Apollo 11 landing site." + ] + }' | jq +``` + +## Production deployment + +Run each container under a process supervisor so it restarts on failure +and boots with the host. A systemd unit that wraps `docker run` in the +foreground (no `-d`) gives you `systemctl start/stop` and `journalctl` +logs, matching how you would run any other serving container: + +```ini +# /etc/systemd/system/tei-embed.service +[Unit] +Description=TEI embeddings (bge-small-en-v1.5) +Requires=docker.service +After=docker.service network-online.target +Wants=network-online.target + +[Service] +# Foreground docker run so systemd owns the lifecycle. --rm cleans up on +# stop; the ExecStartPre clears any stale container a hard crash left behind. +ExecStartPre=-/usr/bin/docker rm -f tei-embed +ExecStart=/usr/bin/docker run --rm --name tei-embed \ + --gpus '"device=0"' \ + -p 8083:80 \ + -v /home/youruser/.cache/huggingface:/data \ + ghcr.io/huggingface/text-embeddings-inference:86-1.9 \ + --model-id BAAI/bge-small-en-v1.5 \ + --port 80 +ExecStop=/usr/bin/docker stop tei-embed +Restart=always +RestartSec=5 +# First boot pulls the image + model before the port opens; raise the +# start timeout so systemd does not kill the unit mid-download. +TimeoutStartSec=300 + +[Install] +WantedBy=multi-user.target +``` + +The reranker unit is the same shape with its own name, port, model, and +the `--auto-truncate false` flag. `systemctl enable --now tei-embed` boots +it and starts it; a second unit does the reranker. + +A few production notes: + +- **Wait for the network.** `network-online.target` in the unit matters: + the first boot pulls the model from HuggingFace, and pairing it with a + generous `TimeoutStartSec` keeps systemd from killing the unit while the + download is still in flight (a small embedder is quick; a large reranker + can take a minute or two cold). +- **GPU pinning.** `--gpus '"device=0"'` binds a container to one GPU by + index. On a multi-GPU host where indices can reorder across reboots, pin + by UUID instead (`--gpus '"device=GPU-..."'`, from `nvidia-smi -L`) so a + container always lands on the intended card. +- **VRAM planning.** Both containers fit comfortably on a single 24 GB + card (a small embedder plus a large reranker is a few GB total), leaving + headroom for other work. Check placement with `nvidia-smi` after start. +- **Cache volume.** Point the `/data` mount at the HuggingFace cache that + holds the weights (an absolute path, since a `multi-user.target` unit + runs as root and `%h` would resolve to `/root`). Keep it on fast local + disk; it is the difference between a sub-second restart and a + re-download. +- **Health checks.** Point your orchestrator's liveness probe at + `GET /health` (200 OK) and readiness at `GET /info` (confirms the model + loaded). + +## Where to next + +- **[Retrieval Providers](index.md)**: the bundled providers, the + protocol contract, and the error categories. +- **[Retrieval concept page](../concepts/retrieval.md)**: embedding, + reranking, `input_type`, chunking, and the nullable-usage contract. +- **[Authoring a Provider](authoring.md)**: implement the protocol for a + backend TEI does not cover. diff --git a/examples/README.md b/examples/README.md index 7d138ba1..9d11ef51 100644 --- a/examples/README.md +++ b/examples/README.md @@ -106,6 +106,25 @@ with a hard turn cap to prevent runaway loops. Demonstrates: `ToolMessage(tool_call_id=...)` round-trip, multi-turn tool-calling loop as a graph cycle. +### Retrieval + +#### [`retrieval-rag/`](./retrieval-rag/main.py) + +Retrieval-augmented answering over a lunar knowledge base, using the +two-stage retrieve-then-rerank pattern before generation. Batch-embeds a +small corpus once into an index (`OpenAIEmbeddingProvider.embed` over a +list, one vector per passage), then per query embeds the question, ranks +the corpus by cosine similarity for a broad shortlist, reranks that +shortlist with a cross-encoder (`CohereRerankProvider.rerank`) for +precision, and grounds an LLM answer in the reranked top passages. The +retrieve and rerank steps run as graph nodes, so their `EmbeddingEvent` / +`RerankEvent` reach any attached observer. Demonstrates: +`OpenAIEmbeddingProvider` and `CohereRerankProvider` from +`openarmature.retrieval`, the `input_type` query/document knob (a wire +no-op on symmetric OpenAI, meaningful on asymmetric providers), mapping +`ScoredDocument` results back by `.index`, and retrieval feeding an +`OpenAIProvider` answer node. Needs `OPENAI_API_KEY` and `COHERE_API_KEY`. + ### Reliability #### [`checkpointing-and-migration/`](./checkpointing-and-migration/main.py) diff --git a/examples/retrieval-rag/main.py b/examples/retrieval-rag/main.py new file mode 100644 index 00000000..ca8255ce --- /dev/null +++ b/examples/retrieval-rag/main.py @@ -0,0 +1,320 @@ +"""Retrieval-augmented answering over a lunar knowledge base. + +A question about the Moon is answered from a small corpus of passages, +using the two-stage retrieval pattern before generation: + +1. **Index** the corpus once, offline: batch-embed every passage into a + vector (``OpenAIEmbeddingProvider.embed`` over a list returns one + vector per passage, in input order). +2. **Retrieve** per query: embed the question, rank the corpus by cosine + similarity, and keep the top handful of candidates. Cheap and broad. +3. **Rerank** those candidates with a cross-encoder + (``CohereRerankProvider.rerank``), which scores each candidate + against the query directly. More accurate than similarity, so it + reorders the shortlist and drops the weakest. +4. **Generate** a grounded answer from the reranked passages with an LLM. + +Retrieval gives recall (find plausibly-relevant passages fast); reranking +gives precision (put the genuinely-relevant ones first). Feeding only the +reranked top few to the model keeps the context tight and on-topic. + +**Demonstrates:** + +- ``OpenAIEmbeddingProvider`` from ``openarmature.retrieval``: batch + ``embed`` for the index, single ``embed`` for the query, one vector + per input in input order. +- ``EmbeddingRuntimeConfig(input_type=...)``: the query/document knob. + On OpenAI it is a wire no-op (the model is symmetric), but the same + call selects the right representation on an asymmetric provider (TEI, + Cohere, Jina), so setting it keeps the pipeline portable. +- ``CohereRerankProvider`` from ``openarmature.retrieval``: ``rerank`` + returns ``ScoredDocument`` results sorted by relevance. Cohere does + not echo the document text, so results are mapped back to the + candidate list by ``ScoredDocument.index`` rather than ``.document``. +- The retrieval providers driven inside graph nodes, so their + ``EmbeddingEvent`` / ``RerankEvent`` reach the attached ``trace`` + observer (which prints them), the same way LLM completions do; the + offline index build runs outside the graph and is not observed. +- An ``OpenAIProvider`` answer node grounded in the reranked passages. + +**Configuration** (env vars): + +- ``OPENAI_API_KEY``: required (used for both embeddings and the answer). +- ``OPENAI_BASE_URL``: host root, defaults to ``https://api.openai.com``. + Point it at any OpenAI-compatible endpoint; the impl appends the + ``/v1`` routes, so do not include ``/v1``. +- ``OPENAI_EMBED_MODEL``: defaults to ``text-embedding-3-small``. +- ``OPENAI_CHAT_MODEL``: defaults to ``gpt-4o-mini``. +- ``COHERE_API_KEY``: required (used for reranking). +- ``COHERE_RERANK_MODEL``: defaults to ``rerank-v3.5``. + +Run with: + + uv sync --group examples + OPENAI_API_KEY=sk-... COHERE_API_KEY=... uv run python examples/retrieval-rag/main.py +""" + +from __future__ import annotations + +import asyncio +import math +import os +from collections.abc import Mapping, Sequence +from typing import Any + +from openarmature.graph import ( + END, + CompiledGraph, + EmbeddingEvent, + GraphBuilder, + ObserverEvent, + RerankEvent, + State, +) +from openarmature.llm import OpenAIProvider, RuntimeConfig, UserMessage +from openarmature.retrieval import ( + CohereRerankProvider, + EmbeddingRuntimeConfig, + OpenAIEmbeddingProvider, + RerankRuntimeConfig, +) + +# A tiny lunar knowledge base. In a real app the corpus is far larger and +# its vectors live in a vector store; here it is inline so the demo is +# self-contained. The passages deliberately overlap in topic so reranking +# has something to disambiguate. +CORPUS: list[str] = [ + "The Moon's south pole holds water ice in permanently shadowed crater floors " + "that never see sunlight, a candidate resource for future crews.", + "Apollo 11 landed in the Sea of Tranquility in July 1969; Armstrong and Aldrin " + "spent about two and a half hours walking on the surface.", + "The lunar maria are vast basaltic plains that formed when ancient impact basins " + "flooded with lava, giving the near side its dark patches.", + "A lunar day lasts about 29.5 Earth days, so any point on the surface sees roughly " + "two weeks of sunlight followed by two weeks of night.", + "The Moon is slowly receding from Earth at about 3.8 centimeters per year, measured " + "by bouncing lasers off the retroreflectors Apollo crews left behind.", + "Regolith, the layer of loose dust and broken rock covering the Moon, is abrasive " + "and clings to everything, a persistent hazard for equipment and spacesuits.", + "Apollo 13 aborted its landing after an oxygen tank ruptured; the crew looped around " + "the Moon and returned safely using the lunar module as a lifeboat.", + "Permanently shadowed regions near the poles are among the coldest places in the " + "solar system, cold enough to trap water ice for billions of years.", +] + +# How many candidates cosine retrieval hands to the reranker, and how many +# survive reranking to ground the answer. Retrieve broad, rerank narrow. +_RETRIEVE_K = 4 +_RERANK_K = 2 + + +# Lazy, module-level singletons: constructed on first use inside a node +# body (or the index build), never at import time, so importing this +# module for inspection does not open an httpx connection pool. Same +# pattern as the other examples. +_embedder: OpenAIEmbeddingProvider | None = None +_reranker: CohereRerankProvider | None = None +_answerer: OpenAIProvider | None = None + +# The offline-built index: corpus passage vectors, in CORPUS order. Built +# once in main() before the graph runs. +_corpus_vectors: list[list[float]] = [] + + +def _get_embedder() -> OpenAIEmbeddingProvider: + global _embedder + if _embedder is None: + _embedder = OpenAIEmbeddingProvider( + base_url=os.environ.get("OPENAI_BASE_URL", "https://api.openai.com"), + model=os.environ.get("OPENAI_EMBED_MODEL", "text-embedding-3-small"), + api_key=os.environ.get("OPENAI_API_KEY") or None, + ) + return _embedder + + +def _get_reranker() -> CohereRerankProvider: + global _reranker + if _reranker is None: + _reranker = CohereRerankProvider( + model=os.environ.get("COHERE_RERANK_MODEL", "rerank-v3.5"), + api_key=os.environ.get("COHERE_API_KEY") or None, + ) + return _reranker + + +def _get_answerer() -> OpenAIProvider: + global _answerer + if _answerer is None: + _answerer = OpenAIProvider( + base_url=os.environ.get("OPENAI_BASE_URL", "https://api.openai.com"), + model=os.environ.get("OPENAI_CHAT_MODEL", "gpt-4o-mini"), + api_key=os.environ.get("OPENAI_API_KEY") or None, + ) + return _answerer + + +def _cosine(a: Sequence[float], b: Sequence[float]) -> float: + # OpenAI embeddings are unit-normalized, so the dot product already + # equals cosine similarity; the full formula is spelled out here so the + # example stays correct for any provider whose vectors are not + # pre-normalized. + dot = sum(x * y for x, y in zip(a, b, strict=True)) + norm_a = math.sqrt(sum(x * x for x in a)) + norm_b = math.sqrt(sum(y * y for y in b)) + if norm_a == 0.0 or norm_b == 0.0: + return 0.0 + return dot / (norm_a * norm_b) + + +class SearchState(State): + query: str + # Corpus indices of the cosine-retrieved candidates, best-first. + candidate_indices: list[int] = [] + # Corpus indices after reranking, best-first (a reordered, trimmed + # subset of candidate_indices). + ranked_indices: list[int] = [] + answer: str | None = None + + +async def retrieve(state: SearchState) -> Mapping[str, Any]: + # The corpus index is process-global state, built once by _build_index() + # before the graph runs (in a real app it would live in a vector store, + # not a module global). Guard the dependency explicitly: without it, + # retrieve would return no candidates and the failure would surface two + # nodes later as an opaque "empty documents" rerank error that never + # names the real cause. + if not _corpus_vectors: + raise RuntimeError("corpus index not built; call _build_index() before invoking the graph") + + # Embed the query and rank the corpus index by cosine similarity. + # input_type="query" is the asymmetric-embedding knob: a no-op on + # OpenAI's symmetric model, but the same code selects the query + # representation on TEI / Cohere / Jina, so the pipeline ports without + # change. embed() takes a list and returns one vector per input; here + # the list is a single query. + response = await _get_embedder().embed( + [state.query], + config=EmbeddingRuntimeConfig(input_type="query"), + ) + query_vector = response.vectors[0] + + scored = sorted( + range(len(_corpus_vectors)), + key=lambda i: _cosine(query_vector, _corpus_vectors[i]), + reverse=True, + ) + return {"candidate_indices": scored[:_RETRIEVE_K]} + + +async def rerank(state: SearchState) -> Mapping[str, Any]: + # Cross-encoder reranking of the retrieved shortlist. The reranker + # scores each candidate against the query directly (more accurate than + # vector similarity) and returns results sorted best-first. Cohere does + # not echo the document text, so results map back to the corpus by + # ScoredDocument.index, not .document. return_documents is set to show + # the knob; it is a no-op on the Cohere wire. + candidates = [CORPUS[i] for i in state.candidate_indices] + response = await _get_reranker().rerank( + state.query, + candidates, + top_k=_RERANK_K, + config=RerankRuntimeConfig(return_documents=True), + ) + # result.index points into `candidates`; translate back to corpus index. + ranked = [state.candidate_indices[result.index] for result in response.results] + return {"ranked_indices": ranked} + + +async def answer(state: SearchState) -> Mapping[str, Any]: + # Ground the answer in the reranked passages only. Passing the tight, + # reordered top-K (rather than everything retrieved) keeps the context + # on-topic and the answer faithful. + context = "\n".join(f"- {CORPUS[i]}" for i in state.ranked_indices) + response = await _get_answerer().complete( + [ + UserMessage( + content=( + f"Answer the question using only these passages. If they do not " + f"contain the answer, say so.\n\nPassages:\n{context}\n\n" + f"Question: {state.query}" + ) + ) + ], + config=RuntimeConfig(temperature=0.0), + ) + return {"answer": response.message.content} + + +async def trace(event: ObserverEvent) -> None: + # Every embed / rerank call made inside a graph node dispatches a typed + # EmbeddingEvent / RerankEvent to any attached observer, the same way LLM + # completions do, so retrieval is a first-class boundary in your telemetry. + # A real deployment routes these to OTel or Langfuse (see the observability + # examples); this one just prints them. The offline index build runs outside + # the graph, so its embed dispatches nothing and does not appear here. + if isinstance(event, EmbeddingEvent): + print(f" [obs] embed {event.node_name}: {event.input_count} in / {event.dimensions}d") + elif isinstance(event, RerankEvent): + print(f" [obs] rerank {event.node_name}: {event.document_count} in / top {event.result_count}") + + +def build_graph() -> CompiledGraph[SearchState]: + return ( + GraphBuilder(SearchState) + .add_node("retrieve", retrieve) + .add_node("rerank", rerank) + .add_node("answer", answer) + .add_edge("retrieve", "rerank") + .add_edge("rerank", "answer") + .add_edge("answer", END) + .set_entry("retrieve") + .compile() + ) + + +async def _build_index() -> None: + # Offline index build: batch-embed the whole corpus once. Runs outside + # any graph, so it dispatches no observer events (unlike the per-query + # embed inside retrieve()). input_type="document" is the counterpart to + # the query knob above. embed() over a list returns one vector per + # passage, in CORPUS order, so the index lines up positionally. + global _corpus_vectors + response = await _get_embedder().embed( + CORPUS, + config=EmbeddingRuntimeConfig(input_type="document"), + ) + _corpus_vectors = response.vectors + + +async def main() -> None: + # build_graph() only compiles the topology and opens no network client, + # so it is safe before the try. Everything that constructs a provider + # (the index build and the per-query runs) goes inside the try, so a + # failure on any of them, including a bad or missing API key on the very + # first embed, still closes the httpx clients in the finally. + graph = build_graph() + graph.attach_observer(trace) + try: + await _build_index() + print(f"indexed {len(_corpus_vectors)} passages\n") + for query in ( + "Why did Apollo 13 not land on the Moon?", + "Where might future crews find water on the Moon?", + ): + final = await graph.invoke(SearchState(query=query)) + print(f"Q: {query}") + print(f" retrieved: {final.candidate_indices}") + print(f" reranked: {final.ranked_indices}") + print(f" A: {final.answer}\n") + finally: + await graph.drain() + if _embedder is not None: + await _embedder.aclose() + if _reranker is not None: + await _reranker.aclose() + if _answerer is not None: + await _answerer.aclose() + + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/mkdocs.yml b/mkdocs.yml index 2825a812..4b7d134a 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -72,6 +72,8 @@ plugins: - examples/*.md Model Providers: - model-providers/*.md + Retrieval Providers: + - retrieval-providers/*.md Reference: - reference/*.md @@ -114,6 +116,7 @@ nav: - Fan-out: concepts/fan-out.md - Parallel branches: concepts/parallel-branches.md - LLMs: concepts/llms.md + - Retrieval: concepts/retrieval.md - Prompts: concepts/prompts.md - Observability: concepts/observability.md - Checkpointing: concepts/checkpointing.md @@ -121,10 +124,15 @@ nav: - model-providers/index.md - Self-hosted vLLM: model-providers/vllm.md - Authoring a Provider: model-providers/authoring.md + - Retrieval Providers: + - retrieval-providers/index.md + - Self-hosting TEI: retrieval-providers/tei.md + - Authoring a Provider: retrieval-providers/authoring.md - Reference: - reference/index.md - openarmature.graph: reference/graph.md - openarmature.llm: reference/llm.md + - openarmature.retrieval: reference/retrieval.md - openarmature.prompts: reference/prompts.md - openarmature.checkpoint: reference/checkpoint.md - openarmature.observability: reference/observability.md diff --git a/src/openarmature/AGENTS.md b/src/openarmature/AGENTS.md index 05cea805..5b62fbb7 100644 --- a/src/openarmature/AGENTS.md +++ b/src/openarmature/AGENTS.md @@ -1617,6 +1617,7 @@ _Runnable example programs shipped in the source tree at `examples/`. The full c - **`examples/observer-hooks/main.py`** — openarmature demo: observer hooks for structured logging, per-call metrics, and OTel spans. - **`examples/parallel-branches/main.py`** — openarmature demo: enrich a lunar-mission news article with several independent analyses running concurrently. - **`examples/production-observability/main.py`** — openarmature demo: production observability with dual OTel + Langfuse observers, caller hooks for trace.input/output, and the canonical TimingMiddleware. +- **`examples/retrieval-rag/main.py`** — Retrieval-augmented answering over a lunar knowledge base. - **`examples/routing-and-subgraphs/main.py`** — openarmature demo: conditional routing + subgraph with a custom projection. - **`examples/tool-use/main.py`** — openarmature demo: a lunar-mission assistant that calls local Python functions as tools to answer fact and physics questions about Apollo / Artemis missions. diff --git a/src/openarmature/graph/__init__.py b/src/openarmature/graph/__init__.py index a65a6eec..8647fc94 100644 --- a/src/openarmature/graph/__init__.py +++ b/src/openarmature/graph/__init__.py @@ -39,6 +39,8 @@ UnreachableNode, ) from .events import ( + EmbeddingEvent, + EmbeddingFailedEvent, FailureIsolatedEvent, InvocationCompletedEvent, InvocationStartedEvent, @@ -47,6 +49,8 @@ LlmRetryAttemptEvent, MetadataAugmentationEvent, NodeEvent, + RerankEvent, + RerankFailedEvent, ) from .fan_out import FanOutConfig, FanOutNode from .middleware import ( @@ -82,6 +86,8 @@ "DegradedUpdate", "DrainSummary", "EdgeException", + "EmbeddingEvent", + "EmbeddingFailedEvent", "EndSentinel", "ExplicitMapping", "FailureIsolatedEvent", @@ -123,6 +129,8 @@ "Reducer", "ReducerError", "RemoveHandle", + "RerankEvent", + "RerankFailedEvent", "RetryConfig", "RetryMiddleware", "RoutingError", diff --git a/tests/test_examples_smoke.py b/tests/test_examples_smoke.py index ea1c0dbe..31034657 100644 --- a/tests/test_examples_smoke.py +++ b/tests/test_examples_smoke.py @@ -43,6 +43,7 @@ "observer-hooks", "langfuse-observability", "production-observability", + "retrieval-rag", ] From a84e6f4b8dfd6d472cf5da0728d672639983a8aa Mon Sep 17 00:00:00 2001 From: chris-colinsky Date: Sat, 18 Jul 2026 10:40:32 -0700 Subject: [PATCH 2/5] Fix implicit string concat and skeleton 404 mapping Address the review on PR #220. Wrap each retrieval-rag CORPUS passage in parentheses so the multi-line string is an explicit grouped expression, not adjacent string literals implicitly concatenated inside the list (a CodeQL missing-comma footgun in a copy-me example). Add the 404 -> ProviderInvalidModel branch and import to the authoring-guide skeleton's _classify(), which contradicted its own contract checklist. --- docs/retrieval-providers/authoring.md | 3 ++ examples/retrieval-rag/main.py | 48 ++++++++++++++++++--------- 2 files changed, 35 insertions(+), 16 deletions(-) diff --git a/docs/retrieval-providers/authoring.md b/docs/retrieval-providers/authoring.md index a60f0ff9..a009ce10 100644 --- a/docs/retrieval-providers/authoring.md +++ b/docs/retrieval-providers/authoring.md @@ -23,6 +23,7 @@ import httpx from openarmature.llm import ( ProviderAuthentication, + ProviderInvalidModel, ProviderInvalidRequest, ProviderInvalidResponse, ProviderRateLimit, @@ -99,6 +100,8 @@ def _classify(resp: httpx.Response) -> Exception: status = resp.status_code if status in (401, 403): return ProviderAuthentication(f"HTTP {status}") + if status == 404: + return ProviderInvalidModel("model not found") if status == 429: return ProviderRateLimit("HTTP 429") if status in (400, 413, 422): diff --git a/examples/retrieval-rag/main.py b/examples/retrieval-rag/main.py index ca8255ce..9206aba3 100644 --- a/examples/retrieval-rag/main.py +++ b/examples/retrieval-rag/main.py @@ -84,22 +84,38 @@ # self-contained. The passages deliberately overlap in topic so reranking # has something to disambiguate. CORPUS: list[str] = [ - "The Moon's south pole holds water ice in permanently shadowed crater floors " - "that never see sunlight, a candidate resource for future crews.", - "Apollo 11 landed in the Sea of Tranquility in July 1969; Armstrong and Aldrin " - "spent about two and a half hours walking on the surface.", - "The lunar maria are vast basaltic plains that formed when ancient impact basins " - "flooded with lava, giving the near side its dark patches.", - "A lunar day lasts about 29.5 Earth days, so any point on the surface sees roughly " - "two weeks of sunlight followed by two weeks of night.", - "The Moon is slowly receding from Earth at about 3.8 centimeters per year, measured " - "by bouncing lasers off the retroreflectors Apollo crews left behind.", - "Regolith, the layer of loose dust and broken rock covering the Moon, is abrasive " - "and clings to everything, a persistent hazard for equipment and spacesuits.", - "Apollo 13 aborted its landing after an oxygen tank ruptured; the crew looped around " - "the Moon and returned safely using the lunar module as a lifeboat.", - "Permanently shadowed regions near the poles are among the coldest places in the " - "solar system, cold enough to trap water ice for billions of years.", + ( + "The Moon's south pole holds water ice in permanently shadowed crater floors " + "that never see sunlight, a candidate resource for future crews." + ), + ( + "Apollo 11 landed in the Sea of Tranquility in July 1969; Armstrong and Aldrin " + "spent about two and a half hours walking on the surface." + ), + ( + "The lunar maria are vast basaltic plains that formed when ancient impact basins " + "flooded with lava, giving the near side its dark patches." + ), + ( + "A lunar day lasts about 29.5 Earth days, so any point on the surface sees roughly " + "two weeks of sunlight followed by two weeks of night." + ), + ( + "The Moon is slowly receding from Earth at about 3.8 centimeters per year, measured " + "by bouncing lasers off the retroreflectors Apollo crews left behind." + ), + ( + "Regolith, the layer of loose dust and broken rock covering the Moon, is abrasive " + "and clings to everything, a persistent hazard for equipment and spacesuits." + ), + ( + "Apollo 13 aborted its landing after an oxygen tank ruptured; the crew looped around " + "the Moon and returned safely using the lunar module as a lifeboat." + ), + ( + "Permanently shadowed regions near the poles are among the coldest places in the " + "solar system, cold enough to trap water ice for billions of years." + ), ] # How many candidates cosine retrieval hands to the reranker, and how many From 21284ca58b412f37822d7b1a43b45104709fb470 Mon Sep 17 00:00:00 2001 From: chris-colinsky Date: Sat, 18 Jul 2026 15:28:20 -0700 Subject: [PATCH 3/5] Use typed default_factory for SearchState list fields Align the retrieval-rag SearchState list defaults with the codebase's Field(default_factory=...) idiom (PR #220 review). Use the typed default_factory=list[int] rather than a bare default_factory=list, since the bare form infers list[Unknown] and trips strict pyright's reportUnknownVariableType on a list[int] field. --- examples/retrieval-rag/main.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/examples/retrieval-rag/main.py b/examples/retrieval-rag/main.py index 9206aba3..07dce3aa 100644 --- a/examples/retrieval-rag/main.py +++ b/examples/retrieval-rag/main.py @@ -62,6 +62,8 @@ from collections.abc import Mapping, Sequence from typing import Any +from pydantic import Field + from openarmature.graph import ( END, CompiledGraph, @@ -185,10 +187,12 @@ def _cosine(a: Sequence[float], b: Sequence[float]) -> float: class SearchState(State): query: str # Corpus indices of the cosine-retrieved candidates, best-first. - candidate_indices: list[int] = [] + # (default_factory takes the typed list[int] so it stays clean under + # strict typing, where a bare default_factory=list infers list[Unknown].) + candidate_indices: list[int] = Field(default_factory=list[int]) # Corpus indices after reranking, best-first (a reordered, trimmed # subset of candidate_indices). - ranked_indices: list[int] = [] + ranked_indices: list[int] = Field(default_factory=list[int]) answer: str | None = None From 362f3af32df50280b625168adc15e0a24f9d96df Mon Sep 17 00:00:00 2001 From: chris-colinsky Date: Sat, 18 Jul 2026 15:51:25 -0700 Subject: [PATCH 4/5] Drop stale example count from README The Next Steps 'Examples' bullet hard-coded 'ten runnable demos', which was already stale and this PR adds another. Remove the fixed number so it stops drifting as examples change (PR #220 review). --- README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/README.md b/README.md index 702379a9..76e1f79b 100644 --- a/README.md +++ b/README.md @@ -180,7 +180,7 @@ A few things to notice: - **Concepts**: typed state, reducers, graphs, composition, fan-out, parallel branches, LLMs, retrieval, prompts, observability, checkpointing. [openarmature.ai/concepts](https://openarmature.ai/concepts/) - **Model Providers**: implement the Provider Protocol for a custom LLM backend. [openarmature.ai/model-providers/authoring](https://openarmature.ai/model-providers/authoring/) - **API reference**: auto-generated from docstrings. [openarmature.ai/reference](https://openarmature.ai/reference/) -- **Examples**: ten runnable demos with walk-throughs. [openarmature.ai/examples](https://openarmature.ai/examples/) (source at [./examples/](./examples/)) +- **Examples**: runnable demos with walk-throughs. [openarmature.ai/examples](https://openarmature.ai/examples/) (source at [./examples/](./examples/)) - **Spec**: behavioral contract this implementation conforms to. [LunarCommand/openarmature-spec](https://github.com/LunarCommand/openarmature-spec) ## For AI agents From 231d352b0631888b4865696c672f65ce0f8b4990 Mon Sep 17 00:00:00 2001 From: chris-colinsky Date: Sat, 18 Jul 2026 16:06:57 -0700 Subject: [PATCH 5/5] Map malformed 200 body to ProviderInvalidResponse in skeleton The authoring-guide embed() skeleton parsed the 200 response body without guarding it, so a malformed body would leak a raw KeyError/ValueError instead of the ProviderInvalidResponse its own contract checklist prescribes. Wrap the parse (PR #220 review). --- docs/retrieval-providers/authoring.md | 13 +++++++++---- 1 file changed, 9 insertions(+), 4 deletions(-) diff --git a/docs/retrieval-providers/authoring.md b/docs/retrieval-providers/authoring.md index a009ce10..8506084a 100644 --- a/docs/retrieval-providers/authoring.md +++ b/docs/retrieval-providers/authoring.md @@ -70,10 +70,15 @@ class MyEmbeddingProvider: if resp.status_code != 200: raise _classify(resp) - payload = resp.json() - # Order the returned data by index so vectors line up with input. - entries = sorted(payload["data"], key=lambda e: e["index"]) - vectors = [[float(x) for x in e["embedding"]] for e in entries] + # A malformed 200 body (bad JSON, missing keys, wrong types) is a + # provider_invalid_response, not a leaked KeyError/ValueError. + try: + payload = resp.json() + # Order the returned data by index so vectors line up with input. + entries = sorted(payload["data"], key=lambda e: e["index"]) + vectors = [[float(x) for x in e["embedding"]] for e in entries] + except (ValueError, KeyError, TypeError) as exc: + raise ProviderInvalidResponse("malformed embeddings response") from exc # Validates the count and uniform dimensionality, returns the dim. dimensions = validate_embedding_response(vectors, len(input_strings))