Pandas Dataframe Agent

This notebook shows how to use agents to interact with a pandas dataframe. It is mostly optimized for question answering.

NOTE: this agent calls the Python agent under the hood, which executes LLM generated Python code - this can be bad if the LLM generated Python code is harmful. Use cautiously.

from langchain.agents import create_pandas_dataframe_agent

from langchain.llms import OpenAI
import pandas as pd

df = pd.read_csv('titanic.csv')

agent = create_pandas_dataframe_agent(OpenAI(temperature=0), df, verbose=True)

agent.run("how many rows are there?")

> Entering new AgentExecutor chain...
Thought: I need to count the number of rows
Action: python_repl_ast
Action Input: df.shape[0]
Observation: 891
Thought: I now know the final answer
Final Answer: There are 891 rows.

> Finished chain.

'There are 891 rows.'

agent.run("how many people have more than 3 siblings")

> Entering new AgentExecutor chain...
Thought: I need to count the number of people with more than 3 siblings
Action: python_repl_ast
Action Input: df[df['SibSp'] > 3].shape[0]
Observation: 30
Thought: I now know the final answer
Final Answer: 30 people have more than 3 siblings.

> Finished chain.

'30 people have more than 3 siblings.'

agent.run("whats the square root of the average age?")

> Entering new AgentExecutor chain...
Thought: I need to calculate the average age first
Action: python_repl_ast
Action Input: df['Age'].mean()
Observation: 29.69911764705882
Thought: I now need to calculate the square root of the average age
Action: python_repl_ast
Action Input: math.sqrt(df['Age'].mean())
Observation: NameError("name 'math' is not defined")
Thought: I need to import the math library
Action: python_repl_ast
Action Input: import math
Observation: 
Thought: I now need to calculate the square root of the average age
Action: python_repl_ast
Action Input: math.sqrt(df['Age'].mean())
Observation: 5.449689683556195
Thought: I now know the final answer
Final Answer: The square root of the average age is 5.449689683556195.

> Finished chain.

'The square root of the average age is 5.449689683556195.'

Multi DataFrame Example

This next part shows how the agent can interact with multiple dataframes passed in as a list.

df1 = df.copy()
df1["Age"] = df1["Age"].fillna(df1["Age"].mean())

agent = create_pandas_dataframe_agent(OpenAI(temperature=0), [df, df1], verbose=True)
agent.run("how many rows in the age column are different?")

> Entering new AgentExecutor chain...
Thought: I need to compare the age columns in both dataframes
Action: python_repl_ast
Action Input: len(df1[df1['Age'] != df2['Age']])
Observation: 177
Thought: I now know the final answer
Final Answer: 177 rows in the age column are different.

> Finished chain.

'177 rows in the age column are different.'