RAG

How RAG works in the SDK: the provider lifecycle and how retrieval is exposed to agents as tools.

What is RAG?

Retrieval-Augmented Generation grounds the model in facts:

1. User asks a question
2. Agent retrieves relevant documents from a memory
3. Documents pass back to the model as retrieved context
4. Model responds, citing numbered sources

Enabling RAG on an Agent

Set the provider

Set rag_provider = RAGProvider() on the agent.

Return the documents

Override get_rag_queryset() (and/or get_rag_documents()) to select what gets indexed.

Return a pipeline

Implement get_rag_pipeline() (or one of the named variants) to return a RAG pipeline.

Link memories to threads

get_rag_tools() builds one retrieval tool per active memory.

from django_ai_sdk.memories.models import Entry
from django_ai_sdk.rags import BM25QueryExpanderRAG, BM25QueryExpanderRAGConfig
from django_ai_sdk.rags.provider import RAGProvider

class MyAgent(Agent):
    rag_provider = RAGProvider()

    async def get_rag_queryset(self, memory_id=None):
        if memory_id:
            return Entry.objects.filter(memory_id=memory_id)
        return Entry.objects.all()

    async def get_rag_pipeline(self, memory_id=None):
        documents = await self.get_rag_documents(memory_id)
        if not documents:
            return None
        return BM25QueryExpanderRAG(
            documents=documents,
            config=BM25QueryExpanderRAGConfig(top_k=5, n_expansions=4),
        )

Wire the retrieval tools into the ToolAgent inside get_pipeline_adapter():

async def get_pipeline_adapter(self, thread_id=None, user=None):
    generator = OpenAIChatGenerator(...)
    citation_registry = self.get_citation_registry()

    rag_tools = await self.get_rag_tools(
        thread_id=thread_id,
        citation_registry=citation_registry,
        citation_formatter=self.get_citation_formatter(),
        user=user,
    )
    tool_agent = ToolAgent(
        config=ToolAgentConfig(
            model=self.get_model(),
            system_prompt=self.get_system_prompt(),
            tools=[*await self.get_tools(thread_id=thread_id, user=user), *rag_tools],
        ),
        generator=generator,
    )
    return Stream(
        pipeline=tool_agent.pipeline(),
        generator=generator,
        storage_adapter=await self.get_storage_adapter(thread_id),
        citation_registry=citation_registry,
    )

RAGProvider

RAGProvider manages the RAG lifecycle: warmup (building indexes), an LRU cache of instances (_MAX_CACHE_SIZE = 100), and conversion of pipelines to tools.

Integration: Tool Calling

RAG is exposed as a per-memory retrieval tool the model decides to call: there is no automatic context injection. Each memory linked to the thread produces one tool named after its ToolSpec:

spec = await memory.get_tool_spec()  # -> ToolSpec(name="search_documents", ...)

Tool names are namespaced per memory so unrelated memories don’t collide; if two memories produce the same name, the later one is suffixed with a short memory-ID hash and a warning is logged. Memories with zero documents are skipped with a warning. Retrieval results stream back as numbered SourceEvents.

Citations

Retrieved sources are rendered to the model as numbered citations and streamed to the frontend:

• get_citation_registry(): a fresh per-turn CitationRegistry so indices reset between turns.
• get_citation_formatter() / citation_formatter_class: how sources render for the model. DefaultCitationFormatter numbers sources so the model can reference them inline.
• Stream emits a SourceEvent(index, title, content, source_id, ...) per source after a retrieval tool runs, in cumulative-index order.
• Protocol handlers translate these into the frontend’s citation format (Vercel source-document parts carry the citation index in providerMetadata).

Wire the same registry and formatter into both get_rag_tools() and the Stream so citations round-trip consistently.

Lifecycle

await MyAgent.warmup(agent, memory_id="mem-1")                 # build once
await MyAgent.reindex(agent, memory_id="mem-1")                # rebuild
MyAgent.clear_rag_cache(agent)                                 # drop cache

Best Practices

Next: Storage, where conversations persist.