Skip to content

agent.md

Latest commit

 

History

History
319 lines (254 loc) · 12.9 KB

agent.md

File metadata and controls

319 lines (254 loc) · 12.9 KB
Raw

rag-agent + agent surface reference (Step 3.6)

Overview

rag-agent is the production agent runtime: an explicit state-machine loop that iterates plan → act → observe → finalize, drives tools (retrieval today, more later), enforces four budget dimensions, and checkpoints every step so a run is resumable. It is the graduation of the Step 2.11 AgentLoopV0 spike (a fixed for loop) into a controller-driven, budgeted, resumable runtime.

The package depends only on rag-core — it is pure orchestration. The two concerns that touch the outside world (retrieval and final-answer egress governance) are the gateway's job and are wired at the boundary, so the loop itself makes no governed SPI calls.

The runtime is exposed over two gateway surfaces that serialise the same event stream frame-for-frame:

Surface Entry point Streaming
REST POST /v1/agent Server-Sent Events (text/event-stream)
gRPC rag.gateway.v1.RagService/Converse server-streaming
CLI ragctl agent "<goal>" consumes the SSE stream in-process

Wire types (AgentRequest, AgentEvent, AgentRunResult, …) live in rag_core.agent_types so SDKs, both surfaces, and the eval harness share one definition without a dependency cycle — the same precedent as rag_core.gateway_types (Step 3.1) and rag_core.openai_types (Step 3.4).

Usage

REST — POST /v1/agent

Post an AgentRequest; receive a stream of AgentEvent frames, each one SSE data: line of JSON, closed by the canonical data: [DONE] sentinel.

curl -N http://127.0.0.1:8000/v1/agent \
  -H 'content-type: application/json' \
  -H 'x-tenant-id: acme' -H 'x-principal-id: alice' \
  -d '{
    "tenant_id": "acme",
    "principal_id": "alice",
    "goal": "What changed in the Q3 security policy?",
    "top_k": 8,
    "max_steps": 6,
    "budget_tokens": 4000
  }'
data: {"kind":"run_started","run_id":"…", …}
data: {"kind":"phase_changed","phase":"planning","step_index":0, …}
data: {"kind":"tool_started","tool_call":{"tool":"retrieve", …}, …}
data: {"kind":"tool_completed","tool_result":{"ok":true, …}, …}
data: {"kind":"checkpoint_saved","step_index":1, …}
data: {"kind":"phase_changed","phase":"finalizing", …}
data: {"kind":"answer_delta","text":"Based on 8 retrieved passage(s), …"}
data: {"kind":"run_completed","result":{…}}
data: [DONE]

gRPC — Converse

import grpc
from rag_gateway._grpc_gen import rag_pb2, rag_pb2_grpc

async with grpc.aio.insecure_channel("127.0.0.1:50051") as channel:
    stub = rag_pb2_grpc.RagServiceStub(channel)
    req = rag_pb2.ConverseRequest(
        tenant_id="acme",
        principal_id="alice",
        goal="What changed in the Q3 security policy?",
        top_k=8,
        max_steps=6,
        budget_tokens=4000,
    )
    async for event in stub.Converse(req):
        kind = rag_pb2.AgentEventKind.Name(event.kind)
        if kind == "AGENT_EVENT_KIND_ANSWER_DELTA":
            print(event.text)
        elif kind == "AGENT_EVENT_KIND_RUN_COMPLETED":
            print(event.result.status, event.result.final_answer)

CLI — ragctl agent

uv run ragctl agent "What is RAG?" --tenant acme --top-k 8 --budget-tokens 4000

Drives POST /v1/agent against an in-process gateway (build_app() with the noop wiring) over FastAPI's TestClient, printing one line per event frame plus the terminal result. Point the same client at a running gateway to drive a real loop.

Wire types (rag_core.agent_types)

AgentRequest

The POST /v1/agent body / ConverseRequest payload.

