"""Dispatch Agent - Handles zone analysis, route assignment, and order dispatching.""" import asyncio import uuid from datetime import datetime, timedelta from typing import Dict, List, Any, Optional from dataclasses import dataclass import redis.asyncio as aioredis from core.agent import SpecializedAgent from core.types import AgentTask, MessageType, Priority, ZoneType from core.logger import logger from core.http_client import api_post from config.system_config import REDIS_HOST, REDIS_PORT, REDIS_PASSWORD, GO_API_BASE_URL, INTERNAL_API_KEY @dataclass class ZoneInfo: zone_id: str zone_name: str pincode_range: tuple hub_id: str vehicle_types: List[str] avg_delivery_time: int current_load: int = 0 max_capacity: int = 100 @dataclass class RouteAssignment: route_id: str order_id: str from_hub: str to_hub: str vehicle_id: Optional[str] estimated_pickup: datetime estimated_delivery: datetime route_type: str class DispatchAgent(SpecializedAgent): """ Dispatch Agent - Orchestrates order dispatching and route assignment. Uses Redis GEO (GEORADIUS on milers:locations) to find nearest milers for a given pickup/delivery coordinate. Falls back to pincode-based zone lookup when no milers are found within range. """ def __init__(self): super().__init__( agent_id="DISPATCH_AGENT", domain="dispatch", description="Handles zone analysis, route assignment, and dispatch coordination" ) self._redis = aioredis.Redis( host=REDIS_HOST, port=REDIS_PORT, password=REDIS_PASSWORD, decode_responses=True, ) self._zones = self._init_zones() self._hubs = self._init_hubs() self._active_routes: Dict[str, RouteAssignment] = {} self._dispatch_queue: asyncio.Queue = asyncio.Queue() def _init_zones(self) -> Dict[str, ZoneInfo]: return { "north_delhi": ZoneInfo("north_delhi", "North Delhi", ("100", "199"), "DL-HUB-01", ["bike", "scooter", "van"], 45), "south_delhi": ZoneInfo("south_delhi", "South Delhi", ("200", "299"), "DL-HUB-02", ["bike", "scooter", "van"], 50), "mumbai_west": ZoneInfo("mumbai_west", "West Mumbai", ("400", "449"), "MU-HUB-01", ["bike", "scooter", "van", "truck"], 60), "mumbai_east": ZoneInfo("mumbai_east", "East Mumbai", ("450", "499"), "MU-HUB-02", ["bike", "scooter", "van", "truck"], 55), "bangalore": ZoneInfo("bangalore", "Bangalore", ("560", "562"), "BL-HUB-01", ["bike", "scooter", "van"], 40), "hyderabad": ZoneInfo("hyderabad", "Hyderabad", ("500", "509"), "HY-HUB-01", ["bike", "scooter", "van", "truck"], 50), "pune": ZoneInfo("pune", "Pune", ("411", "415"), "PU-HUB-01", ["bike", "scooter", "van"], 45), "kolkata": ZoneInfo("kolkata", "Kolkata", ("700", "700"), "KL-HUB-01", ["bike", "scooter", "van"], 55), } def _init_hubs(self) -> Dict[str, Dict]: return { "DL-HUB-01": {"name": "Delhi North Hub", "capacity": 500, "current_load": 150}, "DL-HUB-02": {"name": "Delhi South Hub", "capacity": 400, "current_load": 200}, "MU-HUB-01": {"name": "Mumbai West Hub", "capacity": 600, "current_load": 300}, "MU-HUB-02": {"name": "Mumbai East Hub", "capacity": 550, "current_load": 250}, "BL-HUB-01": {"name": "Bangalore Hub", "capacity": 500, "current_load": 180}, "HY-HUB-01": {"name": "Hyderabad Hub", "capacity": 450, "current_load": 220}, "PU-HUB-01": {"name": "Pune Hub", "capacity": 400, "current_load": 160}, "KL-HUB-01": {"name": "Kolkata Hub", "capacity": 350, "current_load": 140}, } async def handle_task(self, task: AgentTask) -> Dict[str, Any]: handlers = { "analyze_and_assign": self._analyze_and_assign, "find_route": self._find_route, "optimize_routes": self._optimize_routes, "assign_vehicle": self._assign_vehicle, "get_zone_status": self._get_zone_status, "reschedule_route": self._reschedule_route, "cancel_route": self._cancel_route, } handler = handlers.get(task.task_type, self._unknown_task) return await handler(task) # ------------------------------------------------------------------ # # Redis GEO zone lookup # # ------------------------------------------------------------------ # async def _find_zone(self, lat: float, lon: float) -> Optional[Dict[str, Any]]: """Find the nearest miler using GEORADIUS on milers:locations.""" try: results = await self._redis.georadius( "milers:locations", lon, lat, 10, "km", sort="ASC", count=5, ) if not results: return None nearest_miler = results[0] miler_info = await self._redis.hgetall(f"miler:{nearest_miler}") return { "miler_id": nearest_miler, "hub_id": miler_info.get("hub_id", "UNKNOWN"), "zone_id": miler_info.get("zone_id", "unknown"), "avg_delivery_time": int(miler_info.get("avg_delivery_time", 60)), } except Exception as e: logger.warning(f"Redis GEORADIUS error: {e}") return None def _find_zone_by_pincode(self, pincode_prefix: str) -> Optional[ZoneInfo]: for zone in self._zones.values(): start, end = zone.pincode_range if start <= pincode_prefix <= end: return zone return None # ------------------------------------------------------------------ # # Core dispatch logic # # ------------------------------------------------------------------ # async def _analyze_and_assign(self, task: AgentTask) -> Dict[str, Any]: order_data = task.data.get("order", {}) order_id = order_data.get("order_id", "unknown") logger.info(f"Dispatch Agent: Analyzing order {order_id}") pickup_addr = order_data.get("pickup_address", {}) delivery_addr = order_data.get("delivery_address", {}) pickup_lat = pickup_addr.get("lat") or pickup_addr.get("latitude") pickup_lon = pickup_addr.get("lng") or pickup_addr.get("longitude") delivery_lat = delivery_addr.get("lat") or delivery_addr.get("latitude") delivery_lon = delivery_addr.get("lng") or delivery_addr.get("longitude") pickup_zone_geo = await self._find_zone(pickup_lat, pickup_lon) if (pickup_lat and pickup_lon) else None delivery_zone_geo = await self._find_zone(delivery_lat, delivery_lon) if (delivery_lat and delivery_lon) else None pickup_pincode = pickup_addr.get("pincode", "")[:3] delivery_pincode = delivery_addr.get("pincode", "")[:3] pickup_zone_pc = self._find_zone_by_pincode(pickup_pincode) delivery_zone_pc = self._find_zone_by_pincode(delivery_pincode) from_hub = ( pickup_zone_geo.get("hub_id") if pickup_zone_geo else (pickup_zone_pc.hub_id if pickup_zone_pc else "UNKNOWN") ) to_hub = ( delivery_zone_geo.get("hub_id") if delivery_zone_geo else (delivery_zone_pc.hub_id if delivery_zone_pc else "UNKNOWN") ) if pickup_pincode == delivery_pincode: route_type = ZoneType.LAST_MILE.value elif from_hub == to_hub: route_type = ZoneType.HUB_TO_SPOKE.value else: route_type = ZoneType.HUB_TO_HUB.value avg_time = ( (pickup_zone_geo or {}).get("avg_delivery_time") or (pickup_zone_pc.avg_delivery_time if pickup_zone_pc else 60) ) now = datetime.now() estimated_pickup = now + timedelta(minutes=30) estimated_delivery = now + timedelta(minutes=avg_time * 2) route_id = f"RT-{datetime.now().strftime('%Y%m%d')}-{uuid.uuid4().hex[:8].upper()}" route_assignment = RouteAssignment( route_id=route_id, order_id=order_id, from_hub=from_hub, to_hub=to_hub, vehicle_id=None, estimated_pickup=estimated_pickup, estimated_delivery=estimated_delivery, route_type=route_type, ) self._active_routes[route_id] = route_assignment logger.info( f"Route assigned: {from_hub} -> {to_hub} | type={route_type} " f"pickup={estimated_pickup.strftime('%H:%M')} delivery={estimated_delivery.strftime('%H:%M')}" ) # Assign miler in Postgres via Go internal API. # booking_id is the Go API's integer ID (set by order_agent after FIX 4). booking_id = order_data.get("booking_id") or order_id nearest_miler = (pickup_zone_geo or {}).get("miler_id") if nearest_miler: logger.debug(f"Nearest miler from GEORADIUS: {nearest_miler}") assign_result = await api_post( f"{GO_API_BASE_URL}/api/v1/internal/bookings/{booking_id}/reassign", json={"reason": "initial_assignment"}, headers={"X-Internal-Key": INTERNAL_API_KEY}, ) if assign_result is not None: logger.info(f"Miler assigned to booking {booking_id} via Go API") else: logger.warning(f"Go API miler assignment failed for booking {booking_id} — booking may stay unassigned") await self.send_message( recipient="HUB_AGENT", message_type=MessageType.AGENT_TASK, payload={"task_type": "prepare_receiving", "route_id": route_id, "hub_id": to_hub, "order_id": order_id}, correlation_id=order_id, ) await self.send_message( recipient="CUSTOMER_AGENT", message_type=MessageType.NOTIFICATION_SENT, payload={ "order_id": order_id, "message_type": "dispatch_assigned", "data": { "route_id": route_id, "estimated_pickup": estimated_pickup.isoformat(), "estimated_delivery": estimated_delivery.isoformat(), "route_type": route_type, }, }, correlation_id=order_id, ) return { "status": "assigned", "route_id": route_id, "order_id": order_id, "booking_id": booking_id, "from_hub": from_hub, "to_hub": to_hub, "route_type": route_type, "nearest_miler": nearest_miler, "miler_assigned": assign_result is not None, "estimated_pickup": estimated_pickup.isoformat(), "estimated_delivery": estimated_delivery.isoformat(), } async def _find_route(self, task: AgentTask) -> Dict[str, Any]: from_pincode = task.data.get("from_pincode", "")[:3] to_pincode = task.data.get("to_pincode", "")[:3] from_zone = self._find_zone_by_pincode(from_pincode) to_zone = self._find_zone_by_pincode(to_pincode) if not from_zone or not to_zone: return {"status": "error", "message": "Zone not found for one or both pincodes"} route_type = ( ZoneType.HUB_TO_SPOKE.value if from_zone.zone_id == to_zone.zone_id else ZoneType.HUB_TO_HUB.value ) return { "status": "found", "from_hub": from_zone.hub_id, "to_hub": to_zone.hub_id, "route_type": route_type, "estimated_time_minutes": from_zone.avg_delivery_time + to_zone.avg_delivery_time, } async def _optimize_routes(self, task: AgentTask) -> Dict[str, Any]: order_ids = task.data.get("order_ids", []) logger.info(f"Dispatch Agent: Optimizing {len(order_ids)} routes") zone_groups: Dict[str, List[str]] = {} for order_id in order_ids: for route in self._active_routes.values(): if route.order_id == order_id: zone_groups.setdefault(route.from_hub, []).append(order_id) optimized = [ {"zone": zone, "orders": orders, "saved_km": len(orders) * 2.5} for zone, orders in zone_groups.items() ] return {"status": "optimized", "groups": optimized, "total_orders": len(order_ids)} async def _assign_vehicle(self, task: AgentTask) -> Dict[str, Any]: route_id = task.data.get("route_id") if route_id not in self._active_routes: return {"status": "error", "message": f"Route {route_id} not found"} self._active_routes[route_id].vehicle_id = task.data.get("vehicle_id") return {"status": "vehicle_assigned", "route_id": route_id} async def _get_zone_status(self, task: AgentTask) -> Dict[str, Any]: zone_status = [] for zone_id, zone in self._zones.items(): hub = self._hubs.get(zone.hub_id, {}) zone_status.append({ "zone_id": zone_id, "zone_name": zone.zone_name, "hub_id": zone.hub_id, "current_load": hub.get("current_load", 0), "max_capacity": hub.get("capacity", 0), "utilization_percent": (hub.get("current_load", 0) / max(hub.get("capacity", 1), 1)) * 100, }) return {"zones": zone_status, "total_zones": len(zone_status)} async def _reschedule_route(self, task: AgentTask) -> Dict[str, Any]: route_id = task.data.get("route_id") if route_id not in self._active_routes: return {"status": "error", "message": f"Route {route_id} not found"} route = self._active_routes[route_id] if task.data.get("new_pickup_time"): route.estimated_pickup = datetime.fromisoformat(task.data["new_pickup_time"]) if task.data.get("new_delivery_time"): route.estimated_delivery = datetime.fromisoformat(task.data["new_delivery_time"]) await self.send_message( recipient="CUSTOMER_AGENT", message_type=MessageType.NOTIFICATION_SENT, payload={ "order_id": route.order_id, "message_type": "schedule_changed", "data": { "new_pickup": route.estimated_pickup.isoformat(), "new_delivery": route.estimated_delivery.isoformat(), }, }, correlation_id=route_id, ) return { "status": "rescheduled", "route_id": route_id, "new_pickup": route.estimated_pickup.isoformat(), "new_delivery": route.estimated_delivery.isoformat(), } async def _cancel_route(self, task: AgentTask) -> Dict[str, Any]: route_id = task.data.get("route_id") reason = task.data.get("reason", "Route cancelled") if route_id not in self._active_routes: return {"status": "error", "message": f"Route {route_id} not found"} route = self._active_routes.pop(route_id) await self.send_message( recipient="FLEET_AGENT", message_type=MessageType.AGENT_TASK, payload={"task_type": "release_vehicle", "vehicle_id": route.vehicle_id}, ) return {"status": "cancelled", "route_id": route_id, "reason": reason} async def _unknown_task(self, task: AgentTask) -> Dict[str, Any]: return {"status": "error", "message": f"Unknown task: {task.task_type}"} async def think(self, context: str, options: List[str] = None) -> str: return f"[DISPATCH_AGENT reasoning]: {context}"