fix: memory in team

l3vels · Nov 29, 2023 · c28dcc1 · c28dcc1
1 parent 728b8bb
commit c28dcc1
Showing 1 changed file with 22 additions and 18 deletions.
diff --git a/apps/server/agents/agent_simulations/agent/dialogue_agent_with_tools.py b/apps/server/agents/agent_simulations/agent/dialogue_agent_with_tools.py
@@ -40,24 +40,25 @@ def send(self) -> str:
         and returns the message string
         """
 
-        memory: ConversationBufferMemory
-
-        if self.is_memory:
-            memory = ZepMemory(
-                session_id=self.session_id,
-                url=Config.ZEP_API_URL,
-                api_key=Config.ZEP_API_KEY,
-                memory_key="chat_history",
-                return_messages=True,
-            )
-
-            memory.human_name = self.sender_name
-            memory.ai_name = self.agent_with_configs.agent.name
-            memory.auto_save = False
-        else:
-            memory = ConversationBufferMemory(
-                memory_key="chat_history", return_messages=True
-            )
+        memory: ZepMemory
+
+        # FIXME: This is a hack to get the memory working
+        # if self.is_memory:
+        memory = ZepMemory(
+            session_id=self.session_id,
+            url=Config.ZEP_API_URL,
+            api_key=Config.ZEP_API_KEY,
+            memory_key="chat_history",
+            return_messages=True,
+        )
+
+        memory.human_name = self.sender_name
+        memory.ai_name = self.agent_with_configs.agent.name
+        memory.auto_save = False
+        # else:
+        #     memory = ConversationBufferMemory(
+        #         memory_key="chat_history", return_messages=True
+        #     )
 
         callbacks = []
 
@@ -83,6 +84,9 @@ def send(self) -> str:
 
         res = agent.run(input=prompt)
 
+        # FIXME: is memory
+        memory.save_ai_message(res)
+
         message = AIMessage(content=res)
 
         return message.content