03_System_Tools/task_orchestrator.py

#!/usr/bin/env python3
"""
Luuna Task Distribution and Agent Orchestration System
Intelligently distributes tasks to specialized agents based on document analysis
"""

import json
import time
import uuid
from typing import Dict, List, Optional, Callable
from pathlib import Path
import logging
from dataclasses import dataclass, asdict
from enum import Enum

# Configure logging
logger = logging.getLogger(__name__)

class TaskPriority(Enum):
    LOW = "low"
    MEDIUM = "medium"
    HIGH = "high"
    CRITICAL = "critical"

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

@dataclass
class Task:
    """Represents a single task to be executed by an agent"""
    task_id: str
    task_type: str
    description: str
    priority: TaskPriority
    estimated_duration: str
    required_skills: List[str]
    inputs: List[str]
    expected_outputs: List[str]
    status: TaskStatus = TaskStatus.PENDING
    assigned_agent: Optional[str] = None
    created_at: float = None
    started_at: Optional[float] = None
    completed_at: Optional[float] = None
    workspace_path: Optional[str] = None
    
    def __post_init__(self):
        if self.created_at is None:
            self.created_at = time.time()
        if isinstance(self.priority, str):
            self.priority = TaskPriority(self.priority)
        if isinstance(self.status, str):
            self.status = TaskStatus(self.status)

@dataclass
class AgentProfile:
    """Defines an agent's capabilities and configuration"""
    agent_id: str
    name: str
    category: str
    description: str
    skills: List[str]
    capabilities: List[str]
    max_concurrent_tasks: int = 3
    current_load: int = 0
    status: str = "available"  # available, busy, offline
    last_heartbeat: float = None
    
    def __post_init__(self):
        if self.last_heartbeat is None:
            self.last_heartbeat = time.time()

