Merge branch 'dev' into feat/ark_embed_support

Ki-Seki · web-flow · commit 6923c524f972 · 2025-07-13T02:16:42.000+08:00
diff --git a/pyproject.toml b/pyproject.toml
@@ -2,7 +2,7 @@
 
 [tool.poetry]
 name = "MemoryOS"
-version = "0.1.13"
+version = "0.2.0"
 description = "Intelligence Begins with Memory"
 license = "Apache-2.0"
 authors = ["MemTensor <lizy@memtensor.cn>"]
diff --git a/src/memos/__init__.py b/src/memos/__init__.py
@@ -1,4 +1,4 @@
-__version__ = "0.1.13"
+__version__ = "0.2.0"
 
 from memos.configs.mem_cube import GeneralMemCubeConfig
 from memos.configs.mem_os import MOSConfig
diff --git a/src/memos/configs/llm.py b/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):
diff --git a/src/memos/configs/mem_reader.py b/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):
diff --git a/src/memos/graph_dbs/neo4j.py b/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):
diff --git a/src/memos/llms/openai.py b/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,
diff --git a/src/memos/mem_reader/simple_struct.py b/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:
diff --git a/src/memos/memories/textual/tree_text_memory/organize/conflict.py b/src/memos/memories/textual/tree_text_memory/organize/conflict.py
@@ -167,10 +167,12 @@ def _resolve_in_graph(
             if not self.graph_store.edge_exists(new_from, new_to, edge["type"], direction="ANY"):
                 self.graph_store.add_edge(new_from, new_to, edge["type"])
 
-        self.graph_store.delete_node(conflict_a.id)
-        self.graph_store.delete_node(conflict_b.id)
+        self.graph_store.update_node(conflict_a.id, {"status": "archived"})
+        self.graph_store.update_node(conflict_b.id, {"status": "archived"})
+        self.graph_store.add_edge(conflict_a.id, merged.id, type="MERGED_TO")
+        self.graph_store.add_edge(conflict_b.id, merged.id, type="MERGED_TO")
         logger.debug(
-            f"Remove {conflict_a.id} and {conflict_b.id}, and inherit their edges to {merged.id}."
+            f"Archive {conflict_a.id} and {conflict_b.id}, and inherit their edges to {merged.id}."
         )
 
     def _merge_metadata(
diff --git a/src/memos/memories/textual/tree_text_memory/organize/redundancy.py b/src/memos/memories/textual/tree_text_memory/organize/redundancy.py
@@ -30,7 +30,7 @@ def detect(
         self, memory: TextualMemoryItem, top_k: int = 5, scope: str | None = None
     ) -> list[tuple[TextualMemoryItem, TextualMemoryItem]]:
         """
-        Detect redundancy by finding the most similar items in the graph database based on embedding, then use LLM to judge conflict.
+        Detect redundancy by finding the most similar items in the graph database based on embedding, then use LLM to judge redundancy.
         Args:
             memory: The memory item (should have an embedding attribute or field).
             top_k: Number of top similar nodes to retrieve.
@@ -49,15 +49,15 @@ def detect(
             for info in embedding_candidates_info
             if info["score"] >= self.EMBEDDING_THRESHOLD and info["id"] != memory.id
         ]
-        # 3. Judge conflicts using LLM
+        # 3. Judge redundancys using LLM
         embedding_candidates = self.graph_store.get_nodes(embedding_candidates_ids)
         redundant_pairs = []
         for embedding_candidate in embedding_candidates:
             embedding_candidate = TextualMemoryItem.from_dict(embedding_candidate)
             prompt = [
                 {
                     "role": "system",
-                    "content": "You are a conflict detector for memory items.",
+                    "content": "You are a redundancy detector for memory items.",
                 },
                 {
                     "role": "user",
@@ -71,25 +71,25 @@ def detect(
             if "yes" in result.lower():
                 redundant_pairs.append([memory, embedding_candidate])
         if len(redundant_pairs):
-            conflict_text = "\n".join(
+            redundant_text = "\n".join(
                 f'"{pair[0].memory!s}" <==REDUNDANCY==> "{pair[1].memory!s}"'
                 for pair in redundant_pairs
             )
             logger.warning(
-                f"Detected {len(redundant_pairs)} redundancies for memory {memory.id}\n {conflict_text}"
+                f"Detected {len(redundant_pairs)} redundancies for memory {memory.id}\n {redundant_text}"
             )
         return redundant_pairs
 
     def resolve_two_nodes(self, memory_a: TextualMemoryItem, memory_b: TextualMemoryItem) -> None:
         """
         Resolve detected redundancies between two memory items using LLM fusion.
         Args:
-            memory_a: The first conflicting memory item.
-            memory_b: The second conflicting memory item.
+            memory_a: The first redundant memory item.
+            memory_b: The second redundant memory item.
         Returns:
             A fused TextualMemoryItem representing the resolved memory.
         """
-
+        return  # waiting for implementation
         # ———————————— 1. LLM generate fused memory ————————————
         metadata_for_resolve = ["key", "background", "confidence", "updated_at"]
         metadata_1 = memory_a.metadata.model_dump_json(include=metadata_for_resolve)
@@ -115,18 +115,10 @@ def resolve_two_nodes(self, memory_a: TextualMemoryItem, memory_b: TextualMemory
         try:
             answer = re.search(r"<answer>(.*?)</answer>", response, re.DOTALL)
             answer = answer.group(1).strip()
-            # —————— 2.1 Can't resolve conflict, hard update by comparing timestamp ————
-            if len(answer) <= 10 and "no" in answer.lower():
-                logger.warning(
-                    f"Conflict between {memory_a.id} and {memory_b.id} could not be resolved. "
-                )
-                self._hard_update(memory_a, memory_b)
-            # —————— 2.2 Conflict resolved, update metadata and memory ————
-            else:
-                fixed_metadata = self._merge_metadata(answer, memory_a.metadata, memory_b.metadata)
-                merged_memory = TextualMemoryItem(memory=answer, metadata=fixed_metadata)
-                logger.info(f"Resolved result: {merged_memory}")
-                self._resolve_in_graph(memory_a, memory_b, merged_memory)
+            fixed_metadata = self._merge_metadata(answer, memory_a.metadata, memory_b.metadata)
+            merged_memory = TextualMemoryItem(memory=answer, metadata=fixed_metadata)
+            logger.info(f"Resolved result: {merged_memory}")
+            self._resolve_in_graph(memory_a, memory_b, merged_memory)
         except json.decoder.JSONDecodeError:
             logger.error(f"Failed to parse LLM response: {response}")
 
@@ -145,48 +137,37 @@ def resolve_one_node(self, memory: TextualMemoryItem) -> None:
         )
         logger.debug(f"Merged memory: {memory.memory}")
 
-    def _hard_update(self, memory_a: TextualMemoryItem, memory_b: TextualMemoryItem):
-        """
-        Hard update: compare updated_at, keep the newer one, overwrite the older one's metadata.
-        """
-        time_a = datetime.fromisoformat(memory_a.metadata.updated_at)
-        time_b = datetime.fromisoformat(memory_b.metadata.updated_at)
-
-        newer_mem = memory_a if time_a >= time_b else memory_b
-        older_mem = memory_b if time_a >= time_b else memory_a
-
-        self.graph_store.delete_node(older_mem.id)
-        logger.warning(
-            f"Delete older memory {older_mem.id}: <{older_mem.memory}> due to conflict with {newer_mem.id}: <{newer_mem.memory}>"
-        )
-
     def _resolve_in_graph(
         self,
-        conflict_a: TextualMemoryItem,
-        conflict_b: TextualMemoryItem,
+        redundant_a: TextualMemoryItem,
+        redundant_b: TextualMemoryItem,
         merged: TextualMemoryItem,
     ):
-        edges_a = self.graph_store.get_edges(conflict_a.id, type="ANY", direction="ANY")
-        edges_b = self.graph_store.get_edges(conflict_b.id, type="ANY", direction="ANY")
+        edges_a = self.graph_store.get_edges(redundant_a.id, type="ANY", direction="ANY")
+        edges_b = self.graph_store.get_edges(redundant_b.id, type="ANY", direction="ANY")
         all_edges = edges_a + edges_b
 
         self.graph_store.add_node(
             merged.id, merged.memory, merged.metadata.model_dump(exclude_none=True)
         )
 
         for edge in all_edges:
-            new_from = merged.id if edge["from"] in (conflict_a.id, conflict_b.id) else edge["from"]
-            new_to = merged.id if edge["to"] in (conflict_a.id, conflict_b.id) else edge["to"]
+            new_from = (
+                merged.id if edge["from"] in (redundant_a.id, redundant_b.id) else edge["from"]
+            )
+            new_to = merged.id if edge["to"] in (redundant_a.id, redundant_b.id) else edge["to"]
             if new_from == new_to:
                 continue
             # Check if the edge already exists before adding
             if not self.graph_store.edge_exists(new_from, new_to, edge["type"], direction="ANY"):
                 self.graph_store.add_edge(new_from, new_to, edge["type"])
 
-        self.graph_store.delete_node(conflict_a.id)
-        self.graph_store.delete_node(conflict_b.id)
+        self.graph_store.update_node(redundant_a.id, {"status": "archived"})
+        self.graph_store.update_node(redundant_b.id, {"status": "archived"})
+        self.graph_store.add_edge(redundant_a.id, merged.id, type="MERGED_TO")
+        self.graph_store.add_edge(redundant_b.id, merged.id, type="MERGED_TO")
         logger.debug(
-            f"Remove {conflict_a.id} and {conflict_b.id}, and inherit their edges to {merged.id}."
+            f"Archive {redundant_a.id} and {redundant_b.id}, and inherit their edges to {merged.id}."
         )
 
     def _merge_metadata(
diff --git a/src/memos/memories/textual/tree_text_memory/organize/reorganizer.py b/src/memos/memories/textual/tree_text_memory/organize/reorganizer.py
@@ -164,9 +164,9 @@ def handle_add(self, message: QueueMessage):
                 logger.info(f"Resolved conflict between {added_node.id} and {existing_node.id}.")
 
         # ———————— 2. check for redundancy ————————
-        redundancy = self.redundancy.detect(added_node, scope=added_node.metadata.memory_type)
-        if redundancy:
-            for added_node, existing_node in redundancy:
+        redundancies = self.redundancy.detect(added_node, scope=added_node.metadata.memory_type)
+        if redundancies:
+            for added_node, existing_node in redundancies:
                 self.redundancy.resolve_two_nodes(added_node, existing_node)
                 logger.info(f"Resolved redundancy between {added_node.id} and {existing_node.id}.")
 
@@ -176,7 +176,7 @@ def handle_remove(self, message: QueueMessage):
     def handle_merge(self, message: QueueMessage):
         after_node = message.after_node[0]
         logger.debug(f"Handling merge operation: <{after_node.memory}>")
-        self.redundancy_resolver.resolve_one_node(after_node)
+        self.redundancy.resolve_one_node(after_node)
 
     def optimize_structure(
         self,
diff --git a/src/memos/templates/mem_reader_prompts.py b/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."
+}
+
+"""
diff --git a/tests/configs/test_llm.py b/tests/configs/test_llm.py
@@ -47,6 +47,7 @@ def test_openai_llm_config():
             "top_k",
             "api_base",
             "remove_think_prefix",
+            "extra_body",
         ],
     )
 
