pipeline.py

"""
NL2SQL Pipeline
===============
Translates natural-language business questions into validated SQL through an LLM,
with schema-aware prompt engineering and runtime validation.

Pipeline stages:
  1. Schema Loading    — parse DDL or connect to a live database to extract schema
  2. Prompt Assembly   — build a schema-aware prompt with the user's question
  3. SQL Generation    — call the LLM to produce SQL
  4. Validation        — syntax check + schema alignment (tables/columns exist)
  5. Execution         — run the SQL and return results
  6. Verification      — confirm the query executed without errors
"""

from __future__ import annotations

import json
import logging
import os
import re
import sqlite3
from dataclasses import dataclass, field, asdict
from pathlib import Path
from typing import Optional

import requests

logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(name)s] %(message)s")
logger = logging.getLogger("nl2sql")


# ---------------------------------------------------------------------------
# Configuration
# ---------------------------------------------------------------------------

@dataclass
class PipelineConfig:
    """Central configuration."""
    ollama_base_url: str = os.getenv("OLLAMA_BASE_URL", "http://localhost:11434")
    model_name: str = os.getenv("OLLAMA_MODEL", "llama3.2")
    temperature: float = 0.0
    db_path: str = "clinic.db"
    max_retries: int = 2


# ---------------------------------------------------------------------------
# Schema Manager
# ---------------------------------------------------------------------------

@dataclass
class ColumnInfo:
    name: str
    dtype: str
    nullable: bool = True
    primary_key: bool = False
    foreign_key: Optional[str] = None

    def __str__(self):
        parts = [f"{self.name} {self.dtype}"]
        if self.primary_key:
            parts.append("PRIMARY KEY")
        if not self.nullable:
            parts.append("NOT NULL")
        if self.foreign_key:
            parts.append(f"REFERENCES {self.foreign_key}")
        return " ".join(parts)


@dataclass
class TableInfo:
    name: str
    columns: list[ColumnInfo] = field(default_factory=list)

    def column_names(self) -> list[str]:
        return [c.name for c in self.columns]

    def __str__(self):
        cols = "\n  ".join(str(c) for c in self.columns)
        return f"CREATE TABLE {self.name} (\n  {cols}\n);"


class SchemaManager:
    """Loads and manages database schema information."""

    def __init__(self):
        self.tables: dict[str, TableInfo] = {}

    def load_from_sqlite(self, db_path: str):
        """Extract schema from an existing SQLite database."""
        conn = sqlite3.connect(db_path)
        cursor = conn.cursor()

        cursor.execute("SELECT name FROM sqlite_master WHERE type='table' AND name NOT LIKE 'sqlite_%'")
        table_names = [row[0] for row in cursor.fetchall()]

        for table_name in table_names:
            cursor.execute(f"PRAGMA table_info({table_name})")
            columns = []
            for row in cursor.fetchall():
                # row: (cid, name, type, notnull, dflt_value, pk)
                columns.append(ColumnInfo(
                    name=row[1],
                    dtype=row[2] or "TEXT",
                    nullable=not bool(row[3]),
                    primary_key=bool(row[5]),
                ))

            # Get foreign keys
            cursor.execute(f"PRAGMA foreign_key_list({table_name})")
            for fk_row in cursor.fetchall():
                ref_table, from_col, to_col = fk_row[2], fk_row[3], fk_row[4]
                for col in columns:
                    if col.name == from_col:
                        col.foreign_key = f"{ref_table}({to_col})"

            self.tables[table_name] = TableInfo(name=table_name, columns=columns)

        conn.close()
        logger.info("Loaded schema: %d tables from %s", len(self.tables), db_path)

    def load_from_ddl(self, ddl_text: str):
        """Parse CREATE TABLE statements from raw DDL text."""
        # Simple regex-based DDL parser for common patterns
        table_pattern = re.compile(
            r"CREATE\s+TABLE\s+(?:IF\s+NOT\s+EXISTS\s+)?(\w+)\s*\((.*?)\);",
            re.IGNORECASE | re.DOTALL,
        )
        for match in table_pattern.finditer(ddl_text):
            table_name = match.group(1)
            body = match.group(2)
            columns = []

            for line in body.split(","):
                line = line.strip()
                if not line or line.upper().startswith(("PRIMARY KEY", "FOREIGN KEY", "CONSTRAINT", "UNIQUE", "CHECK")):
                    continue
                parts = line.split()
                if len(parts) >= 2:
                    col_name = parts[0]
                    col_type = parts[1]
                    is_pk = "PRIMARY" in line.upper() and "KEY" in line.upper()
                    is_nn = "NOT NULL" in line.upper()
                    fk = None
                    ref_match = re.search(r"REFERENCES\s+(\w+\(\w+\))", line, re.IGNORECASE)
                    if ref_match:
                        fk = ref_match.group(1)
                    columns.append(ColumnInfo(
                        name=col_name, dtype=col_type,
                        nullable=not is_nn, primary_key=is_pk, foreign_key=fk,
                    ))

            self.tables[table_name] = TableInfo(name=table_name, columns=columns)

        logger.info("Parsed %d tables from DDL", len(self.tables))

    def to_schema_text(self) -> str:
        """Produce a human-readable schema description for the LLM prompt."""
        parts = []
        for table in self.tables.values():
            parts.append(str(table))
        return "\n\n".join(parts)

    def validate_sql_references(self, sql: str) -> list[str]:
        """Check that tables and columns referenced in SQL exist in the schema."""
        errors = []
        sql_upper = sql.upper()

        # Extract referenced table names (after FROM, JOIN, UPDATE, INTO)
        table_refs = re.findall(
            r"(?:FROM|JOIN|UPDATE|INTO)\s+(\w+)", sql, re.IGNORECASE
        )
        known_tables = {t.upper() for t in self.tables}
        known_columns = set()
        for table in self.tables.values():
            for col in table.columns:
                known_columns.add(col.name.upper())

        for ref in table_refs:
            if ref.upper() not in known_tables:
                errors.append(f"Unknown table: {ref}")

        return errors


