Merge branch 'dev' into feat/demo-reference-change

Author: fridayL

examples/mem_os/multi_user_memos_example.py

@@ -2,6 +2,8 @@
 Example demonstrating how to use MOSProduct for multi-user scenarios.
 """
 
+import os
+
 from memos.configs.mem_cube import GeneralMemCubeConfig
 from memos.configs.mem_os import MOSConfig
 from memos.mem_cube.general import GeneralMemCube
@@ -16,28 +18,53 @@ def get_config(user_name):
         "top_p": 0.9,
         "top_k": 50,
         "remove_think_prefix": True,
-        "api_key": "your-api-key-here",
-        "api_base": "https://api.openai.com/v1",
+        "api_key": os.getenv("OPENAI_API_KEY"),
+        "api_base": os.getenv("OPENAI_API_BASE"),
     }
     # Create a default configuration
     default_config = MOSConfig(
         user_id="root",
         chat_model={"backend": "openai", "config": openapi_config},
         mem_reader={
-            "backend": "naive",
+            "backend": "simple_struct",
             "config": {
                 "llm": {
                     "backend": "openai",
                     "config": openapi_config,
                 },
                 "embedder": {
-                    "backend": "ollama",
+                    "backend": "universal_api",
+                    "config": {
+                        "provider": os.getenv("MOS_EMBEDDER_PROVIDER", "openai"),
+                        "api_key": os.getenv("MOS_EMBEDDER_API_KEY", "sk-xxxx"),
+                        "model_name_or_path": os.getenv(
+                            "MOS_EMBEDDER_MODEL", "text-embedding-3-large"
+                        ),
+                        "base_url": os.getenv("MOS_EMBEDDER_API_BASE", "http://openai.com"),
+                    },
+                },
+                "chunker": {
+                    "backend": "sentence",
                     "config": {
-                        "model_name_or_path": "nomic-embed-text:latest",
+                        "tokenizer_or_token_counter": "gpt2",
+                        "chunk_size": 512,
+                        "chunk_overlap": 128,
+                        "min_sentences_per_chunk": 1,
                     },
                 },
             },
         },
+        user_manager={
+            "backend": "mysql",
+            "config": {
+                "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"),
+            },
+        },
         enable_textual_memory=True,
         enable_activation_memory=False,
         top_k=5,
@@ -55,17 +82,27 @@ def get_config(user_name):
                     "graph_db": {
                         "backend": "neo4j",
                         "config": {
-                            "uri": "bolt://localhost:7687",
-                            "user": "neo4j",
-                            "password": "12345678",
-                            "db_name": user_name,
+                            "uri": os.getenv("NEO4J_URI", "bolt://localhost:7687"),
+                            "user": os.getenv("NEO4J_USER", "neo4j"),
+                            "password": os.getenv("NEO4J_PASSWORD", "12345678"),
+                            "db_name": os.getenv(
+                                "NEO4J_DB_NAME", "shared-tree-textual-memory-test"
+                            ),
+                            "user_name": f"memos{user_name.replace('-', '')}",
+                            "embedding_dimension": int(os.getenv("EMBEDDING_DIMENSION", 768)),
+                            "use_multi_db": False,
                             "auto_create": True,
                         },
                     },
                     "embedder": {
-                        "backend": "ollama",
+                        "backend": "universal_api",
                         "config": {
-                            "model_name_or_path": "nomic-embed-text:latest",
+                            "provider": os.getenv("MOS_EMBEDDER_PROVIDER", "openai"),
+                            "api_key": os.getenv("MOS_EMBEDDER_API_KEY", "sk-xxxx"),
+                            "model_name_or_path": os.getenv(
+                                "MOS_EMBEDDER_MODEL", "text-embedding-3-large"
+                            ),
+                            "base_url": os.getenv("MOS_EMBEDDER_API_BASE", "http://openai.com"),
                         },
                     },
                 },
@@ -109,7 +146,7 @@ def main():
     print(f"\nSearch result for Alice: {search_result}")
 
     # Search memories for Alice
-    search_result = mos_product.get_all(query="conference", user_id="alice", memory_type="text_mem")
+    search_result = mos_product.get_all(user_id="alice", memory_type="text_mem")
     print(f"\nSearch result for Alice: {search_result}")
 
     # List all users

src/memos/graph_dbs/nebular.py

@@ -1,3 +1,4 @@
+import json
 import traceback
 
 from contextlib import suppress
@@ -35,7 +36,28 @@ def _compose_node(item: dict[str, Any]) -> tuple[str, str, dict[str, Any]]:
 
 @timed
 def _escape_str(value: str) -> str:
-    return value.replace('"', '\\"')
+    out = []
+    for ch in value:
+        code = ord(ch)
+        if ch == "\\":
+            out.append("\\\\")
+        elif ch == '"':
+            out.append('\\"')
+        elif ch == "\n":
+            out.append("\\n")
+        elif ch == "\r":
+            out.append("\\r")
+        elif ch == "\t":
+            out.append("\\t")
+        elif ch == "\b":
+            out.append("\\b")
+        elif ch == "\f":
+            out.append("\\f")
+        elif code < 0x20 or code in (0x2028, 0x2029):
+            out.append(f"\\u{code:04x}")
+        else:
+            out.append(ch)
+    return "".join(out)
 
 
 @timed
@@ -1153,28 +1175,36 @@ def import_graph(self, data: dict[str, Any]) -> None:
             data: A dictionary containing all nodes and edges to be loaded.
         """
         for node in data.get("nodes", []):