diff --git a/tests/mem_reader/test_simple_structure.py b/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 (

Original file line number	Diff line number	Diff line change
`@@ -1,4 +1,4 @@`
`1`		`-__version__ = "0.1.13"`
	`1`	`+__version__ = "0.2.0"`
`2`	`2`
`3`	`3`	`from memos.configs.mem_cube import GeneralMemCubeConfig`
`4`	`4`	`from memos.configs.mem_os import MOSConfig`
Original file line number	Diff line number	Diff line change
`@@ -24,6 +24,7 @@ class OpenAILLMConfig(BaseLLMConfig):`
`24`	`24`	`api_base: str = Field(`
`25`	`25`	`default="https://api.openai.com/v1", description="Base URL for OpenAI API"`
`26`	`26`	`)`
	`27`	`+ extra_body: Any = Field(default=None, description="extra body")`
`27`	`28`
`28`	`29`
`29`	`30`	`class OllamaLLMConfig(BaseLLMConfig):`
Original file line number	Diff line number	Diff line change
`@@ -22,6 +22,10 @@ class BaseMemReaderConfig(BaseConfig):`
`22`	`22`	`chunker: ChunkerConfigFactory = Field(`
`23`	`23`	`..., description="Chunker configuration for the MemReader"`
`24`	`24`	`)`
	`25`	`+ remove_prompt_example: bool = Field(`
	`26`	`+ default=False,`
	`27`	`+ description="whether remove example in memory extraction prompt to save token",`
	`28`	`+ )`
`25`	`29`
`26`	`30`
`27`	`31`	`class SimpleStructMemReaderConfig(BaseMemReaderConfig):`