Skip to content

agent-loop.md

Latest commit

 

History

History
246 lines (205 loc) · 9.42 KB

agent-loop.md

File metadata and controls

246 lines (205 loc) · 9.42 KB
Raw

rag-retrieval — Agent loop (Step 2.11 spike)

Spike artefact. AgentLoopV0 is the Step 2.11 validation-spike orchestrator. It is not a production agent runtime — Phase 3.x will replace its heuristic stop-condition controller with an LLM tool-loop driver in a dedicated rag-agent package. Until then, this module exists primarily to surface Phase 2 design gaps under realistic agent access patterns (see agent-loop-v0-gaps.md).

rag-retrieval 0.3 ships AgentLoopV0 alongside the existing HybridRetriever and RetrievalRouter, composed via a structural QueryUnderstander Protocol so the package stays free of a hard dependency on rag-query.

Overview

Symbol Purpose
AgentLoopV0 Async orchestrator: iterative retrieve with budget enforcement, sticky routing, plan-reuse measurement, shared embedding cache
AgentLoopConfig Per-loop knobs (max_iterations, diminishing-returns threshold, sticky-mode flag, per-iter cost estimates)
AgentLoopResult Frozen aggregate output (iterations, final chunks, stop reason, budget consumed, cache stats)
AgentLoopIterResult Frozen per-iteration record
QueryUnderstander Structural Protocol matching QueryUnderstandingPipeline.understand
UnderstoodQueryLike Structural Protocol over the fields the loop reads from UnderstoodQuery
RefineQuery Callable type for caller-supplied query refinement between iterations
ShouldContinue Callable type for caller-supplied stop callback
StopReason Enum: which condition exited the loop (re-exported from rag-core)
BudgetSpend Cumulative resource use across the run (re-exported from rag-core)
CacheStats Embedding-cache hit/miss accounting (re-exported from rag-core)
AgentLoopError Raised on non-recoverable orchestration failures

Usage

Minimal loop with the existing query pipeline

from rag_core.spi.noop import NoopEmbedder, NoopEmbeddingCache
from rag_query.pipeline import QueryUnderstandingPipeline
from rag_retrieval import (
    AgentLoopConfig,
    AgentLoopV0,
    HybridRetriever,
    RetrievalRouter,
)

# Production wiring: real backends + LLM-backed understander
hybrid = HybridRetriever(
    vector_backend=pg_vector,
    keyword_backend=es_keyword,
    graph_backend=neo4j,
)
router = RetrievalRouter(hybrid=hybrid, embedder=openai_embedder)
understander = QueryUnderstandingPipeline(
    rewriter=llm_rewriter,
    decomposer=llm_decomposer,
    hyde=hyde_generator,
    glossary=tenant_glossary,
)
embedding_cache = redis_embedding_cache  # any rag_core.spi.EmbeddingCache impl

loop = AgentLoopV0(
    understander=understander,
    router=router,
    embedder=openai_embedder,
    embedding_cache=embedding_cache,
    config=AgentLoopConfig(
        max_iterations=5,
        min_new_chunks_per_iter=2,
        allow_route_switch=False,   # sticky: iter 0 picks the route
    ),
)

result = await loop.run(
    ctx,                              # RequestContext — must carry budget
    query_text="explain reranking",
    top_k=10,
)

assert result.stop_reason.value in ("max_iterations", "diminishing_returns",
                                    "budget_tokens", "budget_iter", ...)

Caller-driven refinement

async def refiner(original: str, acc, iter_idx):
    last_chunks = acc[-1].chunks if acc else []
    if not last_chunks:
        return None                   # → StopReason.caller_stop
    return f"more details about {last_chunks[0].chunk_id}"

result = await loop.run(ctx, query_text="X", refine_query=refiner)

CLI smoke

The loop ships with a ragctl subcommand that exercises the full flow against the noop SPIs with no LLM credentials required:

uv run ragctl agent-loop "what is RAG?" \
    --max-iter 4 \
    --sub-query "how does BM25 work" \
    --sub-query "what is reciprocal rank fusion" \
    --min-new-chunks 0
query:        'what is RAG?'
stop_reason:  max_iterations
budget_used:  tokens=0 wall_ms=0 iters=4
cache_stats:  hits=1 misses=3 rate=0.25
final_chunks: 5
iterations:
  #0  q='what is RAG?'             shape=semantic   results=5   new=5  plan_reused=False sticky=False
  #1  q='how does BM25 work'       shape=semantic   results=5   new=0  plan_reused=True  sticky=True
  ...

