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:
+11
-1
@@ -15,7 +15,7 @@
|
||||
| C.5 — Import/Export | `done` | 2026-08-13 | 2026-08-13 | 8 | 8 |
|
||||
| D — Undo/Restore | `done` | 2026-08-13 | 2026-08-13 | 13 | 13 |
|
||||
| E — Search | `done` | 2026-08-14 | 2026-08-14 | 24 | 24 |
|
||||
| F — Agents | `not_started` | — | — | 0 | ~28 |
|
||||
| F — Agents | `in_progress` | 2026-08-17 | — | 1 | ~28 |
|
||||
| G — Workflows | `not_started` | — | — | 0 | ~24 |
|
||||
| H — Knowledge | `not_started` | — | — | 0 | ~18 |
|
||||
| I — Integration & Workstream | `not_started` | — | — | 0 | ~25 |
|
||||
@@ -237,6 +237,16 @@ Detaillierte Task-Listen werden beim Start der jeweiligen Phase eingetragen. Sie
|
||||
|
||||
---
|
||||
|
||||
## Phase F — Agents
|
||||
|
||||
### F.1 ReAct Loop
|
||||
|
||||
| Task | Status | Forgejo Issue | Verifiziert |
|
||||
|------|-------|---------------|------------|
|
||||
| F-LOOP | `done` | — | ✅ `app/ai/agent_loop.py` — ReActStep/ReActResult dataclasses + `run_react_loop()` mit LLM→Tool→Observe Loop, max_steps/timeout Graceful-Stop, ErrorCategory-basiertes Retry (TRANSIENT→retry, PERMANENT→stop, PARTIAL→continue), Cost-Accumulation, `agent.step` Hook. `AgentRunStep` Model + Migration 0121. `agent_runner.py` auf ReAct-Loop umgestellt. 11/11 Tests grün (mocked, kein DB/LLM benötigt) |
|
||||
|
||||
---
|
||||
|
||||
## Blockierte Tasks
|
||||
|
||||
| Task | Grund | Blockiert seit | Lösung |
|
||||
|
||||
@@ -0,0 +1,51 @@
|
||||
"""Create automation_agent_run_steps table for ReAct loop step tracking.
|
||||
|
||||
Revision ID: 0121
|
||||
Revises: 0120
|
||||
"""
|
||||
|
||||
from alembic import op
|
||||
import sqlalchemy as sa
|
||||
from sqlalchemy.dialects.postgresql import JSONB, UUID as PGUUID
|
||||
|
||||
revision = "0121"
|
||||
down_revision = "0120"
|
||||
branch_labels = None
|
||||
depends_on = None
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
op.create_table(
|
||||
"automation_agent_run_steps",
|
||||
sa.Column("id", PGUUID(as_uuid=True), primary_key=True),
|
||||
sa.Column("tenant_id", PGUUID(as_uuid=True), nullable=False, index=True),
|
||||
sa.Column(
|
||||
"agent_run_id",
|
||||
PGUUID(as_uuid=True),
|
||||
sa.ForeignKey("automation_agent_runs.id", ondelete="CASCADE"),
|
||||
nullable=False,
|
||||
index=True,
|
||||
),
|
||||
sa.Column("step_number", sa.Integer, nullable=False),
|
||||
sa.Column("thought", sa.Text, nullable=True),
|
||||
sa.Column("action", sa.String(255), nullable=True),
|
||||
sa.Column("action_input", JSONB, nullable=True),
|
||||
sa.Column("observation", sa.Text, nullable=True),
|
||||
sa.Column("cost_usd", sa.Float, nullable=False, server_default="0.0"),
|
||||
sa.Column(
|
||||
"created_at",
|
||||
sa.DateTime(timezone=True),
|
||||
nullable=False,
|
||||
server_default=sa.func.now(),
|
||||
),
|
||||
)
|
||||
op.create_index(
|
||||
"ix_agent_run_steps_run",
|
||||
"automation_agent_run_steps",
|
||||
["tenant_id", "agent_run_id"],
|
||||
)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
op.drop_index("ix_agent_run_steps_run", table_name="automation_agent_run_steps")
|
||||
op.drop_table("automation_agent_run_steps")
|
||||
@@ -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
|
||||
@@ -5,25 +5,22 @@ Safety features:
|
||||
- Max duration per execution (asyncio timeout)
|
||||
- Auto-stop on infinite loop (same tool called 5x consecutively)
|
||||
- Budget limit per agent (track cumulative cost_usd, stop if over budget)
|
||||
- ReAct loop with structured Thought/Action/Observation step tracking
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
from datetime import UTC, datetime
|
||||
from typing import Any
|
||||
|
||||
from sqlalchemy import func, select
|
||||
|
||||
from app.ai.llm_client import llm_complete
|
||||
from app.ai.agent_loop import ReActResult, run_react_loop
|
||||
from app.core.db import get_session_factory
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Track consecutive tool calls per agent run to detect infinite loops
|
||||
_consecutive_tool_calls: dict[str, dict[str, int]] = {} # run_id -> {tool_name: count}
|
||||
|
||||
|
||||
async def run_agent(
|
||||
ctx: dict[str, Any],
|
||||
@@ -32,16 +29,20 @@ async def run_agent(
|
||||
trigger_data: dict[str, Any] | None = None,
|
||||
) -> dict[str, Any]:
|
||||
"""ARQ job function. Loads AgentDefinition from DB, checks rate limits,
|
||||
gathers context, calls LLM via LiteLLM, executes tool calls via ToolRegistry,
|
||||
saves result to AgentRun. Returns result dict.
|
||||
gathers context, runs the ReAct loop, saves steps and result to AgentRun.
|
||||
|
||||
Safety checks:
|
||||
1. Rate limit: max_executions_per_hour
|
||||
2. Max duration: max_duration_seconds (asyncio.timeout)
|
||||
3. Infinite loop: same tool 5x consecutively
|
||||
3. Infinite loop: same tool 5x consecutively (handled in ReAct loop)
|
||||
4. Budget limit: cumulative cost_usd
|
||||
"""
|
||||
from app.plugins.builtins.automation.models import AgentDefinition, AgentRun
|
||||
from app.plugins.builtins.automation.models import (
|
||||
AgentDefinition,
|
||||
AgentRun,
|
||||
AgentRunStep,
|
||||
)
|
||||
from app.plugins.builtins.ai_assistant.contracts import get_tool_registry
|
||||
|
||||
factory = get_session_factory()
|
||||
|
||||
@@ -61,10 +62,6 @@ async def run_agent(
|
||||
return {"error": "Agent is inactive", "status": "skipped"}
|
||||
|
||||
# ── Safety Check 1: Rate Limit ──
|
||||
# Uses DB-based counting (AgentRun rows in last hour) rather than
|
||||
# check_rate_limit() because this counts actual executions per agent,
|
||||
# not just attempts. This is more accurate for per-agent execution caps
|
||||
# and respects the agent-specific max_executions_per_hour setting.
|
||||
if agent.max_executions_per_hour:
|
||||
async with factory() as db:
|
||||
one_hour_ago = datetime.now(UTC)
|
||||
@@ -102,7 +99,6 @@ async def run_agent(
|
||||
# Gather context
|
||||
context_data: dict[str, Any] = {}
|
||||
if trigger_type == "proactive" or agent.mode == "proactive":
|
||||
# Collect context data (recent contacts, mails, events)
|
||||
try:
|
||||
from app.services.contact_service import list_contacts
|
||||
async with factory() as db:
|
||||
@@ -114,16 +110,17 @@ async def run_agent(
|
||||
try:
|
||||
from app.plugins.builtins.mail.models import Mail
|
||||
from sqlalchemy import select as _select
|
||||
mail_q = await db.execute(
|
||||
_select(Mail)
|
||||
.where(Mail.tenant_id == agent.tenant_id)
|
||||
.order_by(Mail.date.desc())
|
||||
.limit(5)
|
||||
)
|
||||
context_data["recent_mails"] = [
|
||||
{"id": str(m.id), "subject": m.subject, "from": m.sender}
|
||||
for m in mail_q.scalars()
|
||||
]
|
||||
async with factory() as db:
|
||||
mail_q = await db.execute(
|
||||
_select(Mail)
|
||||
.where(Mail.tenant_id == agent.tenant_id)
|
||||
.order_by(Mail.date.desc())
|
||||
.limit(5)
|
||||
)
|
||||
context_data["recent_mails"] = [
|
||||
{"id": str(m.id), "subject": m.subject, "from": m.sender}
|
||||
for m in mail_q.scalars()
|
||||
]
|
||||
except Exception:
|
||||
logger.warning("Failed to collect mails for proactive context")
|
||||
|
||||
@@ -135,7 +132,6 @@ async def run_agent(
|
||||
except Exception:
|
||||
logger.warning("Failed to collect events for proactive context")
|
||||
else:
|
||||
# Reactive mode: use trigger_data as context
|
||||
context_data = trigger_data or {}
|
||||
|
||||
# ── Lifecycle: before run ──
|
||||
@@ -153,145 +149,156 @@ async def run_agent(
|
||||
)
|
||||
await db.commit()
|
||||
|
||||
# Call LLM via LiteLLM
|
||||
# ── Prepare tools ──
|
||||
registry = get_tool_registry()
|
||||
tool_ids: list[str] = list(agent.tool_ids or [])
|
||||
tools = registry.get_by_names(tool_ids) if tool_ids else []
|
||||
tool_schemas = [t.to_openai_schema() for t in tools] if tools else []
|
||||
|
||||
# ── Create AgentRun record ──
|
||||
run_id: uuid.UUID | None = None
|
||||
started_at = datetime.now(UTC)
|
||||
async with factory() as db:
|
||||
run = AgentRun(
|
||||
tenant_id=agent.tenant_id,
|
||||
agent_id=agent.id,
|
||||
status="running",
|
||||
started_at=started_at,
|
||||
trigger_type=trigger_type,
|
||||
trigger_data=context_data,
|
||||
)
|
||||
db.add(run)
|
||||
await db.flush()
|
||||
run_id = run.id
|
||||
await db.commit()
|
||||
|
||||
# ── Run ReAct loop ──
|
||||
result_data: dict[str, Any] = {
|
||||
"agent_id": str(agent.id),
|
||||
"agent_name": agent.name,
|
||||
"trigger_type": trigger_type,
|
||||
"status": "running",
|
||||
"run_id": str(run_id) if run_id else None,
|
||||
"llm_response": None,
|
||||
"tool_calls": [],
|
||||
"cost_usd": 0.0,
|
||||
"steps": [],
|
||||
"error": None,
|
||||
}
|
||||
|
||||
# ── Safety Check 3: Max Duration ──
|
||||
max_duration = agent.max_duration_seconds or 300
|
||||
|
||||
try:
|
||||
async def _run_llm() -> None:
|
||||
"""Inner coroutine for LLM call with tool execution."""
|
||||
# Build system prompt from agent configuration
|
||||
system_prompt = agent.system_prompt or "You are a helpful AI assistant."
|
||||
user_prompt = f"Context: {context_data}"
|
||||
import asyncio
|
||||
import uuid as uuid_mod
|
||||
|
||||
litellm_model = agent.model or "gpt-4o"
|
||||
if agent.provider and agent.provider != "openai":
|
||||
litellm_model = f"{agent.provider}/{litellm_model}"
|
||||
react_result: ReActResult = await asyncio.wait_for(
|
||||
run_react_loop(
|
||||
agent_definition=agent,
|
||||
messages=[{"role": "user", "content": f"Context: {context_data}"}],
|
||||
tools=tool_schemas,
|
||||
tool_registry=registry,
|
||||
db=None, # ReAct loop doesn't need DB session for LLM calls directly
|
||||
tenant_id=agent.tenant_id,
|
||||
user_id=agent.created_by or uuid_mod.uuid4(),
|
||||
agent_run_id=run_id,
|
||||
max_steps=20,
|
||||
timeout_seconds=max_duration,
|
||||
),
|
||||
timeout=max_duration + 10, # Extra buffer beyond loop's own timeout
|
||||
)
|
||||
|
||||
result = await llm_complete(
|
||||
model=litellm_model,
|
||||
messages=[
|
||||
{"role": "system", "content": system_prompt},
|
||||
{"role": "user", "content": user_prompt},
|
||||
],
|
||||
temperature=0.3,
|
||||
max_tokens=agent.max_tokens or 1000,
|
||||
api_key=agent.api_key or None,
|
||||
api_base=agent.api_base or None,
|
||||
)
|
||||
content = result["content"]
|
||||
|
||||
# Track cost
|
||||
result_data["cost_usd"] = result["cost_usd"]
|
||||
|
||||
result_data["llm_response"] = content
|
||||
|
||||
# Execute tool calls if LLM returned function calls
|
||||
raw_response = result["raw_response"]
|
||||
if hasattr(raw_response.choices[0].message, "tool_calls") and raw_response.choices[0].message.tool_calls:
|
||||
from app.plugins.builtins.ai_assistant.contracts import get_tool_registry
|
||||
|
||||
registry = get_tool_registry()
|
||||
tool_call_count: dict[str, int] = {}
|
||||
for tc in raw_response.choices[0].message.tool_calls:
|
||||
tool_name = tc.function.name
|
||||
# ── Safety Check 4: Infinite Loop Detection ──
|
||||
tool_call_count[tool_name] = tool_call_count.get(tool_name, 0) + 1
|
||||
if tool_call_count[tool_name] >= 5:
|
||||
logger.warning(
|
||||
"Infinite loop detected for agent %s: tool '%s' called %d times consecutively",
|
||||
agent.id, tool_name, tool_call_count[tool_name],
|
||||
)
|
||||
result_data["error"] = f"Infinite loop detected: tool '{tool_name}' called 5+ times consecutively"
|
||||
result_data["status"] = "failed"
|
||||
return
|
||||
|
||||
tool = registry.get(tool_name)
|
||||
if tool:
|
||||
try:
|
||||
import json
|
||||
args = json.loads(tc.function.arguments)
|
||||
tool_result = await tool.handler(
|
||||
arguments=args,
|
||||
context={"tenant_id": str(agent.tenant_id)},
|
||||
)
|
||||
result_data["tool_calls"].append({
|
||||
"tool": tool_name,
|
||||
"arguments": args,
|
||||
"result": tool_result,
|
||||
})
|
||||
except Exception as e:
|
||||
result_data["tool_calls"].append({
|
||||
"tool": tool_name,
|
||||
"error": str(e),
|
||||
})
|
||||
|
||||
result_data["status"] = "completed"
|
||||
|
||||
# Run with timeout
|
||||
try:
|
||||
await asyncio.wait_for(_run_llm(), timeout=max_duration)
|
||||
except TimeoutError:
|
||||
logger.warning(
|
||||
"Agent %s execution timed out after %d seconds",
|
||||
agent.id, max_duration,
|
||||
)
|
||||
result_data["status"] = "timed_out"
|
||||
result_data["error"] = f"Execution timed out after {max_duration} seconds"
|
||||
result_data["status"] = react_result.status
|
||||
result_data["llm_response"] = react_result.final_content
|
||||
result_data["cost_usd"] = react_result.total_cost_usd
|
||||
result_data["error"] = react_result.error
|
||||
result_data["steps"] = [
|
||||
{
|
||||
"step_number": s.step_number,
|
||||
"thought": s.thought,
|
||||
"action": s.action,
|
||||
"action_input": s.action_input,
|
||||
"observation": s.observation,
|
||||
"cost_usd": s.cost_usd,
|
||||
}
|
||||
for s in react_result.steps
|
||||
]
|
||||
result_data["tool_calls"] = [
|
||||
{"tool": s.action, "arguments": s.action_input, "result": s.observation}
|
||||
for s in react_result.steps if s.action
|
||||
]
|
||||
|
||||
except TimeoutError:
|
||||
logger.warning("Agent %s execution timed out after %d seconds", agent.id, max_duration)
|
||||
result_data["status"] = "stopped_timeout"
|
||||
result_data["error"] = f"Execution timed out after {max_duration} seconds"
|
||||
except Exception as e:
|
||||
logger.exception("Agent run failed for %s", agent.id)
|
||||
result_data["status"] = "failed"
|
||||
result_data["status"] = "stopped_error"
|
||||
result_data["error"] = str(e)
|
||||
|
||||
# ── Save steps to DB ──
|
||||
completed_at = datetime.now(UTC)
|
||||
duration_seconds = (completed_at - started_at).total_seconds()
|
||||
|
||||
try:
|
||||
async with factory() as db:
|
||||
# Save each step
|
||||
for step_data in result_data.get("steps", []):
|
||||
step = AgentRunStep(
|
||||
tenant_id=agent.tenant_id,
|
||||
agent_run_id=run_id,
|
||||
step_number=step_data["step_number"],
|
||||
thought=step_data.get("thought"),
|
||||
action=step_data.get("action"),
|
||||
action_input=step_data.get("action_input"),
|
||||
observation=step_data.get("observation"),
|
||||
cost_usd=step_data.get("cost_usd", 0.0),
|
||||
)
|
||||
db.add(step)
|
||||
|
||||
# Update AgentRun with final results
|
||||
run_result = await db.execute(
|
||||
select(AgentRun).where(AgentRun.id == run_id)
|
||||
)
|
||||
run = run_result.scalar_one_or_none()
|
||||
if run:
|
||||
run.status = result_data["status"]
|
||||
run.completed_at = completed_at
|
||||
run.duration_seconds = duration_seconds
|
||||
run.result = result_data.get("llm_response")
|
||||
run.error = result_data.get("error")
|
||||
run.cost_usd = result_data.get("cost_usd", 0.0)
|
||||
|
||||
await db.commit()
|
||||
except Exception as e:
|
||||
logger.exception("Failed to save agent run steps for %s", agent.id)
|
||||
result_data["save_error"] = str(e)
|
||||
|
||||
# ── Lifecycle: after run ──
|
||||
from app.core.hooks import do_action
|
||||
await do_action("agent.after_run", agent_id=str(agent.id), tenant_id=str(agent.tenant_id), status=result_data.get("status"), result=result_data)
|
||||
from app.core.outbox import enqueue_outbox_event
|
||||
await do_action(
|
||||
"agent.after_run",
|
||||
agent_id=str(agent.id),
|
||||
tenant_id=str(agent.tenant_id),
|
||||
status=result_data.get("status"),
|
||||
result=result_data,
|
||||
)
|
||||
async with factory() as db:
|
||||
await enqueue_outbox_event(
|
||||
db,
|
||||
agent.tenant_id,
|
||||
'agent.run_completed',
|
||||
{'agent_id': str(agent.id), 'tenant_id': str(agent.tenant_id), 'status': result_data.get('status'), 'cost_usd': result_data.get('cost_usd', 0.0)},
|
||||
{
|
||||
'agent_id': str(agent.id),
|
||||
'tenant_id': str(agent.tenant_id),
|
||||
'status': result_data.get('status'),
|
||||
'cost_usd': result_data.get('cost_usd', 0.0),
|
||||
},
|
||||
aggregate_type='agent',
|
||||
aggregate_id=agent.id,
|
||||
)
|
||||
await db.commit()
|
||||
|
||||
# Save result to AgentRun
|
||||
try:
|
||||
async with factory() as db:
|
||||
run = AgentRun(
|
||||
tenant_id=agent.tenant_id,
|
||||
agent_id=agent.id,
|
||||
trigger_type=trigger_type,
|
||||
status=result_data["status"],
|
||||
input_data=context_data,
|
||||
output_data=result_data.get("llm_response"),
|
||||
tool_calls=result_data.get("tool_calls", []),
|
||||
cost_usd=result_data.get("cost_usd", 0.0),
|
||||
error_message=result_data.get("error"),
|
||||
duration_seconds=None,
|
||||
)
|
||||
db.add(run)
|
||||
await db.flush()
|
||||
result_data["run_id"] = str(run.id)
|
||||
except Exception as e:
|
||||
logger.exception("Failed to save AgentRun for %s", agent.id)
|
||||
result_data["save_error"] = str(e)
|
||||
|
||||
return result_data
|
||||
|
||||
|
||||
|
||||
@@ -16,6 +16,7 @@ from sqlalchemy import (
|
||||
String,
|
||||
Text,
|
||||
UniqueConstraint,
|
||||
func,
|
||||
)
|
||||
from sqlalchemy.dialects.postgresql import JSONB
|
||||
from sqlalchemy.dialects.postgresql import UUID as PGUUID
|
||||
@@ -254,6 +255,34 @@ class AutomationRun(Base, TenantMixin):
|
||||
dry_run: Mapped[bool] = mapped_column(Boolean, nullable=False, default=False)
|
||||
|
||||
|
||||
class AgentRunStep(Base, TenantMixin):
|
||||
"""Individual step in a ReAct loop execution (Thought → Action → Observation)."""
|
||||
|
||||
__tablename__ = "automation_agent_run_steps"
|
||||
__table_args__ = (
|
||||
Index("ix_agent_run_steps_run", "tenant_id", "agent_run_id"),
|
||||
)
|
||||
|
||||
id: Mapped[uuid.UUID] = mapped_column(
|
||||
PGUUID(as_uuid=True), primary_key=True, default=uuid.uuid4
|
||||
)
|
||||
agent_run_id: Mapped[uuid.UUID] = mapped_column(
|
||||
PGUUID(as_uuid=True),
|
||||
ForeignKey("automation_agent_runs.id", ondelete="CASCADE"),
|
||||
nullable=False,
|
||||
index=True,
|
||||
)
|
||||
step_number: Mapped[int] = mapped_column(Integer, nullable=False)
|
||||
thought: Mapped[str | None] = mapped_column(Text, nullable=True)
|
||||
action: Mapped[str | None] = mapped_column(String(255), nullable=True)
|
||||
action_input: Mapped[dict[str, Any] | None] = mapped_column(JSONB, nullable=True)
|
||||
observation: Mapped[str | None] = mapped_column(Text, nullable=True)
|
||||
cost_usd: Mapped[float] = mapped_column(Float, nullable=False, default=0.0)
|
||||
created_at: Mapped[datetime] = mapped_column(
|
||||
DateTime(timezone=True), nullable=False, server_default=func.now()
|
||||
)
|
||||
|
||||
|
||||
class AgentSubtask(Base, TenantMixin):
|
||||
"""A subtask delegated from one agent to another for multi-agent orchestration."""
|
||||
|
||||
|
||||
@@ -0,0 +1,481 @@
|
||||
"""Tests for the ReAct agent loop (app/ai/agent_loop.py).
|
||||
|
||||
All tests mock llm_complete and tool_registry — no real LLM or DB needed.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import uuid
|
||||
from dataclasses import dataclass
|
||||
from typing import Any
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
from app.ai.agent_loop import ReActResult, ReActStep, run_react_loop
|
||||
|
||||
|
||||
# ──────────────────────────────────────────────────────────────────────────
|
||||
# Test helpers
|
||||
# ──────────────────────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
@dataclass
|
||||
class MockAgentDefinition:
|
||||
"""Minimal stand-in for AgentDefinition used by run_react_loop."""
|
||||
|
||||
id: uuid.UUID
|
||||
llm_model: str = "gpt-4o"
|
||||
system_prompt: str = "You are a helpful assistant."
|
||||
api_key: str | None = None
|
||||
api_base: str | None = None
|
||||
provider: str | None = None
|
||||
max_tokens: int = 1000
|
||||
|
||||
|
||||
def _make_tool_call(
|
||||
call_id: str = "call_1",
|
||||
name: str = "search",
|
||||
arguments: dict[str, Any] | None = None,
|
||||
) -> dict[str, Any]:
|
||||
"""Build a mock LiteLLM tool-call object."""
|
||||
args_str = json.dumps(arguments or {})
|
||||
function_mock = MagicMock()
|
||||
function_mock.name = name
|
||||
function_mock.arguments = args_str
|
||||
tc = MagicMock()
|
||||
tc.id = call_id
|
||||
tc.function = function_mock
|
||||
return tc
|
||||
|
||||
|
||||
def _make_llm_response(
|
||||
content: str = "",
|
||||
tool_calls: list[Any] | None = None,
|
||||
cost_usd: float = 0.001,
|
||||
) -> dict[str, Any]:
|
||||
"""Build a mock llm_complete() return value."""
|
||||
message = MagicMock()
|
||||
message.content = content
|
||||
message.tool_calls = tool_calls or None
|
||||
|
||||
choice = MagicMock()
|
||||
choice.message = message
|
||||
|
||||
raw_response = MagicMock()
|
||||
raw_response.choices = [choice]
|
||||
|
||||
return {
|
||||
"content": content,
|
||||
"usage": {"total_tokens": 100},
|
||||
"cost_usd": cost_usd,
|
||||
"model": "gpt-4o",
|
||||
"raw_response": raw_response,
|
||||
}
|
||||
|
||||
|
||||
def _make_tool_registry(tools: dict[str, AsyncMock] | None = None) -> MagicMock:
|
||||
"""Build a mock ToolRegistry."""
|
||||
registry = MagicMock()
|
||||
_tools = tools or {}
|
||||
|
||||
def _get(name: str) -> Any:
|
||||
if name in _tools:
|
||||
tool = MagicMock()
|
||||
tool.handler = _tools[name]
|
||||
return tool
|
||||
return None
|
||||
|
||||
registry.get = _get
|
||||
return registry
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def agent_def() -> MockAgentDefinition:
|
||||
return MockAgentDefinition(id=uuid.uuid4())
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def tenant_id() -> uuid.UUID:
|
||||
return uuid.uuid4()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def user_id() -> uuid.UUID:
|
||||
return uuid.uuid4()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_db() -> AsyncMock:
|
||||
return AsyncMock()
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _mock_hooks():
|
||||
"""Patch do_action so the loop doesn't try to fire real hooks."""
|
||||
with patch("app.core.hooks.do_action", new_callable=AsyncMock):
|
||||
yield
|
||||
|
||||
|
||||
# ──────────────────────────────────────────────────────────────────────────
|
||||
# Tests
|
||||
# ──────────────────────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_basic_react_loop_no_tools(agent_def, tenant_id, user_id, mock_db):
|
||||
"""Single-step loop: LLM returns final response, no tool calls."""
|
||||
registry = _make_tool_registry()
|
||||
|
||||
with patch("app.ai.agent_loop.llm_complete", new_callable=AsyncMock) as mock_llm:
|
||||
mock_llm.return_value = _make_llm_response(
|
||||
content="Hello! How can I help you?",
|
||||
cost_usd=0.002,
|
||||
)
|
||||
|
||||
result = await run_react_loop(
|
||||
agent_definition=agent_def,
|
||||
messages=[{"role": "user", "content": "Hi"}],
|
||||
tools=[],
|
||||
tool_registry=registry,
|
||||
db=mock_db,
|
||||
tenant_id=tenant_id,
|
||||
user_id=user_id,
|
||||
)
|
||||
|
||||
assert result.status == "completed"
|
||||
assert result.final_content == "Hello! How can I help you?"
|
||||
assert result.steps_taken == 1
|
||||
assert len(result.steps) == 1
|
||||
assert result.steps[0].action is None # No tool call
|
||||
assert result.steps[0].thought == "Hello! How can I help you?"
|
||||
assert result.total_cost_usd == pytest.approx(0.002)
|
||||
mock_llm.assert_awaited_once()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_multi_step_loop_with_tool_calls(agent_def, tenant_id, user_id, mock_db):
|
||||
"""Multi-step loop: LLM calls a tool, gets result, then gives final answer."""
|
||||
search_handler = AsyncMock(return_value="Found 3 contacts: Alice, Bob, Charlie")
|
||||
registry = _make_tool_registry({"search": search_handler})
|
||||
|
||||
# Step 1: LLM calls search tool
|
||||
# Step 2: LLM gives final answer
|
||||
responses = [
|
||||
_make_llm_response(
|
||||
content="I'll search for contacts.",
|
||||
tool_calls=[_make_tool_call(name="search", arguments={"query": "contacts"})],
|
||||
cost_usd=0.003,
|
||||
),
|
||||
_make_llm_response(
|
||||
content="I found 3 contacts: Alice, Bob, and Charlie.",
|
||||
cost_usd=0.004,
|
||||
),
|
||||
]
|
||||
|
||||
with patch("app.ai.agent_loop.llm_complete", new_callable=AsyncMock) as mock_llm:
|
||||
mock_llm.side_effect = responses
|
||||
|
||||
result = await run_react_loop(
|
||||
agent_definition=agent_def,
|
||||
messages=[{"role": "user", "content": "Find contacts"}],
|
||||
tools=[{"type": "function", "function": {"name": "search"}}],
|
||||
tool_registry=registry,
|
||||
db=mock_db,
|
||||
tenant_id=tenant_id,
|
||||
user_id=user_id,
|
||||
)
|
||||
|
||||
assert result.status == "completed"
|
||||
assert result.steps_taken == 2
|
||||
assert len(result.steps) == 2
|
||||
assert result.steps[0].action == "search"
|
||||
assert result.steps[0].action_input == {"query": "contacts"}
|
||||
assert result.steps[0].observation == "Found 3 contacts: Alice, Bob, Charlie"
|
||||
assert result.steps[1].action is None # Final response
|
||||
assert result.final_content == "I found 3 contacts: Alice, Bob, and Charlie."
|
||||
assert result.total_cost_usd == pytest.approx(0.007)
|
||||
assert mock_llm.await_count == 2
|
||||
search_handler.assert_awaited_once()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_max_steps_limit(agent_def, tenant_id, user_id, mock_db):
|
||||
"""Loop stops gracefully when max_steps is reached."""
|
||||
search_handler = AsyncMock(return_value="result")
|
||||
registry = _make_tool_registry({"search": search_handler})
|
||||
|
||||
# Every response has a tool call — never gives final answer
|
||||
tool_call = _make_tool_call(name="search", arguments={})
|
||||
response = _make_llm_response(
|
||||
content="Searching...",
|
||||
tool_calls=[tool_call],
|
||||
cost_usd=0.001,
|
||||
)
|
||||
|
||||
with patch("app.ai.agent_loop.llm_complete", new_callable=AsyncMock) as mock_llm:
|
||||
mock_llm.return_value = response
|
||||
|
||||
result = await run_react_loop(
|
||||
agent_definition=agent_def,
|
||||
messages=[{"role": "user", "content": "Keep searching"}],
|
||||
tools=[{"type": "function", "function": {"name": "search"}}],
|
||||
tool_registry=registry,
|
||||
db=mock_db,
|
||||
tenant_id=tenant_id,
|
||||
user_id=user_id,
|
||||
max_steps=3,
|
||||
)
|
||||
|
||||
assert result.status == "stopped_max_steps"
|
||||
assert result.steps_taken == 3
|
||||
assert len(result.steps) == 3
|
||||
assert "max_steps" in (result.error or "")
|
||||
assert mock_llm.await_count == 3
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_timeout_graceful_stop(agent_def, tenant_id, user_id, mock_db):
|
||||
"""Loop stops gracefully when timeout is exceeded."""
|
||||
registry = _make_tool_registry()
|
||||
|
||||
# Simulate slow LLM responses that eventually exceed timeout
|
||||
async def slow_llm(**kwargs: Any) -> dict[str, Any]:
|
||||
await asyncio.sleep(0.5)
|
||||
return _make_llm_response(content="thinking...", tool_calls=None, cost_usd=0.001)
|
||||
|
||||
with patch("app.ai.agent_loop.llm_complete", new_callable=AsyncMock) as mock_llm:
|
||||
mock_llm.side_effect = slow_llm
|
||||
|
||||
result = await run_react_loop(
|
||||
agent_definition=agent_def,
|
||||
messages=[{"role": "user", "content": "Hi"}],
|
||||
tools=[],
|
||||
tool_registry=registry,
|
||||
db=mock_db,
|
||||
tenant_id=tenant_id,
|
||||
user_id=user_id,
|
||||
timeout_seconds=0, # Immediate timeout on first check
|
||||
)
|
||||
|
||||
assert result.status == "stopped_timeout"
|
||||
assert "timeout" in (result.error or "").lower()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_transient_error_retry(agent_def, tenant_id, user_id, mock_db):
|
||||
"""Transient errors are retried, then succeed."""
|
||||
registry = _make_tool_registry()
|
||||
|
||||
# First call raises transient error, second succeeds
|
||||
responses: list[Any] = [
|
||||
Exception("rate limit exceeded"),
|
||||
_make_llm_response(content="Success after retry", cost_usd=0.002),
|
||||
]
|
||||
|
||||
with patch("app.ai.agent_loop.llm_complete", new_callable=AsyncMock) as mock_llm:
|
||||
mock_llm.side_effect = responses
|
||||
with patch("app.ai.agent_loop.asyncio.sleep", new_callable=AsyncMock):
|
||||
result = await run_react_loop(
|
||||
agent_definition=agent_def,
|
||||
messages=[{"role": "user", "content": "Hi"}],
|
||||
tools=[],
|
||||
tool_registry=registry,
|
||||
db=mock_db,
|
||||
tenant_id=tenant_id,
|
||||
user_id=user_id,
|
||||
)
|
||||
|
||||
assert result.status == "completed"
|
||||
assert result.final_content == "Success after retry"
|
||||
assert mock_llm.await_count == 2 # First failed, second succeeded
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_permanent_error_stops(agent_def, tenant_id, user_id, mock_db):
|
||||
"""Permanent errors stop the loop immediately."""
|
||||
registry = _make_tool_registry()
|
||||
|
||||
with patch("app.ai.agent_loop.llm_complete", new_callable=AsyncMock) as mock_llm:
|
||||
mock_llm.side_effect = Exception("authentication error: invalid api key")
|
||||
|
||||
result = await run_react_loop(
|
||||
agent_definition=agent_def,
|
||||
messages=[{"role": "user", "content": "Hi"}],
|
||||
tools=[],
|
||||
tool_registry=registry,
|
||||
db=mock_db,
|
||||
tenant_id=tenant_id,
|
||||
user_id=user_id,
|
||||
)
|
||||
|
||||
assert result.status == "stopped_error"
|
||||
assert "Permanent error" in (result.error or "")
|
||||
assert mock_llm.await_count == 1 # No retry for permanent errors
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cost_accumulation(agent_def, tenant_id, user_id, mock_db):
|
||||
"""Cost is accumulated across multiple LLM calls."""
|
||||
search_handler = AsyncMock(return_value="result")
|
||||
registry = _make_tool_registry({"search": search_handler})
|
||||
|
||||
responses = [
|
||||
_make_llm_response(
|
||||
content="Step 1",
|
||||
tool_calls=[_make_tool_call(name="search")],
|
||||
cost_usd=0.01,
|
||||
),
|
||||
_make_llm_response(
|
||||
content="Step 2",
|
||||
tool_calls=[_make_tool_call(name="search")],
|
||||
cost_usd=0.02,
|
||||
),
|
||||
_make_llm_response(
|
||||
content="Final answer",
|
||||
cost_usd=0.03,
|
||||
),
|
||||
]
|
||||
|
||||
with patch("app.ai.agent_loop.llm_complete", new_callable=AsyncMock) as mock_llm:
|
||||
mock_llm.side_effect = responses
|
||||
|
||||
result = await run_react_loop(
|
||||
agent_definition=agent_def,
|
||||
messages=[{"role": "user", "content": "Search twice"}],
|
||||
tools=[{"type": "function", "function": {"name": "search"}}],
|
||||
tool_registry=registry,
|
||||
db=mock_db,
|
||||
tenant_id=tenant_id,
|
||||
user_id=user_id,
|
||||
)
|
||||
|
||||
assert result.status == "completed"
|
||||
assert result.total_cost_usd == pytest.approx(0.06) # 0.01 + 0.02 + 0.03
|
||||
assert result.steps_taken == 3
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_step_persistence_via_callback(agent_def, tenant_id, user_id, mock_db):
|
||||
"""Steps are passed to on_step callback for persistence."""
|
||||
registry = _make_tool_registry()
|
||||
collected_steps: list[ReActStep] = []
|
||||
|
||||
async def on_step(step: ReActStep) -> None:
|
||||
collected_steps.append(step)
|
||||
|
||||
with patch("app.ai.agent_loop.llm_complete", new_callable=AsyncMock) as mock_llm:
|
||||
mock_llm.return_value = _make_llm_response(content="Done", cost_usd=0.001)
|
||||
|
||||
result = await run_react_loop(
|
||||
agent_definition=agent_def,
|
||||
messages=[{"role": "user", "content": "Hi"}],
|
||||
tools=[],
|
||||
tool_registry=registry,
|
||||
db=mock_db,
|
||||
tenant_id=tenant_id,
|
||||
user_id=user_id,
|
||||
on_step=on_step,
|
||||
)
|
||||
|
||||
assert len(collected_steps) == 1
|
||||
assert collected_steps[0].step_number == 1
|
||||
assert collected_steps[0].thought == "Done"
|
||||
assert result.status == "completed"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_tool_not_found(agent_def, tenant_id, user_id, mock_db):
|
||||
"""When a tool is not found, the observation contains an error message."""
|
||||
registry = _make_tool_registry() # No tools registered
|
||||
|
||||
with patch("app.ai.agent_loop.llm_complete", new_callable=AsyncMock) as mock_llm:
|
||||
mock_llm.side_effect = [
|
||||
_make_llm_response(
|
||||
content="Calling unknown tool",
|
||||
tool_calls=[_make_tool_call(name="nonexistent", arguments={})],
|
||||
cost_usd=0.001,
|
||||
),
|
||||
_make_llm_response(content="OK", cost_usd=0.001),
|
||||
]
|
||||
|
||||
result = await run_react_loop(
|
||||
agent_definition=agent_def,
|
||||
messages=[{"role": "user", "content": "Call unknown tool"}],
|
||||
tools=[{"type": "function", "function": {"name": "nonexistent"}}],
|
||||
tool_registry=registry,
|
||||
db=mock_db,
|
||||
tenant_id=tenant_id,
|
||||
user_id=user_id,
|
||||
)
|
||||
|
||||
assert result.status == "completed"
|
||||
assert result.steps[0].action == "nonexistent"
|
||||
assert "not found" in (result.steps[0].observation or "").lower()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_multiple_tool_calls_per_step(agent_def, tenant_id, user_id, mock_db):
|
||||
"""Multiple tool calls in a single LLM response are all executed."""
|
||||
handler_a = AsyncMock(return_value="Result A")
|
||||
handler_b = AsyncMock(return_value="Result B")
|
||||
registry = _make_tool_registry({"tool_a": handler_a, "tool_b": handler_b})
|
||||
|
||||
with patch("app.ai.agent_loop.llm_complete", new_callable=AsyncMock) as mock_llm:
|
||||
mock_llm.side_effect = [
|
||||
_make_llm_response(
|
||||
content="Calling two tools",
|
||||
tool_calls=[
|
||||
_make_tool_call(call_id="c1", name="tool_a", arguments={}),
|
||||
_make_tool_call(call_id="c2", name="tool_b", arguments={}),
|
||||
],
|
||||
cost_usd=0.005,
|
||||
),
|
||||
_make_llm_response(content="Done with both", cost_usd=0.002),
|
||||
]
|
||||
|
||||
result = await run_react_loop(
|
||||
agent_definition=agent_def,
|
||||
messages=[{"role": "user", "content": "Call two tools"}],
|
||||
tools=[
|
||||
{"type": "function", "function": {"name": "tool_a"}},
|
||||
{"type": "function", "function": {"name": "tool_b"}},
|
||||
],
|
||||
tool_registry=registry,
|
||||
db=mock_db,
|
||||
tenant_id=tenant_id,
|
||||
user_id=user_id,
|
||||
)
|
||||
|
||||
assert result.status == "completed"
|
||||
assert result.steps_taken == 2
|
||||
# Step 1 has two tool calls → two step entries
|
||||
assert len(result.steps) == 3 # 2 tool steps + 1 final step
|
||||
handler_a.assert_awaited_once()
|
||||
handler_b.assert_awaited_once()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_react_result_dataclass_fields():
|
||||
"""ReActResult and ReActStep have correct default values."""
|
||||
step = ReActStep(
|
||||
step_number=1,
|
||||
thought="test",
|
||||
action=None,
|
||||
action_input=None,
|
||||
observation=None,
|
||||
cost_usd=0.0,
|
||||
timestamp="2025-01-01T00:00:00+00:00",
|
||||
)
|
||||
assert step.step_number == 1
|
||||
assert step.thought == "test"
|
||||
|
||||
result = ReActResult(final_content="hello")
|
||||
assert result.final_content == "hello"
|
||||
assert result.steps == []
|
||||
assert result.total_cost_usd == 0.0
|
||||
assert result.steps_taken == 0
|
||||
assert result.status == "completed"
|
||||
assert result.error is None
|
||||
Reference in New Issue
Block a user