""" 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)