Field Type Default Notes
tenant_id TenantId — required
principal_id PrincipalId — required
goal str — the task/question (min length 1)
corpus_ids list[CorpusId] [] restrict retrieval to these corpora
top_k int 8 1–100
max_steps int 6 1–50; config-level hard ceiling (defence-in-depth)
budget_tokens int | None None None = uncapped on that dimension
budget_dollars float | None None
budget_wall_ms int | None None
budget_iter int | None None operator-set per-request iteration cap
request_id RequestId | None None round-trips into the result
resume_run_id str | None None resume a checkpointed run instead of starting fresh

max_steps (config ceiling) and budget_iter (operator cap) are distinct on purpose: max_steps is always present as a safety net, budget_iter is an explicit per-request limit. Whichever bites first sets the StopReason.

AgentEvent — one stream frame

A single flat model (not a discriminated union) keeps the SSE JSON and the gRPC mapping mechanical: kind discriminates, the optional fields carry whatever that kind needs.

Field Populated on
kind every frame (AgentEventKind)
run_id every frame
step_index phase_changed, tool_started, tool_completed, checkpoint_saved
phase phase_changed
tool_call tool_started
tool_result tool_completed
text answer_delta
result run_completed only (the full AgentRunResult)
error_code, message run_failed only
at every frame (UTC timestamp)

AgentEventKind: run_started, phase_changed, tool_started, tool_completed, answer_delta, checkpoint_saved, run_completed, run_failed.

AgentRunResult — the run_completed payload

run_id, request_id, goal, status (StopReason), phase (AgentPhase), final_answer, steps (list[AgentStep]), citations, chunk_refs, budget_consumed (BudgetSpend), elapsed_ms.

AgentStep, AgentToolCall, AgentToolResult

One AgentStep per planning → acting → observing cycle: index, phase, thought, optional tool_call / tool_result, elapsed_ms. An AgentToolResult with ok=False is a recoverable tool failure — the loop folds the error into the transcript so the controller can react; a non-recoverable fault raises AgentLoopError instead.

StopReason (terminal status)

final_answer, budget_tokens, budget_dollars, budget_wall, budget_iter, max_iterations, diminishing_returns, caller_stop.

AgentPhase

planning → acting → observing (tool step) or planning → finalizing (finish step); terminal done / failed.

Runtime API (rag_agent)

from rag_agent import (
    AgentLoop, AgentConfig,
    Controller, ScriptedController, HeuristicController, LLMController,
    Tool, ToolRegistry, ToolSpec, RetrieveTool, Retriever,
    CheckpointStore, InMemoryCheckpointStore,
    AgentSnapshot,
)

AgentLoop

Pure orchestration with three injected collaborators:

loop = AgentLoop(
    controller=HeuristicController(),
    tools=ToolRegistry([RetrieveTool(retriever)]),
    checkpoint_store=InMemoryCheckpointStore(),
    config=AgentConfig(),
)

async for event in loop.stream(ctx, request):   # one AgentEvent per transition
    ...
result = await loop.run(ctx, request)            # drive to completion
result = await loop.resume(ctx, run_id)          # reload a checkpoint + continue

run / resume re-raise AgentLoopError on a failed run (the stream emits a run_failed frame first). A resume of an unknown run_id raises AgentLoopError inside stream before the first frame — the gateway catches it and surfaces an in-band run_failed frame rather than aborting the transport.

AgentConfig

Field Default Purpose
max_steps 6 hard step ceiling
per_step_token_estimate 256 flat token delta charged per step when no measured cost is returned (keeps a token-capped run terminating against cost-free noop backends)
per_step_dollar_estimate 0.0 flat dollar delta per step
checkpoint_each_step True persist a snapshot after every step

Controller

