Agentic AI Systems Design

import asyncio
import json
import time
from typing import Dict, List, Any, Optional, Callable, Union
from dataclasses import dataclass, asdict
from enum import Enum
from abc import ABC, abstractmethod
import uuid
import logging
from datetime import datetime, timedelta

# Agent Types and States
class AgentType(Enum):
    COORDINATOR = "coordinator"
    RESEARCHER = "researcher"
    ANALYST = "analyst"
    EXECUTOR = "executor"
    REVIEWER = "reviewer"

class TaskStatus(Enum):
    PENDING = "pending"
    IN_PROGRESS = "in_progress" 
    COMPLETED = "completed"
    FAILED = "failed"
    DELEGATED = "delegated"

class AgentState(Enum):
    IDLE = "idle"
    WORKING = "working"
    WAITING = "waiting"
    ERROR = "error"

@dataclass
class Task:
    """Represents a task that can be executed by agents"""
    id: str
    type: str
    description: str
    requirements: Dict[str, Any]
    dependencies: List[str]
    assigned_agent: Optional[str] = None
    status: TaskStatus = TaskStatus.PENDING
    result: Optional[Dict[str, Any]] = None
    created_at: datetime = None
    started_at: Optional[datetime] = None
    completed_at: Optional[datetime] = None
    error_message: Optional[str] = None
    
    def __post_init__(self):
        if self.created_at is None:
            self.created_at = datetime.now()

@dataclass
class Message:
    """Inter-agent communication message"""
    id: str
    sender_id: str
    recipient_id: str
    message_type: str
    content: Dict[str, Any]
    timestamp: datetime
    conversation_id: Optional[str] = None

class Agent(ABC):
    """Base class for all agents in the system"""
    
    def __init__(self, agent_id: str, agent_type: AgentType, capabilities: List[str]):
        self.agent_id = agent_id
        self.agent_type = agent_type
        self.capabilities = capabilities
        self.state = AgentState.IDLE
        self.current_task: Optional[Task] = None
        self.memory: Dict[str, Any] = {}
        self.message_queue: List[Message] = []
        self.tools: Dict[str, Callable] = {}
        self.logger = logging.getLogger(f"Agent-{agent_id}")
        
        # Performance metrics
        self.tasks_completed = 0
        self.tasks_failed = 0
        self.avg_task_duration = 0.0
        self.last_activity = datetime.now()
    
    @abstractmethod
    async def process_task(self, task: Task) -> Dict[str, Any]:
        """Process a assigned task and return results"""
        pass
    
    async def receive_message(self, message: Message) -> Optional[Message]:
        """Receive and process a message from another agent"""
        self.message_queue.append(message)
        self.logger.info(f"Received message from {message.sender_id}: {message.message_type}")
        
        # Process message immediately if agent is idle
        if self.state == AgentState.IDLE:
            return await self._process_message(message)
        return None
    
    async def _process_message(self, message: Message) -> Optional[Message]:
        """Process incoming message and potentially send response"""
        if message.message_type == "task_delegation":
            task_data = message.content.get("task")
            if task_data:
                task = Task(**task_data)
                result = await self.execute_task(task)
                
                # Send result back to sender
                response = Message(
                    id=str(uuid.uuid4()),
                    sender_id=self.agent_id,
                    recipient_id=message.sender_id,
                    message_type="task_result",
                    content={"task_id": task.id, "result": result},
                    timestamp=datetime.now(),
                    conversation_id=message.conversation_id
                )
                return response
        
        elif message.message_type == "information_request":
            # Provide information from memory
            requested_info = message.content.get("query")
            if requested_info in self.memory:
                response = Message(
                    id=str(uuid.uuid4()),
                    sender_id=self.agent_id,
                    recipient_id=message.sender_id,
                    message_type="information_response",
                    content={"query": requested_info, "data": self.memory[requested_info]},
                    timestamp=datetime.now(),
                    conversation_id=message.conversation_id
                )
                return response
        
        return None
    
    async def execute_task(self, task: Task) -> Dict[str, Any]:
        """Execute a task and track performance"""
        self.current_task = task
        self.state = AgentState.WORKING
        task.assigned_agent = self.agent_id
        task.started_at = datetime.now()
        
        try:
            self.logger.info(f"Starting task: {task.description}")
            result = await self.process_task(task)
            
            task.status = TaskStatus.COMPLETED
            task.result = result
            task.completed_at = datetime.now()
            
            # Update performance metrics
            duration = (task.completed_at - task.started_at).total_seconds()
            self.tasks_completed += 1
            self.avg_task_duration = (self.avg_task_duration * (self.tasks_completed - 1) + duration) / self.tasks_completed
            
            self.state = AgentState.IDLE
            self.current_task = None
            self.last_activity = datetime.now()
            
            self.logger.info(f"Completed task {task.id} in {duration:.2f}s")
            return result
            
        except Exception as e:
            task.status = TaskStatus.FAILED
            task.error_message = str(e)
            task.completed_at = datetime.now()
            
            self.tasks_failed += 1
            self.state = AgentState.ERROR
            self.logger.error(f"Task {task.id} failed: {e}")
            
            return {"error": str(e), "status": "failed"}
    
    def add_tool(self, name: str, tool_func: Callable):
        """Add a tool that the agent can use"""
        self.tools[name] = tool_func
        self.logger.info(f"Added tool: {name}")
    
    async def use_tool(self, tool_name: str, **kwargs) -> Any:
        """Use a registered tool"""
        if tool_name in self.tools:
            try:
                return await self.tools[tool_name](**kwargs)
            except Exception as e:
                self.logger.error(f"Tool {tool_name} failed: {e}")
                raise
        else:
            raise ValueError(f"Tool {tool_name} not available")

