diff --git a/README.md b/README.md index 9f43e211..76e1f79b 100644 --- a/README.md +++ b/README.md @@ -177,10 +177,10 @@ 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/)) +- **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 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..8506084a --- /dev/null +++ b/docs/retrieval-providers/authoring.md @@ -0,0 +1,242 @@ +# 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, + ProviderInvalidModel, + 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) + + # 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)) + + 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 == 404: + return ProviderInvalidModel("model not found") + 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..07dce3aa --- /dev/null +++ b/examples/retrieval-rag/main.py @@ -0,0 +1,340 @@ +"""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 pydantic import Field + +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. + # (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] = Field(default_factory=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", ]