query_processing.py

"""
Query processing module: normalization, expansion, and rewriting.
Handles query preprocessing to improve retrieval quality.
"""

import re
import unicodedata
from typing import Dict, List, Optional

from config import (
 ABBREVIATION_MAP,
 ENABLE_QUERY_NORMALIZATION,
 ENABLE_QUERY_EXPANSION,
 MAX_QUERY_EXPANSIONS,
)


class QueryNormalizer:
 """Production-grade query normalization pipeline."""
 
 def __init__(self, abbreviations: Dict[str, str] = None):
 self.abbreviations = abbreviations or ABBREVIATION_MAP
 self._whitespace_re = re.compile(r'\s+')
 self._special_chars_re = re.compile(r'[^\w\s\-\'\"?.,!]')
 self._repeated_punct_re = re.compile(r'([?!.,])\1+')
 
 def normalize(self, query: str) -> Dict[str, any]:
 """Full normalization: unicode → whitespace → punctuation → abbreviations."""
 if not query or not query.strip():
 return {
 "original": query,
 "normalized": "",
 "display": "",
 "expansions": []
 }
 
 original = query
 expansions = []
 
 # Unicode normalize + clean
 text = unicodedata.normalize('NFKC', query).strip()
 text = self._whitespace_re.sub(' ', text)
 text = self._special_chars_re.sub('', text)
 text = self._repeated_punct_re.sub(r'\1', text)
 
 display = text
 text_lower = text.lower()
 
 # Expand abbreviations
 words = text_lower.split()
 expanded = []
 for word in words:
 word_clean = word.rstrip('?.,!"\'-')
 trailing = word[len(word_clean):] if len(word_clean) < len(word) else ''
 if word_clean in self.abbreviations:
 exp = self.abbreviations[word_clean]
 expansions.append(f"{word_clean} → {exp}")
 expanded.append(exp + trailing)
 else:
 expanded.append(word)
 
 return {
 "original": original,
 "normalized": ' '.join(expanded),
 "display": display,
 "expansions": expansions,
 }
 
 def quick_normalize(self, query: str) -> str:
 """Quick normalization returning only the normalized string."""
 return self.normalize(query)["normalized"]


class QueryExpander:
 """Generates semantic query variants for improved recall."""
 
 PATTERNS = [
 (r'^what is (the )?(.*?)\??$', ['define {0}', 'explain {0}', '{0} definition']),
 (r'^how (do|does|can|to) (.*?)\??$', ['method for {0}', 'process of {0}']),
 (r'^why (is|are|do|does) (.*?)\??$', ['reason for {0}', 'cause of {0}']),
 ]
 
 STOP_WORDS = {
 'what', 'is', 'are', 'the', 'a', 'an', 'how', 'do', 'does', 'can',
 'why', 'when', 'where', 'which', 'who', 'to', 'for', 'in', 'on',
 'at', 'by', 'with', 'about', 'of', 'from', 'this', 'that', 'it',
 'be', 'been', 'being', 'have', 'has', 'had', 'and', 'or', 'but',
 }
 
 def __init__(self, max_expansions: int = MAX_QUERY_EXPANSIONS):
 self.max_expansions = max_expansions
 
 def expand(self, query: str) -> List[str]:
 """Generate query variants for improved recall."""
 variants = [query]
 query_lower = query.lower().strip()
 
 # Pattern-based expansion
 for pattern, templates in self.PATTERNS:
 match = re.match(pattern, query_lower, re.IGNORECASE)
 if match:
 content = match.groups()[-1]
 for tmpl in templates[:self.max_expansions]:
 variant = tmpl.format(content)
 if variant not in variants:
 variants.append(variant)
 break
 
 # Keyword extraction variant
 keywords = self._extract_keywords(query)
 if keywords and keywords not in variants:
 variants.append(keywords)
 
 return variants[:self.max_expansions + 1]
 
 def _extract_keywords(self, query: str) -> str:
 """Extract meaningful keywords from query."""
 words = re.findall(r'\b\w+\b', query.lower())
 keywords = [w for w in words if w not in self.STOP_WORDS and len(w) > 2]
 return ' '.join(keywords) if keywords else ''


class QueryProcessor:
 """Main query processing orchestrator."""
 
 def __init__(
 self,
 enable_normalization: bool = ENABLE_QUERY_NORMALIZATION,
 enable_expansion: bool = ENABLE_QUERY_EXPANSION,
 ):
 self.normalizer = QueryNormalizer()
 self.expander = QueryExpander()
 self.enable_normalization = enable_normalization
 self.enable_expansion = enable_expansion
 
 def process(self, query: str) -> Dict:
 """Process query: normalize → expand."""
 result = {
 "original": query,
 "normalized": query,
 "variants": [query],
 "expansions": []
 }
 
 if self.enable_normalization:
 norm = self.normalizer.normalize(query)
 result["normalized"] = norm["normalized"]
 result["expansions"] = norm["expansions"]
 
 if self.enable_expansion:
 result["variants"] = self.expander.expand(result["normalized"])
 
 return result


# Global instance
query_processor = QueryProcessor()