feature: mem reader update and some bug fixs (#43)

* fix n4j cypher query

* feat: add llm extra body

* feat: update memory extraction prompt and result parser

* fix: evaluation locomo search

* ci: fix format and update test

* feat: update result json parser

Author: Nyakult

src/memos/configs/llm.py

@@ -24,6 +24,7 @@ class OpenAILLMConfig(BaseLLMConfig):
     api_base: str = Field(
         default="https://api.openai.com/v1", description="Base URL for OpenAI API"
     )
+    extra_body: Any = Field(default=None, description="extra body")
 
 
 class OllamaLLMConfig(BaseLLMConfig):

src/memos/configs/mem_reader.py

@@ -22,6 +22,10 @@ class BaseMemReaderConfig(BaseConfig):
     chunker: ChunkerConfigFactory = Field(
         ..., description="Chunker configuration for the MemReader"
     )
+    remove_prompt_example: bool = Field(
+        default=False,
+        description="whether remove example in memory extraction prompt to save token",
+    )
 
 
 class SimpleStructMemReaderConfig(BaseMemReaderConfig):

src/memos/graph_dbs/neo4j.py

@@ -811,7 +811,7 @@ def drop_database(self) -> None:
 
     def _ensure_database_exists(self):
         with self.driver.session(database="system") as session:
-            session.run(f"CREATE DATABASE {self.db_name} IF NOT EXISTS")
+            session.run(f"CREATE DATABASE $db_name IF NOT EXISTS", db_name=self.db_name)
 
         # Wait until the database is available
         for _ in range(10):

src/memos/llms/openai.py

@@ -22,6 +22,7 @@ def generate(self, messages: MessageList) -> str:
         response = self.client.chat.completions.create(
             model=self.config.model_name_or_path,
             messages=messages,
+            extra_body=self.config.extra_body,
             temperature=self.config.temperature,
             max_tokens=self.config.max_tokens,
             top_p=self.config.top_p,

src/memos/mem_reader/simple_struct.py

@@ -1,7 +1,7 @@
 import concurrent.futures
 import copy
 import json
-
+import re
 from abc import ABC
 from typing import Any
 
@@ -17,6 +17,7 @@
 from memos.templates.mem_reader_prompts import (
     SIMPLE_STRUCT_DOC_READER_PROMPT,
     SIMPLE_STRUCT_MEM_READER_PROMPT,
+    SIMPLE_STRUCT_MEM_READER_EXAMPLE,
 )
 
 
@@ -39,11 +40,11 @@ def __init__(self, config: SimpleStructMemReaderConfig):
         self.chunker = ChunkerFactory.from_config(config.chunker)
 
     def _process_chat_data(self, scene_data_info, info):
-        prompt = (
-            SIMPLE_STRUCT_MEM_READER_PROMPT.replace("${user_a}", "user")
-            .replace("${user_b}", "assistant")
-            .replace("${conversation}", "\n".join(scene_data_info))
+        prompt = SIMPLE_STRUCT_MEM_READER_PROMPT.replace(
+            "${conversation}", "\n".join(scene_data_info)
         )
+        if self.config.remove_prompt_example:
+            prompt = prompt.replace(SIMPLE_STRUCT_MEM_READER_EXAMPLE, "")
 
         messages = [{"role": "user", "content": prompt}]
 
@@ -228,7 +229,11 @@ def _process_doc_data(self, scene_data_info, info):
 
     def parse_json_result(self, response_text):
         try:
-            response_text = response_text.replace("```", "").replace("json", "")
+            json_start = response_text.find("{")
+            response_text = response_text[json_start:]
+            response_text = response_text.replace("```", "").strip()
+            if response_text[-1] != "}":
+                response_text += "}"
             response_json = json.loads(response_text)
             return response_json
         except json.JSONDecodeError as e:

src/memos/templates/mem_reader_prompts.py

@@ -1,44 +1,42 @@
-SIMPLE_STRUCT_MEM_READER_PROMPT = """
-You are a memory extraction expert.
+SIMPLE_STRUCT_MEM_READER_PROMPT = """You are a memory extraction expert.
 
-Your task is to extract memories from the perspective of ${user_a}, based on a conversation between ${user_a} and ${user_b}. This means identifying what ${user_a} would plausibly remember — including their own experiences, thoughts, plans, or relevant statements and actions made by others (such as ${user_b}) that impacted or were acknowledged by ${user_a}.
+Your task is to extract memories from the perspective of user, based on a conversation between user and assistant. This means identifying what user would plausibly remember — including their own experiences, thoughts, plans, or relevant statements and actions made by others (such as assistant) that impacted or were acknowledged by user.
 
 Please perform:
-1. Identify information that reflects ${user_a}'s experiences, beliefs, concerns, decisions, plans, or reactions — including meaningful input from ${user_b} that ${user_a} acknowledged or responded to.
+1. Identify information that reflects user's experiences, beliefs, concerns, decisions, plans, or reactions — including meaningful input from assistant that user acknowledged or responded to.
 2. Resolve all time, person, and event references clearly:
    - Convert relative time expressions (e.g., “yesterday,” “next Friday”) into absolute dates using the message timestamp if possible.
    - Clearly distinguish between event time and message time.
    - If uncertainty exists, state it explicitly (e.g., “around June 2025,” “exact date unclear”).
    - Include specific locations if mentioned.
    - Resolve all pronouns, aliases, and ambiguous references into full names or identities.
    - Disambiguate people with the same name if applicable.
-3. Always write from a third-person perspective, referring to ${user_a} as
+3. Always write from a third-person perspective, referring to user as
 "The user" or by name if name mentioned, rather than using first-person ("I", "me", "my").
 For example, write "The user felt exhausted..." instead of "I felt exhausted...".
-4. Do not omit any information that ${user_a} is likely to remember.
+4. Do not omit any information that user is likely to remember.
    - Include all key experiences, thoughts, emotional responses, and plans — even if they seem minor.
    - Prioritize completeness and fidelity over conciseness.
-   - Do not generalize or skip details that could be personally meaningful to ${user_a}.
+   - Do not generalize or skip details that could be personally meaningful to user.
 
 Return a single valid JSON object with the following structure:
 
 {
   "memory list": [
     {
-      "key": <string, a unique, concise memory title in English>,
+      "key": <string, a unique, concise memory title>,
       "memory_type": <string, Either "LongTermMemory" or "UserMemory">,
       "value": <A detailed, self-contained, and unambiguous memory statement — written in English if the input conversation is in English, or in Chinese if the conversation is in Chinese>,
-      "tags": <A list of relevant English thematic keywords (e.g.,
-      ["deadline", "team", "planning"])>
+      "tags": <A list of relevant thematic keywords (e.g., ["deadline", "team", "planning"])>
     },
     ...
   ],
-  "summary": <a natural paragraph summarizing the above memories from ${user_a}'s perspective, 120–200 words, same language as the input>
+  "summary": <a natural paragraph summarizing the above memories from user's perspective, 120–200 words, same language as the input>
 }
 
 Language rules:
-- The `value` fields and `summary` must match the language of the input conversation.
-- All metadata fields (`key`, `memory_type`, `tags`) must be in English.
+- The `key`, `value`, `tags`, `summary` fields must match the language of the input conversation.
+- Keep `memory_type` in English.
 
 Example:
 Conversation:
@@ -71,8 +69,7 @@
 Conversation:
 ${conversation}
 
-Your Output:
-"""
+Your Output:"""
 
 SIMPLE_STRUCT_DOC_READER_PROMPT = """
 You are an expert text analyst for a search and retrieval system. Your task is to process a document chunk and generate a single, structured JSON object.
@@ -96,3 +93,33 @@
 
 Produce ONLY the JSON object as your response.
 """
+
+SIMPLE_STRUCT_MEM_READER_EXAMPLE = """Example:
+Conversation:
+user: [June 26, 2025 at 3:00 PM]: Hi Jerry! Yesterday at 3 PM I had a meeting with my team about the new project.
+assistant: Oh Tom! Do you think the team can finish by December 15?
+user: [June 26, 2025 at 3:00 PM]: I’m worried. The backend won’t be done until
+December 10, so testing will be tight.
+assistant: [June 26, 2025 at 3:00 PM]: Maybe propose an extension?
+user: [June 26, 2025 at 4:21 PM]: Good idea. I’ll raise it in tomorrow’s 9:30 AM meeting—maybe shift the deadline to January 5.
+
+Output:
+{
+  "memory list": [
+    {
+        "key": "Initial project meeting",
+        "memory_type": "LongTermMemory",
+        "value": "On June 25, 2025 at 3:00 PM, Tom held a meeting with their team to discuss a new project. The conversation covered the timeline and raised concerns about the feasibility of the December 15, 2025 deadline.",
+        "tags": ["project", "timeline", "meeting", "deadline"]
+    },
+    {
+        "key": "Planned scope adjustment",
+        "memory_type": "UserMemory",
+        "value": "Tom planned to suggest in a meeting on June 27, 2025 at 9:30 AM that the team should prioritize features and propose shifting the project deadline to January 5, 2026.",
+        "tags": ["planning", "deadline change", "feature prioritization"]
+    },
+  ],
+  "summary": "Tom is currently focused on managing a new project with a tight schedule. After a team meeting on June 25, 2025, he realized the original deadline of December 15 might not be feasible due to backend delays. Concerned about insufficient testing time, he welcomed Jerry’s suggestion of proposing an extension. Tom plans to raise the idea of shifting the deadline to January 5, 2026 in the next morning’s meeting. His actions reflect both stress about timelines and a proactive, team-oriented problem-solving approach."
+}
+
+"""

tests/configs/test_llm.py

@@ -47,6 +47,7 @@ def test_openai_llm_config():
             "top_k",
             "api_base",
             "remove_think_prefix",
+            "extra_body",
         ],
     )
 

tests/mem_reader/test_simple_structure.py

@@ -19,6 +19,7 @@ def setUp(self):
         self.config.llm = MagicMock()
         self.config.embedder = MagicMock()
         self.config.chunker = MagicMock()
+        self.config.remove_prompt_example = MagicMock()
 
         # Mock dependencies
         with (