From 29c0e591fc979277294c2bc8925cd89e24ec9bc8 Mon Sep 17 00:00:00 2001
From: bdqfork <bdqsky@163.com>
Date: Tue, 20 Feb 2024 00:23:03 +0800
Subject: [PATCH] support redis memory

---
 libs/superagent/.env.example            |  5 +++
 libs/superagent/app/agents/langchain.py | 48 ++++++++++++++++++-------
 2 files changed, 40 insertions(+), 13 deletions(-)

diff --git a/libs/superagent/.env.example b/libs/superagent/.env.example
index ef376275b..3e02966e0 100644
--- a/libs/superagent/.env.example
+++ b/libs/superagent/.env.example
@@ -8,6 +8,11 @@ OPENROUTER_API_KEY=
 # Mandatory for Neon DB
 DATABASE_SHADOW_URL= 
 # Memory (mandatory)
+MEMORY=
+# Redis Memory
+REDIS_URL=
+REDIS_MEMORY_WINDOW=
+# Motorhead Memory
 MEMORY_API_URL=https://memory.superagent.sh
 # NOTE: Vectorstores (one is mandatory if you plan on loading datasources)
 VECTORSTORE=pinecone # `qdrant`, `weaviate` etc.
diff --git a/libs/superagent/app/agents/langchain.py b/libs/superagent/app/agents/langchain.py
index 9423dfaa9..94afde800 100644
--- a/libs/superagent/app/agents/langchain.py
+++ b/libs/superagent/app/agents/langchain.py
@@ -6,7 +6,11 @@
 from decouple import config
 from langchain.agents import AgentType, initialize_agent
 from langchain.chains import LLMChain
-from langchain.memory.motorhead_memory import MotorheadMemory
+from langchain.memory import (
+    ConversationBufferWindowMemory,
+    MotorheadMemory,
+    RedisChatMessageHistory,
+)
 from langchain.prompts import MessagesPlaceholder, PromptTemplate
 from langchain.schema import SystemMessage
 from langchain_openai import AzureChatOpenAI, ChatOpenAI
@@ -181,18 +185,36 @@ async def _get_prompt(self, agent: Agent) -> str:
         return SystemMessage(content=content)
 
     async def _get_memory(self) -> List:
-        memory = MotorheadMemory(
-            session_id=(
-                f"{self.agent_id}-{self.session_id}"
-                if self.session_id
-                else f"{self.agent_id}"
-            ),
-            memory_key="chat_history",
-            url=config("MEMORY_API_URL"),
-            return_messages=True,
-            output_key="output",
-        )
-        await memory.init()
+        memory_type = config("MEMORY", "Motorhead")
+        if memory_type == "redis":
+            memory = ConversationBufferWindowMemory(
+                chat_memory=RedisChatMessageHistory(
+                    session_id=(
+                        f"{self.agent_id}-{self.session_id}"
+                        if self.session_id
+                        else f"{self.agent_id}"
+                    ),
+                    url=config("REDIS_URL"),
+                    key_prefix="superagent:",
+                ),
+                memory_key="chat_history",
+                return_messages=True,
+                output_key="output",
+                k=config("REDIS_MEMORY_WINDOW", 10),
+            )
+        else:
+            memory = MotorheadMemory(
+                session_id=(
+                    f"{self.agent_id}-{self.session_id}"
+                    if self.session_id
+                    else f"{self.agent_id}"
+                ),
+                memory_key="chat_history",
+                url=config("MEMORY_API_URL"),
+                return_messages=True,
+                output_key="output",
+            )
+            await memory.init()
         return memory
 
     async def get_agent(self):