Appendix G: Retrieval for Atlas (the RAG the Book Assumes)

What search_kb actually sits behind - chunking, embeddings, and a retrieval-specific eval distinct from Chapter 21’s agent-level suite - kept minimal on purpose.

Chapter 7 gave search_kb an interface - a query in, an article or a “no match” string out - and treated everything behind that interface as a seeded backend, deliberately not re-taught. The graph reaches it through an adapter that returns scored documents (Doc), which Chapter 12’s select_docs ranks and caps. This book is about orchestration, not retrieval; a chapter-length RAG treatment would have added a chapter without advancing the thesis. This appendix opens the box just far enough to make Atlas reproducible and the retrieval interface’s assumptions honest, not to turn this into a RAG book.

What’s behind the interface

Chapter 7’s search_kb stays exactly what it was there: a plain, substring-matched dict lookup with no external account and no network call, the same seeded fixture every other chapter’s tests run against (Appendix A). What follows is the optional embedding pipeline a deployment would put in its place. It is not wired in by default. Each seeded article is chunked and every chunk embedded; a query gets back the closest chunks by cosine similarity, as the same Docs the retrieve node already expects.

atlas/retrieval.py

import functools
import math
import threading
from collections.abc import Callable
from typing import Protocol

from atlas.helpers import _last_user_text
from atlas.state import Doc
from atlas.tools import _KB, KnowledgeBaseUnavailable

NO_MATCH_BELOW = 0.3  # a starting point; set it from your recall_at_k set


class Embeddings(Protocol):
    """The narrow shape this module needs from an embeddings client - just
    `embed_query`, so a hand-rolled fake and a real `langchain` embeddings
    object both satisfy it without either depending on the other."""

    def embed_query(self, text: str) -> list[float]: ...


def default_embeddings() -> Embeddings:
    """Chapter 13's model, imported and built only when called. Raises at
    once if `OPENAI_API_KEY` is not set."""
    from langchain.embeddings import init_embeddings

    return init_embeddings("openai:text-embedding-3-small")


def chunk_article(text: str, size: int = 400, overlap: int = 50) -> list[str]:
    """Fixed-size windows with overlap: a sentence cut at one boundary
    reads whole in the neighboring chunk."""
    if overlap < 0 or size <= overlap:
        raise ValueError("need 0 <= overlap < size, or the window never advances")
    chunks, start = [], 0
    while start < len(text):
        chunks.append(text[start : start + size])
        start += size - overlap
    return chunks


def cosine_similarity(a: list[float], b: list[float]) -> float:
    """0.0 for an all-zero vector rather than a division by zero."""
    if len(a) != len(b):
        raise ValueError(f"vectors differ in length: {len(a)} vs {len(b)}")
    dot = sum(x * y for x, y in zip(a, b))
    norm = math.sqrt(sum(x * x for x in a)) * math.sqrt(sum(y * y for y in b))
    return dot / norm if norm else 0.0

Fixed-size chunking with overlap is the simplest strategy there is. On the two seeded articles it never splits anything: each is about fifty characters, one chunk at a window of 400. It is there for a real knowledge base, where the overlap keeps a sentence cut at a boundary readable in one of the two chunks. The guards matter more than they look. Without the first, overlap >= size loops forever. Without the length check, two embedding models mixed in one index produce a quietly truncated score instead of an error, because zip stops at the shorter vector.

The index and the node-side adapter:

atlas/retrieval.py

