signalwire-agents 0.1.13__py3-none-any.whl → 1.0.17.dev4__py3-none-any.whl

signalwire_agents/search/query_processor.py

@@ -77,22 +77,87 @@ def load_spacy_model(language: str):
             _spacy_warning_shown = True
         return None
 
-def vectorize_query(query: str):
+# Model cache - stores multiple models by name
+_model_cache = {}  # model_name -> SentenceTransformer instance
+_model_lock = None
+
+def set_global_model(model):
+    """Legacy function - adds model to cache instead of setting globally"""
+    if model and hasattr(model, 'model_name'):
+        _model_cache[model.model_name] = model
+        logger.info(f"Model added to cache: {model.model_name}")
+
+def _get_cached_model(model_name: str = None):
+    """Get or create cached sentence transformer model
+
+    Args:
+        model_name: Optional model name. If not provided, uses default.
+    """
+    global _model_cache, _model_lock
+
+    # Default model
+    if model_name is None:
+        model_name = 'sentence-transformers/all-mpnet-base-v2'
+
+    # Initialize lock if needed
+    if _model_lock is None:
+        import threading
+        _model_lock = threading.Lock()
+
+    # Check if model is already in cache
+    if model_name in _model_cache:
+        return _model_cache[model_name]
+
+    # Load model with lock to prevent race conditions
+    with _model_lock:
+        # Double check in case another thread loaded it
+        if model_name in _model_cache:
+            return _model_cache[model_name]
+
+        try:
+            from sentence_transformers import SentenceTransformer
+            logger.info(f"Loading sentence transformer model: {model_name}")
+            model = SentenceTransformer(model_name)
+            # Store the model name for identification
+            model.model_name = model_name
+            # Add to cache
+            _model_cache[model_name] = model
+            logger.info(f"Successfully loaded and cached model: {model_name}")
+            return model
+        except ImportError:
+            logger.error("sentence-transformers not available. Cannot load model.")
+            return None
+        except Exception as e:
+            logger.error(f"Failed to load model {model_name}: {e}")
+            return None
+
+def vectorize_query(query: str, model=None, model_name: str = None):
     """
     Vectorize query using sentence transformers
     Returns numpy array of embeddings
+    
+    Args:
+        query: Query string to vectorize
+        model: Optional pre-loaded model instance. If not provided, uses cached model.
+        model_name: Optional model name to use if loading a new model
     """
     try:
-        from sentence_transformers import SentenceTransformer
         import numpy as np
         
-        # Use the same model as specified in the architecture
-        model = SentenceTransformer('sentence-transformers/all-mpnet-base-v2')
+        # Use provided model or get cached one
+        if model is None:
+            model = _get_cached_model(model_name)
+            if model is None:
+                return None
+        
         embedding = model.encode(query, show_progress_bar=False)
         return embedding
         
     except ImportError:
-        logger.error("sentence-transformers not available. Cannot vectorize query.")
+        logger.error("numpy not available. Cannot vectorize query.")
+        return None
+    except Exception as e:
+        logger.error(f"Error vectorizing query: {e}")
         return None
 
 # Language to NLTK stopwords mapping
@@ -118,15 +183,28 @@ stopwords_language_map = {
 # Function to ensure NLTK resources are downloaded
 def ensure_nltk_resources():
     """Download required NLTK resources if not already present"""
-    resources = ['punkt', 'wordnet', 'averaged_perceptron_tagger', 'stopwords']
+    resources = ['punkt', 'punkt_tab', 'wordnet', 'averaged_perceptron_tagger', 'stopwords']
     for resource in resources:
         try:
-            nltk.data.find(f'corpora/{resource}')
+            # Try different paths for different resource types
+            if resource in ['punkt', 'punkt_tab']:
+                nltk.data.find(f'tokenizers/{resource}')
+            elif resource in ['wordnet']:
+                nltk.data.find(f'corpora/{resource}')
+            elif resource in ['averaged_perceptron_tagger']:
+                nltk.data.find(f'taggers/{resource}')
+            elif resource in ['stopwords']:
+                nltk.data.find(f'corpora/{resource}')
+            else:
+                nltk.data.find(f'corpora/{resource}')
         except LookupError:
             try:
+                logger.info(f"Downloading NLTK resource '{resource}'...")
                 nltk.download(resource, quiet=True)
+                logger.info(f"Successfully downloaded NLTK resource '{resource}'")
             except Exception as e:
                 logger.warning(f"Failed to download NLTK resource '{resource}': {e}")
+                # Continue without this resource - some functionality may be degraded
 
 # Initialize NLTK resources
 ensure_nltk_resources()
@@ -186,7 +264,9 @@ def remove_duplicate_words(input_string: str) -> str:
 
 def preprocess_query(query: str, language: str = 'en', pos_to_expand: Optional[List[str]] = None, 
                     max_synonyms: int = 5, debug: bool = False, vector: bool = False, 
-                    vectorize_query_param: bool = False, nlp_backend: str = 'nltk') -> Dict[str, Any]:
+                    vectorize_query_param: bool = False, nlp_backend: str = None, 
+                    query_nlp_backend: str = 'nltk', model_name: str = None,
+                    preserve_original: bool = True) -> Dict[str, Any]:
     """
     Advanced query preprocessing with language detection, POS tagging, synonym expansion, and vectorization
     
@@ -198,12 +278,19 @@ def preprocess_query(query: str, language: str = 'en', pos_to_expand: Optional[L
         debug: Enable debug output
         vector: Include vector embedding in output
         vectorize_query_param: If True, just vectorize without other processing
-        nlp_backend: NLP backend to use ('nltk' for fast, 'spacy' for better quality)
+        nlp_backend: DEPRECATED - use query_nlp_backend instead
+        query_nlp_backend: NLP backend for query processing ('nltk' for fast, 'spacy' for better quality)
         
     Returns:
         Dict containing processed query, language, POS tags, and optionally vector
     """
     
