src/parsers/spacy_parser.py

"""
spaCy-based query parser for rule-based keyword extraction.
"""

import re
from typing import Dict, List, Any, Optional, Set
from pathlib import Path

try:
    import spacy
    from spacy.matcher import Matcher, PhraseMatcher
    from spacy.tokens import Doc, Token
    SPACY_AVAILABLE = True
except ImportError:
    SPACY_AVAILABLE = False

from .base_parser import BaseParser, ParsedQuery
from config.settings import settings
from src.utils.logging_config import logger


class SpacyParser(BaseParser):
    """spaCy-based parser for deterministic keyword extraction."""
    
    def __init__(self, config: Optional[Dict[str, Any]] = None):
        """Initialize the spaCy parser."""
        if not SPACY_AVAILABLE:
            raise ImportError("spaCy is not installed. Please install it with: pip install spacy")
        
        super().__init__(config)
    
    def _initialize(self) -> None:
        """Initialize spaCy components."""
        model_name = self.config.get('model', settings.spacy_model)
        
        try:
            self.nlp = spacy.load(model_name)
            logger.info(f"Loaded spaCy model: {model_name}")
        except OSError:
            logger.warning(f"spaCy model '{model_name}' not found. Using blank model.")
            self.nlp = spacy.blank("en")
        
        # Initialize matchers
        self.matcher = Matcher(self.nlp.vocab)
        self.phrase_matcher = PhraseMatcher(self.nlp.vocab, attr="LOWER")
        
        # Set up patterns
        self._setup_patterns()
        
        # Configure processing pipeline
        self.max_keywords = self.config.get('max_keywords', settings.max_keywords)
        self.min_keyword_length = self.config.get('min_keyword_length', 2)
        self.exclude_pos = self.config.get('exclude_pos', {'DET', 'PRON', 'ADP', 'CONJ', 'PUNCT'})
        self.exclude_stop_words = self.config.get('exclude_stop_words', True)
    
    def _setup_patterns(self) -> None:
        """Set up matching patterns for common query structures."""
        
        # API-related patterns
        api_patterns = [
            [{"LOWER": "api"}],
            [{"LOWER": "endpoint"}],
            [{"LOWER": "service"}],
            [{"LOWER": "method"}],
            [{"LOWER": "function"}],
            [{"LOWER": "call"}],
            [{"LOWER": "request"}],
            [{"LOWER": "response"}],
        ]
        
        # Data operation patterns
        data_patterns = [
            [{"LOWER": "get"}, {"POS": "NOUN"}],
            [{"LOWER": "fetch"}, {"POS": "NOUN"}],
            [{"LOWER": "retrieve"}, {"POS": "NOUN"}],
            [{"LOWER": "search"}, {"POS": "NOUN"}],
            [{"LOWER": "find"}, {"POS": "NOUN"}],
            [{"LOWER": "list"}, {"POS": "NOUN"}],
            [{"LOWER": "create"}, {"POS": "NOUN"}],
            [{"LOWER": "update"}, {"POS": "NOUN"}],
            [{"LOWER": "delete"}, {"POS": "NOUN"}],
        ]
        
        # CDMS-specific patterns
        cdms_patterns = [
            [{"LOWER": "cdms"}],
            [{"LOWER": "label"}],
            [{"LOWER": "labels"}],
            [{"LOWER": "metadata"}],
            [{"LOWER": "dataset"}],
            [{"LOWER": "collection"}],
        ]
        
        # Add patterns to matcher
        self.matcher.add("API_TERMS", api_patterns)
        self.matcher.add("DATA_OPS", data_patterns)
        self.matcher.add("CDMS_TERMS", cdms_patterns)
        
        # Technical term phrases for phrase matcher
        technical_terms = [
            "machine learning", "data science", "artificial intelligence",
            "natural language processing", "computer vision", "deep learning",
            "neural network", "classification", "regression", "clustering",
            "feature extraction", "model training", "data preprocessing",
            "API endpoint", "REST API", "GraphQL", "JSON", "XML",
            "database query", "SQL", "NoSQL", "vector database"
        ]
        
        # Convert to spaCy docs for phrase matching
        phrase_docs = [self.nlp(term) for term in technical_terms]
        self.phrase_matcher.add("TECH_TERMS", phrase_docs)
    
    def parse(self, query: str) -> ParsedQuery:
        """
        Parse a natural language query using spaCy.
        
        Args:
            query: Input query string
            
        Returns:
            ParsedQuery with extracted information
        """
        # Preprocess query
        processed_query = self.preprocess_query(query)
        
        # Process with spaCy
        doc = self.nlp(processed_query)
        
        # Extract keywords using multiple methods
        keywords = self._extract_keywords_comprehensive(doc)
        
        # Extract entities
        entities = self._extract_entities(doc)
        
        # Determine intent (basic heuristic)
        intent = self._determine_intent(doc, keywords)
        
        # Calculate confidence
        confidence = self._calculate_confidence(doc, keywords, entities)
        
        # Create parsed query
        parsed_query = ParsedQuery(
            original_query=query,
            keywords=keywords[:self.max_keywords],
            entities=entities,
            intent=intent,
            confidence=confidence,
            metadata={
                'doc_length': len(doc),
                'sentence_count': len(list(doc.sents)),
                'token_count': len([token for token in doc if not token.is_space])
            }
        )
        
        return self.postprocess_results(parsed_query)
    
    def extract_keywords(self, query: str) -> List[str]:
        """Extract keywords from query."""
        doc = self.nlp(self.preprocess_query(query))
        return self._extract_keywords_comprehensive(doc)
    
    def _extract_keywords_comprehensive(self, doc: Doc) -> List[str]:
        """Extract keywords using multiple spaCy-based methods."""
        keywords = []
        
        # Method 1: Pattern matching
        pattern_keywords = self._extract_pattern_keywords(doc)
        keywords.extend(pattern_keywords)
        
        # Method 2: Named entities
        entity_keywords = [ent.text for ent in doc.ents if len(ent.text) >= self.min_keyword_length]
        keywords.extend(entity_keywords)
        
        # Method 3: Noun phrases
        noun_phrase_keywords = self._extract_noun_phrases(doc)
        keywords.extend(noun_phrase_keywords)
        
        # Method 4: Important tokens (nouns, verbs, adjectives)
        token_keywords = self._extract_important_tokens(doc)
        keywords.extend(token_keywords)
        
        # Remove duplicates while preserving order
        seen = set()
        unique_keywords = []
        for keyword in keywords:
            keyword_lower = keyword.lower()
            if keyword_lower not in seen and len(keyword) >= self.min_keyword_length:
                seen.add(keyword_lower)
                unique_keywords.append(keyword)
        
        return unique_keywords
    
    def _extract_pattern_keywords(self, doc: Doc) -> List[str]:
        """Extract keywords using pattern matching."""
        keywords = []
        
        # Get matches from both matchers
        matches = self.matcher(doc)
        phrase_matches = self.phrase_matcher(doc)
        
        # Process pattern matches
        for match_id, start, end in matches:
            span = doc[start:end]
            keywords.append(span.text)
        
        # Process phrase matches
        for match_id, start, end in phrase_matches:
            span = doc[start:end]
            keywords.append(span.text)
        
        return keywords
    
    def _extract_noun_phrases(self, doc: Doc) -> List[str]:
        """Extract meaningful noun phrases."""
        noun_phrases = []
        
        for chunk in doc.noun_chunks:
            # Filter out single pronouns and determiners
            if len(chunk) > 1 or (len(chunk) == 1 and chunk[0].pos_ in {'NOUN', 'PROPN'}):
                # Clean the noun phrase
                cleaned = self._clean_phrase(chunk.text)
                if cleaned and len(cleaned) >= self.min_keyword_length:
                    noun_phrases.append(cleaned)
        
        return noun_phrases
    
    def _extract_important_tokens(self, doc: Doc) -> List[str]:
        """Extract important individual tokens."""
        important_tokens = []
        
        for token in doc:
            if self._is_important_token(token):
                important_tokens.append(token.lemma_)
        
        return important_tokens
    
    def _is_important_token(self, token: Token) -> bool:
        """Check if a token is important for keyword extraction."""
        # Skip if too short
        if len(token.text) < self.min_keyword_length:
            return False
        
        # Skip stop words if configured
        if self.exclude_stop_words and token.is_stop:
            return False
        
        # Skip certain POS tags
        if token.pos_ in self.exclude_pos:
            return False
        
        # Skip punctuation and spaces
        if token.is_punct or token.is_space:
            return False
        
        # Include important POS tags
        if token.pos_ in {'NOUN', 'PROPN', 'VERB', 'ADJ'}:
            return True
        
        # Include tokens that look like technical terms or identifiers
        if re.match(r'^[A-Z][a-zA-Z0-9_]*$', token.text):  # CamelCase
            return True
        
        if re.match(r'^[a-z_][a-z0-9_]*$', token.text):  # snake_case
            return True
        
        return False
    
    def _clean_phrase(self, phrase: str) -> str:
        """Clean and normalize a phrase."""
        # Remove extra whitespace
        phrase = ' '.join(phrase.split())
        
        # Remove leading/trailing articles and prepositions
        words = phrase.split()
        while words and words[0].lower() in {'the', 'a', 'an', 'of', 'in', 'on', 'at', 'by', 'for'}:
            words = words[1:]
        
        while words and words[-1].lower() in {'of', 'in', 'on', 'at', 'by', 'for'}:
            words = words[:-1]
        
        return ' '.join(words)
    
    def _extract_entities(self, doc: Doc) -> List[Dict[str, Any]]:
        """Extract named entities from the document."""
        entities = []
        
        for ent in doc.ents:
            entities.append({
                'text': ent.text,
                'label': ent.label_,
                'description': spacy.explain(ent.label_),
                'start': ent.start_char,
                'end': ent.end_char,
                'confidence': getattr(ent, 'confidence', 1.0)
            })
        
        return entities
    
    def _determine_intent(self, doc: Doc, keywords: List[str]) -> Optional[str]:
        """Determine the intent of the query using simple heuristics."""
        text_lower = doc.text.lower()
        
        # Search/retrieval intent
        if any(word in text_lower for word in ['find', 'search', 'get', 'retrieve', 'fetch', 'show', 'list']):
            return 'search'
        
        # Creation intent
        if any(word in text_lower for word in ['create', 'make', 'add', 'new', 'generate']):
            return 'create'
        
        # Update intent
        if any(word in text_lower for word in ['update', 'modify', 'change', 'edit', 'alter']):
            return 'update'
        
        # Deletion intent
        if any(word in text_lower for word in ['delete', 'remove', 'drop', 'destroy']):
            return 'delete'
        
        # API/technical intent
        if any(word in text_lower for word in ['api', 'endpoint', 'service', 'method', 'function']):
            return 'api'
        
        # Default to search if unclear
        return 'search'
    
    def _calculate_confidence(self, doc: Doc, keywords: List[str], entities: List[Dict[str, Any]]) -> float:
        """Calculate confidence score for the parsing results."""
        confidence = 0.0
        
        # Base confidence from document structure
        if len(doc) > 0:
            confidence += 0.3
        
        # Boost for keywords found
        if keywords:
            confidence += min(0.4, len(keywords) * 0.1)
        
        # Boost for entities found
        if entities:
            confidence += min(0.2, len(entities) * 0.05)
        
        # Boost for well-formed sentences
        sentence_count = len(list(doc.sents))
        if sentence_count > 0:
            confidence += 0.1
        
        return min(1.0, confidence)