Retrieval-Augmented Generation (RAG)

aimu.rag grounds a model’s answers in your own documents: chunk text, store it, retrieve the pieces relevant to a question, and fold them into the prompt. The primitives are plain functions over the MemoryStore interface – there is no retriever / splitter / loader class hierarchy.

What this notebook covers

Split text into chunks (split_text)
Chunk size, overlap, and token-aware splitting
Ingest documents into a store
Retrieve and augment a prompt (the full RAG flow)
Reranking (optional)
RAG as an agent tool
Choosing a store

Setup

The splitting, ingest, and retrieval cells run fully offline (SemanticMemoryStore uses ChromaDB’s built-in embedder by default; DocumentStore uses substring search). The cells that answer a question call a chat model – this notebook uses OpenAI (OPENAI_API_KEY), but any provider:model_id works. Reranking needs the [hf] extra (sentence-transformers).

A small document corpus

We’ll use a few short documents about AIMU as the knowledge base.

documents = [
    (
        "AIMU is a Python library for building AI-powered applications, with language models "
        "as the primary building block. It provides a unified, provider-agnostic interface "
        "across text, image, and audio generation."
    ),
    (
        "AIMU follows six design principles: plain Python, plain data, composability through "
        "uniform interfaces, progressive disclosure, direct paths for common tasks, and "
        "failures that are apparent. The small public surface is itself a feature."
    ),
    (
        "Embeddings map text to fixed-length vectors. AIMU exposes them via "
        "aimu.embedding_client() and aimu.embed(), with OpenAI, Ollama, and local "
        "HuggingFace providers. They back semantic memory and retrieval."
    ),
    (
        "Agentic patterns follow Anthropic's Building Effective Agents taxonomy: an Agent "
        "tool-calling loop, plus code-controlled workflows (Chain, Router, Parallel, "
        "EvaluatorOptimizer) that compose through one Runner interface."
    ),
]
print(len(documents), "documents")

1. Split text into chunks

split_text breaks a long string into overlapping chunks, preferring the largest separator that keeps each chunk under chunk_size (paragraphs → lines → sentences → words → characters).

from aimu.rag import split_text

long_text = "\n\n".join(documents)
chunks = split_text(long_text, chunk_size=200, chunk_overlap=40)
print(f"{len(chunks)} chunks (max len {max(len(c) for c in chunks)})\n")
for i, c in enumerate(chunks):
    print(f"[{i}] {c[:90]}...")

2. Chunk size, overlap, and token-aware splitting

chunk_overlap repeats the tail of one chunk at the head of the next so context isn’t lost at boundaries (it must be smaller than chunk_size). length_function defaults to len (characters); pass a token counter for token-aware chunking.

# Token-aware via a word counter (stand-in for a real tokenizer like tiktoken):
#   import tiktoken; enc = tiktoken.get_encoding("cl100k_base")
#   length_function=lambda s: len(enc.encode(s))
def word_count(s):
    return len(s.split())


token_chunks = split_text(long_text, chunk_size=40, chunk_overlap=8, length_function=word_count)
print("chunk word-counts:", [len(c.split()) for c in token_chunks])

# Custom separators, e.g. Markdown headings first:
md_chunks = split_text(
    "# A\ncontent a\n# B\ncontent b", chunk_size=20, chunk_overlap=0, separators=["\n# ", "\n", " ", ""]
)
print("markdown-ish chunks:", md_chunks)

3. Ingest documents into a store

ingest chunks each document and stores the chunks. Any MemoryStore works; here a SemanticMemoryStore (ChromaDB vector search, default embedder, runs offline). To pick the embedding model, pass embedding_client=aimu.embedding_client(...) (see notebook 11).

from aimu.memory import SemanticMemoryStore
from aimu.rag import ingest

store = SemanticMemoryStore(collection_name="rag_demo")
n_chunks = ingest(store, documents, chunk_size=200, chunk_overlap=40)
print("ingested", n_chunks, "chunks")

4. Retrieve and augment

retrieve pulls the chunks most relevant to a query; format_context joins them for the prompt. Then it’s an ordinary chat() call: RAG is a workflow you compose, not a new object type.

from aimu.rag import retrieve, format_context

question = "How do I turn text into vectors in AIMU?"
hits = retrieve(store, question, n_results=2)
print("retrieved:")
for h in hits:
    print("  -", h[:90], "...")

