backend/core/rag_router.py

"""
Advanced RAG Router - Intelligent Strategy Selection
=====================================================

The RAG Router analyzes queries and selects the BEST RAG technique:
1. Simple queries → Basic RAG
2. Complex reasoning → Self-RAG
3. Long documents → RAPTOR
4. Need freshness → Corrective RAG with web fallback
5. Context-heavy → Contextual Retrieval
6. Multi-hop → GraphRAG

This maximizes quality while minimizing latency and API costs.

Uses FREE APIs only (Groq/Gemini).
"""

import os
import json
import logging
from typing import Dict, List, Any, Optional, Tuple
from dataclasses import dataclass
from enum import Enum

from core.llm import chat

logger = logging.getLogger(__name__)


class RAGStrategy(Enum):
    """Available RAG strategies"""
    BASIC = "basic"                     # Simple retrieval + generation
    SELF_RAG = "self_rag"               # Self-reflective with validation
    CORRECTIVE = "corrective"           # With web fallback
    CONTEXTUAL = "contextual"           # Anthropic's contextual retrieval
    RAPTOR = "raptor"                   # Hierarchical for long docs
    GRAPH = "graph"                     # GraphRAG for relationships
    HYBRID = "hybrid"                   # Combines multiple strategies
    AGENTIC = "agentic"                 # Full agent-based approach


@dataclass
class RAGRoutingDecision:
    """Result of routing decision"""
    strategy: RAGStrategy
    confidence: float
    reasoning: str
    estimated_latency: str
    fallback_strategy: RAGStrategy


class QueryComplexity(Enum):
    """Complexity levels for queries"""
    SIMPLE = "simple"           # Single fact lookup
    MODERATE = "moderate"       # Some reasoning required
    COMPLEX = "complex"         # Multi-step reasoning
    EXPERT = "expert"           # Deep analysis needed


