Data Analyst Agent

A conversational LangGraph agent that answers questions about a 26,872-row customer-support dataset. It routes each question by type, calls typed Pydantic tools in a ReAct loop, remembers conversations across restarts, and is exposed through a CLI, a Streamlit web UI and an MCP server.

How it works

user question
     │
     ▼
 [classify] ── LLM router with structured output (Pydantic)
     │
     ├── Out-of-scope ──► [reject]            polite refusal, no tool calls
     ├── Structured   ──► [ReAct, ≤10 steps]  counts, distributions, lookups
     └── Unstructured ──► [ReAct, ≤20 steps]  themes, tone, "how do agents handle…"
                              │
                              ▼
                 typed tools over the dataset (pandas)

• Routing before reasoning. A cheap classification step decides how much reasoning budget a question gets: 10 ReAct steps for factual lookups, 20 for interpretive ones. Out-of-scope requests never reach the tools.
• Typed tools. Every tool is a StructuredTool with a Pydantic input schema and typed output. search_keyword caps its results at 10 and adds an overflow note, which keeps the context window small.
• Persistent memory. A LangGraph SqliteSaver checkpointer keeps full conversation history per --session across restarts.
• Robustness fixes found in testing:
  • DeepSeek leaked its raw tool-call markup in multi-turn chats, so the agent is now instructed to always use structured tool calls.
  • Runaway loops are caught with GraphRecursionError and answered with a friendly message.
  • Graph construction is deferred so startup stays fast.

Results: router evaluation

The evals/eval_classify.json set has 100 labelled questions across structured, unstructured and out-of-scope. They are evaluated in parallel with evals/eval_classify.py.

	Before prompt work	After
Unstructured accuracy	45%	85%
Overall accuracy	80%	95%

Tools

Tool	Purpose
get_distribution(column)	Value counts for category, intent or flags
get_examples(n, offset, category, intent)	Paginated row lookup with filters
count_rows(category, intent)	Filtered counts
search_keyword(keyword, column)	Case-insensitive search, capped at 10 hits
get_stats()	Row count, columns, unique counts, average message lengths

The MCP server also exposes get_categories() and get_intents(category).

Dataset: Bitext customer support, 11 categories and 27 intents (ACCOUNT, ORDER, REFUND, SHIPPING, …).

Run it

python -m venv .venv && source .venv/bin/activate
pip install -r requirements.txt
cp .env.example .env            # add NEBIUS_API_KEY

Interface	Command
CLI	python main.py --session my_session
Web UI	streamlit run streamlit_app.py, then open http://localhost:8501
MCP server	python mcp_server.py (test with python mcp_client.py)
LangGraph Studio	langgraph dev (graph data_analyst in langgraph.json)
Router eval	python evals/eval_classify.py

Example questions:

• What's the distribution of categories?
• Show me 3 examples of complaints.
• How do agents typically handle refund requests?

Project structure

├── agent.py               # LangGraph graph: classify → reject / ReAct, tools, prompts
├── data_analysis_tools.py # DatasetAnalyzer (pandas over the HF dataset)
├── main.py                # CLI with persistent sessions
├── streamlit_app.py       # chat UI with session management
├── mcp_server.py          # FastMCP server exposing the dataset tools
├── mcp_client.py          # MCP smoke-test client
├── langgraph.json         # LangGraph Platform entry point
└── evals/                 # 100-question router eval

Context

This was an assignment for Nebius Academy: AI Performance Engineering ("From AI Model to AI Agent", Assignment 3).

Authors

Nathanael Bloch · Tomer Porat