src/agents/router.py

#!/usr/bin/env python3
"""
Router Agent for GAIA Question Classification
Analyzes questions and routes them to appropriate specialized agents
"""

import re
import logging
from typing import List, Dict, Any, Tuple
from urllib.parse import urlparse

from agents.state import GAIAAgentState, QuestionType, AgentRole, AgentResult
from models.qwen_client import QwenClient, ModelTier

logger = logging.getLogger(__name__)

class RouterAgent:
    """
    Router agent that classifies GAIA questions and determines processing strategy
    """
    
    def __init__(self, llm_client: QwenClient):
        self.llm_client = llm_client
        
    def route_question(self, state: GAIAAgentState) -> GAIAAgentState:
        """
        Main routing function - analyzes question and updates state with routing decisions
        """
        logger.info(f"Routing question: {state.question[:100]}...")
        state.add_processing_step("Router: Starting question analysis")
        
        # Step 1: Enhanced question classification with multi-type detection
        question_types, primary_type = self._classify_question_types(state.question, state.file_name)
        state.question_type = primary_type
        state.add_processing_step(f"Router: Primary type: {primary_type.value}, All types: {[t.value for t in question_types]}")
        
        # Step 2: Complexity assessment
        complexity = self._assess_complexity(state.question)
        state.complexity_assessment = complexity
        state.add_processing_step(f"Router: Assessed complexity as {complexity}")
        
        # Step 3: Select appropriate agents with sequencing
        selected_agents = self._select_agents_enhanced(question_types, primary_type, state.file_name is not None, complexity)
        state.selected_agents = selected_agents
        state.add_processing_step(f"Router: Selected agents: {[a.value for a in selected_agents]}")
        
        # Step 4: Estimate cost
        estimated_cost = self._estimate_cost(complexity, selected_agents)
        state.estimated_cost = estimated_cost
        state.add_processing_step(f"Router: Estimated cost: ${estimated_cost:.4f}")
        
        # Step 5: Create routing decision summary
        state.routing_decision = {
            "primary_type": primary_type.value,
            "all_types": [t.value for t in question_types],
            "complexity": complexity,
            "agents": [agent.value for agent in selected_agents],
            "estimated_cost": estimated_cost,
            "reasoning": self._get_routing_reasoning(primary_type, complexity, selected_agents, question_types)
        }
        
        # Step 6: Use LLM for complex routing decisions if needed
        if complexity == "complex" or primary_type == QuestionType.UNKNOWN or len(question_types) > 2:
            state = self._llm_enhanced_routing(state)
        
        logger.info(f"✅ Routing complete: {primary_type.value} -> {[a.value for a in selected_agents]}")
        return state
    
    def _classify_question_types(self, question: str, file_name: str = None) -> Tuple[List[QuestionType], QuestionType]:
        """
        Enhanced classification that can detect multiple question types
        Returns: (all_detected_types, primary_type)
        """
        
        question_lower = question.lower()
        detected_types = []
        
        # File processing questions (highest priority when file is present)
        if file_name:
            file_ext = file_name.lower().split('.')[-1] if '.' in file_name else ""
            
            if file_ext in ['jpg', 'jpeg', 'png', 'gif', 'bmp', 'svg']:
                detected_types.append(QuestionType.FILE_PROCESSING)
            elif file_ext in ['mp3', 'wav', 'ogg', 'flac', 'm4a']:
                detected_types.append(QuestionType.FILE_PROCESSING)
            elif file_ext in ['xlsx', 'xls', 'csv']:
                detected_types.append(QuestionType.FILE_PROCESSING)
            elif file_ext in ['py', 'js', 'java', 'cpp', 'c']:
                detected_types.append(QuestionType.CODE_EXECUTION)
            else:
                detected_types.append(QuestionType.FILE_PROCESSING)
        
        # Enhanced URL-based classification
        url_patterns = {
            QuestionType.WIKIPEDIA: [
                r'wikipedia\.org', r'featured article', r'promoted.*wikipedia',
                r'english wikipedia', r'wiki.*article'
            ],
            QuestionType.YOUTUBE: [
                r'youtube\.com', r'youtu\.be', r'watch\?v=', r'video.*youtube',
                r'https://www\.youtube\.com/watch'
            ]
        }
        
        for question_type, patterns in url_patterns.items():
            if any(re.search(pattern, question_lower) for pattern in patterns):
                detected_types.append(question_type)
        
        # Enhanced content-based classification with better patterns
        classification_patterns = {
            QuestionType.MATHEMATICAL: [
                # Counting/quantity questions
                r'\bhow many\b', r'\bhow much\b', r'\bcount\b', r'\bnumber of\b',
                r'\btotal\b', r'\bsum\b', r'\baverage\b', r'\bmean\b',
                # Calculations
                r'\bcalculate\b', r'\bcompute\b', r'\bsolve\b', 
                # Mathematical operations
                r'\d+\s*[\+\-\*/]\s*\d+', r'\bsquare root\b', r'\bpercentage\b',
                # Table analysis
                r'\btable\b.*\bdefining\b', r'\bgiven.*table\b', r'\boperation table\b',
                # Specific math terms
                r'\bequation\b', r'\bformula\b', r'\bratio\b', r'\bfactorial\b',
                # Economic/statistical
                r'\binterest\b', r'\bcompound\b', r'\bstatistics\b'
            ],
            QuestionType.TEXT_MANIPULATION: [
                # Text operations
                r'\breverse\b', r'\bbackwards\b', r'\bencode\b', r'\bdecode\b',
                r'\btransform\b', r'\bconvert\b', r'\buppercase\b', r'\blowercase\b',
                r'\breplace\b', r'\bextract\b', r'\bopposite\b',
                # Pattern recognition for backwards text
                r'[a-z]+\s+[a-z]+\s+[a-z]+.*\.',  # Potential backwards sentence
                # Specific text manipulation clues
                r'\.rewsna\b', r'\bword.*opposite\b'
            ],
            QuestionType.CODE_EXECUTION: [
                r'\bcode\b', r'\bprogram\b', r'\bscript\b', r'\bfunction\b', r'\balgorithm\b',
                r'\bexecute\b', r'\brun.*code\b', r'\bpython\b', r'\bjavascript\b',
                r'\battached.*code\b', r'\bfinal.*output\b', r'\bnumeric output\b'
            ],
            QuestionType.REASONING: [
                # Logical reasoning
                r'\bwhy\b', r'\bexplain\b', r'\banalyze\b', r'\breasoning\b', r'\blogic\b',
                r'\brelationship\b', r'\bcompare\b', r'\bcontrast\b', r'\bconclusion\b',
                # Complex analysis
                r'\bexamine\b', r'\bidentify\b', r'\bdetermine\b', r'\bassess\b',
                r'\bevaluate\b', r'\binterpret\b'
            ],
            QuestionType.WEB_RESEARCH: [
                # General research
                r'\bsearch\b', r'\bfind.*information\b', r'\bresearch\b', r'\blook up\b',
                r'\bwebsite\b', r'\bonline\b', r'\binternet\b',
                # Who/what/when/where questions
                r'\bwho\s+(?:is|was|are|were|did|does)\b',
                r'\bwhat\s+(?:is|was|are|were)\b', r'\bwhen\s+(?:is|was|did|does)\b',
                r'\bwhere\s+(?:is|was|are|were)\b',
                # Factual queries
                r'\bauthor\b', r'\bpublished\b', r'\bhistory\b', r'\bhistorical\b',
                r'\bcentury\b', r'\byear\b', r'\bbiography\b', r'\bwinner\b',
                # Specific research indicators
                r'\bstudio albums\b', r'\brecipient\b', r'\bcompetition\b', r'\bspecimens\b'
            ]
        }
        
        # Score each category with enhanced scoring
        type_scores = {}
        for question_type, patterns in classification_patterns.items():
            score = 0
            for pattern in patterns:
                matches = re.findall(pattern, question_lower)
                score += len(matches)
                # Give extra weight to certain patterns
                if question_type == QuestionType.MATHEMATICAL and pattern in [r'\bhow many\b', r'\bhow much\b']:
                    score += 2  # Boost counting questions
                elif question_type == QuestionType.TEXT_MANIPULATION and any(special in pattern for special in ['opposite', 'reverse', 'backwards']):
                    score += 1  # Reduced further to avoid over-weighting
            if score > 0:
                type_scores[question_type] = score
        
        # Special handling for specific question patterns
        
        # Detect backwards/scrambled text (strong indicator) - only for clearly backwards text
        if re.search(r'\.rewsna\b|etirw\b|dnatsrednu\b', question_lower):
            type_scores[QuestionType.TEXT_MANIPULATION] = type_scores.get(QuestionType.TEXT_MANIPULATION, 0) + 3
            
        # Detect code execution patterns (strong indicator)
        if re.search(r'\bfinal.*output\b|\bnumeric.*output\b|\battached.*code\b', question_lower):
            type_scores[QuestionType.CODE_EXECUTION] = type_scores.get(QuestionType.CODE_EXECUTION, 0) + 4
            
        # Detect mathematical operations with numbers (boost mathematical score)
        if re.search(r'\b\d+.*\b(?:studio albums|between|and)\b.*\d+', question_lower):
            type_scores[QuestionType.MATHEMATICAL] = type_scores.get(QuestionType.MATHEMATICAL, 0) + 3
            
        # Detect table/grid operations
        if re.search(r'\btable.*defining.*\*', question_lower) or '|*|' in question:
            type_scores[QuestionType.MATHEMATICAL] = type_scores.get(QuestionType.MATHEMATICAL, 0) + 4
            
        # Multi-step questions that need research AND calculation
        if ('how many' in question_lower or 'how much' in question_lower) and \
           any(term in question_lower for term in ['between', 'from', 'during', 'published', 'released']):
            type_scores[QuestionType.WEB_RESEARCH] = type_scores.get(QuestionType.WEB_RESEARCH, 0) + 3  # Increased from 2
            type_scores[QuestionType.MATHEMATICAL] = type_scores.get(QuestionType.MATHEMATICAL, 0) + 3  # Increased from 2
            
        # Detect factual research questions (boost web research)
        if any(pattern in question_lower for pattern in ['who is', 'who was', 'who did', 'what is', 'when did', 'where', 'which']):
            type_scores[QuestionType.WEB_RESEARCH] = type_scores.get(QuestionType.WEB_RESEARCH, 0) + 2
            
        # Detect image/file references
        if any(term in question_lower for term in ['image', 'picture', 'photo', 'file', 'attached', 'provided']):
            type_scores[QuestionType.FILE_PROCESSING] = type_scores.get(QuestionType.FILE_PROCESSING, 0) + 4  # Increased from 3
            
        # Detect Wikipedia-specific questions
        if any(term in question_lower for term in ['wikipedia', 'featured article', 'english wikipedia']):
            type_scores[QuestionType.WIKIPEDIA] = type_scores.get(QuestionType.WIKIPEDIA, 0) + 4
        
        # Add detected types based on scores
        for qtype, score in type_scores.items():
            if score > 0 and qtype not in detected_types:
                detected_types.append(qtype)
        
        # If no types detected, default to web research
        if not detected_types:
            detected_types.append(QuestionType.WEB_RESEARCH)
        
        # Determine primary type (highest scoring)
        if type_scores:
            primary_type = max(type_scores.keys(), key=lambda t: type_scores[t])
        else:
            primary_type = detected_types[0] if detected_types else QuestionType.WEB_RESEARCH
        
        return detected_types, primary_type
    
    def _assess_complexity(self, question: str) -> str:
        """Assess question complexity with enhanced logic"""
        
        question_lower = question.lower()
        
        # Complex indicators
        complex_indicators = [
            'multi-step', 'multiple', 'several', 'complex', 'detailed',
            'analyze', 'explain why', 'reasoning', 'relationship',
            'compare and contrast', 'comprehensive', 'thorough',
            'between.*and', 'table.*defining', 'attached.*file'
        ]
        
        # Simple indicators  
        simple_indicators = [
            'what is', 'who is', 'when did', 'where is', 'yes or no',
            'true or false', 'simple', 'quick', 'name'
        ]
        
        complex_score = sum(1 for indicator in complex_indicators if re.search(indicator, question_lower))
        simple_score = sum(1 for indicator in simple_indicators if re.search(indicator, question_lower))
        
        # Additional complexity factors
        if len(question) > 200:
            complex_score += 1
        if len(question.split()) > 30:
            complex_score += 1
        if question.count('?') > 1:  # Multiple questions
            complex_score += 1
        if '|' in question and '*' in question:  # Tables
            complex_score += 2
        if re.search(r'\d+.*between.*\d+', question_lower):  # Date ranges
            complex_score += 1
        
        # Determine complexity
        if complex_score >= 3:
            return "complex"
        elif complex_score >= 1 and simple_score == 0:
            return "medium"
        elif simple_score >= 2 and complex_score == 0:
            return "simple"
        else:
            return "medium"
    
    def _select_agents_enhanced(self, question_types: List[QuestionType], primary_type: QuestionType, 
                              has_file: bool, complexity: str) -> List[AgentRole]:
        """
        Enhanced agent selection that can choose multiple agents for complex workflows
        """
        
        agents = []
        
        # Always include synthesizer at the end for final answer compilation
        # (We'll add it at the end to ensure proper ordering)
        
        # Multi-agent selection based on detected question types
        agent_priorities = {
            QuestionType.FILE_PROCESSING: [AgentRole.FILE_PROCESSOR],
            QuestionType.CODE_EXECUTION: [AgentRole.CODE_EXECUTOR],
            QuestionType.WIKIPEDIA: [AgentRole.WEB_RESEARCHER],
            QuestionType.YOUTUBE: [AgentRole.WEB_RESEARCHER],
            QuestionType.WEB_RESEARCH: [AgentRole.WEB_RESEARCHER],
            QuestionType.MATHEMATICAL: [AgentRole.REASONING_AGENT],
            QuestionType.TEXT_MANIPULATION: [AgentRole.REASONING_AGENT],
            QuestionType.REASONING: [AgentRole.REASONING_AGENT]
        }
        
        # Add agents based on all detected question types
        for qtype in question_types:
            if qtype in agent_priorities:
                for agent in agent_priorities[qtype]:
                    if agent not in agents:
                        agents.append(agent)
        
        # Special combinations for multi-step questions
        
        # For CODE_EXECUTION as primary type, prioritize code executor
        if primary_type == QuestionType.CODE_EXECUTION:
            # Ensure code executor is first, followed by any other needed agents
            ordered_agents = []
            if AgentRole.CODE_EXECUTOR not in ordered_agents:
                ordered_agents.append(AgentRole.CODE_EXECUTOR)
            # Add other agents if needed for multi-type questions
            for agent in agents:
                if agent != AgentRole.CODE_EXECUTOR and agent not in ordered_agents:
                    ordered_agents.append(agent)
            agents = ordered_agents
        
        # Research + Math combinations (e.g., "How many albums between 2000-2009?")
        elif (QuestionType.WEB_RESEARCH in question_types and QuestionType.MATHEMATICAL in question_types):
            # Ensure proper order: Research first, then math
            ordered_agents = []
            if AgentRole.WEB_RESEARCHER not in ordered_agents:
                ordered_agents.append(AgentRole.WEB_RESEARCHER)
            if AgentRole.REASONING_AGENT not in ordered_agents:
                ordered_agents.append(AgentRole.REASONING_AGENT)
            agents = ordered_agents
        
        # File + Analysis combinations
        elif has_file and len(question_types) > 1:
            # File processing should come first
            ordered_agents = []
            if AgentRole.FILE_PROCESSOR not in ordered_agents:
                ordered_agents.append(AgentRole.FILE_PROCESSOR)
            # Then add other agents
            for agent in agents:
                if agent != AgentRole.FILE_PROCESSOR and agent not in ordered_agents:
                    ordered_agents.append(agent)
            agents = ordered_agents
        
        # For complex questions, add reasoning if not already present
        if complexity == "complex" and AgentRole.REASONING_AGENT not in agents:
            agents.append(AgentRole.REASONING_AGENT)
        
        # Fallback for unknown/unclear questions - use multiple agents
        if primary_type == QuestionType.UNKNOWN or not agents:
            agents = [AgentRole.WEB_RESEARCHER, AgentRole.REASONING_AGENT]
        
        # Always add synthesizer at the end
        agents.append(AgentRole.SYNTHESIZER)
        
        # Remove duplicates while preserving order
        seen = set()
        unique_agents = []
        for agent in agents:
            if agent not in seen:
                seen.add(agent)
                unique_agents.append(agent)
                
        return unique_agents
    
    def _estimate_cost(self, complexity: str, agents: List[AgentRole]) -> float:
        """Estimate processing cost based on complexity and agents"""
        
        base_costs = {
            "simple": 0.005,   # Router model mostly
            "medium": 0.015,   # Mix of router and main
            "complex": 0.035   # Include complex model usage
        }
        
        base_cost = base_costs.get(complexity, 0.015)
        
        # Additional cost per agent (more agents = more processing)
        agent_cost = len(agents) * 0.008
        
        return base_cost + agent_cost
    
    def _get_routing_reasoning(self, primary_type: QuestionType, complexity: str, 
                             agents: List[AgentRole], all_types: List[QuestionType]) -> str:
        """Generate human-readable reasoning for routing decision"""
        
        reasons = []
        
        # Primary type reasoning
        type_descriptions = {
            QuestionType.WIKIPEDIA: "References Wikipedia content",
            QuestionType.YOUTUBE: "Involves YouTube video analysis", 
            QuestionType.FILE_PROCESSING: "Requires file processing",
            QuestionType.MATHEMATICAL: "Involves mathematical computation/counting",
            QuestionType.CODE_EXECUTION: "Requires code execution",
            QuestionType.TEXT_MANIPULATION: "Involves text transformation/manipulation",
            QuestionType.REASONING: "Requires logical reasoning/analysis",
            QuestionType.WEB_RESEARCH: "Needs web research for factual information"
        }
        
        if primary_type in type_descriptions:
            reasons.append(type_descriptions[primary_type])
        
        # Multi-type questions
        if len(all_types) > 1:
            other_types = [t for t in all_types if t != primary_type]
            reasons.append(f"Also involves: {', '.join([t.value for t in other_types])}")
        
        # Complexity reasoning
        if complexity == "complex":
            reasons.append("Complex multi-step reasoning required")
        elif complexity == "simple":
            reasons.append("Straightforward question")
        
        # Agent workflow reasoning
        agent_names = [agent.value.replace('_', ' ') for agent in agents]
        if len(agents) > 2:  # More than synthesizer + one agent
            reasons.append(f"Multi-agent workflow: {' → '.join(agent_names)}")
        else:
            reasons.append(f"Single-agent workflow: {', '.join(agent_names)}")
        
        return "; ".join(reasons)
    
    def _llm_enhanced_routing(self, state: GAIAAgentState) -> GAIAAgentState:
        """Use LLM for enhanced routing analysis of complex/unknown questions"""
        
        prompt = f"""
        Analyze this GAIA benchmark question and provide routing guidance:
        
        Question: {state.question}
        File attached: {state.file_name if state.file_name else "None"}
        Detected types: {state.routing_decision.get('all_types', [])}
        Primary classification: {state.question_type.value}
        Current complexity: {state.complexity_assessment}
        Selected agents: {[a.value for a in state.selected_agents]}
        
        Does this question need:
        1. Web research to find factual information?
        2. Mathematical calculation or counting?
        3. Text manipulation or decoding?
        4. File processing or analysis?
        5. Logical reasoning or analysis?
        
        Should the agent selection be adjusted? If so, provide specific recommendations.
        Keep response concise and focused on routing decisions.
        """
        
        try:
            # Use main model (32B) for better routing decisions
            tier = ModelTier.MAIN
            result = self.llm_client.generate(prompt, tier=tier, max_tokens=300)
            
            if result.success:
                state.add_processing_step("Router: Enhanced with LLM analysis")
                state.routing_decision["llm_analysis"] = result.response
                logger.info("✅ LLM enhanced routing completed")
            else:
                state.add_error(f"LLM routing enhancement failed: {result.error}")
                
        except Exception as e:
            state.add_error(f"LLM routing error: {str(e)}")
            logger.error(f"LLM routing failed: {e}")
        
        return state