support redis memory

bdqfork · Feb 19, 2024 · 8ec71ad · 8ec71ad
1 parent 276541b
commit 8ec71ad
Show file tree

Hide file tree

Showing 4 changed files with 66 additions and 20 deletions.
diff --git 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
@@ -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", "redis://localhost:6379/0"),
+                    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):

diff --git a/libs/superagent/poetry.lock b/libs/superagent/poetry.lock
diff --git a/libs/superagent/pyproject.toml b/libs/superagent/pyproject.toml
@@ -63,6 +63,7 @@ langchain-openai = "^0.0.5"
 python-docx = "^1.1.0"
 prisma = "^0.12.0"
 stripe = "^8.2.0"
+redis = "^5.0.1"