context = format_context(hits, numbered=True)
print("\ncontext passed to the model:\n", context[:300])

import aimu

# Augment the prompt with retrieved context and answer (needs a chat model).
answer = aimu.chat(
    f"Use only the context to answer.\n\nContext:\n{context}\n\nQuestion: {question}",
    model="openai:gpt-4o-mini",
)
print(answer)

5. Reranking (optional)

A vector search is fast but coarse. Fetch a wide candidate set, then rerank it with a cross-encoder to keep the best few. Requires the [hf] extra (sentence-transformers); the model loads once and is cached.

from aimu.rag import rerank

candidates = retrieve(store, question, n_results=4)
top = rerank(question, candidates, top_n=2)
for t in top:
    print("-", t[:90], "...")

6. As an agent tool

make_retrieval_tool(store) exposes retrieval as a retrieve_context(query) tool so an agent can fetch background on demand. Like the memory tools, the store is explicit (no env-var singleton).

from aimu.tools.builtin import make_retrieval_tool
from aimu.agents import Agent

client = aimu.client("openai:gpt-4o-mini")
agent = Agent(client, tools=[make_retrieval_tool(store, n_results=3)])
print(agent.run("Using the knowledge base, what are AIMU's design principles?"))

7. Choosing a store

Any MemoryStore works with ingest / retrieve. DocumentStore does case-insensitive substring search with no embeddings, handy for exact-term lookup or when you can’t run an embedding model.

from aimu.memory import DocumentStore

doc_store = DocumentStore()
ingest(doc_store, documents, chunk_size=200, chunk_overlap=0)
print(retrieve(doc_store, "workflows", n_results=1))

Scope

aimu.rag ships chunk / retrieve / rerank only. Document loaders (PDF, HTML, …) are intentionally out of scope: the built-in read_file / get_web_content tools, or any library that returns text, feed ingest directly. Chunks are stored as plain strings (the MemoryStore contract), so per-chunk metadata is not currently carried.