class ResearchAgent(Agent):
    """Agent specialized in research and information gathering"""
    
    def __init__(self, agent_id: str):
        super().__init__(agent_id, AgentType.RESEARCHER, ["web_search", "data_analysis", "report_generation"])
    
    async def process_task(self, task: Task) -> Dict[str, Any]:
        """Process research tasks"""
        query = task.requirements.get("query", "")
        sources = task.requirements.get("sources", ["web", "database"])
        
        # Simulate research process
        await asyncio.sleep(2.0)  # Simulate research time
        
        research_findings = {
            "query": query,
            "sources_searched": sources,
            "findings": [
                {"title": "Market Analysis Report", "relevance": 0.95, "summary": "Key market trends..."},
                {"title": "Competitor Intelligence", "relevance": 0.87, "summary": "Competitive landscape..."},
                {"title": "Industry Forecasts", "relevance": 0.82, "summary": "Future predictions..."}
            ],
            "confidence_score": 0.89,
            "research_duration": 2.0
        }
        
        # Store findings in memory for future use
        self.memory[f"research_{task.id}"] = research_findings
        
        return research_findings

class AnalystAgent(Agent):
    """Agent specialized in data analysis and insights generation"""
    
    def __init__(self, agent_id: str):
        super().__init__(agent_id, AgentType.ANALYST, ["data_processing", "statistical_analysis", "visualization"])
    
    async def process_task(self, task: Task) -> Dict[str, Any]:
        """Process analysis tasks"""
        data_source = task.requirements.get("data_source", "")
        analysis_type = task.requirements.get("analysis_type", "descriptive")
        
        # Simulate analysis process
        await asyncio.sleep(1.5)
        
        analysis_result = {
            "data_source": data_source,
            "analysis_type": analysis_type,
            "insights": [
                {"metric": "Revenue Growth", "value": "23.5%", "trend": "increasing"},
                {"metric": "Customer Churn", "value": "3.2%", "trend": "decreasing"},
                {"metric": "Market Share", "value": "15.7%", "trend": "stable"}
            ],
            "recommendations": [
                "Focus on customer retention strategies",
                "Expand into emerging markets",
                "Optimize pricing strategy"
            ],
            "confidence": 0.91
        }
        
        self.memory[f"analysis_{task.id}"] = analysis_result
        return analysis_result

class ExecutorAgent(Agent):
    """Agent specialized in executing actions and implementing solutions"""
    
    def __init__(self, agent_id: str):
        super().__init__(agent_id, AgentType.EXECUTOR, ["automation", "api_integration", "process_execution"])
    
    async def process_task(self, task: Task) -> Dict[str, Any]:
        """Process execution tasks"""
        action_type = task.requirements.get("action", "")
        parameters = task.requirements.get("parameters", {})
        
        # Simulate execution process
        await asyncio.sleep(1.0)
        
        execution_result = {
            "action": action_type,
            "parameters": parameters,
            "status": "completed",
            "execution_time": 1.0,
            "outputs": {
                "files_processed": 150,
                "api_calls_made": 23,
                "records_updated": 1247
            },
            "success_rate": 0.98
        }
        
        return execution_result

