Blame for agent.py

"""Text-to-SQL agent: question in, safe SQL out, results back.
An NVIDIA NIM model writes the SQL as JSON, Pydantic validates the shape,
a read-only connection enforces safety, and failed queries are sent back to the model to fix."""
import os, re, json, tempfile
import duckdb
from openai import OpenAI  # NVIDIA NIM speaks the OpenAI API
from pydantic import BaseModel, Field
from dotenv import load_dotenv

load_dotenv()
DB_PATH = os.getenv("DB_PATH", "ecommerce.duckdb")
BASE_URL = os.getenv("NVIDIA_BASE_URL", "https://integrate.api.nvidia.com/v1")
MODEL = os.getenv("NVIDIA_MODEL", "nvidia/nemotron-3-super-120b-a12b")
MAX_ATTEMPTS = 3
ROW_LIMIT = 500
MAX_UPLOAD_MB = 5
DEFAULT_NOTE = "Revenue is in Indian rupees."


class SQLAnswer(BaseModel):
    sql: str = Field(description="One read-only DuckDB SELECT statement that answers the question")
    explanation: str = Field(description="One or two plain-English sentences on what the query does")


BLOCKED = re.compile(
    r"\b(insert|update|delete|drop|alter|create|attach|detach|copy|pragma|install|load|export|call|set"
    r"|read_csv|read_csv_auto|read_parquet|read_json|read_text|read_blob|glob)\b", re.I
)


def connect(db_path: str):
    """Read-only connection with file and network access switched off, so generated
    SQL can only ever read the tables inside this one database."""
    return duckdb.connect(db_path, read_only=True, config={"enable_external_access": False})


def get_schema(db_path: str) -> str:
    """Describe every table, with three sample rows so the model sees real values."""
    con = connect(db_path)
    tables = [r[0] for r in con.execute(
        "SELECT table_name FROM information_schema.tables WHERE table_schema='main' ORDER BY 1").fetchall()]
    parts = []
    for t in tables:
        cols = con.execute(f"DESCRIBE {t}").fetchall()
        sample = con.execute(f"SELECT * FROM {t} LIMIT 3").fetchall()
        parts.append(f"{t}({', '.join(f'{c[0]} {c[1]}' for c in cols)})\n  sample rows: {sample}")
    con.close()
    return "\n".join(parts)


def check_safe(sql: str) -> str:
    """Allow a single SELECT (or WITH ... SELECT) and nothing else."""
    sql = sql.strip().rstrip(";").strip()
    if ";" in sql:
        raise ValueError("Only one statement is allowed")
    if not re.match(r"^(select|with)\b", sql, re.I):
        raise ValueError("Only SELECT queries are allowed")
    if BLOCKED.search(sql):
        raise ValueError("Query contains a blocked keyword")
    return sql


def run_sql(sql: str, db_path: str):
    con = connect(db_path)
    try:
        return con.execute(f"SELECT * FROM ({sql}) LIMIT {ROW_LIMIT}").df()
    finally:
        con.close()


def parse_answer(text: str) -> SQLAnswer:
    """Pull the JSON object out of the model's reply (handles code fences and <think> blocks)."""
    text = re.sub(r"<think>.*?</think>", "", text, flags=re.S)
    match = re.search(r"\{.*\}", text, re.S)
    if not match:
        raise ValueError("the reply was not a JSON object")
    return SQLAnswer(**json.loads(match.group(0)))


def ask(question: str, db_path: str = DB_PATH):
    """Returns (SQLAnswer, DataFrame, number_of_failed_attempts)."""
    key = os.getenv("NVIDIA_API_KEY")
    if not key:
        raise RuntimeError("NVIDIA_API_KEY is missing. Add it to your .env file.")
    client = OpenAI(base_url=BASE_URL, api_key=key)
    system = (
        "You are a careful data analyst. Write DuckDB SQL for the user's question using only "
        "the tables below. Use clear column aliases. "
        + (DEFAULT_NOTE if db_path == DB_PATH else "") + "\n\n" + get_schema(db_path)
        + '\n\nReply with ONLY a JSON object and no other text, in exactly this form:\n'
        '{"sql": "<one SELECT statement>", "explanation": "<one or two plain-English sentences>"}'
    )
    messages = [{"role": "system", "content": system}, {"role": "user", "content": question}]
    last_error = None
    for attempt in range(MAX_ATTEMPTS):
        resp = client.chat.completions.create(model=MODEL, messages=messages, temperature=0, max_tokens=1000)
        text = resp.choices[0].message.content or ""
        try:
            answer = parse_answer(text)
            answer.sql = check_safe(answer.sql)
            return answer, run_sql(answer.sql, db_path), attempt
        except Exception as e:  # feed the error back so the model can repair its own reply
            last_error = e
            messages += [
                {"role": "assistant", "content": text},
                {"role": "user", "content": f"That failed: {e}. Reply again with only the corrected JSON object."},
            ]
    raise RuntimeError(f"Could not produce a working query: {last_error}")


def load_csv(csv_path: str):
    """Load an uploaded CSV into its own temporary database. Returns (db_path, summary)."""
    size_mb = os.path.getsize(csv_path) / 1_000_000
    if size_mb > MAX_UPLOAD_MB:
        raise ValueError(f"File is {size_mb:.1f} MB. The limit is {MAX_UPLOAD_MB} MB.")
    name = re.sub(r"\W+", "_", os.path.splitext(os.path.basename(csv_path))[0]).strip("_").lower() or "uploaded_data"
    if name[0].isdigit():
        name = "t_" + name
    db_path = os.path.join(tempfile.mkdtemp(), "upload.duckdb")
    con = duckdb.connect(db_path)  # normal connection, only used here to read the file
    try:
        con.execute(f'CREATE TABLE "{name}" AS SELECT * FROM read_csv_auto(?, sample_size=-1)', [csv_path])
        rows = con.execute(f'SELECT count(*) FROM "{name}"').fetchone()[0]
        cols = [c[0] for c in con.execute(f'DESCRIBE "{name}"').fetchall()]
    except Exception as e:
        raise ValueError(f"Could not read that file as a CSV: {e}")
    finally:
        con.close()
    return db_path, f"Loaded **{name}**: {rows:,} rows, columns: {', '.join(cols)}"