beeai_agent.py

"""
BeeAI Framework SQL Agent for Natural Language to SQL Conversion.

This module implements a SQL generation agent using IBM's BeeAI Framework,
which provides native WatsonX integration with clean async patterns.

Classes:
    BeeAISQLAgent: Main agent class for SQL generation and execution

Features:
    - Native WatsonX integration via BeeAI adapters
    - Embedding-based semantic similarity for few-shot learning
    - Pattern storage for continuous improvement
    - SQL execution with error handling and retry logic
    - Detailed timing breakdown for each processing step

Architecture:
    User Query
        ↓
    Learning Store Lookup (find similar patterns)
        ↓
    Prompt Building (schema + few-shot + error avoidance)
        ↓
    WatsonX LLM Generation (via BeeAI WatsonxChatModel)
        ↓
    SQL Cleanup & Validation
        ↓
    SQL Execution
        ↓
    Pattern Storage (for future learning)

Example:
    >>> from beeai_agent import BeeAISQLAgent
    >>>
    >>> agent = BeeAISQLAgent(model_id="ibm/granite-4-h-small")
    >>> result = agent.run_query("show laptop sales in USA")
    >>>
    >>> if result['success']:
    ...     print(f"Found {result['row_count']} rows")
    ...     print(f"SQL: {result['sql']}")

Dependencies:
    - beeai-framework: IBM BeeAI Framework
    - ibm-watsonx-ai: WatsonX SDK (for credentials)

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()

# Use centralized logging
from logging_config import get_logger

logger = get_logger(__name__)

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

# Learning store for pattern matching
from learning_store import get_learning_store

# Smart result analysis
from query_classifier import analyze_empty_result


class BeeAISQLAgent:
    """
    SQL Agent using BeeAI Framework with WatsonX LLM backend.

    Features:
    - Direct LLM calls with structured prompts
    - Embedding-based semantic similarity for few-shot learning
    - Pattern storage for continuous improvement
    - SQL execution and retry logic
    """

    # Available models
    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, use_learning: bool = True):
        """Initialize the BeeAI SQL 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 in environment variables")

        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)

        self._init_llm()
        self._processing_steps = []

        # Initialize learning store for pattern matching
        self.use_learning = use_learning
        if use_learning:
            self.learning_store = get_learning_store()
            logger.info(f"BeeAI learning enabled (embeddings: {self.learning_store.use_embeddings})")
        else:
            self.learning_store = None

        self._last_pattern_id = None

    def _init_llm(self):
        """Initialize the WatsonX LLM via BeeAI adapter."""
        logger.info(f"Initializing BeeAI WatsonX 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_prompt(
        self, user_query: str, similar_patterns: list[dict] = None, common_errors: list[dict] = None
    ) -> str:
        """Build the SQL generation prompt with self-improving features."""
        schema = self._get_schema()

        # Build errors section (self-improving - learn from mistakes)
        errors_section = ""
        if common_errors:
            errors_section = "\n## ⚠️ COMMON MISTAKES TO AVOID (learn from past errors):\n"
            for err in common_errors:
                errors_section += f"- {err['error_type']}: Occurred {err['count']} times\n"

        # Build few-shot examples section from similar patterns
        examples_section = ""
        if similar_patterns:
            examples_section = "\n## ✅ SIMILAR SUCCESSFUL QUERIES (use these as reference):\n"
            for i, pattern in enumerate(similar_patterns, 1):
                sim_score = pattern.get("similarity", 0)
                method = pattern.get("method", "unknown")
                thumbs = f"👍{pattern.get('thumbs_up', 0)}" if pattern.get("thumbs_up", 0) > 0 else ""
                examples_section += f"""
Example {i} (similarity: {sim_score:.2f}, {method}) {thumbs}:
Question: {pattern['user_query']}
SQL: {pattern['sql']}
"""

        return f"""You are a SQLite expert. Given a question, analyze it step-by-step and then create ONE syntactically correct SQLite query.

CRITICAL:
1. First, OUTPUT a 'Thought:' section where you plan the query logic step-by-step.
2. Then, OUTPUT a 'SQL:' section with the final query.
3. 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;
   - "show all orders" → SELECT * FROM orders;
   - 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 when using sales view - it already has all columns!
   - NO table prefixes (use product_name, NOT products.product_name)

3. COLUMN NAMES - Use exact column names:
   - products table: product_id, product_code, product_name, category, price, stock_quantity
   - customers table: customer_id, customer_name, email, city, country, region
   - orders table: order_id, customer_id, product_id, order_date, quantity, total_amount
   - sales VIEW: product_name, product_category, unit_price, customer_name, country, region, quantity, total_amount, order_date
   - NOTE: products table uses 'category', sales VIEW uses 'product_category'

4. FILTERING:
   - "laptop products" → WHERE product_name LIKE '%laptop%'
   - "from USA" → WHERE country = 'USA'
   - Only add WHERE if user specifies a filter condition!

5. AGGREGATION (only when user asks for totals/counts/averages):
   - Include GROUP BY ONLY when using SUM/COUNT/AVG with other columns
   - "total revenue" → SELECT SUM(total_amount) FROM sales;
   - "revenue by country" → SELECT country, SUM(total_amount) FROM sales GROUP BY country;

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;

7. COMPLEX LOGIC (Ratios, Comparisons):
   - "stock less than 10% of sales":
     - Thought: I need stock from 'products' and total sales from 'sales' view.
     - Thought: Join products with aggregated sales.
     - SQL:
       SELECT p.product_name, p.stock_quantity, s.total_sold
       FROM products p
       JOIN (SELECT product_name, SUM(quantity) as total_sold FROM sales GROUP BY product_name) s
       ON p.product_name = s.product_name
       WHERE p.stock_quantity < (s.total_sold * 0.1);

{examples_section}
Database Schema:
{schema}

Question: {user_query}

Thought:"""

    def _clean_sql(self, response_text: str) -> tuple[str, str | None]:
        """Extract and clean SQL from LLM response (supports CoT)."""
        # Remove "Thought:" prefix if it exists in the raw text (already filtered by regex below but good for safety)
        text = response_text.strip()

        if "SQL:" in text:
            thoughts = text.split("SQL:")[0].replace("Thought:", "").strip()
            logger.info(f"Agent Thoughts:\n{thoughts}")
        else:
            thoughts = None

        # Extract SQL part
        # 1. Look for explicit "SQL:" marker
        parts = text.split("SQL:")
        if len(parts) > 1:
            sql = parts[-1].strip()
        else:
            # 2. Key fallback: Look for "SELECT" if no explicit marker
            match = re.search(r"(SELECT\s+.+)", text, re.IGNORECASE | re.DOTALL)
            if match:
                sql = match.group(1).strip()
            else:
                sql = text

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

        # Extract first SELECT statement (up to semicolon or double newline)
        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
        # The LLM sometimes incorrectly joins orders/customers/products with the sales view
        # when it should just use sales directly
        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, thoughts

    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}")

            # Extract the SELECT and WHERE/GROUP BY/ORDER BY parts
            # 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_cleanup", "Cleanup"),
            ("sql_execution", "SQL"),
            ("pattern_store", "Store"),
            ("format_results", "Format"),
        ]
        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 a nice 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
        header = "| " + " | ".join(columns) + " |"
        separator = "|" + "|".join(["---" for _ in columns]) + "|"
        lines.append(header)
        lines.append(separator)

        # Format rows
        total = 0
        currency_col_idx = None

        for i, row in enumerate(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*")

        # Add total for sum queries
        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:
            logger.info(f"Switching model from {self.current_model_id} to {model_id}")
            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 query(self, user_query: str, max_retries: int = 2) -> dict[str, Any]:
        """Process a natural language query with learning capabilities."""
        self._processing_steps = []
        self._last_pattern_id = None
        self._last_prompt = None  # Store last prompt for debugging
        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 patterns and common errors (self-improving)
            similar_patterns = []
            common_errors = []
            if self.learning_store:
                similar_start = time.time()
                similar_patterns = self.learning_store.find_similar_patterns(user_query, limit=3)
                common_errors = self.learning_store.get_common_errors(limit=3)
                similar_time = (time.time() - similar_start) * 1000
                self._step_timings["learning_lookup"] = similar_time

                if similar_patterns or common_errors:
                    method = similar_patterns[0].get("method", "unknown") if similar_patterns else "n/a"
                    self._processing_steps.append(
                        {
                            "step": f"2. Learning Lookup ({similar_time:.0f}ms)",
                            "content": f"Found {len(similar_patterns)} patterns, {len(common_errors)} errors ({method})",
                            "time_ms": similar_time,
                        }
                    )
                else:
                    self._processing_steps.append(
                        {
                            "step": f"2. Learning Lookup ({similar_time:.0f}ms)",
                            "content": "No patterns found yet (learning...)",
                            "time_ms": similar_time,
                        }
                    )

            # Step 3: Build prompt with self-improving features
            prompt_start = time.time()
            prompt = self._build_prompt(user_query, similar_patterns, common_errors)
            self._last_prompt = prompt  # Store for debugging
            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": prompt,  # Store full prompt for debug display
                    "time_ms": prompt_time,
                }
            )

            # Step 4: Generate SQL using async (LLM call)
            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()

            # Extract SQL
            clean_start = time.time()
            sql_query, thoughts = self._clean_sql(response)
            clean_time = (time.time() - clean_start) * 1000
            self._step_timings["sql_cleanup"] = clean_time

            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

                    if self.learning_store:
                        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="beeai",
                        )
                        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,
                            }
                        )

                    # Format results
                    format_start = time.time()
                    formatted = self._format_results(result, sql_query, user_query)
                    format_time = (time.time() - format_start) * 1000
                    self._step_timings["format_results"] = format_time

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

                    return {
                        "success": True,
                        "answer": formatted,
                        "sql": sql_query,
                        "error": None,
                        "steps": self._processing_steps,
                        "execution_time_ms": total_time,
                        "step_timings": self._step_timings,
                        "pattern_id": self._last_pattern_id,
                        "similar_count": len(similar_patterns),
                        "thoughts": thoughts,
                        "prompt": self._last_prompt,  # Include prompt for debugging
                    }
                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
            if self.learning_store:
                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="beeai",
                )

            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,
            }

    def get_learning_stats(self) -> dict[str, Any]:
        """Get learning statistics."""
        if self.learning_store:
            return self.learning_store.get_statistics()
        return {"enabled": False}

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

    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
- 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 BeeAI SQL Agent...")
    agent = BeeAISQLAgent()

    print(f"Model: {agent.get_current_model()}")
    print(f"Tables: {agent.get_tables()}")

    result = agent.query("how many customers do we have?")
    print(f"\nSuccess: {result['success']}")
    print(f"SQL: {result.get('sql')}")
    print(f"Answer: {result.get('answer')}")