class CoordinatorAgent(Agent):
    """Agent responsible for coordinating other agents and managing workflows"""
    
    def __init__(self, agent_id: str):
        super().__init__(agent_id, AgentType.COORDINATOR, ["task_planning", "agent_management", "workflow_orchestration"])
        self.managed_agents: Dict[str, Agent] = {}
        self.active_workflows: Dict[str, Dict[str, Any]] = {}
    
    def add_managed_agent(self, agent: Agent):
        """Add an agent to be managed by this coordinator"""
        self.managed_agents[agent.agent_id] = agent
        self.logger.info(f"Added managed agent: {agent.agent_id}")
    
    async def process_task(self, task: Task) -> Dict[str, Any]:
        """Process coordination tasks - typically involves delegating to other agents"""
        workflow_id = str(uuid.uuid4())
        
        # Decompose complex task into subtasks
        subtasks = await self._decompose_task(task)
        
        # Execute subtasks in parallel or sequence as needed
        workflow_result = await self._execute_workflow(workflow_id, subtasks)
        
        return {
            "workflow_id": workflow_id,
            "original_task": task.id,
            "subtasks_completed": len(subtasks),
            "workflow_result": workflow_result,
            "coordination_overhead": 0.15  # 15% overhead for coordination
        }
    
    async def _decompose_task(self, task: Task) -> List[Task]:
        """Decompose a complex task into simpler subtasks"""
        # Simulate task decomposition logic
        subtasks = []
        
        if "business_analysis" in task.type:
            # Create research subtask
            research_task = Task(
                id=f"{task.id}_research",
                type="research",
                description=f"Research for {task.description}",
                requirements={"query": task.requirements.get("topic", ""), "sources": ["web", "database"]},
                dependencies=[]
            )
            
            # Create analysis subtask
            analysis_task = Task(
                id=f"{task.id}_analysis",
                type="analysis",
                description=f"Analyze data for {task.description}",
                requirements={"data_source": "research_output", "analysis_type": "comprehensive"},
                dependencies=[research_task.id]
            )
            
            # Create execution subtask
            execution_task = Task(
                id=f"{task.id}_execution",
                type="execution",
                description=f"Execute recommendations for {task.description}",
                requirements={"action": "implement_strategy", "parameters": {}},
                dependencies=[analysis_task.id]
            )
            
            subtasks = [research_task, analysis_task, execution_task]
        
        return subtasks
    
    async def _execute_workflow(self, workflow_id: str, subtasks: List[Task]) -> Dict[str, Any]:
        """Execute a workflow with multiple subtasks"""
        workflow_start = datetime.now()
        completed_tasks = []
        workflow_results = {}
        
        # Simple sequential execution (could be enhanced with parallel execution)
        for subtask in subtasks:
            # Find appropriate agent for the subtask
            assigned_agent = self._find_suitable_agent(subtask.type)
            
            if assigned_agent:
                # Delegate task to agent
                result = await assigned_agent.execute_task(subtask)
                workflow_results[subtask.id] = result
                completed_tasks.append(subtask)
            else:
                self.logger.warning(f"No suitable agent found for task: {subtask.type}")
        
        workflow_duration = (datetime.now() - workflow_start).total_seconds()
        
        return {
            "duration": workflow_duration,
            "tasks_completed": len(completed_tasks),
            "results": workflow_results,
            "success_rate": len(completed_tasks) / len(subtasks)
        }
    
    def _find_suitable_agent(self, task_type: str) -> Optional[Agent]:
        """Find the most suitable agent for a task type"""
        for agent in self.managed_agents.values():
            if agent.state == AgentState.IDLE:
                if task_type == "research" and agent.agent_type == AgentType.RESEARCHER:
                    return agent
                elif task_type == "analysis" and agent.agent_type == AgentType.ANALYST:
                    return agent
                elif task_type == "execution" and agent.agent_type == AgentType.EXECUTOR:
                    return agent
        return None

