"""Dispatch Agent — watches NATS assignment events and escalates coverage gaps.""" import asyncio import json from datetime import datetime from typing import Dict, List, Any, Optional import nats import nats.js.errors import redis.asyncio as aioredis from core.agent import SpecializedAgent from core.types import AgentTask, MessageType from core.logger import logger from config.system_config import ( NATS_HOST, NATS_PORT, NATS_USER, NATS_PASSWORD, REDIS_HOST, REDIS_PORT, REDIS_PASSWORD, ) class DispatchAgent(SpecializedAgent): """ 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="Watches assignment events and escalates coverage gaps" ) self._redis = aioredis.Redis( host=REDIS_HOST, port=REDIS_PORT, password=REDIS_PASSWORD, decode_responses=True, ) self._nats_nc = None self._nats_js = None self._nats_subs: list = [] # ------------------------------------------------------------------ # # Lifecycle # # ------------------------------------------------------------------ # 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, radius_km, "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 # ------------------------------------------------------------------ # # NATS event handlers # # ------------------------------------------------------------------ # 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] booking.assigned — booking={booking_id} " f"miler={miler_id} hub={hub_id} confidence={confidence} " f"reasoning={reasoning!r}" ) 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=str(booking_id), ) except Exception as e: logger.error(f"DispatchAgent booking.assigned handler error: {e}") finally: await msg.ack() 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. 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") logger.warning( f"[DISPATCH] booking.assignment_failed — " f"booking={booking_id} zone={zone_id}" ) # 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 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}" ) # 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) 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" ) except Exception as e: logger.error(f"DispatchAgent booking.assignment_failed handler error: {e}") finally: await msg.ack() # ------------------------------------------------------------------ # # Task handler (no inbound task types remain) # # ------------------------------------------------------------------ # 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: return f"[DISPATCH_AGENT reasoning]: {context}"