# ---------------------------------------------------------------------------
# Prompt Builder
# ---------------------------------------------------------------------------

SYSTEM_PROMPT = """You are a SQL generation assistant. Given a database schema and a natural-language question, generate a valid SQL query that answers the question.

RULES:
- Output ONLY the SQL query, no explanations, no markdown fences, no comments.
- Use only tables and columns that exist in the provided schema.
- Use standard SQL syntax compatible with SQLite.
- For date/time operations, use SQLite date functions (date(), strftime(), etc.).
- If the question is ambiguous, make reasonable assumptions and state them as SQL comments at the top.
- Always use explicit column names, never SELECT *.
"""

QUERY_PROMPT_TEMPLATE = """DATABASE SCHEMA:
{schema}

SAMPLE DATA (first 3 rows per table):
{sample_data}

QUESTION: {question}

SQL:"""


def build_prompt(schema_text: str, question: str, sample_data: str = "") -> list[dict]:
    """Assemble the full prompt for the LLM."""
    user_content = QUERY_PROMPT_TEMPLATE.format(
        schema=schema_text,
        sample_data=sample_data if sample_data else "(no sample data available)",
        question=question,
    )
    return [
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user", "content": user_content},
    ]


# ---------------------------------------------------------------------------
# LLM Client
# ---------------------------------------------------------------------------

class OllamaClient:
    """Calls Ollama for SQL generation."""

    def __init__(self, config: PipelineConfig):
        self.config = config

    def generate(self, messages: list[dict]) -> str:
        """Send messages to Ollama and return the response text."""
        # Convert chat messages to a single prompt for the /api/generate endpoint
        prompt_parts = []
        for msg in messages:
            role = msg["role"].upper()
            prompt_parts.append(f"[{role}]\n{msg['content']}")
        prompt = "\n\n".join(prompt_parts)

        payload = {
            "model": self.config.model_name,
            "prompt": prompt,
            "stream": False,
            "options": {
                "temperature": self.config.temperature,
            },
        }
        try:
            resp = requests.post(
                f"{self.config.ollama_base_url}/api/generate",
                json=payload,
                timeout=120,
            )
            resp.raise_for_status()
            return resp.json().get("response", "").strip()
        except requests.RequestException as e:
            logger.error("Ollama request failed: %s", e)
            raise RuntimeError(f"LLM call failed: {e}") from e


# ---------------------------------------------------------------------------
# SQL Validator
# ---------------------------------------------------------------------------

def clean_sql(raw: str) -> str:
    """Strip markdown fences and extra whitespace from LLM output."""
    # Remove ```sql ... ``` wrappers
    cleaned = re.sub(r"```(?:sql)?\s*", "", raw)
    cleaned = re.sub(r"```", "", cleaned)
    cleaned = cleaned.strip()
    # Remove trailing semicolons for SQLite compatibility
    if cleaned.endswith(";"):
        cleaned = cleaned[:-1]
    return cleaned


def validate_syntax(sql: str, db_path: str) -> tuple[bool, str]:
    """
    Use SQLite EXPLAIN to check syntax without executing.
    Returns (is_valid, error_message).
    """
    conn = sqlite3.connect(db_path)
    try:
        conn.execute(f"EXPLAIN {sql}")
        conn.close()
        return True, ""
    except sqlite3.Error as e:
        conn.close()
        return False, str(e)


# ---------------------------------------------------------------------------
# Query Executor
# ---------------------------------------------------------------------------

def execute_query(sql: str, db_path: str) -> dict:
    """Execute SQL and return results as a dict with columns and rows."""
    conn = sqlite3.connect(db_path)
    try:
        cursor = conn.execute(sql)
        columns = [desc[0] for desc in cursor.description] if cursor.description else []
        rows = cursor.fetchall()
        conn.close()
        return {
            "success": True,
            "columns": columns,
            "rows": [list(r) for r in rows],
            "row_count": len(rows),
        }
    except sqlite3.Error as e:
        conn.close()
        return {
            "success": False,
            "error": str(e),
            "columns": [],
            "rows": [],
            "row_count": 0,
        }


