feat: add further questions for dialogue (#236)

* feat: add further questions for

* fix: remove  print

* feat: add stream for further question

---------

Co-authored-by: CaralHsi <[email protected]>

Author: fridayL

src/memos/api/product_models.py

@@ -161,3 +161,4 @@ class SuggestionRequest(BaseRequest):
 
     user_id: str = Field(..., description="User ID")
     language: Literal["zh", "en"] = Field("zh", description="Language for suggestions")
+    message: list[MessageDict] | None = Field(None, description="List of messages to store.")

src/memos/api/routers/product_router.py

@@ -148,7 +148,9 @@ def get_suggestion_queries_post(suggestion_req: SuggestionRequest):
     try:
         mos_product = get_mos_product_instance()
         suggestions = mos_product.get_suggestion_query(
-            user_id=suggestion_req.user_id, language=suggestion_req.language
+            user_id=suggestion_req.user_id,
+            language=suggestion_req.language,
+            message=suggestion_req.message,
         )
         return SuggestionResponse(
             message="Suggestions retrieved successfully", data={"query": suggestions}

src/memos/mem_os/product.py

@@ -36,7 +36,13 @@
 from memos.memories.textual.item import (
     TextualMemoryItem,
 )
-from memos.templates.mos_prompts import MEMOS_PRODUCT_BASE_PROMPT, MEMOS_PRODUCT_ENHANCE_PROMPT
+from memos.templates.mos_prompts import (
+    FURTHER_SUGGESTION_PROMPT,
+    MEMOS_PRODUCT_BASE_PROMPT,
+    MEMOS_PRODUCT_ENHANCE_PROMPT,
+    SUGGESTION_QUERY_PROMPT_EN,
+    SUGGESTION_QUERY_PROMPT_ZH,
+)
 from memos.types import MessageList
 
 
@@ -641,7 +647,23 @@ 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, language: str = "zh") -> list[str]:
+    def _get_further_suggestion(self, message: MessageList | None = None) -> list[str]:
+        """Get further suggestion prompt."""
+        try:
+            dialogue_info = "\n".join([f"{msg['role']}: {msg['content']}" for msg in message[-2:]])
+            further_suggestion_prompt = FURTHER_SUGGESTION_PROMPT.format(dialogue=dialogue_info)
+            message_list = [{"role": "system", "content": further_suggestion_prompt}]
+            response = self.chat_llm.generate(message_list)
+            clean_response = clean_json_response(response)
+            response_json = json.loads(clean_response)
+            return response_json["query"]
+        except Exception as e:
+            logger.error(f"Error getting further suggestion: {e}", exc_info=True)
+            return []
+
+    def get_suggestion_query(
+        self, user_id: str, language: str = "zh", message: MessageList | None = None
+    ) -> list[str]:
         """Get suggestion query from LLM.
         Args:
             user_id (str): User ID.
@@ -650,37 +672,13 @@ def get_suggestion_query(self, user_id: str, language: str = "zh") -> list[str]:
         Returns:
             list[str]: The suggestion query list.
         """
-
+        if message:
+            further_suggestion = self._get_further_suggestion(message)
+            return further_suggestion
         if language == "zh":
-            suggestion_prompt = """
-            你是一个有用的助手，可以帮助用户生成建议查询。
-            我将获取用户最近的一些记忆，
-            你应该生成一些建议查询，这些查询应该是用户想要查询的内容，
-            用户最近的记忆是：
-            {memories}
-            请生成3个建议查询用中文，
-            输出应该是json格式，键是"query"，值是一个建议查询列表。
-
-            示例：
-            {{
-                "query": ["查询1", "查询2", "查询3"]
-            }}
-            """
+            suggestion_prompt = SUGGESTION_QUERY_PROMPT_ZH
         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}
-            if the user recently memories is empty, 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"]
-            }}
-            """
+            suggestion_prompt = SUGGESTION_QUERY_PROMPT_EN
         text_mem_result = super().search("my recently memories", user_id=user_id, top_k=3)[
             "text_mem"
         ]
@@ -844,6 +842,11 @@ def chat_with_references(
         total_time = round(float(time_end - time_start), 1)
 
         yield f"data: {json.dumps({'type': 'time', 'data': {'total_time': total_time, 'speed_improvement': f'{speed_improvement}%'}})}\n\n"
+        # get further suggestion
+        current_messages.append({"role": "assistant", "content": full_response})
+        further_suggestion = self._get_further_suggestion(current_messages)
+        logger.info(f"further_suggestion: {further_suggestion}")
+        yield f"data: {json.dumps({'type': 'suggestion', 'data': further_suggestion})}\n\n"
         yield f"data: {json.dumps({'type': 'end'})}\n\n"
 
         logger.info(f"user_id: {user_id}, cube_id: {cube_id}, current_messages: {current_messages}")

src/memos/templates/mos_prompts.py

@@ -185,3 +185,57 @@
 {dialogue}
 Current question: {query}
 Answer:"""
+
+SUGGESTION_QUERY_PROMPT_ZH = """
+你是一个有用的助手，可以帮助用户生成建议查询。
+我将获取用户最近的一些记忆，
+你应该生成一些建议查询，这些查询应该是用户想要查询的内容，
+用户最近的记忆是：
+{memories}
+请生成3个建议查询用中文，如果用户最近的记忆是空，请直接随机生成3个建议查询用中文，不要有多余解释。
+输出应该是json格式，键是"query"，值是一个建议查询列表。
+
+示例：
+{{
+    "query": ["查询1", "查询2", "查询3"]
+}}
+"""
+
+SUGGESTION_QUERY_PROMPT_EN = """
+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}
+if the user recently memories is empty, please generate 3 suggestion query in English,do not generate any other text,
+output should be a json format, the key is "query", the value is a list of suggestion query.
+
+example:
+{{
+    "query": ["query1", "query2", "query3"]
+}}
+"""
+
+FURTHER_SUGGESTION_PROMPT = """
+You are a helpful assistant.
+You are given a dialogue between a user and a assistant.
+You need to suggest a further question based on the dialogue.
+Requirements:
+1. The further question should be related to the dialogue.
+2. The further question should be concise and accurate.
+3. You must return ONLY a valid JSON object. Do not include any other text, explanations, or formatting.
+the lastest dialogue is:
+{dialogue}
+output should be a json format, the key is "query", the value is a list of suggestion query.
+if dialogue is chinese,the quersuggestion query should be in chinese,if dialogue is english,the suggestion query should be in english.
+please do not generate any other text.
+
+example english:
+{{
+    "query": ["query1", "query2", "query3"]
+}}
+example chinese:
+{{
+    "query": ["问题1", "问题2", "问题3"]
+}}
+"""