Doormile/AI_engine
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

refactor: DispatchAgent watcher, ExceptionAgent consumer fix, JARVIS cleanup

Browse Source
This commit is contained in:
Suriya
2026-07-01 13:07:54 +05:30
parent f49193ee73
commit 377ee710a7
3 changed files with 183 additions and 318 deletions
Download Patch File Download Diff File Expand all files Collapse all files

488
agents/dispatch_agent.py
View File

@@ -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 asyncio
import uuid import json
from datetime import datetime, timedelta from datetime import datetime
from typing import Dict, List, Any, Optional from typing import Dict, List, Any, Optional
from dataclasses import dataclass
import nats
import nats.js.errors
import redis.asyncio as aioredis import redis.asyncio as aioredis
from core.agent import SpecializedAgent 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.logger import logger
from core.http_client import api_post from config.system_config import (
from config.system_config import REDIS_HOST, REDIS_PORT, REDIS_PASSWORD, GO_API_BASE_URL, INTERNAL_API_KEY NATS_HOST, NATS_PORT, NATS_USER, NATS_PASSWORD,
REDIS_HOST, REDIS_PORT, REDIS_PASSWORD,
)
@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): class DispatchAgent(SpecializedAgent):
""" """
Dispatch Agent - Orchestrates order dispatching and route assignment. Watches NATS ASSIGNMENTS stream for booking.assigned and
booking.assignment_failed events published by the Go backend
Uses Redis GEO (GEORADIUS on milers:locations) to find nearest milers after routemate AI or fallback assignment. Does not trigger
for a given pickup/delivery coordinate. Falls back to pincode-based assignment itself.
zone lookup when no milers are found within range.
""" """
def __init__(self): def __init__(self):
super().__init__( super().__init__(
agent_id="DISPATCH_AGENT", agent_id="DISPATCH_AGENT",
domain="dispatch", domain="dispatch",
description="Handles zone analysis, route assignment, and dispatch coordination" description="Watches assignment events and escalates coverage gaps"
) )
self._redis = aioredis.Redis( self._redis = aioredis.Redis(
@@ -61,59 +39,79 @@ class DispatchAgent(SpecializedAgent):
decode_responses=True, decode_responses=True,
) )
self._zones = self._init_zones() self._nats_nc = None
self._hubs = self._init_hubs() self._nats_js = None
self._active_routes: Dict[str, RouteAssignment] = {} self._nats_subs: list = []
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 # # Lifecycle #
# ------------------------------------------------------------------ # # ------------------------------------------------------------------ #
async def _find_zone(self, lat: float, lon: float) -> Optional[Dict[str, Any]]: async def start(self):
"""Find the nearest miler using GEORADIUS on milers:locations.""" 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: try:
results = await self._redis.georadius( results = await self._redis.georadius(
"milers:locations", "milers:locations",
lon, lat, lon, lat,
10, "km", radius_km, "km",
sort="ASC", sort="ASC",
count=5, count=5,
) )
@@ -133,246 +131,126 @@ class DispatchAgent(SpecializedAgent):
logger.warning(f"Redis GEORADIUS error: {e}") logger.warning(f"Redis GEORADIUS error: {e}")
return None 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]: async def _on_nats_booking_assigned(self, msg):
order_data = task.data.get("order", {}) """
order_id = order_data.get("order_id", "unknown") 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", {}) await self.send_message(
delivery_addr = order_data.get("delivery_address", {}) recipient="HUB_AGENT",
message_type=MessageType.AGENT_TASK,
pickup_lat = pickup_addr.get("lat") or pickup_addr.get("latitude") payload={
pickup_lon = pickup_addr.get("lng") or pickup_addr.get("longitude") "task_type": "prepare_receiving",
delivery_lat = delivery_addr.get("lat") or delivery_addr.get("latitude") "booking_id": booking_id,
delivery_lon = delivery_addr.get("lng") or delivery_addr.get("longitude") "miler_id": miler_id,
"hub_id": hub_id,
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=str(booking_id),
correlation_id=order_id, )
)
return { except Exception as e:
"status": "assigned", logger.error(f"DispatchAgent booking.assigned handler error: {e}")
"route_id": route_id, finally:
"order_id": order_id, await msg.ack()
"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]: async def _on_nats_booking_assignment_failed(self, msg):
from_pincode = task.data.get("from_pincode", "")[:3] """
to_pincode = task.data.get("to_pincode", "")[:3] 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) Widens the GEORADIUS search to gauge coverage depth, tracks daily
to_zone = self._find_zone_by_pincode(to_pincode) 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: logger.warning(
return {"status": "error", "message": "Zone not found for one or both pincodes"} f"[DISPATCH] booking.assignment_failed — "
f"booking={booking_id} zone={zone_id}"
)
route_type = ( # Wider sweeps to gauge how far the nearest miler actually is.
ZoneType.HUB_TO_SPOKE.value has_coords = lat is not None and lon is not None
if from_zone.zone_id == to_zone.zone_id found_20km = await self._find_zone(lat, lon, radius_km=20) if has_coords else None
else ZoneType.HUB_TO_HUB.value found_30km = await self._find_zone(lat, lon, radius_km=30) if has_coords else None
)
return { if found_20km:
"status": "found", logger.info(
"from_hub": from_zone.hub_id, f"[DISPATCH] Nearest miler within 20 km: {found_20km['miler_id']}"
"to_hub": to_zone.hub_id, )
"route_type": route_type, elif found_30km:
"estimated_time_minutes": from_zone.avg_delivery_time + to_zone.avg_delivery_time, 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]: # Increment daily failure counter (48 h TTL covers day rollover).
order_ids = task.data.get("order_ids", []) today = datetime.utcnow().strftime("%Y-%m-%d")
logger.info(f"Dispatch Agent: Optimizing {len(order_ids)} routes") 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]] = {} if count >= 3:
for order_id in order_ids: logger.warning(
for route in self._active_routes.values(): f"[DISPATCH] Coverage gap detected: zone {zone_id} has "
if route.order_id == order_id: f"{count} failed assignments today — escalating"
zone_groups.setdefault(route.from_hub, []).append(order_id) )
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 = [ except Exception as e:
{"zone": zone, "orders": orders, "saved_km": len(orders) * 2.5} logger.error(f"DispatchAgent booking.assignment_failed handler error: {e}")
for zone, orders in zone_groups.items() 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]: async def handle_task(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}"} return {"status": "error", "message": f"Unknown task: {task.task_type}"}
async def think(self, context: str, options: List[str] = None) -> str: async def think(self, context: str, options: List[str] = None) -> str:

4
agents/exception_agent.py
View File

@@ -495,10 +495,6 @@ class ExceptionAgent(SpecializedAgent):
detected_at=datetime.now(), resolved_at=None, resolution=None, detected_at=datetime.now(), resolved_at=None, resolution=None,
assigned_to=None, status="resolving", 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( await self.send_message(
recipient="CUSTOMER_AGENT", message_type=MessageType.NOTIFICATION_SENT, recipient="CUSTOMER_AGENT", message_type=MessageType.NOTIFICATION_SENT,
payload={"order_id": order_id, "message_type": "rescheduled", payload={"order_id": order_id, "message_type": "rescheduled",

9
core/agent.py
View File

@@ -178,15 +178,6 @@ class MasterAgent(Agent):
data={"order": order_data} 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") customer_agent = self._sub_agents.get("CUSTOMER_AGENT")
if customer_agent: if customer_agent:
await customer_agent.submit_task(AgentTask( await customer_agent.submit_task(AgentTask(