353 lines
15 KiB
Python
353 lines
15 KiB
Python
"""Dispatch Agent — watches NATS assignment events and escalates coverage gaps."""
|
|
import asyncio
|
|
import json
|
|
import os
|
|
from datetime import datetime, timezone
|
|
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 core.llm import decide_assignment_failure, build_assignment_failure_context
|
|
from config.system_config import (
|
|
NATS_HOST, NATS_PORT, NATS_USER, NATS_PASSWORD,
|
|
REDIS_HOST, REDIS_PORT, REDIS_PASSWORD,
|
|
)
|
|
|
|
# Notifying a customer is outward-facing, so it is gated: without autonomy the
|
|
# agent proposes it as an internal ops alert instead of messaging the customer.
|
|
DISPATCH_AGENT_AUTONOMOUS = os.getenv("DISPATCH_AGENT_AUTONOMOUS", "false").lower() == "true"
|
|
|
|
# The Go backend owns the JetStream stream carrying booking.* events; this agent
|
|
# is a pure consumer and must not create it. Leave empty to auto-discover the
|
|
# stream by subject, or pin it to a stream name if discovery isn't desired.
|
|
DISPATCH_STREAM = os.getenv("DISPATCH_STREAM", "")
|
|
|
|
|
|
class DispatchAgent(SpecializedAgent):
|
|
"""
|
|
Consumes booking.assigned and booking.assignment_failed events published by
|
|
the Go backend (after routemate AI or fallback assignment) by binding to the
|
|
backend-owned JetStream stream that carries them. It is a pure consumer — it
|
|
does not create the stream, and 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()
|
|
|
|
# Bind consumers to the *existing* backend-owned stream. Do NOT create
|
|
# a stream here: booking.* is published by the Go backend, so creating
|
|
# a second stream over the same subjects raises an overlap error that
|
|
# (previously swallowed) left the consumer unbound. Mirrors the
|
|
# explicit stream= binding ExceptionAgent uses for the TRACKING stream.
|
|
sub_assigned = await self._bind_consumer(
|
|
"booking.assigned", "dispatch-booking-assigned",
|
|
self._on_nats_booking_assigned,
|
|
)
|
|
sub_failed = await self._bind_consumer(
|
|
"booking.assignment_failed", "dispatch-booking-assignment-failed",
|
|
self._on_nats_booking_assignment_failed,
|
|
)
|
|
self._nats_subs.extend(s for s in (sub_assigned, sub_failed) if s is not None)
|
|
if self._nats_subs:
|
|
logger.info(f"DispatchAgent bound {len(self._nats_subs)} booking-event consumer(s)")
|
|
|
|
except Exception as e:
|
|
logger.error(f"DispatchAgent NATS connect error: {e}")
|
|
|
|
async def _bind_consumer(self, subject, durable, cb):
|
|
"""Bind a durable push consumer to whichever existing stream carries
|
|
`subject`. Returns the subscription, or None (with a clear, actionable
|
|
log) if no stream carries it — never a swallowed or opaque error."""
|
|
try:
|
|
stream = DISPATCH_STREAM or await self._nats_js.find_stream_name_by_subject(subject)
|
|
sub = await self._nats_js.subscribe(subject, durable=durable, stream=stream, cb=cb)
|
|
logger.info(f"DispatchAgent bound '{subject}' on stream '{stream}' (durable={durable})")
|
|
return sub
|
|
except nats.js.errors.NotFoundError:
|
|
logger.error(
|
|
f"DispatchAgent: no JetStream stream carries '{subject}'. The publisher "
|
|
f"(Go backend) must create a stream covering it, or set DISPATCH_STREAM to "
|
|
f"the stream name. Not subscribing to '{subject}'."
|
|
)
|
|
return None
|
|
except Exception as e:
|
|
logger.error(f"DispatchAgent failed to bind '{subject}': {e}")
|
|
return None
|
|
|
|
# ------------------------------------------------------------------ #
|
|
# 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.
|
|
|
|
Gathers coverage context (nearest-rider sweep + daily failure counter),
|
|
then asks Claude to choose: monitor / notify_customer / ops_alert /
|
|
escalate. Customer notification is gated behind autonomy. Falls back to
|
|
the previous count>=3 heuristic if the LLM is unavailable.
|
|
"""
|
|
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 — booking={booking_id} zone={zone_id}")
|
|
|
|
facts, count = await self._gather_assignment_facts(zone_id, lat, lon)
|
|
logger.info(f"[DISPATCH] booking={booking_id} context facts: {facts}")
|
|
|
|
decision = await decide_assignment_failure(build_assignment_failure_context(facts))
|
|
|
|
if decision is None:
|
|
logger.warning(f"No LLM decision for booking {booking_id}; falling back to count>=3 heuristic")
|
|
if count >= 3:
|
|
await self._ops_alert(
|
|
zone_id, booking_id,
|
|
f"Zone {zone_id} has had {count} failed assignments today — recommend onboarding milers here.",
|
|
)
|
|
else:
|
|
logger.info(f"[DISPATCH] zone {zone_id} failure count {count}/3 — transient, no escalation")
|
|
return
|
|
|
|
logger.info(
|
|
f"[DISPATCH] booking={booking_id} zone={zone_id} decision={decision.action} "
|
|
f"confidence={decision.confidence:.2f} reasoning={decision.reasoning!r}"
|
|
)
|
|
|
|
if decision.action == "monitor":
|
|
pass # transient — the backend will retry
|
|
|
|
elif decision.action == "notify_customer":
|
|
if DISPATCH_AGENT_AUTONOMOUS:
|
|
await self._notify_customer_delay(booking_id)
|
|
else:
|
|
# Outward-facing action not authorised — raise an internal proposal instead.
|
|
await self._ops_alert(
|
|
zone_id, booking_id,
|
|
f"[proposed: notify_customer, conf={decision.confidence:.2f}] {decision.reasoning}",
|
|
)
|
|
|
|
elif decision.action == "ops_alert":
|
|
await self._ops_alert(zone_id, booking_id, decision.reasoning)
|
|
|
|
else: # "escalate"
|
|
await self._escalate_dispatch(zone_id, booking_id, decision)
|
|
|
|
except Exception as e:
|
|
logger.error(f"DispatchAgent booking.assignment_failed handler error: {e}")
|
|
finally:
|
|
await msg.ack()
|
|
|
|
# ------------------------------------------------------------------ #
|
|
# Context gathering + actions #
|
|
# ------------------------------------------------------------------ #
|
|
|
|
async def _gather_assignment_facts(self, zone_id, lat, lon):
|
|
"""Read-only coverage context plus the daily failure counter for this
|
|
zone. Returns (facts, failures_today). Best-effort — any failure just
|
|
yields fewer facts, never raises."""
|
|
facts: Dict[str, Any] = {"zone_id": zone_id}
|
|
has_coords = lat is not None and lon is not None
|
|
facts["has_coordinates"] = has_coords
|
|
|
|
nearest_km = None
|
|
if has_coords:
|
|
for radius in (10, 20, 30):
|
|
try:
|
|
if await self._find_zone(lat, lon, radius_km=radius):
|
|
nearest_km = radius
|
|
break
|
|
except Exception as e:
|
|
logger.warning(f"assignment facts: GEORADIUS {radius}km failed: {e}")
|
|
break
|
|
|
|
if not has_coords:
|
|
facts["nearest_miler_within_km"] = "unknown (no coordinates)"
|
|
elif nearest_km is None:
|
|
facts["nearest_miler_within_km"] = "none within 30km"
|
|
else:
|
|
facts["nearest_miler_within_km"] = nearest_km
|
|
|
|
count = 0
|
|
try:
|
|
today = datetime.now(timezone.utc).strftime("%Y-%m-%d")
|
|
key = f"failed_assignment:{zone_id}:{today}"
|
|
count = await self._redis.incr(key)
|
|
await self._redis.expire(key, 172800) # 48h covers day rollover
|
|
except Exception as e:
|
|
logger.warning(f"assignment facts: failure counter failed for zone {zone_id}: {e}")
|
|
facts["failures_today"] = count
|
|
|
|
return facts, count
|
|
|
|
async def _ops_alert(self, zone_id, booking_id, reasoning):
|
|
await self.send_message(
|
|
recipient="CUSTOMER_AGENT",
|
|
message_type=MessageType.AGENT_TASK,
|
|
payload={"task_type": "ops_alert", "zone_id": zone_id, "reasoning": reasoning},
|
|
correlation_id=str(booking_id),
|
|
)
|
|
|
|
async def _notify_customer_delay(self, booking_id):
|
|
await self.send_message(
|
|
recipient="CUSTOMER_AGENT",
|
|
message_type=MessageType.AGENT_TASK,
|
|
payload={
|
|
"task_type": "customer_notify",
|
|
"booking_id": booking_id,
|
|
"message": "We're finding the right rider for your delivery — thanks for your patience.",
|
|
},
|
|
correlation_id=str(booking_id),
|
|
)
|
|
|
|
async def _escalate_dispatch(self, zone_id, booking_id, decision):
|
|
logger.warning(
|
|
f"[DISPATCH] Escalating booking {booking_id} (zone {zone_id}) to a human "
|
|
f"(proposed={decision.action}, confidence={decision.confidence:.2f}): {decision.reasoning}"
|
|
)
|
|
await self.send_message(
|
|
recipient="JARVIS",
|
|
message_type=MessageType.AGENT_TASK,
|
|
payload={
|
|
"task_type": "human_review",
|
|
"reason": "assignment_failed",
|
|
"zone_id": zone_id,
|
|
"booking_id": booking_id,
|
|
"proposed_action": decision.action,
|
|
"confidence": decision.confidence,
|
|
"reasoning": decision.reasoning,
|
|
},
|
|
correlation_id=str(booking_id),
|
|
)
|
|
|
|
# ------------------------------------------------------------------ #
|
|
# 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}"
|