"""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:]