From 9e983c9500d52a9e1476917722ad15a63ec56b19 Mon Sep 17 00:00:00 2001 From: Nikita Pokidyshev Date: Thu, 25 Apr 2024 01:51:58 +0300 Subject: [PATCH] langchain[patch]: fix agent_token_buffer_memory not working with openai tools (#20708) - **Description:** fix a bug in the agent_token_buffer_memory - **Issue:** agent_token_buffer_memory was not working with openai tools - **Dependencies:** None - **Twitter handle:** @pokidyshef --- .../agent_token_buffer_memory.py | 13 ++++++++++--- 1 file changed, 10 insertions(+), 3 deletions(-) diff --git a/libs/langchain/langchain/agents/openai_functions_agent/agent_token_buffer_memory.py b/libs/langchain/langchain/agents/openai_functions_agent/agent_token_buffer_memory.py index fe976ddb286..21725ea4faa 100644 --- a/libs/langchain/langchain/agents/openai_functions_agent/agent_token_buffer_memory.py +++ b/libs/langchain/langchain/agents/openai_functions_agent/agent_token_buffer_memory.py @@ -4,8 +4,9 @@ from typing import Any, Dict, List from langchain_core.language_models import BaseLanguageModel from langchain_core.messages import BaseMessage, get_buffer_string -from langchain.agents.format_scratchpad.openai_functions import ( +from langchain.agents.format_scratchpad import ( format_to_openai_function_messages, + format_to_tool_messages, ) from langchain.memory.chat_memory import BaseChatMemory @@ -18,11 +19,12 @@ class AgentTokenBufferMemory(BaseChatMemory): llm: BaseLanguageModel memory_key: str = "history" max_token_limit: int = 12000 - """The max number of tokens to keep in the buffer. + """The max number of tokens to keep in the buffer. Once the buffer exceeds this many tokens, the oldest messages will be pruned.""" return_messages: bool = True output_key: str = "output" intermediate_steps_key: str = "intermediate_steps" + format_as_tools: bool = False @property def buffer(self) -> List[BaseMessage]: @@ -53,7 +55,12 @@ class AgentTokenBufferMemory(BaseChatMemory): """Save context from this conversation to buffer. Pruned.""" input_str, output_str = self._get_input_output(inputs, outputs) self.chat_memory.add_user_message(input_str) - steps = format_to_openai_function_messages(outputs[self.intermediate_steps_key]) + format_to_messages = ( + format_to_tool_messages + if self.format_as_tools + else format_to_openai_function_messages + ) + steps = format_to_messages(outputs[self.intermediate_steps_key]) for msg in steps: self.chat_memory.add_message(msg) self.chat_memory.add_ai_message(output_str)