04 — The OpenAI agent loop with SQL tool

Duration: 20 min Prerequisites: chapter 03 (you’ve seen the side-by-side).

The sequence (identical shape, OpenAI client)

sequenceDiagram
participant U as You
participant S as Streamlit
participant O as OpenAI API
participant D as SQLite

U->>S: "What is my biggest category?"
S->>O: chat.completions.create(messages, tools=[SQL_TOOL])
O->>S: choices[0].message.tool_calls=[tc1]
S->>D: safe_query(tc1.function.arguments.sql)
D->>S: DataFrame [(Housing, 30895.29)]
S->>O: chat.completions.create(messages + tool result)
O->>S: "Your biggest category is Housing with $30,895.29."
S->>U: rendered answer + SQL expander

Same two LLM round-trips as Project 1. Cost: ~$0.0003. Latency: ~2 s end-to-end.

The tool spec (recap, unchanged)

SQL_TOOL = {
    "type": "function",
    "function": {
        "name": "query_sql",
        "description": (
            "Run a read-only SELECT query on a SQLite table called `transactions`. "
            "Columns: profile (text), date (text YYYY-MM-DD), card (text), "
            "description (text, the merchant), category (text), amount (real, "
            "positive = expense, negative = refund). "
            "Always filter by `profile = 'data1'`."
        ),
        "parameters": {
            "type": "object",
            "properties": {
                "sql": {
                    "type": "string",
                    "description": "A SELECT statement. No INSERT/UPDATE/DELETE."
                }
            },
            "required": ["sql"],
        },
    },
}

This dict is the contract the model sees. OpenAI invented this JSON-schema shape (called Function Calling); Ollama copied it. So our spec is portable.

Reading the `tool_calls` from the response

response = client.chat.completions.create(
    model="gpt-4o-mini",
    messages=messages,
    tools=[SQL_TOOL],
)

msg = response.choices[0].message
# msg is a pydantic-like object with attributes:
#   .role        → "assistant"
#   .content     → None (if it decided to call a tool first)
#   .tool_calls  → list of ChatCompletionMessageToolCall objects

if msg.tool_calls:
    for tc in msg.tool_calls:
        print(tc.id)                          # "call_abc123" — KEEP THIS
        print(tc.type)                        # "function"
        print(tc.function.name)               # "query_sql"
        print(tc.function.arguments)          # '{"sql": "SELECT ..."}'  (string!)

Appending the assistant message and tool result

The shape OpenAI expects in the next call:

messages = [
    {"role": "system",    "content": SYSTEM_PROMPT},
    {"role": "user",      "content": "What is my biggest category?"},
    # The assistant message that contained the tool_call:
    {
        "role": "assistant",
        "content": None,
        "tool_calls": [
            {"id": "call_abc123",
             "type": "function",
             "function": {"name": "query_sql",
                          "arguments": '{"sql": "SELECT ..."}'}},
        ],
    },
    # Your tool execution result, MUST reference tool_call_id:
    {
        "role": "tool",
        "tool_call_id": "call_abc123",        # ← MUST match the tc.id above
        "content": "category,total\nHousing,30895.29\n",
    },
]

If you skip tool_call_id or pass the wrong one → 400 Bad Request: “messages with role ‘tool’ must be a response to a preceding message with ‘tool_calls’.”

A clean way to append the assistant message:

messages.append(msg.model_dump(exclude_none=True))

model_dump converts the pydantic object back to a dict with the right shape. exclude_none=True strips content: None if you’d rather (OpenAI accepts both).

Executing the SQL and feeding back

import json
import sys
sys.path.insert(0, "../csv-llm-shared")
import db

for tc in (msg.tool_calls or []):
    args = json.loads(tc.function.arguments)   # ← parse the JSON string
    sql = args["sql"]
    try:
        result_df = db.safe_query(conn, sql)
        tool_content = result_df.to_csv(index=False)
    except ValueError as e:
        tool_content = f"ERROR: {e}"

    messages.append({
        "role": "tool",
        "tool_call_id": tc.id,
        "content": tool_content,
    })