class EmbeddedKB:
    """`_KB`'s articles, chunked and embedded on first search, then cached."""

    def __init__(self, embeddings: Embeddings | None = None) -> None:
        self._embeddings = embeddings  # None: default_embeddings() on first use
        self._index: list[tuple[str, str, list[float]]] | None = None
        self._lock = threading.RLock()  # one build, however many first requests

    def _embed(self, text: str) -> list[float]:
        try:
            with self._lock:
                if self._embeddings is None:
                    self._embeddings = default_embeddings()
            return self._embeddings.embed_query(text)
        except Exception as exc:  # any provider failure: escalate, never retry
            raise KnowledgeBaseUnavailable(f"embedding provider: {exc}") from exc

    def _build_index(self) -> list[tuple[str, str, list[float]]]:
        with self._lock:  # a failed build leaves None, so the next one retries
            if self._index is None:
                self._index = [
                    (f"kb:{article_id}#{n}", chunk, self._embed(chunk))
                    for article_id, text in _KB.items()
                    for n, chunk in enumerate(chunk_article(text))
                ]
            return self._index

    def search(self, query: str, k: int = 3) -> list[Doc]:
        """The top k chunks by cosine similarity, as the `Doc`s the
        `retrieve` node expects. Below `NO_MATCH_BELOW`, nothing."""
        query_vec = self._embed(query)
        docs = [
            Doc(id=doc_id, text=chunk, score=cosine_similarity(query_vec, vec))
            for doc_id, chunk, vec in self._build_index()
        ]
        top = sorted(docs, key=lambda doc: doc["score"], reverse=True)[:k]
        return [doc for doc in top if doc["score"] > NO_MATCH_BELOW]


_default_kb: EmbeddedKB | None = None
_default_kb_lock = threading.Lock()


def shared_kb() -> EmbeddedKB:
    """The one process-wide `EmbeddedKB`, created on first use."""
    global _default_kb
    with _default_kb_lock:
        if _default_kb is None:
            _default_kb = EmbeddedKB()
        return _default_kb


def search_kb_embedded(messages: list, kb: EmbeddedKB | None = None) -> list[Doc]:
    """`atlas.helpers.search_kb` over embeddings: same messages in, same
    `Doc`s out, so `retrieve`, `select_docs` and the escalation path are
    unchanged. Searches the shared index unless given a `kb`."""
    query = _last_user_text(messages)
    if not query:
        return []
    return (kb if kb is not None else shared_kb()).search(query)


def make_search_kb(kb: EmbeddedKB | None = None) -> Callable[[list], list[Doc]]:
    """The drop-in: `search(messages)`, with `kb` bound."""
    return functools.partial(search_kb_embedded, kb=kb)

Nothing is embedded until the first search, so importing the module, or constructing an EmbeddedKB, needs no key. A Doc’s id names the article and the chunk number. Its score is the cosine similarity, which is what select_docs sorts on.

One detail is new here: the dict-backed adapter Atlas runs by default scores every hit 1.0. Behind this one, Chapter 12’s ranking sorts real scores.

make_search_kb() is the drop-in. It returns a one-argument search(messages), the signature atlas.helpers.search_kb has, with the index bound: the shared one, or the EmbeddedKB you pass. build_graph takes no search argument. retrieve looks up the module-level name search_kb in atlas.graph on every call, so putting the drop-in in place means rebinding that one name before the graph takes traffic, the same swap the companion’s retrieval tests make:

import atlas.graph
from atlas.retrieval import make_search_kb

atlas.graph.search_kb = make_search_kb()  # the shared EmbeddedKB

The shared index is created, and each index built, under a lock. Under ainvoke the runtime runs the sync retrieve on a worker thread, and without the lock concurrent first requests would each embed the whole knowledge base.

Reusing Chapter 13’s text-embedding-3-small (1536 dimensions) rather than introducing a second embedding model is deliberate. The “keep the embedding model fixed across write and query” discipline Chapter 13 argued for semantic memory applies equally here, and a real deployment gains nothing from running two embedding models for two retrieval jobs.

What happens on a miss, and on an outage

The 0.3 cut decides which of Atlas’s existing paths a question takes. Below it, search returns no documents, and Chapter 6’s bounded retry takes over: retrieve runs up to three times (MAX_RETRIEVE_ATTEMPTS), then the run escalates. Above it, select_docs keeps the best chunks that fit the 2,000-token retrieved slice, and they reach the model only through reference_text, where Chapter 23’s screen_untrusted tags each one as untrusted content. Treat 0.3 as a starting point and move it with your recall_at_k set, not by eye.

