"""Hub Agent - Manages hub operations, transit, and capacity.""" import uuid from datetime import datetime, timedelta from typing import Dict, List, Any, Optional from dataclasses import dataclass from collections import deque from core.agent import SpecializedAgent from core.types import AgentTask, MessageType, Hub from core.logger import logger @dataclass class TransitRecord: record_id: str order_id: str entry_time: datetime expected_exit: datetime actual_exit: Optional[datetime] destination: str status: str # waiting, processing, in_transit, delivered items_count: int @dataclass class HubInventory: hub_id: str current_load: int capacity: int incoming_transits: int outgoing_transits: int pending_processing: int class HubAgent(SpecializedAgent): """Hub Agent - Manages hub operations and transit flow.""" def __init__(self): super().__init__( agent_id="HUB_AGENT", domain="hub_management", description="Manages hub operations, transit flow, and capacity" ) self._hubs = self._init_hubs() self._in_transit: Dict[str, List[TransitRecord]] = {h: [] for h in self._hubs.keys()} self._processing_queues: Dict[str, deque] = {h: deque() for h in self._hubs.keys()} self._metrics: Dict[str, Dict] = {} # BUG FIX: metrics init was dead code inside _init_hubs (after return). # Moved here so it actually runs after _hubs is populated. for hub_id in self._hubs.keys(): self._metrics[hub_id] = { "processed_today": 0, "average_processing_time": 15, "efficiency": 95.0, "bottlenecks": [], } def _init_hubs(self) -> Dict[str, Dict]: return { "DL-HUB-01": { "name": "Delhi North Hub", "location": {"lat": 28.6139, "lng": 77.2090}, "address": "Sector 12, Narela, Delhi", "capacity": 500, "current_load": 150, "spokes": ["DL-SP-01", "DL-SP-02", "DL-SP-03"], "connected_hubs": ["DL-HUB-02", "MU-HUB-01", "KL-HUB-01"], "processing_rate": 100, "operating_hours": "24/7", }, "DL-HUB-02": { "name": "Delhi South Hub", "location": {"lat": 28.5355, "lng": 77.2100}, "address": "Nehru Place, Delhi", "capacity": 400, "current_load": 200, "spokes": ["DL-SP-04", "DL-SP-05"], "connected_hubs": ["DL-HUB-01", "BL-HUB-01"], "processing_rate": 80, "operating_hours": "24/7", }, "MU-HUB-01": { "name": "Mumbai West Hub", "location": {"lat": 19.0760, "lng": 72.8777}, "address": "Andheri West, Mumbai", "capacity": 600, "current_load": 300, "spokes": ["MU-SP-01", "MU-SP-02", "MU-SP-03", "MU-SP-04"], "connected_hubs": ["MU-HUB-02", "DL-HUB-01", "PU-HUB-01"], "processing_rate": 120, "operating_hours": "24/7", }, "MU-HUB-02": { "name": "Mumbai East Hub", "location": {"lat": 19.1650, "lng": 72.8500}, "address": "Thane, Mumbai", "capacity": 550, "current_load": 250, "spokes": ["MU-SP-05", "MU-SP-06"], "connected_hubs": ["MU-HUB-01", "PU-HUB-01"], "processing_rate": 90, "operating_hours": "24/7", }, "BL-HUB-01": { "name": "Bangalore Hub", "location": {"lat": 12.9716, "lng": 77.5946}, "address": "Whitefield, Bangalore", "capacity": 500, "current_load": 180, "spokes": ["BL-SP-01", "BL-SP-02", "BL-SP-03"], "connected_hubs": ["HY-HUB-01", "DL-HUB-02"], "processing_rate": 100, "operating_hours": "24/7", }, "HY-HUB-01": { "name": "Hyderabad Hub", "location": {"lat": 17.3850, "lng": 78.4867}, "address": "Hi-Tech City, Hyderabad", "capacity": 450, "current_load": 220, "spokes": ["HY-SP-01", "HY-SP-02"], "connected_hubs": ["BL-HUB-01", "KL-HUB-01"], "processing_rate": 85, "operating_hours": "24/7", }, "PU-HUB-01": { "name": "Pune Hub", "location": {"lat": 18.5204, "lng": 73.8567}, "address": "Hinjewadi, Pune", "capacity": 400, "current_load": 160, "spokes": ["PU-SP-01", "PU-SP-02"], "connected_hubs": ["MU-HUB-01", "MU-HUB-02"], "processing_rate": 75, "operating_hours": "24/7", }, "KL-HUB-01": { "name": "Kolkata Hub", "location": {"lat": 22.5726, "lng": 88.3639}, "address": "Salt Lake, Kolkata", "capacity": 350, "current_load": 140, "spokes": ["KL-SP-01", "KL-SP-02"], "connected_hubs": ["DL-HUB-01", "HY-HUB-01"], "processing_rate": 70, "operating_hours": "24/7", }, } async def handle_task(self, task: AgentTask) -> Dict[str, Any]: handlers = { "prepare_receiving": self._prepare_receiving, "process_incoming": self._process_incoming, "dispatch_outgoing": self._dispatch_outgoing, "get_hub_status": self._get_hub_status, "transfer_between_hubs": self._transfer_between_hubs, "get_transit_status": self._get_transit_status, "optimize_hub_operations": self._optimize_hub_operations, "escalate_overflow": self._escalate_overflow, } handler = handlers.get(task.task_type, self._unknown_task) return await handler(task) async def _prepare_receiving(self, task: AgentTask) -> Dict[str, Any]: hub_id = task.data.get("hub_id") order_id = task.data.get("order_id") route_id = task.data.get("route_id") if hub_id not in self._hubs: return {"status": "error", "message": f"Hub {hub_id} not found"} hub = self._hubs[hub_id] available_capacity = hub["capacity"] - hub["current_load"] logger.info(f"Hub Agent: Preparing {hub['name']} to receive order {order_id}") if available_capacity < 10: await self.send_message( recipient="JARVIS", message_type=MessageType.EXCEPTION_DETECTED, payload={ "type": "capacity_warning", "hub_id": hub_id, "current_load": hub["current_load"], "capacity": hub["capacity"], "order_id": order_id, }, ) logger.warning(f"Hub {hub_id} near capacity: {hub['current_load']}/{hub['capacity']}") self._processing_queues[hub_id].append({ "order_id": order_id, "route_id": route_id, "arrival_time": datetime.now(), "priority": "normal", }) return { "status": "prepared", "hub_id": hub_id, "hub_name": hub["name"], "queue_position": len(self._processing_queues[hub_id]), "estimated_wait": len(self._processing_queues[hub_id]) * 5, } async def _process_incoming(self, task: AgentTask) -> Dict[str, Any]: hub_id = task.data.get("hub_id") order_id = task.data.get("order_id") if hub_id not in self._hubs: return {"status": "error", "message": f"Hub {hub_id} not found"} hub = self._hubs[hub_id] hub["current_load"] += 1 record_id = f"TR-{hub_id}-{uuid.uuid4().hex[:6].upper()}" record = TransitRecord( record_id=record_id, order_id=order_id, entry_time=datetime.now(), expected_exit=datetime.now() + timedelta(minutes=15), # 15 min avg processing actual_exit=None, destination=task.data.get("destination", "unknown"), status="processing", items_count=task.data.get("items_count", 1), ) self._in_transit[hub_id].append(record) self._metrics[hub_id]["processed_today"] += 1 logger.info(f"Order {order_id} received at {hub['name']} | load={hub['current_load']}/{hub['capacity']}") await self.send_message( recipient="CUSTOMER_AGENT", message_type=MessageType.HUB_STATUS_UPDATE, payload={ "order_id": order_id, "hub_id": hub_id, "hub_name": hub["name"], "status": "arrived_at_hub", "arrival_time": datetime.now().isoformat(), }, ) return {"status": "processed", "record_id": record_id, "hub_id": hub_id, "current_load": hub["current_load"]} async def _dispatch_outgoing(self, task: AgentTask) -> Dict[str, Any]: hub_id = task.data.get("hub_id") order_id = task.data.get("order_id") if hub_id not in self._hubs: return {"status": "error", "message": f"Hub {hub_id} not found"} hub = self._hubs[hub_id] record = None for r in self._in_transit[hub_id]: if r.order_id == order_id and r.status in ["processing", "waiting"]: r.actual_exit = datetime.now() r.status = "in_transit" record = r break if record: hub["current_load"] -= 1 logger.info(f"Order {order_id} dispatched from {hub['name']}") return { "status": "dispatched", "record_id": record.record_id, "hub_id": hub_id, "exit_time": record.actual_exit.isoformat(), "destination": record.destination, } return {"status": "error", "message": "Transit record not found"} async def _get_hub_status(self, task: AgentTask) -> Dict[str, Any]: hub_id = task.data.get("hub_id") if hub_id: if hub_id not in self._hubs: return {"status": "error", "message": f"Hub {hub_id} not found"} return self._format_hub_status(hub_id) return { "hubs": {hid: self._format_hub_status(hid) for hid in self._hubs.keys()}, "total_hubs": len(self._hubs), } def _format_hub_status(self, hub_id: str) -> Dict[str, Any]: hub = self._hubs[hub_id] metrics = self._metrics[hub_id] utilization = (hub["current_load"] / hub["capacity"]) * 100 if hub["capacity"] > 0 else 0 return { "hub_id": hub_id, "name": hub["name"], "location": hub["location"], "capacity": hub["capacity"], "current_load": hub["current_load"], "utilization_percent": round(utilization, 1), "spokes": hub["spokes"], "connected_hubs": hub["connected_hubs"], "processing_queue": len(self._processing_queues[hub_id]), "in_transit": len([r for r in self._in_transit[hub_id] if r.status == "in_transit"]), "metrics": { "processed_today": metrics["processed_today"], "avg_processing_time": metrics["average_processing_time"], "efficiency": metrics["efficiency"], }, "status": "operational" if utilization < 90 else "high_load", } async def _transfer_between_hubs(self, task: AgentTask) -> Dict[str, Any]: from_hub = task.data.get("from_hub") to_hub = task.data.get("to_hub") order_id = task.data.get("order_id") if from_hub not in self._hubs or to_hub not in self._hubs: return {"status": "error", "message": "Invalid hub ID(s)"} if to_hub not in self._hubs[from_hub]["connected_hubs"]: return { "status": "no_direct_route", "message": f"No direct route from {from_hub} to {to_hub}", "suggested_route": self._find_route(from_hub, to_hub), } transfer_id = f"XF-{uuid.uuid4().hex[:8].upper()}" logger.info(f"Hub transfer {transfer_id}: {self._hubs[from_hub]['name']} -> {self._hubs[to_hub]['name']} | order={order_id}") await self.send_message( recipient="ROUTE_OPTIMIZER", message_type=MessageType.AGENT_TASK, payload={ "task_type": "plan_inter_hub_route", "transfer_id": transfer_id, "from_hub": from_hub, "to_hub": to_hub, "order_id": order_id, }, ) return { "status": "transfer_initiated", "transfer_id": transfer_id, "from_hub": from_hub, "to_hub": to_hub, "estimated_time": 180, "route": [from_hub, to_hub], } async def _get_transit_status(self, task: AgentTask) -> Dict[str, Any]: order_id = task.data.get("order_id") transit_info = [] for hub_id, records in self._in_transit.items(): for record in records: if record.order_id == order_id: transit_info.append({ "record_id": record.record_id, "hub_id": hub_id, "hub_name": self._hubs[hub_id]["name"], "status": record.status, "entry_time": record.entry_time.isoformat(), "expected_exit": record.expected_exit.isoformat(), "actual_exit": record.actual_exit.isoformat() if record.actual_exit else None, "destination": record.destination, }) return { "order_id": order_id, "transit_records": transit_info, "current_location": transit_info[-1]["hub_name"] if transit_info else "unknown", } async def _optimize_hub_operations(self, task: AgentTask) -> Dict[str, Any]: hub_id = task.data.get("hub_id") if hub_id not in self._hubs: return {"status": "error", "message": f"Hub {hub_id} not found"} hub = self._hubs[hub_id] metrics = self._metrics[hub_id] recommendations = [] utilization = (hub["current_load"] / hub["capacity"]) * 100 if utilization > 80: recommendations.append({ "type": "capacity", "priority": "high", "message": f"High utilization at {utilization:.0f}% - consider overflow routing", "action": "request_additional_vehicle", }) queue_length = len(self._processing_queues[hub_id]) if queue_length > 20: recommendations.append({ "type": "queue", "priority": "medium", "message": f"Long processing queue ({queue_length} items)", "action": "increase_processing_capacity", }) if metrics["average_processing_time"] > 30: recommendations.append({ "type": "efficiency", "priority": "medium", "message": "Above average processing time", "action": "review_sorting_operations", }) logger.info(f"Hub {hub['name']}: util={utilization:.1f}% queue={queue_length} recommendations={len(recommendations)}") return { "hub_id": hub_id, "recommendations": recommendations, "potential_savings": { "time_minutes": len(recommendations) * 10, "cost_percent": len(recommendations) * 2, }, } async def _escalate_overflow(self, task: AgentTask) -> Dict[str, Any]: from_hub = task.data.get("from_hub") order_id = task.data.get("order_id") if from_hub not in self._hubs: return {"status": "error", "message": f"Hub {from_hub} not found"} hub = self._hubs[from_hub] alternatives = [] for connected in hub["connected_hubs"]: alt_hub = self._hubs[connected] utilization = (alt_hub["current_load"] / alt_hub["capacity"]) * 100 if utilization < 70: alternatives.append({ "hub_id": connected, "name": alt_hub["name"], "utilization": utilization, }) if not alternatives: return {"status": "no_alternatives", "message": "All connected hubs at capacity"} best_alt = alternatives[0] logger.warning(f"Hub overflow: routing to {best_alt['name']}") return await self._transfer_between_hubs(AgentTask( task_id=f"overflow_{order_id}", agent_type="hub", task_type="transfer_between_hubs", data={"from_hub": from_hub, "to_hub": best_alt["hub_id"], "order_id": order_id}, )) async def _unknown_task(self, task: AgentTask) -> Dict[str, Any]: return {"status": "error", "message": f"Unknown task: {task.task_type}"} def _find_route(self, from_hub: str, to_hub: str) -> List[str]: """BFS route between hubs.""" if from_hub == to_hub: return [from_hub] visited = {from_hub} queue = [(from_hub, [from_hub])] while queue: current, path = queue.pop(0) for neighbor in self._hubs[current]["connected_hubs"]: if neighbor == to_hub: return path + [to_hub] if neighbor not in visited: visited.add(neighbor) queue.append((neighbor, path + [neighbor])) return [] async def think(self, context: str, options: List[str] = None) -> str: return f"[HUB_AGENT reasoning]: {context}"