"""Recursive agent — reusable agentic step with compaction and recursion. A ``RecursiveAgent`` wraps a PydanticAI agent with: - Two-tier compaction (microcompaction + full summarization) - Depth-based recursion via a ``recurse`` tool - Budget management (token + request limits) - Sync/async execution Both the RR (Recursive Reflector) and the agentic SkillManager build on this. Callers provide their own tools, output type, and prompts. Usage:: from ace.core.recursive_agent import RecursiveAgent, AgenticConfig agent = RecursiveAgent( model="gpt-4o-mini", output_type=MyOutput, system_prompt="You are a ...", config=AgenticConfig(max_requests=20), tools=[my_tool_registrar], # list of (agent) -> None functions tool_names_to_compact=("my_tool",), ) output, metadata = agent.run(prompt="Analyze this", deps=my_deps) """ from __future__ import annotations import asyncio import concurrent.futures import copy import logging from contextlib import nullcontext from dataclasses import dataclass from typing import Any, Awaitable, Callable, Sequence, Type from pydantic_ai import Agent as PydanticAgent try: import logfire _logfire: Any = logfire except ImportError: _logfire = None def _rr_span(name: str, **attrs: Any): """Open a logfire span if logfire is installed, else a no-op context.""" if _logfire is not None: return _logfire.span(name, **attrs) return nullcontext() from pydantic_ai.exceptions import UsageLimitExceeded from pydantic_ai.messages import ( ModelRequest, ModelResponse, TextPart, ToolReturnPart, UserPromptPart, ) from pydantic_ai.models import Model as PydanticModel from pydantic_ai.settings import ModelSettings from pydantic_ai.usage import RequestUsage, UsageLimits from pydantic_ai import ModelRetry, RunContext from .metered_model import MeteredModel from .sandbox import TraceSandbox from ..providers.pydantic_ai import resolve_model UsageCallback = Callable[[RequestUsage, str], None] logger = logging.getLogger(__name__) # ------------------------------------------------------------------ # Default tools # ------------------------------------------------------------------ def register_execute_code(agent: PydanticAgent[AgenticDeps, Any]) -> None: """Register the generic ``execute_code`` tool. Expects ``deps.sandbox`` (a :class:`TraceSandbox` or compatible) and ``deps.config.timeout`` / ``deps.config.max_output_chars``. """ @agent.tool(retries=3) def execute_code(ctx: RunContext[AgenticDeps], code: str) -> str: """Execute Python code in the sandbox. Variables persist across calls. Pre-loaded modules: ``json``, ``re``, ``collections``, ``datetime``. Built-in helper: ``register_helper(name, source, description)`` defines a reusable Python function in this sandbox AND auto-injects it into every child you later spawn via ``recurse`` — register extraction/scoring logic once, reuse it across children. Args: code: Python code to execute. Returns: Captured stdout/stderr from execution. """ ctx.deps.iteration += 1 if ctx.deps.sandbox is None: return "(no sandbox configured)" sandbox = ctx.deps.sandbox timeout = ctx.deps.config.timeout max_output = ctx.deps.config.max_output_chars result = sandbox.execute(code, timeout=timeout) if result.exception: error_msg = f"{type(result.exception).__name__}: {result.exception}" stdout_ctx = "" if result.stdout: stdout_ctx = f"stdout before error:\n{result.stdout[:max_output]}\n\n" raise ModelRetry( f"{stdout_ctx}Code error:\n{error_msg}\n\nFix the bug and try again." ) parts: list[str] = [] if result.stdout: parts.append(result.stdout) if result.stderr: parts.append(f"stderr: {result.stderr}") output = "\n".join(parts) if parts else "(no output)" if len(output) > max_output: remaining = len(output) - max_output output = ( f"{output[:max_output]}\n" f"[TRUNCATED: {remaining} chars remaining]" ) return output def register_recurse(agent: PydanticAgent[AgenticDeps, Any]) -> None: """Register the generic ``recurse`` tool for depth-based decomposition. Expects ``deps.run_session_fn`` to be set (done by :meth:`RecursiveAgent.run`). """ @agent.tool async def recurse( ctx: RunContext[AgenticDeps], prompt: str, context_code: str = "", ) -> str: """Spawn a child session to investigate a sub-problem in isolation. Use this to keep your context lean: the child works through bulky data in its own context window and returns only a text summary. The child inherits a copy of your sandbox variables and any helpers you've registered via `register_helper`. It does NOT see your conversation, so `prompt` must be self-contained. Calling `recurse` multiple times in a single assistant turn dispatches the children in parallel. Args: prompt: Self-contained instructions. Name the sandbox variables to inspect and say what to return. context_code: Optional Python run once in the child's sandbox before it starts (e.g. ``chunk = traces[5:10]``). Returns: Text summary of the child's structured output. """ deps = ctx.deps if deps.run_session_fn is None: return "(recurse unavailable — no session runner configured)" if deps.sandbox is None: return "(recurse unavailable — no sandbox on deps)" sandbox = deps.sandbox # Create child sandbox inheriting parent's injected data child_sandbox = TraceSandbox(trace=None) for key, value in sandbox.namespace.items(): if not key.startswith("_") and not callable(value): child_sandbox.inject(key, value) # Inherit registered helpers parent_registry = sandbox.namespace.get("helper_registry", {}) if isinstance(parent_registry, dict): timeout = deps.config.timeout for hname, meta in parent_registry.items(): if isinstance(meta, dict) and isinstance(meta.get("source"), str): try: child_sandbox.execute(meta["source"], timeout=timeout) child_registry = child_sandbox.namespace.setdefault( "helper_registry", {} ) child_registry[hname] = { "description": meta.get("description", ""), "source": meta["source"], } except Exception: pass # Run optional context_code if context_code.strip(): result = child_sandbox.execute(context_code, timeout=deps.config.timeout) if result.exception: raise ModelRetry( f"context_code failed: {result.exception}\n" "Fix the code and try again." ) # Compute child budget cfg = deps.config remaining = max(0, cfg.max_tokens - deps.parent_usage_tokens) child_token_budget = max(10_000, int(remaining * cfg.child_budget_fraction)) # Build child deps (same type as parent) child_deps = deps.__class__( **{ **{ f.name: getattr(deps, f.name) for f in deps.__dataclass_fields__.values() }, "sandbox": child_sandbox, "depth": deps.depth + 1, "iteration": 0, "parent_usage_tokens": 0, } ) try: output, _ = await deps.run_session_fn( deps=child_deps, prompt=prompt, depth=deps.depth + 1, ) # Serialize child output to text if hasattr(output, "model_dump"): d = output.model_dump(exclude={"raw"}, exclude_defaults=True) parts = [f"{k}: {v}" for k, v in d.items() if v] return "\n".join(parts) if parts else "(empty output)" return str(output) if output else "(empty output)" except Exception as e: return f"(child session failed: {e})" # ------------------------------------------------------------------ # Configuration # ------------------------------------------------------------------ DEFAULT_COMPACTION_SUMMARY_PROMPT = """\ Summarize your progress so far. Structure your response with these sections: 1. **What you've done**: Steps completed, tools used, key decisions made. 2. **Findings so far**: Concrete results, computed values, identified patterns. 3. **Remaining work**: What hasn't been done yet. 4. **Current direction**: What you were investigating when this summary was requested. Be concise but preserve all concrete results and variable names.""" @dataclass class AgenticConfig: """Base configuration for agentic steps with compaction and recursion. Subclass to add step-specific fields (e.g. sandbox timeout). """ # Budget (wired to PydanticAI UsageLimits) max_tokens: int = 500_000 max_requests: int = 50 context_window: int = 128_000 # Recursion max_depth: int = 2 child_budget_fraction: float = 0.5 # Compaction max_compactions: int = 3 microcompact_keep_recent: int = 3 # Sandbox execution timeout: float = 60.0 max_output_chars: int = 50_000 # Metering — fired once per completed pydantic-ai model request # (orchestrator turn, child session, compaction summary). Exceptions # inside the callback are swallowed by MeteredModel so a broken # meter never crashes a run. usage_callback: UsageCallback | None = None def build_usage_limits(self, remaining_tokens: int | None = None) -> UsageLimits: base = remaining_tokens or self.max_tokens return UsageLimits( total_tokens_limit=base, request_limit=self.max_requests, ) # ------------------------------------------------------------------ # Dependency container # ------------------------------------------------------------------ @dataclass class AgenticDeps: """Base dependencies for agentic steps. Subclass to add step-specific deps (trace data, etc.). """ config: AgenticConfig sandbox: Any = None # TraceSandbox or compatible depth: int = 0 max_depth: int = 2 iteration: int = 0 run_session_fn: Callable[..., Awaitable[tuple[Any, Any]]] | None = None parent_usage_tokens: int = 0 # ------------------------------------------------------------------ # Exceptions # ------------------------------------------------------------------ class BudgetExhausted(Exception): """Raised when the agent's token or request budget is fully spent.""" def __init__(self, compaction_count: int = 0, usage: Any = None) -> None: self.compaction_count = compaction_count self.usage = usage super().__init__("Agent budget exhausted") # ------------------------------------------------------------------ # Compaction utilities # ------------------------------------------------------------------ def cost_equivalent_tokens(usage: Any) -> int: """Cost-equivalent token count for budget purposes. Anthropic Bedrock pricing (input side): - fresh: 1.0x base - cache_write: 1.25x base - cache_read: 0.10x base PydanticAI's Bedrock wrapper sets ``input_tokens = fresh + cache_write + cache_read``. We rebuild the cost-equivalent: subtract 0.90x of cache_read (since it should weigh 0.10 not 1.0) and add 0.25x of cache_write (since it should weigh 1.25 not 1.0). Output is counted at 1.0x. """ input_tokens = getattr(usage, "input_tokens", 0) or 0 cache_read = getattr(usage, "cache_read_tokens", 0) or 0 cache_write = getattr(usage, "cache_write_tokens", 0) or 0 output = getattr(usage, "output_tokens", 0) or 0 cost_input = input_tokens - 0.90 * cache_read + 0.25 * cache_write return int(cost_input + output) def is_budget_exhausted( limits: UsageLimits, usage: Any, cost_budget: int | None = None, ) -> bool: """True if cost-equivalent or request budget is spent. When ``cost_budget`` is provided, the *cost-equivalent* token count (``cost_equivalent_tokens``) is checked against it. The gross ``total_tokens_limit`` on ``limits`` is treated as a coarse outer cap and is normally inflated relative to the real budget, so it should rarely trip first when caching is in play. """ if cost_budget is not None and cost_equivalent_tokens(usage) >= cost_budget: return True if limits.total_tokens_limit and usage.total_tokens >= limits.total_tokens_limit: return True if limits.request_limit and usage.requests >= limits.request_limit: return True return False def microcompact( messages: list, keep_recent: int, tool_names: tuple[str, ...], placeholder: str = "[cleared — use tools to re-inspect if needed]", ) -> list: """Tier 1: Clear old tool results from message history. Returns the **same list object** if nothing was cleared — caller uses identity check to detect whether compaction did anything. """ tool_result_positions = [] for msg_idx, msg in enumerate(messages): if isinstance(msg, ModelRequest): for part_idx, part in enumerate(msg.parts): if isinstance(part, ToolReturnPart) and part.tool_name in tool_names: tool_result_positions.append((msg_idx, part_idx)) if len(tool_result_positions) <= keep_recent: return messages to_clear = ( tool_result_positions[:-keep_recent] if keep_recent > 0 else tool_result_positions ) compacted = copy.deepcopy(messages) for msg_idx, part_idx in to_clear: compacted[msg_idx].parts[part_idx].content = placeholder return compacted async def summarize_and_compact( agent: PydanticAgent[AgenticDeps, Any], messages: list, deps: Any, compaction_count: int, summary_prompt: str = DEFAULT_COMPACTION_SUMMARY_PROMPT, continuation_message: str = "", ) -> list: """Tier 2: Full summarization — LLM summarizes, history pruned.""" summary_result = await agent.run( summary_prompt, message_history=messages, deps=deps, output_type=str, ) summary = summary_result.output if not continuation_message: continuation_message = ( f"Your conversation was compacted ({compaction_count} time(s)). " "Do NOT repeat work already completed. Continue." ) return [ ModelResponse( parts=[ TextPart(content=f"[Compaction summary #{compaction_count}]\n{summary}") ] ), ModelRequest(parts=[UserPromptPart(content=continuation_message)]), ] # ------------------------------------------------------------------ # Async runner # ------------------------------------------------------------------ async def run_agent_with_compaction( agent: PydanticAgent[AgenticDeps, Any], *, deps: AgenticDeps, prompt: str, usage_limits: UsageLimits, config: AgenticConfig, tool_names_to_compact: tuple[str, ...] = (), compaction_summary_prompt: str = DEFAULT_COMPACTION_SUMMARY_PROMPT, compaction_continuation: str = "", microcompact_placeholder: str = "[cleared — use tools to re-inspect if needed]", on_compaction: Callable[[AgenticDeps, int, list], None] | None = None, span_label: str = "rr", ) -> tuple[Any, dict]: """Run a PydanticAI agent with two-tier compaction. Returns ``(output, metadata)``. Raises :class:`BudgetExhausted` when budget is fully spent. """ message_history = None compaction_count = 0 user_prompt = prompt cumulative_usage = None last_run: Any = None span_name = ( f"{span_label}.session" if deps.depth == 0 else f"{span_label}.session.child" ) with _rr_span(span_name, depth=deps.depth): while True: try: async with agent.iter( user_prompt, deps=deps, message_history=message_history, usage_limits=usage_limits, usage=cumulative_usage, ) as agent_run: last_run = agent_run async for _node in agent_run: deps.parent_usage_tokens = agent_run.usage().total_tokens or 0 assert agent_run.result is not None output = agent_run.result.output usage = agent_run.result.usage() metadata = { "usage": { "input_tokens": usage.input_tokens, "output_tokens": usage.output_tokens, "total_tokens": usage.total_tokens, "requests": usage.requests, "cache_read_tokens": getattr(usage, "cache_read_tokens", 0), "cache_write_tokens": getattr( usage, "cache_write_tokens", 0 ), }, "compactions": compaction_count, "depth": deps.depth, "iterations": deps.iteration, "timed_out": False, } return output, metadata except UsageLimitExceeded: messages = last_run.all_messages() cumulative_usage = last_run.usage() if is_budget_exhausted( usage_limits, cumulative_usage, cost_budget=config.max_tokens, ): raise BudgetExhausted( compaction_count=compaction_count, usage=cumulative_usage, ) compacted = microcompact( messages, config.microcompact_keep_recent, tool_names_to_compact, placeholder=microcompact_placeholder, ) if compacted is messages: compaction_count += 1 if compaction_count > config.max_compactions: raise BudgetExhausted( compaction_count=compaction_count, usage=cumulative_usage, ) if on_compaction: on_compaction(deps, compaction_count, messages) compacted = await summarize_and_compact( agent, messages, deps, compaction_count, summary_prompt=compaction_summary_prompt, continuation_message=compaction_continuation, ) message_history = compacted user_prompt = "Continue your analysis." # ------------------------------------------------------------------ # Sync wrapper # ------------------------------------------------------------------ def run_agent_sync( agent: PydanticAgent[AgenticDeps, Any], **kwargs: Any ) -> tuple[Any, dict]: """Synchronous wrapper around :func:`run_agent_with_compaction`.""" coro = run_agent_with_compaction(agent, **kwargs) try: loop = asyncio.get_running_loop() except RuntimeError: loop = None if loop and loop.is_running(): with concurrent.futures.ThreadPoolExecutor(max_workers=1) as pool: return pool.submit(asyncio.run, coro).result() else: return asyncio.run(coro) # ------------------------------------------------------------------ # RecursiveAgent — high-level API # ------------------------------------------------------------------ ToolRegistrar = Callable[..., None] class RecursiveAgent: """A PydanticAI agent with compaction, recursion, and budget management. This is the high-level API. Callers provide: - ``output_type``: The structured output schema - ``tools``: List of tool registrar functions ``(agent) -> None`` - ``system_prompt``: The system prompt - ``config``: Budget, compaction, and recursion settings The agent handles compaction and child session spawning automatically. Example:: agent = RecursiveAgent( model="gpt-4o-mini", output_type=ReflectorOutput, system_prompt="You are a trace analyst...", tools=[register_execute_code, register_analysis_tools], tool_names_to_compact=("execute_code", "analyze"), ) output, metadata = agent.run(prompt="Analyze...", deps=my_deps) """ def __init__( self, model: str | PydanticModel, *, output_type: Type, system_prompt: str, config: AgenticConfig | None = None, model_settings: ModelSettings | None = None, tools: Sequence[ToolRegistrar] = (), tool_names_to_compact: tuple[str, ...] = (), compaction_summary_prompt: str = DEFAULT_COMPACTION_SUMMARY_PROMPT, compaction_continuation: str = "", microcompact_placeholder: str = "[cleared — use tools to re-inspect if needed]", on_compaction: Callable[[AgenticDeps, int, list], None] | None = None, span_label: str = "rr", ) -> None: self.config = config or AgenticConfig() self._model = model self._model_settings = model_settings self._output_type = output_type self._system_prompt = system_prompt self._tools = list(tools) self._tool_names_to_compact = tool_names_to_compact self._compaction_summary_prompt = compaction_summary_prompt self._compaction_continuation = compaction_continuation self._microcompact_placeholder = microcompact_placeholder self._on_compaction = on_compaction self._span_label = span_label # Build root agent (depth=0) self._agent = self._create_agent(depth=0) def _create_agent(self, depth: int = 0) -> PydanticAgent[AgenticDeps, Any]: """Create a PydanticAI agent for the given recursion depth. The root (depth 0) uses the configured ``output_type`` (typically a structured Pydantic model). Children return free-form text: they exist to investigate one sub-problem and report a focused answer, not to produce a full reflection. """ if isinstance(self._model, PydanticModel): resolved = self._model else: resolved = resolve_model(self._model) if self.config.usage_callback is not None: resolved = MeteredModel(resolved, self.config.usage_callback) output_type = self._output_type if depth == 0 else str agent: PydanticAgent[AgenticDeps, Any] = PydanticAgent( resolved, output_type=output_type, system_prompt=self._system_prompt, retries=3, model_settings=self._model_settings, defer_model_check=True, ) # Default tools: execute_code + recurse (if not at max depth) register_execute_code(agent) if depth < self.config.max_depth: register_recurse(agent) # Additional caller-provided tools for registrar in self._tools: registrar(agent) return agent async def _run_child_session( self, *, deps: AgenticDeps, prompt: str, depth: int = 0, ) -> tuple[Any, AgenticDeps]: """Run a child session with its own agent and budget.""" child_agent = self._create_agent(depth=depth) remaining = getattr(deps, "_remaining_tokens", None) try: output, metadata = await run_agent_with_compaction( child_agent, deps=deps, prompt=prompt, usage_limits=self.config.build_usage_limits(remaining_tokens=remaining), config=self.config, tool_names_to_compact=self._tool_names_to_compact, compaction_summary_prompt=self._compaction_summary_prompt, compaction_continuation=self._compaction_continuation, microcompact_placeholder=self._microcompact_placeholder, on_compaction=self._on_compaction, span_label=self._span_label, ) return output, deps except BudgetExhausted: return None, deps def run( self, *, deps: AgenticDeps, prompt: str, remaining_tokens: int | None = None, ) -> tuple[Any, dict]: """Run the agent synchronously with compaction. Args: deps: Agent dependencies. prompt: Initial prompt. remaining_tokens: Override token budget (for child sessions). Returns: Tuple of (output, metadata_dict). Raises: BudgetExhausted: When budget is fully spent. """ # Wire up child session runner deps.run_session_fn = self._run_child_session return run_agent_sync( self._agent, deps=deps, prompt=prompt, usage_limits=self.config.build_usage_limits( remaining_tokens=remaining_tokens ), config=self.config, tool_names_to_compact=self._tool_names_to_compact, compaction_summary_prompt=self._compaction_summary_prompt, compaction_continuation=self._compaction_continuation, microcompact_placeholder=self._microcompact_placeholder, on_compaction=self._on_compaction, span_label=self._span_label, ) # ------------------------------------------------------------------ # Sandbox helpers # ------------------------------------------------------------------ def create_sandbox( self, *, trace: Any = None, variables: dict[str, Any] | None = None, ) -> TraceSandbox: """Create a sandbox and inject variables. Args: trace: Optional trace object passed to TraceSandbox constructor. variables: Dict of ``{name: value}`` to inject into the sandbox namespace. Returns: A ready-to-use :class:`TraceSandbox`. """ sandbox = TraceSandbox(trace=trace, llm_query_fn=None) if variables: for name, value in variables.items(): sandbox.inject(name, value) return sandbox @staticmethod def on_compaction(deps: AgenticDeps, compaction_count: int, messages: list) -> None: """Default compaction callback — save metadata to sandbox history. Subclasses can override or pass a different callback via ``on_compaction`` in ``__init__``. """ sandbox = getattr(deps, "sandbox", None) if sandbox is not None: history = sandbox.namespace.get("history", []) history.append( { "compaction_round": compaction_count, "message_count": len(messages), } ) sandbox.namespace["history"] = history