Skip to content

Retrieval

Retrieval allows you to provide context to an agent or to do general queries on a previously ingested knowledge base. After initliazation of your runtime to connect to the vectorized knowledge base, runtime.retrieve(...) will return RetrievalResult which contains a list of RetrievedChunks


import asyncio

from railtracks.retrieval import RetrievalRuntime


async def query(runtime: RetrievalRuntime):
    result = await runtime.retrieve(
        "What is the refund policy?",
        top_k=5,
    )

    print(f"query={result.query}")
    for hit in result.chunks:
        print(f"  [{hit.score:.3f}] {hit.chunk.content[:120]}")

RetrievalResult is {query, chunks} where each RetrievedChunk carries a score in [0, 1], a 0-indexed rank, and the original Chunk (content + metadata + document_id).

The runtime caches the embedder's reported model name on the first successful ingest. If a later retrieve uses a runtime whose embedder reports a different model, you get EmbeddingModelMismatchError signaling corruption of the vector spaces.


Metadata filters

Any scalar in Chunk.metadata is filterable. The runtime also writes two well-known keys automatically (source_path (from Document.source) and content_hash) so you can filter by source without setting metadata yourself:

async def filtered(runtime: RetrievalRuntime):
    # Filters are flat equality on Chunk.metadata.
    # The runtime also writes source_path and content_hash automatically.
    result = await runtime.retrieve(
        "flood coverage exclusions",
        top_k=5,
        metadata_filters={
            "jurisdiction": "TX",
            "doc_type": "regulation",
        },
    )
    return result

Filters are flat equality only

Currently we only support field==value for filtering. For other common filtering needs (or, is_in, etc) please open an issue or a PR at https://github.com/RailtownAI/railtracks/issues


Per-call scope

When the same runtime serves multiple users, pass StoreScope on each retrieve() (and ingest()) call to isolate their results:

from railtracks.retrieval.stores import StoreScope  # noqa: E402


async def scoped_query(runtime: RetrievalRuntime):
    # Scope is per-call. One runtime serves any number of tenants —
    # each call filters reads (and tags writes) by the scope you pass.
    alice_hits = await runtime.retrieve(
        "policy update",
        scope=StoreScope(labels={"user_id": "alice"}),
    )
    bob_hits = await runtime.retrieve(
        "policy update",
        scope=StoreScope(labels={"user_id": "bob"}),
    )
    return alice_hits, bob_hits

Scope is enforced at the store layer, not the runtime and even VectorStore.nearest_neighbors() (the lower-level bypass method) honors it. And unlike metadata filters, scope is stamped onto entries at write time, not just applied at read time.


Audit hook

on_retrieve=callback fires synchronously after every retrieve() call, with (query: str, result: RetrievalResult). Use it for query logs, hit-rate metrics, or feeding queries into an evaluation harness:

from railtracks.retrieval import RetrievalRuntime  # noqa: E402, F811
from railtracks.retrieval.chunking import RecursiveCharacterChunker  # noqa: E402
from railtracks.retrieval.embedding import OpenAIEmbedding  # noqa: E402
from railtracks.retrieval.stores import InMemoryVectorBackend, VectorStore  # noqa: E402


def log_retrieve(query: str, result):
    # Synchronous; runs after the retrieve call returns. Use for query-side
    # audit logging, hit-rate metrics, or feeding an evaluation harness.
    print(f"[retrieve] {query!r}{len(result.chunks)} hits")


def build_with_hook():
    return RetrievalRuntime(
        chunker=RecursiveCharacterChunker(chunk_size=800),
        embedder=OpenAIEmbedding(),
        store=VectorStore(InMemoryVectorBackend()),
        on_retrieve=log_retrieve,
    )

Like on_ingest, it runs on the request path. Push to a queue for any non-trivial work.


Wiring retrieval into an agent

runtime.retrieve() is the only primitive and there's no built-in "RAG mode" you have to opt into. Two patterns cover most uses:

As a tool the agent calls deliberately. The agent decides when it needs context and what to search for. Best when only some turns need retrieval (general chat with occasional doc lookups), or when you want the model to refine queries before each search:

import railtracks as rt  # noqa: E402


def build_docs_bot(runtime: RetrievalRuntime):
    # The LLM decides when to call this and what to search for.
    # Best when only some turns need retrieval, or you want the model to
    # refine the query before each search.
    @rt.function_node
    async def search_docs(query: str, top_k: int = 5) -> str:
        """Search the documentation. Returns the most relevant chunks, separated by ---."""
        result = await runtime.retrieve(query, top_k=top_k)
        return "\n\n---\n\n".join(hit.chunk.content for hit in result.chunks)

    return rt.agent_node(
        name="DocsBot",
        llm=rt.llm.OpenAILLM("gpt-4o"),
        system_message=(
            "Use search_docs to ground every factual answer. "
            "Cite retrieved chunks verbatim."
        ),
        tool_nodes=[search_docs],
    )

As a pre-invoke step that always injects context. Retrieval runs on every agent invocation, transparent to the model. Best when the LLM should ground answers in the corpus on every turn (docs bots, support assistants):

async def ask_with_context(agent_cls, runtime: RetrievalRuntime, question: str):
    # Run retrieval yourself, then call the agent with the augmented prompt.
    # Best when the LLM should ground answers in the corpus on every turn
    # without having to decide whether to search.
    result = await runtime.retrieve(question, top_k=5)
    context = "\n\n---\n\n".join(hit.chunk.content for hit in result.chunks)

    return await rt.call(
        agent_cls,
        user_input=f"Context:\n{context}\n\nQuestion: {question}",
    )

Neither pattern hides the retrieval call; you control the query, the top_k, the filters, and where the chunks land in the prompt. Swap embedders, chunkers, or stores by passing a different runtime to the agent.

RagConfig has been removed

Earlier versions exposed a RagConfig shortcut that wired the pre-invoke pattern in one line. It has been removed; the patterns above are the supported way to attach retrieval to an agent.


  • Ingestion: the write path: streaming events, re-ingest, multi-tenant writes, sanitization, token guards.
  • Components → Stores: the Store protocol, StoreQuery, and the low-level nearest_neighbors bypass.
  • Components → Design: embedding-model guard, staleness check, what's not in scope yet (boolean filters, hybrid search, built-in reranker).