Why isn't a single retrieval-plus-generate call enough for an AI agent?

Plain RAG handles one retrieval and one generation pass well, but breaks down for queries that need multiple steps — calling an internal API for one fact and a search tool for another, then reasoning over both. An agent loop lets the model decide, step by step, whether to retrieve more, call a tool, or answer, rather than forcing every query through exactly one retrieval-then-generate shape.

What stops an agent loop from running forever?

An explicit step budget — a maximum number of loop iterations set before the interaction starts. Without one, a model that keeps deciding 'I need one more piece of information' has no natural stopping point, and both latency and cost become unbounded. The budget is a deliberate trade-off: more steps can mean better answers on complex queries, at the cost of slower and more expensive responses.

Why is prompt injection a bigger concern in agent systems than in plain chatbots?

Because an agent doesn't just read retrieved text and summarize it — it can act on tool outputs and retrieved content by feeding them back into its own reasoning loop. If a retrieved document or a tool's response contains text designed to look like an instruction ('ignore previous instructions and...'), an agent that treats all input as equally trustworthy can be hijacked into taking unintended actions. Retrieved and tool content has to be treated as untrusted data, never as instructions, throughout the loop.

What's the detail that separates a senior answer from a mid-level one here?

Recognizing that failure containment has to happen at the level of an individual step, not the whole session — a single tool timeout or a bad retrieval shouldn't crash an entire multi-turn conversation. The failure needs to be fed back into the model's own context as information it can reason about (retry, try a different approach, answer with a caveat), rather than the orchestrator simply giving up.

RAG and AI Agent Platform Design

A chatbot that answers from one retrieval pass is a good demo. A system that can decide it needs to check an internal API, notice the API failed, try a different approach, and still return a grounded, safe answer — that's a genuinely different design problem, and it's fast becoming as standard an interview topic as any classic system design prompt. This post walks through that design the way an interview conversation actually unfolds.

Requirements and Goals of the System

Functional Requirements

Answer a user query by retrieving relevant context from a knowledge base and generating a grounded response
Support multi-step reasoning where the model can call tools (search, internal APIs, calculators) and use the results to continue
Maintain conversation memory across multiple turns in a session
Validate outputs — for groundedness and safety — before returning them to the user

Below the line (out of scope)

The LLM inference platform itself — this system is a client of it (see the previous post in this series)
The vector database's internals (indexing, ANN search) — likewise, a dependency this system calls, not something it implements
Document ingestion and chunking — assume a knowledge base of already-chunked, already-embedded documents exists

Non-Functional Requirements

Scale: ~10,000 agent sessions/second initiated, each involving multiple underlying LLM calls, not just one
Latency: an end-to-end multi-step interaction should complete in a few seconds, even though each individual step (an LLM call, a tool call) can itself take hundreds of milliseconds to seconds
Reliability: a single tool-call failure mid-chain must not crash the entire interaction
Groundedness and safety: outputs need to be checked against retrieved context and safety policy before reaching the user, not trusted blindly

Capacity Estimation and Constraints

LLM call volume: at 10,000 sessions/second with an average of ~3 LLM calls per session (an initial reasoning step, a tool-call decision, a final synthesis), that's roughly 30,000 LLM calls/second directed at the underlying inference platform — meaningfully higher than the session rate itself, which is the number that matters for capacity-planning the dependency, not the session count alone.

Retrieval load: if most sessions trigger at least one retrieval, that's on the order of 10,000 vector-search queries/second, each typically pulling a top-k of 10-20 chunks — a real but much more modest load on the vector store compared to the LLM-serving side.

Context budget per call: if each generation call packs a few thousand tokens of retrieved context alongside the conversation history, that context consumption directly multiplies the compute cost of every call on the inference platform — which is why session memory compaction (below) isn't just a UX nicety, it's a cost control.

System APIs

http

POST /v1/agent/sessions
{ "user_id", "initial_query" }
→ { "session_id", "response" (streamed), "steps": [...] }

POST /v1/agent/sessions/{id}/messages
{ "message" }
→ { "response" (streamed), "steps": [...] }

High-Level Design

The core requirement to design around: answer queries that may require multiple steps of retrieval and tool use, within a bounded latency and cost budget, with grounded and safe output.

Bad solution: stuff the user's query and as much of the knowledge base as fits into one giant prompt, and make a single LLM call to generate the final answer directly.

This doesn't scale — context windows are finite and expensive, most of a large knowledge base is irrelevant to any given query, and there's no way for the model to take an intermediate action (call an API, run a calculation) if the answer needs something the stuffed-in text alone can't provide.

Good solution: a standard retrieval-augmented generation (RAG) pipeline — embed the query, retrieve the top-k relevant chunks from a vector store, insert them into the prompt, and make one LLM call to generate a grounded answer.

This handles knowledge grounding far better and keeps context bounded to what's actually relevant. It breaks down, though, for queries that need more than one retrieval or an action beyond retrieval — "compare our Q3 revenue to the industry average" needs an internal API call for one fact and a search for the other, then reasoning over both, which a single retrieve-then-generate pass structurally can't do.

