271 lines
11 KiB
Python
271 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)
|
||
|
|
"""
|
||
|
|
|
||
|
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import asyncio
|
||
|
|
import logging
|
||
|
|
from datetime import UTC, datetime
|
||
|
|
from typing import Any
|
||
|
|
|
||
|
|
from sqlalchemy import func, select
|
||
|
|
|
||
|
|
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],
|
||
|
|
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, calls LLM via LiteLLM, executes tool calls via ToolRegistry,
|
||
|
|
saves result to AgentRun. Returns result dict.
|
||
|
|
|
||
|
|
Safety checks:
|
||
|
|
1. Rate limit: max_executions_per_hour
|
||
|
|
2. Max duration: max_duration_seconds (asyncio.timeout)
|
||
|
|
3. Infinite loop: same tool 5x consecutively
|
||
|
|
4. Budget limit: cumulative cost_usd
|
||
|
|
"""
|
||
|
|
from app.plugins.builtins.automation.models import AgentDefinition, AgentRun
|
||
|
|
|
||
|
|
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":
|
||
|
|
# Collect context data (recent contacts, mails, events)
|
||
|
|
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.core.cache import get_cached_mail_summary
|
||
|
|
context_data["recent_mails"] = get_cached_mail_summary(agent.tenant_id) or []
|
||
|
|
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:
|
||
|
|
# Reactive mode: use trigger_data as context
|
||
|
|
context_data = trigger_data or {}
|
||
|
|
|
||
|
|
# Call LLM via LiteLLM
|
||
|
|
result_data: dict[str, Any] = {
|
||
|
|
"agent_id": str(agent.id),
|
||
|
|
"agent_name": agent.name,
|
||
|
|
"trigger_type": trigger_type,
|
||
|
|
"status": "running",
|
||
|
|
"llm_response": None,
|
||
|
|
"tool_calls": [],
|
||
|
|
"cost_usd": 0.0,
|
||
|
|
"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."""
|
||
|
|
from app.ai.llm_client import LLMClient
|
||
|
|
|
||
|
|
llm = LLMClient(
|
||
|
|
model=agent.model or None,
|
||
|
|
api_key=agent.api_key or None,
|
||
|
|
api_base=agent.api_base or None,
|
||
|
|
provider=agent.provider or None,
|
||
|
|
)
|
||
|
|
|
||
|
|
# Build system prompt from agent configuration
|
||
|
|
system_prompt = agent.system_prompt or "You are a helpful AI assistant."
|
||
|
|
user_prompt = f"Context: {context_data}"
|
||
|
|
|
||
|
|
# Use litellm directly for more control
|
||
|
|
import litellm
|
||
|
|
|
||
|
|
litellm_model = agent.model or "gpt-4o"
|
||
|
|
if agent.provider and agent.provider != "openai":
|
||
|
|
litellm_model = f"{agent.provider}/{litellm_model}"
|
||
|
|
|
||
|
|
kwargs: dict[str, Any] = {
|
||
|
|
"model": litellm_model,
|
||
|
|
"messages": [
|
||
|
|
{"role": "system", "content": system_prompt},
|
||
|
|
{"role": "user", "content": user_prompt},
|
||
|
|
],
|
||
|
|
"temperature": 0.3,
|
||
|
|
"max_tokens": agent.max_tokens or 1000,
|
||
|
|
}
|
||
|
|
|
||
|
|
if agent.api_key:
|
||
|
|
kwargs["api_key"] = agent.api_key
|
||
|
|
if agent.api_base:
|
||
|
|
kwargs["api_base"] = agent.api_base
|
||
|
|
|
||
|
|
response = await litellm.acompletion(**kwargs)
|
||
|
|
content = response.choices[0].message.content
|
||
|
|
|
||
|
|
# Track cost
|
||
|
|
if hasattr(response, "usage") and response.usage:
|
||
|
|
usage = response.usage
|
||
|
|
# Estimate cost (simplified)
|
||
|
|
input_cost = (usage.prompt_tokens or 0) * 0.00001 / 1000
|
||
|
|
output_cost = (usage.completion_tokens or 0) * 0.00003 / 1000
|
||
|
|
result_data["cost_usd"] = round(input_cost + output_cost, 6)
|
||
|
|
|
||
|
|
result_data["llm_response"] = content
|
||
|
|
|
||
|
|
# Execute tool calls if LLM returned function calls
|
||
|
|
if hasattr(response.choices[0].message, "tool_calls") and response.choices[0].message.tool_calls:
|
||
|
|
from app.plugins.builtins.ai_assistant.tool_registry import get_tool_registry
|
||
|
|
|
||
|
|
registry = get_tool_registry()
|
||
|
|
tool_call_count: dict[str, int] = {}
|
||
|
|
for tc in 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 asyncio.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"
|
||
|
|
|
||
|
|
except Exception as e:
|
||
|
|
logger.exception("Agent run failed for %s", agent.id)
|
||
|
|
result_data["status"] = "failed"
|
||
|
|
result_data["error"] = str(e)
|
||
|
|
|
||
|
|
# 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
|