feat: update zh en and mem

Author: fridayL

src/memos/api/product_models.py

@@ -150,3 +150,10 @@ class SearchRequest(BaseRequest):
     user_id: str = Field(..., description="User ID")
     query: str = Field(..., description="Search query")
     mem_cube_id: str | None = Field(None, description="Cube ID to search in")
+
+
+class SuggestionRequest(BaseRequest):
+    """Request model for getting suggestion queries."""
+
+    user_id: str = Field(..., description="User ID")
+    language: Literal["zh", "en"] = Field("zh", description="Language for suggestions")

src/memos/api/routers/product_router.py

@@ -15,6 +15,7 @@
     SearchRequest,
     SearchResponse,
     SimpleResponse,
+    SuggestionRequest,
     SuggestionResponse,
     UserRegisterRequest,
     UserRegisterResponse,
@@ -86,7 +87,6 @@ async def register_user(user_req: UserRegisterRequest):
         logger.error(f"Failed to register user: {traceback.format_exc()}")
         raise HTTPException(status_code=500, detail=str(traceback.format_exc())) from err
 
-
 @router.get(
     "/suggestions/{user_id}", summary="Get suggestion queries", response_model=SuggestionResponse
 )
@@ -105,6 +105,25 @@ async def get_suggestion_queries(user_id: str):
         raise HTTPException(status_code=500, detail=str(traceback.format_exc())) from err
 
 
+@router.post("/suggestions", summary="Get suggestion queries with language", response_model=SuggestionResponse)
+async def get_suggestion_queries_post(suggestion_req: SuggestionRequest):
+    """Get suggestion queries for a specific user with language preference."""
+    try:
+        mos_product = get_mos_product_instance()
+        suggestions = mos_product.get_suggestion_query(
+            user_id=suggestion_req.user_id, 
+            language=suggestion_req.language
+        )
+        return SuggestionResponse(
+            message="Suggestions retrieved successfully", data={"query": suggestions}
+        )
+    except ValueError as err:
+        raise HTTPException(status_code=404, detail=str(traceback.format_exc())) from err
+    except Exception as err:
+        logger.error(f"Failed to get suggestions: {traceback.format_exc()}")
+        raise HTTPException(status_code=500, detail=str(traceback.format_exc())) from err
+
+
 @router.post("/get_all", summary="Get all memories for user", response_model=MemoryResponse)
 async def get_all_memories(memory_req: GetMemoryRequest):
     """Get all memories for a specific user."""

src/memos/mem_os/product.py

@@ -579,29 +579,46 @@ def user_register(
         except Exception as e:
             return {"status": "error", "message": f"Failed to register user: {e!s}"}
 
-    def get_suggestion_query(self, user_id: str) -> list[str]:
+    def get_suggestion_query(self, user_id: str, language: str = "zh") -> list[str]:
         """Get suggestion query from LLM.
         Args:
             user_id (str): User ID.
+            language (str): Language for suggestions ("zh" or "en").
 
         Returns:
             list[str]: The suggestion query list.
         """
 
-        suggestion_prompt = """
-        You are a helpful assistant that can help users to generate suggestion query
-        I will get some user recently memories,
-        you should generate some suggestion query  , the query should be user what to query,
-        user recently memories is :
-        {memories}
-        please generate 3 suggestion query,
-        output should be a json format, the key is "query", the value is a list of suggestion query.
-
-        example:
-        {{
-            "query": ["query1", "query2", "query3"]
-        }}
-        """
+        if language == "zh":
+            suggestion_prompt = """
+            你是一个有用的助手，可以帮助用户生成建议查询。
+            我将获取用户最近的一些记忆，
+            你应该生成一些建议查询，这些查询应该是用户想要查询的内容，
+            用户最近的记忆是：
+            {memories}
+            请生成3个建议查询用中文，
+            输出应该是json格式，键是"query"，值是一个建议查询列表。
+
+            示例：
+            {{
+                "query": ["查询1", "查询2", "查询3"]
+            }}
+            """
+        else:  # English
+            suggestion_prompt = """
+            You are a helpful assistant that can help users to generate suggestion query.
+            I will get some user recently memories,
+            you should generate some suggestion query, the query should be user what to query,
+            user recently memories is:
+            {memories}
+            please generate 3 suggestion query in English,
+            output should be a json format, the key is "query", the value is a list of suggestion query.
+
+            example:
+            {{
+                "query": ["query1", "query2", "query3"]
+            }}
+            """
         text_mem_result = super().search("my recently memories", user_id=user_id, top_k=10)["text_mem"]
         if text_mem_result:
             memories = "\n".join(
@@ -842,7 +859,7 @@ def get_all(
                     "LongTermMemory": 0.40,
                     "UserMemory": 0.40,
                 }
-                tree_result = convert_graph_to_tree_forworkmem(
+                tree_result, node_type_count = convert_graph_to_tree_forworkmem(
                     memories, target_node_count=150, type_ratios=custom_type_ratios
                 )
                 memories_filtered = filter_nodes_by_tree_ids(tree_result, memories)
@@ -851,7 +868,7 @@ def get_all(
                 tree_result["children"] = children_sort
                 memories_filtered["tree_structure"] = tree_result
                 reformat_memory_list.append(
-                    {"cube_id": memory["cube_id"], "memories": [memories_filtered]}
+                    {"cube_id": memory["cube_id"], "memories": [memories_filtered], "memory_statistics": node_type_count}
                 )
         elif memory_type == "act_mem":
             reformat_memory_list.append(
@@ -915,7 +932,7 @@ def get_subgraph(
         for memory in memory_list:
             memories = remove_embedding_recursive(memory["memories"])
             custom_type_ratios = {"WorkingMemory": 0.20, "LongTermMemory": 0.40, "UserMemory": 0.4}
-            tree_result = convert_graph_to_tree_forworkmem(
+            tree_result, node_type_count = convert_graph_to_tree_forworkmem(
                 memories, target_node_count=150, type_ratios=custom_type_ratios
             )
             memories_filtered = filter_nodes_by_tree_ids(tree_result, memories)
@@ -924,7 +941,7 @@ def get_subgraph(
             tree_result["children"] = children_sort
             memories_filtered["tree_structure"] = tree_result
             reformat_memory_list.append(
-                {"cube_id": memory["cube_id"], "memories": [memories_filtered]}
+                {"cube_id": memory["cube_id"], "memories": [memories_filtered], "memory_statistics": node_type_count}
             )
 
         return reformat_memory_list

src/memos/mem_os/utils/format_utils.py

@@ -508,6 +508,10 @@ def convert_graph_to_tree_forworkmem(
     for original_edge in original_edges:
         if original_edge["type"] == "PARENT":
             filter_original_edges.append(original_edge)
+    node_type_count = {}
+    for node in original_nodes:
+        node_type = node.get("metadata", {}).get("memory_type", "Unknown")
+        node_type_count[node_type] = node_type_count.get(node_type, 0) + 1
     original_edges = filter_original_edges
     # Use enhanced type-balanced sampling
     if len(original_nodes) > target_node_count:
@@ -625,7 +629,7 @@ def build_tree(node_id: str) -> dict[str, Any]:
         "frequency": 0,
     }
 
-    return result
+    return result, node_type_count
 
 
 def print_tree_structure(node: dict[str, Any], level: int = 0, max_level: int = 5):