self_improving_agent.py

"""
Self-Improving SQL Agent with Maximum Learning Capabilities.

This agent combines the BeeAI Framework with comprehensive learning features
to continuously improve SQL generation accuracy over time.

Classes:
    SelfImprovingAgent: Main agent class with full learning pipeline

Features:
    - Dynamic few-shot learning from semantically similar past queries
    - Automatic pattern storage for all successful queries
    - Error pattern tracking to avoid repeating mistakes
    - User feedback integration (thumbs up/down ratings)
    - Detailed processing step tracking with timing

Learning Flow:
    1. User submits query
    2. Agent searches learning store for similar patterns (via embeddings)
    3. Top-3 similar patterns are included as few-shot examples
    4. Common errors for similar queries are identified
    5. LLM generates SQL with enhanced context
    6. On success: pattern stored for future use
    7. On failure: error pattern stored for avoidance
    8. User feedback updates pattern quality scores

Example:
    >>> from self_improving_agent import SelfImprovingAgent
    >>>
    >>> agent = SelfImprovingAgent(model_id="ibm/granite-4-h-small")
    >>> result = agent.run_query("show laptop sales")
    >>>
    >>> # View learning stats
    >>> stats = agent.get_learning_stats()
    >>> print(f"Patterns learned: {stats['total_patterns']}")
    >>>
    >>> # Provide feedback
    >>> agent.record_feedback(result.get('pattern_id'), is_positive=True)

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

import asyncio
import os
import re
import sqlite3
import time
from typing import Any

from dotenv import load_dotenv

load_dotenv()

from learning_store import get_learning_store
from logging_config import get_logger
from query_classifier import analyze_empty_result

logger = get_logger(__name__)

# BeeAI Framework imports
from beeai_framework.adapters.watsonx import WatsonxChatModel
from beeai_framework.backend.message import UserMessage


class SelfImprovingAgent:
    """
    Self-improving SQL Agent that learns from successful queries and user feedback.

    Features:
    - Dynamic few-shot learning from similar past queries
    - Automatic pattern storage for successful queries
    - Error pattern tracking for continuous improvement
    - User feedback integration (thumbs up/down)
    """

    AVAILABLE_MODELS = {
        "ibm/granite-3-8b-instruct": {
            "name": "Granite 3 8B",
            "description": "IBM Granite 3 - Fast, good for simple queries",
        },
        "ibm/granite-3-3-8b-instruct": {
            "name": "Granite 3.3 8B",
            "description": "IBM Granite 3.3 - Latest Granite",
        },
        "ibm/granite-4-h-small": {
            "name": "Granite 4 Small",
            "description": "IBM Granite 4 - Newest generation, best accuracy",
        },
        "meta-llama/llama-3-3-70b-instruct": {
            "name": "Llama 3.3 70B",
            "description": "Meta Llama 3.3 70B - Powerful reasoning",
        },
        "mistralai/mistral-large": {
            "name": "Mistral Large",
            "description": "Mistral Large - Powerful multilingual model",
        },
    }

    def __init__(self, model_id: str = None, db_path: str = None):
        """Initialize the self-improving agent."""
        self.api_key = os.getenv("WATSONX_API_KEY")
        self.url = os.getenv("WATSONX_URL")
        self.project_id = os.getenv("WATSONX_PROJECT_ID")

        if not all([self.api_key, self.url, self.project_id]):
            raise ValueError("Missing WatsonX credentials")

        self.current_model_id = model_id or os.getenv("WATSONX_MODEL_ID", "ibm/granite-4-h-small")

        self.db_path = db_path or os.getenv("DATABASE_PATH", "./data/database.db")
        if not os.path.isabs(self.db_path):
            self.db_path = os.path.abspath(self.db_path)

        # Initialize components
        self._init_llm()
        self.learning_store = get_learning_store()

        # Track processing for debugging
        self._processing_steps = []
        self._last_pattern_id = None

    def _init_llm(self):
        """Initialize the WatsonX LLM."""
        logger.info(f"Initializing Self-Improving Agent LLM: {self.current_model_id}")
        self.llm = WatsonxChatModel(
            model_id=self.current_model_id,
            api_key=self.api_key,
            project_id=self.project_id,
            url=self.url,
        )

    def _get_schema(self) -> str:
        """Get enriched database schema with descriptions and sample values."""
        try:
            from schema_loader import get_enriched_schema_string

            return get_enriched_schema_string()
        except Exception as e:
            # Fallback to basic schema if enriched version fails
            try:
                conn = sqlite3.connect(self.db_path)
                cursor = conn.cursor()
                cursor.execute("SELECT name FROM sqlite_master WHERE type='table'")
                tables = cursor.fetchall()
                schema_info = ""
                for (table_name,) in tables:
                    if table_name.startswith("sqlite"):
                        continue
                    cursor.execute(f"PRAGMA table_info({table_name})")
                    columns = cursor.fetchall()
                    col_names = [col[1] for col in columns]
                    schema_info += f"Table {table_name}: {', '.join(col_names)}\n"
                conn.close()
                return schema_info
            except Exception as e2:
                return f"Error: {str(e2)}"

    def _build_dynamic_prompt(self, user_query: str, similar_patterns: list[dict]) -> str:
        """Build prompt with dynamic few-shot examples from learning store."""
        schema = self._get_schema()

        # Build few-shot examples section
        examples_section = ""
        if similar_patterns:
            examples_section = "\n## SUCCESSFUL EXAMPLES (learn from these):\n"
            for i, pattern in enumerate(similar_patterns, 1):
                examples_section += f"""