--- title: "Retrieval-Augmented Generation (RAG)" jupyter: python3 --- `aimu.rag` grounds a model's answers in your own documents: chunk text, store it, retrieve the pieces relevant to a question, and fold them into the prompt. The primitives are **plain functions over the [`MemoryStore`](12-memory.qmd) interface** -- there is no retriever / splitter / loader class hierarchy. ## What this notebook covers 1. Split text into chunks (`split_text`) 2. Chunk size, overlap, and token-aware splitting 3. Ingest documents into a store 4. Retrieve and augment a prompt (the full RAG flow) 5. Reranking (optional) 6. RAG as an agent tool 7. Choosing a store ## Setup The splitting, ingest, and retrieval cells run fully offline (`SemanticMemoryStore` uses ChromaDB's built-in embedder by default; `DocumentStore` uses substring search). The cells that *answer* a question call a chat model -- this notebook uses OpenAI (`OPENAI_API_KEY`), but any `provider:model_id` works. Reranking needs the `[hf]` extra (`sentence-transformers`). ## A small document corpus We'll use a few short documents about AIMU as the knowledge base. ```{python} documents = [ ( "AIMU is a Python library for building AI-powered applications, with language models " "as the primary building block. It provides a unified, provider-agnostic interface " "across text, image, and audio generation." ), ( "AIMU follows six design principles: plain Python, plain data, composability through " "uniform interfaces, progressive disclosure, direct paths for common tasks, and " "failures that are apparent. The small public surface is itself a feature." ), ( "Embeddings map text to fixed-length vectors. AIMU exposes them via " "aimu.embedding_client() and aimu.embed(), with OpenAI, Ollama, and local " "HuggingFace providers. They back semantic memory and retrieval." ), ( "Agentic patterns follow Anthropic's Building Effective Agents taxonomy: an Agent " "tool-calling loop, plus code-controlled workflows (Chain, Router, Parallel, " "EvaluatorOptimizer) that compose through one Runner interface." ), ] print(len(documents), "documents") ``` ## 1. Split text into chunks `split_text` breaks a long string into overlapping chunks, preferring the largest separator that keeps each chunk under `chunk_size` (paragraphs → lines → sentences → words → characters). ```{python} from aimu.rag import split_text long_text = "\n\n".join(documents) chunks = split_text(long_text, chunk_size=200, chunk_overlap=40) print(f"{len(chunks)} chunks (max len {max(len(c) for c in chunks)})\n") for i, c in enumerate(chunks): print(f"[{i}] {c[:90]}...") ``` ## 2. Chunk size, overlap, and token-aware splitting `chunk_overlap` repeats the tail of one chunk at the head of the next so context isn't lost at boundaries (it must be smaller than `chunk_size`). `length_function` defaults to `len` (characters); pass a token counter for token-aware chunking. ```{python} # Token-aware via a word counter (stand-in for a real tokenizer like tiktoken): # import tiktoken; enc = tiktoken.get_encoding("cl100k_base") # length_function=lambda s: len(enc.encode(s)) def word_count(s): return len(s.split()) token_chunks = split_text(long_text, chunk_size=40, chunk_overlap=8, length_function=word_count) print("chunk word-counts:", [len(c.split()) for c in token_chunks]) # Custom separators, e.g. Markdown headings first: md_chunks = split_text( "# A\ncontent a\n# B\ncontent b", chunk_size=20, chunk_overlap=0, separators=["\n# ", "\n", " ", ""] ) print("markdown-ish chunks:", md_chunks) ``` ## 3. Ingest documents into a store `ingest` chunks each document and stores the chunks. Any `MemoryStore` works; here a `SemanticMemoryStore` (ChromaDB vector search, default embedder, runs offline). To pick the embedding model, pass `embedding_client=aimu.embedding_client(...)` (see [notebook 11](11-embeddings.qmd)). ```{python} from aimu.memory import SemanticMemoryStore from aimu.rag import ingest store = SemanticMemoryStore(collection_name="rag_demo") n_chunks = ingest(store, documents, chunk_size=200, chunk_overlap=40) print("ingested", n_chunks, "chunks") ``` ## 4. Retrieve and augment `retrieve` pulls the chunks most relevant to a query; `format_context` joins them for the prompt. Then it's an ordinary `chat()` call: RAG is a workflow you compose, not a new object type. ```{python} from aimu.rag import retrieve, format_context question = "How do I turn text into vectors in AIMU?" hits = retrieve(store, question, n_results=2) print("retrieved:") for h in hits: print(" -", h[:90], "...") context = format_context(hits, numbered=True) print("\ncontext passed to the model:\n", context[:300]) ``` ```{python} import aimu # Augment the prompt with retrieved context and answer (needs a chat model). answer = aimu.chat( f"Use only the context to answer.\n\nContext:\n{context}\n\nQuestion: {question}", model="openai:gpt-4o-mini", ) print(answer) ``` ## 5. Reranking (optional) A vector search is fast but coarse. Fetch a wide candidate set, then rerank it with a cross-encoder to keep the best few. Requires the `[hf]` extra (`sentence-transformers`); the model loads once and is cached. ```{python} from aimu.rag import rerank candidates = retrieve(store, question, n_results=4) top = rerank(question, candidates, top_n=2) for t in top: print("-", t[:90], "...") ``` ## 6. As an agent tool `make_retrieval_tool(store)` exposes retrieval as a `retrieve_context(query)` tool so an agent can fetch background on demand. Like the memory tools, the store is explicit (no env-var singleton). ```{python} from aimu.tools.builtin import make_retrieval_tool from aimu.agents import Agent client = aimu.client("openai:gpt-4o-mini") agent = Agent(client, tools=[make_retrieval_tool(store, n_results=3)]) print(agent.run("Using the knowledge base, what are AIMU's design principles?")) ``` ## 7. Choosing a store Any `MemoryStore` works with `ingest` / `retrieve`. `DocumentStore` does case-insensitive substring search with no embeddings, handy for exact-term lookup or when you can't run an embedding model. ```{python} from aimu.memory import DocumentStore doc_store = DocumentStore() ingest(doc_store, documents, chunk_size=200, chunk_overlap=0) print(retrieve(doc_store, "workflows", n_results=1)) ``` ## Scope `aimu.rag` ships chunk / retrieve / rerank only. Document **loaders** (PDF, HTML, …) are intentionally out of scope: the built-in `read_file` / `get_web_content` tools, or any library that returns text, feed `ingest` directly. Chunks are stored as plain strings (the `MemoryStore` contract), so per-chunk metadata is not currently carried.