Files
AI_engine/core/agent.py
Suriyakumarvijayanayagam 4283c602f6 fix(dispatch): liveness-aware coverage check, escalation rate limit, real alert sink
Prod findings (2026-09-22): a backend retry sweep failed 1,001 bookings in
60s; the agent called each a "coverage gap" because GEORADIUS found a miler
in the geo index (last seen in June), then sent 1,001 ops_alert tasks to
CUSTOMER_AGENT, which has no such handler.

- _find_zone filters GEORADIUS candidates by the backend's miler_status:<id>
  key; only status=Available counts. Facts now carry
  nearest_available_miler_within_km plus milers_in_geo_index_within_30km so
  the decision can separate "no riders here" from "riders exist, none on duty".
- Rate limit per zone per day: after the first alert, further failures only
  bump the counter (no LLM call); a summary re-alert goes out every
  DISPATCH_REALERT_EVERY (default 100).
- _ops_alert / _escalate_dispatch send EXCEPTION_DETECTED to JARVIS (the path
  that is actually handled); customer delay notice uses CUSTOMER_AGENT's real
  send_notification contract.
- JARVIS: escalation inbox (_escalations, pending_escalations()) and
  human_review/ops_alert task types are recorded instead of dropped.
- ExceptionAgent pull loops: also catch asyncio.TimeoutError (distinct from
  nats.errors.TimeoutError on 3.11) and log the exception type — the blank
  "pull loop error:" lines.
- Prompt + eval cases updated for the renamed facts; new case for the
  observed index-full/nobody-on-duty pattern. Tests for liveness filtering,
  burst suppression, fallback heuristic, sinks, and the JARVIS inbox.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_012AJLYcbTHCe45fyFnMfEin
2026-09-22 16:26:27 +05:30

370 lines
14 KiB
Python

