Support set Memory from UI

libs/superagent/Dockerfile

@@ -31,11 +31,11 @@ ENV PORT="8080"
 
 COPY --from=builder /app/.venv /app/.venv
 
-COPY . ./
-
 # Improve grpc error messages
 RUN pip install grpcio-status
 
+COPY . ./
+
 # Enable prisma migrations
 RUN prisma generate
 

libs/superagent/app/agents/base.py

@@ -3,7 +3,7 @@
 from app.models.request import LLMParams
 from app.utils.callbacks import CustomAsyncIteratorCallbackHandler
 from prisma.enums import AgentType
-from prisma.models import Agent
+from prisma.models import Agent, MemoryDb
 
 DEFAULT_PROMPT = (
     "You are a helpful AI Assistant, answer the users questions to "
@@ -21,6 +21,7 @@ def __init__(
         callbacks: List[CustomAsyncIteratorCallbackHandler] = [],
         llm_params: Optional[LLMParams] = {},
         agent_config: Agent = None,
+        memory_config: Optional[MemoryDb] = None,
     ):
         self.agent_id = agent_id
         self.session_id = session_id
@@ -29,6 +30,7 @@ def __init__(
         self.callbacks = callbacks
         self.llm_params = llm_params
         self.agent_config = agent_config
+        self.memory_config = memory_config
 
     async def _get_tools(
         self,
@@ -60,6 +62,7 @@ async def get_agent(self):
                 callbacks=self.callbacks,
                 llm_params=self.llm_params,
                 agent_config=self.agent_config,
+                memory_config=self.memory_config,
             )
 
         elif self.agent_config.type == AgentType.LLM:
@@ -72,6 +75,7 @@ async def get_agent(self):
                 callbacks=self.callbacks,
                 llm_params=self.llm_params,
                 agent_config=self.agent_config,
+                memory_config=self.memory_config,
             )
 
         else:
@@ -85,6 +89,7 @@ async def get_agent(self):
                 callbacks=self.callbacks,
                 llm_params=self.llm_params,
                 agent_config=self.agent_config,
+                memory_config=self.memory_config,
             )
 
         return await agent.get_agent()

libs/superagent/app/agents/langchain.py

@@ -1,7 +1,8 @@
 import datetime
 import json
+import logging
 import re
-from typing import Any, List
+from typing import Any, List, Optional
 
 from decouple import config
 from langchain.agents import AgentType, initialize_agent
@@ -23,8 +24,12 @@
 from app.models.tools import DatasourceInput
 from app.tools import TOOL_TYPE_MAPPING, create_pydantic_model_from_object, create_tool
 from app.tools.datasource import DatasourceTool, StructuredDatasourceTool
+from app.utils.helpers import get_first_non_null
 from app.utils.llm import LLM_MAPPING
-from prisma.models import LLM, Agent, AgentDatasource, AgentTool
+from prisma.enums import LLMProvider, MemoryDbProvider
+from prisma.models import LLM, Agent, AgentDatasource, AgentTool, MemoryDb
+
+logger = logging.getLogger(__name__)
 
 DEFAULT_PROMPT = (
     "You are a helpful AI Assistant, answer the users questions to "
@@ -148,7 +153,7 @@ async def _get_llm(self, llm: LLM, model: str) -> Any:
             **(self.llm_params.dict() if self.llm_params else {}),
         }
 
-        if llm.provider == "OPENAI":
+        if llm.provider == LLMProvider.OPENAI:
             return ChatOpenAI(
                 model=LLM_MAPPING[model],
                 openai_api_key=llm.apiKey,
@@ -157,7 +162,7 @@ async def _get_llm(self, llm: LLM, model: str) -> Any:
                 **(llm.options if llm.options else {}),
                 **(llm_params),
             )
