AI Data Agent

A LangGraph multi-agent system that turns plain-English requests into either SQL analysis against a PostgreSQL database or ETL work (pull data from an HTTP API, transform it with Pandas, write it to disk).

A top-level router agent reads your question, decides whether it is a sql or an etl job, and hands it to the matching specialist sub-agent.

---

What it actually does

Ask a question about the data → the SQL agent inspects the live database schema, writes a Postgres query, checks it is genuinely read-only (a deterministic parse plus an LLM judge), runs it, and summarises the result in prose.

"What are the different types of payment methods we have?"
  → SELECT DISTINCT payment_method FROM public.payments LIMIT 10;
  → "Your database records five payment methods: credit card, debit card, ..."

Ask for data movement → the ETL agent picks a tool, calls an API or reads a local file, and (for transforms) writes the Pandas code itself and executes it.

"Extract https://pokeapi.co/api/v2/pokemon and save it to data/extract as csv"
  → data/extract/pokemon.csv

The bundled sample dataset is a synthetic ride-hailing business (users, vehicles, rides, payments, ratings) — roughly 10k users and 20k rides across Canadian provinces. See docs/data-model.md.

---

Documentation

Document	What's in it
docs/setup.md	Install, configure .env, seed the database, run the agents
docs/architecture.md	The three graphs, node-by-node, state schemas, control flow
docs/data-model.md	Database tables, relationships, CSV seed files
docs/known-issues.md	Safety caveats and design trade-offs — worth reading before pointing this at real data

---

Quick start

# 1. Python 3.12+, from the project root
uv sync

# 2. Configure credentials
Copy-Item .env.example .env   # then edit it: API keys + Postgres connection

# 3. Create the tables and load the sample CSVs
python feed_db.py

# 4. Ask it something
python main.py "How many rides were cancelled?"

main.py with no argument runs a built-in example. From your own code:

from agents.data_agent import ask

response = ask("How many rides were cancelled?")
print(response["route_response"])        # "sql" or "etl"
print(response["messages"][-1].content)  # the answer

Either sub-agent also runs on its own:

python agents/sql_analyst.py "which city has the most drivers?"
python agents/etl_analyst.py "extract https://pokeapi.co/api/v2/berry to data/extract as csv"

feed_db.py refuses to load on top of existing data. Pass --reload to TRUNCATE the five tables and reload from scratch — that flag is destructive, so point it at a scratch database.

---

Project layout

AI-Data-Agent/
├── main.py                  Entry point — invokes the router agent once
├── feed_db.py               Creates the Postgres schema and COPYs the CSVs in
│
├── agents/
│   ├── data_agent.py        Router graph: classifies sql vs etl, dispatches
│   ├── sql_analyst.py       7-node SQL graph: curate → schema → generate → judge → run → answer
│   └── etl_analyst.py       ReAct-style tool loop + the two ETL tool definitions
│
├── Models/
│   └── schema.py            All Pydantic state/output schemas
│
├── utils/
│   ├── database.py          DatabaseUtil — psycopg2 schema introspection + read-only queries
│   ├── etl_tools.py         ETLTools — HTTP extract, file preview, exec() of generated code
│   ├── llm_pick.py          pick_llm("low"|"medium"|"high"|"claude") model selector
│   └── sql_safety.py        Deterministic read-only check + markdown fence stripping
│
├── data/
│   ├── *.csv                Seed data for the five tables
│   ├── extract/             ETL extraction output
│   └── transform/           ETL transformation output
│
└── docs/                    This documentation

---

How the pieces fit

graph TD
    U[User question] --> R[router_node<br/>Claude + RouterSchema]
    R -->|"sql"| S[sql_analyst graph]
    R -->|"etl"| E[etl_analyst graph]
    S --> DB[(PostgreSQL)]
    E --> API[HTTP API]
    E --> FS[data/extract<br/>data/transform]

Loading

Full node-by-node breakdown in docs/architecture.md.

---

Model selection

utils/llm_pick.py maps a complexity label to a chat model. Cheap steps (question rewriting, final prose) use "low"; SQL generation and the safety judge use "medium"; routing and Pandas code generation use "claude".

Label	Provider	Default model	Override with
low	OpenAI	gpt-5.6-luna	LLM_MODEL_LOW
medium	OpenAI	gpt-5.6-terra	LLM_MODEL_MEDIUM
high	OpenAI	gpt-5.6-sol (defined but never called)	LLM_MODEL_HIGH
claude	Anthropic	claude-sonnet-5	LLM_MODEL_CLAUDE

If your account cannot reach one of these ids, set the matching LLM_MODEL_* variable in .env rather than editing code. Anything other than these four labels raises ValueError.

---

Requirements

• Python 3.12+
• A reachable PostgreSQL instance
• An OpenAI API key and an Anthropic API key (both providers are used in a single run)

Install with uv sync (the lockfile is committed); pip install -r requirements.txt and pip install -e . both work too. See docs/setup.md for details.

Tip: point the user in .env at a Postgres role with only SELECT privileges. The SQL agent is read-only by design, and a restricted role makes that a hard guarantee.