380 lines
16 KiB
Python
380 lines
16 KiB
Python
"""Dispatch Agent - Handles zone analysis, route assignment, and order dispatching."""
|
|
import asyncio
|
|
import uuid
|
|
from datetime import datetime, timedelta
|
|
from typing import Dict, List, Any, Optional
|
|
from dataclasses import dataclass
|
|
|
|
import redis.asyncio as aioredis
|
|
|
|
from core.agent import SpecializedAgent
|
|
from core.types import AgentTask, MessageType, Priority, ZoneType
|
|
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
|
|
|
|
|
|
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.
|
|
"""
|
|
|
|
def __init__(self):
|
|
super().__init__(
|
|
agent_id="DISPATCH_AGENT",
|
|
domain="dispatch",
|
|
description="Handles zone analysis, route assignment, and dispatch coordination"
|
|
)
|
|
|
|
self._redis = aioredis.Redis(
|
|
host=REDIS_HOST,
|
|
port=REDIS_PORT,
|
|
password=REDIS_PASSWORD,
|
|
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)
|
|
|
|
# ------------------------------------------------------------------ #
|
|
# Redis GEO zone lookup #
|
|
# ------------------------------------------------------------------ #
|
|
|
|
async def _find_zone(self, lat: float, lon: float) -> Optional[Dict[str, Any]]:
|
|
"""Find the nearest miler using GEORADIUS on milers:locations."""
|
|
try:
|
|
results = await self._redis.georadius(
|
|
"milers:locations",
|
|
lon, lat,
|
|
10, "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
|
|
|
|
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 #
|
|
# ------------------------------------------------------------------ #
|
|
|
|
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")
|
|
|
|
logger.info(f"Dispatch Agent: Analyzing order {order_id}")
|
|
|
|
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,
|
|
},
|
|
},
|
|
correlation_id=order_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(),
|
|
}
|
|
|
|
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]
|
|
|
|
from_zone = self._find_zone_by_pincode(from_pincode)
|
|
to_zone = self._find_zone_by_pincode(to_pincode)
|
|
|
|
if not from_zone or not to_zone:
|
|
return {"status": "error", "message": "Zone not found for one or both pincodes"}
|
|
|
|
route_type = (
|
|
ZoneType.HUB_TO_SPOKE.value
|
|
if from_zone.zone_id == to_zone.zone_id
|
|
else ZoneType.HUB_TO_HUB.value
|
|
)
|
|
|
|
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,
|
|
}
|
|
|
|
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")
|
|
|
|
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)
|
|
|
|
optimized = [
|
|
{"zone": zone, "orders": orders, "saved_km": len(orders) * 2.5}
|
|
for zone, orders in zone_groups.items()
|
|
]
|
|
|
|
return {"status": "optimized", "groups": optimized, "total_orders": len(order_ids)}
|
|
|
|
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]:
|
|
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}"
|