search_optimizer.py

"""
Search Optimization Module for Atlas Intelligent Search Management

This module contains all search decision logic and optimization algorithms to reduce
unnecessary web searches by analyzing conversation context and user intent patterns.

Key Features:
- Rule-based search decision patterns
- AI-powered search necessity analysis  
- Conversation history context analysis
- Hybrid decision engine combining rules and AI
- Search term extraction and processing utilities

Authors: Atlas Development Team
Version: 1.0.0
"""

try:
    import google.generativeai as genai
except ImportError:
    genai = None

try:
    import spacy
except ImportError:
    spacy = None

try:
    from rake_nltk import Rake
except ImportError:
    Rake = None

import asyncio
import logging
import time
import json
import re
from typing import Optional, List, Dict, Any, Tuple
from functools import wraps

# Configure logging
logger = logging.getLogger(__name__)

# AI Decision Cache for performance optimization
ai_decision_cache: Dict[str, Tuple[Dict[str, Any], float]] = {}


class SearchOptimizer:
    """
    Main search optimization class that encapsulates all search decision logic.
    
    This class provides a clean interface for search optimization functionality
    while maintaining state for NLP models and caching.
    """
    
    def __init__(self, nlp_model, rake_instance, gemini_model):
        """
        Initialize the SearchOptimizer with required dependencies.
        
        Args:
            nlp_model: spaCy language model instance
            rake_instance: RAKE keyword extraction instance  
            gemini_model: Google Generative AI model instance
        """
        self.nlp = nlp_model
        self.rake = rake_instance
        self.model = gemini_model
        logger.info("SearchOptimizer initialized successfully")


def extract_search_terms(text: str, nlp_model, rake_instance) -> List[str]:
    """
    Extract enhanced search terms using combined NER, syntax, and keywords.
    
    This function uses multiple NLP techniques to identify the most relevant
    search terms from user input:
    1. Named Entity Recognition (NER) for proper nouns
    2. Noun phrase extraction via syntactic analysis
    3. Keyword extraction using RAKE algorithm
    4. Question focus detection through dependency parsing
    
    Args:
        text (str): Input text to extract search terms from
        nlp_model: spaCy language model instance
        rake_instance: RAKE keyword extraction instance
        
    Returns:
        List[str]: Cleaned and deduplicated list of search terms
        
    Example:
        >>> extract_search_terms("What is machine learning?", nlp, rake)
        ['machine learning', 'machine', 'learning']
    """
    try:
        doc = nlp_model(text)
        
        # 1. Extract named entities
        entities = [ent.text for ent in doc.ents]
        
        # 2. Extract noun phrases through syntactic analysis
        noun_phrases = list(doc.noun_chunks)
        
        # 3. Extract question focus using dependency parsing
        focus_phrase = extract_focus_phrase(doc)
        
        # 4. Get keywords using RAKE
        rake_instance.extract_keywords_from_text(text)
        keywords = rake_instance.get_ranked_phrases()[:3]  # Top 3 keywords
        
        # Combine and filter terms
        terms = entities + [np.text for np in noun_phrases] + keywords
        if focus_phrase:
            terms.append(focus_phrase)
        
        # Clean and deduplicate
        return clean_terms(terms, nlp_model)
        
    except Exception as e:
        logger.error(f"Error extracting search terms: {e}")
        # Fallback to simple word extraction
        words = text.lower().split()
        return [word for word in words if len(word) > 2][:5]


def extract_focus_phrase(doc) -> str:
    """
    Extract main question focus using dependency parse tree analysis.
    
    This function identifies the primary focus of a question by analyzing
    the dependency relationships in the parse tree, looking for attributes,
    subjects, and objects related to the root verb.
    
    Args:
        doc: spaCy Doc object with parsed dependencies
        
    Returns:
        str: The focused phrase or empty string if none found
        
    Example:
        For "What is machine learning?", this might return "machine learning"
    """
    try:
        for token in doc:
            if token.dep_ == "ROOT":
                for child in token.children:
                    if child.dep_ in ("attr", "nsubj", "dobj"):
                        return " ".join([t.text for t in child.subtree])
        return ""
    except Exception as e:
        logger.warning(f"Error extracting focus phrase: {e}")
        return ""


