Retrieval-Augmented
Generation

pip install chromadb anthropic

touch rag_agent.py

import os
import chromadb
import anthropic

client = anthropic.Anthropic(api_key=os.environ["ANTHROPIC_API_KEY"])

# Our corpus: each entry is one "document chunk"
CORPUS = [
    {
        "id": "chunk_01",
        "text": (
            "ReAct agents interleave reasoning and action in a loop. "
            "The model produces a Thought, selects an Action (tool call), "
            "receives an Observation, and repeats until it emits a final answer. "
            "This was introduced by Yao et al. (2022) in arXiv:2210.03629."
        ),
        "source": "section-05-agents",
    },
    {
        "id": "chunk_02",
        "text": (
            "Memory in agents is classified into four types: in-context (the active prompt window), "
            "external semantic (vector stores), episodic (structured logs of past sessions), "
            "and procedural (persistent rules embedded in the system prompt). "
            "This taxonomy comes from Lilian Weng's 2023 survey of LLM-powered agents."
        ),
        "source": "section-06-memory",
    },
    {
        "id": "chunk_03",
        "text": (
            "Plan-and-execute agents separate planning from execution. A planner LLM generates "
            "a full step list upfront. An executor LLM works through each step using tools. "
            "When a step fails, the planner is called again with a context summary to replan. "
            "This pattern is documented in Wang et al. (2023), arXiv:2305.04091."
        ),
        "source": "section-07-planning",
    },
    {
        "id": "chunk_04",
        "text": (
            "Prompt injection is the top LLM-specific vulnerability in the OWASP LLM Top 10. "
            "An attacker embeds instructions in user-controlled content (e.g., a webpage or file) "
            "that override the agent's system prompt and redirect its behavior. "
            "Defences include input delimiting, privilege separation, and instruction anchoring."
        ),
        "source": "section-08-prompting",
    },
    {
        "id": "chunk_05",
        "text": (
            "RLHF has three stages: SFT (supervised fine-tuning on demonstration data), "
            "reward model training (on human preference rankings of response pairs), "
            "and PPO fine-tuning (updating the policy to maximize the reward model's score). "
            "A KL-divergence penalty prevents the model from drifting too far from the SFT baseline."
        ),
        "source": "section-09-rl",
    },
    {
        "id": "chunk_06",
        "text": (
            "Constitutional AI (Bai et al., 2022) replaces the human preference labeling step in RLHF "
            "with AI-generated feedback guided by a written set of principles. "
            "The AI critiques and revises its own responses according to the constitution, "
            "and AI-generated preference rankings replace human annotators in the RLAIF stage."
        ),
        "source": "section-09-rl",
    },
    {
        "id": "chunk_07",
        "text": (
            "DPO (Direct Preference Optimization, Rafailov et al. 2023) eliminates the separate "
            "reward model and PPO training loop from RLHF. It reparametrizes the reward in terms "
            "of the policy itself, training directly on preference pairs via binary cross-entropy. "
            "DPO is now the dominant alignment fine-tuning method for open-weight models."
        ),
        "source": "section-09-rl",
    },
    {
        "id": "chunk_08",
        "text": (
            "Transformers use self-attention to process all tokens in a sequence in parallel. "
            "Every token attends to every other token simultaneously, computing a weighted sum "
            "of value vectors based on query-key dot products. This replaced recurrent (RNN) processing "
            "and is the architectural foundation of all modern large language models."
        ),
        "source": "section-04-llms",
    },
]

def build_index() -> chromadb.Collection:
    """Build an in-memory ChromaDB collection from the corpus."""
    db = chromadb.Client()  # ephemeral in-memory client
    collection = db.create_collection(
        name="course_knowledge",
        metadata={"hnsw:space": "cosine"}  # cosine similarity
    )

    collection.add(
        ids=[chunk["id"] for chunk in CORPUS],
        documents=[chunk["text"] for chunk in CORPUS],
        metadatas=[{"source": chunk["source"]} for chunk in CORPUS],
    )

    print(f"Index built: {collection.count()} chunks indexed.")
    return collection

TOP_K = 3

RAG_SYSTEM = """You are a precise question-answering assistant.
Answer the user's question using ONLY the information in the provided context blocks.
Each context block is labeled with a source ID.
If the answer is not in the context, say: "I don't have that information in the provided context."
Do not use outside knowledge. Cite the source ID when you use it."""


def retrieve(collection: chromadb.Collection, query: str, top_k: int = TOP_K) -> list[dict]:
    """Retrieve the top-k most relevant chunks for a query."""
    results = collection.query(
        query_texts=[query],
        n_results=top_k,
        include=["documents", "metadatas", "distances"],
    )
    chunks = []
    for i in range(len(results["ids"][0])):
        chunks.append({
            "id":       results["ids"][0][i],
            "text":     results["documents"][0][i],
            "source":   results["metadatas"][0][i]["source"],
            "distance": results["distances"][0][i],
        })
    return chunks


