Teaching AI to Analyze Data: PydanticAI Takes on LangChain!

Discover how PydanticAI transforms GenAI agents with custom tools, retries, and a schema-first approach—perfect for data pros and AI enthusiasts alike.

A PydanticAI agent analyzing DataFrames and data about car sales

Overview

What if you could teach your AI agent to analyze data like a pro? Imagine asking it to crunch numbers, extract insights, or letting non-technical users in your organization ask it questions in natural language—all without breaking a sweat.

In this post, we’ll explore how we took a LangChain-powered data analysis notebook and rewrote it using the shiny, all-new AI framework, PydanticAI. From crafting tools that actually understand your DataFrame to handling agent retries like a pro, this adventure is packed with insights and practical tips. Ready to level up your GenAI game? Let’s dive in!

If you’ve been following our recent posts on GenAI agents, you might recall our previous adventure in translating a LangChain-based notebook into a PydanticAI-powered one. If you missed it, don’t worry—this post stands perfectly well on its own. But, if you’re curious, you can check out that earlier post to see how we first “PydanticAI-fied” a conversational agent.

Both this and the previous post are based on notebooks from the excellent GenAI_Agents repository by Nir Diamant. If you haven’t checked it out yet, it’s a treasure trove of examples showcasing how to build and deploy Large Language Model (LLM)-based agents for a variety of tasks. The code is clean, the concepts are well-explained, and it’s one of the best resources for learning how to compose and chain GenAI agents. Plus, you’ll find all our new PydanticAI-powered notebooks there, so you can easily use the repository as a source of inspiration—whether you prefer LangChain or PydanticAI!

---

The Task at Hand: Data Analysis Simple Agent

Nir’s repo includes a cool LangChain-based notebook called “Data Analysis Simple Agent”. In this notebook, an agent can query a Pandas DataFrame, figure out its columns, calculate average prices, count rows, and basically do the kind of data gymnastics you might otherwise perform manually. It’s neat, tidy, and highlights one of LangChain’s strengths: easy integration with tools like Pandas DataFrame queries.

But what if you wanted to achieve the same thing in PydanticAI? In my previous post, we translated a simple conversational LangChain agent into PydanticAI. This time, we’re tackling a slightly more complex agent. Along the way, we’ll dive into PydanticAI’s features, explore what makes it unique, and highlight some key differences between it and LangChain.

---

A pedantic AI Agent thinking how to function outside of LangChain

The Challenges: Replacing a Built-in Feature

Here’s the twist: LangChain already provides a neat and tidy agent specifically designed for handling Pandas DataFrames:

from langchain_experimental.agents.agent_toolkits import create_pandas_dataframe_agent

...

agent = create_pandas_dataframe_agent(
    ChatOpenAI(model="gpt-4o", temperature=0),
    df,
    verbose=True,
    allow_dangerous_code=True,
    agent_type=AgentType.OPENAI_FUNCTIONS,
)

This snippet from the LangChain notebook showcases one of LangChain’s greatest strengths: its rich ecosystem of integrations and ready-made tools that can be effortlessly incorporated into your project.

PydanticAI, on the other hand, as a younger framework, doesn’t yet offer this level of built-in support. For example, there’s no official “Pandas DataFrame Agent” at the moment. This meant rolling up our sleeves and crafting a custom tool for querying the DataFrame from scratch.

Here’s what the initial tool definition looked like:

async def df_query(ctx: RunContext[Deps], query: str) -> str:
    """A tool for running queries on the `pandas.DataFrame`. Use this tool to interact with the DataFrame.

    `query` will be executed using `pd.eval(query, target=df)`, so it must contain syntax compatible with
    `pandas.eval`.
    """

    return str(pd.eval(query, target=ctx.deps.df))

It seems straightforward, but there’s a lot going on here. Let’s break it down:

Function Arguments:

• The function takes two arguments: a RunContext[Deps] object (ctx) and a query string.

• RunContext is a PydanticAI class that provides context and information about the current call. It’s also how we access dependencies via dependency injection. Think of it as a bridge between the generic function definition and the specific runtime call.

• query is the command the agent “wants” us to execute on the DataFrame. This is the core argument the agent provides to interact with the DataFrame.

Docstring:

We’ve included a clear and detailed docstring for the function and its arguments. This is crucial because PydanticAI uses both the function signature and the docstring to build a schema for tool calls. This ensures the agent knows exactly what’s expected of it when interacting with the tool.

Implementation:

• The function uses pd.eval to execute the query on the DataFrame (ctx.deps.df).

• pd.eval is deliberately limited to a subset of operations to reduce the risk of malicious code execution.

• We return the results as a string, which ensures the output is serializable. While other return types are possible, serializability is required.

Now, let’s see how this tool ties into the agent itself:

@dataclass
class Deps:
    df: pd.DataFrame

agent = Agent(
    model='openai:gpt-4o-mini',
    system_prompt="""You are an AI assistant that helps extract information from a pandas DataFrame.
    If asked about columns, be sure to check the column names first.
    Be concise in your answers.""",
    deps_type=Deps,
    tools=[df_query],
)

Here, we define both the agent and its dependencies:

Dependencies (Deps):

The only dependency required in this example is the DataFrame itself, df.

Agent Definition:

• The deps_type argument links the dependencies to the agent.

• The tools argument ties the df_query function to the agent, making it available for use.

With this setup, the agent is ready to query the DataFrame and handle the tools we’ve defined, bringing PydanticAI’s structured approach to life!

---

Ready, Set, Go… fail.