Great solution: an explicit agent loop. On each iteration, the model is given the current context (query, retrieved chunks so far, tool results so far) and decides whether to retrieve more, call a registered tool, or produce a final answer. Each decision that involves a tool goes through a tool registry with a defined schema — the model's requested call is validated against that schema before anything actually executes, rather than being run blindly. The loop continues until the model produces a final answer or a step budget (a hard cap on iterations) is exhausted, which bounds worst-case latency and cost. Before the final answer reaches the user, it passes through a guardrail/validation layer.

Database Design

AgentSession — one conversation, spanning multiple turns, holding accumulated memory/state
Step — one iteration within a session: a tool call or an LLM generation, with its inputs and outputs
Tool — a registered capability the agent can invoke, with a defined argument schema
RetrievedChunk — a piece of context pulled from the vector store for a given step
GuardrailResult — the pass/fail outcome (and reason) for a step's output

The modeling choice worth explaining if asked: Step is logged as its own durable record, not just an ephemeral in-memory value inside a loop — every tool call, every retrieval, and every intermediate generation is individually reconstructable after the fact. This is what makes debugging a bad agent answer tractable (which step went wrong) and what makes resuming a crashed session possible (below), instead of a session being an opaque black box that either fully succeeds or has to restart from scratch.

Detailed Component Design

How do we bound latency and cost across a multi-step chain?

The step budget is the primary lever — a hard cap on loop iterations set before the interaction starts, trading off answer quality on complex queries against speed and cost. Within a single step, if the model requests multiple independent pieces of information (two unrelated tool calls, say), running them concurrently rather than serially meaningfully cuts that step's latency, since the two calls don't depend on each other's results.

What happens when a tool call fails or times out mid-chain?

The failure needs to be caught by the orchestrator and fed back into the model's own context as information — "this tool call failed, here's why" — rather than crashing the session outright. This lets the model itself decide whether to retry, fall back to a different tool, or answer with an explicit caveat about the missing information. The same layer also validates the model's requested tool arguments against the tool's registered schema before execution, so a malformed or unsafe request never reaches the actual tool.

How do we know the final answer is actually grounded, not hallucinated?

A guardrail step checks the generated answer against the chunks that were actually retrieved — often via a lighter-weight consistency/entailment check, sometimes a smaller model call dedicated to just this check — and either passes the answer through, flags it, triggers a regeneration, or annotates it with citations pointing back to specific source chunks. The point isn't to trust the generation blindly just because retrieval happened somewhere upstream in the chain.

How does session memory work across a long conversation without blowing the context window?

Recent turns are kept verbatim, but older turns need to be compacted — summarized into a condensed representation — once the conversation grows past a length where full verbatim history would consume an unreasonable share of the context budget on every subsequent call. This compaction runs periodically rather than growing the raw history unbounded, since context consumption directly multiplies the cost of every call on the underlying inference platform.

Fault Tolerance and Replication

Failure	Impact	Response
Orchestrator crash mid-session	Session state could be lost	Session state (steps completed so far) is persisted durably after each step, not just held in memory — a new orchestrator instance resumes from the last completed step instead of restarting the whole session
LLM serving platform outage	This system cannot function at all — it's a hard dependency	Worth stating explicitly rather than pretending otherwise; the honest mitigation is the inference platform's own redundancy (see the previous post), not anything this layer can compensate for
Vector store outage	Retrieval-dependent steps degrade	The agent can fall back to a "no retrieval" mode for queries that don't strictly require grounding, with reduced answer quality flagged to the user, rather than failing the whole session
Tool provider outage	One tool becomes unavailable mid-chain	Handled the same way as any tool failure — surfaced to the model as context it can reason around, not a hard stop

Extended Requirements

Observability & Monitoring

Signal	Why it matters	Alert on
Per-step latency breakdown	Multi-step latency issues are invisible in an end-to-end number alone	Any single step type (retrieval, tool call, generation) regressing vs. its own baseline
Tool call success/failure rate, per tool	Flags a specific failing dependency fast	Sharp rise vs. baseline for one tool
Guardrail flag rate	A rising rate can mean the underlying model or retrieval quality has regressed	Sustained rise vs. baseline
Average steps per session	A cost proxy — more steps means more LLM calls means more spend	Sustained upward drift, which may indicate the model is looping unnecessarily

Security & Compliance

Tool execution is effectively the model requesting an action be taken on its behalf, which means every tool call should run with least privilege and validated arguments — never assume a model-generated request is automatically safe to execute as-is. The more distinctive risk here is prompt injection through retrieved or tool content: if a document in the knowledge base, or a tool's response, contains text crafted to look like an instruction, an agent that treats all input uniformly can be hijacked into taking unintended actions. The mitigation is architectural, not just a filter — retrieved and tool-returned content must be treated as untrusted data the model reasons about, never as instructions it reasons from.

Leveled Expectations

Mid-level: Designs a standard RAG pipeline (retrieve once, generate once), and can be prompted toward a multi-step agent loop once a query requiring multiple actions is raised.

Senior: Independently proposes the agent loop with an explicit step budget; addresses tool-call failure containment (surfacing failures into the model's context) without prompting.

Staff+: Raises prompt injection from retrieved/tool content as a first-class threat model unprompted, not a generic security afterthought; designs grounding/hallucination checks as an explicit guardrail step rather than trusting retrieval-implies-grounded; reasons about session memory compaction as a cost-control mechanism tied directly to the underlying inference platform's economics, not just a context-window technicality.

RAG / AI Agent Orchestration Platform Design