CHKIM79/multi-ai-agentic-system

1"""2Planning Agent - Analyzes options and creates strategies3"""4import asyncio5from typing import Dict, Any6from core.base_agent import BaseAgent, AgentMessage, TaskResult, TaskStatus7 8class PlanningAgent(BaseAgent):9    """Agent specialized in planning and strategy creation"""10    11    def __init__(self):12        super().__init__("planning_agent", ["planning", "analysis", "strategy"])13        self.flight_database = self._initialize_flight_data()14    15    def _initialize_flight_data(self) -> Dict[str, Any]:16        """Mock flight database for demonstration"""17        return {18            "routes": {19                "NYC": {20                    "LAX": [21                        {"airline": "Delta", "price": 350, "duration": "6h 15m", "stops": 0},22                        {"airline": "United", "price": 320, "duration": "6h 30m", "stops": 0},23                        {"airline": "JetBlue", "price": 280, "duration": "7h 45m", "stops": 1}24                    ],25                    "SFO": [26                        {"airline": "American", "price": 380, "duration": "6h 45m", "stops": 0},27                        {"airline": "United", "price": 340, "duration": "6h 20m", "stops": 0}28                    ]29                }30            }31        }32    33    async def process_task(self, message: AgentMessage) -> TaskResult:34        """Process planning-related tasks"""35        start_time = asyncio.get_event_loop().time()36        37        try:38            if message.message_type == "planning":39                result = await self._plan_flight_options(message.data)40            elif message.message_type == "analysis":41                result = await self._analyze_request(message.data)42            else:43                raise ValueError(f"Unknown task type: {message.message_type}")44            45            return TaskResult(46                task_id=message.task_id,47                agent_id=self.agent_id,48                status=TaskStatus.COMPLETED,49                result=result,50                execution_time=asyncio.get_event_loop().time() - start_time51            )52            53        except Exception as e:54            return TaskResult(55                task_id=message.task_id,56                agent_id=self.agent_id,57                status=TaskStatus.FAILED,58                result={},59                error_message=str(e),60                execution_time=asyncio.get_event_loop().time() - start_time61            )62    63    async def _plan_flight_options(self, data: Dict[str, Any]) -> Dict[str, Any]:64        """Plan optimal flight options based on request"""65        request = data.get("request", "").lower()66        67        # Extract destination from request (simplified parsing)68        destination = None69        if "new york" in request or "nyc" in request:70            destination = "NYC"71        elif "los angeles" in request or "lax" in request:72            destination = "LAX"73        elif "san francisco" in request or "sfo" in request:74            destination = "SFO"75        76        if not destination:77            return {78                "findings": ["Could not determine destination from request"],79                "options": [],80                "recommendation": "Please specify a clear destination"81            }82        83        # Find available routes (mock logic)84        available_routes = []85        for origin, destinations in self.flight_database["routes"].items():86            if destination in destinations:87                for flight in destinations[destination]:88                    available_routes.append({89                        "route": f"{origin} → {destination}",90                        **flight91                    })92        93        # Sort by price (could implement more sophisticated ranking)94        available_routes.sort(key=lambda x: x["price"])95        96        # Create recommendations97        best_value = available_routes[0] if available_routes else None98        fastest = min(available_routes, key=lambda x: self._parse_duration(x["duration"])) if available_routes else None99        100        return {101            "findings": [102                f"Found {len(available_routes)} flight options to {destination}",103                f"Price range: ${min(r['price'] for r in available_routes)} - ${max(r['price'] for r in available_routes)}" if available_routes else "No flights found"104            ],105            "options": available_routes[:5],  # Top 5 options106            "recommendation": {107                "best_value": best_value,108                "fastest": fastest,109                "reasoning": "Recommendations based on price and duration optimization"110            }111        }112    113    async def _analyze_request(self, data: Dict[str, Any]) -> Dict[str, Any]:114        """Analyze user request for planning purposes"""115        request = data.get("request", "")116        117        # Extract key information from request118        analysis = {119            "request_type": "travel" if any(word in request.lower() for word in ["flight", "travel", "book"]) else "general",120            "urgency": "high" if any(word in request.lower() for word in ["urgent", "asap", "immediately"]) else "normal",121            "budget_conscious": "yes" if any(word in request.lower() for word in ["cheap", "budget", "affordable"]) else "unknown",122            "extracted_entities": []123        }124        125        # Simple entity extraction126        cities = ["new york", "los angeles", "san francisco", "chicago", "miami"]127        for city in cities:128            if city in request.lower():129                analysis["extracted_entities"].append({"type": "destination", "value": city})130        131        return {132            "findings": [133                f"Request type: {analysis['request_type']}",134                f"Urgency level: {analysis['urgency']}",135                f"Entities found: {len(analysis['extracted_entities'])}"136            ],137            "analysis": analysis,138            "recommendation": "Proceed with specialized planning based on request type"139        }140    141    def _parse_duration(self, duration_str: str) -> int:142        """Parse duration string to minutes for comparison"""143        # Simple parser for "6h 15m" format144        parts = duration_str.replace("h", "").replace("m", "").split()145        hours = int(parts[0]) if len(parts) > 0 else 0146        minutes = int(parts[1]) if len(parts) > 1 else 0147        return hours * 60 + minutes148    149    def get_agent_info(self) -> Dict[str, Any]:150        """Return planning agent information"""151        return {152            "agent_id": self.agent_id,153            "type": "planning",154            "capabilities": self.capabilities,155            "specialization": "Flight planning and travel optimization",156            "database_size": len(self.flight_database["routes"])157        }158