Agent Zero
457 lines
18 KiB
Python
457 lines
18 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
|
|
import uuid
|
|
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.ai_assistant.contracts import get_tool_registry
|
|
from app.plugins.builtins.automation.models import (
|
|
AgentDefinition,
|
|
AgentRun,
|
|
AgentRunStep,
|
|
)
|
|
|
|
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.contracts import get_contract
|
|
mail_contract = get_contract("mail")
|
|
if mail_contract and hasattr(mail_contract, "Mail"):
|
|
from sqlalchemy import select as _select
|
|
Mail = mail_contract.Mail
|
|
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 []
|
|
|
|
# ── Resolve agent permissions (Punkt 2+3 der Audit) ──
|
|
from app.ai.agent_permissions import resolve_agent_permissions
|
|
from app.ai.agent_tools import get_agent_tools
|
|
from app.ai.skill_registry import get_skill_registry
|
|
|
|
async with factory() as db:
|
|
perm_ctx = await resolve_agent_permissions(
|
|
db=db,
|
|
tenant_id=agent.tenant_id,
|
|
user_id=agent.created_by or uuid.uuid4(),
|
|
agent_definition=agent,
|
|
)
|
|
|
|
# Use permission-filtered tools instead of raw tool_ids
|
|
skill_reg = get_skill_registry()
|
|
tool_schemas, _skills = get_agent_tools(
|
|
agent_definition=agent,
|
|
tool_registry=registry,
|
|
skill_registry=skill_reg,
|
|
user_permissions=perm_ctx.user_permissions,
|
|
)
|
|
|
|
# ── 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
|
|
|
|
# ── Build agent context via context_builder (Punkt 1 der Audit) ──
|
|
from app.ai.context_builder import build_agent_context
|
|
|
|
# Sanitize context_data to remove sensitive fields (Punkt 4: data_policy)
|
|
from app.core.sensitive_data import sanitize_dict
|
|
safe_context_data = sanitize_dict(context_data)
|
|
|
|
# Build the user message from sanitized context
|
|
user_message = f"Context: {safe_context_data}" if safe_context_data else "No additional context provided."
|
|
|
|
# Build full message list (system prompt + context + user message)
|
|
messages = await build_agent_context(
|
|
agent_definition=agent,
|
|
user_message=user_message,
|
|
db=None, # No DB session available here; context_builder handles gracefully
|
|
tenant_id=agent.tenant_id,
|
|
user_id=agent.created_by or uuid_mod.uuid4(),
|
|
)
|
|
|
|
# ── Enforce data policy: filter sensitive fields from messages (Punkt 4) ──
|
|
from app.ai.data_policy import enforce_data_policy
|
|
messages = await enforce_data_policy(
|
|
db=None,
|
|
tenant_id=agent.tenant_id,
|
|
messages=messages,
|
|
agent_definition=agent,
|
|
)
|
|
|
|
react_result: ReActResult = await asyncio.wait_for(
|
|
run_react_loop(
|
|
agent_definition=agent,
|
|
messages=messages,
|
|
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,
|
|
require_approval=bool(getattr(agent, "require_approval", False)),
|
|
approval_tools=getattr(agent, "approval_tools", None),
|
|
),
|
|
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
|
|
|
|
# ── Mark result as AI-generated (Punkt 6: transparency) ──
|
|
from app.ai.transparency import mark_as_ai_generated
|
|
if react_result.final_content:
|
|
ai_metadata = mark_as_ai_generated(
|
|
react_result.final_content,
|
|
metadata={
|
|
"model": getattr(agent, "llm_model", "unknown"),
|
|
"provider": getattr(agent, "provider", "unknown"),
|
|
"agent_id": str(agent.id),
|
|
"agent_name": agent.name,
|
|
"run_id": str(run_id) if run_id else None,
|
|
},
|
|
)
|
|
result_data["ai_generated"] = True
|
|
result_data["ai_metadata"] = ai_metadata.get("ai_metadata", {})
|
|
|
|
# ── Create oversight decision record (Punkt 5: oversight) ──
|
|
from app.ai.oversight import DecisionRecord, create_decision_record
|
|
try:
|
|
async with factory() as db:
|
|
record = DecisionRecord(
|
|
agent_run_id=run_id or uuid_mod.uuid4(),
|
|
recommendation=react_result.final_content,
|
|
evidence={
|
|
"steps": len(react_result.steps),
|
|
"cost_usd": react_result.total_cost_usd,
|
|
"status": react_result.status,
|
|
},
|
|
)
|
|
await create_decision_record(db, agent.tenant_id, record)
|
|
await db.commit()
|
|
except Exception as e:
|
|
logger.warning("Failed to create oversight decision record: %s", e)
|
|
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,
|
|
)
|
|
|
|
# ── Post agent result to Communication (F-COMM) ──
|
|
try:
|
|
from app.plugins.builtins.contracts import get_contract_registry
|
|
komm = get_contract_registry().get("kommunikation")
|
|
if komm:
|
|
async with factory() as db:
|
|
# Find or create agent conversation room via contract
|
|
# (find_locked_room_id matches create_plugin_room semantics)
|
|
room_title = f"Agent: {agent.name}"
|
|
conv_id = await komm.find_locked_room_id(
|
|
db=db,
|
|
tenant_id=agent.tenant_id,
|
|
plugin_name="automation",
|
|
title=room_title,
|
|
)
|
|
if not conv_id:
|
|
room = await komm.create_plugin_room(
|
|
db=db,
|
|
tenant_id=agent.tenant_id,
|
|
user_id=agent.created_by,
|
|
plugin_name="automation",
|
|
title=room_title,
|
|
participant_type="agent",
|
|
)
|
|
conv_id = uuid.UUID(room["conversation_id"])
|
|
|
|
# Post result as message with action_card block
|
|
status = result_data.get("status", "unknown")
|
|
result_text = result_data.get("llm_response", result_data.get("error", "No result"))
|
|
await komm.send_message(
|
|
db=db,
|
|
tenant_id=agent.tenant_id,
|
|
conversation_id=conv_id,
|
|
sender_id=agent.id,
|
|
sender_type="agent",
|
|
content=f"Agent '{agent.name}' completed with status: {status}",
|
|
content_format="text",
|
|
blocks=[
|
|
{
|
|
"block_type": "action_card",
|
|
"block_data": {
|
|
"title": f"Agent Result: {agent.name}",
|
|
"description": result_text[:500] if result_text else "No result",
|
|
"actions": [
|
|
{"label": "View Details", "action": "view_agent_run", "data": {"run_id": str(run_id)}},
|
|
],
|
|
},
|
|
"sort_order": 0,
|
|
}
|
|
],
|
|
metadata={"agent_id": str(agent.id), "run_id": str(run_id), "status": status},
|
|
)
|
|
await db.commit()
|
|
logger.info("Posted agent result to conversation %s", conv_id)
|
|
except Exception as e:
|
|
logger.warning("Failed to post agent result to communication: %s", e)
|
|
|
|
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)
|