def clean_terms(terms: List[str], nlp_model) -> List[str]:
    """
    Remove duplicates and irrelevant terms from extracted search terms.
    
    This function performs several cleaning operations:
    1. Removes stopwords and single characters
    2. Filters out punctuation-only terms
    3. Removes redundant subphrases
    4. Deduplicates the final list
    
    Args:
        terms (List[str]): Raw list of extracted terms
        nlp_model: spaCy language model for stopword detection
        
    Returns:
        List[str]: Cleaned and deduplicated list of search terms
    """
    try:
        # Remove stopwords and single characters
        cleaned = [
            t for t in terms 
            if len(t) > 1 and not all(token.is_stop for token in nlp_model(t))
        ]
        
        # Remove redundant subphrases
        final_terms = []
        for term in sorted(cleaned, key=len, reverse=True):
            if not any(term in other for other in final_terms):
                final_terms.append(term)
        
        return final_terms[:10]  # Limit to top 10 terms
        
    except Exception as e:
        logger.error(f"Error cleaning terms: {e}")
        return terms[:5]  # Fallback to first 5 terms


def format_search_context(results: List[Dict[str, Any]]) -> str:
    """
    Create expanded context from combined search results with richer information.
    
    This function formats search results into a readable context string
    that can be used by the AI model for generating responses.
    
    Args:
        results (List[Dict[str, Any]]): List of search result dictionaries
            Each result should have 'source', 'title', and 'body' keys
            
    Returns:
        str: Formatted context string with source attribution
        
    Example:
        >>> results = [{"source": "Brave", "title": "AI Guide", "body": "AI is..."}]
        >>> format_search_context(results)
        '[Brave] AI Guide:\nAI is...'
    """
    try:
        if not results:
            return ""
        
        formatted_results = []
        for res in results[:10]:
            # Handle None or non-dict entries gracefully
            if not isinstance(res, dict):
                continue
                
            source = res.get('source', 'Unknown')
            title = res.get('title', 'No Title')
            body = res.get('body', 'No Content')
            
            # Ensure body is a string and truncate safely
            if body:
                body_str = str(body)[:1200]
            else:
                body_str = 'No Content'
                
            formatted_results.append(f"[{source}] {title}:\n{body_str}")
        
        return "\n\n".join(formatted_results)
        
    except Exception as e:
        logger.error(f"Error formatting search context: {e}")
        return ""


