Merge branch 'dev' into feat/rewrite-query

Author: CaralHsi

src/memos/api/config.py

@@ -1,3 +1,4 @@
+import json
 import os
 
 from typing import Any
@@ -218,6 +219,32 @@ def get_neo4j_shared_config(user_id: str | None = None) -> dict[str, Any]:
             "embedding_dimension": 3072,
         }
 
+    @staticmethod
+    def get_nebular_config(user_id: str | None = None) -> dict[str, Any]:
+        """Get Nebular configuration."""
+        return {
+            "uri": json.loads(os.getenv("NEBULAR_HOSTS", '["localhost"]')),
+            "user": os.getenv("NEBULAR_USER", "root"),
+            "password": os.getenv("NEBULAR_PASSWORD", "xxxxxx"),
+            "space": os.getenv("NEBULAR_SPACE", "shared-tree-textual-memory"),
+            "user_name": f"memos{user_id.replace('-', '')}",
+            "use_multi_db": False,
+            "auto_create": True,
+            "embedding_dimension": 3072,
+        }
+
+    @staticmethod
+    def get_mysql_config() -> dict[str, Any]:
+        """Get MySQL configuration."""
+        return {
+            "host": os.getenv("MYSQL_HOST", "localhost"),
+            "port": int(os.getenv("MYSQL_PORT", "3306")),
+            "username": os.getenv("MYSQL_USERNAME", "root"),
+            "password": os.getenv("MYSQL_PASSWORD", "12345678"),
+            "database": os.getenv("MYSQL_DATABASE", "memos_users"),
+            "charset": os.getenv("MYSQL_CHARSET", "utf8mb4"),
+        }
+
     @staticmethod
     def get_scheduler_config() -> dict[str, Any]:
         """Get scheduler configuration."""
@@ -294,6 +321,7 @@ def get_product_default_config() -> dict[str, Any]:
             "vllm": vllm_config,
         }
         backend = os.getenv("MOS_CHAT_MODEL_PROVIDER", "openai")
