当OpenAI模型max_tokens超限，如何获取LangChain Agent中间步骤？

问题场景 #

使用LangChain的create_pandas_dataframe_agent构建数据查询Agent时，Agent生成的Python代码逻辑正确，但执行后返回的DataFrame结果过大，超出OpenAI模型的token上限，触发错误后无法获取中间步骤里的Action Input（即生成的Python代码）。常规的callbacks和intermediate steps方法在这种场景下失效，因为token超限会导致agent.run()中断，无法完整保留中间步骤。

解决方案 #

方案1：自定义Python执行工具，提前保存代码 #

替换默认的python_repl_ast工具，在执行代码前先将生成的代码保存到变量中，无论后续结果是否超限，代码都会被留存。

from langchain.llms import OpenAI
from langchain.agents import create_pandas_dataframe_agent, Tool
from langchain.tools.python.tool import PythonAstREPLTool
import pandas as pd

# 用于存储生成的Python代码
generated_code = ""

def custom_python_exec(input: str) -> str:
    global generated_code
    # 保存生成的代码
    generated_code = input
    # 调用原工具逻辑执行代码
    return PythonAstREPLTool(locals={"df": df}).run(input)

# 定义自定义工具，替换默认工具
custom_tool = Tool(
    name="python_repl_ast",
    func=custom_python_exec,
    description="A Python shell. Use this to execute python commands. Input should be a valid python command. When using this tool, sometimes output is abbreviated - make sure it does not look abbreviated before using it in your answer."
)

# 加载数据并初始化Agent
df = pd.read_excel("your_data.xlsx")
agent = create_pandas_dataframe_agent(
    OpenAI(temperature=0, model_kwargs={"model": 'text-davinci-003'}),
    df,
    verbose=True,
    max_iterations=3,
    tools=[custom_tool]
)

# 构造返回DataFrame的查询
query = "查询2003年美国地区摩托车产品线的所有数据"

try:
    agent.run(query)
except Exception as e:
    # 即使触发token超限错误，也能获取到生成的代码
    print("生成的Python代码：")
    print(generated_code)

方案2：用自定义Callback捕获Action Input #

利用LangChain的Callback机制，在Agent执行Action的阶段就捕获代码，不需要等到结果返回，避免因后续错误导致代码丢失。

from langchain.llms import OpenAI
from langchain.agents import create_pandas_dataframe_agent
from langchain.callbacks.base import BaseCallbackHandler
import pandas as pd

class CodeCaptureCallback(BaseCallbackHandler):
    def __init__(self):
        self.captured_code = ""

    def on_agent_action(self, action, **kwargs):
        # 当工具为python_repl_ast时，记录Action Input
        if action.tool == "python_repl_ast":
            self.captured_code = action.tool_input

# 加载数据并初始化回调
df = pd.read_excel("your_data.xlsx")
code_callback = CodeCaptureCallback()

# 初始化Agent并绑定回调
agent = create_pandas_dataframe_agent(
    OpenAI(temperature=0, model_kwargs={"model": 'text-davinci-003'}),
    df,
    verbose=True,
    max_iterations=3,
    callbacks=[code_callback]
)

# 构造查询
query = "查询2003年美国地区摩托车产品线的所有数据"

try:
    agent.run(query)
except Exception as e:
    print("生成的Python代码：")
    print(code_callback.captured_code)

方案3：截断结果避免token超限，正常获取中间步骤 #

修改工具的执行结果，将过大的DataFrame截断为预览内容，减少返回的token数，让agent.run()能正常完成，从而从intermediate_steps中提取代码。

from langchain.llms import OpenAI
from langchain.agents import create_pandas_dataframe_agent, Tool
from langchain.tools.python.tool import PythonAstREPLTool
import pandas as pd

def truncated_python_exec(input: str) -> str:
    result = PythonAstREPLTool(locals={"df": df}).run(input)
    # 对DataFrame结果进行截断处理
    if isinstance(result, pd.DataFrame):
        if len(result) > 10:
            return f"数据预览（前10行）：\n{result.head(10)}\n\n总数据行数：{len(result)}"
    return result

# 定义截断后的工具
custom_tool = Tool(
    name="python_repl_ast",
    func=truncated_python_exec,
    description="A Python shell. Use this to execute python commands. Input should be a valid python command. When using this tool, sometimes output is abbreviated - make sure it does not look abbreviated before using it in your answer."
)

# 初始化Agent
df = pd.read_excel("your_data.xlsx")
agent = create_pandas_dataframe_agent(
    OpenAI(temperature=0, model_kwargs={"model": 'text-davinci-003'}),
    df,
    verbose=True,
    max_iterations=3,
    tools=[custom_tool]
)

# 执行查询
query = "查询2003年美国地区摩托车产品线的所有数据"
agent.run(query)

# 从中间步骤提取生成的代码
for step in agent.intermediate_steps:
    if step.action.tool == "python_repl_ast":
        print("生成的Python代码：")
        print(step.action.tool_input)

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