使用LangChain+本地Llama2处理大CSV数据集QA时的问题求助

问题背景 #

使用本地Llama2模型结合LangChain的create_csv_agent处理14万行、18列的CSV数据集，遇到两个核心问题：

• 请求返回10条VAX_TYPE为COVID19的记录时，仅返回5条
• 查询数据集总行数时得到错误结果5
• 已尝试create_pandas_dataframe_agent、数据嵌入向量数据库方案，结果仍不准确

核心原因及解决方法 #

1. 强制加载全量数据（解决抽样问题） #

LangChain的CSV Agent默认会对大文件进行抽样（通常仅加载前几行），这是导致结果与“5”相关的直接原因。手动加载完整DataFrame再传入Agent，避免默认抽样：

import pandas as pd
from langchain_experimental.agents.agent_toolkits import create_csv_agent

# 手动加载全量CSV，绕过Agent的默认抽样逻辑
df = pd.read_csv("MLdata.csv")

agent = create_csv_agent(
    local_llm,
    df,  # 传入完整DataFrame而非文件路径
    verbose=True,
    agent_type=AgentType.ZERO_SHOT_REACT_DESCRIPTION,
    handle_parsing_errors=True
)

# 测试查询
print(agent.run("请返回10条VAX_TYPE为COVID19的记录。"))
print(agent.run("请告诉我数据集的总行数。"))

2. 自定义提示词强化全量查询要求 #

本地Llama2模型可能对默认提示词理解不足，导致生成的查询逻辑仅针对抽样数据。自定义提示词明确要求基于全量数据处理：

from langchain.prompts import PromptTemplate

custom_prompt = PromptTemplate(
    input_variables=["input", "agent_scratchpad"],
    template="""你是处理大型CSV数据集的专家，必须基于完整数据集执行查询，禁止使用抽样数据。
用户问题：{input}
思考步骤：
1. 确认查询需要覆盖的全量数据范围
2. 生成可执行的Python代码完成查询
3. 执行代码并返回准确结果
思考过程：{agent_scratchpad}
"""
)

agent = create_csv_agent(
    local_llm,
    df,
    verbose=True,
    agent_type=AgentType.ZERO_SHOT_REACT_DESCRIPTION,
    handle_parsing_errors=True,
    prompt=custom_prompt
)

3. 优化Agent类型与参数 #

切换到更适合表格数据的Agent类型，同时调整迭代次数确保复杂查询完成：

from langchain.agents.agent_types import AgentType
from langchain_experimental.agents.agent_toolkits import create_pandas_dataframe_agent

agent = create_pandas_dataframe_agent(
    local_llm,
    df,
    verbose=True,
    agent_type=AgentType.OPENAI_FUNCTIONS,  # 支持函数调用的类型更稳定
    handle_parsing_errors=True,
    max_iterations=20  # 增加迭代次数，避免中途终止
)

4. 先验证数据加载正确性 #

排除CSV文件本身的加载问题（如编码、分隔符错误导致行合并）：

# 手动验证数据完整性
print(f"实际总行数：{len(df)}")
print(f"VAX_TYPE为COVID19的实际记录数：{len(df[df['VAX_TYPE'] == 'COVID19'])}")

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