-            id, memory, metadata = _compose_node(node)
+            try:
+                id, memory, metadata = _compose_node(node)
 
-            if not self.config.use_multi_db and self.config.user_name:
-                metadata["user_name"] = self.config.user_name
+                if not self.config.use_multi_db and self.config.user_name:
+                    metadata["user_name"] = self.config.user_name
 
-            metadata = self._prepare_node_metadata(metadata)
-            metadata.update({"id": id, "memory": memory})
-            properties = ", ".join(f"{k}: {self._format_value(v, k)}" for k, v in metadata.items())
-            node_gql = f"INSERT OR IGNORE (n@Memory {{{properties}}})"
-            self.execute_query(node_gql)
+                metadata = self._prepare_node_metadata(metadata)
+                metadata.update({"id": id, "memory": memory})
+                properties = ", ".join(
+                    f"{k}: {self._format_value(v, k)}" for k, v in metadata.items()
+                )
+                node_gql = f"INSERT OR IGNORE (n@Memory {{{properties}}})"
+                self.execute_query(node_gql)
+            except Exception as e:
+                logger.error(f"Fail to load node: {node}, error: {e}")
 
         for edge in data.get("edges", []):
-            source_id, target_id = edge["source"], edge["target"]
-            edge_type = edge["type"]
-            props = ""
-            if not self.config.use_multi_db and self.config.user_name:
-                props = f'{{user_name: "{self.config.user_name}"}}'
-            edge_gql = f'''
-               MATCH (a@Memory {{id: "{source_id}"}}), (b@Memory {{id: "{target_id}"}})
-               INSERT OR IGNORE (a) -[e@{edge_type} {props}]-> (b)
-           '''
-            self.execute_query(edge_gql)
+            try:
+                source_id, target_id = edge["source"], edge["target"]
+                edge_type = edge["type"]
+                props = ""
+                if not self.config.use_multi_db and self.config.user_name:
+                    props = f'{{user_name: "{self.config.user_name}"}}'
+                edge_gql = f'''
+                   MATCH (a@Memory {{id: "{source_id}"}}), (b@Memory {{id: "{target_id}"}})
+                   INSERT OR IGNORE (a) -[e@{edge_type} {props}]-> (b)
+               '''
+                self.execute_query(edge_gql)
+            except Exception as e:
+                logger.error(f"Fail to load edge: {edge}, error: {e}")
 
     @timed
     def get_all_memory_items(self, scope: str, include_embedding: bool = False) -> (list)[dict]:
@@ -1555,6 +1585,7 @@ def _prepare_node_metadata(self, metadata: dict[str, Any]) -> dict[str, Any]:
         # Normalize embedding type
         embedding = metadata.get("embedding")
         if embedding and isinstance(embedding, list):
+            metadata.pop("embedding")
             metadata[self.dim_field] = _normalize([float(x) for x in embedding])
 
         return metadata
@@ -1563,26 +1594,41 @@ def _prepare_node_metadata(self, metadata: dict[str, Any]) -> dict[str, Any]:
     def _format_value(self, val: Any, key: str = "") -> str:
         from nebulagraph_python.py_data_types import NVector
 
+        # None
+        if val is None:
+            return "NULL"
+        # bool
+        if isinstance(val, bool):
+            return "true" if val else "false"
+        # str
         if isinstance(val, str):
             return f'"{_escape_str(val)}"'
+        # num
         elif isinstance(val, (int | float)):
             return str(val)
+        # time
         elif isinstance(val, datetime):
             return f'datetime("{val.isoformat()}")'
+        # list
         elif isinstance(val, list):
             if key == self.dim_field:
                 dim = len(val)
                 joined = ",".join(str(float(x)) for x in val)
                 return f"VECTOR<{dim}, FLOAT>([{joined}])"
             else:
                 return f"[{', '.join(self._format_value(v) for v in val)}]"
+        # NVector
         elif isinstance(val, NVector):
             if key == self.dim_field:
                 dim = len(val)
                 joined = ",".join(str(float(x)) for x in val)
                 return f"VECTOR<{dim}, FLOAT>([{joined}])"
-        elif val is None:
-            return "NULL"
+            else:
+                logger.warning("Invalid NVector")
+        # dict
+        if isinstance(val, dict):
+            j = json.dumps(val, ensure_ascii=False, separators=(",", ":"))
+            return f'"{_escape_str(j)}"'
         else:
             return f'"{_escape_str(str(val))}"'
 

src/memos/graph_dbs/neo4j.py