A Protocol — decide(ctx, *, snapshot, tools) -> AgentAction returning CallTool(tool, args, thought) or Finish(answer, thought). Three impls:

  • ScriptedController — replays a fixed action queue; the backbone of the unit suite (deterministic, no LLM). Finalises once the script is exhausted.
  • HeuristicController — credential-free default: retrieve once for the goal, then finalise from what came back. Makes the surface curl-able with zero credentials (CI, smoke tests, the noop gateway profile). Not for answer quality.
  • LLMController — production: renders snapshot + tool specs into a prompt, asks the wired LLM for a JSON action, parses it (tolerant of Markdown code fences). A parse failure degrades to Finish(raw_text) so the loop always produces an answer.

Tool / ToolRegistry / RetrieveTool / Retriever

A Tool is a named, JSON-schema-described capability (spec + invoke). Tools are the only way the loop touches the outside world. The built-in RetrieveTool (name="retrieve", default_top_k=8, max_top_k=50) wraps a structural Retriever Protocol:

async def retrieve(
    self, ctx, *, query: str, top_k: int, corpus_ids: Sequence[CorpusId] | None,
) -> list[ChunkRef]: ...

The gateway adapts its policed CorpusRouter / RetrievalRouter to this shape (_GatewayRetriever), so read_chunk governance inside HybridRetriever applies unchanged. A RetrievalError from the backend is recoverable — RetrieveTool returns ok=False so the controller can retry or finalise rather than abort the whole run.

CheckpointStore / AgentSnapshot

AgentSnapshot is the full serialisable state of a run at a step boundary — the single source of truth, so resume is just "load, then keep stepping". A CheckpointStore is save / load / delete, last-write-wins per run_id (a snapshot already encodes the full state). InMemoryCheckpointStore is the process-local default, keyed by (tenant_id, run_id) so run state can't leak across tenants; a Postgres/Redis-backed store is a drop-in swap.

Governance boundary

The loop performs no policy enforcement itself:

  • Retrieval governance happens inside the wired Retriever — the gateway adapts a policed router whose HybridRetriever is the read_chunk PDP call site, identical to /v1/query.
  • Final-answer egress governance is the gateway's job at the boundary: both surfaces consult PolicyEngine with egress_text over the final answer once (at the answer_delta frame), caching the decision. A deny withholds the text (the denial notice goes on the wire instead); a transform substitutes a redacted answer. Because the check runs before the answer_delta frame is emitted, denied text never reaches the wire, and the cached decision is replayed onto the terminal run_completed frame for a consistent answer + citation view.

LLMController calls llm.complete, so controller.py is allowlisted in tests/policy/coverage.py (same posture as rag_query.rewriter and rag_graphrag.summarizer). The loop names no governance-relevant SPI method, so the coverage linter does not flag it and no allowlist entry is needed for loop.py.

Observability

Spans:

  • agent.run — the full loop (REST + direct callers). Attributes: rag.tenant_id, rag.agent.run_id, rag.agent.max_steps, rag.agent.steps, rag.agent.stop_reason, rag.agent.final_chunks_n, rag.agent.budget_tokens_used.
  • grpc.converse — the gRPC Converse flow (wraps agent.run).

Structured-log events via rag-observability:

  • agent.step_complete — one per completed tool step.
  • agent.run_failed — non-recoverable orchestration fault.
  • gateway.agent_complete / grpc.converse_complete — one per successful run on the respective surface.
  • gateway.agent.policy_denied — egress_text deny over the final answer.

Extension points

  • Controller — implement the Controller Protocol for a custom planner; production swaps HeuristicController for LLMController sourced from rag.yaml.
  • Tools — subclass Tool and register it; the controller can call any tool in the registry by name (web fetch, calculators, sub-agents are future steps).
  • Checkpoints — implement CheckpointStore (Postgres/Redis) for durable, cross-restart resume; the loop is unchanged.
  • Retriever — supply any object satisfying the Retriever Protocol via RetrieveTool; the gateway's _GatewayRetriever is one such adapter.

See also