"""
ollama_sql_assistant.py
Local AI SQL Assistant — Ollama + SQLite
Ask questions in plain English, get answers from your database.

"""

import json
import re
import sqlite3
import requests

# --- CONFIG -------------------------------------------------------------------

OLLAMA_URL = "http://localhost:11434/api/chat"
MODEL      = "qwen2.5-coder"    # works well with SQL tool-calling

QUESTIONS = [
    "Who are my top 3 customers by total spend? Include their names and total amount spent.",
    "Which product category has the highest average order value?",
    "How many orders were placed in the last 3 months?",
]

# --- DATABASE -----------------------------------------------------------------

_conn = None

def get_conn():
    global _conn
    if _conn is None:
        _conn = sqlite3.connect(":memory:", check_same_thread=False)
        _conn.row_factory = sqlite3.Row
        _conn.executescript("""
            CREATE TABLE customers (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                name TEXT, email TEXT, city TEXT, joined_date TEXT
            );
            CREATE TABLE products (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                name TEXT, category TEXT, price REAL
            );
            CREATE TABLE orders (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                customer_id INTEGER, product_id INTEGER,
                quantity INTEGER, order_date TEXT, total_amount REAL
            );
            CREATE TABLE order_items (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                order_id INTEGER, product_id INTEGER,
                quantity INTEGER, unit_price REAL, subtotal REAL
            );

            INSERT INTO customers (name, email, city, joined_date) VALUES
                ('Sarah Johnson', 'sarah@example.com', 'New York',  '2022-01-15'),
                ('Raj Patel',     'raj@example.com',   'Chicago',   '2021-11-03'),
                ('Emma Williams', 'emma@example.com',  'San Jose',  '2023-03-22'),
                ('James Chen',    'james@example.com', 'Austin',    '2022-07-10'),
                ('Priya Sharma',  'priya@example.com', 'Seattle',   '2021-09-30');

            INSERT INTO products (name, category, price) VALUES
                ('Wireless Headphones', 'Electronics', 129.99),
                ('Running Shoes',       'Footwear',     89.99),
                ('Coffee Maker',        'Appliances',   59.99),
                ('Yoga Mat',            'Fitness',      34.99),
                ('Mechanical Keyboard', 'Electronics', 149.99);

            INSERT INTO orders (customer_id, product_id, quantity, order_date, total_amount) VALUES
                (1, 1, 2, date('now', '-10 days'),  259.98),
                (1, 5, 1, date('now', '-30 days'),  149.99),
                (2, 2, 3, date('now', '-45 days'),  269.97),
                (3, 3, 2, date('now', '-60 days'),  119.98),
                (4, 4, 4, date('now', '-75 days'),  139.96),
                (5, 1, 1, date('now', '-100 days'), 129.99),
                (2, 5, 2, date('now', '-110 days'), 299.98),
                (1, 3, 1, date('now', '-120 days'),  59.99);

            INSERT INTO order_items (order_id, product_id, quantity, unit_price, subtotal)
            SELECT o.id, o.product_id, o.quantity, p.price, o.total_amount
            FROM orders o JOIN products p ON o.product_id = p.id;
        """)
        _conn.commit()
    return _conn

# --- TOOLS --------------------------------------------------------------------

def list_tables() -> str:
    cur = get_conn().cursor()
    cur.execute("SELECT name FROM sqlite_master WHERE type='table';")
    return json.dumps({"tables": [r[0] for r in cur.fetchall()]})

def describe_table(table_name: str) -> str:
    cur = get_conn().cursor()
    cur.execute(f"PRAGMA table_info(`{table_name}`);")
    columns = [{"field": r["name"], "type": r["type"]} for r in cur.fetchall()]
    return json.dumps({"table": table_name, "columns": columns})

def execute_query(sql: str) -> str:
    if not sql.strip().upper().startswith("SELECT"):
        return json.dumps({"error": "Only SELECT queries are allowed."})
    cur = get_conn().cursor()
    try:
        cur.execute(sql)
    except Exception as e:
        return json.dumps({
            "error": str(e),
            "hint": "SQLite syntax only. Use date('now', '-90 days') not INTERVAL."
        })
    rows = [dict(r) for r in cur.fetchmany(50)]
    return json.dumps({"rows": rows, "count": len(rows)})

TOOLS = {
    "list_tables":    {"fn": list_tables,    "args": []},
    "describe_table": {"fn": describe_table, "args": ["table_name"]},
    "execute_query":  {"fn": execute_query,  "args": ["sql"]},
}

# --- SYSTEM PROMPT ------------------------------------------------------------

SYSTEM_PROMPT = """
You are a SQL assistant connected to a SQLite e-commerce database.
You do NOT know the schema. Discover it using tools before writing any SQL.

Output ONLY raw JSON when calling a tool — no prose, no backticks:
{"tool": "<name>", "args": {<key: value>}}

Tools:
- list_tables()                        — lists all tables
- describe_table({"table_name": "X"})  — shows columns and types
- execute_query({"sql": "SELECT ..."}) — runs a SELECT query

Steps for every question:
  1. list_tables()
  2. describe_table() for each table you need
  3. execute_query() with correct column names
  4. Answer in plain English — only after you have real query results

SQLite rules:
- Date math: date('now', '-90 days')  not INTERVAL syntax
- Aggregation: SUM(), AVG(), COUNT(*) with GROUP BY
- "how many orders" = COUNT(*) FROM orders
"""

# --- OLLAMA -------------------------------------------------------------------

def chat(messages: list) -> str:
    resp = requests.post(
        OLLAMA_URL,
        json={"model": MODEL, "messages": messages, "stream": False},
        timeout=120
    )
    resp.raise_for_status()
    return resp.json()["message"]["content"].strip()

