AI Agents with LangGraph for Data Analysis

pip install langgraph langchain langchain-openai pandas numpy

export OPENAI_API_KEY="sk-..."

from typing import TypedDict, Annotated, Sequence
from langchain_core.messages import BaseMessage
import operator

class AgentState(TypedDict):
    messages: Annotated[Sequence[BaseMessage], operator.add]
    data_loaded: bool
    dataset_summary: str

import pandas as pd
import numpy as np
from langchain_core.tools import tool

# In-memory "database" for the demo
_df: pd.DataFrame | None = None

@tool
def load_dataset(path: str) -> str:
    """Load a CSV dataset into memory. Returns a summary of columns and shape."""
    global _df
    _df = pd.read_csv(path)
    summary = (
        f"Dataset loaded: {_df.shape[0]} rows × {_df.shape[1]} columns.\n"
        f"Columns: {', '.join(_df.columns)}\n"
        f"Dtypes:\n{_df.dtypes.to_string()}"
    )
    return summary

@tool
def run_query(query: str) -> str:
    """
    Run a pandas query on the loaded dataset.
    Use standard pandas expressions, e.g.:
      - df[df['churn'] == 1]['monthly_charges'].mean()
      - df.groupby('plan')['churn'].value_counts()
    The variable 'df' refers to the loaded dataset.
    """
    global _df
    if _df is None:
        return "No dataset loaded. Call load_dataset first."
    try:
        result = eval(query, {'df': _df, 'pd': pd, 'np': np})
        return str(result)
    except Exception as e:
        return f"Query error: {e}"

@tool
def describe_column(column_name: str) -> str:
    """Return descriptive statistics for a specific column in the loaded dataset."""
    global _df
    if _df is None:
        return "No dataset loaded."
    if column_name not in _df.columns:
        return f"Column '{column_name}' not found. Available: {list(_df.columns)}"
    stats = _df[column_name].describe().to_string()
    null_pct = _df[column_name].isnull().mean() * 100
    return f"{stats}\nNull values: {null_pct:.1f}%"

tools = [load_dataset, run_query, describe_column]

from langchain_openai import ChatOpenAI
from langchain_core.messages import SystemMessage
from langgraph.graph import StateGraph, END
from langgraph.prebuilt import ToolNode

# LLM with tools bound
llm = ChatOpenAI(model='gpt-4o-mini', temperature=0)
llm_with_tools = llm.bind_tools(tools)

SYSTEM_PROMPT = """You are a data analysis assistant. You have access to tools that 
let you load CSV datasets, run pandas queries, and describe columns.

When a user asks a question:
1. Load the dataset if it has not been loaded yet.
2. Explore the data to understand its structure.
3. Run the specific queries needed to answer the question.
4. Summarize your findings clearly, with numbers.

Always verify your queries worked before drawing conclusions."""

def agent_node(state: AgentState) -> AgentState:
    """The reasoning node: calls the LLM to decide the next action."""
    messages = [SystemMessage(content=SYSTEM_PROMPT)] + list(state['messages'])
    response = llm_with_tools.invoke(messages)
    return {'messages': [response]}

def should_continue(state: AgentState) -> str:
    """Routing function: continue the loop or end?"""
    last_message = state['messages'][-1]
    if hasattr(last_message, 'tool_calls') and last_message.tool_calls:
        return 'tools'
    return END

# Build the graph
tool_node = ToolNode(tools)

graph = StateGraph(AgentState)
graph.add_node('agent', agent_node)
graph.add_node('tools', tool_node)

graph.set_entry_point('agent')
graph.add_conditional_edges('agent', should_continue)
graph.add_edge('tools', 'agent')  # After tool call, go back to agent

agent = graph.compile()

from langchain_core.messages import HumanMessage

def run_agent(question: str, data_path: str | None = None):
    initial_message = HumanMessage(
        content=f"Dataset path: {data_path}\n\nQuestion: {question}"
        if data_path else question
    )

    result = agent.invoke(
        {'messages': [initial_message], 'data_loaded': False, 'dataset_summary': ''},
        config={'recursion_limit': 15},
    )

    final_message = result['messages'][-1]
    return final_message.content

# Example: analyze a sample dataset
answer = run_agent(
    question="What is the churn rate by number of products? Which group has the highest churn?",
    data_path="data/churn.csv",
)
print(answer)

I loaded the dataset (2000 rows × 6 columns) and ran the analysis.

Churn rate by number of products:
  1 product:  42.1%  ← highest churn
  2 products: 31.7%
  3 products: 22.4%
  4 products: 15.8%
  5 products: 9.3%   ← lowest churn

Customers with only 1 product have the highest churn rate at 42.1%, 
which is 4.5× higher than customers with 5 products. This suggests 
that cross-selling additional products significantly improves retention.

from langchain_core.messages import AIMessage, ToolMessage

result = agent.invoke({
    'messages': [HumanMessage(content="Load data/churn.csv and tell me the average monthly charges for churned vs non-churned customers.")],
    'data_loaded': False,
    'dataset_summary': '',
})

for msg in result['messages']:
    if isinstance(msg, AIMessage) and msg.tool_calls:
        for tc in msg.tool_calls:
            print(f"→ Tool call: {tc['name']}({tc['args']})")
    elif isinstance(msg, ToolMessage):
        print(f"← Tool result: {msg.content[:200]}")
    elif isinstance(msg, AIMessage):
        print(f"✓ Final answer: {msg.content}")

→ Tool call: load_dataset({'path': 'data/churn.csv'})
← Tool result: Dataset loaded: 2000 rows × 6 columns. Columns: tenure_months, monthly_charges, ...
→ Tool call: run_query({'query': "df.groupby('churn')['monthly_charges'].mean()"})
← Tool result: churn\n0    58.3\n1    79.2\ndtype: float64
✓ Final answer: Churned customers pay significantly more on average ($79.2/month) compared to 
  retained customers ($58.3/month)...