class AgentRegistry:
    """Manages available agents and their capabilities"""
    
    def __init__(self):
        self.agents: Dict[str, AgentProfile] = {}
        self._initialize_default_agents()
        
    def _initialize_default_agents(self):
        """Initialize default agent profiles"""
        default_agents = [
            AgentProfile(
                agent_id="crypto_market_analyzer_001",
                name="CryptoMarketAnalyzer",
                category="Crypto Agents",
                description="Specialized agent for cryptocurrency market analysis and trend detection",
                skills=["market_analysis", "data_processing", "trend_detection", "price_prediction"],
                capabilities=["analyze_market_data", "detect_trends", "generate_reports", "risk_assessment"]
            ),
            AgentProfile(
                agent_id="risk_assessment_agent_001",
                name="RiskAssessmentAgent",
                category="Business Agents",
                description="Agent for financial risk assessment and compliance evaluation",
                skills=["risk_analysis", "financial_modeling", "compliance", "regulatory_analysis"],
                capabilities=["assess_investment_risks", "evaluate_compliance", "generate_risk_reports"]
            ),
            AgentProfile(
                agent_id="technical_implementation_agent_001",
                name="TechnicalImplementationAgent",
                category="Developer Tools Agents",
                description="Agent for software development and technical implementation tasks",
                skills=["software_development", "system_design", "coding", "api_integration"],
                capabilities=["implement_features", "write_code", "design_architecture", "debug_issues"]
            ),
            AgentProfile(
                agent_id="business_strategy_agent_001",
                name="BusinessStrategyAgent",
                category="Business Agents",
                description="Agent for business strategy development and market analysis",
                skills=["business_analysis", "strategic_planning", "market_research", "competitive_analysis"],
                capabilities=["develop_strategies", "analyze_markets", "create_business_plans"]
            ),
            AgentProfile(
                agent_id="research_analysis_agent_001",
                name="ResearchAnalysisAgent",
                category="Research Agents",
                description="Agent for academic and technical research analysis",
                skills=["research_analysis", "literature_review", "data_interpretation", "statistical_analysis"],
                capabilities=["analyze_research", "summarize_findings", "identify_gaps", "generate_insights"]
            ),
            AgentProfile(
                agent_id="document_processing_agent_001",
                name="DocumentProcessingAgent",
                category="Document Agents",
                description="Agent for document processing and content extraction",
                skills=["document_analysis", "content_extraction", "text_processing", "information_retrieval"],
                capabilities=["process_documents", "extract_information", "categorize_content", "generate_summaries"]
            )
        ]
        
        for agent in default_agents:
            self.register_agent(agent)
            
    def register_agent(self, agent: AgentProfile):
        """Register a new agent"""
        self.agents[agent.agent_id] = agent
        logger.info(f"Registered agent: {agent.name} ({agent.agent_id})")
        
    def get_available_agents(self) -> List[AgentProfile]:
        """Get list of currently available agents"""
        return [agent for agent in self.agents.values() 
                if agent.status == "available" and agent.current_load < agent.max_concurrent_tasks]
                
    def get_agents_by_skill(self, required_skill: str) -> List[AgentProfile]:
        """Find agents that possess a specific skill"""
        return [agent for agent in self.agents.values() 
                if required_skill in agent.skills]
                
    def get_agents_by_category(self, category: str) -> List[AgentProfile]:
        """Find agents in a specific category"""
        return [agent for agent in self.agents.values() 
                if agent.category == category]
                
    def assign_task_to_agent(self, agent_id: str, task_id: str) -> bool:
        """Assign a task to an agent"""
        if agent_id in self.agents:
            agent = self.agents[agent_id]
            if agent.current_load < agent.max_concurrent_tasks:
                agent.current_load += 1
                if agent.current_load >= agent.max_concurrent_tasks:
                    agent.status = "busy"
                logger.info(f"Assigned task {task_id} to agent {agent.name}")
                return True
        return False
        
    def complete_task_for_agent(self, agent_id: str, task_id: str) -> bool:
        """Mark a task as completed for an agent"""
        if agent_id in self.agents:
            agent = self.agents[agent_id]
            if agent.current_load > 0:
                agent.current_load -= 1
                agent.status = "available" if agent.current_load < agent.max_concurrent_tasks else "busy"
                agent.last_heartbeat = time.time()
                logger.info(f"Completed task {task_id} for agent {agent.name}")
                return True
        return False

