Agent Zero
c760b5961c
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
309 lines
11 KiB
Python
309 lines
11 KiB
Python
"""Agent runner — executes AI agents as ARQ jobs with safety checks.
|
|
|
|
Safety features:
|
|
- Max executions per agent per hour (rate limiting)
|
|
- 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 logging
|
|
from datetime import UTC, datetime
|
|
from typing import Any
|
|
|
|
from sqlalchemy import func, select
|
|
|
|
from app.ai.agent_loop import ReActResult, run_react_loop
|
|
from app.core.db import get_session_factory
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
async def run_agent(
|
|
ctx: dict[str, Any],
|
|
agent_id: str,
|
|
trigger_type: str = "manual",
|
|
trigger_data: dict[str, Any] | None = None,
|
|
) -> dict[str, Any]:
|
|
"""ARQ job function. Loads AgentDefinition from DB, checks rate limits,
|
|
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 (handled in ReAct loop)
|
|
4. Budget limit: cumulative cost_usd
|
|
"""
|
|
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()
|
|
|
|
# Load agent definition
|
|
async with factory() as db:
|
|
result = await db.execute(
|
|
select(AgentDefinition).where(AgentDefinition.id == agent_id)
|
|
)
|
|
agent = result.scalar_one_or_none()
|
|
|
|
if agent is None:
|
|
logger.error("AgentDefinition %s not found", agent_id)
|
|
return {"error": f"AgentDefinition {agent_id} not found", "status": "failed"}
|
|
|
|
if not agent.is_active:
|
|
logger.warning("AgentDefinition %s is inactive", agent_id)
|
|
return {"error": "Agent is inactive", "status": "skipped"}
|
|
|
|
# ── Safety Check 1: Rate Limit ──
|
|
if agent.max_executions_per_hour:
|
|
async with factory() as db:
|
|
one_hour_ago = datetime.now(UTC)
|
|
count_result = await db.execute(
|
|
select(func.count())
|
|
.select_from(AgentRun)
|
|
.where(
|
|
AgentRun.agent_id == agent.id,
|
|
AgentRun.created_at >= one_hour_ago,
|
|
)
|
|
)
|
|
recent_runs = count_result.scalar() or 0
|
|
if recent_runs >= agent.max_executions_per_hour:
|
|
logger.warning(
|
|
"Rate limit hit for agent %s: %d runs in last hour (max %d)",
|
|
agent.id, recent_runs, agent.max_executions_per_hour,
|
|
)
|
|
return {"error": "Rate limit exceeded", "status": "rate_limited"}
|
|
|
|
# ── Safety Check 2: Budget Limit ──
|
|
if agent.budget_limit_usd > 0:
|
|
async with factory() as db:
|
|
cost_result = await db.execute(
|
|
select(func.coalesce(func.sum(AgentRun.cost_usd), 0.0))
|
|
.where(AgentRun.agent_id == agent.id)
|
|
)
|
|
total_cost = float(cost_result.scalar() or 0.0)
|
|
if total_cost >= agent.budget_limit_usd:
|
|
logger.warning(
|
|
"Budget limit hit for agent %s: $%.4f total cost (limit $%.2f)",
|
|
agent.id, total_cost, agent.budget_limit_usd,
|
|
)
|
|
return {"error": "Budget limit exceeded", "status": "budget_exceeded"}
|
|
|
|
# Gather context
|
|
context_data: dict[str, Any] = {}
|
|
if trigger_type == "proactive" or agent.mode == "proactive":
|
|
try:
|
|
from app.services.contact_service import list_contacts
|
|
async with factory() as db:
|
|
contacts = await list_contacts(db, agent.tenant_id, page=1, page_size=10)
|
|
context_data["recent_contacts"] = contacts.get("items", [])
|
|
except Exception:
|
|
logger.warning("Failed to collect contacts for proactive context")
|
|
|
|
try:
|
|
from app.plugins.builtins.mail.models import Mail
|
|
from sqlalchemy import select as _select
|
|
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")
|
|
|
|
try:
|
|
from app.services.workflow_service import list_instances
|
|
async with factory() as db:
|
|
events = await list_instances(db, agent.tenant_id, page=1, page_size=10)
|
|
context_data["recent_events"] = events.get("items", [])
|
|
except Exception:
|
|
logger.warning("Failed to collect events for proactive context")
|
|
else:
|
|
context_data = trigger_data or {}
|
|
|
|
# ── Lifecycle: before run ──
|
|
from app.core.hooks import do_action
|
|
await do_action("agent.before_run", agent_id=str(agent.id), tenant_id=str(agent.tenant_id), trigger_type=trigger_type)
|
|
from app.core.outbox import enqueue_outbox_event
|
|
async with factory() as db:
|
|
await enqueue_outbox_event(
|
|
db,
|
|
agent.tenant_id,
|
|
'agent.run_started',
|
|
{'agent_id': str(agent.id), 'tenant_id': str(agent.tenant_id), 'trigger_type': trigger_type},
|
|
aggregate_type='agent',
|
|
aggregate_id=agent.id,
|
|
)
|
|
await db.commit()
|
|
|
|
# ── 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,
|
|
}
|
|
|
|
max_duration = agent.max_duration_seconds or 300
|
|
|
|
try:
|
|
import asyncio
|
|
import uuid as uuid_mod
|
|
|
|
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_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"] = "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 ──
|
|
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),
|
|
},
|
|
aggregate_type='agent',
|
|
aggregate_id=agent.id,
|
|
)
|
|
await db.commit()
|
|
|
|
return result_data
|
|
|
|
|
|
# Register all job functions with the job registry
|
|
from app.core.job_registry import register_job # noqa: E402
|
|
|
|
register_job("run_agent", run_agent)
|