CoolFace
Apppublic

aniket47/document-intelligence-chatbot

sourceHugging Facemitupdated 11mo agoView on Hugging Face
0likes
query_router.py304 linesDownload Raw Back to components
1import re2from typing import Dict, List, Tuple, Optional3from enum import Enum4 5class QueryType(Enum):6    DOCUMENT_ONLY = "document_only"7    WEB_SEARCH = "web_search"8    HYBRID = "hybrid"9 10class QueryRouter:11    """12    Smart query routing logic to determine whether to use document search,13    web search, or both based on query characteristics14    """15    16    def __init__(self):17        # Keywords that trigger web search18        self.web_search_keywords = {19            'temporal': [20                'latest', 'recent', 'current', 'now', 'today', 'this year', 21                '2024', '2025', 'new', 'updated', 'modern', 'contemporary'22            ],23            'explanatory': [24                'explain', 'how does', 'how to', 'what is', 'what are',25                'why does', 'why is', 'tell me about', 'describe'26            ],27            'comparative': [28                'vs', 'versus', 'compare', 'comparison', 'difference between',29                'alternatives to', 'better than', 'similar to', 'like'30            ],31            'current_data': [32                'price', 'cost', 'stock', 'trend', 'trending', 'popular',33                'market', 'value', 'rate', 'statistics', 'data'34            ],35            'specifications': [36                'specs', 'specifications', 'features', 'details', 'technical',37                'performance', 'benchmark', 'review'38            ],39            'superlatives': [40                'slowest', 'biggest', 'smallest', 'best', 'worst',41                'most', 'least', 'highest', 'lowest', 'top', 'bottom',42                'largest', 'tallest', 'strongest', 'weakest'43            ],44            'factual_queries': [45                'world record', 'world', 'global', 'worldwide', 'international',46                'country', 'countries', 'nation', 'capital', 'population'47            ]48        }49        50        # Keywords that strongly suggest document search51        self.document_keywords = [52            'according to', 'in the document', 'from the file', 'mentioned',53            'stated', 'written', 'document says', 'file contains',54            'pdf', 'pdf about', 'this pdf', 'document about', 'file about',55            'resume', 'cv', 'uploaded', 'this document', 'this file'56        ]57        58        # General knowledge keywords that might need web search59        self.general_knowledge_keywords = [60            'definition', 'meaning', 'concept', 'theory', 'principle',61            'history', 'background', 'overview', 'introduction'62        ]63    64    def analyze_query(self, query: str) -> Dict:65        """66        Analyze query to determine routing strategy67        68        Args:69            query: User query string70            71        Returns:72            Dictionary with routing analysis73        """74        query_lower = query.lower()75        76        # Initialize analysis77        analysis = {78            'query': query,79            'web_indicators': [],80            'document_indicators': [],81            'confidence_scores': {82                'web_search': 0.0,83                'document_search': 0.084            },85            'suggested_route': QueryType.DOCUMENT_ONLY,86            'reasoning': []87        }88        89        # Check for web search indicators90        web_score = 091        for category, keywords in self.web_search_keywords.items():92            for keyword in keywords:93                if keyword in query_lower:94                    analysis['web_indicators'].append(f"{keyword} ({category})")95                    web_score += self._get_keyword_weight(category)96        97        # Check for document indicators98        doc_score = 099        for keyword in self.document_keywords:100            if keyword in query_lower:101                analysis['document_indicators'].append(keyword)102                doc_score += 2.0  # High weight for explicit document references103        104        # Check for general knowledge that might need web search105        for keyword in self.general_knowledge_keywords:106            if keyword in query_lower:107                analysis['web_indicators'].append(f"{keyword} (general_knowledge)")108                web_score += 0.5109        110        # Question word analysis111        question_words = ['how', 'what', 'why', 'when', 'where', 'who', 'which']112        question_count = sum(1 for word in question_words if word in query_lower.split())113        if question_count > 0:114            web_score += 0.3 * question_count115        116        # Length analysis (longer queries often need more context)117        if len(query.split()) > 10:118            web_score += 0.2119        120        # Normalize scores121        max_possible_score = 10.0122        analysis['confidence_scores']['web_search'] = min(web_score / max_possible_score, 1.0)123        analysis['confidence_scores']['document_search'] = min(doc_score / max_possible_score, 1.0)124        125        # If no explicit document indicators, boost document search slightly126        if doc_score == 0:127            analysis['confidence_scores']['document_search'] = 0.3128        129        # Determine routing strategy130        web_confidence = analysis['confidence_scores']['web_search']131        doc_confidence = analysis['confidence_scores']['document_search']132        133        if doc_confidence > 0.7:  # Strong document indicators134            analysis['suggested_route'] = QueryType.DOCUMENT_ONLY135            analysis['reasoning'].append("Strong document reference indicators")136        elif web_confidence > 0.35:  # Even lower threshold for web search137            analysis['suggested_route'] = QueryType.WEB_SEARCH138            analysis['reasoning'].append("Web search indicators detected")139        elif web_confidence > 0.25 and doc_confidence > 0.3:  # Mixed signals140            analysis['suggested_route'] = QueryType.HYBRID141            analysis['reasoning'].append("Mixed indicators suggest hybrid approach")142        else:  # Default to document search when documents are available143            analysis['suggested_route'] = QueryType.DOCUMENT_ONLY144            analysis['reasoning'].append("Default to document search - prefer uploaded documents")145        146        return analysis147    148    def _get_keyword_weight(self, category: str) -> float:149        """Get weight for different keyword categories"""150        weights = {151            'temporal': 1.5,        # Strong indicator for web search152            'explanatory': 0.8,     # Medium indicator153            'comparative': 1.2,     # Strong indicator  154            'current_data': 1.5,    # Strong indicator155            'specifications': 1.0,  # Medium indicator156            'superlatives': 1.8,    # Very strong indicator for web search157            'factual_queries': 1.6  # Strong indicator for web search158        }159        return weights.get(category, 0.5)160    161    def should_use_web_search(self, query: str, document_results: List = None) -> Tuple[bool, str]:162        """163        Determine if web search should be used based on query and document results164        165        Args:166            query: User query167            document_results: Results from document search (if any)168            169        Returns:170            Tuple of (should_use_web, reasoning)171        """172        analysis = self.analyze_query(query)173        174        # Always use web search if suggested route is WEB_SEARCH175        if analysis['suggested_route'] == QueryType.WEB_SEARCH:176            return True, "Query indicates need for web search"177        178        # For hybrid queries, be more conservative - prefer documents when available179        if analysis['suggested_route'] == QueryType.HYBRID:180            if not document_results or len(document_results) == 0:181                return True, "Hybrid query with no document results"182            elif len(document_results) > 0:183                # Check quality of document results - lowered threshold to prefer documents184                best_score = max([r.get('score', 0) for r in document_results])185                if best_score < 0.05:  # Very low similarity scores only186                    return True, "Hybrid query with very low-quality document results"187        188        # For document-only queries, almost never use web search189        if analysis['suggested_route'] == QueryType.DOCUMENT_ONLY:190            # Only use web search if absolutely no document results191            if document_results is not None and len(document_results) == 0:192                return True, "No document results found, falling back to web search"193        194        return False, "Document search should be sufficient"195    196    def get_routing_explanation(self, query: str) -> str:197        """198        Get human-readable explanation of routing decision199        200        Args:201            query: User query202            203        Returns:204            Explanation string205        """206        analysis = self.analyze_query(query)207        208        explanation = f"**Query Analysis for:** {query}\n\n"209        210        if analysis['web_indicators']:211            explanation += "**Web Search Indicators Found:**\n"212            for indicator in analysis['web_indicators'][:3]:  # Show top 3213                explanation += f"- {indicator}\n"214            explanation += "\n"215        216        if analysis['document_indicators']:217            explanation += "**Document Search Indicators Found:**\n"218            for indicator in analysis['document_indicators']:219                explanation += f"- {indicator}\n"220            explanation += "\n"221        222        explanation += f"**Suggested Strategy:** {analysis['suggested_route'].value}\n\n"223        224        if analysis['reasoning']:225            explanation += "**Reasoning:** " + ", ".join(analysis['reasoning'])226        227        return explanation228    229    def analyze_query_semantic(self, query: str, vector_store=None, similarity_threshold: float = 0.15) -> Dict:230        """231        Semantic-based query routing using embedding similarity to determine232        if the query is relevant to indexed documents233        234        Args:235            query: User's input query236            vector_store: VectorStore instance with indexed documents237            similarity_threshold: Minimum similarity score to prefer documents (0.0-1.0)238            239        Returns:240            Dict with routing decision and reasoning241        """242        try:243            # If no vector store or no documents, default to web search244            if not vector_store or not hasattr(vector_store, 'search') or len(getattr(vector_store, 'documents', [])) == 0:245                return {246                    'suggested_route': QueryType.WEB_SEARCH,247                    'reasoning': ['No documents available - using web search'],248                    'similarity_score': 0.0249                }250            251            # Still check for strong temporal indicators that should always use web search252            temporal_keywords = ['latest', 'recent', 'current', 'now', 'today', 'this year', '2024', '2025', 'breaking', 'news']253            query_lower = query.lower()254            255            for keyword in temporal_keywords:256                if keyword in query_lower:257                    return {258                        'suggested_route': QueryType.WEB_SEARCH,259                        'reasoning': [f'Temporal keyword "{keyword}" detected - using web search for current information'],260                        'similarity_score': 0.0261                    }262            263            # Get semantic similarity with documents264            try:265                # Search for similar documents266                results = vector_store.search(query, k=3)267                268                if not results:269                    return {270                        'suggested_route': QueryType.WEB_SEARCH,271                        'reasoning': ['No document matches found - using web search'],272                        'similarity_score': 0.0273                    }274                275                # Get the best similarity score276                best_score = max([r.get('score', 0) for r in results])277                278                print(f"DEBUG: Semantic routing - Query: '{query[:50]}...', Best similarity: {best_score:.3f}, Threshold: {similarity_threshold}")279                280                if best_score >= similarity_threshold:281                    return {282                        'suggested_route': QueryType.DOCUMENT_ONLY,283                        'reasoning': [f'High document relevance (score: {best_score:.3f}) - using document search'],284                        'similarity_score': best_score285                    }286                else:287                    return {288                        'suggested_route': QueryType.WEB_SEARCH,289                        'reasoning': [f'Low document relevance (score: {best_score:.3f}) - using web search'],290                        'similarity_score': best_score291                    }292                    293            except Exception as search_error:294                print(f"DEBUG: Semantic search failed: {search_error}")295                return {296                    'suggested_route': QueryType.WEB_SEARCH,297                    'reasoning': ['Document search failed - using web search'],298                    'similarity_score': 0.0299                }300                301        except Exception as e:302            print(f"DEBUG: Semantic routing error: {e}")303            # Fallback to keyword-based routing304            return self.analyze_query(query)