class TaskDistributor:
    """Intelligent task distribution system"""
    
    def __init__(self):
        self.agent_registry = AgentRegistry()
        self.task_assignment_history: List[Dict] = []
        
    def distribute_tasks(self, tasks: List[Task], workspace_path: str) -> List[Dict]:
        """Distribute tasks to appropriate agents"""
        logger.info(f"Distributing {len(tasks)} tasks for workspace: {workspace_path}")
        
        assignments = []
        
        for task in tasks:
            # Find the best agent for this task
            assigned_agent = self._find_best_agent_for_task(task)
            
            if assigned_agent:
                # Update task with assignment
                task.assigned_agent = assigned_agent.agent_id
                task.status = TaskStatus.ASSIGNED
                task.workspace_path = workspace_path
                
                # Assign task to agent
                self.agent_registry.assign_task_to_agent(assigned_agent.agent_id, task.task_id)
                
                # Record assignment
                assignment = {
                    'task_id': task.task_id,
                    'task_description': task.description,
                    'assigned_agent': assigned_agent.name,
                    'agent_id': assigned_agent.agent_id,
                    'agent_category': assigned_agent.category,
                    'assignment_time': time.time(),
                    'workspace_path': workspace_path
                }
                
                assignments.append(assignment)
                self.task_assignment_history.append(assignment)
                
                logger.info(f"Task '{task.task_id}' assigned to {assigned_agent.name}")
            else:
                logger.warning(f"No suitable agent found for task: {task.description}")
                
        return assignments
        
    def _find_best_agent_for_task(self, task: Task) -> Optional[AgentProfile]:
        """Find the most suitable agent for a given task"""
        # Strategy 1: Exact skill match
        best_agents = []
        
        # Find agents with all required skills
        for skill in task.required_skills:
            skilled_agents = self.agent_registry.get_agents_by_skill(skill)
            if not best_agents:
                best_agents = skilled_agents
            else:
                # Intersect with previous results
                best_agents = [agent for agent in best_agents if agent in skilled_agents]
                
        # If we found agents with all required skills, pick the least loaded one
        if best_agents:
            available_agents = [agent for agent in best_agents 
                              if agent.status == "available"]
            if available_agents:
                # Sort by current load (ascending) and last heartbeat (descending)
                available_agents.sort(key=lambda x: (x.current_load, -x.last_heartbeat))
                return available_agents[0]
                
        # Strategy 2: Category-based assignment
        category_agents = self.agent_registry.get_agents_by_category("General Purpose")
        if category_agents:
            available_general_agents = [agent for agent in category_agents 
                                      if agent.status == "available"]
            if available_general_agents:
                available_general_agents.sort(key=lambda x: (x.current_load, -x.last_heartbeat))
                return available_general_agents[0]
                
        # Strategy 3: Any available agent
        available_agents = self.agent_registry.get_available_agents()
        if available_agents:
            available_agents.sort(key=lambda x: (x.current_load, -x.last_heartbeat))
            return available_agents[0]
            
        return None
        
    def generate_tasks_from_document_analysis(self, document_analysis: Dict) -> List[Task]:
        """Generate tasks based on document analysis results"""
        tasks = []
        
        # Extract information from document analysis
        content_category = document_analysis.get('content_category', 'general')
        key_topics = document_analysis.get('key_topics', [])
        action_items = document_analysis.get('action_items', [])
        requirements = document_analysis.get('requirements', [])
        
        # Generate tasks based on content category
        if content_category == 'crypto_finance':
            tasks.extend(self._generate_crypto_tasks(key_topics, requirements))
        elif content_category == 'technical_docs':
            tasks.extend(self._generate_technical_tasks(key_topics, requirements))
        elif content_category == 'business_strategy':
            tasks.extend(self._generate_business_tasks(key_topics, requirements))
        elif content_category == 'research_analysis':
            tasks.extend(self._generate_research_tasks(key_topics, requirements))
            
        # Add action item tasks
        for action_item in action_items:
            task = Task(
                task_id=f"action_{uuid.uuid4().hex[:8]}",
                task_type="action_item",
                description=action_item.get('description', 'Process action item'),
                priority=TaskPriority(action_item.get('priority', 'medium')),
                estimated_duration="1-2 hours",
                required_skills=["general_analysis"],
                inputs=["action_description"],
                expected_outputs=["completed_action", "status_report"]
            )
            tasks.append(task)
            
        return tasks
        
    def _generate_crypto_tasks(self, topics: List[str], requirements: List[str]) -> List[Task]:
        """Generate cryptocurrency-related tasks"""
        tasks = []
        
        if any(topic in ['market', 'trading', 'price'] for topic in topics):
            tasks.append(Task(
                task_id=f"crypto_analysis_{uuid.uuid4().hex[:8]}",
                task_type="market_analysis",
                description="Analyze cryptocurrency market trends and price movements",
                priority=TaskPriority.HIGH,
                estimated_duration="2-3 hours",
                required_skills=["market_analysis", "data_processing"],
                inputs=["market_data", "historical_prices"],
                expected_outputs=["market_analysis_report", "trend_predictions"]
            ))
            
        if any(topic in ['wallet', 'security', 'private_key'] for topic in topics):
            tasks.append(Task(
                task_id=f"crypto_security_{uuid.uuid4().hex[:8]}",
                task_type="security_assessment",
                description="Assess cryptocurrency wallet security and best practices",
                priority=TaskPriority.MEDIUM,
                estimated_duration="1-2 hours",
                required_skills=["security_analysis", "blockchain_knowledge"],
                inputs=["wallet_information", "security_protocols"],
                expected_outputs=["security_assessment", "recommendations"]
            ))
            
        return tasks
        
    def _generate_technical_tasks(self, topics: List[str], requirements: List[str]) -> List[Task]:
        """Generate technical implementation tasks"""
        tasks = []
        
        if any(topic in ['api', 'integration'] for topic in topics):
            tasks.append(Task(
                task_id=f"api_integration_{uuid.uuid4().hex[:8]}",
                task_type="api_development",
                description="Implement API integration based on technical specifications",
                priority=TaskPriority.HIGH,
                estimated_duration="4-6 hours",
                required_skills=["api_development", "software_engineering"],
                inputs=["api_specifications", "technical_requirements"],
                expected_outputs=["working_api", "integration_tests"]
            ))
            
        if any(topic in ['architecture', 'design'] for topic in topics):
            tasks.append(Task(
                task_id=f"system_design_{uuid.uuid4().hex[:8]}",
                task_type="system_design",
                description="Design system architecture based on requirements",
                priority=TaskPriority.MEDIUM,
                estimated_duration="3-4 hours",
                required_skills=["system_design", "software_architecture"],
                inputs=["requirements_document", "technical_constraints"],
                expected_outputs=["system_design_document", "architecture_diagram"]
            ))
            
        return tasks
        
    def _generate_business_tasks(self, topics: List[str], requirements: List[str]) -> List[Task]:
        """Generate business strategy tasks"""
        tasks = []
        
        if any(topic in ['market', 'competition'] for topic in topics):
            tasks.append(Task(
                task_id=f"market_analysis_{uuid.uuid4().hex[:8]}",
                task_type="market_research",
                description="Conduct market analysis and competitive research",
                priority=TaskPriority.HIGH,
                estimated_duration="3-4 hours",
                required_skills=["market_research", "competitive_analysis"],
                inputs=["market_data", "competitor_information"],
                expected_outputs=["market_analysis_report", "competitive_intelligence"]
            ))
            
        if any(topic in ['strategy', 'plan'] for topic in topics):
            tasks.append(Task(
                task_id=f"strategy_development_{uuid.uuid4().hex[:8]}",
                task_type="strategy_development",
                description="Develop business strategy and implementation plan",
                priority=TaskPriority.HIGH,
                estimated_duration="4-5 hours",
                required_skills=["strategic_planning", "business_analysis"],
                inputs=["business_requirements", "market_analysis"],
                expected_outputs=["business_strategy", "implementation_plan"]
            ))
            
        return tasks
        
    def _generate_research_tasks(self, topics: List[str], requirements: List[str]) -> List[Task]:
        """Generate research and analysis tasks"""
        tasks = []
        
        tasks.append(Task(
            task_id=f"literature_review_{uuid.uuid4().hex[:8]}",
            task_type="research_analysis",
            description="Conduct literature review and research synthesis",
            priority=TaskPriority.MEDIUM,
            estimated_duration="3-4 hours",
            required_skills=["research_analysis", "literature_review"],
            inputs=["research_papers", "academic_sources"],
            expected_outputs=["literature_review", "research_summary"]
        ))
        
        return tasks