def should_perform_search(prompt: str, history: Optional[List[Dict[str, str]]], 
                         search_decision_mode: str = "balanced") -> Dict[str, Any]:
    """
    Determine if web search should be performed based on conversation context and prompt patterns.
    
    This function implements rule-based search decision logic by analyzing:
    1. Conversation history presence and quality
    2. Follow-up question patterns (elaboration, clarification, referential)
    3. New information request indicators
    4. Context sufficiency for answering the question
    
    Args:
        prompt (str): User's current question/prompt
        history (Optional[List[Dict[str, str]]]): Conversation history
        search_decision_mode (str): Decision sensitivity ("conservative", "balanced", "aggressive")
        
    Returns:
        Dict[str, Any]: Dictionary containing:
            - should_search (bool): Whether to perform web search
            - reason (str): Explanation for the decision
            - confidence (float): Confidence score (0.0-1.0)
            
    Example:
        >>> should_perform_search("Tell me more about that", [{"user": "What is AI?", "assistant": "AI is..."}])
        {"should_search": False, "reason": "Follow-up question detected", "confidence": 0.8}
    """
    
    try:
        # Configuration based on search decision mode
        sensitivity_config = {
            "conservative": {
                "elaboration_threshold": 0.8,
                "referential_threshold": 0.7,
                "history_weight": 0.9
            },
            "balanced": {
                "elaboration_threshold": 0.6,
                "referential_threshold": 0.5,
                "history_weight": 0.7
            },
            "aggressive": {
                "elaboration_threshold": 0.4,
                "referential_threshold": 0.3,
                "history_weight": 0.5
            }
        }
        
        config = sensitivity_config.get(search_decision_mode, sensitivity_config["balanced"])
        prompt_lower = prompt.lower().strip()
        
        # If no history, always search (unless it's a greeting)
        if not history or len(history) == 0:
            if any(greeting in prompt_lower for greeting in ["hello", "hi", "hey", "good morning", "good afternoon"]):
                return {
                    "should_search": False,
                    "reason": "Simple greeting detected",
                    "confidence": 0.9
                }
            return {
                "should_search": True,
                "reason": "No conversation history available",
                "confidence": 1.0
            }
        
        # Pattern detection arrays
        elaboration_patterns = [
            "elaborate", "explain more", "tell me more", "expand on", "go deeper",
            "more details", "can you explain", "give me more", "detail", "expand"
        ]
        
        clarification_patterns = [
            "what do you mean", "can you clarify", "i don't understand", "unclear",
            "confusing", "what does that mean", "could you explain", "i'm confused"
        ]
        
        referential_patterns = [
            "this", "that", "it", "the previous", "above mentioned", "earlier",
            "you said", "you mentioned", "from before", "the last"
        ]
        
        continuation_patterns = [
            "and what about", "what else", "continue", "also", "additionally",
            "furthermore", "what other", "anything else", "more on"
        ]
        
        # Score patterns
        elaboration_score = sum(1 for pattern in elaboration_patterns if pattern in prompt_lower)
        clarification_score = sum(1 for pattern in clarification_patterns if pattern in prompt_lower)
        referential_score = sum(1 for pattern in referential_patterns if pattern in prompt_lower)
        continuation_score = sum(1 for pattern in continuation_patterns if pattern in prompt_lower)
        
        # Calculate total follow-up score
        total_followup_score = elaboration_score + clarification_score + referential_score + continuation_score
        
        # Analyze recent conversation history for context relevance
        history_context_score = 0
        if history:
            recent_entries = history[-3:]  # Look at last 3 exchanges
            for entry in recent_entries:
                if "role" in entry and "content" in entry:
                    if entry["role"] == "assistant":
                        content = entry["content"].lower()
                        # Check if recent assistant responses contain substantial information
                        if len(content.split()) > 20:  # Substantial response
                            history_context_score += 1
                elif "assistant" in entry:
                    content = entry["assistant"].lower()
                    if len(content.split()) > 20:
                        history_context_score += 1
        
        # Decision logic
        if total_followup_score >= 2:  # Strong follow-up indicators
            confidence = min(0.9, 0.5 + (total_followup_score * 0.2))
            return {
                "should_search": False,
                "reason": f"Follow-up question detected (score: {total_followup_score})",
                "confidence": confidence
            }
        
        if referential_score >= 1 and history_context_score >= 1:
            confidence = config["referential_threshold"] + (referential_score * 0.1)
            return {
                "should_search": False,
                "reason": "Referential question with sufficient context",
                "confidence": min(0.9, confidence)
            }
        
        if elaboration_score >= 1 and history_context_score >= 1:
            confidence = config["elaboration_threshold"]
            if elaboration_score >= 2:
                confidence += 0.2
            return {
                "should_search": False,
                "reason": "Elaboration request with existing context",
                "confidence": min(0.9, confidence)
            }
        
        # Check for new information requests
        new_info_patterns = [
            "what is", "who is", "when", "where", "how", "why", "latest", "recent",
            "current", "update", "news", "today", "now", "2024", "2025"
        ]
        
        new_info_score = sum(1 for pattern in new_info_patterns if pattern in prompt_lower)
        
        if new_info_score >= 2:
            return {
                "should_search": True,
                "reason": f"New information request detected (score: {new_info_score})",
                "confidence": 0.8
            }
        
        # Default: search for new topics
        return {
            "should_search": True,
            "reason": "New topic or insufficient context patterns",
            "confidence": 0.6
        }
        
    except Exception as e:
        logger.error(f"Error in rule-based search decision: {e}")
        return {
            "should_search": True,
            "reason": f"Error in analysis, defaulting to search: {str(e)[:50]}",
            "confidence": 0.5
        }