def find_tool_call(text: str) -> dict | None:
    text = text.strip()
    # scan for every { and try to parse a complete JSON object
    for i, ch in enumerate(text):
        if ch == "{":
            depth = end = 0
            for j, c in enumerate(text[i:], i):
                if c == "{": depth += 1
                elif c == "}":
                    depth -= 1
                    if depth == 0:
                        end = j + 1
                        break
            if end:
                try:
                    data = json.loads(text[i:end])
                    if "tool" in data:
                        return data
                except json.JSONDecodeError:
                    pass
    return None

def dispatch(call: dict) -> str:
    name = call.get("tool")
    args = call.get("args", {})
    if name not in TOOLS:
        return json.dumps({"error": f"Unknown tool: {name}"})
    tool = TOOLS[name]
    if tool["args"]:
        missing = [k for k in tool["args"] if not args.get(k)]
        if missing:
            return json.dumps({"error": f"Missing args: {missing}", "example": '{"tool": "describe_table", "args": {"table_name": "orders"}}'})
        return tool["fn"](*[args[k] for k in tool["args"]])
    return tool["fn"]()

# --- ASK ----------------------------------------------------------------------

def ask(question: str) -> str:
    conversation = [
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user",   "content": question},
    ]
    executed = False
    for hop in range(10):
        reply = chat(conversation)
        tool_call = find_tool_call(reply)
        if tool_call:
            name    = tool_call.get("tool", "?")
            args    = tool_call.get("args", {})
            arg_str = f"({', '.join(str(v) for v in args.values())})" if args else "()"
            print(f"    >> [{hop+1}] {name}{arg_str}")
            if name == "execute_query":
                executed = True
            result   = dispatch(tool_call)
            feedback = (
                f"Tool result:\n{result}\n\n"
                + ("Now answer in plain English using this data only." if executed
                   else "Continue. Call execute_query next to get actual data.")
            )
            conversation += [
                {"role": "assistant", "content": reply},
                {"role": "user",      "content": feedback},
            ]
        else:
            if not executed:
                conversation += [
                    {"role": "assistant", "content": reply},
                    {"role": "user",      "content": "You must call execute_query before answering."},
                ]
                continue
            return reply
    return "[no answer — try a different model or rephrase the question]"

# --- MAIN ---------------------------------------------------------------------

if __name__ == "__main__":
    print(f"\n{'='*50}")
    print(f"  Local AI SQL Assistant  |  {MODEL}")
    print(f"{'='*50}")

    try:
        requests.get("http://localhost:11434", timeout=3)
    except requests.exceptions.ConnectionError:
        print("\n  Ollama is not running. Start with:  ollama serve\n")
        exit(1)

    for i, q in enumerate(QUESTIONS, 1):
        print(f"\nQ{i}: {q}")
        print("-" * 50)
        print(f"A:  {ask(q)}")

    print(f"\n{'='*50}\n")

#-------------------------Output-------------------------------------------------

python3 test_sql_assistant.py
/usr/lib/python3/dist-packages/requests/__init__.py:87: RequestsDependencyWarning: urllib3 (2.2.2) or chardet (5.2.0) doesn't match a supported version!
  warnings.warn("urllib3 ({}) or chardet ({}) doesn't match a supported "

==================================================
  Local AI SQL Assistant  |  qwen2.5-coder
==================================================

Q1: Who are my top 3 customers by total spend? Include their names and total amount spent.
--------------------------------------------------
    >> [1] list_tables()
    >> [2] describe_table(orders)
    >> [3] describe_table(customers)
    >> [4] describe_table(order_items)
    >> [5] execute_query(SELECT TOP 3 c.name, SUM(oi.subtotal) as total_spent FROM customers c JOIN orders o ON c.id = o.customer_id JOIN order_items oi ON o.id = oi.order_id GROUP BY c.id ORDER BY total_spent DESC )
A:  The top 3 customers by total spend are:

1. Raj Patel with $569.95 spent.
2. Sarah Johnson with $469.96 spent.
3. James Chen with $139.96 spent.

Q2: Which product category has the highest average order value?
--------------------------------------------------
    >> [1] list_tables()
    >> [2] describe_table(orders)
    >> [3] describe_table(products)
    >> [4] execute_query(SELECT p.category, AVG(o.total_amount) AS avg_order_value FROM orders o JOIN products p ON o.product_id = p.id GROUP BY p.category ORDER BY avg_order_value DESC )
A:  The product category with the highest average order value is Footwear.

Q3: How many orders were placed in the last 3 months?
--------------------------------------------------
    >> [1] list_tables()
    >> [2] describe_table(orders)
    >> [3] execute_query(SELECT COUNT(*) as Computed  FROM orders WHERE order_date >= date('now', '-3 months'))
A:  5 orders were placed in the last 3 months.

==================================================

Three tool calls, one plain English answer. That is the whole loop.

What's Next?

This script is intentionally minimal. Think of it as a starting point, not a finished product. A few things one can add further:

Query result caching. If the user asks the same question twice, no reason to hit the database again.

Schema caching. Call list_tables() and describe_table() once at startup and cache the result. Feed it into the system prompt so the model skips those tool hops entirely.

Streaming output. Ollama supports streaming. For longer answers it makes the experience feel much more responsive.

Wrapping Up

I have been building AI tools for a few years now, and the thing that keeps surprising me is how much you can do without reaching for a framework. This assistant is around 200 lines of Python. No LangChain. No vector store. No paid API.

The three-tool pattern (list, describe, query) is simple enough that you can reason about it, debug it, and extend it. And because everything runs locally, you can point it at a real database without worrying about what leaves your machine.

Kumar Abhishek is an Engineering Manager specializing in quality engineering and AI-powered developer tooling. Connect me on LinkedIn: https://www.linkedin.com/in/kr0abhishek/