Build a "talk to your data" AI agent

# 1) Install Ollama (free) from https://ollama.com/download, then open a new terminal:
ollama --version
ollama pull qwen2.5-coder:1.5b
ollama run qwen2.5-coder:1.5b "Write SQL to count rows in a table called orders"
# type /bye to leave the chat

# 2) Project
cd ai-data-agent-v2
python -m venv .venv
.venv\Scripts\activate            # macOS/Linux: source .venv/bin/activate
pip install -r requirements.txt

# 3) Safety net: tests use a fake model, so they pass even before the AI is involved
pytest -q

"""Load the CSV files into an in-memory DuckDB database that is locked down for AI-written SQL."""
from pathlib import Path

import duckdb
import pandas as pd

DATE_SUFFIXES = ("_date", "_time", "_timestamp")


def load_tables(data_dir: Path) -> dict[str, pd.DataFrame]:
    """Every CSV becomes one table named after the file. Date-like columns are parsed as dates."""
    tables = {}
    for path in sorted(Path(data_dir).glob("*.csv")):
        df = pd.read_csv(path)
        for col in df.columns:
            if col.endswith(DATE_SUFFIXES):
                df[col] = pd.to_datetime(df[col], errors="coerce")
        tables[path.stem] = df
    if not tables:
        raise FileNotFoundError(f"No CSV files found in {data_dir}")
    return tables


def connect(tables: dict[str, pd.DataFrame]) -> duckdb.DuckDBPyConnection:
    """Read-only style connection: no file, network or extension access, settings locked."""
    con = duckdb.connect(config={"enable_external_access": False})
    for name, df in tables.items():
        columns = ", ".join(
            f'CAST("{c}" AS TIMESTAMP) AS "{c}"' if pd.api.types.is_datetime64_any_dtype(df[c]) else f'"{c}"'
            for c in df.columns
        )
        con.register("_staging", df)                               # in-memory pandas data, not a file
        con.execute(f'CREATE TABLE "{name}" AS SELECT {columns} FROM _staging')
        con.unregister("_staging")
    con.execute("SET lock_configuration = true")                   # SQL can no longer change settings
    return con


def describe(con: duckdb.DuckDBPyConnection, max_values: int = 10) -> str:
    """Schema text for the prompt: columns, types, and the allowed values of small text columns."""
    lines = []
    for (table,) in con.execute(
        "SELECT table_name FROM information_schema.tables ORDER BY table_name"
    ).fetchall():
        rows = con.execute(f'SELECT COUNT(*) FROM "{table}"').fetchone()[0]
        columns = []
        for col, col_type, *_ in con.execute(f'DESCRIBE "{table}"').fetchall():
            text = f"{col} {col_type}"
            if col_type == "VARCHAR":
                values = [v for (v,) in con.execute(
                    f'SELECT DISTINCT "{col}" FROM "{table}" WHERE "{col}" IS NOT NULL LIMIT {max_values + 1}'
                ).fetchall()]
                if len(values) <= max_values:
                    text += f" (values: {', '.join(sorted(map(str, values)))})"
            columns.append(text)
        lines.append(f"{table} ({rows:,} rows): " + "; ".join(columns))
    return "\n".join(lines)


def run_query(con: duckdb.DuckDBPyConnection, sql: str, max_rows: int) -> pd.DataFrame:
    """Run SQL and return at most max_rows rows, keeping the query's own ORDER BY."""
    cursor = con.execute(sql)
    names = [d[0] for d in cursor.description]
    return pd.DataFrame(cursor.fetchmany(max_rows), columns=names)

# explore.py - look at the data before involving any AI
from config import DATA_DIR
from database import connect, describe, load_tables, run_query

con = connect(load_tables(DATA_DIR))
print(describe(con))
print(run_query(con, "SELECT status, COUNT(*) AS total FROM filings GROUP BY 1 ORDER BY 1", 100))