class AgenticSystem:
    """Main system orchestrating all agents"""
    
    def __init__(self):
        self.agents: Dict[str, Agent] = {}
        self.task_queue: List[Task] = []
        self.completed_tasks: List[Task] = []
        self.system_metrics = {
            "total_tasks": 0,
            "completed_tasks": 0,
            "failed_tasks": 0,
            "avg_completion_time": 0.0,
            "system_utilization": 0.0
        }
        self.logger = logging.getLogger("AgenticSystem")
    
    def add_agent(self, agent: Agent):
        """Add an agent to the system"""
        self.agents[agent.agent_id] = agent
        self.logger.info(f"Added agent {agent.agent_id} of type {agent.agent_type.value}")
    
    async def submit_task(self, task: Task) -> str:
        """Submit a task to be executed by the system"""
        self.task_queue.append(task)
        self.system_metrics["total_tasks"] += 1
        self.logger.info(f"Task {task.id} submitted: {task.description}")
        
        # Immediately try to assign task if suitable agent is available
        await self._process_task_queue()
        
        return task.id
    
    async def _process_task_queue(self):
        """Process pending tasks in the queue"""
        for task in self.task_queue[:]:  # Copy list to avoid modification during iteration
            if task.status == TaskStatus.PENDING:
                suitable_agent = self._find_available_agent(task)
                
                if suitable_agent:
                    self.task_queue.remove(task)
                    task.status = TaskStatus.IN_PROGRESS
                    
                    # Execute task asynchronously
                    asyncio.create_task(self._monitor_task_execution(task, suitable_agent))
    
    def _find_available_agent(self, task: Task) -> Optional[Agent]:
        """Find an available agent suitable for the task"""
        # Prefer coordinator for complex tasks
        if "business_analysis" in task.type or len(task.dependencies) > 0:
            for agent in self.agents.values():
                if agent.agent_type == AgentType.COORDINATOR and agent.state == AgentState.IDLE:
                    return agent
        
        # Otherwise find specialized agent
        for agent in self.agents.values():
            if agent.state == AgentState.IDLE:
                if task.type in agent.capabilities or task.type in [cap.lower() for cap in agent.capabilities]:
                    return agent
        
        return None
    
    async def _monitor_task_execution(self, task: Task, agent: Agent):
        """Monitor task execution and update system metrics"""
        start_time = datetime.now()
        
        try:
            result = await agent.execute_task(task)
            
            execution_time = (datetime.now() - start_time).total_seconds()
            
            # Update system metrics
            self.system_metrics["completed_tasks"] += 1
            self.system_metrics["avg_completion_time"] = (
                (self.system_metrics["avg_completion_time"] * (self.system_metrics["completed_tasks"] - 1) + execution_time) 
                / self.system_metrics["completed_tasks"]
            )
            
            self.completed_tasks.append(task)
            self.logger.info(f"Task {task.id} completed successfully in {execution_time:.2f}s")
            
        except Exception as e:
            self.system_metrics["failed_tasks"] += 1
            self.logger.error(f"Task {task.id} failed: {e}")
        
        # Continue processing queue
        await self._process_task_queue()
    
    def get_system_status(self) -> Dict[str, Any]:
        """Get current system status and metrics"""
        active_agents = len([a for a in self.agents.values() if a.state != AgentState.IDLE])
        total_agents = len(self.agents)
        
        self.system_metrics["system_utilization"] = active_agents / total_agents if total_agents > 0 else 0
        
        return {
            "agents": {
                "total": total_agents,
                "active": active_agents,
                "idle": total_agents - active_agents
            },
            "tasks": {
                "pending": len(self.task_queue),
                "total_completed": len(self.completed_tasks)
            },
            "metrics": self.system_metrics
        }

# Demo Usage
async def demo_agentic_system():
    """Demonstrate the agentic AI system"""
    
    # Initialize system
    system = AgenticSystem()
    
    # Create agents
    coordinator = CoordinatorAgent("coord_001")
    researcher = ResearchAgent("research_001")
    analyst = AnalystAgent("analyst_001")
    executor = ExecutorAgent("executor_001")
    
    # Set up coordinator with managed agents
    coordinator.add_managed_agent(researcher)
    coordinator.add_managed_agent(analyst)
    coordinator.add_managed_agent(executor)
    
    # Add agents to system
    system.add_agent(coordinator)
    system.add_agent(researcher)
    system.add_agent(analyst)
    system.add_agent(executor)
    
    print("🤖 Agentic AI System Demo Started\n")
    
    # Submit a complex business analysis task
    complex_task = Task(
        id="business_analysis_001",
        type="business_analysis",
        description="Analyze market opportunities for expanding into Southeast Asian markets",
        requirements={
            "topic": "Southeast Asian market expansion",
            "analysis_scope": "comprehensive",
            "timeline": "Q2 2025"
        },
        dependencies=[]
    )
    
    task_id = await system.submit_task(complex_task)
    print(f"📋 Submitted complex task: {task_id}")
    
    # Wait for task completion
    await asyncio.sleep(5)
    
    # Check system status
    status = system.get_system_status()
    print(f"\n📊 System Status:")
    print(f"   Agents: {status['agents']['active']}/{status['agents']['total']} active")
    print(f"   Tasks: {status['tasks']['total_completed']} completed, {status['tasks']['pending']} pending")
    print(f"   Avg completion time: {status['metrics']['avg_completion_time']:.2f}s")
    print(f"   System utilization: {status['metrics']['system_utilization']*100:.1f}%")
    
    print("\n✅ Demo completed successfully!")

# Run the demo
if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO)
    asyncio.run(demo_agentic_system())