Ask a data team how they calculate their on-time delivery rate, and the answer rarely comes from the database. The columns are all there: promised dates, delivery dates, order status. The rule that turns them into one number lives somewhere else, in a wiki page, a dashboard, or a senior analyst's head.
SQL agents run into this wall all the time. Here's a real example from LiveSQLBench, a public benchmark for agents that turn questions into SQL. One of its databases tracks disaster relief, and one of its questions asks:
I need to analyze all distribution hubs based on their Resource Utilization Ratio. Please show the hub registry ID, the calculated RUR value, and their Resource Utilization Classification. Sort the results by RUR from highest to lowest.
The database has a distributionhubs table, and three of its columns are clearly involved (only the relevant columns are shown):
distributionhubs (hubregistry, hubutilpct, storecapm3, storeavailm3, ...)
But no column is called "Resource Utilization Ratio", and nothing in the schema says how to combine these three. Nothing says where "High Utilization" starts, either. Both answers live in a separate file of business rules:
RUR = (hubutilpct / 100) × (storecapm3 / (storeavailm3 + 1))
High Utilization RUR > 5
Moderate Utilization 2 ≤ RUR ≤ 5
Low Utilization RUR < 2
An agent that sees only the schema has to invent the formula and the thresholds. Whatever it invents, the query runs without an error and returns a neat table. The answer is wrong, and nothing in the output tells you it was a guess.
So the schema isn't enough. The agent also needs the business knowledge, and it needs to find the right piece without reading every rule for every question.
In this part, we build a SQL agent that looks up business knowledge before it writes a query. The knowledge lives in an OKF bundle: a folder of small markdown files, one per table and one per business rule, with an index file in each folder. The agent starts at the top index, opens only the files the question needs, and then queries the database. By the end, you'll have:
- One terminal command that answers a question about a real PostgreSQL database.
- A read-only database user, so the SQL a model writes can't change or damage your data.
- A trace for every answer: which files the agent opened, every SQL attempt and error, and the tokens and cost it took.
- A side-by-side run of the same question with the bundle and with all the raw docs pasted into the prompt.
We use free open models hosted on NVIDIA's API. Switching to Ollama, Groq, OpenAI, Gemini or Anthropic is a change in the config file.
If you followed the previous part, the bundle you compiled there is exactly what this agent reads. If you're starting here, the tutorial guide builds all the bundles with one command.
Set up the project for the agent
We add a second package, okf_agent, next to the compiler, and make the whole project installable. Installing it is what gives us the okf-agent terminal command at the end of this post.
Here's where we're heading. Each file is written, in full, in the step that needs it:
okf-sql-knowledge/
├── data/
│ └── livesqlbench-base-lite/
├── bundles/
│ └── compiled/ the OKF bundles from the compiler
├── okf_compiler/ the compiler
├── okf_agent/
│ ├── __init__.py
│ ├── config.py read the settings file
│ ├── llm.py call the model through LiteLLM
│ ├── tools.py the tools the agent can call
│ ├── agent.py the agent loop
│ ├── trace.py record every step of a run
│ ├── database.py run SQL safely
│ ├── prepare.py start the database, create the read-only user
│ └── cli.py the okf-agent command
├── runs/ one JSONL line per question asked
├── config.yaml which model to use, and other settings
├── .env API key and passwords (never committed)
├── .env.example the template for .env
├── docker-compose.yml the benchmark database
├── PROVIDERS.md how to use each model provider
├── pyproject.toml
└── README.md
Create the okf_agent/ folder with an empty __init__.py inside it. Then install the project as described in the tutorial guide. It adds three dependencies: LiteLLM, our one connection to every model provider; psycopg, which connects Python to PostgreSQL; and python-dotenv, which reads the API key and passwords from a .env file.
Connect to a model and get a test answer
An agent is a model that can ask for tools: "open this file", "run this query". Not every model can, and the ones that can don't all report what each call cost. So before we build anything, we check that our model can do both.
We use free open models hosted on NVIDIA's API. Create a free key there and add it to your .env file as NVIDIA_NIM_API_KEY, as shown in the provider guide. The same guide shows how to use Ollama, Groq, OpenAI, Gemini or Anthropic instead.
The settings file
Create config.yaml in the project root:
model:
name: nvidia_nim/nvidia/nemotron-3-super-120b-a12b
temperature: 0.0
max_tokens: 4096
callbacks: [] # optional trace UI: [langfuse] or [langsmith]
agent:
bundles: bundles/compiled
raw_docs: data/livesqlbench-base-lite
max_steps: 15
runs: runs
database:
host: localhost
port: 5432
agent_user: agent_ro
statement_timeout_s: 30
max_rows: 20
The model name has two parts: the provider (nvidia_nim) and the model's ID at that provider. To switch provider or model, you change this one line. callbacks can also send every model call to a trace UI such as Langfuse or LangSmith. It's off by default, and the provider guide shows how to turn it on. The agent and database settings are used in later steps.
Notice what's not in this file: the API key and the database passwords. config.yaml is committed to git and gets shared, so every secret lives in .env, which git ignores (see the tutorial guide).
Create okf_agent/config.py to read both:
"""Read the settings file, and the secrets from the environment."""
import os
from dataclasses import dataclass
from pathlib import Path
import yaml
from dotenv import load_dotenv
@dataclass
class ModelConfig:
name: str
temperature: float
max_tokens: int
callbacks: list[str]
@dataclass
class AgentConfig:
bundles: Path
raw_docs: Path
max_steps: int
runs: Path
@dataclass
class DatabaseConfig:
host: str
port: int
agent_user: str
statement_timeout_s: int
max_rows: int
@property
def admin_user(self) -> str:
return secret("POSTGRES_USER")
@property
def admin_password(self) -> str:
return secret("POSTGRES_PASSWORD")
@property
def agent_password(self) -> str:
return secret("AGENT_DB_PASSWORD")
@dataclass
class Config:
model: ModelConfig
agent: AgentConfig
database: DatabaseConfig
def secret(name: str) -> str:
value = os.environ.get(name)
if not value:
raise SystemExit(f"{name} is not set. Add it to your .env file (see .env.example).")
return value
def load_config(path: str | Path = "config.yaml") -> Config:
load_dotenv() # read .env into the environment
raw = yaml.safe_load(Path(path).read_text(encoding="utf-8"))
agent = raw["agent"]
return Config(
model=ModelConfig(**raw["model"]),
agent=AgentConfig(
bundles=Path(agent["bundles"]),
raw_docs=Path(agent["raw_docs"]),
max_steps=agent["max_steps"],
runs=Path(agent["runs"]),
),
database=DatabaseConfig(**raw["database"]),
)
One function for every model call
Create okf_agent/llm.py. Every call the agent makes to the model goes through chat(), so this is the one place where we measure tokens, cost and time:
"""Call the model through LiteLLM, and measure every call."""
import json
import time
from dataclasses import dataclass, field
import litellm
from okf_agent.config import ModelConfig
litellm.drop_params = True # skip settings a provider doesn't support
@dataclass
class Reply:
message: dict
tool_calls: list[dict] = field(default_factory=list)
input_tokens: int = 0
output_tokens: int = 0
reasoning_tokens: int = 0
cost_usd: float | None = None
seconds: float = 0.0
waited_seconds: float = 0.0
def chat(model: ModelConfig, messages: list[dict], tools: list[dict] | None = None,
max_retries: int = 5) -> Reply:
litellm.success_callback = model.callbacks
waited = 0.0
for attempt in range(max_retries + 1):
start = time.perf_counter()
try:
response = litellm.completion(
model=model.name,
messages=messages,
tools=tools or None,
temperature=model.temperature,
max_tokens=model.max_tokens,
)
break
except litellm.RateLimitError:
if attempt == max_retries:
raise
pause = 5 * 2**attempt
time.sleep(pause)
waited += pause
seconds = time.perf_counter() - start
message = response.choices[0].message
tool_calls = [
{"id": c.id, "name": c.function.name, "arguments": json.loads(c.function.arguments or "{}")}
for c in (message.tool_calls or [])
]
usage = response.usage
details = getattr(usage, "completion_tokens_details", None)
try:
cost = litellm.completion_cost(completion_response=response)
except Exception:
cost = None # no price known for this model
return Reply(
message=message.model_dump(exclude_none=True),
tool_calls=tool_calls,
input_tokens=usage.prompt_tokens,
output_tokens=usage.completion_tokens,
reasoning_tokens=getattr(details, "reasoning_tokens", 0) or 0,
cost_usd=cost,
seconds=round(seconds, 2),
waited_seconds=waited,
)
if __name__ == "__main__":
from okf_agent.config import load_config
config = load_config()
test_tool = {
"type": "function",
"function": {
"name": "count_rows",
"description": "Count the rows in a database table.",
"parameters": {
"type": "object",
"properties": {"table": {"type": "string"}},
"required": ["table"],
},
},
}
question = [{"role": "user", "content": "How many rows are in the distributionhubs table?"}]
reply = chat(config.model, question, tools=[test_tool])
print(f"model: {config.model.name}")
print(f"tool calls: {[(c['name'], c['arguments']) for c in reply.tool_calls] or 'none'}")
print(f"text: {reply.message.get('content') or '-'}")
print(f"tokens: {reply.input_tokens} in, {reply.output_tokens} out "
f"({reply.reasoning_tokens} reasoning)")
print(f"cost: {'unknown' if reply.cost_usd is None else f'${reply.cost_usd:.6f}'}")
print(f"time: {reply.seconds}s")
Notice three things about what chat() measures:
- Token counts come from the provider's response, not from an estimate. That's how we'll measure what each answer really costs.
- Reasoning tokens are counted separately. Some models "think" before they answer, and those tokens are billed too, so a reasoning model can use many more tokens for the same answer.
- Time spent waiting for a rate limit is kept apart from the call itself. NVIDIA's free tier limits how many requests you can send per minute. When we hit that limit, we pause and retry, and the pause shouldn't count as the model being slow.
Run the test
The block at the bottom of llm.py offers the model one made-up tool, count_rows, and asks how many rows a table has. Nothing gets executed. We only check that the model asks for the tool instead of guessing a number:
uv run -m okf_agent.llm
Output:
model: nvidia_nim/nvidia/nemotron-3-super-120b-a12b
tool calls: [('count_rows', {'table': 'distributionhubs'})]
text: -
tokens: 279 in, 91 out (62 reasoning)
cost: unknown
time: 1.6s
The model asked for count_rows with the table name, which is exactly what our agent needs: it will ask for files and queries the same way. If tool calls says none, the model can't call tools; pick another model in config.yaml.
Notice the reasoning tokens: 62 of the 91 output tokens were spent thinking before the model asked for the tool. A model that thinks first can use several times more tokens for the same answer, which is why we count them separately.
The cost shows as unknown because LiteLLM has no price for a free model. With a paid provider, the same line shows the price of the call.
Build the tools: read a bundle file, submit an answer
The model can't open a file or run a query itself. It can only ask for a tool by name, and our code runs it and sends back the result. So every tool has two parts: a description the model reads, and the Python that runs when the model asks for it.
We start with two tools. read_file reads one file from the bundle. submit_answer is how the agent hands in its final SQL and answer. The tool that runs SQL comes later, once the database is running.
Create okf_agent/tools.py:
"""The tools the agent can call: what the model is told about them, and the code behind them."""
import sys
from pathlib import Path
import yaml
READ_FILE = {
"type": "function",
"function": {
"name": "read_file",
"description": (
"Read one file from the knowledge bundle and return its text. The path is relative "
"to the bundle root, for example 'index.md' or 'tables/distributionhubs.md'."
),
"parameters": {
"type": "object",
"properties": {"path": {"type": "string", "description": "Path of the file to read."}},
"required": ["path"],
},
},
}
SUBMIT_ANSWER = {
"type": "function",
"function": {
"name": "submit_answer",
"description": (
"Submit your final answer. Call it once, at the end, with the final SQL query "
"and a short answer in plain language. If the request can't be answered with a "
"SQL query over this database, leave sql empty and explain why in answer."
),
"parameters": {
"type": "object",
"properties": {
"sql": {"type": "string", "description": "The final PostgreSQL query, or ''."},
"answer": {"type": "string", "description": "A short answer in plain language."},
},
"required": ["sql", "answer"],
},
},
}
class Toolbox:
"""The tools for one run. Without a bundle, the agent gets no read_file."""
def __init__(self, bundle: Path | None):
self.bundle = bundle.resolve() if bundle else None
self.answer: dict | None = None
def schemas(self) -> list[dict]:
return [READ_FILE, SUBMIT_ANSWER] if self.bundle else [SUBMIT_ANSWER]
def run(self, name: str, arguments: dict) -> str:
handlers = {"submit_answer": self.submit_answer}
if self.bundle:
handlers["read_file"] = self.read_file
if name not in handlers:
return f"Error: there is no tool called {name}."
try:
return handlers[name](**arguments)
except TypeError as error:
return f"Error: wrong arguments for {name}: {error}"
def read_file(self, path: str) -> str:
target = (self.bundle / path.lstrip("/")).resolve()
if not target.is_relative_to(self.bundle):
return f"Error: {path} is not a file in the bundle."
if target.is_dir():
target = target / "index.md"
if target.suffix != ".md":
return f"Error: {path} is not a file in the bundle."
if target.is_file():
return target.read_text(encoding="utf-8")
if target.name == "index.md" and target.parent.is_dir():
return self.list_folder(target.parent) # OKF index files are optional
return f"Error: {path} does not exist. Check the folder's index.md for the right name."
def list_folder(self, folder: Path) -> str:
"""Build an index from the frontmatter, for a folder that has no index.md."""
name = "the bundle root" if folder == self.bundle else folder.relative_to(self.bundle).as_posix()
lines = [f"(No index.md here. Files in {name}:)"]
for item in sorted(folder.iterdir()):
if item.is_dir():
lines.append(f"* [{item.name}/]({item.name}/index.md) - folder")
elif item.suffix == ".md":
text = item.read_text(encoding="utf-8")
meta = yaml.safe_load(text.split("---", 2)[1]) if text.startswith("---") else {}
title = (meta or {}).get("title", item.stem)
lines.append(f"* [{title}]({item.name}) - {(meta or {}).get('description', '')}")
return "\n".join(lines)
def submit_answer(self, sql: str, answer: str) -> str:
self.answer = {"sql": sql, "answer": answer}
return "Answer received."
if __name__ == "__main__":
from okf_agent.config import load_config
database, path = sys.argv[1], sys.argv[2]
toolbox = Toolbox(load_config().agent.bundles / database)
print(toolbox.run("read_file", {"path": path}))
Notice five things:
-
Each description says only what the tool does and what to pass it. How to find your way around the bundle (start at
index.md, follow the links) is not the tool's job. It goes in the agent's instructions, which we write in the next step. -
read_filecan only read markdown files inside this database's bundle. The path is resolved first, so a path like../../.envcan't climb out of the bundle folder. -
OKF makes index files optional, so
read_filecopes without them. If a folder has noindex.md, asking for it returns a list built from each file's frontmatter: the same title and one-line description an index would show. The agent still sees one folder at a time. - Errors go back to the model as text, not as Python exceptions. If the model asks for a file that doesn't exist, it reads the error and can try again, instead of the run crashing.
-
Without a bundle, there is no
read_file. Later we ask the same question without the bundle, and the sameToolboxgives the agent only the other tools.
submit_answer only stores the answer. An empty sql means the agent decided the request can't be answered with a query. Ending the run is the agent loop's job, which we write next.
Try the tool yourself
The block at the bottom of tools.py runs read_file the same way the agent will. Read the bundle's front door first:
uv run -m okf_agent.tools disaster index.md
Output:
---
okf_version: "0.2"
---
# disaster database
* [Tables](tables/index.md) - PostgreSQL tables, with column meanings, JSON fields and joins.
* [Knowledge](knowledge/index.md) - Business rules: calculations, definitions and value illustrations.
Now try to leave the bundle:
uv run -m okf_agent.tools disaster ../../.env
Output:
Error: ../../.env is not a file in the bundle.
The tool refuses, and says so in a sentence the model can read. Whatever path a model makes up, read_file only ever returns knowledge from the bundle.
Build the agent loop
An agent is a loop. We send the model the question and the tools. If it asks for a tool, we run it, add the result to the conversation, and ask the model again. We stop when it calls submit_answer, or when it has used up its steps.
There's no database yet, so the agent can't run its SQL. It reads the bundle, writes a query and hands it in. That's enough to see whether it finds the right knowledge.
Create okf_agent/agent.py:
"""The agent loop: ask the model, run the tools it asks for, repeat until it submits."""
import json
import sys
from collections.abc import Callable
from dataclasses import dataclass, field
from datetime import datetime
from pathlib import Path
from okf_agent.config import Config, load_config
from okf_agent.llm import Reply, chat
from okf_agent.tools import Toolbox
INSTRUCTIONS = """You answer questions about the `{database}` PostgreSQL database by writing one SQL query.
Business terms in the questions, such as ratios, scores and classifications, have exact definitions: formulas, thresholds and the columns they use. Never guess a definition.
Only answer questions about this database's data. If the request is anything else, such as a greeting or a general request for help, call submit_answer right away with an empty sql and an answer that says what you can do. If you find that the database doesn't hold the data a question needs, do the same and say what's missing.
{knowledge}
When you know the query, call submit_answer with the final SQL and a one-sentence answer."""
BUNDLE_KNOWLEDGE = """The database's documentation is a knowledge bundle that you read with read_file.
1. Start with index.md, then read the index of the folder you need. Links in an index are relative to its folder; links that start with '/' are relative to the bundle root.
2. Open only the files the question needs: the business rules it names, the rules those depend on, and the tables whose columns they use.
3. Write the query with the exact formulas, thresholds, table names and column names you read."""
RAW_KNOWLEDGE = """The database's documentation is below: the schema, the meaning of each column, and the business rules. Write the query with the exact formulas, thresholds, table names and column names they give.
{docs}"""
SQL_NOTE = "\n\nBefore you submit, run your query with run_sql and fix any errors."
@dataclass
class ToolCall:
name: str
arguments: dict
result: str
@dataclass
class Step:
reply: Reply
calls: list[ToolCall] = field(default_factory=list)
@dataclass
class Run:
question: str
database: str
mode: str
model: str
started_at: str
steps: list[Step] = field(default_factory=list)
answer: dict | None = None
stop_reason: str = ""
error: str = ""
def raw_docs(folder: Path, database: str) -> str:
parts = []
for suffix in ("schema.txt", "column_meaning_base.json", "kb.jsonl"):
path = folder / database / f"{database}_{suffix}"
parts.append(f"## {path.name}\n\n{path.read_text(encoding='utf-8')}")
return "\n\n".join(parts)
def system_prompt(config: Config, database: str, mode: str, toolbox: Toolbox) -> str:
if mode == "bundle":
knowledge = BUNDLE_KNOWLEDGE
else:
knowledge = RAW_KNOWLEDGE.format(docs=raw_docs(config.agent.raw_docs, database))
names = [tool["function"]["name"] for tool in toolbox.schemas()]
if "run_sql" in names:
knowledge += SQL_NOTE
return INSTRUCTIONS.format(database=database, knowledge=knowledge)
def assistant_message(reply: Reply) -> dict:
"""The model's turn, in the shape every provider accepts back."""
message = {"role": "assistant", "content": reply.message.get("content") or ""}
if reply.tool_calls:
message["tool_calls"] = [
{"id": c["id"], "type": "function",
"function": {"name": c["name"], "arguments": json.dumps(c["arguments"])}}
for c in reply.tool_calls
]
return message
def ask(config: Config, database: str, question: str, mode: str = "bundle",
toolbox: Toolbox | None = None,
on_call: Callable[[int, ToolCall | None], None] | None = None) -> Run:
folder = config.agent.bundles if mode == "bundle" else config.agent.raw_docs
if not (folder / database).is_dir():
raise ValueError(f"No {mode} docs for a database called {database!r} in {folder}")
if toolbox is None:
toolbox = Toolbox(folder / database if mode == "bundle" else None)
run = Run(question, database, mode, config.model.name,
started_at=datetime.now().astimezone().isoformat(timespec="seconds"))
messages = [
{"role": "system", "content": system_prompt(config, database, mode, toolbox)},
{"role": "user", "content": question},
]
try:
for _ in range(config.agent.max_steps):
reply = chat(config.model, messages, tools=toolbox.schemas())
step = Step(reply)
run.steps.append(step)
messages.append(assistant_message(reply))
if not reply.tool_calls:
if on_call:
on_call(len(run.steps), None)
messages.append({"role": "user", "content":
"Keep using the tools. When you are done, call submit_answer."})
continue
for call in reply.tool_calls:
result = toolbox.run(call["name"], call["arguments"])
step.calls.append(ToolCall(call["name"], call["arguments"], result))
if on_call:
on_call(len(run.steps), step.calls[-1])
messages.append({"role": "tool", "tool_call_id": call["id"], "content": result})
if toolbox.answer:
run.answer = toolbox.answer
run.stop_reason = "submitted" if toolbox.answer["sql"].strip() else "declined"
return run
run.stop_reason = "step_limit"
except Exception as error: # a failed API call, a malformed reply, ...
run.stop_reason = "error"
run.error = f"{type(error).__name__}: {error}"
return run
def show_call(number: int, call: ToolCall | None) -> None:
"""Print one step as soon as it happens."""
if call is None:
detail, name = "(answered in text, nudged to use a tool)", ""
elif call.name == "submit_answer":
detail, name = "(final answer)", call.name
else:
name = call.name
detail = call.arguments.get("path") or call.arguments.get("sql", "").split("\n")[0][:60]
if call.result.startswith("Error"):
detail += " -> error"
print(f"step {number}: {name:<14} {detail}", flush=True)
def load_task(config: Config, instance_id: str) -> dict:
path = config.agent.raw_docs / "livesqlbench_data.jsonl"
for line in path.read_text(encoding="utf-8").splitlines():
task = json.loads(line)
if task["instance_id"] == instance_id:
return task
raise SystemExit(f"No task called {instance_id} in {path}")
def question_from_args(config: Config, args: list[str]) -> tuple[str, str, str | None]:
"""A benchmark task ID (disaster_1), or a database and your own question."""
if len(args) == 1:
task = load_task(config, args[0])
return task["selected_database"], task["query"], task["instance_id"]
return args[0], args[1], None
if __name__ == "__main__":
config = load_config()
database, question, _ = question_from_args(
config, [a for a in sys.argv[1:] if not a.startswith("--")])
mode = "raw" if "--raw" in sys.argv else "bundle"
run = ask(config, database, question, mode, on_call=show_call)
print(f"\nstopped: {run.stop_reason} {run.error}".rstrip())
if run.answer:
print(f"\nanswer: {run.answer['answer']}\n\n{run.answer['sql']}".rstrip())
Notice seven things:
-
A step is one call to the model, not one tool call. A model can ask for several files in one turn, and
max_steps(15 inconfig.yaml) counts turns. That's what costs time and tokens. -
The instructions are the same in both modes. Only the knowledge part changes. In
bundlemode, the agent getsread_fileand is told how to navigate: start atindex.md, read one folder index, open only what the question needs. The agent never gets a list of all 64 files. That's progressive disclosure, and these few lines are what ask for it. Inrawmode, the three raw files are pasted into the prompt andread_fileis gone. We userawmode later to compare the two. - The agent is told what's out of scope. A greeting or a general request for help gets an answer straight away, with no files read and no SQL. So does a question about data the database doesn't hold.
- If the model answers in plain text instead of calling a tool, we nudge it and the loop goes on. The step still counts.
-
Nothing crashes the run. A failed API call ends it with
stop_reasonset toerrorand the message saved. Every run ends assubmitted,declined(an answer without SQL),step_limitorerror. A database name with no bundle stops before the first model call. -
You see each step as it happens.
ask()takes an optionalon_callfunction and calls it after every tool call. The command passesshow_call, which prints one line. Without it, the terminal stays empty until the run ends. -
Runkeeps everything: each model reply with its token counts, and each tool call with its result. The next step turns that into a trace.
The block at the bottom takes either a benchmark task ID, such as disaster_1, or a database name and your own question. We'll ask the question from the opening in the next step, once every run leaves a record.
When the request isn't a data question
First, check that the agent knows when not to start. Ask it something that isn't a question about the data:
uv run -m okf_agent.agent disaster "can you help me with an sql generation task"
Output:
step 1: read_file index.md
step 2: submit_answer (final answer)
stopped: declined
answer: I can help you generate SQL queries for the disaster database. Please provide a specific question about the data, such as requesting certain information, calculations, or classifications defined in the knowledge base.
The agent read the root index once to see what the database covers, then declined in its second step and said what it can help with. The instructions say to decline without reading anything, so that one read is a small, visible cost. Without the out-of-scope rule, nothing tells the agent to stop, and it can spend all its steps reading files in search of something to answer.
For a real question, the agent would now hand in SQL that nobody has run, so it may not even be valid. Before we give the agent a database, we make every run leave a record we can read and compare.
Record every step
The steps on screen disappear when you close the terminal. To compare runs, such as the same question with and without the bundle, or one model against another, every run has to leave a record. We save each run as one line of JSON in runs/runs.jsonl.
Create okf_agent/trace.py:
"""Turn a finished run into a record: what it read, what it ran, what it cost."""
import json
import sys
from collections import Counter
from pathlib import Path
from okf_agent.agent import Run, ask, question_from_args, show_call
from okf_agent.config import load_config
def summarize(run: Run, **extra) -> dict:
replies = [step.reply for step in run.steps]
calls = [call for step in run.steps for call in step.calls]
reads = [c for c in calls if c.name == "read_file"]
queries = [c for c in calls if c.name == "run_sql"]
costs = [r.cost_usd for r in replies if r.cost_usd is not None]
return {
**extra,
"database": run.database,
"mode": run.mode,
"model": run.model,
"question": run.question,
"started_at": run.started_at,
"stop_reason": run.stop_reason,
"error": run.error,
"llm_calls": len(replies),
"tool_calls": dict(Counter(c.name for c in calls)),
"files_opened": [c.arguments.get("path") for c in reads if not c.result.startswith("Error")],
"failed_reads": sum(c.result.startswith("Error") for c in reads),
"sql_attempts": len(queries),
"sql_errors": sum(c.result.startswith("Error") for c in queries),
"input_tokens": sum(r.input_tokens for r in replies),
"output_tokens": sum(r.output_tokens for r in replies),
"reasoning_tokens": sum(r.reasoning_tokens for r in replies),
"cost_usd": sum(costs) if costs else None,
"model_seconds": round(sum(r.seconds for r in replies), 2),
"rate_limit_wait_seconds": sum(r.waited_seconds for r in replies),
"answer": run.answer,
}
def save(record: dict, folder: Path) -> Path:
folder.mkdir(parents=True, exist_ok=True)
path = folder / "runs.jsonl"
with path.open("a", encoding="utf-8") as file:
file.write(json.dumps(record, ensure_ascii=False) + "\n")
return path
def report(record: dict) -> str:
tools = ", ".join(f"{name} {count}" for name, count in record["tool_calls"].items())
cost = "unknown" if record["cost_usd"] is None else f"${record['cost_usd']:.4f}"
return "\n".join([
f"stopped: {record['stop_reason']} {record['error']}".rstrip(),
f"model calls: {record['llm_calls']}",
f"tool calls: {tools or 'none'}",
f"files opened: {len(record['files_opened'])} ({record['failed_reads']} failed reads)",
f"sql attempts: {record['sql_attempts']} ({record['sql_errors']} errors)",
f"tokens: {record['input_tokens']:,} in, {record['output_tokens']:,} out "
f"({record['reasoning_tokens']:,} reasoning)",
f"cost: {cost}",
f"model time: {record['model_seconds']}s "
f"(+{record['rate_limit_wait_seconds']}s waiting for rate limits)",
])
if __name__ == "__main__":
config = load_config()
database, question, task_id = question_from_args(
config, [a for a in sys.argv[1:] if not a.startswith("--")])
mode = "raw" if "--raw" in sys.argv else "bundle"
run = ask(config, database, question, mode, on_call=show_call)
record = summarize(run, task_id=task_id)
path = save(record, config.agent.runs)
print(f"\n{report(record)}\n\nrecord saved to {path}")
Notice five things:
- Tokens are added up over every model call, and input tokens grow with each step. Every call sends the whole conversation again: the instructions, the tool descriptions and everything the agent has read so far. An extra step doesn't just cost its own reply; it resends all of that.
-
files_openedlists only the reads that worked. Failed reads, like a wrong guess at a file name, are counted separately, so a run's wrong turns stay visible. -
sql_attemptsandsql_errorsare zero for now. The fields are already there, so every record has the same shape once the agent can run SQL. - Model time leaves out time spent waiting for rate limits. It still depends on the provider's load, so compare it only between runs on the same model and provider.
-
The record is plain JSON Lines. Each run is appended as one line, and any tool (Python, pandas,
jq) can read the file later. If you also want a trace UI, the provider guide shows how to send every model call to Langfuse or LangSmith as well.
The block at the bottom runs a question with live steps, saves the record and prints a summary. Ask the question from the opening of this post, by its ID:
uv run -m okf_agent.trace disaster_1
Output:
step 1: read_file index.md
step 2: read_file knowledge/index.md
step 3: read_file knowledge/calculation/resource-utilization-ratio.md -> error
step 4: read_file knowledge/calculation/index.md -> error
step 5: read_file knowledge/index.md
step 6: read_file knowledge/calculation/resource-utilization-ratio.md -> error
step 7: read_file tables/index.md
step 8: read_file knowledge/resource-utilization-ratio.md
step 9: read_file knowledge/resource-utilization-classification.md
step 10: read_file tables/distributionhubs.md
step 11: submit_answer (final answer)
stopped: submitted
model calls: 11
tool calls: read_file 10, submit_answer 1
files opened: 7 (3 failed reads)
sql attempts: 0 (0 errors)
tokens: 39,280 in, 1,775 out (1,197 reasoning)
cost: unknown
model time: 15.1s (+0.0s waiting for rate limits)
record saved to runs\runs.jsonl
This run took 11 model calls, and three of its reads failed. All three looked for a knowledge/calculation/ folder that doesn't exist: the agent took the # Calculation heading in the knowledge index for a folder, and went back to the same guess after the first error. The record shows exactly what that wrong turn cost.
Compare the input tokens with the reading cost we estimated for this question in the previous part: about 3,600 tokens for the files on the shortest path. This run sent 39,280 input tokens, about eleven times more. The files are only part of it. Every call resends the instructions, the tool descriptions and everything read so far, and each wrong turn adds one more call that resends it all again.
Every run now leaves a record, but the SQL in it has never been run. Next, we give the agent a database, safely.
Start the database with a read-only login for the agent
Until now, the agent's SQL was only a proposal. To run it, the agent needs to log in to the database. The benchmark's Docker setup has a single login, root, and it's a superuser: it can drop tables, read files on the database server and even run shell commands there. SQL written by a model must never run as that user.
So before the agent touches the database, we create a second user, agent_ro, that can only read.
You need Docker installed and running; see the tutorial guide for the setup. Then create docker-compose.yml in the project root:
services:
postgresql:
image: docker.io/shawnxxh/bird-interact-postgresql@sha256:1ae45d7aa5d64dd8eb82e4058f56b4b9625d5035b9b9dc0d2afa0295d9d3053c
container_name: livesqlbench_postgresql
environment:
POSTGRES_USER: ${POSTGRES_USER}
POSTGRES_PASSWORD: ${POSTGRES_PASSWORD}
volumes:
- pgdata:/var/lib/postgresql/data
ports:
- "5432:5432"
volumes:
pgdata:
The image is pinned by its digest, the exact version this post was tested with. The login comes from your .env file, so no password is written in a committed file. POSTGRES_USER must be root, because the image's loading script logs in under that name.
Create okf_agent/prepare.py:
"""Start the benchmark database, and give the agent a login that can only read."""
import subprocess
import time
import psycopg
from psycopg import sql
from okf_agent.config import Config, load_config
CHECKS = [
("read a table", "SELECT count(*) FROM distributionhubs"),
("drop a table", "DROP TABLE distributionhubs"),
("create a table", "CREATE TABLE notes (note text)"),
("read a server file", "SELECT pg_read_file('/etc/passwd')"),
]
def database_names(config: Config) -> list[str]:
"""One database per benchmark folder, named like the folder."""
folder = config.agent.raw_docs
return sorted(p.name for p in folder.iterdir() if (p / f"{p.name}_schema.txt").exists())
def connect(config: Config, dbname: str, admin: bool = False) -> psycopg.Connection:
db = config.database
user, password = ((db.admin_user, db.admin_password) if admin
else (db.agent_user, db.agent_password))
return psycopg.connect(host=db.host, port=db.port, dbname=dbname, user=user,
password=password, connect_timeout=5)
def start_database(config: Config, names: list[str], wait_minutes: int = 30) -> None:
try:
subprocess.run(["docker", "compose", "up", "-d"], check=True)
except (FileNotFoundError, subprocess.CalledProcessError):
raise SystemExit("Couldn't start the database with Docker. Is Docker installed "
"and running? See TUTORIAL.md.")
print(f"waiting until all {len(names)} databases are loaded...", flush=True)
deadline = time.monotonic() + wait_minutes * 60
while time.monotonic() < deadline:
try:
with connect(config, "postgres", admin=True) as conn:
found = {row[0] for row in conn.execute("SELECT datname FROM pg_database")}
if set(names) <= found:
return
except psycopg.OperationalError as error:
if "password authentication failed" in str(error):
raise SystemExit("The database rejected POSTGRES_PASSWORD from .env. "
"See 'Change the database password' in TUTORIAL.md.")
time.sleep(5)
raise SystemExit("The databases didn't finish loading. Check: docker compose logs")
def create_agent_user(config: Config, names: list[str]) -> None:
user = sql.Identifier(config.database.agent_user)
with connect(config, "postgres", admin=True) as conn:
conn.autocommit = True
exists = conn.execute("SELECT 1 FROM pg_roles WHERE rolname = %s",
[config.database.agent_user]).fetchone()
conn.execute(sql.SQL("{} ROLE {} LOGIN PASSWORD {}").format(
sql.SQL("ALTER" if exists else "CREATE"), user,
sql.Literal(config.database.agent_password)))
for name in names:
with connect(config, name, admin=True) as conn:
conn.autocommit = True
database = sql.Identifier(name)
conn.execute("REVOKE CREATE ON SCHEMA public FROM PUBLIC")
conn.execute(sql.SQL("REVOKE TEMPORARY ON DATABASE {} FROM PUBLIC").format(database))
conn.execute(sql.SQL("GRANT CONNECT ON DATABASE {} TO {}").format(database, user))
conn.execute(sql.SQL("GRANT USAGE ON SCHEMA public TO {}").format(user))
conn.execute(sql.SQL("GRANT SELECT ON ALL TABLES IN SCHEMA public TO {}").format(user))
def check_agent_user(config: Config, database: str = "disaster") -> None:
"""Try a few things as the agent's user. Every attempt is rolled back."""
for label, query in CHECKS:
conn = connect(config, database)
try:
cursor = conn.execute(query)
result = cursor.fetchone()[0] if cursor.description else "done"
print(f" allowed {label}: {result}")
except psycopg.Error as error:
print(f" blocked {label}: {str(error).strip().splitlines()[0]}")
finally:
conn.rollback()
conn.close()
if __name__ == "__main__":
config = load_config()
names = database_names(config)
start_database(config, names)
create_agent_user(config, names)
print(f"user {config.database.agent_user} can now read {len(names)} databases\n")
print("what the agent's user can do in disaster:")
check_agent_user(config)
Notice four things:
-
start_databasestarts the container and waits until every database exists. On its first start, the image loads all 18 benchmark databases. For each one it keeps a template copy, such asdisaster_template, and creates a plain database,disaster, which is the one the agent reads. -
create_agent_usergivesagent_rothree rights per database: to connect, to see thepublicschema, and toSELECTfrom its tables. Running it again only updates the password, so it's safe to repeat. -
The two
REVOKElines close gaps that PostgreSQL 14 leaves open. By default, every user may create tables in thepublicschema and create temporary tables. We take both away, so reading is all that's left. -
check_agent_usertries a few things asagent_ro, and rolls back every attempt. Even if a right were set up wrongly, nothing would change in the database.
Run it:
uv run -m okf_agent.prepare
Output:
[+] Running 1/1
✔ Container livesqlbench_postgresql Running
waiting until all 18 databases are loaded...
user agent_ro can now read 18 databases
what the agent's user can do in disaster:
allowed read a table: 999
blocked drop a table: must be owner of table distributionhubs
blocked create a table: permission denied for schema public
blocked read a server file: permission denied for function pg_read_file
The agent's user can read the table, and that's all. Dropping it fails because agent_ro doesn't own it. Creating a table fails because it may not write to the schema. Reading a file on the server fails because only superusers may do that.
These refusals come from PostgreSQL itself, not from our code. Whatever SQL a model writes, the database checks it against this user's rights before running it. Next, we give the agent a tool to run its queries as this user.
Run SQL safely
The database already protects itself: the agent's user can only read. That's the first layer. Before the agent gets a tool to run queries, we add three more in our own code, so that a mistake in one layer is caught by another.
Create okf_agent/database.py:
"""Run one SQL query as the read-only user, with limits, and return text the model can read."""
import sys
from decimal import Decimal
import psycopg
from okf_agent.config import DatabaseConfig, load_config
def cell(value) -> str:
if isinstance(value, (Decimal, float)):
return f"{float(value):.6g}" # short numbers keep the result small
text = "NULL" if value is None else str(value)
return text if len(text) <= 100 else text[:97] + "..."
def run_query(db: DatabaseConfig, database: str, query: str) -> str:
try:
with psycopg.connect(host=db.host, port=db.port, dbname=database, user=db.agent_user,
password=db.agent_password, connect_timeout=5) as conn:
conn.read_only = True # every transaction starts as READ ONLY
conn.execute(f"SET statement_timeout = {db.statement_timeout_s * 1000}")
cursor = conn.execute(query)
if cursor.description is None:
conn.rollback()
return "The statement ran, but returned no rows."
columns = [column.name for column in cursor.description]
total = cursor.rowcount
rows = cursor.fetchmany(db.max_rows)
conn.rollback() # nothing is ever kept
except psycopg.errors.QueryCanceled:
return f"Error: the query ran longer than {db.statement_timeout_s} seconds and was stopped."
except psycopg.Error as error:
hint = getattr(error.diag, "message_hint", None)
return f"Error: {str(error).strip()}" + (f"\nHint: {hint}" if hint else "")
lines = [" | ".join(columns)]
lines += [" | ".join(cell(value) for value in row) for row in rows]
shown = f"{len(rows)} of {total} rows" if total > len(rows) else f"{total} row{'s' * (total != 1)}"
return "\n".join(lines) + f"\n({shown})"
if __name__ == "__main__":
print(run_query(load_config().database, sys.argv[1], sys.argv[2]))
Notice four things:
- Every query runs in a read-only transaction, and is then rolled back. If the agent's user ever had more rights than it should, a statement that changes data would still fail, and nothing it did would be kept.
-
A query that runs longer than 30 seconds is stopped. A model can easily write a query that joins every row to every other row. The limit comes from
statement_timeout_sinconfig.yaml. - The model sees at most 20 rows, plus the total row count. That's enough to check that a query does what it should, without filling the model's context, and paying for it in tokens, with thousands of rows. Numbers are shortened to six significant digits for the same reason.
- Errors come back as text, with PostgreSQL's position marker and hint. The model reads them and fixes its query, just as it does with a wrong file name.
Run a query as the agent's user:
uv run -m okf_agent.database disaster "SELECT hubregistry, (hubutilpct / 100.0) * (storecapm3 / (storeavailm3 + 1)) AS rur FROM distributionhubs ORDER BY rur DESC"
Output:
hubregistry | rur
HUB_XVHV | 460.207
HUB_J1JI | 457.992
HUB_V8QF | 349.307
HUB_XT44 | 200.055
HUB_NHLC | 187.676
...
(20 of 999 rows)
The query ran as agent_ro, and the model would see 20 rows plus the total: 999 hubs. (The output is shortened here.)
Now try to change data:
uv run -m okf_agent.database disaster "DELETE FROM distributionhubs"
Output:
Error: cannot execute DELETE in a read-only transaction
The read-only transaction stops the DELETE before the database even checks the user's rights.
Give the agent the tool
The agent gets a third tool, run_sql. Replace okf_agent/tools.py:
"""The tools the agent can call: what the model is told about them, and the code behind them."""
import sys
from collections.abc import Callable
from pathlib import Path
import yaml
READ_FILE = {
"type": "function",
"function": {
"name": "read_file",
"description": (
"Read one file from the knowledge bundle and return its text. The path is relative "
"to the bundle root, for example 'index.md' or 'tables/distributionhubs.md'."
),
"parameters": {
"type": "object",
"properties": {"path": {"type": "string", "description": "Path of the file to read."}},
"required": ["path"],
},
},
}
RUN_SQL = {
"type": "function",
"function": {
"name": "run_sql",
"description": (
"Run one read-only PostgreSQL query and return the result: the column names, "
"at most 20 rows, and the total number of rows."
),
"parameters": {
"type": "object",
"properties": {"sql": {"type": "string", "description": "The query to run."}},
"required": ["sql"],
},
},
}
SUBMIT_ANSWER = {
"type": "function",
"function": {
"name": "submit_answer",
"description": (
"Submit your final answer. Call it once, at the end, with the final SQL query "
"and a short answer in plain language. If the request can't be answered with a "
"SQL query over this database, leave sql empty and explain why in answer."
),
"parameters": {
"type": "object",
"properties": {
"sql": {"type": "string", "description": "The final PostgreSQL query, or ''."},
"answer": {"type": "string", "description": "A short answer in plain language."},
},
"required": ["sql", "answer"],
},
},
}
class Toolbox:
"""The tools for one run. Without a bundle there is no read_file; without a database, no run_sql."""
def __init__(self, bundle: Path | None, run_sql: Callable[[str], str] | None = None):
self.bundle = bundle.resolve() if bundle else None
self.run_sql = run_sql
self.answer: dict | None = None
def schemas(self) -> list[dict]:
tools = [READ_FILE] if self.bundle else []
tools += [RUN_SQL] if self.run_sql else []
return tools + [SUBMIT_ANSWER]
def run(self, name: str, arguments: dict) -> str:
handlers = {"submit_answer": self.submit_answer}
if self.bundle:
handlers["read_file"] = self.read_file
if self.run_sql:
handlers["run_sql"] = self.run_sql
if name not in handlers:
return f"Error: there is no tool called {name}."
try:
return handlers[name](**arguments)
except TypeError as error:
return f"Error: wrong arguments for {name}: {error}"
def read_file(self, path: str) -> str:
target = (self.bundle / path.lstrip("/")).resolve()
if not target.is_relative_to(self.bundle):
return f"Error: {path} is not a file in the bundle."
if target.is_dir():
target = target / "index.md"
if target.suffix != ".md":
return f"Error: {path} is not a file in the bundle."
if target.is_file():
return target.read_text(encoding="utf-8")
if target.name == "index.md" and target.parent.is_dir():
return self.list_folder(target.parent) # OKF index files are optional
return f"Error: {path} does not exist. Check the folder's index.md for the right name."
def list_folder(self, folder: Path) -> str:
"""Build an index from the frontmatter, for a folder that has no index.md."""
name = "the bundle root" if folder == self.bundle else folder.relative_to(self.bundle).as_posix()
lines = [f"(No index.md here. Files in {name}:)"]
for item in sorted(folder.iterdir()):
if item.is_dir():
lines.append(f"* [{item.name}/]({item.name}/index.md) - folder")
elif item.suffix == ".md":
text = item.read_text(encoding="utf-8")
meta = yaml.safe_load(text.split("---", 2)[1]) if text.startswith("---") else {}
title = (meta or {}).get("title", item.stem)
lines.append(f"* [{title}]({item.name}) - {(meta or {}).get('description', '')}")
return "\n".join(lines)
def submit_answer(self, sql: str, answer: str) -> str:
self.answer = {"sql": sql, "answer": answer}
return "Answer received."
if __name__ == "__main__":
from okf_agent.config import load_config
database, path = sys.argv[1], sys.argv[2]
toolbox = Toolbox(load_config().agent.bundles / database)
print(toolbox.run("read_file", {"path": path}))
What changed compared with the previous version:
-
A new tool,
run_sql. Its description tells the model what it gets back, including the 20-row limit, so the model doesn't mistake 20 rows for the whole result. -
Toolboxtakes a function that runs SQL. The toolbox doesn't know about databases or passwords. It only calls the function it's given, which keeps each module to one job. -
Both modes get
run_sql. With the bundle or with the raw docs, the agent can test its query. The only difference between the two remains how the agent gets its knowledge.
Then make three small changes in okf_agent/agent.py. Import the new module at the top:
from okf_agent.config import Config, load_config
from okf_agent.database import run_query
In ask(), give the toolbox a function that runs SQL on this question's database:
if toolbox is None:
toolbox = Toolbox(folder / database if mode == "bundle" else None,
run_sql=lambda sql: run_query(config.database, database, sql))
And in show_call(), print a query on one line:
detail = call.arguments.get("path") or " ".join(call.arguments.get("sql", "").split())[:70]
Nothing else changes. The instructions already tell the agent to run its query before it submits, as soon as run_sql is among its tools.
Ask the question again, now with a database:
uv run -m okf_agent.trace disaster_1
Output:
step 1: read_file index.md
step 2: read_file knowledge/index.md
step 3: read_file knowledge/resource-utilization-ratio.md
step 4: read_file knowledge/resource-utilization-classification.md
step 5: read_file tables/distributionhubs.md
step 6: run_sql WITH rur_calc AS ( SELECT hubregistry, (hubutilpct / 100.0) * (storeca
step 7: submit_answer (final answer)
stopped: submitted
model calls: 7
tool calls: read_file 5, run_sql 1, submit_answer 1
files opened: 5 (0 failed reads)
sql attempts: 1 (0 errors)
tokens: 18,809 in, 1,268 out (710 reasoning)
cost: unknown
model time: 13.53s (+0.0s waiting for rate limits)
record saved to runs\runs.jsonl
This time the agent took the shortest path: the root index, the knowledge index, the two rules, then the table. It ran its query once, without an error, and handed it in. With no wrong turns, it used 18,809 input tokens, less than half of the earlier run.
Same question, same bundle, a different path. One run shows what can happen, not what usually happens. To know how often the agent finds the right knowledge, and how often its SQL is correct, it has to answer many questions, and its answers have to be checked against the benchmark's answer key.
Put it all behind one command
Every piece now works on its own: prepare starts the database, the agent answers, the trace records. But nobody wants to remember a list of commands. okf-agent works the way coding agents like Claude Code do: you type its name, it gets everything ready, and you ask your question.
For a readable terminal, we add one small package, rich. It draws the spinner while the model thinks, colors the SQL and frames the result. Add it to the dependencies in pyproject.toml and install it:
"rich==15.0.0",
uv sync
Create okf_agent/cli.py:
"""The okf-agent command: get everything ready, then answer questions."""
import argparse
import litellm
from rich.console import Console
from rich.panel import Panel
from rich.syntax import Syntax
from okf_agent import prepare
from okf_agent.agent import ask, load_task, show_call
from okf_agent.config import Config, load_config
from okf_agent.database import run_query
from okf_agent.trace import save, summarize
from okf_compiler.cli import compile_database
console = Console()
HELP = "Type a question about the data. Commands: /db <name>, /raw, /bundle, /exit"
def get_ready(config: Config) -> None:
"""Everything the agent needs. Safe to run every time."""
needed = litellm.validate_environment(model=config.model.name)["missing_keys"]
missing = [key for key in needed if key.endswith("_KEY")]
if missing:
raise SystemExit(f"Add {', '.join(missing)} to your .env file (see PROVIDERS.md).")
names = prepare.database_names(config)
for name in names:
if not (config.agent.bundles / name / "index.md").exists():
with console.status(f"compiling the {name} bundle..."):
compile_database(name)
with console.status("starting the database..."):
prepare.start_database(config, names)
prepare.create_agent_user(config, names)
def answer(config: Config, database: str, question: str, mode: str,
task_id: str | None = None) -> None:
with console.status("thinking..."):
run = ask(config, database, question, mode, on_call=show_call)
record = summarize(run, task_id=task_id)
path = save(record, config.agent.runs)
if run.answer:
console.print(f"\n[bold]answer:[/] {run.answer['answer']}")
if run.answer["sql"].strip():
console.print(Syntax(run.answer["sql"], "sql", word_wrap=True))
result = run_query(config.database, database, run.answer["sql"])
console.print(Panel(result, title="result", title_align="left"))
if run.stop_reason not in ("submitted", "declined"):
console.print(f"[red]stopped: {run.stop_reason} {run.error}[/]")
calls, tokens = record["llm_calls"], record["input_tokens"] + record["output_tokens"]
console.print(f"[dim]{calls} model call{'s' * (calls != 1)} · {tokens:,} tokens · "
f"{record['sql_errors']} SQL errors · saved to {path}[/]\n")
def chat(config: Config, database: str, mode: str) -> None:
console.print(f"[bold]okf-agent[/] · {config.model.name}\n{HELP}\n")
while True:
try:
line = console.input(f"[bold cyan]{database}[/] ({mode})> ").strip()
except (EOFError, KeyboardInterrupt):
break
if line in ("/exit", "/quit"):
break
elif line.startswith("/db "):
database = line.split(maxsplit=1)[1]
elif line in ("/raw", "/bundle"):
mode = line[1:]
elif line.startswith("/") or not line:
console.print(HELP)
else:
try:
answer(config, database, line, mode)
except ValueError as error: # for example, a database with no bundle
console.print(f"[red]{error}[/]")
def main() -> None:
parser = argparse.ArgumentParser(prog="okf-agent", description="Ask a database a question.")
parser.add_argument("--db", default="disaster", help="database to ask (default: disaster)")
parser.add_argument("--raw", action="store_true", help="paste the raw docs instead of the bundle")
commands = parser.add_subparsers(dest="command")
one = commands.add_parser("ask", help="ask one question and exit")
one.add_argument("question", nargs="?", help="your question")
one.add_argument("--db", default=argparse.SUPPRESS, help="database to ask")
one.add_argument("--raw", action="store_true", default=argparse.SUPPRESS,
help="paste the raw docs instead of the bundle")
one.add_argument("--task", help="a benchmark task ID, such as disaster_1")
args = parser.parse_args()
if args.command == "ask" and not (args.task or args.question):
parser.error("ask needs a question or --task")
config = load_config()
get_ready(config)
mode = "raw" if args.raw else "bundle"
if args.command != "ask":
chat(config, args.db, mode)
elif args.task:
task = load_task(config, args.task)
answer(config, task["selected_database"], task["query"], mode, task_id=args.task)
else:
answer(config, args.db, args.question, mode)
if __name__ == "__main__":
main()
Notice four things:
-
get_readyruns every time, and is safe to repeat. It checks that the model's API key is set (LiteLLM knows which key each provider needs), compiles any missing bundle with the compiler from the previous part, starts the database and sets up the read-only user. On later runs, each step finds its work already done. -
The steps print above a spinner.
answerruns the agent insideconsole.status, andshow_call, the same function as before, prints each step as it happens. - After the agent submits, its SQL runs once more, to show you the result. The record is saved before that, so this extra query isn't counted in the agent's metrics.
-
Every question is saved to
runs/runs.jsonl, exactly as with the trace command. The prompt andaskare only two ways into the sameask()function.
You can use it in two ways. Start the prompt and type your questions:
uv run okf-agent
The prompt always shows two things: the database you're asking about and where the agent gets its knowledge, for example disaster (bundle)>. Anything you type is a question, unless it starts with /:
| Command | What it does |
|---|---|
/db <name> |
Switches the database your next questions are about. The name is one of the 18 benchmark databases, the same as the folder names in data/livesqlbench-base-lite, for example /db news or /db credit. |
/raw |
The agent answers without the bundle: the database's three raw documentation files are pasted into its instructions, and it has no read_file tool. Use it to compare the two ways of giving an agent knowledge on the same question. |
/bundle |
Back to the default: the agent finds its knowledge in the OKF bundle with read_file. |
/exit |
Quits the prompt. |
You can also ask one question and exit, which is handy in scripts. --db and --raw do the same as /db and /raw:
uv run okf-agent ask "Which hubs have the highest Resource Utilization Ratio?" --db disaster
uv run okf-agent ask --task disaster_1
See every option:
uv run okf-agent --help
Output:
usage: okf-agent [-h] [--db DB] [--raw] {ask} ...
Ask a database a question.
positional arguments:
{ask}
ask ask one question and exit
options:
-h, --help show this help message and exit
--db DB database to ask (default: disaster)
--raw paste the raw docs instead of the bundle
Ask your first question
Everything is in place. Start the prompt:
uv run okf-agent
Then type the question from the opening of this post:
disaster (bundle)> I need to analyze all distribution hubs based on their Resource Utilization Ratio. Please show the hub registry ID, the calculated RUR value, and their Resource Utilization Classification. Sort the results by RUR from highest to lowest.
Output (the result table is shortened here):
step 1: read_file index.md
step 2: read_file knowledge/index.md
step 3: read_file knowledge/calculation/resource-utilization-ratio.md -> error
step 4: read_file knowledge/calculation/index.md -> error
step 5: read_file knowledge/index.md
step 6: read_file knowledge/resource-utilization-ratio.md
step 7: read_file knowledge/resource-utilization-classification.md
step 8: read_file tables/distributionhubs.md
step 9: run_sql WITH rur_calc AS ( SELECT hubregistry, (hubutilpct/100.0) * (storecapm
step 10: submit_answer (final answer)
answer: The query returns all distribution hubs with their Resource Utilization Ratio (RUR) and classification, sorted from highest to lowest RUR.
WITH rur_calc AS (
SELECT
hubregistry,
(hubutilpct/100.0) * (storecapm3 / (storeavailm3 + 1.0)) AS rur
FROM distributionhubs
)
SELECT
hubregistry,
rur,
CASE
WHEN rur > 5 THEN 'High Utilization'
WHEN rur >= 2 AND rur <= 5 THEN 'Moderate Utilization'
ELSE 'Low Utilization'
END AS resource_utilization_classification
FROM rur_calc
ORDER BY rur DESC;
╭─ result ─────────────────────────────────────────────────╮
│ hubregistry | rur | resource_utilization_classification │
│ HUB_XVHV | 460.207 | High Utilization │
│ HUB_J1JI | 457.992 | High Utilization │
│ HUB_V8QF | 349.307 | High Utilization │
│ ... │
│ (20 of 999 rows) │
╰──────────────────────────────────────────────────────────╯
10 model calls · 37,273 tokens · 0 SQL errors · saved to runs\runs.jsonl
Read the answer from top to bottom:
- The steps show which files the agent opened, and in what order. The knowledge this question needs is in two business rules and one table.
-
The SQL can be checked against those two rules: the formula
(hubutilpct / 100) × (storecapm3 / (storeavailm3 + 1)), and the thresholds 5 and 2 for the three classes. - The result is real rows from the database. Check one by hand: a hub with an RUR above 5 must be labeled High Utilization.
- The last line says what the answer cost and where its record is saved.
In this run, the SQL uses rule 10's formula exactly, and rule 50's thresholds: above 5 is High, 2 to 5 is Moderate, below 2 is Low. The top hub's RUR is 460.207, far above 5, so High Utilization is right.
The run also shows the same wrong turn we saw earlier: in steps 3 and 4, the agent looked for a knowledge/calculation/ folder that doesn't exist, then went back to the knowledge index and found the right files. Those two failed reads cost two extra model calls, and they show up in the 37,273 tokens on the last line.
This is the question the schema alone couldn't answer. No column is called "Resource Utilization Ratio", and nothing in the schema says where "High Utilization" starts. The agent looked both up in the bundle, wrote the query, and ran it before it answered.
But does the bundle make the difference, or would the model do just as well with the raw documentation? That's the next question.
Ask the same question without the bundle
The usual alternative to a knowledge layer is to paste all the documentation into the prompt. Our agent can do exactly that. In raw mode, everything stays the same: the instructions, the model, the database and the run_sql tool. Only one thing changes: instead of reading the bundle with read_file, the agent gets the database's three raw documentation files pasted into its instructions.
That makes raw mode a fair comparison. If anything, it gets slightly more information than the bundle: the raw schema file also contains three sample rows per table, which the bundle leaves out.
In the same prompt, switch to raw mode and ask the same question:
disaster (bundle)> /raw
disaster (raw)> I need to analyze all distribution hubs based on their Resource Utilization Ratio. Please show the hub registry ID, the calculated RUR value, and their Resource Utilization Classification. Sort the results by RUR from highest to lowest.
Output (the result table is shortened here):
step 1: run_sql SELECT hubregistry, (hubutilpct/100.0) * (storecapm3 / (storeavailm3 +
step 2: submit_answer (final answer)
answer: The query returns each distribution hub's ID, Resource Utilization Ratio (RUR), and classification (High, Moderate, or Low Utilization) sorted by RUR from highest to lowest.
SELECT
hubregistry,
(hubutilpct/100.0) * (storecapm3 / (storeavailm3 + 1)) AS rur,
CASE
WHEN (hubutilpct/100.0) * (storecapm3 / (storeavailm3 + 1)) > 5 THEN 'High Utilization'
WHEN (hubutilpct/100.0) * (storecapm3 / (storeavailm3 + 1)) >= 2 AND (hubutilpct/100.0) * (storecapm3 / (storeavailm3 + 1)) <= 5 THEN 'Moderate Utilization'
ELSE 'Low Utilization'
END AS resource_utilization_classification
FROM distributionhubs
ORDER BY rur DESC;
╭─ result ─────────────────────────────────────────────────╮
│ hubregistry | rur | resource_utilization_classification │
│ HUB_XVHV | 460.207 | High Utilization │
│ HUB_J1JI | 457.992 | High Utilization │
│ HUB_V8QF | 349.307 | High Utilization │
│ ... │
│ (20 of 999 rows) │
╰──────────────────────────────────────────────────────────╯
2 model calls · 36,973 tokens · 0 SQL errors · saved to runs\runs.jsonl
Same formula, same thresholds, same rows. With all the documentation already in its instructions, the agent didn't need to look anything up: it wrote the query, ran it once, and handed it in.
Here is how the runs compare. The middle column is the bundle run from the step where the agent first ran SQL, when it took the shortest path:
| Bundle (previous step) | Bundle, shortest path | Raw docs | |
|---|---|---|---|
| Model calls | 10 | 7 | 2 |
| Failed reads | 2 | 0 | – |
| Total tokens | 37,273 | 20,077 | 36,973 |
| SQL errors | 0 | 0 | 0 |
For this question, the bundle didn't win. Raw mode found the same answer with about the same tokens as our latest bundle run, in a fifth of the model calls. The bundle saved tokens only when the agent took the shortest path, and then it used about half.
The two modes spend their tokens differently. In raw mode, the full documentation, about 15,800 tokens for disaster by the estimate from the previous part, is sent with every model call, whatever the question needs. Two calls make it cheap here. With the bundle, the agent pays only for the files it opens, but every extra step resends everything read so far, so wrong turns are expensive.
Which way wins depends on the size of the documentation and on how well the agent navigates, and one question can't settle that. disaster is a small database: its docs fit easily in a prompt. A company's documentation often doesn't, and then pasting everything stops being an option at all. To compare the two fairly, we need many questions, across databases of different sizes, checked against the answer key.
What's next
You now have a SQL agent you start with one command. It looks up business knowledge in an OKF bundle before it writes a query, runs that query as a read-only user, and records what every answer cost. Along the way, we saw three things:
- The knowledge layer works, but navigation costs. In every run, the agent found the two rules and the table the question needs. Its wrong turns, though, cost full model calls, and the index layout invited the same wrong guess more than once.
- Guard the database, not the SQL text. A read-only user, a read-only transaction, a timeout and a row limit stopped every write we tried, without inspecting a single query.
- One run is an anecdote. The same question took 2 to 11 model calls depending on the run and the mode, and on this small database, pasting the raw docs was just as cheap.
In the next part, we put numbers on it. We run many benchmark questions through the same ask() function, with the bundle and with the raw docs, check every answer against the benchmark's answer key, and look closely at where each approach fails.
The code for this series is in the okf-sql-knowledge repository.
Sources
- Open Knowledge Format (OKF) specification, v0.2 — Google Cloud. Concept files, frontmatter, index files (optional in a bundle).
- Your Agents' Knowledge Needs a File Format. Meet OKF — the introduction to OKF and its frontmatter fields.
-
LiveSQLBench and its GitHub repository — BIRD team at HKU and Google Cloud. The
evaluation/folder documents the PostgreSQL Docker image: it loads the databases on first start, logs in asroot, and keeps a template and a plain copy of each database. -
LiveSQLBench-Base-Lite dataset — the 18 databases and the
disaster_1question used in this post (CC BY-SA 4.0). - LiteLLM documentation — one interface to many model providers, token usage and cost; NVIDIA NIM provider.
- NVIDIA API catalog — the free hosted open models used in this post.
-
PostgreSQL 15 release notes — PostgreSQL 15 removed
PUBLIC's right to create objects in thepublicschema; the benchmark runs PostgreSQL 14, where it still exists. -
PostgreSQL 14: privileges and client settings, including
statement_timeout. - psycopg 3 — the PostgreSQL driver (version 3.3.6).
-
Docker Compose: variable interpolation — how Compose reads
.env. -
uv, python-dotenv and rich — project tooling, secrets from
.env, and the terminal interface. - okf-sql-knowledge — the code and guides for this series.
Top comments (0)