We always append a tool message — even on error. The model needs to see the ERROR: ... string to know what went wrong and (sometimes) recover with a different query.

The second LLM call

response = client.chat.completions.create(
    model="gpt-4o-mini",
    messages=messages,
    tools=[SQL_TOOL],   # keep tools available; the model may need a second SQL
)
final_answer = response.choices[0].message.content

In ~95% of our 4 canonical questions, this second call returns a clean content string — no further tool_calls. The agent loop is done.

A complete `ask()` function

import os, json
from dotenv import load_dotenv
from openai import OpenAI
import sys
sys.path.insert(0, "../csv-llm-shared")
import db

load_dotenv()

def ask(question: str, conn,
        model: str = "gpt-4o-mini",
        api_key: str | None = None) -> dict:
    client = OpenAI(api_key=api_key or os.getenv("OPENAI_API_KEY"))
    messages = [
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user",   "content": question},
    ]
    sql_calls = []

    response = client.chat.completions.create(
        model=model, messages=messages, tools=[SQL_TOOL]
    )
    msg = response.choices[0].message
    messages.append(msg.model_dump(exclude_none=True))

    for tc in (msg.tool_calls or []):
        sql = json.loads(tc.function.arguments)["sql"]
        try:
            result_df = db.safe_query(conn, sql)
            tool_content = result_df.to_csv(index=False)
        except ValueError as e:
            tool_content = f"ERROR: {e}"
        sql_calls.append({"sql": sql, "result": tool_content})
        messages.append({
            "role": "tool",
            "tool_call_id": tc.id,
            "content": tool_content,
        })

    if sql_calls:
        response = client.chat.completions.create(
            model=model, messages=messages, tools=[SQL_TOOL]
        )

    return {
        "answer":    response.choices[0].message.content,
        "sql_calls": sql_calls,
    }

Live test

import sys
sys.path.insert(0, "../csv-llm-shared")
import db
from llm import ask

conn = db.open_db("csv-llm-openai/.cache/csv-llm-openai.sqlite")

out = ask("What is my biggest spending category?", conn)
print(out["answer"])
for i, c in enumerate(out["sql_calls"], 1):
    print(f"\n--- SQL #{i} ---")
    print(c["sql"])
    print(c["result"])

Real output (latency ~2 s):

Your biggest spending category is Housing with a total of $30,895.29.

--- SQL #1 ---
SELECT category, SUM(amount) AS total
FROM transactions
WHERE profile = 'data1' AND amount > 0
GROUP BY category
ORDER BY total DESC
LIMIT 1
category,total
Housing,30895.29

The SQL is generated by the model. The number comes from SQLite. Identical to Project 1, 5× faster.

Common pitfalls

Pitfall	Symptom	Fix
Forgot `json.loads(tc.function.arguments)`	`TypeError: string indices must be integers`	Parse the JSON string
Forgot `tool_call_id` on the tool message	400 “must be a response to…”	Always pass `tool_call_id=tc.id`
Stale `openai` package (< 1.40)	`AttributeError: 'ChatCompletion' has no attribute 'choices[0].message.tool_calls'`	`pip install --upgrade openai`
Sent a `dict` where a `BaseModel` was expected	`Object of type ... is not JSON serializable`	Use `msg.model_dump(exclude_none=True)`
API key not loaded	`AuthenticationError`	`load_dotenv()` at module top, then `os.getenv("OPENAI_API_KEY")`

Takeaways

The loop is the same shape as Project 1 — two LLM calls with tool execution in between.
tc.function.arguments is a string — json.loads() it.
tool_call_id is mandatory on the tool message.
msg.model_dump(exclude_none=True) is the clean way to push the assistant message back into the conversation.

Next: Running the app step-by-step →