Example {i}:
Question: {pattern['user_query']}
SQL: {pattern['sql']}
"""

        # Get common error patterns to avoid
        common_errors = self.learning_store.get_common_errors(limit=3)
        errors_section = ""
        if common_errors:
            errors_section = "\n## COMMON MISTAKES TO AVOID:\n"
            for err in common_errors:
                errors_section += f"- {err['error_type']}: Occurred {err['count']} times\n"

        return f"""You are a SQLite expert. Given a question, create ONE syntactically correct SQLite query.

CRITICAL: Return ONLY ONE SQL query. No explanations, no examples, no other text.
{errors_section}
## RULES:
1. SIMPLE QUERIES - For "show all X" or "list X" without filters:
   - "show all products" → SELECT * FROM products;
   - "list customers" → SELECT * FROM customers;
   - DO NOT add WHERE clauses unless user specifies a filter!
   - DO NOT add GROUP BY for simple SELECT statements!

2. SALES/REVENUE QUERIES - Use the 'sales' VIEW (pre-joined data):
   - "show sales" → SELECT * FROM sales;
   - "revenue by country" → SELECT country, SUM(total_amount) FROM sales GROUP BY country;
   - "top products by sales" → SELECT product_name, SUM(total_amount) AS total FROM sales GROUP BY product_name ORDER BY total DESC LIMIT 5;
   - NO JOINs needed, NO table prefixes when using sales view.

3. COLUMN NAMES:
   - products table: product_id, product_code, product_name, category, price, stock_quantity
   - customers table: customer_id, customer_name, email, city, country, region
   - sales VIEW: product_name, product_category, unit_price, customer_name, country, region, quantity, total_amount

4. FILTERING (only when user specifies):
   - "laptop products" → WHERE product_name LIKE '%laptop%'
   - "from USA" → WHERE country = 'USA'

5. AGGREGATION:
   - "by product" means GROUP BY product_name (NOT product_category!)
   - Include GROUP BY ONLY when using SUM/COUNT/AVG with other columns

6. COUNT vs QUANTITY - IMPORTANT DISTINCTION:
   - "product count", "how many sold", "units sold" → SUM(quantity) - total units/items sold
   - "number of orders", "order count", "transactions" → COUNT(*) - number of rows/transactions
   - When user asks "how many [product]" or "product count" they mean UNITS SOLD, use SUM(quantity)
   - Example: "product count by country" → SELECT country, product_name, SUM(quantity) AS product_count FROM sales GROUP BY country, product_name;
{examples_section}
## DATABASE SCHEMA:
{schema}

Question: {user_query}

