File size: 3,678 Bytes
a67fbf4 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 | import asyncio
import time
from datetime import datetime, timezone
from typing import Any, Dict, Optional
from .base import AgentContext, AgentResult, AgentStatus, BaseAgent
from src.config import settings
from src.telemetry import collector
class AgentExecutor:
def __init__(self):
self._semaphore = asyncio.Semaphore(settings.max_concurrent_agents)
self._active_runs: Dict[str, asyncio.Task] = {}
async def run(
self,
agent: BaseAgent,
context: Optional[AgentContext] = None,
timeout: Optional[int] = None,
**kwargs,
) -> AgentResult:
if context is None:
context = AgentContext(agent_id=agent.agent_id)
timeout = timeout or settings.agent_timeout
started_at = datetime.now(timezone.utc)
t0 = time.monotonic()
result = AgentResult(
run_id=context.run_id,
agent_id=agent.agent_id,
status=AgentStatus.PENDING,
started_at=started_at,
)
async with self._semaphore:
result.status = AgentStatus.RUNNING
collector.agent_start(agent.agent_id, context.run_id, agent.agent_type)
try:
await agent.on_start(context)
output = await asyncio.wait_for(
agent.execute(context, **kwargs),
timeout=timeout,
)
result.output = output
result.status = AgentStatus.COMPLETED
except asyncio.TimeoutError:
result.status = AgentStatus.TIMEOUT
result.error = f"Agent timed out after {timeout}s"
collector.agent_failed(agent.agent_id, context.run_id, result.error)
await agent.on_error(context, TimeoutError(result.error))
except asyncio.CancelledError:
result.status = AgentStatus.CANCELLED
result.error = "Agent execution was cancelled"
raise
except Exception as exc:
result.status = AgentStatus.FAILED
result.error = str(exc)
collector.agent_failed(agent.agent_id, context.run_id, result.error)
await agent.on_error(context, exc)
finally:
result.completed_at = datetime.now(timezone.utc)
result.duration_ms = (time.monotonic() - t0) * 1000
if result.status == AgentStatus.COMPLETED:
collector.agent_complete(agent.agent_id, context.run_id, result.duration_ms)
await agent.on_complete(context, result)
return result
async def run_with_retry(
self,
agent: BaseAgent,
context: Optional[AgentContext] = None,
retries: Optional[int] = None,
**kwargs,
) -> AgentResult:
max_retries = retries if retries is not None else settings.retry_limit
last_result: Optional[AgentResult] = None
for attempt in range(max_retries + 1):
ctx = context or AgentContext(agent_id=agent.agent_id)
if attempt > 0:
ctx.metadata["retry_attempt"] = attempt
last_result = await self.run(agent, ctx, **kwargs)
if last_result.status == AgentStatus.COMPLETED:
return last_result
if last_result.status == AgentStatus.CANCELLED:
break
return
async def cancel(self, run_id: str) -> bool:
task = self._active_runs.get(run_id)
if task and not task.done():
task.cancel()
return True
return False
executor = AgentExecutor()
|