async def analyze_search_necessity(prompt: str, history: Optional[List[Dict[str, str]]] = None, 
                                  conversation_context: str = "", gemini_model=None) -> Dict[str, Any]:
    """
    AI-based search necessity analysis using Gemini for intelligent decision making.
    
    This function uses AI to analyze whether a web search is necessary by examining:
    1. Question type classification (new info vs clarification)
    2. Information sufficiency in conversation history
    3. Topic continuity and semantic relationships
    4. Recency requirements for the requested information
    
    Args:
        prompt (str): User's current question
        history (Optional[List[Dict[str, str]]]): Conversation history
        conversation_context (str): Formatted conversation context
        gemini_model: Google Generative AI model instance
        
    Returns:
        Dict[str, Any]: Dictionary containing:
            - should_search (bool): AI decision on search necessity
            - confidence (float): AI confidence score (0.0-1.0)
            - reason (str): Brief explanation of the decision
            - analysis (dict): Detailed analysis breakdown
            
    Example:
        >>> await analyze_search_necessity("What's the weather like?", [], "", model)
        {"should_search": True, "confidence": 0.9, "reason": "Requires current information", ...}
    """
    try:
        if not gemini_model:
            raise ValueError("Gemini model instance required for AI analysis")
            
        # Create cache key for performance optimization
        cache_key = f"{hash(prompt)}_{hash(str(history))}"
        
        # Check cache first (cache expires after 5 minutes for this session)
        current_time = time.time()
        if cache_key in ai_decision_cache:
            cached_result, timestamp = ai_decision_cache[cache_key]
            if current_time - timestamp < 300:  # 5 minute cache
                logger.info("Using cached AI search decision")
                return cached_result
        
        # Format conversation history for AI analysis
        from app import format_conversation_history  # Import to avoid circular dependency
        history_text = format_conversation_history(history, max_entries=5) if history else "No previous conversation"
        
        # Create AI prompt for search decision analysis
        analysis_prompt = f"""
        Analyze whether a web search is necessary for the following user question, considering the conversation history.

        **Conversation History:**
        {history_text}

        **Current Question:** {prompt}

        **Context:** {conversation_context[:500] if conversation_context else "No additional context"}

        Please analyze:
        1. **Question Type**: Is this asking for new information, clarification, elaboration, or continuation?
        2. **Information Sufficiency**: Does the conversation history contain enough information to answer this question?
        3. **Topic Continuity**: Is this question related to the previous conversation topics?
        4. **Recency Requirements**: Does this question require current/recent information that might not be in the history?
        5. **Semantic Relationship**: How semantically similar is this question to previous exchanges?

        Based on your analysis, determine if a web search is needed. Respond with a JSON object:
        {{
            "should_search": true/false,
            "confidence": 0.0-1.0,
            "reason": "Brief explanation of the decision",
            "analysis": {{
                "question_type": "new_information|clarification|elaboration|continuation",
                "information_sufficient": true/false,
                "topic_continuity": true/false,
                "requires_recent_info": true/false,
                "semantic_similarity": 0.0-1.0
            }}
        }}

        **Guidelines:**
        - If the question asks for new information not covered in history: should_search = true
        - If asking for clarification/elaboration of existing history content: should_search = false  
        - If asking for recent/current information (dates, news, updates): should_search = true
        - If question is semantically very similar to recent history: should_search = false
        - Confidence should reflect how certain you are about the decision
        """

        # Make AI call with timeout
        try:
            ai_response = await asyncio.wait_for(
                gemini_model.generate_content_async(analysis_prompt),
                timeout=5.0  # 5 second timeout for AI decision
            )
            
            # Parse AI response
            response_text = ai_response.text.strip()
            
            # Extract JSON from response (handle cases where AI adds extra text)
            json_match = re.search(r'\{.*\}', response_text, re.DOTALL)
            if json_match:
                json_str = json_match.group()
                result = json.loads(json_str)
                
                # Validate required fields
                required_fields = ['should_search', 'confidence', 'reason']
                if all(field in result for field in required_fields):
                    # Ensure confidence is within valid range
                    result['confidence'] = max(0.0, min(1.0, float(result['confidence'])))
                    
                    # Cache the result
                    ai_decision_cache[cache_key] = (result, current_time)
                    
                    logger.info(f"AI search decision: {result['should_search']} (confidence: {result['confidence']:.2f}) - {result['reason']}")
                    return result
                else:
                    raise ValueError("Missing required fields in AI response")
            else:
                raise ValueError("No valid JSON found in AI response")
                
        except asyncio.TimeoutError:
            logger.warning("AI search decision timed out")
            raise
        except Exception as e:
            logger.error(f"AI search decision parsing error: {e}")
            raise
            
    except Exception as e:
        logger.error(f"AI search analysis failed: {e}")
        # Return fallback decision indicating AI failure
        return {
            "should_search": True,  # Conservative fallback
            "confidence": 0.3,
            "reason": f"AI analysis failed: {str(e)[:100]}",
            "analysis": {"ai_failed": True}
        }