SQL:"""

    def _clean_sql(self, response_text: str) -> str:
        """Extract and clean SQL from LLM response."""
        sql = response_text.strip()

        # Remove markdown code blocks
        sql = re.sub(r"```sql\s*", "", sql)
        sql = re.sub(r"```\s*", "", sql)

        # Extract first SELECT statement
        match = re.search(r"(SELECT\s+.+?;)", sql, re.IGNORECASE | re.DOTALL)
        if match:
            sql = match.group(1)
        else:
            match = re.search(r"(SELECT\s+.+?)(?:\n\n|$)", sql, re.IGNORECASE | re.DOTALL)
            if match:
                sql = match.group(1)

        sql = sql.strip()
        if not sql.endswith(";"):
            sql += ";"

        # Fix common column name errors
        sql = re.sub(r"\bcategory\b", "product_category", sql, flags=re.IGNORECASE)

        # Clean up incorrect JOINs with sales view
        sql = self._cleanup_sales_view_joins(sql)

        # Clean for sales view - remove table prefixes
        if "FROM SALES" in sql.upper() or "FROM sales" in sql:
            sql = re.sub(r"\bcustomers\.(\w+)", r"\1", sql, flags=re.IGNORECASE)
            sql = re.sub(r"\bproducts\.(\w+)", r"\1", sql, flags=re.IGNORECASE)
            sql = re.sub(r"\borders\.(\w+)", r"\1", sql, flags=re.IGNORECASE)
            sql = re.sub(r"\bsales\.(\w+)", r"\1", sql, flags=re.IGNORECASE)

        return sql

    def _cleanup_sales_view_joins(self, sql: str) -> str:
        """
        Fix incorrect JOINs involving the sales view.

        The 'sales' view already contains all joined data from orders, products, and customers.
        If the LLM incorrectly JOINs these tables with sales, convert to use sales directly.
        """
        sql_upper = sql.upper()

        # Detect incorrect patterns: JOIN sales ON ... or FROM orders JOIN sales
        incorrect_patterns = [
            r"FROM\s+orders\s+JOIN\s+sales\s+ON",
            r"FROM\s+sales\s+JOIN\s+orders\s+ON",
            r"FROM\s+customers\s+JOIN\s+sales\s+ON",
            r"FROM\s+products\s+JOIN\s+sales\s+ON",
            r"JOIN\s+sales\s+ON\s+\w+\.order_id",
        ]

        needs_fix = any(re.search(pattern, sql_upper) for pattern in incorrect_patterns)

        if needs_fix:
            logger.info(f"Detected incorrect JOIN with sales view, fixing: {sql}")

            # Replace FROM ... JOIN ... ON ... with FROM sales
            sql = re.sub(
                r"FROM\s+\w+\s+JOIN\s+sales\s+ON\s+[^WHERE;]+",
                "FROM sales ",
                sql,
                flags=re.IGNORECASE,
            )
            sql = re.sub(
                r"FROM\s+sales\s+JOIN\s+\w+\s+ON\s+[^WHERE;]+",
                "FROM sales ",
                sql,
                flags=re.IGNORECASE,
            )

            # Clean up any remaining table prefixes
            sql = re.sub(r"\borders\.(\w+)", r"\1", sql, flags=re.IGNORECASE)
            sql = re.sub(r"\bcustomers\.(\w+)", r"\1", sql, flags=re.IGNORECASE)
            sql = re.sub(r"\bproducts\.(\w+)", r"\1", sql, flags=re.IGNORECASE)
            sql = re.sub(r"\bsales\.(\w+)", r"\1", sql, flags=re.IGNORECASE)

            logger.info(f"Fixed sales view query: {sql}")

        return sql

    def _execute_sql(self, sql: str) -> tuple[bool, Any]:
        """Execute SQL and return (success, result_or_error)."""
        try:
            conn = sqlite3.connect(self.db_path)
            cursor = conn.cursor()
            cursor.execute(sql)
            results = cursor.fetchall()
            columns = [desc[0] for desc in cursor.description]
            conn.close()
            return True, {"columns": columns, "rows": results}
        except Exception as e:
            return False, str(e)

    def _format_timing_summary(self) -> str:
        """Format a summary of step timings."""
        parts = []
        timing_order = [
            ("learning_lookup", "Learning"),
            ("prompt_build", "Prompt"),
            ("llm_generation", "LLM"),
            ("sql_execution", "SQL"),
            ("pattern_store", "Store"),
        ]
        for key, label in timing_order:
            if key in self._step_timings:
                ms = self._step_timings[key]
                parts.append(f"{label}: {ms:.0f}ms")
        return " | ".join(parts)

    def _format_results(self, data: dict, sql: str, user_query: str = "") -> str:
        """Format query results into markdown table with smart analysis for empty results."""
        rows = data["rows"]
        columns = data["columns"]

        if not rows:
            # Provide smart analysis for zero results
            if user_query:
                analysis = analyze_empty_result(user_query, sql)
                return f"No results found.\n\n{analysis}"
            return "No results found."

        lines = [f"Found **{len(rows)}** result(s):\n"]

        # Build table
        header = "| " + " | ".join(columns) + " |"
        separator = "|" + "|".join(["---" for _ in columns]) + "|"
        lines.append(header)
        lines.append(separator)

        total = 0
        currency_col_idx = None

        for row in rows[:20]:
            formatted_vals = []
            for j, val in enumerate(row):
                if isinstance(val, float):
                    if currency_col_idx is None and val > 1:
                        currency_col_idx = j
                    formatted_vals.append(f"${val:,.2f}")
                    if j == currency_col_idx:
                        total += val
                elif isinstance(val, int):
                    formatted_vals.append(f"{val:,}")
                elif val is None:
                    formatted_vals.append("-")
                else:
                    formatted_vals.append(str(val))
            lines.append("| " + " | ".join(formatted_vals) + " |")

        if len(rows) > 20:
            lines.append(f"\n*... and {len(rows) - 20} more rows*")

        if currency_col_idx is not None and "sum" in sql.lower():
            total = sum(row[currency_col_idx] for row in rows if isinstance(row[currency_col_idx], int | float))
            lines.append(f"\n**Total:** ${total:,.2f}")

        return "\n".join(lines)

    def set_model(self, model_id: str) -> tuple[bool, str]:
        """Change the LLM model."""
        if model_id == self.current_model_id:
            return True, f"Already using {model_id}"

        try:
            self.current_model_id = model_id
            self._init_llm()
            return True, f"Switched to {model_id}"
        except Exception as e:
            return False, f"Failed: {str(e)[:200]}"

    def get_current_model(self) -> str:
        return self.current_model_id

    def get_model_info(self) -> dict:
        return self.AVAILABLE_MODELS.get(
            self.current_model_id,
            {
                "name": self.current_model_id,
                "description": "Custom model",
            },
        )

    def get_tables(self) -> list:
        try:
            conn = sqlite3.connect(self.db_path)
            cursor = conn.cursor()
            cursor.execute("SELECT name FROM sqlite_master WHERE type='table'")
            tables = [row[0] for row in cursor.fetchall() if not row[0].startswith("sqlite")]
            conn.close()
            return tables
        except:
            return []

    def get_learning_stats(self) -> dict[str, Any]:
        """Get statistics about learned patterns."""
        return self.learning_store.get_statistics()

    def record_feedback(self, is_positive: bool) -> bool:
        """Record user feedback for the last query."""
        if self._last_pattern_id:
            return self.learning_store.record_feedback(self._last_pattern_id, is_positive)
        return False

    def query(self, user_query: str, max_retries: int = 2) -> dict[str, Any]:
        """Process a natural language query with self-improving capabilities."""
        self._processing_steps = []
        self._last_pattern_id = None
        self._step_timings = {}  # Track individual step timings
        start_time = time.time()

        self._processing_steps.append({"step": "1. User Query", "content": user_query, "time_ms": 0})

        try:
            # Step 2: Find similar successful patterns
            similar_start = time.time()
            similar_patterns = self.learning_store.find_similar_patterns(user_query, limit=3)
            similar_time = (time.time() - similar_start) * 1000
            self._step_timings["learning_lookup"] = similar_time

            self._processing_steps.append(
                {
                    "step": f"2. Learning Lookup ({similar_time:.0f}ms)",
                    "content": f"Found {len(similar_patterns)} similar patterns"
                    + (f": {[p['user_query'][:40] for p in similar_patterns]}" if similar_patterns else ""),
                    "time_ms": similar_time,
                }
            )

            # Step 3: Build dynamic prompt with examples
            prompt_start = time.time()
            prompt = self._build_dynamic_prompt(user_query, similar_patterns)
            prompt_time = (time.time() - prompt_start) * 1000
            self._step_timings["prompt_build"] = prompt_time

            self._processing_steps.append(
                {
                    "step": f"3. Prompt Built ({prompt_time:.0f}ms)",
                    "content": f"{len(prompt)} chars, {len(similar_patterns)} few-shot examples",
                    "time_ms": prompt_time,
                }
            )

            # Step 4: Generate SQL
            loop = asyncio.new_event_loop()
            asyncio.set_event_loop(loop)

            try:
                gen_start = time.time()
                response = loop.run_until_complete(self._generate_sql(prompt))
                gen_time = (time.time() - gen_start) * 1000
                self._step_timings["llm_generation"] = gen_time
            finally:
                loop.close()

            sql_query = self._clean_sql(response)

            self._processing_steps.append(
                {"step": f"4. LLM Generated ({gen_time:.0f}ms)", "content": sql_query, "time_ms": gen_time}
            )

            # Step 5: Execute with retries
            last_error = None
            for attempt in range(max_retries + 1):
                exec_start = time.time()
                success, result = self._execute_sql(sql_query)
                exec_time = (time.time() - exec_start) * 1000
                self._step_timings["sql_execution"] = exec_time

                if success:
                    result_count = len(result["rows"])

                    self._processing_steps.append(
                        {
                            "step": f"5. SQL Executed ({exec_time:.0f}ms)",
                            "content": f"{result_count} rows returned",
                            "time_ms": exec_time,
                        }
                    )

                    # Step 6: Store successful pattern for learning
                    total_time = (time.time() - start_time) * 1000
                    store_start = time.time()
                    self._last_pattern_id = self.learning_store.store_successful_pattern(
                        user_query=user_query,
                        generated_sql=sql_query,
                        result_count=result_count,
                        execution_time_ms=total_time,
                        model_id=self.current_model_id,
                        mode="self_improving",
                    )
                    store_time = (time.time() - store_start) * 1000
                    self._step_timings["pattern_store"] = store_time

                    self._processing_steps.append(
                        {
                            "step": f"6. Pattern Stored ({store_time:.0f}ms)",
                            "content": f"Saved as pattern #{self._last_pattern_id}",
                            "time_ms": store_time,
                        }
                    )

                    # Add timing summary
                    self._processing_steps.append(
                        {
                            "step": f"⏱️ Total: {total_time:.0f}ms",
                            "content": self._format_timing_summary(),
                            "time_ms": total_time,
                        }
                    )

                    formatted = self._format_results(result, sql_query, user_query)

                    return {
                        "success": True,
                        "answer": formatted,
                        "sql": sql_query,
                        "error": None,
                        "steps": self._processing_steps,
                        "step_timings": self._step_timings,
                        "pattern_id": self._last_pattern_id,
                        "similar_count": len(similar_patterns),
                        "execution_time_ms": total_time,
                    }
                else:
                    last_error = result
                    self._processing_steps.append(
                        {"step": f"5. Error (attempt {attempt + 1})", "content": last_error, "time_ms": exec_time}
                    )

                    if attempt < max_retries:
                        # Try to fix
                        loop = asyncio.new_event_loop()
                        asyncio.set_event_loop(loop)
                        try:
                            response = loop.run_until_complete(self._fix_sql(sql_query, last_error, user_query))
                            sql_query = self._clean_sql(response)
                        finally:
                            loop.close()

            # Store error pattern for learning
            self.learning_store.store_error_pattern(
                user_query=user_query,
                failed_sql=sql_query,
                error_message=last_error,
                model_id=self.current_model_id,
                mode="self_improving",
            )

            return {
                "success": False,
                "answer": None,
                "sql": sql_query,
                "error": f"Failed after {max_retries + 1} attempts: {last_error}",
                "steps": self._processing_steps,
                "pattern_id": None,
            }

        except Exception as e:
            return {
                "success": False,
                "answer": None,
                "sql": None,
                "error": str(e),
                "steps": self._processing_steps,
                "pattern_id": None,
            }

    async def _generate_sql(self, prompt: str) -> str:
        """Generate SQL using the LLM."""
        messages = [UserMessage(prompt)]
        response = await self.llm.run(messages)

        if response.output:
            content = response.output[0].content
            if content and len(content) > 0:
                return content[0].text
        return ""

    async def _fix_sql(self, sql: str, error: str, user_query: str) -> str:
        """Try to fix SQL based on error."""
        fix_prompt = f"""Fix this SQL query that failed:

Question: {user_query}
Failed SQL: {sql}
Error: {error}

RULES:
- For 'sales' view: NO table prefixes, NO JOINs
- Use product_name LIKE '%keyword%' for product filtering
- "by product" means GROUP BY product_name
- Include GROUP BY with aggregates

Return ONLY the fixed SQL:"""

        messages = [UserMessage(fix_prompt)]
        response = await self.llm.run(messages)

        if response.output:
            content = response.output[0].content
            if content and len(content) > 0:
                return content[0].text
        return sql


# Test
if __name__ == "__main__":
    print("Testing Self-Improving Agent...")
    agent = SelfImprovingAgent()

    print(f"Model: {agent.get_current_model()}")
    print(f"Learning Stats: {agent.get_learning_stats()}")

    # Test query
    result = agent.query("show laptop sales by country")
    print(f"\nSuccess: {result['success']}")
    print(f"SQL: {result.get('sql')}")
    print(f"Similar patterns used: {result.get('similar_count', 0)}")
    print(f"Pattern ID: {result.get('pattern_id')}")

    # Simulate feedback
    if result["success"]:
        agent.record_feedback(is_positive=True)
        print("Recorded positive feedback")

    print(f"\nUpdated Learning Stats: {agent.get_learning_stats()}")