Retrieve and generate with citations

The last two nodes are the simplest. Retrieve packages the current section as a chunk of context. Generate sends the query plus the retrieved context to the LLM with strict instructions to cite every claim.

retriever.py
python
def _make_retrieve():
    def retrieve(state: RetrievalState) -> dict:
        node: TreeNode = state["current_node"]
        chunk = (
            f"=== {node.title} "
            f"(Pages {node.page_start}-{node.page_end}) ===\n"
            f"{node.content}"
        )
        return {"retrieved_content": [chunk]}
    return retrieve

Retrieve packages the current node's full content with its title and page range as a header. The reducer appends the chunk to retrieved_content so multiple calls could accumulate.

retriever.py
python
def _make_generate(client, model):
    def generate_answer(state: RetrievalState) -> dict:
        context = "\n\n---\n\n".join(state["retrieved_content"])

        prompt = f"""You are an expert on distributed systems and database engineering.
Answer the question using ONLY the retrieved document sections below.
Cite the section title and page range for every claim you make.
If the context is insufficient, say so clearly, do not guess.

Question: {state['query']}

Retrieved sections:
{context}

Answer:"""

        raw, elapsed = _call_llm(client, model, prompt, "answer", call_num)
        return {"final_answer": raw}
    return generate_answer

The generation prompt is the only place full section text appears. Two guardrails matter: ONLY use retrieved context, and cite every claim with a section title and page range.

retriever.py
python
def retrieve(query: str, tree: TreeNode, client, model: str | None = None):
    graph = _build_graph(client, model)
    result = graph.invoke({
        "query":             query,
        "current_node":      None,
        "tree":              tree,
        "path_taken":        [],
        "retrieved_content": [],
        "reasoning":         "",
        "confidence":        0.0,
        "should_descend":    True,
        "target_child_id":   None,
        "depth":             0,
        "final_answer":      None,
        "call_log":          [],
    })
    return {
        "answer":     result.get("final_answer"),
        "path":       result.get("path_taken", []),
        "reasoning":  result.get("reasoning", ""),
        "confidence": result.get("confidence", 0.0),
        "sources":    result.get("retrieved_content", []),
        "call_log":   result.get("call_log", []),
    }

The public API hands in the query and tree, invokes the compiled graph, and unpacks the final state into a clean response with the path, confidence, and sources for the caller.

Path_taken shows how the agent decided, not just where it ended up. Users and operators use it differently. A product UI might show the final sources as citations. A debugging dashboard shows the full path so you can see why the agent descended one way instead of another. Both views come from the same run.

Quiz: Quiz

Loading practice…

Checkpoint: Agent checkpoint

Loading practice…