learning_store.py

"""
Learning Store for Self-Improving SQL Generation System.

This module provides persistent storage and retrieval of query patterns
to enable continuous learning and improvement of SQL generation.

Classes:
    LearningStore: SQLite-based pattern storage with embedding search

Features:
    - Storage for successful query→SQL pattern mappings
    - User feedback tracking (thumbs up/down ratings)
    - Error pattern tracking for mistake avoidance
    - Semantic similarity search using WatsonX embeddings (preferred)
    - Keyword-based similarity search (fallback when embeddings unavailable)
    - Pattern normalization for deduplication

Database Schema:
    query_patterns:
        - user_query: Original natural language query
        - generated_sql: Successfully executed SQL
        - embedding: 384-dimensional vector (IBM Slate)
        - thumbs_up/down: User feedback counts
        - execution_time_ms: Performance tracking

    error_patterns:
        - user_query: Query that caused error
        - attempted_sql: Failed SQL
        - error_message: Error description
        - occurrence_count: How often this error happens

Usage:
    >>> from learning_store import get_learning_store
    >>>
    >>> store = get_learning_store()  # Singleton instance
    >>>
    >>> # Find similar patterns
    >>> patterns = store.find_similar_patterns("show laptop revenue")
    >>> for p in patterns:
    ...     print(f"{p['similarity']:.2f}: {p['user_query']}")
    >>>
    >>> # Store successful pattern
    >>> pattern_id = store.store_success_pattern(
    ...     "show laptops",
    ...     "SELECT * FROM products WHERE product_name LIKE '%laptop%'"
    ... )
    >>>
    >>> # Record feedback
    >>> store.record_feedback(pattern_id, is_positive=True)

Author: Markus van Kempen (markus.van.kempen@gmail.com)
"""

import asyncio
import json
import math
import os
import re
import sqlite3
from typing import Any

from dotenv import load_dotenv

load_dotenv()

from logging_config import get_logger

logger = get_logger(__name__)

# Try to import embedding model
try:
    from beeai_framework.adapters.watsonx import WatsonxEmbeddingModel

    EMBEDDINGS_AVAILABLE = True
except ImportError:
    EMBEDDINGS_AVAILABLE = False
    logger.warning("Embeddings not available, using keyword matching")


