Update chatbox.py

chatbox.py

@@ -12,7 +12,7 @@ import inspect
 
 from langchain.agents.initialize import initialize_agent
 from langchain.agents.tools import Tool
-from langchain.memory import ConversationBufferMemory
+from langchain.chains.conversation.memory import ConversationBufferMemory
 from langchain_community.chat_models import ChatOpenAI
 import torch
 from PIL import Image, ImageDraw, ImageOps
@@ -141,7 +141,7 @@ class ConversationBot:
     def __init__(self, tools, api_key=""):
         # load_dict = {'VisualQuestionAnswering':'cuda:0', 'ImageCaptioning':'cuda:1',...}
         print("chatbot api",api_key)
-        llm = ChatOpenAI(model_name="gpt-4o", temperature=0.7, openai_api_key=api_key, model_kwargs={"api_version": "2020-11-07"})
+        llm = ChatOpenAI(model_name="gpt-4o", temperature=0.7, openai_api_key=api_key)
         self.llm = llm
         self.memory = ConversationBufferMemory(memory_key="chat_history", output_key='output')
         self.tools = tools
@@ -172,11 +172,17 @@ class ConversationBot:
         return ans
     
     def run_text(self, text, state, aux_state):
-        self.agent.memory.buffer = cut_dialogue_history(self.agent.memory.buffer, keep_last_n_words=500)
+        memory_str = self.agent.memory.buffer_as_str
+        trimmed_memory_str = cut_dialogue_history(memory_str, keep_last_n_words=500)
+        trimmed_messages = self.memory.buffer_as_messages[:len(trimmed_memory_str.split())]
+        # self.agent.memory.buffer = cut_dialogue_history(self.agent.memory.buffer, keep_last_n_words=500)
+        self.memory.chat_memory.messages = trimmed_messages
+        print("done")
         if self.point_prompt != "":
             Human_prompt = f'\nHuman: {self.point_prompt}\n'
             AI_prompt = 'Ok'
-            self.agent.memory.buffer = self.agent.memory.buffer + Human_prompt + 'AI: ' + AI_prompt
+            # self.agent.memory.buffer = self.agent.memory.buffer + Human_prompt + 'AI: ' + AI_prompt
+            self.agent.memory.save_context({'input': Human_prompt}, {'output': AI_prompt})
             self.point_prompt = ""
         res = self.agent({"input": text})
         res['output'] = res['output'].replace("\\", "/")