@@ -323,12 +323,11 @@ def edge_exists(
             return result.single() is not None
 
     # Graph Query & Reasoning
-    def get_node(self, id: str, include_embedding: bool = True) -> dict[str, Any] | None:
+    def get_node(self, id: str, **kwargs) -> dict[str, Any] | None:
         """
         Retrieve the metadata and memory of a node.
         Args:
             id: Node identifier.
-            include_embedding (bool): Whether to include the large embedding field.
         Returns:
             Dictionary of node fields, or None if not found.
         """
@@ -345,12 +344,11 @@ def get_node(self, id: str, include_embedding: bool = True) -> dict[str, Any] |
             record = session.run(query, params).single()
             return self._parse_node(dict(record["n"])) if record else None
 
-    def get_nodes(self, ids: list[str], include_embedding: bool = True) -> list[dict[str, Any]]:
+    def get_nodes(self, ids: list[str], **kwargs) -> list[dict[str, Any]]:
         """
         Retrieve the metadata and memory of a list of nodes.
         Args:
             ids: List of Node identifier.
-            include_embedding (bool): Whether to include the large embedding field.
         Returns:
         list[dict]: Parsed node records containing 'id', 'memory', and 'metadata'.
 
@@ -833,7 +831,7 @@ def clear(self) -> None:
             logger.error(f"[ERROR] Failed to clear database '{self.db_name}': {e}")
             raise
 
-    def export_graph(self, include_embedding: bool = True) -> dict[str, Any]:
+    def export_graph(self, **kwargs) -> dict[str, Any]:
         """
         Export all graph nodes and edges in a structured form.
 
@@ -914,13 +912,12 @@ def import_graph(self, data: dict[str, Any]) -> None:
                     target_id=edge["target"],
                 )
 
-    def get_all_memory_items(self, scope: str, include_embedding: bool = True) -> list[dict]:
+    def get_all_memory_items(self, scope: str, **kwargs) -> list[dict]:
         """
         Retrieve all memory items of a specific memory_type.
 
         Args:
             scope (str): Must be one of 'WorkingMemory', 'LongTermMemory', or 'UserMemory'.
-            include_embedding (bool): Whether to include the large embedding field.
         Returns:
 
         Returns:
@@ -946,9 +943,7 @@ def get_all_memory_items(self, scope: str, include_embedding: bool = True) -> li
             results = session.run(query, params)
             return [self._parse_node(dict(record["n"])) for record in results]
 
-    def get_structure_optimization_candidates(
-        self, scope: str, include_embedding: bool = True
-    ) -> list[dict]:
+    def get_structure_optimization_candidates(self, scope: str, **kwargs) -> list[dict]:
         """
         Find nodes that are likely candidates for structure optimization:
         - Isolated nodes, nodes with empty background, or nodes with exactly one child.

src/memos/graph_dbs/neo4j_community.py

@@ -169,14 +169,12 @@ def search_by_embedding(
         # Return consistent format
         return [{"id": r.id, "score": r.score} for r in results]
 
-    def get_all_memory_items(self, scope: str, include_embedding: bool = True) -> list[dict]:
+    def get_all_memory_items(self, scope: str, **kwargs) -> list[dict]:
         """
         Retrieve all memory items of a specific memory_type.
 
         Args:
             scope (str): Must be one of 'WorkingMemory', 'LongTermMemory', or 'UserMemory'.
-            include_embedding (bool): Whether to include the large embedding field.
-
         Returns:
             list[dict]: Full list of memory items under this scope.
         """

src/memos/mem_os/core.py

@@ -971,7 +971,7 @@ def get_user_info(self) -> dict[str, Any]:
         return {
             "user_id": user.user_id,
             "user_name": user.user_name,
-            "role": user.role.value,
+            "role": user.role.value if hasattr(user.role, "value") else user.role,
             "created_at": user.created_at.isoformat(),
             "accessible_cubes": [
                 {

src/memos/mem_user/mysql_user_manager.py

@@ -55,7 +55,9 @@ class User(Base):
 
     user_id = Column(String(255), primary_key=True, default=lambda: str(uuid.uuid4()))
     user_name = Column(String(255), unique=True, nullable=False)
-    role = Column(String(20), default=UserRole.USER.value, nullable=False)
+    role = Column(
+        String(20), default=UserRole.USER.value, nullable=False
+    )  # for sqlite backend this is SQLEnum
     created_at = Column(DateTime, default=datetime.now, nullable=False)
     updated_at = Column(DateTime, default=datetime.now, onupdate=datetime.now, nullable=False)
     is_active = Column(Boolean, default=True, nullable=False)
@@ -65,7 +67,7 @@ class User(Base):
     owned_cubes = relationship("Cube", back_populates="owner", cascade="all, delete-orphan")
 
     def __repr__(self):
-        return f"<User(user_id='{self.user_id}', user_name='{self.user_name}', role='{self.role.value}')>"
+        return f"<User(user_id='{self.user_id}', user_name='{self.user_name}', role='{self.role}')>"
 
 
 class Cube(Base):