diff --git a/agents/dispatch_agent.py b/agents/dispatch_agent.py index cb42df3..a1160fa 100644 --- a/agents/dispatch_agent.py +++ b/agents/dispatch_agent.py @@ -1,57 +1,35 @@ -"""Dispatch Agent - Handles zone analysis, route assignment, and order dispatching.""" +"""Dispatch Agent — watches NATS assignment events and escalates coverage gaps.""" import asyncio -import uuid -from datetime import datetime, timedelta +import json +from datetime import datetime from typing import Dict, List, Any, Optional -from dataclasses import dataclass +import nats +import nats.js.errors import redis.asyncio as aioredis from core.agent import SpecializedAgent -from core.types import AgentTask, MessageType, Priority, ZoneType +from core.types import AgentTask, MessageType 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 +from config.system_config import ( + NATS_HOST, NATS_PORT, NATS_USER, NATS_PASSWORD, + REDIS_HOST, REDIS_PORT, REDIS_PASSWORD, +) 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. + Watches NATS ASSIGNMENTS stream for booking.assigned and + booking.assignment_failed events published by the Go backend + after routemate AI or fallback assignment. Does not trigger + assignment itself. """ def __init__(self): super().__init__( agent_id="DISPATCH_AGENT", domain="dispatch", - description="Handles zone analysis, route assignment, and dispatch coordination" + description="Watches assignment events and escalates coverage gaps" ) self._redis = aioredis.Redis( @@ -61,59 +39,79 @@ class DispatchAgent(SpecializedAgent): 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) + self._nats_nc = None + self._nats_js = None + self._nats_subs: list = [] # ------------------------------------------------------------------ # - # Redis GEO zone lookup # + # Lifecycle # # ------------------------------------------------------------------ # - async def _find_zone(self, lat: float, lon: float) -> Optional[Dict[str, Any]]: - """Find the nearest miler using GEORADIUS on milers:locations.""" + async def start(self): + await self._connect_nats() + await super().start() + + async def stop(self): + await super().stop() + for sub in self._nats_subs: + try: + await sub.unsubscribe() + except Exception: + pass + if self._nats_nc: + await self._nats_nc.drain() + + async def _connect_nats(self): + try: + self._nats_nc = await nats.connect( + servers=[f"nats://{NATS_HOST}:{NATS_PORT}"], + user=NATS_USER, + password=NATS_PASSWORD, + max_reconnect_attempts=10, + ) + self._nats_js = self._nats_nc.jetstream() + + try: + await self._nats_js.add_stream( + name="ASSIGNMENTS", + subjects=["booking.>"], + ) + logger.info("NATS stream 'ASSIGNMENTS' created") + except nats.js.errors.BadRequestError: + pass # stream already exists + + sub_assigned = await self._nats_js.subscribe( + "booking.assigned", + durable="dispatch-booking-assigned", + cb=self._on_nats_booking_assigned, + ) + sub_failed = await self._nats_js.subscribe( + "booking.assignment_failed", + durable="dispatch-booking-assignment-failed", + cb=self._on_nats_booking_assignment_failed, + ) + self._nats_subs.extend([sub_assigned, sub_failed]) + logger.info( + "DispatchAgent subscribed to booking.assigned " + "and booking.assignment_failed" + ) + + except Exception as e: + logger.error(f"DispatchAgent NATS connect error: {e}") + + # ------------------------------------------------------------------ # + # Redis GEO lookup # + # ------------------------------------------------------------------ # + + async def _find_zone( + self, lat: float, lon: float, radius_km: int = 10 + ) -> Optional[Dict[str, Any]]: + """GEORADIUS sweep on milers:locations; returns nearest miler or None.""" try: results = await self._redis.georadius( "milers:locations", lon, lat, - 10, "km", + radius_km, "km", sort="ASC", count=5, ) @@ -133,246 +131,126 @@ class DispatchAgent(SpecializedAgent): 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 # + # NATS event handlers # # ------------------------------------------------------------------ # - 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") + async def _on_nats_booking_assigned(self, msg): + """ + booking.assigned — Go publishes this after a successful AI or + fallback assignment. Forward a prepare_receiving task to HUB_AGENT + using the real miler_id/hub_id from the event. + """ + try: + data = json.loads(msg.data.decode()) + booking_id = data.get("booking_id") + miler_id = data.get("miler_id") + hub_id = data.get("hub_id") + confidence = data.get("confidence") + reasoning = data.get("reasoning", "") - logger.info(f"Dispatch Agent: Analyzing order {order_id}") + logger.info( + f"[DISPATCH] booking.assigned — booking={booking_id} " + f"miler={miler_id} hub={hub_id} confidence={confidence} " + f"reasoning={reasoning!r}" + ) - 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, + await self.send_message( + recipient="HUB_AGENT", + message_type=MessageType.AGENT_TASK, + payload={ + "task_type": "prepare_receiving", + "booking_id": booking_id, + "miler_id": miler_id, + "hub_id": hub_id, }, - }, - correlation_id=order_id, - ) + correlation_id=str(booking_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(), - } + except Exception as e: + logger.error(f"DispatchAgent booking.assigned handler error: {e}") + finally: + await msg.ack() - 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] + async def _on_nats_booking_assignment_failed(self, msg): + """ + booking.assignment_failed — Go publishes this when decide-assignment + returns escalate=true or both AI and fallback find nobody. - from_zone = self._find_zone_by_pincode(from_pincode) - to_zone = self._find_zone_by_pincode(to_pincode) + Widens the GEORADIUS search to gauge coverage depth, tracks daily + failures per zone in Redis, and escalates to CUSTOMER_AGENT as an + ops_alert when a zone hits 3 failures in a day. + """ + try: + data = json.loads(msg.data.decode()) + booking_id = data.get("booking_id") + zone_id = data.get("zone_id", "unknown") + lat = data.get("lat") + lon = data.get("lon") - if not from_zone or not to_zone: - return {"status": "error", "message": "Zone not found for one or both pincodes"} + logger.warning( + f"[DISPATCH] booking.assignment_failed — " + f"booking={booking_id} zone={zone_id}" + ) - route_type = ( - ZoneType.HUB_TO_SPOKE.value - if from_zone.zone_id == to_zone.zone_id - else ZoneType.HUB_TO_HUB.value - ) + # Wider sweeps to gauge how far the nearest miler actually is. + has_coords = lat is not None and lon is not None + found_20km = await self._find_zone(lat, lon, radius_km=20) if has_coords else None + found_30km = await self._find_zone(lat, lon, radius_km=30) if has_coords else None - 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, - } + if found_20km: + logger.info( + f"[DISPATCH] Nearest miler within 20 km: {found_20km['miler_id']}" + ) + elif found_30km: + logger.info( + f"[DISPATCH] Nearest miler within 30 km: {found_30km['miler_id']}" + ) + else: + logger.warning( + f"[DISPATCH] No miler found within 30 km of zone {zone_id}" + ) - 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") + # Increment daily failure counter (48 h TTL covers day rollover). + today = datetime.utcnow().strftime("%Y-%m-%d") + counter_key = f"failed_assignment:{zone_id}:{today}" + count = await self._redis.incr(counter_key) + await self._redis.expire(counter_key, 172800) - 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) + if count >= 3: + logger.warning( + f"[DISPATCH] Coverage gap detected: zone {zone_id} has " + f"{count} failed assignments today — escalating" + ) + await self.send_message( + recipient="CUSTOMER_AGENT", + message_type=MessageType.AGENT_TASK, + payload={ + "task_type": "ops_alert", + "zone_id": zone_id, + "reasoning": ( + f"Zone {zone_id} has had {count} failed assignments today " + f"— recommend onboarding milers here." + ), + }, + correlation_id=str(booking_id), + ) + else: + logger.info( + f"[DISPATCH] zone {zone_id} failure count {count}/3 — " + f"transient, no escalation" + ) - optimized = [ - {"zone": zone, "orders": orders, "saved_km": len(orders) * 2.5} - for zone, orders in zone_groups.items() - ] + except Exception as e: + logger.error(f"DispatchAgent booking.assignment_failed handler error: {e}") + finally: + await msg.ack() - return {"status": "optimized", "groups": optimized, "total_orders": len(order_ids)} + # ------------------------------------------------------------------ # + # Task handler (no inbound task types remain) # + # ------------------------------------------------------------------ # - 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]: + async def handle_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: diff --git a/agents/exception_agent.py b/agents/exception_agent.py index 76dbdc9..7a54d32 100644 --- a/agents/exception_agent.py +++ b/agents/exception_agent.py @@ -495,10 +495,6 @@ class ExceptionAgent(SpecializedAgent): detected_at=datetime.now(), resolved_at=None, resolution=None, assigned_to=None, status="resolving", ) - await self.send_message( - recipient="DISPATCH_AGENT", message_type=MessageType.AGENT_TASK, - payload={"task_type": "reschedule_route", "order_id": order_id, "new_date": new_date}, - ) await self.send_message( recipient="CUSTOMER_AGENT", message_type=MessageType.NOTIFICATION_SENT, payload={"order_id": order_id, "message_type": "rescheduled", diff --git a/core/agent.py b/core/agent.py index f0bee02..9b2bbed 100644 --- a/core/agent.py +++ b/core/agent.py @@ -178,15 +178,6 @@ class MasterAgent(Agent): data={"order": order_data} )) - dispatch_agent = self._sub_agents.get("DISPATCH_AGENT") - if dispatch_agent: - await dispatch_agent.submit_task(AgentTask( - task_id=f"{order_id}_dispatch", - agent_type="dispatch", - task_type="analyze_and_assign", - data={"order": order_data} - )) - customer_agent = self._sub_agents.get("CUSTOMER_AGENT") if customer_agent: await customer_agent.submit_task(AgentTask(