The retry is less useful here than on Chapter 7’s backend. Re-running an embedding search on the same question returns the same miss, so the two extra attempts spend two embedding calls on nothing; Chapter 6 built that retry for an empty result that might fill on a second look.

An outage is the case the wrapper in _embed exists for. retrieve catches KnowledgeBaseUnavailable and records the error, and the run escalates: Chapter 6’s “tool failure goes to a human, not a fake answer.” An embedding provider’s own errors are plain Exception subclasses, which retrieve’s RetryPolicy would run three times before failing the run with no escalation at all. Translating them at the one place every provider call passes through keeps the embedded backend on the same failure path as the seeded one. The same design hides a misconfiguration: a missing key or a dead provider shows up only as every retrieval escalating to a human, so check health at startup with one search before the server takes traffic.

Figure 1 opens the box Chapter 6’s bounded retry loop treated as closed, with both backends behind the one call.

Figure 1: The retrieval path from retrieve to an answer or an escalation. retrieve calls search_kb(messages), served by the dict-backed default or the optional embedded drop-in, which raises KnowledgeBaseUnavailable on a provider error. route_after_retrieve escalates on an error, answers when documents came back, and otherwise retries, up to three attempts.

What’s deliberately not built

Reranking. A cross-encoder rerank pass over the top-N cosine results is a real production technique this appendix names and does not implement. The seeded KB is small enough that similarity search alone is a reasonable baseline, and adding a reranker to a fixture built to be “simple enough to read in full” (Appendix A) would cost more clarity than it buys accuracy on two seeded articles.

A real vector database. The index is an in-memory list, not pgvector or a dedicated vector store. That is correct for a seeded, no-external-accounts fixture (Appendix A), and it is the first thing to swap for PostgresStore’s own vector index (Chapter 13) or a dedicated vector database once a KB is too large to hold in memory.

A retrieval-specific eval, separate from Chapter 21’s

Chapter 21’s eval suite tests whether Atlas uses retrieved content well - did the right route fire, did the answer reflect what was retrieved. It does not test whether retrieval itself found the right content; a suite that only checks downstream behavior can pass every example while search_kb quietly returns the third-best chunk instead of the best one, because the model wrote a plausible answer from mediocre context regardless. Retrieval quality is a precondition those evals assume holds, not something they’re built to catch failing.

atlas/retrieval.py

def search_kb_impl(query: str, k: int = 3, kb: EmbeddedKB | None = None) -> list[Doc]:
    """The appendix's name for a search against one shared, lazily built
    index (or the `kb` passed in)."""
    return (kb if kb is not None else shared_kb()).search(query, k=k)


def recall_at_k(
    query: str, expected_article_id: str, k: int = 3, kb: EmbeddedKB | None = None
) -> bool:
    """Did a chunk of the expected article place in the top k? Checked on
    the article id each `Doc` carries, not the chunk text: `_KB`'s ids
    ("refund window") never appear in their own articles."""
    results = search_kb_impl(query, k=k, kb=kb)
    return any(doc["id"].startswith(f"kb:{expected_article_id}#") for doc in results)

search_kb_impl takes a query string rather than messages, and runs it against the shared index or against the kb you pass. With a hand-labeled pair such as (“How long are refunds available after purchase?”, "refund window"), the check passes on the seeded KB; the companion’s tests run both seeded articles through it with a deterministic stand-in for the embedding model.

Run this against a small, hand-labeled set of (query, expected article) pairs - the retrieval equivalent of Chapter 21’s path-coverage discipline, one example per query shape the knowledge base actually needs to answer, not per plausible-sounding question. A regression here explains an Atlas failure Chapter 21’s own suite would otherwise misattribute to the model, the prompt, or the route - when the real cause never left the retrieval step.