265 lines
9.4 KiB
Python
265 lines
9.4 KiB
Python
"""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}
|
|
))
|
|
|
|
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:]
|