Agentic RAG Knowledge Assistant

Before any real files, this is the shape of the whole system in ten lines — a map you'll recognise in every file that follows. Don't run it; read it as pseudocode that names the four moving parts: plan, retrieve, judge_sufficiency, and reformulate.

1. def answer(question, max_steps=4): — the whole assistant is one function. max_steps=4 is the leash: it can search at most four times, then it must stop. That cap is a safety feature, not a detail.
2. subqs = plan(question) breaks the question into sub-questions; context = [] is the empty notebook where retrieved facts pile up.
3. for step in range(max_steps): repeats the search cycle up to four times. Each pass: retrieve(...) fetches hits, context += hits adds them to the notebook, and judge_sufficiency(...) asks 'is this enough?'.
4. if grade.sufficient: return generate_answer(...) — the moment the evidence is good enough, it writes a cited answer and leaves the loop. Otherwise reformulate(question, grade.gap) rewrites the query to chase the missing piece and the loop goes round again.
5. If all four tries pass without ever being sufficient, the last line return abstain(...) runs — an honest 'I don't have enough to answer'.

What the output means: Nothing prints — this is a skeleton to read. The real, runnable version is loop.py in Step 5, which follows this exact shape.

Try this: Find the four verbs — plan, retrieve, judge, reformulate — in the picture above. This code is that diagram written out. Everything after here just fills in each verb with real Python.

This file builds the three workers the loop bosses around: retrieve (find documents), judge (the grader — decide if we have enough), and generate (write the final answer). Each has a fast mock version for learning offline and a real version that calls Claude — switched by one environment variable, so the loop's logic never changes.

1. CORPUS is a tiny fake library: two documents, a "cur"rent refund policy and an "old" 2024 one. Answering "compare them" needs both — that's the multi-hop case the whole project is built to handle.
2. @dataclass class Grade: is a small record with two fields: sufficient (a yes/no) and gap (what's still missing). The grader hands one of these back so the loop knows whether to stop or search again.
3. retrieve(query) is deliberately naive: it only returns the current policy if the query mentions "current"/"refund", and only the 2024 policy if it mentions "2024"/"old"/"previous". So one query naturally finds only half a compare question — forcing the loop to go back for the rest.
4. judge(question, context) is the grader. Its mock rule: for a "compare" question it needs both the current and the 2024 fact; if only one is present it returns Grade(False, gap=...) naming what's missing. Any other question is satisfied by any context.
5. generate(...) writes the answer (mock: just joins the facts and tags them [grounded]). The _real_judge/_real_generate helpers do the same jobs by calling client.messages.create(...) — the grader asks Claude for JSON {sufficient, gap}, and only runs when USE_REAL_API is set to "1".

What the output means: Nothing runs yet — this file only defines the three pieces. The loop in Step 5 imports and orchestrates them.

Try this: Read retrieve and predict: for the query "what is the current refund policy?", which of the two docs comes back? (Just the current one — which is exactly why a single lookup can't answer a compare question.)

This is the real, runnable version of the shape snippet — the agentic loop itself. It retrieves, grades, and reformulates from the gap, stopping either when the answer is sufficient or when it hits the step cap and abstains. Notice it takes the three pieces as arguments with defaults, so tests can swap in scripted fakes.

1. MAX_STEPS = 4 and ABSTAIN = "..." are the leash and the honest fallback message, defined once at the top where they're easy to find and change.
2. def answer(question, retrieve=_retrieve, judge=_judge, generate=_generate): — the pieces default to the real ones from pieces.py, but a caller (like a test) can pass its own. This is dependency injection: same loop, swappable parts.
3. subquery = question starts the search with the raw question; context collects facts and trace records what happened each step (for debugging and observability).
4. Inside for step in range(MAX_STEPS): — retrieve hits, add only new ones to context (if h not in context avoids duplicates), grade, and append a trace row. If grade.sufficient, return the generated answer plus steps/trace. Otherwise rebuild subquery as f"{question} (need: {grade.gap})" — the gap steers the next search.
5. If the for loop finishes without returning, all tries failed: the final return gives back the ABSTAIN message with "abstained": True. The if __name__ == "__main__" block runs a demo compare question and prints the trace.

What the output means: Run directly, it prints one trace line per step and then the final answer — you'll see this in Step 6.

Try this: Change MAX_STEPS to 1 and re-run the demo. The compare question now abstains, because one search can't gather both policies — proof that the cap really does bound the work.

This is what python loop.py prints for the compare question — the loop reasoning out loud. Each dict is one pass through the search cycle, so you can literally watch it notice a gap and fix it.

1. Step 0: subquery is the raw question; 'hits': 1 means retrieve found only one document (the current policy); 'sufficient': False — the grader correctly says 'this compare question isn't fully answered yet'.
2. Step 1: the subquery now ends with (need: the 2024 policy) — the loop reformulated using the gap the grader reported. That extra phrase makes retrieve find the second document, and now 'sufficient': True.
3. ANSWER line: both facts are joined together and tagged [grounded] — the complete answer that needed two lookups to assemble.
4. steps: 2 — it used two of its four allowed tries. Efficient, and well under the leash.

What the output means: The two trace lines show the gap being detected then filled; the final two lines are the assembled grounded answer and the step count (2).

Try this: This is the payoff of the whole project: a plain RAG system would have stopped at Step 0 and answered with only half the story. Re-read Step 0's 'sufficient': False — that single 'no' is what saved the answer.

These six tests prove the loop's control flow without ever calling the API. The trick: each test passes its own tiny retrieve/judge/generate as arguments (using lambda, a one-line throwaway function), so the loop runs against scripted behaviour that's fast, free, and predictable.

1. test_answers_when_first_retrieval_suffices — a judge that always says Grade(True) means the loop should finish in exactly 1 step and not abstain. This proves it doesn't waste loops when the answer is already there.
2. test_reformulates_then_answers — a judge that returns False the first time and True the second (counting with calls["n"]) proves the gap drives a second retrieval: steps == 2.
3. test_stops_at_max_steps and test_abstains_when_never_sufficient use a judge that always says False. The loop must stop at MAX_STEPS and return the ABSTAIN message — this is the safety leash under test.
4. test_trace_records_every_step checks the trace has one entry per step (observability). test_real_multihop_uses_two_steps uses the real built-in mocks (no scripting) to confirm the genuine compare flow takes 2 steps.

What the output means: Run with pytest, all six should report PASSED and end with 6 passed — see the next lab.

Try this: A lambda q: ["fact"] is just a mini-function: 'given a query q, return the list ["fact"]'. Try adding a seventh test where the judge says True on the third call, and assert steps == 3.

This flips the same loop from the offline mock to the real Claude model — without changing a line of loop.py. Two environment variables do it, which is exactly why the pieces were built with a mock/real switch earlier.

1. export ANTHROPIC_API_KEY="sk-ant-..." puts your secret key in the environment so the anthropic client can authenticate. It never appears in your code.
2. export USE_REAL_API=1 is the switch the pieces check: with it set, judge and generate call their _real_* versions (Claude) instead of the mocks.
3. python loop.py then runs the identical loop — same retrieve/grade/reformulate logic — but now the grading and answering are done by the real model.

What the output means: The loop behaves the same, but judge() and generate() now hit the Anthropic API. Costs real tokens; the loop structure is unchanged.

Try this: Run unset USE_REAL_API to go back to the free, offline mock. Being able to flip between mock and real with one variable — and keeping tests always on the mock — is a pattern you'll reuse in every LLM project.