def get_sample_data(db_path: str, tables: dict) -> str:
    """Fetch first 3 rows per table for context in the prompt."""
    conn = sqlite3.connect(db_path)
    parts = []
    for table_name in tables:
        try:
            cursor = conn.execute(f"SELECT * FROM {table_name} LIMIT 3")
            cols = [d[0] for d in cursor.description]
            rows = cursor.fetchall()
            if rows:
                header = " | ".join(cols)
                row_strs = [" | ".join(str(v) for v in row) for row in rows]
                parts.append(f"{table_name}:\n{header}\n" + "\n".join(row_strs))
        except sqlite3.Error:
            pass
    conn.close()
    return "\n\n".join(parts)


# ---------------------------------------------------------------------------
# Pipeline Orchestrator
# ---------------------------------------------------------------------------

@dataclass
class QueryResult:
    """Full trace of a single query through the pipeline."""
    question: str
    generated_sql: str
    cleaned_sql: str
    validation_errors: list[str]
    syntax_valid: bool
    execution_result: dict
    retries_used: int
    category: str = ""

    @property
    def success(self) -> bool:
        return self.syntax_valid and self.execution_result.get("success", False)

    def __str__(self):
        status = "SUCCESS" if self.success else "FAILED"
        result_preview = ""
        if self.success:
            cols = self.execution_result["columns"]
            rows = self.execution_result["rows"]
            result_preview = f"\nColumns: {cols}\nRows ({self.execution_result['row_count']}):"
            for row in rows[:5]:
                result_preview += f"\n  {row}"
            if len(rows) > 5:
                result_preview += f"\n  ... ({len(rows) - 5} more rows)"
        else:
            err = self.execution_result.get("error", "unknown")
            result_preview = f"\nError: {err}"

        return (
            f"[{status}] {self.question}\n"
            f"SQL: {self.cleaned_sql}\n"
            f"Validation: {self.validation_errors or 'OK'}\n"
            f"Retries: {self.retries_used}"
            f"{result_preview}"
        )


class NL2SQLPipeline:
    """
    End-to-end pipeline: question → schema-aware prompt → LLM → SQL → validate → execute.

    Supports retry with error feedback if validation or execution fails.
    """

    def __init__(self, config: Optional[PipelineConfig] = None):
        self.config = config or PipelineConfig()
        self.schema_manager = SchemaManager()
        self.llm = OllamaClient(self.config)
        self._initialized = False

    def initialize(self, db_path: Optional[str] = None, ddl_path: Optional[str] = None):
        """Load schema from database or DDL file."""
        if db_path:
            self.config.db_path = db_path
            self.schema_manager.load_from_sqlite(db_path)
        elif ddl_path:
            with open(ddl_path) as f:
                self.schema_manager.load_from_ddl(f.read())
        else:
            raise ValueError("Provide db_path or ddl_path")
        self._initialized = True

    def query(self, question: str, category: str = "") -> QueryResult:
        """Run a natural-language question through the full pipeline."""
        if not self._initialized:
            raise RuntimeError("Call initialize() first")

        schema_text = self.schema_manager.to_schema_text()
        sample_data = get_sample_data(self.config.db_path, self.schema_manager.tables)
        retries = 0
        last_error = ""

        for attempt in range(1 + self.config.max_retries):
            # Build prompt (include previous error on retries)
            messages = build_prompt(schema_text, question, sample_data)
            if last_error and attempt > 0:
                messages.append({
                    "role": "user",
                    "content": (
                        f"The previous SQL had an error: {last_error}\n"
                        f"Please fix the query and output only the corrected SQL."
                    ),
                })

            # Generate SQL
            logger.info("Attempt %d: generating SQL for '%s'", attempt + 1, question[:60])
            raw_sql = self.llm.generate(messages)
            cleaned = clean_sql(raw_sql)

            # Validate against schema
            schema_errors = self.schema_manager.validate_sql_references(cleaned)

            # Validate syntax
            syntax_ok, syntax_err = validate_syntax(cleaned, self.config.db_path)

            if not syntax_ok:
                last_error = syntax_err
                retries = attempt
                continue

            # Execute
            exec_result = execute_query(cleaned, self.config.db_path)

            if not exec_result["success"]:
                last_error = exec_result["error"]
                retries = attempt
                continue

            # Success
            return QueryResult(
                question=question,
                generated_sql=raw_sql,
                cleaned_sql=cleaned,
                validation_errors=schema_errors,
                syntax_valid=True,
                execution_result=exec_result,
                retries_used=attempt,
                category=category,
            )

        # All attempts exhausted
        return QueryResult(
            question=question,
            generated_sql=raw_sql,
            cleaned_sql=cleaned,
            validation_errors=schema_errors,
            syntax_valid=syntax_ok,
            execution_result=exec_result if 'exec_result' in dir() else {"success": False, "error": last_error},
            retries_used=retries,
            category=category,
        )