如何使Agent对包装前后的LLM输出做一致处理？

问题描述 #

原始代码使用OpenAI LLM初始化Agent：

llm = OpenAI(temperature = 0)
agent = initialize_agent(
    tools,
    llm,
    agent = 'conversational-react-description',
    verbose = True,
    memory = memory,
    return_intermediate_steps = True
)
response = agent(prompt)

运行后得到的response内容为：

Thought: Do I need to use a tool? Yes
Action: Python Evaluator
Action Input: sin(3)/cos(5)
Observation: 0.4974931989239781

为LLM添加自定义包装类后：

class wrapper(LLM):

    llm: langchain.llms.openai.OpenAI
    log: str

    @property
    def _llm_type(self) -> str:
        return "custom"

    def _call(self, prompt: str, stop: Optional[List[str]] = None) -> str:
        reply = self.llm(prompt)
        f = open(self.log, 'a')
        f.write('\n{\n' + prompt + '\n}{\n' + reply + '\n}\n' + '.'*100)
        f.close()
        return reply

    @property
    def _identifying_params(self) -> Mapping[str, Any]:
        """Get the identifying parameters."""
        return {"llm": self.llm}

wllm = wrapper(llm = llm, log = 'output.log')

agent = initialize_agent(
    tools,
    wllm,
    agent = 'conversational-react-description',
    verbose = True,
    memory = memory,
    return_intermediate_steps = True
)
response = agent(prompt)

实际得到的response出现了重复的Observation行：

Thought: Do I need to use a tool? Yes
Action: Python Evaluator
Action Input: sin(3)/cos(5)
Observation: 0.8390075430348345
Observation: 0.4974931989239781

问题原因 #

1. LLM类型不匹配：自定义包装类的_llm_type返回custom，而原生OpenAI LLM的_llm_type是openai。Agent会针对不同LLM类型使用不同的输出解析逻辑，custom类型会触发通用解析规则，容易出现格式识别错误，导致重复记录Observation。
2. 未传递stop参数：包装类的_call方法没有将stop参数传递给底层LLM调用。原生OpenAI LLM会通过stop参数截断输出，确保内容符合Agent预期格式；缺少该参数可能导致LLM返回多余内容，被错误解析为额外的Observation。

解决方法 #

1. 对齐LLM类型标识 #

修改包装类的_llm_type属性，和原生OpenAI LLM保持一致：

@property
def _llm_type(self) -> str:
    return "openai"

2. 正确传递stop参数 #

在_call方法中把stop参数传递给底层LLM调用，保证输出符合Agent的截断规则：

def _call(self, prompt: str, stop: Optional[List[str]] = None) -> str:
    reply = self.llm(prompt, stop=stop)
    f = open(self.log, 'a')
    f.write('\n{\n' + prompt + '\n}{\n' + reply + '\n}\n' + '.'*100)
    f.close()
    return reply

3. 优化日志写入（可选） #

使用上下文管理器替代频繁的文件打开关闭操作，提升稳定性：

def _call(self, prompt: str, stop: Optional[List[str]] = None) -> str:
    reply = self.llm(prompt, stop=stop)
    with open(self.log, 'a') as f:
        f.write('\n{\n' + prompt + '\n}{\n' + reply + '\n}\n' + '.'*100)
    return reply

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