feat(F-LOOP): true ReAct loop with structured Thought/Action/Observation step tracking
Check Cross-Plugin Imports / check (push) Has been cancelled
Check Cross-Plugin Imports / check (push) Has been cancelled
- app/ai/agent_loop.py: ReActStep + ReActResult dataclasses, run_react_loop() with LLM→Tool→Observe loop, max_steps/timeout graceful stop, ErrorCategory retry (TRANSIENT→retry, PERMANENT→stop, PARTIAL→continue), cost accumulation, agent.step hook, on_step callback - app/plugins/builtins/automation/models.py: AgentRunStep model - alembic/versions/0121_agent_run_steps.py: migration for agent run steps table - app/plugins/builtins/automation/agent_runner.py: refactored to use run_react_loop(), saves steps to DB, updates AgentRun with cost/status/duration - tests/test_agent_loop.py: 11 tests (all passing, mocked, no DB/LLM needed) - PROGRESS.md: Phase F started, F-LOOP marked done
This commit is contained in:
@@ -0,0 +1,389 @@
|
||||
"""Core ReAct (Reasoning + Acting) loop for AI agents.
|
||||
|
||||
Implements a true ReAct loop that alternates between LLM reasoning and tool
|
||||
execution. Each step records the thought (LLM content), action (tool name),
|
||||
action_input (tool arguments), and observation (tool result).
|
||||
|
||||
The loop terminates when:
|
||||
- The LLM returns a final response without tool calls (completed)
|
||||
- max_steps is reached (stopped_max_steps)
|
||||
- timeout is exceeded (stopped_timeout)
|
||||
- A permanent error occurs (stopped_error)
|
||||
|
||||
Error handling uses ``ErrorCategory`` from ``app.core.error_codes``:
|
||||
- TRANSIENT → retry the LLM call (up to 3 retries per step)
|
||||
- PERMANENT → stop the loop immediately
|
||||
- PARTIAL → continue with partial results
|
||||
|
||||
Usage::
|
||||
|
||||
from app.ai.agent_loop import run_react_loop
|
||||
|
||||
result = await run_react_loop(
|
||||
agent_definition=agent,
|
||||
messages=[{"role": "user", "content": "Summarize recent emails"}],
|
||||
tools=tool_schemas,
|
||||
tool_registry=registry,
|
||||
db=db_session,
|
||||
tenant_id=tenant_id,
|
||||
user_id=user_id,
|
||||
)
|
||||
print(result.final_content, result.total_cost_usd, result.steps_taken)
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
import time
|
||||
import uuid
|
||||
from dataclasses import dataclass, field
|
||||
from datetime import UTC, datetime
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
from app.ai.llm_client import llm_complete
|
||||
from app.core.error_codes import ErrorCategory, classify_exception
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import Callable
|
||||
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
|
||||
from app.plugins.builtins.ai_assistant.tool_registry import ToolRegistry
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Maximum retries for transient errors per LLM step
|
||||
_MAX_TRANSIENT_RETRIES = 3
|
||||
|
||||
|
||||
# ──────────────────────────────────────────────────────────────────────────
|
||||
# Data structures
|
||||
# ──────────────────────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
@dataclass
|
||||
class ReActStep:
|
||||
"""A single step in the ReAct loop (Thought → Action → Observation)."""
|
||||
|
||||
step_number: int
|
||||
thought: str # LLM content before tool calls
|
||||
action: str | None # Tool name (None if final response)
|
||||
action_input: dict[str, Any] | None # Tool arguments
|
||||
observation: str | None # Tool result
|
||||
cost_usd: float
|
||||
timestamp: str # ISO format
|
||||
|
||||
|
||||
@dataclass
|
||||
class ReActResult:
|
||||
"""Final result of the ReAct loop."""
|
||||
|
||||
final_content: str
|
||||
steps: list[ReActStep] = field(default_factory=list)
|
||||
total_cost_usd: float = 0.0
|
||||
steps_taken: int = 0
|
||||
status: str = "completed" # completed | stopped_max_steps | stopped_timeout | stopped_error
|
||||
error: str | None = None
|
||||
|
||||
|
||||
# ──────────────────────────────────────────────────────────────────────────
|
||||
# Core loop
|
||||
# ──────────────────────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
def _extract_tool_calls(raw_response: Any) -> list[dict[str, Any]]:
|
||||
"""Extract tool calls from a LiteLLM raw response.
|
||||
|
||||
Returns a list of dicts with keys: ``id``, ``name``, ``arguments``.
|
||||
"""
|
||||
tool_calls: list[dict[str, Any]] = []
|
||||
try:
|
||||
msg = raw_response.choices[0].message
|
||||
if hasattr(msg, "tool_calls") and msg.tool_calls:
|
||||
for tc in msg.tool_calls:
|
||||
tool_calls.append({
|
||||
"id": tc.id or "",
|
||||
"name": tc.function.name if tc.function else "",
|
||||
"arguments": tc.function.arguments if tc.function and tc.function.arguments else "{}",
|
||||
})
|
||||
except (AttributeError, IndexError, TypeError) as exc:
|
||||
logger.debug("Failed to extract tool calls from response: %s", exc)
|
||||
return tool_calls
|
||||
|
||||
|
||||
async def _execute_tool(
|
||||
tool_registry: ToolRegistry,
|
||||
tool_name: str,
|
||||
arguments: dict[str, Any],
|
||||
context: dict[str, Any],
|
||||
) -> str:
|
||||
"""Execute a single tool call via the registry.
|
||||
|
||||
Returns the tool result as a string, or an error message.
|
||||
"""
|
||||
tool = tool_registry.get(tool_name)
|
||||
if tool is None:
|
||||
return f"Error: Tool '{tool_name}' not found"
|
||||
|
||||
try:
|
||||
result = await tool.handler(arguments=arguments, context=context)
|
||||
return result if isinstance(result, str) else json.dumps(result)
|
||||
except Exception as exc:
|
||||
logger.exception("Tool '%s' execution failed", tool_name)
|
||||
return f"Error: {exc}"
|
||||
|
||||
|
||||
async def run_react_loop(
|
||||
agent_definition: Any, # AgentDefinition from automation models
|
||||
messages: list[dict[str, Any]],
|
||||
tools: list[dict[str, Any]],
|
||||
tool_registry: ToolRegistry,
|
||||
db: AsyncSession,
|
||||
tenant_id: uuid.UUID,
|
||||
user_id: uuid.UUID,
|
||||
agent_run_id: uuid.UUID | None = None,
|
||||
max_steps: int = 20,
|
||||
timeout_seconds: int = 300,
|
||||
trace_id: str | None = None,
|
||||
on_step: Callable | None = None,
|
||||
) -> ReActResult:
|
||||
"""Execute a ReAct loop: LLM reasoning → tool execution → repeat.
|
||||
|
||||
Args:
|
||||
agent_definition: AgentDefinition with llm_model, system_prompt, etc.
|
||||
messages: Initial chat messages (without system prompt).
|
||||
tools: OpenAI-format tool schemas for function calling.
|
||||
tool_registry: ToolRegistry instance for tool execution.
|
||||
db: Async DB session.
|
||||
tenant_id: Tenant ID for multi-tenancy.
|
||||
user_id: User ID for permission context.
|
||||
agent_run_id: Optional AgentRun ID for step persistence.
|
||||
max_steps: Maximum loop iterations (default 20).
|
||||
timeout_seconds: Overall timeout (default 300).
|
||||
trace_id: Optional trace ID for correlation.
|
||||
on_step: Optional async callback fired after each step.
|
||||
|
||||
Returns:
|
||||
ReActResult with final content, steps, cost, and status.
|
||||
"""
|
||||
from app.core.hooks import do_action
|
||||
|
||||
result = ReActResult(final_content="", status="completed")
|
||||
start_time = time.monotonic()
|
||||
|
||||
# Build LLM parameters from agent definition
|
||||
litellm_model = getattr(agent_definition, "llm_model", None) or "gpt-4o"
|
||||
system_prompt = getattr(agent_definition, "system_prompt", "") or "You are a helpful AI assistant."
|
||||
api_key = getattr(agent_definition, "api_key", None)
|
||||
api_base = getattr(agent_definition, "api_base", None)
|
||||
provider = getattr(agent_definition, "provider", None)
|
||||
max_tokens = getattr(agent_definition, "max_tokens", None) or 1000
|
||||
|
||||
# Build the full message list with system prompt prepended
|
||||
full_messages: list[dict[str, Any]] = [
|
||||
{"role": "system", "content": system_prompt},
|
||||
*messages,
|
||||
]
|
||||
|
||||
tool_context: dict[str, Any] = {
|
||||
"tenant_id": str(tenant_id),
|
||||
"user_id": str(user_id),
|
||||
"db": db,
|
||||
}
|
||||
|
||||
for step_num in range(1, max_steps + 1):
|
||||
# ── Timeout check ──
|
||||
elapsed = time.monotonic() - start_time
|
||||
if elapsed >= timeout_seconds:
|
||||
result.status = "stopped_timeout"
|
||||
result.error = f"Timeout after {elapsed:.1f}s (limit {timeout_seconds}s)"
|
||||
logger.warning("ReAct loop timed out at step %d: %s", step_num, result.error)
|
||||
break
|
||||
|
||||
# ── LLM call with transient retry ──
|
||||
llm_result: dict[str, Any] | None = None
|
||||
last_error: str | None = None
|
||||
|
||||
for retry in range(_MAX_TRANSIENT_RETRIES + 1):
|
||||
try:
|
||||
llm_result = await llm_complete(
|
||||
model=litellm_model,
|
||||
messages=full_messages,
|
||||
tools=tools if tools else None,
|
||||
temperature=0.3,
|
||||
max_tokens=max_tokens,
|
||||
api_key=api_key,
|
||||
api_base=api_base,
|
||||
provider=provider,
|
||||
trace_id=trace_id,
|
||||
tenant_id=tenant_id,
|
||||
db=db,
|
||||
)
|
||||
break
|
||||
except Exception as exc:
|
||||
last_error = str(exc)
|
||||
category = classify_exception(exc)
|
||||
|
||||
if category == ErrorCategory.PERMANENT:
|
||||
result.status = "stopped_error"
|
||||
result.error = f"Permanent error at step {step_num}: {exc}"
|
||||
logger.error("ReAct loop permanent error: %s", result.error)
|
||||
return result
|
||||
|
||||
if category == ErrorCategory.TRANSIENT and retry < _MAX_TRANSIENT_RETRIES:
|
||||
backoff = 2 ** retry
|
||||
logger.warning(
|
||||
"Transient error at step %d (retry %d/%d): %s — retrying in %ds",
|
||||
step_num, retry + 1, _MAX_TRANSIENT_RETRIES, exc, backoff,
|
||||
)
|
||||
await asyncio.sleep(backoff)
|
||||
continue
|
||||
|
||||
# PARTIAL or exhausted retries
|
||||
if category == ErrorCategory.PARTIAL:
|
||||
logger.warning("Partial error at step %d: %s — continuing", step_num, exc)
|
||||
last_error = str(exc)
|
||||
break
|
||||
|
||||
# Exhausted transient retries
|
||||
result.status = "stopped_error"
|
||||
result.error = f"Error after {retry + 1} retries at step {step_num}: {exc}"
|
||||
logger.error("ReAct loop error: %s", result.error)
|
||||
return result
|
||||
|
||||
if llm_result is None:
|
||||
result.status = "stopped_error"
|
||||
result.error = f"LLM call failed at step {step_num}: {last_error}"
|
||||
return result
|
||||
|
||||
# ── Extract response data ──
|
||||
content = llm_result.get("content", "")
|
||||
cost_usd = llm_result.get("cost_usd", 0.0)
|
||||
result.total_cost_usd += cost_usd
|
||||
|
||||
tool_calls = _extract_tool_calls(llm_result.get("raw_response"))
|
||||
|
||||
# ── No tool calls → final response ──
|
||||
if not tool_calls:
|
||||
step = ReActStep(
|
||||
step_number=step_num,
|
||||
thought=content,
|
||||
action=None,
|
||||
action_input=None,
|
||||
observation=None,
|
||||
cost_usd=cost_usd,
|
||||
timestamp=datetime.now(UTC).isoformat(),
|
||||
)
|
||||
result.steps.append(step)
|
||||
result.final_content = content
|
||||
result.steps_taken = step_num
|
||||
|
||||
# Fire hook
|
||||
await do_action(
|
||||
"agent.step",
|
||||
agent_id=str(getattr(agent_definition, "id", "")),
|
||||
step_number=step_num,
|
||||
thought=content,
|
||||
action=None,
|
||||
observation=None,
|
||||
cost_usd=cost_usd,
|
||||
agent_run_id=str(agent_run_id) if agent_run_id else None,
|
||||
trace_id=trace_id,
|
||||
)
|
||||
|
||||
# Callback
|
||||
if on_step:
|
||||
try:
|
||||
await on_step(step)
|
||||
except Exception:
|
||||
logger.debug("on_step callback failed", exc_info=True)
|
||||
|
||||
break
|
||||
|
||||
# ── Execute tool calls ──
|
||||
# Append assistant message with tool calls to conversation
|
||||
full_messages.append({
|
||||
"role": "assistant",
|
||||
"content": content,
|
||||
"tool_calls": [
|
||||
{
|
||||
"id": tc["id"],
|
||||
"type": "function",
|
||||
"function": {"name": tc["name"], "arguments": tc["arguments"]},
|
||||
}
|
||||
for tc in tool_calls
|
||||
],
|
||||
})
|
||||
|
||||
# Execute each tool call and collect observations
|
||||
observations: list[str] = []
|
||||
for tc in tool_calls:
|
||||
tool_name = tc["name"]
|
||||
try:
|
||||
args = json.loads(tc["arguments"]) if tc["arguments"] else {}
|
||||
except json.JSONDecodeError:
|
||||
args = {}
|
||||
logger.warning("Invalid JSON arguments for tool '%s': %s", tool_name, tc["arguments"])
|
||||
|
||||
observation = await _execute_tool(tool_registry, tool_name, args, tool_context)
|
||||
observations.append(observation)
|
||||
|
||||
# Feed tool result back into conversation
|
||||
full_messages.append({
|
||||
"role": "tool",
|
||||
"tool_call_id": tc["id"],
|
||||
"content": observation,
|
||||
})
|
||||
|
||||
# Record step
|
||||
step = ReActStep(
|
||||
step_number=step_num,
|
||||
thought=content,
|
||||
action=tool_name,
|
||||
action_input=args,
|
||||
observation=observation,
|
||||
cost_usd=cost_usd / len(tool_calls) if tool_calls else cost_usd,
|
||||
timestamp=datetime.now(UTC).isoformat(),
|
||||
)
|
||||
result.steps.append(step)
|
||||
|
||||
# Fire hook
|
||||
await do_action(
|
||||
"agent.step",
|
||||
agent_id=str(getattr(agent_definition, "id", "")),
|
||||
step_number=step_num,
|
||||
thought=content,
|
||||
action=tool_name,
|
||||
observation=observation,
|
||||
cost_usd=cost_usd,
|
||||
agent_run_id=str(agent_run_id) if agent_run_id else None,
|
||||
trace_id=trace_id,
|
||||
)
|
||||
|
||||
# Callback
|
||||
if on_step:
|
||||
try:
|
||||
await on_step(step)
|
||||
except Exception:
|
||||
logger.debug("on_step callback failed", exc_info=True)
|
||||
|
||||
result.steps_taken = step_num
|
||||
|
||||
# If this was the last allowed step, stop gracefully
|
||||
if step_num >= max_steps:
|
||||
result.status = "stopped_max_steps"
|
||||
result.error = f"Reached max_steps limit ({max_steps})"
|
||||
result.final_content = content
|
||||
logger.warning("ReAct loop stopped at max_steps=%d", max_steps)
|
||||
break
|
||||
|
||||
# If loop completed without a final response (e.g. all steps had tool calls)
|
||||
if not result.final_content and result.steps:
|
||||
result.final_content = result.steps[-1].thought or ""
|
||||
|
||||
if result.status == "completed" and not result.final_content:
|
||||
result.final_content = ""
|
||||
|
||||
return result
|
||||
Reference in New Issue
Block a user