+        mysql_config = APIConfig.get_mysql_config()
         config = {
             "user_id": os.getenv("MOS_USER_ID", "root"),
             "chat_model": {"backend": backend, "config": backend_model[backend]},
@@ -330,6 +358,13 @@ def get_product_default_config() -> dict[str, Any]:
         else:
             config["enable_mem_scheduler"] = False
 
+        # Add user manager configuration if enabled
+        if os.getenv("MOS_USER_MANAGER_BACKEND", "sqlite").lower() == "mysql":
+            config["user_manager"] = {
+                "backend": "mysql",
+                "config": mysql_config,
+            }
+
         return config
 
     @staticmethod
@@ -372,6 +407,7 @@ def create_user_config(user_name: str, user_id: str) -> tuple[MOSConfig, General
         openai_config = APIConfig.get_openai_config()
         qwen_config = APIConfig.qwen_config()
         vllm_config = APIConfig.vllm_config()
+        mysql_config = APIConfig.get_mysql_config()
         backend = os.getenv("MOS_CHAT_MODEL_PROVIDER", "openai")
         backend_model = {
             "openai": openai_config,
@@ -417,10 +453,18 @@ def create_user_config(user_name: str, user_id: str) -> tuple[MOSConfig, General
         else:
             config_dict["enable_mem_scheduler"] = False
 
+        # Add user manager configuration if enabled
+        if os.getenv("MOS_USER_MANAGER_BACKEND", "sqlite").lower() == "mysql":
+            config_dict["user_manager"] = {
+                "backend": "mysql",
+                "config": mysql_config,
+            }
+
         default_config = MOSConfig(**config_dict)
 
         neo4j_community_config = APIConfig.get_neo4j_community_config(user_id)
         neo4j_config = APIConfig.get_neo4j_config(user_id)
+        nebular_config = APIConfig.get_nebular_config(user_id)
         internet_config = (
             APIConfig.get_internet_config()
             if os.getenv("ENABLE_INTERNET", "false").lower() == "true"
@@ -429,6 +473,7 @@ def create_user_config(user_name: str, user_id: str) -> tuple[MOSConfig, General
         graph_db_backend_map = {
             "neo4j-community": neo4j_community_config,
             "neo4j": neo4j_config,
+            "nebular": nebular_config,
         }
         graph_db_backend = os.getenv("NEO4J_BACKEND", "neo4j-community").lower()
         if graph_db_backend in graph_db_backend_map:
@@ -475,9 +520,11 @@ def get_default_cube_config() -> GeneralMemCubeConfig | None:
         openai_config = APIConfig.get_openai_config()
         neo4j_community_config = APIConfig.get_neo4j_community_config(user_id="default")
         neo4j_config = APIConfig.get_neo4j_config(user_id="default")
+        nebular_config = APIConfig.get_nebular_config(user_id="default")
         graph_db_backend_map = {
             "neo4j-community": neo4j_community_config,
             "neo4j": neo4j_config,
+            "nebular": nebular_config,
         }
         internet_config = (
             APIConfig.get_internet_config()

src/memos/api/product_models.py

@@ -88,7 +88,7 @@ class ChatRequest(BaseRequest):
 
 class UserCreate(BaseRequest):
     user_name: str | None = Field(None, description="Name of the user")
-    role: str = Field("user", description="Role of the user")
+    role: str = Field("USER", description="Role of the user")
     user_id: str = Field(..., description="User ID")
 
 
@@ -143,6 +143,8 @@ class MemoryCreateRequest(BaseRequest):
     memory_content: str | None = Field(None, description="Memory content to store")
     doc_path: str | None = Field(None, description="Path to document to store")
     mem_cube_id: str | None = Field(None, description="Cube ID")
+    source: str | None = Field(None, description="Source of the memory")
+    user_profile: bool = Field(False, description="User profile memory")
 
 
 class SearchRequest(BaseRequest):

src/memos/api/routers/product_router.py

@@ -198,6 +198,8 @@ def create_memory(memory_req: MemoryCreateRequest):
             messages=memory_req.messages,
             doc_path=memory_req.doc_path,
             mem_cube_id=memory_req.mem_cube_id,
+            source=memory_req.source,
+            user_profile=memory_req.user_profile,
         )
         return SimpleResponse(message="Memory created successfully")
 

src/memos/mem_cube/utils.py

@@ -92,9 +92,20 @@ def merge_config_with_default(
                     merged_graph_config["db_name"] = default_graph_config.get("db_name")
                 else:
                     logger.info("use_multi_db is already False, no need to change")
-
+            if "neo4j" not in default_text_config["graph_db"]["backend"]:
+                if "db_name" in merged_graph_config:
+                    merged_graph_config.pop("db_name")
+                    logger.info("neo4j is not supported, remove db_name")
+                else:
+                    logger.info("db_name is not in merged_graph_config, no need to remove")
+            else:
+                if "space" in merged_graph_config:
+                    merged_graph_config.pop("space")
+                    logger.info("neo4j is not supported, remove db_name")
+                else:
+                    logger.info("space is not in merged_graph_config, no need to remove")
             preserved_graph_db = {
-                "backend": existing_text_config["graph_db"]["backend"],
+                "backend": default_text_config["graph_db"]["backend"],
                 "config": merged_graph_config,
             }
 

src/memos/mem_os/core.py

@@ -1,5 +1,7 @@
 import json
 import os
+import time
+import uuid
 
 from datetime import datetime
 from pathlib import Path
@@ -579,6 +581,7 @@ def search(
                 and (mem_cube.text_mem is not None)
                 and self.config.enable_textual_memory
             ):
+                time_start = time.time()
                 memories = mem_cube.text_mem.search(
                     query,
                     top_k=top_k if top_k else self.config.top_k,
@@ -594,6 +597,10 @@ def search(
                 logger.info(
                     f"🧠 [Memory] Searched memories from {mem_cube_id}:\n{self._str_memories(memories)}\n"
                 )
+                search_time_end = time.time()
+                logger.info(
+                    f"time search graph: search graph time user_id: {target_user_id} time is: {search_time_end - time_start}"
+                )
         return result
 
     def add(