-        elif llm.provider == "AZURE_OPENAI":
+        elif llm.provider == LLMProvider.AZURE_OPENAI:
             return AzureChatOpenAI(
                 api_key=llm.apiKey,
                 streaming=self.enable_streaming,
@@ -193,33 +198,59 @@ async def _get_prompt(self, agent: Agent) -> str:
             content = f"{content}" f"\n\n{datetime.datetime.now().strftime('%Y-%m-%d')}"
         return SystemMessage(content=content)
 
-    async def _get_memory(self) -> List:
-        memory_type = config("MEMORY", "motorhead")
-        if memory_type == "redis":
+    async def _get_memory(self, memory_db: Optional[MemoryDb]) -> List:
+        logger.debug(f"Use memory config: {memory_db}")
+        if memory_db is None:
+            memory_provider = config("MEMORY", "motorhead")
+            options = {}
+        else:
+            memory_provider = memory_db.provider
+            options = memory_db.options
+
+        memory_provider = memory_provider.upper()
+        logger.info(f"Using memory provider: {memory_provider}")
+
+        if memory_provider == MemoryDbProvider.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_MEMORY_URL", "redis://localhost:6379/0"),
+                    url=get_first_non_null(
+                        options.get("REDIS_MEMORY_URL"),
+                        config("REDIS_MEMORY_URL", "redis://localhost:6379/0"),
+                    ),
                     key_prefix="superagent:",
                 ),
                 memory_key="chat_history",
                 return_messages=True,
                 output_key="output",
-                k=config("REDIS_MEMORY_WINDOW", 10),
+                k=get_first_non_null(
+                    options.get("REDIS_MEMORY_WINDOW"),
+                    config("REDIS_MEMORY_WINDOW", 10),
+                ),
             )
-        else:
+        elif memory_provider == MemoryDbProvider.MOTORHEAD:
+            url = get_first_non_null(
+                options.get("MEMORY_API_URL"),
+                config("MEMORY_API_URL"),
+            )
+
+            if not url:
+                raise ValueError(
+                    "Memory API URL is required for Motorhead memory provider"
+                )
+
             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"),
+                url=url,
                 return_messages=True,
                 output_key="output",
             )
@@ -235,7 +266,7 @@ async def get_agent(self):
             agent_tools=self.agent_config.tools,
         )
         prompt = await self._get_prompt(agent=self.agent_config)