Stop conditions

Checked in this order at the top of each iteration:

  1. Budget exhaustion — any of the four Budget dimensions (max_tokens, max_dollars, max_wall_ms, max_iter) reaching zero produces the matching StopReason.budget_*. ⚠️ 0 is treated as exhausted, not "uncapped" — None is the canonical "uncapped" sentinel. See G-03 in the gap-list for an open semantic question.

Checked at the end of each iteration:

  1. should_continue callback returning False → StopReason.caller_stop.
  2. Diminishing returns — new chunks per iter < min_new_chunks_per_iter → StopReason.diminishing_returns. Never fires on iter 0. Setting the threshold to 0 disables this check.
  3. refine_query returning None → StopReason.caller_stop.

Reaching AgentLoopConfig.max_iterations produces StopReason.max_iterations.

Telemetry

OTel span agent_loop.run

Attribute Type
rag.tenant_id str
rag.agent_loop.max_iterations int
rag.agent_loop.allow_route_switch bool
rag.agent_loop.iterations int
rag.agent_loop.stop_reason str (StopReason value)
rag.agent_loop.budget_tokens_used int
rag.agent_loop.budget_wall_ms_used int
rag.agent_loop.cache_hits_embedding int
rag.agent_loop.cache_misses_embedding int
rag.agent_loop.plan_reuse_rate float (0.0 – 1.0)
rag.agent_loop.sticky_route_rate float (0.0 – 1.0)
rag.agent_loop.final_chunks_n int
rag.agent_loop.elapsed_ms float

Structured-log event agent_loop.iteration_complete

Emitted once per iteration via rag_observability.get_logger:

{
  "event_kind": "agent_loop.iteration_complete",
  "iteration": 0,
  "query_chars": 23,
  "decision_shape": "semantic",
  "decision_reason": "shape=SEMANTIC",
  "results_n": 5,
  "new_chunks_n": 5,
  "plan_reused": false,
  "route_sticky": false,
  "cache_hits": 0,
  "cache_misses": 1,
  "elapsed_ms": 1
}

Internals

Cross-package isolation

The loop integrates with rag-query via the structural Protocol QueryUnderstander (and UnderstoodQueryLike for the understanding output). This keeps rag-retrieval's pyproject dependency list intact (rag-core only). Test code uses simple dataclass-shaped stubs; production wires QueryUnderstandingPipeline from rag-query.

Cache-key alignment

_stable_text_hash normalises text (lowercase + whitespace collapse) before hashing, matching the rag-packer dedup normalisation so embedding-cache keys stay aligned with downstream dedup decisions.

PolicyEngine boundary

AgentLoopV0 is not a PDP consumer. Every retrieval delegation goes through RetrievalRouter → HybridRetriever (the canonical read_chunk PDP call site). The tests/policy/coverage.py linter passes naturally because agent_loop.py makes no governed SPI calls itself.

Extension points

  • Custom understander — implement the two attributes / one method of UnderstoodQueryLike + QueryUnderstander. Useful for offline / golden-data harnesses.
  • Custom refiner — pass refine_query=... to drive iteration text from outside the loop. The refiner sees every prior AgentLoopIterResult, including chunks and decisions, so smart strategies (LLM-as-refiner, deterministic golden refiner) plug in without orchestrator changes.
  • Stop callback — pass should_continue=... for caller-driven early exit. Exceptions are wrapped as AgentLoopError(iter=…).
  • Cache impl — any EmbeddingCache works. Wrap with SemanticEmbeddingCache for Jaccard-similarity lookup above the exact-match layer (Step 2.11 G-01 fix; doubles overall hit rate on agent-style traffic). See caching.md for the broader three-cache split (Embedding / Retrieval / Answer).
  • Cost attribution (Step 2.11 G-04) — pass cost_estimator=DefaultCostEstimator() (or any CostEstimator-shaped object) on AgentLoopConfig to replace the flat per_iter_token_estimate heuristic with per-stage attribution (understand / embed / retrieve / rerank). Cache hits attribute no embed cost. Production callers override DefaultCostEstimator with provider-reported bills. Full per-SPI Cost-return shape is deferred to Phase 3.x — see the ADR-0008 backlink.

Related