The first few rows of the DataFrame we’ll be using in the following examples

With everything set up, it’s time to run our agent and see how it performs. Let’s start with a simple question:

res = agent.run_sync(
    'What are the column names in this dataset?',
    deps=Deps(df=df),
)
print(res.new_messages()[-1].content)

The agent responds:

The column names in the dataset are: 

- Date
- Make
- Model
- ...

So far, so good—it got the answer right on the first try! Things are looking promising.

Now, let’s turn up the difficulty just a little:

res = agent.run_sync('How many rows are in this dataset?', ...)
print(res.new_messages()[-1].content)

And then…

------------------------------------------------------------------------
...
ValueError                             Traceback (most recent call last)
...
ValueError: "len" is not a supported function

Oops. Not so great.

What happened here? It looks like the agent tried to use len(df), which is perfectly valid Python syntax but not supported by pd.eval. While pd.eval is useful for safe execution, it only supports a subset of operations. Unfortunately, len isn’t one of them.

Ready, set, go.. Fail! An AI bot not quite succeeding at the task

---

Second Time’s a Charm

There are a few ways we could address the issue with the unsupported syntax. For example, we could tweak the prompt, provide a few usage examples (few-shot learning), or even allow for broader syntax. But instead, let’s take advantage of PydanticAI’s retry mechanism and see if the agent can learn from its mistakes.

Here’s the updated tool definition:

@agent.tool
async def df_query(ctx: RunContext[Deps], query: str) -> str:
    """..."""

    # Debug print - so we can see what the agent is trying to execute.
    print(f'Running query: `{query}`')

    try:
        return str(pd.eval(query, target=ctx.deps.df))
    except Exception as e:
        raise ModelRetry(f'query: `{query}` is not a valid query. Reason: `{e}`') from e

Let’s break it down:

• The @agent.tool Decorator:
  While not directly related to retries, this decorator is a convenient way to define tools. By using @agent.tool, we don’t need to declare the tool separately during the agent definition—it’s automatically connected to the agent.

• Retry Mechanism:
  When an error occurs during pd.eval, we raise a ModelRetry exception. This exception informs the agent that it should retry the operation and provides details about what went wrong. This gives the agent an opportunity to correct its query and try again.

• Debugging with Print Statements:
  We’ve added a print statement to log the queries the agent is trying. This makes it easier to debug and understand the agent’s behavior.

Next, here’s the updated agent definition:

agent = Agent(
    model='openai:gpt-4o-mini',
    system_prompt="""...""",
    deps_type=Deps,
    retries=10,
)

What Changed?

• No More tools Argument: Since we’re using the @agent.tool decorator, the tools argument is no longer necessary. The tool is already connected to the agent via the decorator.

• Retries: We’ve added retries=10, which allows the agent to retry up to 10 times before giving up.

---

Testing the Fix

Let’s see if the retry mechanism works:

res = agent.run_sync('How many rows are in this dataset?', ...)
print(res.new_messages()[-1].content)

The output:

Running query: `len(df)`
Running query: `df.shape[0]`
There are 1000 rows in this dataset.

Voilà! It worked. The agent tried len(df), realized it wasn’t supported, and then corrected itself to df.shape[0].

Now, let’s test it with another question:

res = agent.run_sync('What is the average price of cars sold?', ...)
print(res.new_messages()[-1].content)

And here’s the output:

Running query: `cars['price'].mean()`
Running query: `df['price'].mean()`
Running query: `df.columns`
Running query: `df['Price'].mean()`
The average price of cars sold is approximately $51,145.36.

Once again, we see the agent needed a few retries to get everything right. It initially used the wrong DataFrame name, then the wrong column name. Finally, it queried the column names, adjusted its query, and arrived at the correct answer on its fourth attempt.

AI Finally winning the day and managing to get the correct answers!

---

Wrapping Up

In this post, we explored how to take a LangChain-based data analysis agent and rebuild it using PydanticAI, showcasing the flexibility and power of this new framework. Along the way, we encountered challenges like unsupported syntax in pd.eval, tackled them using PydanticAI’s retry mechanism, and ultimately created an agent that can adapt, learn from its mistakes, and deliver accurate results.

By building our own tools and defining dependencies, we not only replicated the functionality of LangChain but also got a deeper understanding of how PydanticAI operates. This schema-first approach gives you more control, predictable behaviors, and the satisfaction of knowing exactly what’s happening under the hood.

Here’s what we accomplished:

• Custom Tool Creation: We built a tool for querying Pandas DataFrames, using PydanticAI’s dependency injection system to cleanly tie the DataFrame to the agent.

• Retry Logic: By leveraging ModelRetry exceptions, we enabled the agent to gracefully handle errors and refine its queries until it succeeded.

• Agent Definition: We streamlined the agent setup using the @agent.tool decorator and demonstrated how retries can improve performance in ambiguous scenarios.

This journey highlights the differences between LangChain’s plug-and-play ecosystem and PydanticAI’s more hands-on, schema-driven design. The takeaway here isn’t just about choosing between LangChain and PydanticAI—it’s about understanding how different tools can serve different purposes. While LangChain makes it easy to get started with pre-built integrations, PydanticAI shines when you need precise control, structure, and a lightweight framework to build tools that fit your specific needs.

Finally, feel free to explore the full code in this notebook, as well as other code examples, in Nir Diamant’s repository.

If you need help with your own AI or Data Engineering needs, reach out to us at contact@hipposys.com. We’d love to talk with you!

Until next time, happy coding! 🚀