class AdvancedRAGRouter:
    """
    Intelligently routes queries to the best RAG strategy.
    
    This optimizes for:
    - Quality (best answer)
    - Speed (lowest latency)
    - Cost (fewer API calls)
    
    Uses FREE APIs (Groq/Gemini).
    """
    
    def __init__(self):
        self.routing_history: List[Dict[str, Any]] = []
        
        # Strategy characteristics
        self.strategy_info = {
            RAGStrategy.BASIC: {
                "latency": "fast",
                "quality": "good",
                "best_for": ["simple questions", "fact lookup", "definitions"]
            },
            RAGStrategy.SELF_RAG: {
                "latency": "medium",
                "quality": "excellent",
                "best_for": ["accuracy-critical", "needs verification", "complex answers"]
            },
            RAGStrategy.CORRECTIVE: {
                "latency": "slow",
                "quality": "excellent",
                "best_for": ["incomplete context", "needs external info", "current events"]
            },
            RAGStrategy.CONTEXTUAL: {
                "latency": "fast",
                "quality": "very good",
                "best_for": ["context-dependent", "multiple documents", "ambiguous chunks"]
            },
            RAGStrategy.RAPTOR: {
                "latency": "medium",
                "quality": "excellent",
                "best_for": ["long documents", "summarization", "overview questions"]
            },
            RAGStrategy.GRAPH: {
                "latency": "slow",
                "quality": "excellent",
                "best_for": ["relationships", "comparisons", "entity connections"]
            },
            RAGStrategy.HYBRID: {
                "latency": "slow",
                "quality": "best",
                "best_for": ["complex multi-part questions", "research tasks"]
            },
            RAGStrategy.AGENTIC: {
                "latency": "slowest",
                "quality": "best",
                "best_for": ["actions needed", "multi-step plans", "tool use"]
            }
        }
    
    def _analyze_query_type(self, query: str) -> Dict[str, Any]:
        """
        Analyze query characteristics using heuristics.
        Fast local analysis before LLM if needed.
        """
        
        query_lower = query.lower()
        
        analysis = {
            "is_simple_lookup": False,
            "needs_comparison": False,
            "needs_calculation": False,
            "needs_visualization": False,
            "needs_current_info": False,
            "is_multi_part": False,
            "needs_action": False,
            "document_scope": "unknown"
        }
        
        # Simple lookup patterns
        simple_patterns = ["what is", "who is", "when was", "where is", "define", "list"]
        if any(query_lower.startswith(p) for p in simple_patterns):
            analysis["is_simple_lookup"] = True
        
        # Comparison patterns
        if any(p in query_lower for p in ["compare", "vs", "versus", "difference between", "better"]):
            analysis["needs_comparison"] = True
        
        # Calculation patterns
        if any(p in query_lower for p in ["calculate", "total", "sum", "average", "percentage", "growth"]):
            analysis["needs_calculation"] = True
        
        # Visualization patterns
        if any(p in query_lower for p in ["chart", "graph", "visualize", "plot", "show me", "trend"]):
            analysis["needs_visualization"] = True
        
        # Current info patterns
        if any(p in query_lower for p in ["latest", "current", "today", "now", "this week", "recent"]):
            analysis["needs_current_info"] = True
        
        # Multi-part patterns
        if any(p in query_lower for p in [" and ", " also ", " then ", "first,", "second,"]):
            analysis["is_multi_part"] = True
        
        # Action patterns
        if any(p in query_lower for p in ["send", "email", "notify", "create", "update", "schedule"]):
            analysis["needs_action"] = True
        
        # Document scope
        if any(p in query_lower for p in ["all documents", "across", "summary of all", "overview"]):
            analysis["document_scope"] = "multi"
        elif any(p in query_lower for p in ["this document", "this file", "the report"]):
            analysis["document_scope"] = "single"
        
        return analysis
    
    def _estimate_complexity(self, query: str, analysis: Dict[str, Any]) -> QueryComplexity:
        """
        Estimate query complexity.
        """
        
        score = 0
        
        if analysis["is_simple_lookup"]:
            score -= 2
        if analysis["needs_comparison"]:
            score += 1
        if analysis["needs_calculation"]:
            score += 1
        if analysis["is_multi_part"]:
            score += 2
        if analysis["needs_action"]:
            score += 3
        if len(query.split()) > 20:
            score += 1
        
        if score <= 0:
            return QueryComplexity.SIMPLE
        elif score <= 2:
            return QueryComplexity.MODERATE
        elif score <= 4:
            return QueryComplexity.COMPLEX
        else:
            return QueryComplexity.EXPERT
    
    def route(
        self,
        query: str,
        context_available: bool = True,
        document_count: int = 1,
        prefer_speed: bool = False,
        prefer_quality: bool = False
    ) -> RAGRoutingDecision:
        """
        Route query to the best RAG strategy.
        
        Uses heuristics for fast routing (no LLM call!).
        
        Args:
            query: User's question
            context_available: Whether we have document context
            document_count: Number of documents in context
            prefer_speed: Prioritize speed over quality
            prefer_quality: Prioritize quality over speed
        """
        
        # Analyze query
        analysis = self._analyze_query_type(query)
        complexity = self._estimate_complexity(query, analysis)
        
        logger.info(f"RAG Router: Complexity={complexity.value}, Analysis={analysis}")
        
        # Decision logic
        if analysis["needs_action"]:
            strategy = RAGStrategy.AGENTIC
            reasoning = "Query requires taking action, using full agentic approach"
            
        elif not context_available or analysis["needs_current_info"]:
            strategy = RAGStrategy.CORRECTIVE
            reasoning = "Need external information or context incomplete, using corrective RAG with web fallback"
            
        elif analysis["needs_comparison"] or complexity == QueryComplexity.EXPERT:
            strategy = RAGStrategy.SELF_RAG
            reasoning = "Complex comparison or expert-level question, using self-reflective RAG for accuracy"
            
        elif document_count > 3 or analysis["document_scope"] == "multi":
            strategy = RAGStrategy.RAPTOR
            reasoning = "Multiple documents or need overview, using hierarchical RAPTOR"
            
        elif analysis["is_multi_part"]:
            strategy = RAGStrategy.HYBRID
            reasoning = "Multi-part question, using hybrid approach"
            
        elif prefer_speed or (analysis["is_simple_lookup"] and complexity == QueryComplexity.SIMPLE):
            strategy = RAGStrategy.BASIC
            reasoning = "Simple lookup, using basic fast RAG"
            
        elif prefer_quality:
            strategy = RAGStrategy.SELF_RAG
            reasoning = "Quality-first mode, using self-reflective RAG"
            
        else:
            # Default: contextual retrieval as good balance
            strategy = RAGStrategy.CONTEXTUAL
            reasoning = "Standard query, using contextual retrieval for balanced performance"
        
        # Determine fallback
        fallback_map = {
            RAGStrategy.BASIC: RAGStrategy.CONTEXTUAL,
            RAGStrategy.SELF_RAG: RAGStrategy.CORRECTIVE,
            RAGStrategy.CORRECTIVE: RAGStrategy.AGENTIC,
            RAGStrategy.CONTEXTUAL: RAGStrategy.SELF_RAG,
            RAGStrategy.RAPTOR: RAGStrategy.SELF_RAG,
            RAGStrategy.GRAPH: RAGStrategy.SELF_RAG,
            RAGStrategy.HYBRID: RAGStrategy.AGENTIC,
            RAGStrategy.AGENTIC: RAGStrategy.CORRECTIVE
        }
        
        # Confidence based on match quality
        confidence = 0.85
        if prefer_speed and strategy in [RAGStrategy.BASIC, RAGStrategy.CONTEXTUAL]:
            confidence = 0.95
        if prefer_quality and strategy in [RAGStrategy.SELF_RAG, RAGStrategy.AGENTIC]:
            confidence = 0.95
        
        # Latency estimate
        latency_map = {
            RAGStrategy.BASIC: "~1-2s",
            RAGStrategy.CONTEXTUAL: "~2-3s",
            RAGStrategy.SELF_RAG: "~3-5s",
            RAGStrategy.CORRECTIVE: "~5-8s",
            RAGStrategy.RAPTOR: "~3-5s",
            RAGStrategy.GRAPH: "~5-10s",
            RAGStrategy.HYBRID: "~8-12s",
            RAGStrategy.AGENTIC: "~10-20s"
        }
        
        decision = RAGRoutingDecision(
            strategy=strategy,
            confidence=confidence,
            reasoning=reasoning,
            estimated_latency=latency_map.get(strategy, "unknown"),
            fallback_strategy=fallback_map.get(strategy, RAGStrategy.BASIC)
        )
        
        # Record decision
        self.routing_history.append({
            "query": query[:100],
            "strategy": strategy.value,
            "confidence": confidence
        })
        
        return decision
    
    async def route_with_llm(
        self,
        query: str,
        context_preview: str = ""
    ) -> RAGRoutingDecision:
        """
        Use LLM for more accurate routing (slower but smarter).
        
        Only use for complex queries where heuristics might fail.
        """
        
        prompt = f"""You are a RAG strategy selector. Choose the best retrieval strategy for this query.

QUERY: {query}

CONTEXT PREVIEW (if any):
{context_preview[:500]}

AVAILABLE STRATEGIES:
1. basic - Fast, simple retrieval for straightforward questions
2. self_rag - Self-reflective, validates answers, great for accuracy
3. corrective - Falls back to web search if context insufficient
4. contextual - Adds document context to chunks, good for multi-doc
5. raptor - Hierarchical summarization, great for long documents
6. graph - Knowledge graph-based, great for relationships
7. hybrid - Combines multiple strategies for complex queries
8. agentic - Full agent approach with tools, for actions

Select ONE strategy and explain briefly.

Respond with JSON:
{{"strategy": "strategy_name", "confidence": 0.0-1.0, "reasoning": "why"}}"""

        try:
            result = chat(prompt, temperature=0.1)
            
            # Parse response
            import re
            json_match = re.search(r'\{[^{}]*\}', result, re.DOTALL)
            if json_match:
                data = json.loads(json_match.group())
                strategy = RAGStrategy(data.get("strategy", "basic"))
                confidence = float(data.get("confidence", 0.8))
                reasoning = data.get("reasoning", "")
                
                return RAGRoutingDecision(
                    strategy=strategy,
                    confidence=confidence,
                    reasoning=reasoning,
                    estimated_latency="varies",
                    fallback_strategy=RAGStrategy.BASIC
                )
                
        except Exception as e:
            logger.warning(f"LLM routing failed: {e}, using heuristics")
        
        # Fall back to heuristic routing
        return self.route(query)
    
    def get_strategy_executor(self, strategy: RAGStrategy):
        """
        Get the executor function for a strategy.
        Import here to avoid circular dependencies.
        """
        
        executors = {
            RAGStrategy.BASIC: "core.rag_engine.query",
            RAGStrategy.SELF_RAG: "core.self_rag.self_rag_query",
            RAGStrategy.CORRECTIVE: "core.corrective_rag.correctiverag_pipeline",
            RAGStrategy.CONTEXTUAL: "core.contextual_retrieval.contextualize_chunks",
            RAGStrategy.RAPTOR: "core.raptor_rag.build_raptor_tree",
            RAGStrategy.GRAPH: "core.graph_rag.query",
            RAGStrategy.AGENTIC: "core.agentic_rag.AgenticRAG.run"
        }
        
        return executors.get(strategy, "core.rag_engine.query")