def has_meaningful_conversation_history(history: Optional[List[Dict[str, str]]] = None) -> bool:
    """
    Detect if request has meaningful conversation history for context-aware flow routing.
    
    This function analyzes conversation history to determine if there's sufficient
    context for making informed search decisions. It handles both conversation
    formats and validates content quality.
    
    Args:
        history (Optional[List[Dict[str, str]]]): Conversation history list
            Supports formats: [{"role": "user/assistant", "content": "..."}]
            or [{"user": "...", "assistant": "..."}]
            
    Returns:
        bool: True if conversation has meaningful history, False otherwise
        
    Example:
        >>> has_meaningful_conversation_history([{"user": "Hi", "assistant": "Hello there!"}])
        True
        >>> has_meaningful_conversation_history([])
        False
    """
    try:
        if not history or len(history) == 0:
            return False
        
        # Check for malformed history entries
        meaningful_entries = 0
        for entry in history:
            if isinstance(entry, dict):
                # Handle both formats: {"role": "user/assistant", "content": "..."} 
                # and {"user": "...", "assistant": "..."}
                if ("role" in entry and "content" in entry and 
                    entry.get("content", "").strip() and 
                    len(entry["content"].strip()) > 5):  # Minimum meaningful content
                    meaningful_entries += 1
                elif ("user" in entry and "assistant" in entry and
                      entry.get("user", "").strip() and entry.get("assistant", "").strip() and
                      len(entry["user"].strip()) > 5 and len(entry["assistant"].strip()) > 5):
                    meaningful_entries += 1
        
        # Consider history meaningful if we have at least one substantive exchange
        return meaningful_entries >= 1
        
    except Exception as e:
        logger.warning(f"Error detecting conversation history: {e}")
        return False  # Conservative fallback


