如何让Langchain Pandas Agent过滤并返回/保存DataFrame结果？

问题分析 #

默认的create_pandas_dataframe_agent在处理保存文件或返回DataFrame实例的请求时，倾向于返回方法说明而非实际执行操作——这是因为agent的默认prompt更偏向解释而非执行代码，且run()方法默认返回文本结果而非对象。

1. 让Agent实际保存DataFrame到CSV #

需要明确指令让agent生成并执行保存代码，而非仅解释方法。示例代码：

# 明确要求智能体执行保存操作
agent.run("编写并执行Python代码，将df_tmp保存为名为'output.csv'的CSV文件")

执行后，当前目录会生成output.csv文件，可直接用pd.read_csv('output.csv')读取。

2. 获取过滤后的DataFrame实例 #

由于agent.run()默认返回文本，无法直接返回DataFrame对象，可通过以下两种方法实现：

方法1：提取并执行Agent生成的过滤代码 #

让agent生成过滤数据的Python代码，提取后自行执行得到DataFrame：

# 让智能体仅返回过滤代码，不添加解释
code_response = agent.run("生成Python代码，过滤df_tmp中name为'Pandas'的行，将结果存入df_filtered。只输出代码，不要解释。")

# 执行代码
exec(code_response)

# 验证结果
print(len(df_filtered))  # 输出：1

方法2：自定义工具返回JSON格式结果 #

创建自定义工具，让agent调用工具返回过滤后的DataFrame的JSON字符串，再解析为DataFrame：

from langchain.tools import tool

@tool
def filter_dataframe(filter_condition: str) -> str:
    """
    根据给定的条件字符串过滤df_tmp。
    条件需为有效的pandas布尔索引表达式，例如："name == 'Pandas'"
    返回过滤后的DataFrame的JSON字符串。
    """
    filtered_df = df_tmp.query(filter_condition)
    return filtered_df.to_json(orient='records')

# 重新创建agent，添加自定义工具
agent = create_pandas_dataframe_agent(
    ChatOpenAI(temperature=0, model="gpt-3.5-turbo-0613"),
    df_tmp,
    verbose=True,
    agent_type=AgentType.OPENAI_FUNCTIONS,
    extra_tools=[filter_dataframe]
)

# 调用工具获取结果并解析
json_response = agent.run("使用filter_dataframe工具，过滤df_tmp中name为'Pandas'的行，返回JSON结果")
df_filtered = pd.read_json(json_response, orient='records')

print(len(df_filtered))  # 输出：1

内容的提问来源于stack exchange，提问作者baswenneker