258 lines
9.5 KiB
Python
258 lines
9.5 KiB
Python
"""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}"
|