def analyze_conversation_context(prompt: str, history: Optional[List[Dict[str, str]]] = None, 
                                nlp_model=None) -> Dict[str, Any]:
    """
    Analyze conversation context for semantic similarity and topic continuity.
    
    This function performs sophisticated context analysis using NLP techniques:
    1. Semantic similarity analysis using spaCy word vectors
    2. Topic continuity assessment through keyword overlap
    3. Information coverage evaluation based on history richness
    4. Context quality scoring for decision confidence
    
    Args:
        prompt (str): Current user question
        history (Optional[List[Dict[str, str]]]): Conversation history
        nlp_model: spaCy language model for semantic analysis
        
    Returns:
        Dict[str, Any]: Dictionary containing context analysis metrics:
            - topic_continuity (float): Topic overlap score (0.0-1.0)
            - semantic_similarity (float): Average semantic similarity (0.0-1.0)
            - information_coverage (float): History coverage score (0.0-1.0)
            - context_richness (float): Overall context quality (0.0-1.0)
            
    Example:
        >>> analyze_conversation_context("Tell me more", [{"assistant": "AI is..."}], nlp)
        {"topic_continuity": 0.7, "semantic_similarity": 0.8, ...}
    """
    try:
        if not nlp_model:
            raise ValueError("NLP model required for context analysis")
            
        if not history or len(history) == 0:
            return {
                "topic_continuity": 0.0,
                "semantic_similarity": 0.0,
                "information_coverage": 0.0,
                "context_richness": 0.0
            }
        
        # Use spaCy to analyze semantic similarity
        prompt_doc = nlp_model(prompt.lower())
        
        # Analyze recent conversation entries
        recent_entries = history[-3:] if len(history) > 3 else history
        
        similarity_scores = []
        topic_keywords = set()
        total_context_length = 0
        
        for entry in recent_entries:
            if "role" in entry and "content" in entry and entry["role"] == "assistant":
                content = entry["content"]
                content_doc = nlp_model(content.lower())
                
                # Calculate semantic similarity
                similarity = prompt_doc.similarity(content_doc)
                similarity_scores.append(similarity)
                
                # Extract topic keywords
                for token in content_doc:
                    if not token.is_stop and not token.is_punct and len(token.text) > 2:
                        topic_keywords.add(token.lemma_)
                
                total_context_length += len(content.split())
            
            elif "assistant" in entry:
                content = entry["assistant"]
                content_doc = nlp_model(content.lower())
                
                similarity = prompt_doc.similarity(content_doc)
                similarity_scores.append(similarity)
                
                for token in content_doc:
                    if not token.is_stop and not token.is_punct and len(token.text) > 2:
                        topic_keywords.add(token.lemma_)
                
                total_context_length += len(content.split())
        
        # Calculate metrics
        avg_similarity = sum(similarity_scores) / len(similarity_scores) if similarity_scores else 0.0
        
        # Topic continuity based on keyword overlap
        prompt_keywords = set()
        for token in prompt_doc:
            if not token.is_stop and not token.is_punct and len(token.text) > 2:
                prompt_keywords.add(token.lemma_)
        
        topic_overlap = len(prompt_keywords.intersection(topic_keywords)) / max(len(prompt_keywords), 1)
        
        # Information coverage (how much context is available)
        context_richness = min(1.0, total_context_length / 100)  # Normalize to 0-1
        
        return {
            "topic_continuity": topic_overlap,
            "semantic_similarity": avg_similarity,
            "information_coverage": len(recent_entries) / 3.0,  # Normalized to max 3 entries
            "context_richness": context_richness
        }
        
    except Exception as e:
        logger.error(f"Context analysis failed: {e}")
        return {
            "topic_continuity": 0.0,
            "semantic_similarity": 0.0,
            "information_coverage": 0.0,
            "context_richness": 0.0
        }


