Initial commit
This commit is contained in:
24
core/__init__.py
Normal file
24
core/__init__.py
Normal file
@@ -0,0 +1,24 @@
|
||||
# Core module for LogiFlow AI
|
||||
from core.agent import Agent, MasterAgent, SpecializedAgent
|
||||
from core.message_bus import message_bus, MessageBus
|
||||
from core.types import *
|
||||
|
||||
__all__ = [
|
||||
'Agent',
|
||||
'MasterAgent',
|
||||
'SpecializedAgent',
|
||||
'message_bus',
|
||||
'MessageBus',
|
||||
'AgentTask',
|
||||
'AgentMessage',
|
||||
'AgentState',
|
||||
'MessageType',
|
||||
'Priority',
|
||||
'OrderStatus',
|
||||
'ZoneType',
|
||||
'Order',
|
||||
'Hub',
|
||||
'Vehicle',
|
||||
'Address',
|
||||
'OrderItem'
|
||||
]
|
||||
273
core/agent.py
Normal file
273
core/agent.py
Normal file
@@ -0,0 +1,273 @@
|
||||
"""Base Agent class - All agents inherit from this."""
|
||||
import asyncio
|
||||
import uuid
|
||||
from datetime import datetime
|
||||
from typing import Dict, List, Optional, Any
|
||||
from abc import ABC, abstractmethod
|
||||
|
||||
from core.types import AgentMessage, MessageType, AgentState, AgentTask, Priority
|
||||
from core.message_bus import message_bus
|
||||
from core.logger import logger
|
||||
|
||||
|
||||
class Agent(ABC):
|
||||
"""
|
||||
Base class for all AI agents in the logistics system.
|
||||
Each agent has its own task queue and can communicate via the message bus.
|
||||
"""
|
||||
|
||||
def __init__(self, agent_id: str, agent_type: str, description: str = ""):
|
||||
self.agent_id = agent_id
|
||||
self.agent_type = agent_type
|
||||
self.description = description
|
||||
self.state = AgentState(
|
||||
agent_id=agent_id,
|
||||
agent_type=agent_type,
|
||||
status="idle"
|
||||
)
|
||||
self._task_queue: asyncio.Queue = asyncio.Queue()
|
||||
self._running = False
|
||||
self._task_handlers: Dict[str, callable] = {}
|
||||
|
||||
# Register with message bus
|
||||
message_bus.register_agent(self)
|
||||
|
||||
async def start(self):
|
||||
"""Start the agent's processing loop."""
|
||||
self._running = True
|
||||
self.state.status = "idle"
|
||||
logger.info(f"Agent started: {self.agent_id}")
|
||||
|
||||
while self._running:
|
||||
try:
|
||||
try:
|
||||
task = await asyncio.wait_for(
|
||||
self._task_queue.get(),
|
||||
timeout=1.0
|
||||
)
|
||||
await self._process_task(task)
|
||||
except asyncio.TimeoutError:
|
||||
await self._heartbeat()
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error in agent {self.agent_id}: {e}")
|
||||
self.state.status = "error"
|
||||
await asyncio.sleep(5)
|
||||
|
||||
async def stop(self):
|
||||
"""Stop the agent."""
|
||||
self._running = False
|
||||
self.state.status = "stopped"
|
||||
logger.info(f"Agent stopped: {self.agent_id}")
|
||||
|
||||
async def _process_task(self, task: AgentTask):
|
||||
"""Process a task from the queue."""
|
||||
self.state.status = "working"
|
||||
self.state.current_task = task.task_id
|
||||
|
||||
try:
|
||||
if task.task_type in self._task_handlers:
|
||||
handler = self._task_handlers[task.task_type]
|
||||
result = await handler(task)
|
||||
task.status = "completed"
|
||||
task.result = result
|
||||
else:
|
||||
result = await self.handle_task(task)
|
||||
task.status = "completed"
|
||||
task.result = result
|
||||
|
||||
self.state.tasks_completed += 1
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Task {task.task_id} failed in {self.agent_id}: {e}")
|
||||
task.status = "failed"
|
||||
task.error = str(e)
|
||||
self.state.tasks_failed += 1
|
||||
|
||||
finally:
|
||||
self.state.current_task = None
|
||||
self.state.last_active = datetime.now()
|
||||
self.state.status = "idle"
|
||||
|
||||
@abstractmethod
|
||||
async def handle_task(self, task: AgentTask) -> Dict[str, Any]:
|
||||
"""Override this method to handle tasks. Return a dict with results."""
|
||||
pass
|
||||
|
||||
async def _heartbeat(self):
|
||||
"""Called periodically when idle. Override for custom behavior."""
|
||||
pass
|
||||
|
||||
def register_task_handler(self, task_type: str, handler: callable):
|
||||
self._task_handlers[task_type] = handler
|
||||
|
||||
async def submit_task(self, task: AgentTask):
|
||||
await self._task_queue.put(task)
|
||||
|
||||
async def send_message(self, recipient: str, message_type: MessageType, payload: Dict[str, Any], correlation_id: Optional[str] = None) -> str:
|
||||
return await message_bus.send_to_agent(
|
||||
sender=self.agent_id,
|
||||
recipient=recipient,
|
||||
message_type=message_type,
|
||||
payload=payload,
|
||||
correlation_id=correlation_id
|
||||
)
|
||||
|
||||
async def broadcast(self, message_type: MessageType, payload: Dict[str, Any], correlation_id: Optional[str] = None) -> str:
|
||||
return await message_bus.broadcast(
|
||||
sender=self.agent_id,
|
||||
message_type=message_type,
|
||||
payload=payload,
|
||||
correlation_id=correlation_id
|
||||
)
|
||||
|
||||
async def receive_messages(self) -> List[AgentMessage]:
|
||||
return await message_bus.get_messages(self.agent_id)
|
||||
|
||||
def subscribe_to(self, message_type: MessageType, callback: callable):
|
||||
message_bus.subscribe(message_type, callback)
|
||||
|
||||
async def think(self, context: str, options: List[str] = None) -> str:
|
||||
return context
|
||||
|
||||
|
||||
class MasterAgent(Agent):
|
||||
"""
|
||||
Master Agent (JARVIS) - Orchestrates all other agents.
|
||||
Makes high-level decisions and assigns work.
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
super().__init__(
|
||||
agent_id="JARVIS",
|
||||
agent_type="master",
|
||||
description="Central orchestrator for all logistics operations"
|
||||
)
|
||||
self._sub_agents: Dict[str, Agent] = {}
|
||||
self._active_orders: Dict[str, Dict] = {}
|
||||
self._decision_log: List[Dict] = []
|
||||
|
||||
def register_sub_agent(self, agent: Agent):
|
||||
self._sub_agents[agent.agent_id] = agent
|
||||
logger.info(f"JARVIS registered sub-agent: {agent.agent_id}")
|
||||
|
||||
async def handle_task(self, task: AgentTask) -> Dict[str, Any]:
|
||||
if task.task_type == "orchestrate_order":
|
||||
return await self._orchestrate_order(task)
|
||||
elif task.task_type == "monitor_operations":
|
||||
return await self._monitor_operations(task)
|
||||
elif task.task_type == "handle_exception":
|
||||
return await self._handle_exception(task)
|
||||
elif task.task_type == "generate_report":
|
||||
return await self._generate_report(task)
|
||||
else:
|
||||
return {"status": "unknown_task", "task_type": task.task_type}
|
||||
|
||||
async def _orchestrate_order(self, task: AgentTask) -> Dict[str, Any]:
|
||||
order_data = task.data.get("order", {})
|
||||
order_id = order_data.get("order_id", "unknown")
|
||||
|
||||
logger.info(f"JARVIS: Orchestrating order {order_id}")
|
||||
|
||||
order_agent = self._sub_agents.get("ORDER_AGENT")
|
||||
if order_agent:
|
||||
await order_agent.submit_task(AgentTask(
|
||||
task_id=f"{order_id}_validate",
|
||||
agent_type="order",
|
||||
task_type="validate_order",
|
||||
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")
|
||||
if customer_agent:
|
||||
await customer_agent.submit_task(AgentTask(
|
||||
task_id=f"{order_id}_notify",
|
||||
agent_type="customer",
|
||||
task_type="send_confirmation",
|
||||
data={"order": order_data}
|
||||
))
|
||||
|
||||
self._decision_log.append({
|
||||
"timestamp": datetime.now(),
|
||||
"action": "orchestrate_order",
|
||||
"order_id": order_id,
|
||||
"agents_involved": list(self._sub_agents.keys())
|
||||
})
|
||||
|
||||
return {
|
||||
"status": "orchestrated",
|
||||
"order_id": order_id,
|
||||
"assigned_agents": list(self._sub_agents.keys())
|
||||
}
|
||||
|
||||
async def _monitor_operations(self, task: AgentTask) -> Dict[str, Any]:
|
||||
status_report = {
|
||||
"timestamp": datetime.now().isoformat(),
|
||||
"total_agents": len(self._sub_agents),
|
||||
"agent_statuses": {},
|
||||
"active_orders": len(self._active_orders),
|
||||
"recent_decisions": self._decision_log[-10:]
|
||||
}
|
||||
|
||||
for agent_id, agent in self._sub_agents.items():
|
||||
status_report["agent_statuses"][agent_id] = {
|
||||
"status": agent.state.status,
|
||||
"tasks_completed": agent.state.tasks_completed,
|
||||
"tasks_failed": agent.state.tasks_failed,
|
||||
"current_task": agent.state.current_task
|
||||
}
|
||||
|
||||
return status_report
|
||||
|
||||
async def _handle_exception(self, task: AgentTask) -> Dict[str, Any]:
|
||||
exception_type = task.data.get("type", "unknown")
|
||||
logger.warning(f"JARVIS: Handling exception - {exception_type}")
|
||||
|
||||
exception_agent = self._sub_agents.get("EXCEPTION_AGENT")
|
||||
if exception_agent:
|
||||
await exception_agent.submit_task(task)
|
||||
return {"status": "delegated", "exception_type": exception_type}
|
||||
|
||||
return {"status": "no_exception_agent"}
|
||||
|
||||
async def _generate_report(self, task: AgentTask) -> Dict[str, Any]:
|
||||
return {
|
||||
"report_type": task.data.get("type", "summary"),
|
||||
"generated_at": datetime.now().isoformat(),
|
||||
"agents": [a.state for a in self._sub_agents.values()],
|
||||
"decisions": self._decision_log[-50:]
|
||||
}
|
||||
|
||||
async def _heartbeat(self):
|
||||
for agent_id, agent in self._sub_agents.items():
|
||||
if agent.state.status == "error":
|
||||
logger.warning(f"Agent {agent_id} in error state")
|
||||
|
||||
|
||||
class SpecializedAgent(Agent):
|
||||
"""Base class for specialized domain agents."""
|
||||
|
||||
def __init__(self, agent_id: str, domain: str, description: str):
|
||||
super().__init__(agent_id, domain, description)
|
||||
self._domain_knowledge: Dict[str, Any] = {}
|
||||
self._learning_history: List[Dict] = []
|
||||
|
||||
async def think(self, context: str, options: List[str] = None) -> str:
|
||||
return f"[{self.agent_type} reasoning]: {context}"
|
||||
|
||||
def learn_from(self, experience: Dict[str, Any]):
|
||||
self._learning_history.append({
|
||||
"timestamp": datetime.now(),
|
||||
"experience": experience
|
||||
})
|
||||
if len(self._learning_history) > 100:
|
||||
self._learning_history = self._learning_history[-100:]
|
||||
72
core/http_client.py
Normal file
72
core/http_client.py
Normal file
@@ -0,0 +1,72 @@
|
||||
"""Shared aiohttp session with connection pooling and exponential-backoff retry."""
|
||||
import asyncio
|
||||
from typing import Optional, Dict, Any
|
||||
|
||||
import aiohttp
|
||||
|
||||
from core.logger import logger
|
||||
|
||||
_session: Optional[aiohttp.ClientSession] = None
|
||||
|
||||
|
||||
def _new_session() -> aiohttp.ClientSession:
|
||||
connector = aiohttp.TCPConnector(limit=100, limit_per_host=20, keepalive_timeout=30)
|
||||
return aiohttp.ClientSession(
|
||||
connector=connector,
|
||||
timeout=aiohttp.ClientTimeout(total=30),
|
||||
)
|
||||
|
||||
|
||||
async def get_session() -> aiohttp.ClientSession:
|
||||
global _session
|
||||
if _session is None or _session.closed:
|
||||
_session = _new_session()
|
||||
return _session
|
||||
|
||||
|
||||
async def close_session():
|
||||
global _session
|
||||
if _session and not _session.closed:
|
||||
await _session.close()
|
||||
_session = None
|
||||
|
||||
|
||||
async def api_get(url: str, **kwargs) -> Optional[Dict[str, Any]]:
|
||||
return await _request("get", url, **kwargs)
|
||||
|
||||
|
||||
async def api_post(url: str, **kwargs) -> Optional[Dict[str, Any]]:
|
||||
return await _request("post", url, **kwargs)
|
||||
|
||||
|
||||
async def api_patch(url: str, **kwargs) -> Optional[Dict[str, Any]]:
|
||||
return await _request("patch", url, **kwargs)
|
||||
|
||||
|
||||
async def _request(method: str, url: str, **kwargs) -> Optional[Dict[str, Any]]:
|
||||
"""Execute an HTTP request with up to 3 attempts and exponential backoff."""
|
||||
max_retries = 3
|
||||
backoff = 1.0
|
||||
session = await get_session()
|
||||
|
||||
for attempt in range(max_retries):
|
||||
try:
|
||||
async with getattr(session, method)(url, **kwargs) as resp:
|
||||
if resp.content_type == "application/json":
|
||||
return await resp.json()
|
||||
return {"status_code": resp.status}
|
||||
except aiohttp.ClientError as exc:
|
||||
if attempt < max_retries - 1:
|
||||
logger.warning(
|
||||
f"{method.upper()} {url} attempt {attempt + 1} failed: {exc} "
|
||||
f"— retrying in {backoff:.0f}s"
|
||||
)
|
||||
await asyncio.sleep(backoff)
|
||||
backoff *= 2
|
||||
else:
|
||||
logger.error(f"{method.upper()} {url} failed after {max_retries} attempts: {exc}")
|
||||
except Exception as exc:
|
||||
logger.error(f"{method.upper()} {url} unexpected error: {exc}")
|
||||
break
|
||||
|
||||
return None
|
||||
27
core/logger.py
Normal file
27
core/logger.py
Normal file
@@ -0,0 +1,27 @@
|
||||
"""Structured logging for LogiFlow AI — single loguru setup imported by all modules."""
|
||||
import sys
|
||||
import os
|
||||
from pathlib import Path
|
||||
from loguru import logger
|
||||
|
||||
Path("logs").mkdir(exist_ok=True)
|
||||
|
||||
logger.remove()
|
||||
|
||||
logger.add(
|
||||
sys.stderr,
|
||||
format="{time:YYYY-MM-DD HH:mm:ss} | {level: <8} | {name}:{function}:{line} - {message}",
|
||||
level=os.getenv("LOG_LEVEL", "INFO"),
|
||||
colorize=False,
|
||||
)
|
||||
|
||||
logger.add(
|
||||
"logs/logiflow_{time:YYYY-MM-DD}.log",
|
||||
rotation="100 MB",
|
||||
retention="30 days",
|
||||
compression="gz",
|
||||
level="DEBUG",
|
||||
format="{time:YYYY-MM-DD HH:mm:ss} | {level: <8} | {name}:{function}:{line} - {message}",
|
||||
)
|
||||
|
||||
__all__ = ["logger"]
|
||||
299
core/message_bus.py
Normal file
299
core/message_bus.py
Normal file
@@ -0,0 +1,299 @@
|
||||
"""Agent Communication Bus - backed by NATS JetStream with local fallback."""
|
||||
import asyncio
|
||||
import json
|
||||
import uuid
|
||||
from datetime import datetime
|
||||
from typing import Dict, List, Callable, Optional, Any
|
||||
from collections import defaultdict
|
||||
from dataclasses import dataclass
|
||||
from enum import Enum
|
||||
|
||||
import nats
|
||||
import nats.js.errors
|
||||
|
||||
from core.types import AgentMessage, MessageType
|
||||
from core.logger import logger
|
||||
from config.system_config import NATS_HOST, NATS_PORT, NATS_USER, NATS_PASSWORD
|
||||
|
||||
|
||||
class QueuePriority(str, Enum):
|
||||
HIGH = "high"
|
||||
NORMAL = "normal"
|
||||
LOW = "low"
|
||||
|
||||
|
||||
@dataclass
|
||||
class QueuedMessage:
|
||||
message: AgentMessage
|
||||
priority: QueuePriority = QueuePriority.NORMAL
|
||||
retry_count: int = 0
|
||||
max_retries: int = 3
|
||||
|
||||
|
||||
class MessageBus:
|
||||
"""
|
||||
Central message bus backed by NATS JetStream.
|
||||
|
||||
Call `await message_bus.connect()` once at startup before starting agents.
|
||||
Falls back to in-process local dispatch when not connected.
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
self._nc = None # NATS client
|
||||
self._js = None # JetStream context
|
||||
self._nats_subs: list = [] # hold sub refs to prevent GC
|
||||
|
||||
self._subscribers: Dict[MessageType, List[Callable]] = defaultdict(list)
|
||||
self._queues: Dict[str, List[QueuedMessage]] = defaultdict(list)
|
||||
self._agents: Dict[str, Any] = {}
|
||||
self._history: List[AgentMessage] = []
|
||||
self._hooks: Dict[str, List[Callable]] = defaultdict(list)
|
||||
self._lock = asyncio.Lock()
|
||||
|
||||
# ------------------------------------------------------------------ #
|
||||
# Connection #
|
||||
# ------------------------------------------------------------------ #
|
||||
|
||||
async def connect(self):
|
||||
"""Connect to NATS and create the logistics JetStream stream."""
|
||||
self._nc = await nats.connect(
|
||||
servers=[f"nats://{NATS_HOST}:{NATS_PORT}"],
|
||||
user=NATS_USER,
|
||||
password=NATS_PASSWORD,
|
||||
error_cb=self._on_nats_error,
|
||||
reconnected_cb=self._on_reconnect,
|
||||
max_reconnect_attempts=10,
|
||||
)
|
||||
self._js = self._nc.jetstream()
|
||||
|
||||
try:
|
||||
await self._js.add_stream(name="logistics", subjects=["logistics.>"])
|
||||
logger.info("NATS stream 'logistics' created")
|
||||
except nats.js.errors.BadRequestError:
|
||||
pass # stream already exists
|
||||
except Exception as e:
|
||||
logger.warning(f"Could not create NATS stream: {e}")
|
||||
|
||||
# Set up push subscriptions for all pre-registered message types
|
||||
for msg_type in list(self._subscribers.keys()):
|
||||
await self._setup_type_sub(msg_type)
|
||||
|
||||
# Set up direct-message subscriptions for all pre-registered agents
|
||||
for agent_id in list(self._agents.keys()):
|
||||
await self._setup_agent_sub(agent_id)
|
||||
|
||||
logger.info("MessageBus connected to NATS JetStream")
|
||||
|
||||
async def _on_nats_error(self, err):
|
||||
logger.warning(f"NATS error: {err}")
|
||||
|
||||
async def _on_reconnect(self):
|
||||
logger.info("NATS reconnected")
|
||||
|
||||
async def disconnect(self):
|
||||
"""Drain and close NATS connection."""
|
||||
if self._nc:
|
||||
await self._nc.drain()
|
||||
|
||||
# ------------------------------------------------------------------ #
|
||||
# Agent registry #
|
||||
# ------------------------------------------------------------------ #
|
||||
|
||||
def register_agent(self, agent):
|
||||
self._agents[agent.agent_id] = agent
|
||||
logger.debug(f"Agent registered: {agent.agent_id} ({agent.agent_type})")
|
||||
if self._js is not None:
|
||||
asyncio.create_task(self._setup_agent_sub(agent.agent_id))
|
||||
|
||||
def unregister_agent(self, agent_id: str):
|
||||
self._agents.pop(agent_id, None)
|
||||
logger.debug(f"Agent unregistered: {agent_id}")
|
||||
|
||||
# ------------------------------------------------------------------ #
|
||||
# Pub / Sub #
|
||||
# ------------------------------------------------------------------ #
|
||||
|
||||
def subscribe(self, message_type: MessageType, callback: Callable):
|
||||
"""Subscribe to a broadcast message type."""
|
||||
self._subscribers[message_type].append(callback)
|
||||
if self._js is not None:
|
||||
asyncio.create_task(self._setup_type_sub(message_type))
|
||||
|
||||
def unsubscribe(self, message_type: MessageType, callback: Callable):
|
||||
if message_type in self._subscribers:
|
||||
try:
|
||||
self._subscribers[message_type].remove(callback)
|
||||
except ValueError:
|
||||
pass
|
||||
|
||||
async def publish(self, message: AgentMessage):
|
||||
"""Publish a message — via NATS if connected, otherwise local dispatch."""
|
||||
async with self._lock:
|
||||
self._history.append(message)
|
||||
if len(self._history) > 1000:
|
||||
self._history = self._history[-500:]
|
||||
|
||||
if self._js is not None:
|
||||
await self._nats_publish(message)
|
||||
else:
|
||||
await self._local_dispatch(message)
|
||||
|
||||
await self._trigger_hook(f"on_{message.message_type.value}", message)
|
||||
logger.debug(f"[{message.sender}] -> [{message.recipient}]: {message.message_type.value}")
|
||||
|
||||
async def send_to_agent(
|
||||
self,
|
||||
sender: str,
|
||||
recipient: str,
|
||||
message_type: MessageType,
|
||||
payload: Dict[str, Any],
|
||||
correlation_id: Optional[str] = None,
|
||||
) -> str:
|
||||
message_id = str(uuid.uuid4())
|
||||
message = AgentMessage(
|
||||
message_id=message_id,
|
||||
sender=sender,
|
||||
recipient=recipient,
|
||||
message_type=message_type,
|
||||
payload=payload,
|
||||
timestamp=datetime.now(),
|
||||
correlation_id=correlation_id or message_id,
|
||||
)
|
||||
await self.publish(message)
|
||||
return message_id
|
||||
|
||||
async def broadcast(
|
||||
self,
|
||||
sender: str,
|
||||
message_type: MessageType,
|
||||
payload: Dict[str, Any],
|
||||
correlation_id: Optional[str] = None,
|
||||
) -> str:
|
||||
return await self.send_to_agent(sender, "ALL", message_type, payload, correlation_id)
|
||||
|
||||
async def get_messages(self, agent_id: str) -> List[AgentMessage]:
|
||||
async with self._lock:
|
||||
queued = self._queues.pop(agent_id, [])
|
||||
return [q.message for q in queued]
|
||||
|
||||
async def peek_messages(self, agent_id: str) -> List[AgentMessage]:
|
||||
return [q.message for q in self._queues.get(agent_id, [])]
|
||||
|
||||
# ------------------------------------------------------------------ #
|
||||
# Hooks #
|
||||
# ------------------------------------------------------------------ #
|
||||
|
||||
def add_hook(self, event: str, callback: Callable):
|
||||
self._hooks[event].append(callback)
|
||||
|
||||
async def _trigger_hook(self, event: str, message: AgentMessage):
|
||||
for callback in self._hooks.get(event, []):
|
||||
try:
|
||||
if asyncio.iscoroutinefunction(callback):
|
||||
await callback(message)
|
||||
else:
|
||||
callback(message)
|
||||
except Exception as e:
|
||||
logger.error(f"Hook error [{event}]: {e}")
|
||||
|
||||
# ------------------------------------------------------------------ #
|
||||
# Agent / history helpers #
|
||||
# ------------------------------------------------------------------ #
|
||||
|
||||
def get_agent(self, agent_id: str):
|
||||
return self._agents.get(agent_id)
|
||||
|
||||
def get_all_agents(self) -> Dict[str, Any]:
|
||||
return self._agents.copy()
|
||||
|
||||
def get_messages_by_type(self, message_type: MessageType) -> List[AgentMessage]:
|
||||
return [m for m in self._history if m.message_type == message_type]
|
||||
|
||||
def get_messages_by_sender(self, sender: str) -> List[AgentMessage]:
|
||||
return [m for m in self._history if m.sender == sender]
|
||||
|
||||
def clear_history(self):
|
||||
self._history = []
|
||||
|
||||
# ------------------------------------------------------------------ #
|
||||
# Internal: NATS helpers #
|
||||
# ------------------------------------------------------------------ #
|
||||
|
||||
async def _nats_publish(self, message: AgentMessage):
|
||||
if message.recipient != "ALL":
|
||||
subject = f"logistics.direct.{message.recipient}"
|
||||
else:
|
||||
subject = f"logistics.{message.message_type.value}"
|
||||
try:
|
||||
await self._js.publish(subject, message.to_json().encode())
|
||||
except Exception as e:
|
||||
logger.warning(f"NATS publish error: {e} — falling back to local dispatch")
|
||||
await self._local_dispatch(message)
|
||||
|
||||
async def _setup_type_sub(self, message_type: MessageType):
|
||||
"""Create a JetStream push subscription for a broadcast message type."""
|
||||
subject = f"logistics.{message_type.value}"
|
||||
durable = f"logiflow-{message_type.value}"
|
||||
callbacks = self._subscribers[message_type]
|
||||
|
||||
async def handler(msg):
|
||||
try:
|
||||
agent_msg = AgentMessage.from_json(msg.data.decode())
|
||||
for cb in list(callbacks):
|
||||
try:
|
||||
if asyncio.iscoroutinefunction(cb):
|
||||
await cb(agent_msg)
|
||||
else:
|
||||
cb(agent_msg)
|
||||
except Exception as e:
|
||||
logger.error(f"Subscriber callback error: {e}")
|
||||
except Exception as e:
|
||||
logger.error(f"NATS type-sub decode error [{message_type.value}]: {e}")
|
||||
finally:
|
||||
await msg.ack()
|
||||
|
||||
try:
|
||||
sub = await self._js.subscribe(subject, durable=durable, cb=handler)
|
||||
self._nats_subs.append(sub)
|
||||
except Exception as e:
|
||||
logger.warning(f"Could not subscribe to NATS subject {subject}: {e}")
|
||||
|
||||
async def _setup_agent_sub(self, agent_id: str):
|
||||
"""Create a JetStream push subscription for directed messages to an agent."""
|
||||
subject = f"logistics.direct.{agent_id}"
|
||||
durable = f"logiflow-direct-{agent_id}"
|
||||
|
||||
async def handler(msg):
|
||||
try:
|
||||
agent_msg = AgentMessage.from_json(msg.data.decode())
|
||||
async with self._lock:
|
||||
self._queues[agent_id].append(QueuedMessage(agent_msg))
|
||||
except Exception as e:
|
||||
logger.error(f"NATS agent-sub decode error [{agent_id}]: {e}")
|
||||
finally:
|
||||
await msg.ack()
|
||||
|
||||
try:
|
||||
sub = await self._js.subscribe(subject, durable=durable, cb=handler)
|
||||
self._nats_subs.append(sub)
|
||||
except Exception as e:
|
||||
logger.warning(f"Could not subscribe to NATS subject {subject}: {e}")
|
||||
|
||||
async def _local_dispatch(self, message: AgentMessage):
|
||||
"""In-process dispatch used when NATS is not connected."""
|
||||
if message.recipient != "ALL":
|
||||
async with self._lock:
|
||||
self._queues[message.recipient].append(QueuedMessage(message))
|
||||
else:
|
||||
for callback in list(self._subscribers.get(message.message_type, [])):
|
||||
try:
|
||||
if asyncio.iscoroutinefunction(callback):
|
||||
await callback(message)
|
||||
else:
|
||||
callback(message)
|
||||
except Exception as e:
|
||||
logger.error(f"Local dispatch callback error: {e}")
|
||||
|
||||
|
||||
# Global message bus instance
|
||||
message_bus = MessageBus()
|
||||
210
core/types.py
Normal file
210
core/types.py
Normal file
@@ -0,0 +1,210 @@
|
||||
"""Core types and message protocols for agent communication."""
|
||||
from enum import Enum
|
||||
from typing import Optional, Dict, Any, List
|
||||
from datetime import datetime
|
||||
from pydantic import BaseModel, Field
|
||||
from dataclasses import dataclass, asdict
|
||||
import json
|
||||
|
||||
|
||||
class MessageType(str, Enum):
|
||||
"""Types of messages agents can send"""
|
||||
ORDER_RECEIVED = "ORDER_RECEIVED"
|
||||
ORDER_VALIDATED = "ORDER_VALIDATED"
|
||||
ORDER_ASSIGNED = "ORDER_ASSIGNED"
|
||||
ORDER_IN_TRANSIT = "ORDER_IN_TRANSIT"
|
||||
ORDER_DELIVERED = "ORDER_DELIVERED"
|
||||
ORDER_CANCELLED = "ORDER_CANCELLED"
|
||||
ORDER_DELAYED = "ORDER_DELAYED"
|
||||
HUB_STATUS_UPDATE = "HUB_STATUS_UPDATE"
|
||||
VEHICLE_ASSIGNED = "VEHICLE_ASSIGNED"
|
||||
ROUTE_OPTIMIZED = "ROUTE_OPTIMIZED"
|
||||
NOTIFICATION_SENT = "NOTIFICATION_SENT"
|
||||
EXCEPTION_DETECTED = "EXCEPTION_DETECTED"
|
||||
AGENT_TASK = "AGENT_TASK"
|
||||
AGENT_RESPONSE = "AGENT_RESPONSE"
|
||||
HEARTBEAT = "HEARTBEAT"
|
||||
STATUS_QUERY = "STATUS_QUERY"
|
||||
|
||||
|
||||
class Priority(str, Enum):
|
||||
"""Order priority levels"""
|
||||
LOW = "low"
|
||||
MEDIUM = "medium"
|
||||
HIGH = "high"
|
||||
URGENT = "urgent"
|
||||
|
||||
|
||||
class OrderStatus(str, Enum):
|
||||
"""Order lifecycle status"""
|
||||
RECEIVED = "received"
|
||||
VALIDATING = "validating"
|
||||
VALIDATED = "validated"
|
||||
PENDING_DISPATCH = "pending_dispatch"
|
||||
ASSIGNED = "assigned"
|
||||
PICKED_UP = "picked_up"
|
||||
IN_TRANSIT = "in_transit"
|
||||
AT_HUB = "at_hub"
|
||||
OUT_FOR_DELIVERY = "out_for_delivery"
|
||||
DELIVERED = "delivered"
|
||||
FAILED = "failed"
|
||||
CANCELLED = "cancelled"
|
||||
RETURNED = "returned"
|
||||
|
||||
|
||||
class ZoneType(str, Enum):
|
||||
"""Zone classification"""
|
||||
HUB_TO_HUB = "hub_to_hub"
|
||||
HUB_TO_SPOKE = "hub_to_spoke"
|
||||
SPOKE_TO_HUB = "spoke_to_hub"
|
||||
LAST_MILE = "last_mile"
|
||||
|
||||
|
||||
@dataclass
|
||||
class Address:
|
||||
"""Address information"""
|
||||
street: str
|
||||
city: str
|
||||
state: str
|
||||
pincode: str
|
||||
landmark: Optional[str] = None
|
||||
latitude: Optional[float] = None
|
||||
longitude: Optional[float] = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class OrderItem:
|
||||
"""Individual item in an order"""
|
||||
sku: str
|
||||
name: str
|
||||
quantity: int
|
||||
weight: float # in kg
|
||||
dimensions: Optional[Dict[str, float]] = None # LxWxH in cm
|
||||
|
||||
|
||||
@dataclass
|
||||
class Order:
|
||||
"""Complete order data"""
|
||||
order_id: str
|
||||
customer_name: str
|
||||
customer_phone: str
|
||||
customer_email: Optional[str]
|
||||
pickup_address: Address
|
||||
delivery_address: Address
|
||||
items: List[OrderItem]
|
||||
priority: Priority = Priority.MEDIUM
|
||||
special_instructions: Optional[str] = None
|
||||
created_at: datetime = None
|
||||
estimated_delivery: Optional[datetime] = None
|
||||
|
||||
def __post_init__(self):
|
||||
if self.created_at is None:
|
||||
self.created_at = datetime.now()
|
||||
|
||||
def to_dict(self) -> Dict:
|
||||
return {
|
||||
'order_id': self.order_id,
|
||||
'customer_name': self.customer_name,
|
||||
'customer_phone': self.customer_phone,
|
||||
'customer_email': self.customer_email,
|
||||
'pickup_address': asdict(self.pickup_address),
|
||||
'delivery_address': asdict(self.delivery_address),
|
||||
'items': [asdict(item) for item in self.items],
|
||||
'priority': self.priority.value,
|
||||
'special_instructions': self.special_instructions,
|
||||
'created_at': self.created_at.isoformat(),
|
||||
'estimated_delivery': self.estimated_delivery.isoformat() if self.estimated_delivery else None
|
||||
}
|
||||
|
||||
|
||||
@dataclass
|
||||
class Hub:
|
||||
"""Hub information"""
|
||||
hub_id: str
|
||||
name: str
|
||||
location: str
|
||||
pincode_prefix: str # First 3 digits of pincode
|
||||
capacity: int
|
||||
current_load: int = 0
|
||||
zones_covered: List[str] = None
|
||||
|
||||
def __post_init__(self):
|
||||
if self.zones_covered is None:
|
||||
self.zones_covered = []
|
||||
|
||||
|
||||
@dataclass
|
||||
class Vehicle:
|
||||
"""Vehicle information"""
|
||||
vehicle_id: str
|
||||
vehicle_type: str # bike, scooter, van, truck
|
||||
capacity_kg: float
|
||||
capacity_volume: float # in cubic meters
|
||||
current_location: str
|
||||
current_load: float = 0
|
||||
status: str = "available" # available, in_transit, maintenance, offline
|
||||
assigned_hub: Optional[str] = None
|
||||
current_route: Optional[List[str]] = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class AgentMessage:
|
||||
"""Standard message format for agent communication"""
|
||||
message_id: str
|
||||
sender: str
|
||||
recipient: str # "ALL" for broadcast
|
||||
message_type: MessageType
|
||||
payload: Dict[str, Any]
|
||||
timestamp: datetime
|
||||
correlation_id: Optional[str] = None # For tracing related messages
|
||||
|
||||
def to_json(self) -> str:
|
||||
data = {
|
||||
'message_id': self.message_id,
|
||||
'sender': self.sender,
|
||||
'recipient': self.recipient,
|
||||
'message_type': self.message_type.value,
|
||||
'payload': self.payload,
|
||||
'timestamp': self.timestamp.isoformat(),
|
||||
'correlation_id': self.correlation_id
|
||||
}
|
||||
return json.dumps(data)
|
||||
|
||||
@classmethod
|
||||
def from_json(cls, json_str: str) -> 'AgentMessage':
|
||||
data = json.loads(json_str)
|
||||
return cls(
|
||||
message_id=data['message_id'],
|
||||
sender=data['sender'],
|
||||
recipient=data['recipient'],
|
||||
message_type=MessageType(data['message_type']),
|
||||
payload=data['payload'],
|
||||
timestamp=datetime.fromisoformat(data['timestamp']),
|
||||
correlation_id=data.get('correlation_id')
|
||||
)
|
||||
|
||||
|
||||
class AgentTask(BaseModel):
|
||||
"""Task assigned to an agent"""
|
||||
task_id: str
|
||||
agent_type: str
|
||||
task_type: str
|
||||
description: str = ""
|
||||
priority: Priority = Priority.MEDIUM
|
||||
data: Dict[str, Any] = Field(default_factory=dict)
|
||||
created_at: datetime = Field(default_factory=datetime.now)
|
||||
status: str = "pending" # pending, in_progress, completed, failed
|
||||
result: Optional[Dict] = None
|
||||
error: Optional[str] = None
|
||||
|
||||
|
||||
class AgentState(BaseModel):
|
||||
"""Current state of an agent"""
|
||||
agent_id: str
|
||||
agent_type: str
|
||||
status: str # idle, working, waiting, error
|
||||
current_task: Optional[str] = None
|
||||
tasks_completed: int = 0
|
||||
tasks_failed: int = 0
|
||||
last_active: datetime = Field(default_factory=datetime.now)
|
||||
performance_metrics: Dict[str, Any] = Field(default_factory=dict)
|
||||
Reference in New Issue
Block a user