"""Base Agent class - All agents inherit from this."""
import asyncio
import time
import uuid
from datetime import datetime
from typing import Dict, List, Optional, Any
from abc import ABC, abstractmethod
from core.types import AgentMessage, MessageType, AgentState, AgentTask, Priority
from core.message_bus import message_bus
from core.logger import logger
class Agent(ABC):
"""
Base class for all AI agents in the logistics system.
Each agent has its own task queue and can communicate via the message bus.
"""
def __init__(self, agent_id: str, agent_type: str, description: str = ""):
self.agent_id = agent_id
self.agent_type = agent_type
self.description = description
self.state = AgentState(
agent_id=agent_id,
agent_type=agent_type,
status="idle"
)
self._task_queue: asyncio.Queue = asyncio.Queue()
self._running = False
self._task_handlers: Dict[str, callable] = {}
self._last_state_emit = 0.0
# Register with message bus
message_bus.register_agent(self)
async def start(self):
"""Start the agent's processing loop."""
self._running = True
self.state.status = "idle"
logger.info(f"Agent started: {self.agent_id}")
while self._running:
try:
try:
task = await asyncio.wait_for(
self._task_queue.get(),
timeout=1.0
)
await self._process_task(task)
except asyncio.TimeoutError:
await self._heartbeat()
if time.monotonic() - self._last_state_emit >= 5.0:
await self._emit_state()
except Exception as e:
logger.error(f"Error in agent {self.agent_id}: {e}")
self.state.status = "error"
await asyncio.sleep(5)
async def stop(self):
"""Stop the agent."""
self._running = False
self.state.status = "stopped"
logger.info(f"Agent stopped: {self.agent_id}")
async def _process_task(self, task: AgentTask):
"""Process a task from the queue."""
self.state.status = "working"
self.state.current_task = task.task_id
started = time.monotonic()
await self._emit_state()
try:
if task.task_type in self._task_handlers:
handler = self._task_handlers[task.task_type]
result = await handler(task)
task.status = "completed"
task.result = result
else:
result = await self.handle_task(task)
task.status = "completed"
task.result = result
self.state.tasks_completed += 1
except Exception as e:
logger.error(f"Task {task.task_id} failed in {self.agent_id}: {e}")
task.status = "failed"
task.error = str(e)
self.state.tasks_failed += 1
finally:
duration_ms = int((time.monotonic() - started) * 1000)
await message_bus.publish_telemetry("task", {
"agent_id": self.agent_id,
"task_id": task.task_id,
"task_type": task.task_type,
"status": task.status,
"error": task.error,
"duration_ms": duration_ms,
})
self.state.current_task = None
self.state.last_active = datetime.now()
self.state.status = "idle"
await self._emit_state()
@abstractmethod
async def handle_task(self, task: AgentTask) -> Dict[str, Any]:
"""Override this method to handle tasks. Return a dict with results."""
pass
async def _heartbeat(self):
"""Called periodically when idle. Override for custom behavior."""
pass
async def _emit_state(self):
"""Publish this agent's live state as telemetry (for the command center)."""
self._last_state_emit = time.monotonic()
await message_bus.publish_telemetry("agent", {
"agent_id": self.agent_id,
"agent_type": self.agent_type,
"status": self.state.status,
"current_task": self.state.current_task,
"tasks_completed": self.state.tasks_completed,
"tasks_failed": self.state.tasks_failed,
})
def register_task_handler(self, task_type: str, handler: callable):
self._task_handlers[task_type] = handler
async def submit_task(self, task: AgentTask):
await self._task_queue.put(task)
async def send_message(self, recipient: str, message_type: MessageType, payload: Dict[str, Any], correlation_id: Optional[str] = None) -> str:
return await message_bus.send_to_agent(
sender=self.agent_id,
recipient=recipient,
message_type=message_type,
payload=payload,
correlation_id=correlation_id
)
async def broadcast(self, message_type: MessageType, payload: Dict[str, Any], correlation_id: Optional[str] = None) -> str:
return await message_bus.broadcast(
sender=self.agent_id,
message_type=message_type,
payload=payload,
correlation_id=correlation_id
)
async def receive_messages(self) -> List[AgentMessage]:
return await message_bus.get_messages(self.agent_id)
async def deliver(self, message: AgentMessage):
"""Consume an inbound directed message.
This is what makes agent-to-agent messaging actually work: the message
bus calls it for every directed message addressed to this agent. Messages
that carry a ``task_type`` in their payload are enqueued as tasks so
``handle_task`` processes them exactly like a submitted task; everything
else is handed to ``handle_message`` so the agent can react. Without this,
directed messages sat unread in the bus queue forever.
"""
payload = message.payload if isinstance(message.payload, dict) else {}
task_type = payload.get("task_type")
if task_type:
await self._task_queue.put(AgentTask(
task_id=f"{message.correlation_id or message.message_id}:{task_type}",
agent_type=self.agent_type,
task_type=task_type,
data=payload,
))
else:
try:
await self.handle_message(message)
except Exception as e:
logger.error(f"{self.agent_id} handle_message error: {e}")
async def handle_message(self, message: AgentMessage):
"""React to a non-task directed message (a notification). Default just
logs it; agents override this to act on events like EXCEPTION_DETECTED."""
logger.debug(
f"{self.agent_id} received {message.message_type.value} "
f"from {message.sender} (no task_type; not handled)"
)
def subscribe_to(self, message_type: MessageType, callback: callable):
message_bus.subscribe(message_type, callback)
async def think(self, context: str, options: List[str] = None) -> str:
return context
class MasterAgent(Agent):
"""
Master Agent (JARVIS) - Orchestrates all other agents.
Makes high-level decisions and assigns work.
"""
def __init__(self):
super().__init__(
agent_id="JARVIS",
agent_type="master",
description="Central orchestrator for all logistics operations"
)
self._sub_agents: Dict[str, Agent] = {}
self._active_orders: Dict[str, Dict] = {}
self._decision_log: List[Dict] = []
# Human-review inbox: every escalation/ops alert a sub-agent raised.
# The only sink for proposals the agents aren't authorised to act on;
# surfaced by pending_escalations() for the Command Center.
self._escalations: List[Dict] = []
def register_sub_agent(self, agent: Agent):
self._sub_agents[agent.agent_id] = agent
logger.info(f"JARVIS registered sub-agent: {agent.agent_id}")
async def handle_message(self, message: AgentMessage):
"""Surface notifications from sub-agents. Escalations (proposals a human
must decide) are logged at WARNING and recorded so they are visible
rather than silently dropped — the endpoint of the human-review path."""
if message.message_type == MessageType.EXCEPTION_DETECTED:
self._record_escalation(message.sender, message.payload)
else:
logger.info(f"JARVIS: {message.message_type.value} from {message.sender}")
self._decision_log.append({
"timestamp": datetime.now(),
"action": "received_notification",
"from": message.sender,
"type": message.message_type.value,
"payload": message.payload,
})
if len(self._decision_log) > 500:
self._decision_log = self._decision_log[-500:]
async def handle_task(self, task: AgentTask) -> Dict[str, Any]:
if task.task_type == "orchestrate_order":
return await self._orchestrate_order(task)
elif task.task_type == "monitor_operations":
return await self._monitor_operations(task)
elif task.task_type == "handle_exception":
return await self._handle_exception(task)
elif task.task_type == "generate_report":
return await self._generate_report(task)
elif task.task_type in ("human_review", "ops_alert"):
# Escalations sent as tasks land here instead of dead-lettering.
self._record_escalation(task.data.get("source", "unknown"), task.data)
return {"status": "recorded", "escalations_pending": len(self._escalations)}
else:
logger.warning(f"JARVIS: unknown task type {task.task_type!r} — dropped")
return {"status": "unknown_task", "task_type": task.task_type}
def _record_escalation(self, sender: str, payload: Dict[str, Any]):
logger.warning(f"JARVIS: escalation from {sender}: {payload}")
self._escalations.append({
"timestamp": datetime.now(),
"from": sender,
"exception_type": payload.get("exception_type"),
"severity": payload.get("severity"),
"payload": payload,
})
if len(self._escalations) > 500:
self._escalations = self._escalations[-500:]
def pending_escalations(self, limit: int = 50) -> List[Dict[str, Any]]:
"""Most recent escalations first — the human-review inbox."""
return list(reversed(self._escalations[-limit:]))
async def _orchestrate_order(self, task: AgentTask) -> Dict[str, Any]:
order_data = task.data.get("order", {})
order_id = order_data.get("order_id", "unknown")
logger.info(f"JARVIS: Orchestrating order {order_id}")
order_agent = self._sub_agents.get("ORDER_AGENT")
if order_agent:
await order_agent.submit_task(AgentTask(
task_id=f"{order_id}_validate",
agent_type="order",
task_type="validate_order",
data={"order": order_data}
))
customer_agent = self._sub_agents.get("CUSTOMER_AGENT")
if customer_agent:
await customer_agent.submit_task(AgentTask(
task_id=f"{order_id}_notify",
agent_type="customer",
task_type="send_confirmation",
data={"order": order_data}
))
self._decision_log.append({
"timestamp": datetime.now(),
"action": "orchestrate_order",
"order_id": order_id,
"agents_involved": list(self._sub_agents.keys())
})
return {
"status": "orchestrated",
"order_id": order_id,
"assigned_agents": list(self._sub_agents.keys())
}
async def _monitor_operations(self, task: AgentTask) -> Dict[str, Any]:
status_report = {
"timestamp": datetime.now().isoformat(),
"total_agents": len(self._sub_agents),
"agent_statuses": {},
"active_orders": len(self._active_orders),
"recent_decisions": self._decision_log[-10:]
}
for agent_id, agent in self._sub_agents.items():
status_report["agent_statuses"][agent_id] = {
"status": agent.state.status,
"tasks_completed": agent.state.tasks_completed,
"tasks_failed": agent.state.tasks_failed,
"current_task": agent.state.current_task
}
return status_report
async def _handle_exception(self, task: AgentTask) -> Dict[str, Any]:
exception_type = task.data.get("type", "unknown")
logger.warning(f"JARVIS: Handling exception - {exception_type}")
exception_agent = self._sub_agents.get("EXCEPTION_AGENT")
if exception_agent:
await exception_agent.submit_task(task)
return {"status": "delegated", "exception_type": exception_type}
return {"status": "no_exception_agent"}
async def _generate_report(self, task: AgentTask) -> Dict[str, Any]:
return {
"report_type": task.data.get("type", "summary"),
"generated_at": datetime.now().isoformat(),
"agents": [a.state for a in self._sub_agents.values()],
"decisions": self._decision_log[-50:]
}
async def _heartbeat(self):
for agent_id, agent in self._sub_agents.items():
if agent.state.status == "error":
logger.warning(f"Agent {agent_id} in error state")
class SpecializedAgent(Agent):
"""Base class for specialized domain agents."""
def __init__(self, agent_id: str, domain: str, description: str):
super().__init__(agent_id, domain, description)
self._domain_knowledge: Dict[str, Any] = {}
self._learning_history: List[Dict] = []
async def think(self, context: str, options: List[str] = None) -> str:
return f"[{self.agent_type} reasoning]: {context}"
def learn_from(self, experience: Dict[str, Any]):
self._learning_history.append({
"timestamp": datetime.now(),
"experience": experience
})
if len(self._learning_history) > 100:
self._learning_history = self._learning_history[-100:]