DataPilot — AI Data Analyst Agent

Multi-agent system: ask a question about a dataset in plain language. Agents plan, generate/execute SQL (or route to a Python tool when SQL genuinely doesn't fit), validate the result, optionally chart it, and explain it back honestly — including saying "I can't compute that" instead of guessing. Built with LangGraph orchestration, real MCP tool servers, guardrails proven independent of LLM behavior, a self-correction retry loop, ChromaDB memory, and an eval harness with ground-truth checks.

Builds on the same orchestration/MCP/guardrails skeleton as the Job/Resume Intel Agent project — domain swapped, agent pattern reused.

All 10 steps of the build are complete. See TESTING.md for the full run-by-run debugging journal (12 real test runs, every bug found and how it was fixed) — that document is arguably more informative than this one about how the system actually behaves.

Architecture

                         ┌─────────────┐
                         │   Planner   │  keyword routing (correlation → Python,
                         │             │  chart request → Chart Agent) + memory
                         └──────┬──────┘  retrieval (ChromaDB business terms)
                                │
                  ┌─────────────┴─────────────┐
                  ▼                            ▼
          ┌───────────────┐            ┌───────────────┐
          │   SQL Agent    │            │  Python Agent  │
          │ (MCP sql_tool) │            │(MCP python_tool)│
          └───────┬───────┘            └───────┬───────┘
                  │                            │
                  └─────────────┬──────────────┘
                                 ▼
                          ┌─────────────┐
                          │  Validator  │──── retry (SQL path only,
                          │             │      capped, on genuine errors
                          └──────┬──────┘      or column mismatches)
                                 │
                       validated?│
                    ┌────────────┴────────────┐
                    ▼                          ▼
            ┌───────────────┐          ┌─────────────┐
            │  Chart Agent   │          │ Interpreter │
            │(MCP chart_tool)│─────────▶│             │
            └───────────────┘          └─────────────┘

• Planner — decides SQL vs Python (deterministic keyword check, not LLM judgment — see Known Limitations for why), whether a chart was asked for, and retrieves relevant business-term context from memory.

• SQL Agent — fetches schema + generates + executes SQL, entirely through MCP tool calls. Can respond NOT_COMPUTABLE when business context says a question isn't answerable from this data (e.g. profit).

• Python Agent — handles correlation questions SQL fights with (aggregate + raw columns "in one row"). No arbitrary code execution — fixed, safe pandas operations only.

• Validator — catches SQL errors, zero-row/NULL-only results, and column-mismatches (SQL that runs fine but ignores a column the question asked about). Drives the retry loop.

• Chart Agent — generates a bar chart when asked, picking axes by data type (not SQL column order — a real bug found and fixed in testing). Fails gracefully on shapes that don't fit.

• Interpreter — compiles the final plain-language answer, or explains a failure honestly instead of hallucinating from bad data.

MCP layer: sql_tool, python_tool, chart_tool — three real MCP servers over stdio, not hardcoded functions. Agents talk to them the same way they'd talk to any MCP-compliant tool.

Guardrails: read-only SQL enforcement lives server-side in sql_tool_server.py (blocks DROP/DELETE/UPDATE/INSERT/ALTER, rejects multi-statement SQL) — and is proven to work independent of the LLM's behavior via tests/test_sql_guardrail.py (7/7 malicious queries blocked, 3/3 safe queries pass), not just by an adversarial test question that the model might happen to dodge on its own.

Memory: app/memory.py — ChromaDB-backed retrieval of 5 business-term definitions (revenue, profit, discount, region, order), filtered by a distance threshold so irrelevant questions get no injected context. Concrete payoff: this dataset has no cost data, so the SQL Agent can correctly decline profit questions instead of guessing.

Observability: app/observability.py — every run's full state is logged to outputs/traces/*.json (always works, no external account). Langfuse is wired too, centrally in get_llm(), but only activates if credentials are set in .env — never required to run the project.

Eval harness: tests/eval_harness.py checks numeric results against ground truth computed independently with pandas (not the agent's own output), plus a row-count consistency check that scans the Interpreter's answers for stated counts and flags mismatches against the real data.

Related MCP server: mcp-sql

Build order

1. ✅ Domain + sample dataset (data/sales.csv) + fixed test queries

2. ✅ State schema + baseline graph

3. ✅ MCP sql_tool server (DuckDB, read-only enforced)

4. ✅ Guardrail proven (tests/test_sql_guardrail.py, LLM-bypassing)

5. ✅ Self-correction retry loop (bad SQL → re-plan → re-execute)

6. ✅ Python analysis tool (MCP-wrapped) for non-SQL-expressible questions

7. ✅ Chart tool (MCP-wrapped) + visualization logic

8. ✅ Memory (ChromaDB for business-term RAG) + observability (Langfuse)

9. ✅ Eval harness

10. ✅ UI (Streamlit)

Setup

Two terminals — Ollama needs to stay running while anything else runs.

Terminal 1 (leave running):

ollama serve

Terminal 2:

pip install -r requirements.txt
cp .env.example .env
ollama pull llama3.2:3b   # only needed once

Pinned to mcp<2.0.0 — MCP 2.x renamed FastMCP to MCPServer and broke this project's server code. If you see ModuleNotFoundError: No module named 'mcp.server.fastmcp', run pip install "mcp<2.0.0" --force-reinstall.

Running it

CLI — runs the fixed test queries and prints full trace/debug info:

python -m app.main

UI — chat-style interface, example questions in the sidebar, charts rendered inline, an expandable panel per answer showing the SQL/Python used, business context, and full agent trace:

streamlit run streamlit_app.py

Guardrail proof — bypasses the LLM entirely, sends known-malicious SQL straight to the MCP tool:

python -m tests.test_sql_guardrail

Eval harness — ground-truth + row-count consistency checks:

python -m tests.eval_harness

Set DEBUG_MEMORY=1 before any of the above to print ChromaDB's actual retrieval distances per question (useful for tuning _MAX_DISTANCE in app/memory.py).

Fixed test queries

1. Simple aggregation — total revenue

2. Group + filter — sales by region, Electronics only

3. Correlation — discount vs quantity (routes to Python Agent)

4. Adversarial — SQL-injection-style instruction embedded in the question

5. Nonsense / no-match — question the dataset can't answer

6. SQL-error-prone — tests the retry loop directly

7. Chart request — grouped data, should produce a real bar chart

8. Chart on a shape that may or may not reduce to a single aggregate

9. Business-term memory — profit margin (not computable from this data)

Project structure

data/
  sales.csv                 # sample dataset used by all test queries
app/
  state.py                   # shared GraphState
  llm.py                     # centralized LLM client (Ollama default) + Langfuse wiring point
  mcp_client.py               # generalized MCP client (works with any server script)
  memory.py                   # ChromaDB business-term retrieval
  observability.py             # local trace logging + optional Langfuse handler
  utils.py                    # shared helpers (e.g. mentioned_columns)
  graph.py                    # LangGraph wiring
  main.py                     # CLI entrypoint, runs fixed test queries
  agents/
    planner.py                  # routing + memory retrieval
    sql_agent.py                # calls MCP sql_tool
    python_agent.py              # calls MCP python_tool (correlation questions)
    validator.py
    chart_agent.py               # calls MCP chart_tool
    interpreter.py
  mcp_servers/
    sql_tool_server.py           # FastMCP server, DuckDB-backed, read-only enforced
    python_tool_server.py         # FastMCP server, pandas-backed, fixed safe stats only
    chart_tool_server.py           # FastMCP server, matplotlib-backed, fixed bar chart only
tests/
  test_sql_guardrail.py       # bypasses the LLM, proves the read-only guardrail directly
  eval_harness.py               # benchmark: ground-truth checks + row-count consistency check
streamlit_app.py               # UI entrypoint

Generated charts are written to outputs/ and run traces to outputs/traces/ (both gitignored — regenerate by running the app). ChromaDB's local data lives in chroma_data/ (also gitignored, regenerated automatically on first run).

Known limitations (honest)

• Small-model self-correction is limited. With llama3.2:3b, a genuinely malformed query sometimes gets retried with the identical broken SQL rather than a fix — the retry cap still kicks in and the system fails honestly rather than hallucinating, but it doesn't always self-correct. A larger model would likely do better; this is a model-capability limit, not a guardrail or architecture gap.

• No arbitrary Python code execution. By design — a real sandboxed run_python(code, df) tool with Docker isolation is future work. exec()-ing LLM-generated code without a sandbox in place was deliberately avoided rather than "temporarily" allowed.

• Routing between SQL and Python is keyword-based, not learned. Deliberately conservative (SQL is the default; only a correlation keyword match routes to Python) after the LLM's own routing judgment proved unreliable twice in testing — once nearly skipping SQL entirely with no fallback. A broader set of "SQL can't express this well" cases (regression, forecasting) isn't covered.

• Chart Agent and Validator assume this dataset's specific shape. Chart Agent only produces one shape (category + numeric value → bar chart); the Validator's checks are written against this table's structure. Generalizing to arbitrary uploaded datasets would need both to reason about shape/schema dynamically.

• Python Agent's column discovery is dynamic (fixed from an earlier hardcoded list) — fetched from the SQL tool's schema at runtime, filtered to numeric types, so a different dataset's numeric columns work without editing the file.

• The Interpreter's narration isn't fully fact-checked against the data. Seen once in testing: a correct 20-row SQL result got summarized as "13 unique orders." Partly addressed by tests/eval_harness.py's row-count consistency check, which hasn't yet caught a live violation on this benchmark — only proven to not false-positive.

• Business-term memory covers 5 fixed, hand-written terms. Adding new ones means editing _SEED_TERMS in app/memory.py directly. Retrieval distance threshold (_MAX_DISTANCE) is an empirical starting point, confirmed reasonable via DEBUG_MEMORY=1 runs but not rigorously tuned.

• Langfuse integration is wired but not verified against real credentials — only local trace logging has been confirmed working.

• This is a single-dataset prototype, not a general data-analyst tool. It demonstrates the architecture pattern — multi-agent orchestration, MCP tool boundaries, guardrails proven independent of LLM behavior, retry/self-correction, memory, eval — on one fixed CSV. Supporting arbitrary uploaded datasets (the original project plan's Phase 5) would need the points above addressed, not just more test questions against this one.

Setup notes for the UI specifically

The Streamlit app is a thin layer over the same build_graph() used by the CLI — no separate agent logic. Each question still spawns fresh MCP subprocess calls per tool use (same as the CLI), so expect the same per-query latency you'd see running python -m app.main, just rendered in a browser instead of a terminal.