class LearningStore:
    """
    SQLite-based learning store for self-improving SQL generation.

    Tables:
    - query_patterns: Successful query→SQL mappings with feedback
    - error_patterns: Failed queries and their errors
    - prompt_rules: Learned rules extracted from patterns

    Features:
    - Embedding-based semantic similarity search (preferred)
    - Keyword-based similarity search (fallback)
    """

    def __init__(self, db_path: str = None, use_embeddings: bool = True):
        """Initialize the learning store."""
        self.db_path = db_path or os.getenv("LEARNING_DB_PATH", "./data/learning.db")
        if not os.path.isabs(self.db_path):
            self.db_path = os.path.abspath(self.db_path)

        # Ensure directory exists
        os.makedirs(os.path.dirname(self.db_path), exist_ok=True)

        # Embedding model
        self.use_embeddings = use_embeddings and EMBEDDINGS_AVAILABLE
        self._embedding_model = None

        if self.use_embeddings:
            self._init_embedding_model()

        self._init_database()

    def _init_embedding_model(self):
        """Initialize the embedding model for semantic search."""
        try:
            self._embedding_model = WatsonxEmbeddingModel(
                model_id="ibm/slate-125m-english-rtrvr-v2",
                api_key=os.getenv("WATSONX_API_KEY"),
                project_id=os.getenv("WATSONX_PROJECT_ID"),
                url=os.getenv("WATSONX_URL"),
            )
            logger.info("Embedding model initialized for semantic search")
        except Exception as e:
            logger.warning(f"Failed to init embedding model: {e}. Using keyword matching.")
            self._embedding_model = None
            self.use_embeddings = False

    def _init_database(self):
        """Initialize database tables."""
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()

        # Successful query patterns
        cursor.execute(
            """
            CREATE TABLE IF NOT EXISTS query_patterns (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                user_query TEXT NOT NULL,
                generated_sql TEXT NOT NULL,
                normalized_query TEXT,
                query_type TEXT,
                success BOOLEAN DEFAULT 1,
                thumbs_up INTEGER DEFAULT 0,
                thumbs_down INTEGER DEFAULT 0,
                execution_time_ms REAL,
                result_count INTEGER,
                model_id TEXT,
                mode TEXT,
                embedding TEXT,
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
            )
        """
        )

        # Error patterns for learning what NOT to do
        cursor.execute(
            """
            CREATE TABLE IF NOT EXISTS error_patterns (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                user_query TEXT NOT NULL,
                failed_sql TEXT,
                error_message TEXT,
                error_type TEXT,
                fixed_sql TEXT,
                model_id TEXT,
                mode TEXT,
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
            )
        """
        )

        # Learned prompt rules
        cursor.execute(
            """
            CREATE TABLE IF NOT EXISTS prompt_rules (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                rule_type TEXT NOT NULL,
                pattern TEXT NOT NULL,
                replacement TEXT,
                description TEXT,
                success_count INTEGER DEFAULT 0,
                failure_count INTEGER DEFAULT 0,
                is_active BOOLEAN DEFAULT 1,
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
            )
        """
        )

        # Query keywords index for similarity search
        cursor.execute(
            """
            CREATE TABLE IF NOT EXISTS query_keywords (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                pattern_id INTEGER,
                keyword TEXT NOT NULL,
                weight REAL DEFAULT 1.0,
                FOREIGN KEY (pattern_id) REFERENCES query_patterns(id)
            )
        """
        )

        # Add embedding column if it doesn't exist (migration)
        try:
            cursor.execute("ALTER TABLE query_patterns ADD COLUMN embedding TEXT")
            logger.info("Added embedding column to query_patterns")
        except sqlite3.OperationalError:
            pass  # Column already exists

        # Create indexes
        cursor.execute("CREATE INDEX IF NOT EXISTS idx_patterns_query ON query_patterns(normalized_query)")
        cursor.execute("CREATE INDEX IF NOT EXISTS idx_keywords ON query_keywords(keyword)")
        cursor.execute("CREATE INDEX IF NOT EXISTS idx_patterns_success ON query_patterns(success, thumbs_up)")

        conn.commit()
        conn.close()
        logger.info(f"Learning store initialized at {self.db_path}")

    def _normalize_query(self, query: str) -> str:
        """Normalize a query for comparison."""
        # Lowercase
        normalized = query.lower().strip()
        # Remove extra whitespace
        normalized = re.sub(r"\s+", " ", normalized)
        # Remove punctuation except apostrophes
        normalized = re.sub(r"[^\w\s\']", "", normalized)
        return normalized

    def _extract_keywords(self, query: str) -> list[str]:
        """Extract meaningful keywords from a query."""
        # Common SQL-related words to keep
        important_words = {
            "sales",
            "revenue",
            "total",
            "sum",
            "count",
            "average",
            "avg",
            "product",
            "products",
            "customer",
            "customers",
            "order",
            "orders",
            "country",
            "region",
            "city",
            "category",
            "laptop",
            "phone",
            "smartphone",
            "tablet",
            "keyboard",
            "mouse",
            "usa",
            "uk",
            "germany",
            "france",
            "japan",
            "top",
            "best",
            "highest",
            "lowest",
            "most",
            "least",
            "by",
            "per",
            "each",
            "group",
            "filter",
            "last",
            "month",
            "year",
            "week",
            "today",
            "yesterday",
            "compare",
            "between",
            "and",
            "or",
        }

        # Stopwords to remove
        stopwords = {
            "show",
            "me",
            "the",
            "a",
            "an",
            "of",
            "in",
            "to",
            "for",
            "with",
            "what",
            "how",
            "is",
            "are",
            "our",
            "my",
        }

        words = self._normalize_query(query).split()
        keywords = [w for w in words if w in important_words or (len(w) > 2 and w not in stopwords)]

        return keywords

    def _get_embedding(self, text: str) -> list[float] | None:
        """Get embedding vector for a text using WatsonX."""
        if not self._embedding_model:
            return None

        try:
            loop = asyncio.new_event_loop()
            asyncio.set_event_loop(loop)
            try:
                result = loop.run_until_complete(self._embedding_model.create([text]))
                return result.embeddings[0] if result.embeddings else None
            finally:
                loop.close()
        except Exception as e:
            logger.warning(f"Failed to get embedding: {e}")
            return None

    def _cosine_similarity(self, a: list[float], b: list[float]) -> float:
        """Calculate cosine similarity between two vectors."""
        dot = sum(x * y for x, y in zip(a, b))
        norm_a = math.sqrt(sum(x * x for x in a))
        norm_b = math.sqrt(sum(x * x for x in b))
        if norm_a == 0 or norm_b == 0:
            return 0.0
        return dot / (norm_a * norm_b)

    def store_successful_pattern(
        self,
        user_query: str,
        generated_sql: str,
        result_count: int = 0,
        execution_time_ms: float = 0,
        model_id: str = None,
        mode: str = None,
    ) -> int:
        """Store a successful query pattern with embedding."""
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()

        normalized = self._normalize_query(user_query)
        keywords = self._extract_keywords(user_query)

        # Determine query type
        query_type = self._classify_query(user_query, generated_sql)

        # Get embedding for semantic search
        embedding = self._get_embedding(user_query)
        embedding_json = json.dumps(embedding) if embedding else None

        # Check if similar pattern exists
        cursor.execute(
            """
            SELECT id FROM query_patterns
            WHERE normalized_query = ? AND success = 1
            LIMIT 1
        """,
            (normalized,),
        )

        existing = cursor.fetchone()

        if existing:
            # Update existing pattern
            pattern_id = existing[0]
            cursor.execute(
                """
                UPDATE query_patterns
                SET thumbs_up = thumbs_up + 1,
                    updated_at = CURRENT_TIMESTAMP
                WHERE id = ?
            """,
                (pattern_id,),
            )
        else:
            # Insert new pattern with embedding
            cursor.execute(
                """
                INSERT INTO query_patterns
                (user_query, generated_sql, normalized_query, query_type,
                 success, result_count, execution_time_ms, model_id, mode, embedding)
                VALUES (?, ?, ?, ?, 1, ?, ?, ?, ?, ?)
            """,
                (
                    user_query,
                    generated_sql,
                    normalized,
                    query_type,
                    result_count,
                    execution_time_ms,
                    model_id,
                    mode,
                    embedding_json,
                ),
            )

            pattern_id = cursor.lastrowid

            # Store keywords (for fallback search)
            for keyword in keywords:
                cursor.execute(
                    """
                    INSERT INTO query_keywords (pattern_id, keyword)
                    VALUES (?, ?)
                """,
                    (pattern_id, keyword),
                )

        conn.commit()
        conn.close()

        logger.debug(
            f"Stored pattern {pattern_id} with {'embedding' if embedding else 'keywords'}: {user_query[:50]}..."
        )
        return pattern_id

    def store_error_pattern(
        self,
        user_query: str,
        failed_sql: str,
        error_message: str,
        fixed_sql: str = None,
        model_id: str = None,
        mode: str = None,
    ) -> int:
        """Store a failed query pattern for learning."""
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()

        # Classify error type
        error_type = self._classify_error(error_message)

        cursor.execute(
            """
            INSERT INTO error_patterns
            (user_query, failed_sql, error_message, error_type, fixed_sql, model_id, mode)
            VALUES (?, ?, ?, ?, ?, ?, ?)
        """,
            (user_query, failed_sql, error_message, error_type, fixed_sql, model_id, mode),
        )

        pattern_id = cursor.lastrowid
        conn.commit()
        conn.close()

        logger.debug(f"Stored error pattern {pattern_id}: {error_type}")
        return pattern_id

    def record_feedback(self, pattern_id: int, is_positive: bool) -> bool:
        """Record user feedback (thumbs up/down)."""
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()

        if is_positive:
            cursor.execute(
                """
                UPDATE query_patterns
                SET thumbs_up = thumbs_up + 1, updated_at = CURRENT_TIMESTAMP
                WHERE id = ?
            """,
                (pattern_id,),
            )
        else:
            cursor.execute(
                """
                UPDATE query_patterns
                SET thumbs_down = thumbs_down + 1, updated_at = CURRENT_TIMESTAMP
                WHERE id = ?
            """,
                (pattern_id,),
            )

        success = cursor.rowcount > 0
        conn.commit()
        conn.close()

        logger.debug(f"Recorded {'positive' if is_positive else 'negative'} feedback for pattern {pattern_id}")
        return success

    def find_similar_patterns(self, user_query: str, limit: int = 3, min_score: float = 0.3) -> list[dict[str, Any]]:
        """Find similar successful patterns using embeddings (preferred) or keywords (fallback)."""

        # Try embedding-based search first
        if self.use_embeddings and self._embedding_model:
            results = self._find_similar_by_embedding(user_query, limit, min_score)
            if results:
                logger.debug(f"Found {len(results)} similar patterns via embeddings")
                return results

        # Fallback to keyword matching
        return self._find_similar_by_keywords(user_query, limit, min_score)

    def _find_similar_by_embedding(
        self, user_query: str, limit: int = 3, min_score: float = 0.6  # Higher threshold for embeddings
    ) -> list[dict[str, Any]]:
        """Find similar patterns using embedding similarity."""
        try:
            # Get embedding for query
            query_embedding = self._get_embedding(user_query)
            if not query_embedding:
                return []

            conn = sqlite3.connect(self.db_path)
            cursor = conn.cursor()

            # Get all patterns with embeddings
            cursor.execute(
                """
                SELECT id, user_query, generated_sql, query_type,
                       thumbs_up, thumbs_down, embedding
                FROM query_patterns
                WHERE success = 1 AND embedding IS NOT NULL
                  AND (thumbs_up - thumbs_down) >= 0
            """
            )

            results = []
            for row in cursor.fetchall():
                stored_embedding = json.loads(row[6])
                similarity = self._cosine_similarity(query_embedding, stored_embedding)

                if similarity >= min_score:
                    results.append(
                        {
                            "id": row[0],
                            "user_query": row[1],
                            "sql": row[2],
                            "query_type": row[3],
                            "thumbs_up": row[4],
                            "thumbs_down": row[5],
                            "similarity": similarity,
                            "method": "embedding",
                        }
                    )

            conn.close()

            # Sort by similarity and return top N
            results.sort(key=lambda x: x["similarity"], reverse=True)
            return results[:limit]

        except Exception as e:
            logger.error(f"Error in embedding search: {e}")
            return []

    def _find_similar_by_keywords(
        self, user_query: str, limit: int = 3, min_score: float = 0.3
    ) -> list[dict[str, Any]]:
        """Find similar patterns using keyword matching (fallback)."""
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()

        keywords = self._extract_keywords(user_query)

        if not keywords:
            conn.close()
            return []

        try:
            placeholders = ",".join(["?" for _ in keywords])
            cursor.execute(
                f"""
                SELECT
                    p.id, p.user_query, p.generated_sql, p.query_type,
                    p.thumbs_up, p.thumbs_down,
                    COUNT(k.keyword) as match_count,
                    (p.thumbs_up - p.thumbs_down) as score
                FROM query_patterns p
                JOIN query_keywords k ON p.id = k.pattern_id
                WHERE k.keyword IN ({placeholders})
                  AND p.success = 1
                  AND (p.thumbs_up - p.thumbs_down) >= 0
                GROUP BY p.id, p.user_query, p.generated_sql, p.query_type, p.thumbs_up, p.thumbs_down
                HAVING match_count >= ?
                ORDER BY match_count DESC, score DESC
                LIMIT ?
            """,
                (*keywords, max(1, len(keywords) // 3), limit),
            )

            results = []
            for row in cursor.fetchall():
                similarity = row[6] / len(keywords)
                if similarity >= min_score:
                    results.append(
                        {
                            "id": row[0],
                            "user_query": row[1],
                            "sql": row[2],
                            "query_type": row[3],
                            "thumbs_up": row[4],
                            "thumbs_down": row[5],
                            "similarity": similarity,
                            "method": "keywords",
                        }
                    )

            conn.close()
            return results
        except Exception as e:
            logger.error(f"Error finding similar patterns: {e}")
            conn.close()
            return []

    def get_best_examples(self, query_type: str = None, limit: int = 5) -> list[dict[str, Any]]:
        """Get highest-rated examples for few-shot learning."""
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()

        if query_type:
            cursor.execute(
                """
                SELECT id, user_query, generated_sql, query_type, thumbs_up, thumbs_down
                FROM query_patterns
                WHERE success = 1 AND query_type = ?
                ORDER BY (thumbs_up - thumbs_down) DESC, thumbs_up DESC
                LIMIT ?
            """,
                (query_type, limit),
            )
        else:
            cursor.execute(
                """
                SELECT id, user_query, generated_sql, query_type, thumbs_up, thumbs_down
                FROM query_patterns
                WHERE success = 1
                ORDER BY (thumbs_up - thumbs_down) DESC, thumbs_up DESC
                LIMIT ?
            """,
                (limit,),
            )

        results = [
            {
                "id": row[0],
                "user_query": row[1],
                "sql": row[2],
                "query_type": row[3],
                "thumbs_up": row[4],
                "thumbs_down": row[5],
            }
            for row in cursor.fetchall()
        ]

        conn.close()
        return results

    def get_common_errors(self, limit: int = 10) -> list[dict[str, Any]]:
        """Get common error patterns for prompt improvement."""
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()

        try:
            cursor.execute(
                """
                SELECT error_type, COUNT(*) as count,
                       GROUP_CONCAT(error_message, ' | ') as examples
                FROM error_patterns
                GROUP BY error_type
                ORDER BY count DESC
                LIMIT ?
            """,
                (limit,),
            )

            results = [{"error_type": row[0], "count": row[1], "examples": row[2]} for row in cursor.fetchall()]

            conn.close()
            return results
        except Exception as e:
            logger.error(f"Error getting common errors: {e}")
            conn.close()
            return []

    def get_statistics(self) -> dict[str, Any]:
        """Get learning store statistics."""
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()

        stats = {}

        # Total patterns
        cursor.execute("SELECT COUNT(*) FROM query_patterns WHERE success = 1")
        stats["total_patterns"] = cursor.fetchone()[0]

        # Patterns with positive feedback
        cursor.execute("SELECT COUNT(*) FROM query_patterns WHERE thumbs_up > thumbs_down")
        stats["positive_patterns"] = cursor.fetchone()[0]

        # Total errors
        cursor.execute("SELECT COUNT(*) FROM error_patterns")
        stats["total_errors"] = cursor.fetchone()[0]

        # Query types distribution
        cursor.execute(
            """
            SELECT query_type, COUNT(*) as count
            FROM query_patterns
            WHERE success = 1
            GROUP BY query_type
        """
        )
        stats["query_types"] = {row[0]: row[1] for row in cursor.fetchall()}

        # Error types distribution
        cursor.execute(
            """
            SELECT error_type, COUNT(*) as count
            FROM error_patterns
            GROUP BY error_type
        """
        )
        stats["error_types"] = {row[0]: row[1] for row in cursor.fetchall()}

        conn.close()
        return stats

    def _classify_query(self, user_query: str, sql: str) -> str:
        """Classify the type of query."""
        query_lower = user_query.lower()
        sql_upper = sql.upper()

        if "count" in query_lower or "how many" in query_lower:
            return "count"
        elif "sum" in sql_upper or "total" in query_lower or "revenue" in query_lower:
            return "aggregation"
        elif "group by" in sql_upper:
            if "where" in sql_upper:
                return "filtered_aggregation"
            return "grouped"
        elif "where" in sql_upper:
            return "filtered"
        elif "order by" in sql_upper and "limit" in sql_upper:
            return "top_n"
        elif "compare" in query_lower or "vs" in query_lower:
            return "comparison"
        else:
            return "simple"

    def _classify_error(self, error_message: str) -> str:
        """Classify the type of error."""
        error_lower = error_message.lower()

        if "no such column" in error_lower:
            return "invalid_column"
        elif "no such table" in error_lower:
            return "invalid_table"
        elif "syntax error" in error_lower:
            return "syntax"
        elif "incomplete" in error_lower:
            return "incomplete"
        elif "ambiguous" in error_lower:
            return "ambiguous"
        else:
            return "other"


# Singleton instance
_learning_store = None


def get_learning_store() -> LearningStore:
    """Get the singleton learning store instance."""
    global _learning_store
    if _learning_store is None:
        _learning_store = LearningStore()
    return _learning_store


# Test
if __name__ == "__main__":
    store = LearningStore()

    # Test storing a pattern
    pattern_id = store.store_successful_pattern(
        user_query="show me laptop sales by country",
        generated_sql="SELECT country, SUM(total_amount) FROM sales WHERE product_name LIKE '%laptop%' GROUP BY country",
        result_count=6,
        execution_time_ms=500,
        model_id="test",
        mode="test",
    )
    print(f"Stored pattern: {pattern_id}")

    # Test finding similar
    similar = store.find_similar_patterns("show laptop sales by region")
    print(f"Similar patterns: {similar}")

    # Test stats
    stats = store.get_statistics()
    print(f"Statistics: {stats}")