class TaskOrchestrator:
    """Orchestrates task execution and monitors progress"""
    
    def __init__(self):
        self.distributor = TaskDistributor()
        self.active_tasks: Dict[str, Task] = {}
        self.completed_tasks: List[Task] = []
        self.failed_tasks: List[Task] = []
        
    def process_workspace_documents(self, workspace_path: str, document_results: Dict) -> Dict:
        """Process workspace documents and orchestrate task execution"""
        logger.info(f"Processing workspace documents: {workspace_path}")
        
        # Generate tasks from document analysis
        all_tasks = []
        for file_result in document_results.get('processed_files', []):
            doc_analysis = file_result['result'].get('analysis_results', {})
            tasks = self.distributor.generate_tasks_from_document_analysis(doc_analysis)
            all_tasks.extend(tasks)
            
        # Distribute tasks to agents
        assignments = self.distributor.distribute_tasks(all_tasks, workspace_path)
        
        # Track active tasks
        for task in all_tasks:
            self.active_tasks[task.task_id] = task
            
        # Generate orchestration summary
        summary = {
            'workspace_path': workspace_path,
            'total_tasks_generated': len(all_tasks),
            'tasks_assigned': len(assignments),
            'agent_assignments': assignments,
            'processing_timestamp': time.time(),
            'status': 'tasks_distributed'
        }
        
        return summary
        
    def get_task_status(self, task_id: str) -> Optional[Dict]:
        """Get current status of a specific task"""
        if task_id in self.active_tasks:
            task = self.active_tasks[task_id]
            return {
                'task_id': task.task_id,
                'status': task.status.value,
                'assigned_agent': task.assigned_agent,
                'created_at': task.created_at,
                'started_at': task.started_at,
                'progress': self._calculate_task_progress(task)
            }
        return None
        
    def _calculate_task_progress(self, task: Task) -> float:
        """Calculate task progress percentage"""
        if task.status == TaskStatus.COMPLETED:
            return 100.0
        elif task.status == TaskStatus.IN_PROGRESS:
            # Simple time-based progress estimation
            if task.started_at:
                elapsed = time.time() - task.started_at
                estimated_total = self._parse_duration_hours(task.estimated_duration) * 3600
                return min((elapsed / estimated_total) * 100, 90.0)  # Cap at 90% until completion
        return 0.0
        
    def _parse_duration_hours(self, duration_str: str) -> float:
        """Parse duration string to hours"""
        # Simple parser for duration strings like "2-3 hours" or "4 hours"
        import re
        match = re.search(r'(\d+)(?:-(\d+))?\s*hours?', duration_str)
        if match:
            start = int(match.group(1))
            end = int(match.group(2)) if match.group(2) else start
            return (start + end) / 2.0
        return 1.0  # Default 1 hour
        
    def get_workspace_summary(self, workspace_path: str) -> Dict:
        """Get summary of all tasks for a workspace"""
        workspace_tasks = [task for task in self.active_tasks.values() 
                          if task.workspace_path == workspace_path]
                          
        completed_count = len([t for t in workspace_tasks if t.status == TaskStatus.COMPLETED])
        pending_count = len([t for t in workspace_tasks if t.status == TaskStatus.PENDING])
        in_progress_count = len([t for t in workspace_tasks if t.status == TaskStatus.IN_PROGRESS])
        
        return {
            'workspace_path': workspace_path,
            'total_tasks': len(workspace_tasks),
            'completed_tasks': completed_count,
            'pending_tasks': pending_count,
            'in_progress_tasks': in_progress_count,
            'completion_percentage': (completed_count / len(workspace_tasks) * 100) if workspace_tasks else 0
        }

# Example usage
if __name__ == "__main__":
    # Test the task distribution system
    orchestrator = TaskOrchestrator()
    
    # Simulate document analysis results
    sample_analysis = {
        'content_category': 'crypto_finance',
        'key_topics': ['market', 'trading', 'wallet'],
        'action_items': [
            {'description': 'Analyze Bitcoin price trends', 'priority': 'high'},
            {'description': 'Review wallet security protocols', 'priority': 'medium'}
        ],
        'requirements': ['Market data analysis', 'Security assessment']
    }
    
    # Generate tasks
    tasks = orchestrator.distributor.generate_tasks_from_document_analysis(sample_analysis)
    print(f"Generated {len(tasks)} tasks:")
    for task in tasks:
        print(f"- {task.description} (Priority: {task.priority.value})")