def rag_answer(collection: chromadb.Collection, query: str) -> str:
    """Full RAG pipeline: retrieve → augment → generate."""
    # 1. Retrieve
    chunks = retrieve(collection, query)

    print(f"\n  Retrieved {len(chunks)} chunks:")
    for c in chunks:
        print(f"    [{c['id']}] distance={c['distance']:.3f} source={c['source']}")
        print(f"    {c['text'][:80]}...")

    # 2. Augment — assemble context block
    context_block = "\n\n".join(
        f"[Source: {c['id']} | {c['source']}]\n{c['text']}"
        for c in chunks
    )
    augmented_query = f"Context:\n{context_block}\n\nQuestion: {query}"

    # 3. Generate
    response = client.messages.create(
        model="claude-haiku-4-5-20251001",
        max_tokens=512,
        system=RAG_SYSTEM,
        messages=[{"role": "user", "content": augmented_query}],
    )
    return response.content[0].text

TEST_QUERIES = [
    "How does a ReAct agent loop work?",
    "What are the four types of memory in agents?",
    "How does DPO differ from RLHF?",
    "What is prompt injection and how is it defended against?",
    "What year was the iPhone 15 released?",  # out-of-corpus query
]


if __name__ == "__main__":
    collection = build_index()

    for query in TEST_QUERIES:
        print(f"\n{'='*60}")
        print(f"QUERY: {query}")
        answer = rag_answer(collection, query)
        print(f"\nANSWER:\n{answer}")

python rag_agent.py

Index built: 8 chunks indexed.

============================================================
QUERY: How does a ReAct agent loop work?

  Retrieved 3 chunks:
    [chunk_01] distance=0.081 source=section-05-agents
    The ReAct agent interleaves reasoning and action in a loop...
    [chunk_03] distance=0.412 source=section-07-planning
    Plan-and-execute agents separate planning from execution...
    [chunk_02] distance=0.451 source=section-06-memory
    Memory in agents is classified into four types...

ANSWER:
According to [chunk_01 | section-05-agents], a ReAct agent loop works
by interleaving reasoning and action: the model produces a Thought,
selects an Action (a tool call), receives an Observation, and repeats
this cycle until it emits a final answer. This approach was introduced
by Yao et al. (2022).

============================================================
QUERY: What year was the iPhone 15 released?

  Retrieved 3 chunks:
    [chunk_04] distance=0.682 ...
    [chunk_01] distance=0.711 ...
    [chunk_07] distance=0.731 ...

ANSWER:
I don't have that information in the provided context.

RERANKER_SYSTEM = """You are a relevance ranker.
Given a query and a list of numbered passages, return ONLY a JSON array
of passage numbers in order of relevance to the query, most relevant first.
Example: [2, 5, 1] means passage 2 is most relevant, then 5, then 1.
Return ONLY the JSON array, no other text."""


def llm_rerank(query: str, chunks: list[dict], keep: int = 3) -> list[dict]:
    """Use the LLM to rerank a candidate chunk list."""
    passages = "\n\n".join(
        f"[{i+1}] {c['text'][:200]}" for i, c in enumerate(chunks)
    )
    prompt = f"Query: {query}\n\nPassages:\n{passages}"

    response = client.messages.create(
        model="claude-haiku-4-5-20251001",
        max_tokens=64,
        system=RERANKER_SYSTEM,
        messages=[{"role": "user", "content": prompt}],
    )
    import json
    ranked_indices = json.loads(response.content[0].text.strip())
    # Convert 1-indexed to 0-indexed, cap at available chunks
    reranked = [chunks[i - 1] for i in ranked_indices if 1 <= i <= len(chunks)]
    return reranked[:keep]


def rag_answer_reranked(collection: chromadb.Collection, query: str) -> str:
    """RAG with LLM reranking: retrieve 6 → rerank → take top 3 → generate."""
    candidates = retrieve(collection, query, top_k=6)
    top_chunks = llm_rerank(query, candidates, keep=3)

    print(f"  After reranking, using chunks: {[c['id'] for c in top_chunks]}")

    context_block = "\n\n".join(
        f"[Source: {c['id']} | {c['source']}]\n{c['text']}"
        for c in top_chunks
    )
    augmented_query = f"Context:\n{context_block}\n\nQuestion: {query}"

    response = client.messages.create(
        model="claude-haiku-4-5-20251001",
        max_tokens=512,
        system=RAG_SYSTEM,
        messages=[{"role": "user", "content": augmented_query}],
    )
    return response.content[0].text