A LangGraph-based ReAct agent that answers questions about the Bitext Customer Support dataset (26,872 customer/agent interactions across 11 categories and 27 intents).
The agent handles three kinds of queries:
- Structured — concrete questions answered by filtering/counting/sampling the data ("How many refund requests?", "Show me 5 SHIPPING examples").
- Unstructured — open-ended questions answered by reading and summarizing text content ("Summarize the FEEDBACK category").
- Out-of-scope — anything unrelated to the dataset is politely declined ("Who is the president of France?").
This repo covers Tasks 1, 2, 3, and Bonuses A + B — the full agent, persistent memory, an MCP server exposing the same tools to external clients, a Streamlit chat UI, and an interactive query recommender.
- Python 3.10 or newer (tested on 3.12).
- A Nebius Token Factory API key.
# 1. Clone or unzip the repo, then enter the directory
cd "naomi submission"
# 2. Create a virtual environment and install dependencies
python3 -m venv venv
source venv/bin/activate # on Windows: venv\Scripts\activate
pip install --upgrade pip
pip install -r requirements.txt
# 3. Set your Nebius API key
cp .env.example .env
# Open .env and replace the placeholder with your real key:
# NEBIUS_API_KEY=your_real_key_herepython main.pyOn the first run the agent will download the Bitext CSV (~20 MB) into
data/bitext_dataset.csv. Subsequent runs read from that cache.
You should land in an interactive prompt:
Customer Service Data Analyst Agent
Ask me anything about the Bitext customer-support dataset:
- "What categories exist?"
- "How many refund requests?"
- "Summarize the FEEDBACK category."
Commands: /help, /exit, /quit
You: ▮
Type /exit (or Ctrl+D) to quit.
[START]
│
▼
[ Router ] ← classifies the query
│
┌──────┼──────┐
▼ ▼
[ Decline ] [ ReAct Agent ] ← out_of_scope vs. (un)structured
│ │
└──────┬──────┘
▼
[END]
- Router node classifies the incoming question as
structured/unstructured/out_of_scopeusing a small, fast model with a Pydantic-enforced output schema. - If
out_of_scope, the Decline node returns a polite refusal without consulting any tools or general knowledge. - Otherwise, the ReAct agent (LangGraph's
create_react_agent) takes over: it picks a tool, observes the result, may chain to another tool, and produces a final natural-language answer.
A recursion_limit is set on the outer graph (MAX_ITERATIONS * 2 + 4).
If the agent doesn't reach a final answer in time, a graceful fallback
message is returned instead of an infinite loop.
Two models are used, each suited to its role:
| Role | Model | Why |
|---|---|---|
| Router | Qwen/Qwen3-30B-A3B-Instruct-2507 |
Mixture-of-Experts with only 3B active parameters — very low latency, plenty smart for a 3-way classification. |
| Agent | Qwen/Qwen3-32B |
Strong instruction-following and reliable tool-calling. In practice it follows the "stop after a successful tool call" rule more reliably than Llama 3.3 70B, which tended to loop on get_examples when filtering by category. |
Both are accessed via Nebius's OpenAI-compatible API
(https://api.studio.nebius.com/v1/).
Six tools, all with Pydantic input schemas and detailed WHEN-TO-USE
docstrings (see src/tools.py):
| Tool | Returns | When to use |
|---|---|---|
list_categories |
list of category names | "What categories exist?" |
list_intents |
list of intent names, optionally per category | Discover the right intent name before filtering. |
count_rows |
{count, filters_applied} |
"How many X?" questions. |
get_examples |
sample rows (instruction + response) | "Show me N examples of X". |
intent_distribution |
intent → count for one category | "What's the breakdown of X?". |
get_texts_for_summary |
a batch of rows for the LLM to summarize | Open-ended "summarize / how do agents respond" questions. |
Each tool's docstring also tells the model to OMIT optional arguments
(rather than passing the string 'null'), and the schemas use
field_validators that coerce 'null'/'none'/'' to None as a
safety net.
The toolset is intentionally small and composable. Per the assignment:
"A few well-designed tools beat many poorly described ones." The
example multi-step path the assignment hints at —
list_intents('REFUND') → get_examples(intent='get_refund') — works
out of the box (see Test 5 below).
The CLI streams the agent's reasoning to the terminal as it happens:
[Router] structured — reason— the classification.[Tool call] count_rows(intent='get_refund')— each tool invocation.[Result of count_rows] {...}— the tool's output (truncated).🤖 ...— the final answer.
This satisfies the "print reasoning steps, not just the final answer" requirement.
The agent has two complementary memory layers, both persistent across restarts.
Powered by LangGraph's SqliteSaver checkpointer (file:
checkpoints.sqlite). Pass --session <id> and the same id will
restore the same conversation, even after restarting the CLI:
python main.py --session naomi
# > Show me 3 examples from the REFUND category
# > /exit
python main.py --session naomi # new process, same session
# > Show me 3 more ← agent knows "more" means REFUNDIf you omit --session, the CLI uses the session default — memory is
always on; different --session values produce independent threads
with no shared history.
The router is conversation-aware: it sees the recent turns when classifying the latest question, so follow-ups like "what about refunds?" or "total of the last two?" are correctly classified as in-scope and the agent does arithmetic over earlier answers without re-querying the data.
A separate Markdown file per user under context/<session>.md (e.g.
context/naomi.md). After every turn, a small "summary node" reads the
current profile + latest exchange and asks Qwen/Qwen3-30B-A3B-Instruct-2507
(the cheap router model) whether anything new and durable should be
added. The profile is then injected into the agent's system prompt on
subsequent turns, so questions like "What do you remember about me?"
are answered from the profile without calling any data tools.
Profiles capture durable facts only (name, role, recurring interests, preferences) — not a replay of past messages.
Example profile after introducing yourself:
- Name: Naomi
- Role: Data Analyst
- Company: Nebius
- Long-term interest: refund patterns in customer-support dataBoth checkpoints.sqlite and context/ are .gitignored (runtime
state, not source).
naomi submission/
├── data/
│ └── bitext_dataset.csv # downloaded on first run
├── context/ # per-user profile MD files (Task 2b)
├── src/
│ ├── __init__.py
│ ├── config.py # model names + ChatOpenAI factory
│ ├── data_loader.py # downloads / caches the CSV
│ ├── tools.py # 6 tools + Pydantic input schemas
│ ├── router.py # query classifier (conversation-aware)
│ ├── agent.py # LangGraph wiring + checkpointer + run_agent()
│ ├── memory.py # profile load/save + summary node (Task 2b)
│ └── cli.py # interactive REPL with reasoning trace
├── main.py # CLI entry point (python main.py [--session id])
├── mcp_server.py # MCP server entry point (Task 3)
├── streamlit_app.py # Streamlit chat UI entry point (Bonus A)
├── checkpoints.sqlite # created at runtime (Task 2a)
├── requirements.txt
├── .env.example
├── .gitignore
├── tests_output.txt # captured run of all 8 Task 1 example queries
├── tests_output_task2.txt # captured runs of all 3 Task 2 scenarios
├── PLAN.md # planning document (kept for reference)
└── README.md
These are the eight example queries from the assignment, each of
which the agent answers correctly. The full captured trace is in
tests_output.txt.
| # | Query | Expected route | Notes |
|---|---|---|---|
| 1 | What categories exist in the dataset? | structured | One tool call (list_categories). |
| 2 | How many refund requests did we get? | structured | count_rows(intent='get_refund') → 997. |
| 3 | Show me 5 examples of the SHIPPING category. | structured | get_examples(n=5, category='SHIPPING'). |
| 4 | Summarize how agents respond to complaint intents. | unstructured | get_texts_for_summary(intent='complaint', n=30) then the LLM summarizes. |
| 5 | Show me examples of people wanting their money back. | structured | Multi-step: list_intents('REFUND') → get_examples(intent='get_refund'). |
| 6 | What is the distribution of intents in the ACCOUNT category? | structured | One call to intent_distribution. |
| 7 | What's the best CRM software for handling complaints? | out_of_scope | Declined politely, no tools called. |
| 8 | Who is the president of France? | out_of_scope | Declined politely, no tools called. |
To re-run them yourself:
python main.py
# then paste each question, one per lineTo see the max-iterations fallback message in action, temporarily
lower MAX_ITERATIONS in src/config.py to 2
and re-run a complex query like "How many refund requests did we get?".
You should see:
🤖 I couldn't reach a final answer within the iteration limit
(2 steps). Could you rephrase your question or break it into
smaller parts?
A standalone FastMCP server in mcp_server.py exposes
all six data-analyst tools over the Model Context Protocol, so any
MCP-compatible client (Claude Desktop, Cursor, a custom Python client,
etc.) can call them.
Exposed tools (the same 6 the agent uses, re-registered with
@mcp.tool):
list_categorieslist_intentscount_rowsget_examplesintent_distributionget_texts_for_summary
The MCP server is a SEPARATE process from the CLI agent and speaks
STDIO transport — the MCP default that Claude Desktop, Cursor, and
the FastMCP Client use:
python mcp_server.pyYou usually don't run the server manually — MCP clients launch it as a subprocess on demand. The example below shows exactly that.
This Python snippet uses FastMCP's Client to talk to the server.
It works whether you've already started the server or not — with
STDIO, the client launches the server itself.
import asyncio
from fastmcp import Client
async def main():
# STDIO transport: Client launches `python mcp_server.py` for you.
async with Client("mcp_server.py") as client:
# 1. Discover what tools the server exposes
tools = await client.list_tools()
print(f"{len(tools)} tools available:")
for t in tools:
print(f" - {t.name}")
# 2. Call one of them
result = await client.call_tool(
"count_rows", {"intent": "get_refund"}
)
print(f"\ncount_rows(intent='get_refund') -> {result.data}")
# Expected: {'count': 997, 'filters_applied': {'category': None,
# 'intent': 'get_refund'}}
asyncio.run(main())Save the snippet as mcp_client_demo.py (or paste it into a REPL)
and run it with python mcp_client_demo.py. You should see the six
tool names followed by the refund count.
Instead of (or in addition to) the Python client above, you can expose the same six tools directly to Claude Desktop as a "native" MCP server.
-
Open Claude Desktop → Settings → Developer → Edit Config. This opens
claude_desktop_config.json. -
Add (or merge into) the
mcpServersblock:{ "mcpServers": { "bitext-analyst": { "command": "/ABSOLUTE/PATH/TO/venv/bin/python", "args": [ "/ABSOLUTE/PATH/TO/naomi submission/mcp_server.py" ] } } }Replace both paths with the absolute paths on your machine. The
commandmust point at the Python inside this project's venv (so the server has access to the right dependencies); on Windows it would be...\venv\Scripts\python.exe. -
Quit Claude Desktop fully and reopen it. The six tools (
list_categories,count_rows,get_examples,intent_distribution,list_intents,get_texts_for_summary) now appear under the "🔌" tools menu and Claude can call them in conversation, e.g.:"Using the bitext-analyst tools, how many refund requests are in the dataset?"
A browser chat wrapper around the same agent, in
streamlit_app.py.
streamlit run streamlit_app.pyStreamlit opens http://localhost:8501 automatically. The page has a
chat input at the bottom, a session-ID box in the sidebar, and renders
the full reasoning trace (router decision, tool calls, tool
results) inline as the agent works — not just the final answer.
- Live reasoning trace. Each turn opens a collapsible status box
("🧠 Thinking…") that streams the router's decision (
🔀), each tool call (🔧), and each tool result (📊, JSON-pretty-printed) as they happen. When the agent finishes the box collapses to "✅ Reasoning complete" — click to expand and review. - Session ID in the sidebar. Type any session id; switching it saves the in-progress chat back to its thread and reloads the new session's history from the LangGraph SQLite checkpointer. A toast ("📂 Resumed session X (N prior messages)") confirms the load. Brand-new ids start with an empty chat.
- Same memory as the CLI. The Streamlit UI and
python main.pysharecheckpoints.sqlite, so a conversation you started in the CLI is fully recoverable in the browser (and vice versa). - Per-user profile. The summary node still runs after every turn,
so
context/<session>.mdkeeps growing across both UIs.
When the user asks for ideas ("what should I query next?", "any
suggestions?", "got any ideas?"), the agent enters a recommendation
flow defined entirely in the agent's system prompt
(AGENT_SYSTEM_PROMPT — "RECOMMENDATION MODE"):
- It reads the conversation history and the per-user profile to find topics the user cares about.
- It proposes ONE concrete query in natural language and ends with a confirmation question — no tool is called yet.
- If the user refines the suggestion (e.g. "I'd rather see examples instead"), the agent revises and asks again — still no tool.
- Only an unambiguous yes ("yes", "go ahead", "do it") triggers the tool call. A decline ("no thanks") drops the suggestion politely.
Works in both the CLI and the Streamlit UI, because both go through the same graph. The exact assignment example flow is reproduced:
You: What should I query next?
🤖 (suggestion — e.g. "list of categories might be a good start;
want me to show it?") ← no tool call
You: I'd rather see examples instead.
🤖 (refined suggestion — e.g. "5 random examples from the ACCOUNT
category; should I go ahead?") ← no tool call
You: Yes, do it.
🤖 [Tool call] get_examples(n=3, category='ACCOUNT')
…results… ← tool call now
The router was also updated so suggestion-seeking phrases are
classified as structured (in-scope) instead of being rejected.
AuthenticationError: 401— yourNEBIUS_API_KEYis missing, expired, or for a different Nebius product. Generate a new one at tokenfactory.nebius.com.SSL: CERTIFICATE_VERIFY_FAILEDon a corporate network —truststore(in requirements.txt) is included specifically to fix this; make surepip install -r requirements.txtran to completion.- Hugging Face download hangs at 0% — handled. We download the CSV
directly via
pandas.read_csv(URL)instead of through thedatasetslibrary, which avoids HF Hub's rate limit on unauthenticated requests.