rag_system/retriever.py

"""
SQL Retriever for RAG System
Intelligent retrieval of relevant SQL examples based on question similarity and table schema analysis.
"""

import re
from typing import List, Dict, Any, Optional, Tuple
from collections import defaultdict
import numpy as np
from loguru import logger

from .vector_store import VectorStore

class SQLRetriever:
    """Intelligent SQL example retriever with schema-aware filtering."""
    
    def __init__(self, vector_store: VectorStore):
        """
        Initialize the SQL retriever.
        
        Args:
            vector_store: Initialized vector store instance
        """
        self.vector_store = vector_store
        self.schema_cache = {}  # Cache for table schema analysis
        
    def retrieve_examples(self, 
                         question: str, 
                         table_headers: List[str], 
                         top_k: int = 5,
                         use_schema_filtering: bool = True) -> List[Dict[str, Any]]:
        """
        Retrieve relevant SQL examples using multiple retrieval strategies.
        
        Args:
            question: Natural language question
            table_headers: List of table column names
            top_k: Number of examples to retrieve
            use_schema_filtering: Whether to use schema-aware filtering
            
        Returns:
            List of retrieved examples with relevance scores
        """
        # Strategy 1: Vector similarity search
        vector_results = self.vector_store.search_similar(
            query=question,
            table_headers=table_headers,
            top_k=top_k * 2,  # Get more for filtering
            similarity_threshold=0.6
        )
        
        if not vector_results:
            logger.warning("No vector search results found")
            return []
        
        # Strategy 2: Schema-aware filtering and ranking
        if use_schema_filtering:
            filtered_results = self._apply_schema_filtering(
                vector_results, question, table_headers
            )
        else:
            filtered_results = vector_results
        
        # Strategy 3: Question type classification and boosting
        enhanced_results = self._enhance_with_question_analysis(
            filtered_results, question, table_headers
        )
        
        # Strategy 4: Final ranking and selection
        final_results = self._final_ranking(
            enhanced_results, question, table_headers, top_k
        )
        
        logger.info(f"Retrieved {len(final_results)} relevant examples")
        return final_results
    
    def _apply_schema_filtering(self, 
                               results: List[Dict[str, Any]], 
                               question: str, 
                               table_headers: List[str]) -> List[Dict[str, Any]]:
        """Apply schema-aware filtering to improve relevance."""
        filtered_results = []
        
        # Analyze current table schema
        current_schema = self._analyze_schema(table_headers)
        
        for result in results:
            # Analyze example table schema
            example_headers = result["table_headers"]
            if isinstance(example_headers, str):
                example_headers = [h.strip() for h in example_headers.split(",")]
            
            example_schema = self._analyze_schema(example_headers)
            
            # Calculate schema similarity
            schema_similarity = self._calculate_schema_similarity(
                current_schema, example_schema
            )
            
            # Boost score based on schema similarity
            result["schema_similarity"] = schema_similarity
            result["enhanced_score"] = (
                result["similarity_score"] * 0.7 + 
                schema_similarity * 0.3
            )
            
            # Filter out examples with very low schema similarity
            if schema_similarity > 0.3:
                filtered_results.append(result)
        
        return filtered_results
    
    def _analyze_schema(self, table_headers: List[str]) -> Dict[str, Any]:
        """Analyze table schema for intelligent matching."""
        if not table_headers:
            return {}
        
        schema_info = {
            "column_count": len(table_headers),
            "column_types": {},
            "has_numeric": False,
            "has_text": False,
            "has_date": False,
            "has_boolean": False,
            "primary_key_candidates": [],
            "foreign_key_candidates": []
        }
        
        for header in table_headers:
            header_lower = header.lower()
            
            # Detect column types based on naming patterns
            if any(word in header_lower for word in ['id', 'key', 'pk', 'fk']):
                if 'id' in header_lower:
                    schema_info["primary_key_candidates"].append(header)
                if 'fk' in header_lower or 'foreign' in header_lower:
                    schema_info["foreign_key_candidates"].append(header)
            
            # Detect data types
            if any(word in header_lower for word in ['age', 'count', 'number', 'price', 'salary', 'amount']):
                schema_info["has_numeric"] = True
                schema_info["column_types"][header] = "numeric"
            
            if any(word in header_lower for word in ['name', 'title', 'description', 'text', 'comment']):
                schema_info["has_text"] = True
                schema_info["column_types"][header] = "text"
            
            if any(word in header_lower for word in ['date', 'time', 'created', 'updated', 'birth']):
                schema_info["has_date"] = True
                schema_info["column_types"][header] = "date"
            
            if any(word in header_lower for word in ['is_', 'has_', 'active', 'enabled', 'status']):
                schema_info["has_boolean"] = True
                schema_info["column_types"][header] = "boolean"
        
        return schema_info
    
    def _calculate_schema_similarity(self, 
                                   schema1: Dict[str, Any], 
                                   schema2: Dict[str, Any]) -> float:
        """Calculate similarity between two table schemas."""
        if not schema1 or not schema2:
            return 0.0
        
        # Column count similarity
        count_diff = abs(schema1.get("column_count", 0) - schema2.get("column_count", 0))
        count_similarity = max(0, 1 - (count_diff / max(schema1.get("column_count", 1), 1)))
        
        # Data type similarity
        type_similarity = 0.0
        if schema1.get("has_numeric") == schema2.get("has_numeric"):
            type_similarity += 0.25
        if schema1.get("has_text") == schema2.get("has_text"):
            type_similarity += 0.25
        if schema1.get("has_date") == schema2.get("has_date"):
            type_similarity += 0.25
        if schema1.get("has_boolean") == schema2.get("has_boolean"):
            type_similarity += 0.25
        
        # Primary key similarity
        pk_similarity = 0.0
        if (schema1.get("primary_key_candidates") and 
            schema2.get("primary_key_candidates")):
            pk_similarity = 0.2
        
        # Weighted combination
        final_similarity = (
            count_similarity * 0.4 +
            type_similarity * 0.4 +
            pk_similarity * 0.2
        )
        
        return final_similarity
    
    def _enhance_with_question_analysis(self, 
                                       results: List[Dict[str, Any]], 
                                       question: str, 
                                       table_headers: List[str]) -> List[Dict[str, Any]]:
        """Enhance results with question type analysis."""
        # Analyze question type
        question_type = self._classify_question_type(question)
        
        for result in results:
            # Boost examples that match question type
            if question_type in result.get("category", "").lower():
                result["enhanced_score"] *= 1.2
            
            # Boost examples with similar complexity
            question_complexity = self._assess_question_complexity(question)
            example_complexity = self._assess_question_complexity(result["question"])
            
            complexity_match = 1 - abs(question_complexity - example_complexity) / max(question_complexity, 1)
            result["enhanced_score"] *= (0.9 + complexity_match * 0.1)
        
        return results
    
    def _classify_question_type(self, question: str) -> str:
        """Classify the type of SQL question."""
        question_lower = question.lower()
        
        if any(word in question_lower for word in ['count', 'how many', 'number of']):
            return "aggregation"
        elif any(word in question_lower for word in ['average', 'mean', 'sum', 'total']):
            return "aggregation"
        elif any(word in question_lower for word in ['group by', 'grouped', 'by department', 'by category']):
            return "grouping"
        elif any(word in question_lower for word in ['join', 'combine', 'merge', 'connect']):
            return "join"
        elif any(word in question_lower for word in ['order by', 'sort', 'rank', 'top', 'highest', 'lowest']):
            return "sorting"
        elif any(word in question_lower for word in ['where', 'filter', 'condition']):
            return "filtering"
        else:
            return "general"
    
    def _assess_question_complexity(self, question: str) -> float:
        """Assess the complexity of a question (0-1 scale)."""
        complexity_score = 0.0
        
        # Length complexity
        if len(question.split()) > 20:
            complexity_score += 0.3
        elif len(question.split()) > 10:
            complexity_score += 0.2
        
        # Keyword complexity
        complex_keywords = ['join', 'group by', 'having', 'subquery', 'union', 'intersect']
        for keyword in complex_keywords:
            if keyword in question.lower():
                complexity_score += 0.15
        
        # Question type complexity
        if '?' in question:
            complexity_score += 0.1
        
        return min(1.0, complexity_score)
    
    def _final_ranking(self, 
                       results: List[Dict[str, Any]], 
                       question: str, 
                       table_headers: List[str], 
                       top_k: int) -> List[Dict[str, Any]]:
        """Final ranking and selection of examples."""
        if not results:
            return []
        
        # Sort by enhanced score
        results.sort(key=lambda x: x.get("enhanced_score", 0), reverse=True)
        
        # Ensure diversity in results
        diverse_results = []
        seen_categories = set()
        
        for result in results:
            if len(diverse_results) >= top_k:
                break
            
            category = result.get("category", "general")
            if category not in seen_categories or len(diverse_results) < top_k // 2:
                diverse_results.append(result)
                seen_categories.add(category)
        
        # Fill remaining slots with highest scoring examples
        remaining_slots = top_k - len(diverse_results)
        if remaining_slots > 0:
            for result in results:
                if result not in diverse_results and len(diverse_results) < top_k:
                    diverse_results.append(result)
        
        # Final formatting
        for result in diverse_results:
            result["final_score"] = result.get("enhanced_score", result.get("similarity_score", 0))
            # Remove internal scoring fields
            result.pop("enhanced_score", None)
            result.pop("schema_similarity", None)
        
        return diverse_results[:top_k]
    
    def get_retrieval_stats(self) -> Dict[str, Any]:
        """Get statistics about the retrieval system."""
        vector_stats = self.vector_store.get_statistics()
        
        return {
            "vector_store_stats": vector_stats,
            "schema_cache_size": len(self.schema_cache),
            "retrieval_strategies": [
                "vector_similarity",
                "schema_filtering", 
                "question_analysis",
                "diversity_ranking"
            ]
        }