async def hybrid_search_decision(prompt: str, history: Optional[List[Dict[str, str]]] = None,
                               search_decision_mode: str = "balanced", nlp_model=None, 
                               gemini_model=None) -> Dict[str, Any]:
    """
    Hybrid search decision combining rule-based and AI-based analysis.
    
    This function implements the core hybrid decision engine that combines:
    1. Fast rule-based pattern matching for obvious cases
    2. AI analysis for ambiguous scenarios requiring deeper understanding
    3. Context analysis for semantic relationship assessment
    4. Confidence-based decision weighting and fallback mechanisms
    
    Args:
        prompt (str): User's current question
        history (Optional[List[Dict[str, str]]]): Conversation history
        search_decision_mode (str): Decision mode ("conservative", "balanced", "aggressive")
        nlp_model: spaCy language model for context analysis
        gemini_model: Google Generative AI model for intelligent analysis
        
    Returns:
        Dict[str, Any]: Comprehensive decision dictionary containing:
            - should_search (bool): Final search decision
            - confidence (float): Overall confidence score
            - reason (str): Explanation of decision logic
            - rule_decision (dict): Rule-based analysis results
            - ai_decision (dict, optional): AI analysis results if used
            - context_analysis (dict): Semantic context metrics
            - decision_method (str): Method used ("rule_based", "hybrid", "fallback_rule")
            
    Example:
        >>> await hybrid_search_decision("What else can you tell me?", history, "balanced", nlp, ai)
        {"should_search": False, "confidence": 0.85, "reason": "Rule-based: Follow-up detected", ...}
    """
    try:
        # Step 1: Get rule-based decision
        rule_decision = should_perform_search(prompt, history, search_decision_mode)
        
        # Step 2: Analyze conversation context
        context_analysis = analyze_conversation_context(prompt, history, nlp_model)
        
        # Step 3: Determine if AI analysis is needed
        ai_threshold = {
            "conservative": 0.8,
            "balanced": 0.6, 
            "aggressive": 0.4
        }.get(search_decision_mode, 0.6)
        
        # Use AI for ambiguous cases (low confidence rule decisions)
        if rule_decision["confidence"] < ai_threshold:
            logger.info(f"Rule confidence {rule_decision['confidence']:.2f} below threshold {ai_threshold}, using AI analysis")
            
            # Get AI decision
            from app import format_conversation_history  # Import to avoid circular dependency
            conversation_context = format_conversation_history(history, max_entries=3)
            ai_decision = await analyze_search_necessity(prompt, history, conversation_context, gemini_model)
            
            # Combine decisions with weighted confidence
            rule_weight = rule_decision["confidence"]
            ai_weight = ai_decision["confidence"]
            total_weight = rule_weight + ai_weight
            
            if total_weight > 0:
                # Weighted decision
                final_should_search = (
                    (rule_decision["should_search"] * rule_weight + 
                     ai_decision["should_search"] * ai_weight) / total_weight
                ) > 0.5
                
                final_confidence = (rule_decision["confidence"] + ai_decision["confidence"]) / 2
            else:
                # Fallback to rule decision
                final_should_search = rule_decision["should_search"]
                final_confidence = rule_decision["confidence"]
            
            return {
                "should_search": final_should_search,
                "confidence": final_confidence,
                "reason": f"Hybrid: Rule={rule_decision['reason'][:50]}..., AI={ai_decision['reason'][:50]}...",
                "rule_decision": rule_decision,
                "ai_decision": ai_decision,
                "context_analysis": context_analysis,
                "decision_method": "hybrid"
            }
        else:
            # High confidence rule decision, no need for AI
            logger.info(f"Rule confidence {rule_decision['confidence']:.2f} above threshold, using rule-based decision")
            return {
                "should_search": rule_decision["should_search"],
                "confidence": rule_decision["confidence"],
                "reason": f"Rule-based: {rule_decision['reason']}",
                "rule_decision": rule_decision,
                "context_analysis": context_analysis,
                "decision_method": "rule_based"
            }
            
    except Exception as e:
        logger.error(f"Hybrid search decision failed: {e}")
        # Fallback to rule-based decision
        rule_decision = should_perform_search(prompt, history, search_decision_mode)
        return {
            "should_search": rule_decision["should_search"],
            "confidence": rule_decision["confidence"],
            "reason": f"Fallback to rules due to error: {str(e)[:50]}",
            "rule_decision": rule_decision,
            "decision_method": "fallback_rule",
            "error": str(e)
        }


# Convenience functions for maintaining backward compatibility
def get_search_optimizer_instance(nlp_model, rake_instance, gemini_model) -> SearchOptimizer:
    """
    Factory function to create SearchOptimizer instance with dependencies.
    
    Args:
        nlp_model: spaCy language model instance
        rake_instance: RAKE keyword extraction instance
        gemini_model: Google Generative AI model instance
        
    Returns:
        SearchOptimizer: Configured optimizer instance
    """
    return SearchOptimizer(nlp_model, rake_instance, gemini_model)