-        memory = await self._get_memory()
+        memory = await self._get_memory(memory_db=self.memory_config)
 
         if len(tools) > 0:
             agent = initialize_agent(

libs/superagent/app/api/agents.py

@@ -452,6 +452,10 @@ async def invoke(
     if not model and metadata.get("model"):
         model = metadata.get("model")
 
+    memory_config = await prisma.memorydb.find_first(
+        where={"provider": agent_config.memory, "apiUserId": api_user.id},
+    )
+
     def track_agent_invocation(result):
         intermediate_steps_to_obj = [
             {
@@ -571,6 +575,7 @@ async def send_message(
         callbacks=monitoring_callbacks,
         llm_params=body.llm_params,
         agent_config=agent_config,
+        memory_config=memory_config,
     )
     agent = await agent_base.get_agent()
 

libs/superagent/app/api/memory_dbs.py

@@ -0,0 +1,104 @@
+import json
+
+import segment.analytics as analytics
+from decouple import config
+from fastapi import APIRouter, Depends
+
+from app.models.request import MemoryDb as MemoryDbRequest
+from app.models.response import MemoryDb as MemoryDbResponse
+from app.models.response import MemoryDbList as MemoryDbListResponse
+from app.utils.api import get_current_api_user, handle_exception
+from app.utils.prisma import prisma
+from prisma import Json
+
+SEGMENT_WRITE_KEY = config("SEGMENT_WRITE_KEY", None)
+
+router = APIRouter()
+analytics.write_key = SEGMENT_WRITE_KEY
+
+
+@router.post(
+    "/memory-db",
+    name="create",
+    description="Create a new Memory Database",
+    response_model=MemoryDbResponse,
+)
+async def create(body: MemoryDbRequest, api_user=Depends(get_current_api_user)):
+    """Endpoint for creating a Memory Database"""
+    if SEGMENT_WRITE_KEY:
+        analytics.track(api_user.id, "Created Memory Database")
+
+    data = await prisma.memorydb.create(
+        {
+            **body.dict(),
+            "apiUserId": api_user.id,
+            "options": json.dumps(body.options),
+        }
+    )
+    data.options = json.dumps(data.options)
+    return {"success": True, "data": data}
+
+
+@router.get(
+    "/memory-dbs",
+    name="list",
+    description="List all Memory Databases",
+    response_model=MemoryDbListResponse,
+)
+async def list(api_user=Depends(get_current_api_user)):
+    """Endpoint for listing all Memory Databases"""
+    try:
+        data = await prisma.memorydb.find_many(
+            where={"apiUserId": api_user.id}, order={"createdAt": "desc"}
+        )
+        # Convert options to string
+        for item in data:
+            item.options = json.dumps(item.options)
+        return {"success": True, "data": data}
+    except Exception as e:
+        handle_exception(e)
+
+
+@router.get(
+    "/memory-dbs/{memory_db_id}",
+    name="get",
+    description="Get a single Memory Database",
+    response_model=MemoryDbResponse,
+)
+async def get(memory_db_id: str, api_user=Depends(get_current_api_user)):
+    """Endpoint for getting a single Memory Database"""
+    try:
+        data = await prisma.memorydb.find_first(
+            where={"id": memory_db_id, "apiUserId": api_user.id}
+        )
+        data.options = json.dumps(data.options)
+        return {"success": True, "data": data}
+    except Exception as e:
+        handle_exception(e)
+
+
+@router.patch(
+    "/memory-dbs/{memory_db_id}",
+    name="update",
+    description="Patch a Memory Database",
+    response_model=MemoryDbResponse,
+)
+async def update(
+    memory_db_id: str, body: MemoryDbRequest, api_user=Depends(get_current_api_user)
+):
+    """Endpoint for patching a Memory Database"""
+    try:
+        if SEGMENT_WRITE_KEY:
+            analytics.track(api_user.id, "Updated Memory Database")
+        data = await prisma.memorydb.update(
+            where={"id": memory_db_id},
+            data={
+                **body.dict(exclude_unset=True),
+                "apiUserId": api_user.id,
+                "options": Json(body.options),
+            },
+        )
+        data.options = json.dumps(data.options)
+        return {"success": True, "data": data}
+    except Exception as e:
+        handle_exception(e)

libs/superagent/app/main.py

@@ -1,4 +1,5 @@
 import logging
+import os
 import time
 
 import colorlog
@@ -26,7 +27,7 @@
 console_handler.setFormatter(formatter)
 
 logging.basicConfig(
-    level=logging.INFO,
+    level=os.environ.get("LOG_LEVEL", "INFO"),
     format="%(levelname)s: %(message)s",
     handlers=[console_handler],
     force=True,

libs/superagent/app/models/request.py

@@ -3,7 +3,7 @@
 from openai.types.beta.assistant_create_params import Tool as OpenAiAssistantTool
 from pydantic import BaseModel
 
-from prisma.enums import AgentType, LLMProvider, VectorDbProvider
+from prisma.enums import AgentType, LLMProvider, MemoryDbProvider, VectorDbProvider
 
 
 class ApiUser(BaseModel):
@@ -40,6 +40,7 @@ class AgentUpdate(BaseModel):
     initialMessage: Optional[str]
     prompt: Optional[str]
     llmModel: Optional[str]
+    memory: Optional[str]
     description: Optional[str]
     avatar: Optional[str]
     type: Optional[str]
@@ -132,3 +133,8 @@ class WorkflowInvoke(BaseModel):
 class VectorDb(BaseModel):
     provider: VectorDbProvider
     options: Dict
+
+
+class MemoryDb(BaseModel):
+    provider: MemoryDbProvider
+    options: Dict

libs/superagent/app/models/response.py

@@ -20,6 +20,9 @@
 from prisma.models import (
     Datasource as DatasourceModel,
 )
+from prisma.models import (
+    MemoryDb as MemoryDbModel,
+)
 from prisma.models import (
     Tool as ToolModel,
 )
@@ -141,3 +144,13 @@ class VectorDb(BaseModel):
 class VectorDbList(BaseModel):
     success: bool
     data: Optional[List[VectorDbModel]]
+
+
+class MemoryDb(BaseModel):
+    success: bool
+    data: Optional[MemoryDbModel]
+
+
+class MemoryDbList(BaseModel):
+    success: bool
+    data: Optional[List[MemoryDbModel]]

libs/superagent/app/routers.py

@@ -5,6 +5,7 @@
     api_user,
     datasources,
     llms,
+    memory_dbs,
     tools,
     vector_dbs,
     workflows,
@@ -24,3 +25,4 @@
     workflow_configs.router, tags=["Workflow Config"], prefix=api_prefix
 )
 router.include_router(vector_dbs.router, tags=["Vector Database"], prefix=api_prefix)
+router.include_router(memory_dbs.router, tags=["Memory Database"], prefix=api_prefix)