# Windows PowerShell: call the local model's REST API directly
$body = @{
  model = "qwen2.5-coder:1.5b"; stream = $false
  messages = @(@{ role = "user"; content = "SQL to count rows in table orders. Only SQL." })
} | ConvertTo-Json -Depth 4
(Invoke-RestMethod -Uri http://localhost:11434/api/chat -Method Post -Body $body -ContentType "application/json").message.content

"""Talk to a local model through Ollama's REST API. Free, offline, data never leaves the laptop."""
import httpx

from config import LLM_TIMEOUT_SECONDS, MODEL, OLLAMA_URL

SYSTEM_MESSAGE = "You are a DuckDB SQL expert. Reply with one SQL query only."


class LLMError(RuntimeError):
    """Ollama is not reachable or the model is missing."""


def ask_ollama(prompt: str, model: str = MODEL) -> str:
    try:
        response = httpx.post(
            f"{OLLAMA_URL}/api/chat",
            json={
                "model": model,
                "messages": [
                    {"role": "system", "content": SYSTEM_MESSAGE},
                    {"role": "user", "content": prompt},
                ],
                "stream": False,
                "options": {"temperature": 0, "num_predict": 512},
            },
            timeout=LLM_TIMEOUT_SECONDS,
        )
        response.raise_for_status()
    except httpx.ConnectError as err:
        raise LLMError(f"Cannot reach Ollama at {OLLAMA_URL}. Open the Ollama app or run: ollama serve") from err
    except httpx.TimeoutException as err:
        raise LLMError(f"The model took longer than {LLM_TIMEOUT_SECONDS}s. Try a smaller model.") from err
    except httpx.HTTPStatusError as err:
        if err.response.status_code == 404:
            raise LLMError(f"Model '{model}' is not installed. Run: ollama pull {model}") from err
        raise LLMError(f"Ollama error {err.response.status_code}: {err.response.text}") from err
    return response.json()["message"]["content"]

from llm import ask_ollama
print(ask_ollama("SQL to count rows in table orders"))

"""Business knowledge the model cannot guess from column names (a small 'semantic layer').

Edit this file, not the code, when the data model or business rules change.
"""

RELATIONSHIPS = [
    "customers.region_id = regions.region_id",
    "employees.region_id = regions.region_id",
    "orders.customer_id = customers.customer_id",
    "order_items.order_id = orders.order_id",
    "order_items.product_id = products.product_id",
    "payments.customer_id = customers.customer_id",
    "filings.customer_id = customers.customer_id",
    "visits.customer_id = customers.customer_id",
    "support_tickets.customer_id = customers.customer_id",
]

GLOSSARY = {
    "revenue / sales": "SUM(order_items.quantity * products.price). There is no revenue column.",
    "order value": "SUM(order_items.quantity * products.price) for one order_id.",
    "customer name": "customers.name (there is no customer_name column).",
    "visitors / visits / traffic": "rows in visits, dated by visits.visit_timestamp.",
    "payments": "payments.amount. Payments have no date column.",
}

# Worked examples (few-shot). Keep them DIFFERENT from the questions in eval/golden.json,
# otherwise the evaluation only measures copying.
EXAMPLES = [
    ("total filings per day",
     "SELECT CAST(filing_date AS DATE) AS day, COUNT(*) AS total_filings FROM filings GROUP BY 1 ORDER BY 1"),
    ("total payments by customer",
     "SELECT c.customer_id, c.name, SUM(p.amount) AS total_amount FROM customers c "
     "JOIN payments p ON c.customer_id = p.customer_id GROUP BY 1, 2 ORDER BY 3 DESC"),
    ("filings per month",
     "SELECT DATE_TRUNC('month', filing_date) AS month, COUNT(*) AS total FROM filings GROUP BY 1 ORDER BY 1"),
    ("support tickets by status",
     "SELECT status, COUNT(*) AS total FROM support_tickets GROUP BY 1 ORDER BY 2 DESC"),
    ("revenue by product",
     "SELECT p.product_name, SUM(oi.quantity * p.price) AS revenue FROM order_items oi "
     "JOIN products p ON oi.product_id = p.product_id GROUP BY 1 ORDER BY 2 DESC"),
]

"""Prompt templates. The model sees: rules, schema, relationships, glossary, examples, question."""
from knowledge import EXAMPLES, GLOSSARY, RELATIONSHIPS

RULES = """Rules:
1. Write ONE DuckDB SELECT query. Return only the SQL: no markdown, no explanation.
2. Use only the tables and columns listed in the schema. Never invent columns.
3. Use the exact text values listed in the schema (for example status = 'pending').
4. Give every calculated column a clear alias, e.g. COUNT(*) AS total.
5. Add ORDER BY for rankings and time series."""


def _context(schema: str) -> str:
    relationships = "\n".join(f"- {r}" for r in RELATIONSHIPS)
    glossary = "\n".join(f"- {term}: {meaning}" for term, meaning in GLOSSARY.items())
    return f"Schema:\n{schema}\n\nRelationships:\n{relationships}\n\nBusiness terms:\n{glossary}"


def build_prompt(schema: str, question: str) -> str:
    examples = "\n\n".join(f"Q: {q}\nSQL: {sql}" for q, sql in EXAMPLES)
    return f"{RULES}\n\n{_context(schema)}\n\nExamples:\n\n{examples}\n\nQ: {question}\nSQL:"


def build_retry_prompt(schema: str, question: str, failed_sql: str, error: str) -> str:
    return (
        f"Your previous SQL failed. Fix it.\n\n{RULES}\n\n{_context(schema)}\n\n"
        f"Question: {question}\nFailed SQL: {failed_sql}\nError: {error}\n\nCorrected SQL:"
    )

Exercise: prove that business knowledge matters
1. python cli.py  ->  ask: total revenue by region      (note the SQL and the numbers)
2. In knowledge.py, comment out the "revenue / sales" line in GLOSSARY. Restart cli.py.
3. Ask the same question again. Does the model still multiply quantity * price?
4. Put the line back.

"""Turn the model's reply into one safe, read-only SQL query - or refuse it.

Layer 1 of 2. Layer 2 is the locked DuckDB connection in database.py.
"""
import re

BLOCKED_KEYWORDS = {
    "insert", "update", "delete", "drop", "create", "alter", "truncate",
    "copy", "attach", "detach", "install", "load", "pragma", "set", "reset",
    "export", "import", "call", "checkpoint", "vacuum", "grant", "revoke",
}
BLOCKED_FUNCTIONS = {
    "read_csv", "read_csv_auto", "read_parquet", "parquet_scan", "read_json",
    "read_json_auto", "read_ndjson", "read_text", "read_blob", "glob", "sniff_csv",
}


class UnsafeSQLError(ValueError):
    """The query is not a single read-only SELECT."""


def extract_sql(reply: str) -> str:
    """Remove markdown fences and chatter such as 'Here is the query:'."""
    text = re.sub(r"```(?:sql)?", "", reply, flags=re.IGNORECASE).replace("`", '"').strip()
    match = re.search(r"\b(with|select)\b", text, flags=re.IGNORECASE)
    if match:
        text = text[match.start():]
    return text.strip().rstrip(";").strip()


def _code_only(sql: str) -> str:
    """Blank out comments and quoted text so words inside them are not mistaken for commands."""
    sql = re.sub(r"--[^\n]*", " ", sql)
    sql = re.sub(r"/\*.*?\*/", " ", sql, flags=re.DOTALL)
    sql = re.sub(r"'(?:[^']|'')*'", "''", sql)
    sql = re.sub(r'"(?:[^"]|"")*"', '""', sql)
    return sql.lower()


def validate(sql: str) -> str:
    if not sql:
        raise UnsafeSQLError("The model did not return a SQL query.")
    code = _code_only(sql)
    if ";" in code:
        raise UnsafeSQLError("Only one SQL statement is allowed.")
    first_word = re.match(r"\s*\(*\s*(\w+)", code)
    if not first_word or first_word.group(1) not in {"select", "with"}:
        raise UnsafeSQLError("Only SELECT queries are allowed.")
    words = set(re.findall(r"[a-z_]+", code))
    if words & BLOCKED_KEYWORDS:
        raise UnsafeSQLError(f"Blocked keyword: {', '.join(sorted(words & BLOCKED_KEYWORDS))}")
    called = set(re.findall(r"([a-z_]+)\s*\(", code))
    if called & BLOCKED_FUNCTIONS:
        raise UnsafeSQLError(f"File access is not allowed: {', '.join(sorted(called & BLOCKED_FUNCTIONS))}")
    return sql

"""pytest -q   (no Ollama and no DuckDB needed for this file)"""
import sys
from pathlib import Path

import pytest

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
from sql_guard import UnsafeSQLError, extract_sql, validate  # noqa: E402


@pytest.mark.parametrize("reply, expected", [
    ("```sql\nSELECT 1;\n```", "SELECT 1"),
    ("Here is the query:\nSELECT name FROM customers", "SELECT name FROM customers"),
    ("SELECT `name` FROM customers", 'SELECT "name" FROM customers'),
])
def test_extract_sql(reply, expected):
    assert extract_sql(reply) == expected


@pytest.mark.parametrize("sql", [
    "SELECT status, COUNT(*) AS total FROM filings GROUP BY 1",
    "WITH t AS (SELECT * FROM orders) SELECT COUNT(*) FROM t",
    "SELECT REPLACE(name, '_', ' ') AS clean_name FROM customers",
    "SELECT * FROM support_tickets WHERE status = 'delete me'",
    "SELECT order_id FROM orders LIMIT 10 OFFSET 5",
])
def test_allows_read_only_queries(sql):
    assert validate(sql) == sql


@pytest.mark.parametrize("sql", [
    "DELETE FROM customers",
    "DROP TABLE customers",
    "SELECT 1; DROP TABLE customers",
    "COPY customers TO 'stolen.csv'",
    "SELECT * FROM read_csv('C:/Users/me/secrets.csv')",
    "INSTALL httpfs",
    "SET enable_external_access = true",
    "ATTACH 'other.db'",
    "",
])
def test_blocks_unsafe_queries(sql):
    with pytest.raises(UnsafeSQLError):
        validate(sql)

python cli.py
Ask: delete all customers
Ask: show me the file C:/Windows/win.ini using read_csv
Ask: how many filings are pending          <- still works

"""The agent loop: question -> prompt -> model -> guard -> DuckDB -> answer, with self-correction."""
import time
from dataclasses import dataclass, field
from pathlib import Path
from typing import Callable

import duckdb
import pandas as pd

from config import DATA_DIR, MAX_RETRIES, MAX_ROWS, MODEL
from database import connect, describe, load_tables, run_query
from llm import LLMError, ask_ollama
from prompts import build_prompt, build_retry_prompt
from sql_guard import UnsafeSQLError, extract_sql, validate


@dataclass
class Attempt:
    sql: str
    error: str | None = None


@dataclass
class Answer:
    question: str
    sql: str | None = None
    data: pd.DataFrame | None = None
    error: str | None = None
    attempts: list[Attempt] = field(default_factory=list)
    seconds: float = 0.0

    @property
    def ok(self) -> bool:
        return self.data is not None


class DataAgent:
    def __init__(self, data_dir: Path = DATA_DIR, model: str = MODEL,
                 llm: Callable[[str, str], str] = ask_ollama):
        self.tables = load_tables(data_dir)
        self.con = connect(self.tables)
        self.schema = describe(self.con)
        self.model = model
        self.llm = llm                       # swap in a fake for tests

    def ask(self, question: str) -> Answer:
        answer = Answer(question=question)
        start = time.perf_counter()
        prompt = build_prompt(self.schema, question)
        try:
            for _ in range(MAX_RETRIES + 1):
                sql = extract_sql(self.llm(prompt, self.model))
                try:
                    validate(sql)
                except UnsafeSQLError as err:          # never retry a write or file-access attempt
                    answer.attempts.append(Attempt(sql, str(err)))
                    answer.error = f"Refused: {err} This assistant is read-only."
                    break
                try:
                    answer.data = run_query(self.con, sql, MAX_ROWS)
                    answer.sql = sql
                    answer.attempts.append(Attempt(sql))
                    break
                except duckdb.Error as err:            # wrong column, bad syntax: let the model fix it
                    answer.attempts.append(Attempt(sql, str(err)))
                    answer.error = str(err)
                    prompt = build_retry_prompt(self.schema, question, sql, str(err))
            if answer.ok:
                answer.error = None
        except LLMError as err:
            answer.error = str(err)
        answer.seconds = round(time.perf_counter() - start, 2)
        return answer

python cli.py

Ask: how many filings are pending
Ask: total revenue by region
Ask: orders per month
Ask: top 5 customers by total payments
Ask: schema          <- shows exactly what the model sees

streamlit run app.py
# opens http://localhost:8501

[
  {"question": "how many filings are pending",
   "reference_sql": "SELECT COUNT(*) AS pending_filings FROM filings WHERE status = 'pending'"},
  {"question": "number of open support tickets",
   "reference_sql": "SELECT COUNT(*) AS open_tickets FROM support_tickets WHERE status = 'open'"},
  {"question": "what is the total revenue",
   "reference_sql": "SELECT SUM(oi.quantity * p.price) AS revenue FROM order_items oi JOIN products p ON oi.product_id = p.product_id"},
  {"question": "total revenue by region",
   "reference_sql": "SELECT r.region_name, SUM(oi.quantity * p.price) AS revenue FROM order_items oi JOIN products p ON oi.product_id = p.product_id JOIN orders o ON oi.order_id = o.order_id JOIN customers c ON o.customer_id = c.customer_id JOIN regions r ON c.region_id = r.region_id GROUP BY 1 ORDER BY 2 DESC"},
  {"question": "average payment amount",
   "reference_sql": "SELECT AVG(amount) AS avg_payment FROM payments"},
  {"question": "orders per month",
   "reference_sql": "SELECT DATE_TRUNC('month', order_date) AS month, COUNT(*) AS orders FROM orders GROUP BY 1 ORDER BY 1"},
  {"question": "number of customers in each region",
   "reference_sql": "SELECT r.region_name, COUNT(*) AS customers FROM customers c JOIN regions r ON c.region_id = r.region_id GROUP BY 1 ORDER BY 1"},
  {"question": "employees per region",
   "reference_sql": "SELECT r.region_name, COUNT(*) AS employees FROM employees e JOIN regions r ON e.region_id = r.region_id GROUP BY 1 ORDER BY 1"},
  {"question": "which customers have more than 4 filings",
   "reference_sql": "SELECT customer_id, COUNT(*) AS filings FROM filings GROUP BY 1 HAVING COUNT(*) > 4 ORDER BY 1"},
  {"question": "average order value",
   "reference_sql": "SELECT AVG(order_value) AS avg_order_value FROM (SELECT oi.order_id, SUM(oi.quantity * p.price) AS order_value FROM order_items oi JOIN products p ON oi.product_id = p.product_id GROUP BY 1)"}
]

python evaluate.py                       # default: qwen2.5-coder:1.5b

ollama pull qwen2.5-coder:7b             # about 4.7 GB
$env:MODEL = "qwen2.5-coder:7b"          # macOS/Linux: export MODEL=qwen2.5-coder:7b
python evaluate.py