+    # Handle backward compatibility
+    if nlp_backend is not None:
+        query_nlp_backend = nlp_backend
+        if debug:
+            logger.info(f"Using deprecated 'nlp_backend' parameter, please use 'query_nlp_backend' instead")
+    
     if vectorize_query_param:
         # Vectorize the query directly
         vectorized_query = vectorize_query(query)
@@ -226,18 +313,32 @@ def preprocess_query(query: str, language: str = 'en', pos_to_expand: Optional[L
     
     # Load spaCy model based on the language and backend choice
     nlp = None
-    if nlp_backend == 'spacy':
+    if query_nlp_backend == 'spacy':
         nlp = load_spacy_model(language)
         if nlp is None and debug:
             logger.info("spaCy backend requested but not available, falling back to NLTK")
-    elif nlp_backend == 'nltk':
+    elif query_nlp_backend == 'nltk':
         if debug:
-            logger.info("Using NLTK backend for NLP processing")
+            logger.info("Using NLTK backend for query processing")
     else:
-        logger.warning(f"Unknown NLP backend '{nlp_backend}', using NLTK")
+        logger.warning(f"Unknown query NLP backend '{query_nlp_backend}', using NLTK")
+        query_nlp_backend = 'nltk'
     
     # Tokenization and stop word removal
-    tokens = nltk.word_tokenize(query)
+    try:
+        tokens = nltk.word_tokenize(query)
+    except LookupError as e:
+        # If tokenization fails, try to download punkt resources
+        logger.warning(f"NLTK tokenization failed: {e}")
+        try:
+            nltk.download('punkt', quiet=True)
+            nltk.download('punkt_tab', quiet=True)
+            tokens = nltk.word_tokenize(query)
+        except Exception as fallback_error:
+            # If all else fails, use simple split as fallback
+            logger.warning(f"NLTK tokenization fallback failed: {fallback_error}. Using simple word splitting.")
+            tokens = query.split()
+    
     nltk_language = stopwords_language_map.get(language, 'english')
     
     try:
@@ -258,7 +359,7 @@ def preprocess_query(query: str, language: str = 'en', pos_to_expand: Optional[L
     lemmas = []
     pos_tags = {}
 
-    if nlp and nlp_backend == 'spacy':
+    if nlp and query_nlp_backend == 'spacy':
         # Use spaCy for better POS tagging
         doc = nlp(" ".join(tokens))
         for token in doc:
@@ -270,27 +371,51 @@ def preprocess_query(query: str, language: str = 'en', pos_to_expand: Optional[L
             logger.info(f"POS Tagging Results (spaCy): {pos_tags}")
     else:
         # Use NLTK (default or fallback)
-        nltk_pos_tags = nltk.pos_tag(tokens)
-        for token, pos_tag in nltk_pos_tags:
-            lemma = lemmatizer.lemmatize(token, get_wordnet_pos(pos_tag)).lower()
-            stemmed = stemmer.stem(lemma)
-            lemmas.append((token.lower(), stemmed))
-            pos_tags[token.lower()] = pos_tag
-        if debug:
-            logger.info(f"POS Tagging Results (NLTK): {pos_tags}")
+        try:
+            nltk_pos_tags = nltk.pos_tag(tokens)
+            for token, pos_tag in nltk_pos_tags:
+                try:
+                    lemma = lemmatizer.lemmatize(token, get_wordnet_pos(pos_tag)).lower()
+                except Exception:
+                    # Fallback if lemmatization fails
+                    lemma = token.lower()
+                stemmed = stemmer.stem(lemma)
+                lemmas.append((token.lower(), stemmed))
+                pos_tags[token.lower()] = pos_tag
+            if debug:
+                logger.info(f"POS Tagging Results (NLTK): {pos_tags}")
+        except Exception as pos_error:
+            # Fallback if POS tagging fails completely
+            logger.warning(f"NLTK POS tagging failed: {pos_error}. Using basic token processing.")
+            for token in tokens:
+                lemma = token.lower()
+                stemmed = stemmer.stem(lemma)
+                lemmas.append((token.lower(), stemmed))
+                pos_tags[token.lower()] = 'NN'  # Default to noun
+            if debug:
+                logger.info(f"Using fallback token processing for: {tokens}")
 
     # Expanding query with synonyms
     expanded_query_set = set()
     expanded_query = []
     
+    # If preserve_original is True, always include the original query first
+    if preserve_original:
+        # Add original query terms first (maintains exact phrases)
+        original_tokens = query.lower().split()
+        for token in original_tokens:
+            if token not in expanded_query_set:
+                expanded_query.append(token)
+                expanded_query_set.add(token)
+    
     for original, lemma in lemmas:
         if original not in expanded_query_set:
             expanded_query.append(original)
             expanded_query_set.add(original)
-        if lemma not in expanded_query_set:
+        if lemma not in expanded_query_set and not preserve_original:  # Only add lemmas if not preserving original
             expanded_query.append(lemma)
             expanded_query_set.add(lemma)
-        if pos_tags.get(original) in pos_to_expand:
+        if pos_tags.get(original) in pos_to_expand and max_synonyms > 0:
             synonyms = get_synonyms(lemma, pos_tags[original], max_synonyms)
             for synonym in synonyms:
                 if synonym not in expanded_query_set:
@@ -303,19 +428,19 @@ def preprocess_query(query: str, language: str = 'en', pos_to_expand: Optional[L
 
     if debug:
         logger.info(f"Expanded Query: {final_query_str}")
-        logger.info(f"NLP Backend Used: {nlp_backend if nlp or nlp_backend == 'nltk' else 'nltk (fallback)'}")
+        logger.info(f"NLP Backend Used: {query_nlp_backend if nlp or query_nlp_backend == 'nltk' else 'nltk (fallback)'}")
     
     formatted_output = {
         'input': final_query_str,
         'enhanced_text': final_query_str,  # Alias for compatibility
         'language': language,
         'POS': pos_tags,
-        'nlp_backend_used': nlp_backend if nlp or nlp_backend == 'nltk' else 'nltk'
+        'nlp_backend_used': query_nlp_backend if nlp or query_nlp_backend == 'nltk' else 'nltk'
     }
     
     # Vectorize query if requested
     if vector:
-        vectorized_query = vectorize_query(final_query_str)
+        vectorized_query = vectorize_query(final_query_str, model_name=model_name)
         if vectorized_query is not None:
             formatted_output['vector'] = vectorized_query.tolist()
         else:
@@ -323,19 +448,25 @@ def preprocess_query(query: str, language: str = 'en', pos_to_expand: Optional[L
 
     return formatted_output
 
-def preprocess_document_content(content: str, language: str = 'en', nlp_backend: str = 'nltk') -> Dict[str, Any]:
+def preprocess_document_content(content: str, language: str = 'en', nlp_backend: str = None, 
+                               index_nlp_backend: str = 'nltk') -> Dict[str, Any]:
     """
     Preprocess document content for better searchability
     
     Args:
         content: Document content to process
         language: Language code for processing
-        nlp_backend: NLP backend to use ('nltk' for fast, 'spacy' for better quality)
+        nlp_backend: DEPRECATED - use index_nlp_backend instead
+        index_nlp_backend: NLP backend for document processing ('nltk' for fast, 'spacy' for better quality)
         
     Returns:
         Dict containing enhanced text and extracted keywords
     """
     
+    # Handle backward compatibility
+    if nlp_backend is not None:
+        index_nlp_backend = nlp_backend
+    
     # Use existing preprocessing but adapted for documents
     processed = preprocess_query(
         content,
@@ -344,7 +475,7 @@ def preprocess_document_content(content: str, language: str = 'en', nlp_backend:
         max_synonyms=2,  # Fewer synonyms for documents
         debug=False,
         vector=False,
-        nlp_backend=nlp_backend
+        query_nlp_backend=index_nlp_backend
     )
     
     # Extract key terms for keyword search