# Global router instance
_router = None

def get_router() -> AdvancedRAGRouter:
    """Get or create the global router"""
    global _router
    if _router is None:
        _router = AdvancedRAGRouter()
    return _router


def route_query(
    query: str,
    context_available: bool = True,
    prefer_speed: bool = False,
    prefer_quality: bool = False
) -> RAGRoutingDecision:
    """
    Simple interface to route a query.
    """
    router = get_router()
    return router.route(query, context_available, prefer_speed=prefer_speed, prefer_quality=prefer_quality)


# Test
if __name__ == "__main__":
    router = AdvancedRAGRouter()
    
    test_queries = [
        "What is the total revenue?",
        "Compare Q3 vs Q4 performance and explain the difference",
        "What's the latest news about our industry?",
        "Show me a chart of monthly sales trend",
        "First calculate growth rate, then identify top customers, and create a report",
        "Send email to the team with the weekly summary",
        "Give me an overview of all our financial documents"
    ]
    
    print("RAG Routing Tests:\n")
    for query in test_queries:
        decision = router.route(query)
        print(f"Query: {query[:50]}...")
        print(f"  Strategy: {decision.strategy.value}")
        print(f"  Latency: {decision.estimated_latency}")
        print(f"  Reasoning: {decision.reasoning}")
        print()