fix: fix strategy reader input; code reformat (#457)

* update reader and search strategy

* set strategy reader and search config

* fix install problem

* fix

* fix test

* turn off graph recall

* turn off graph recall

* turn off graph recall

* fix Searcher input bug

* fix Searcher

* fix Search

* fix bug

* adjust strategy reader

* adjust strategy reader

* adjust search config input

* reformat code

* re pr

* format repair

---------

Co-authored-by: 黑布林 <[email protected]>
Co-authored-by: CaralHsi <[email protected]>
Co-authored-by: chunyu li <[email protected]>

Author: 4 people

README.md

@@ -57,7 +57,7 @@
 MemOS demonstrates significant improvements over baseline memory solutions in multiple memory tasks,
 showcasing its capabilities in **information extraction**, **temporal and cross-session reasoning**, and **personalized preference responses**.
 
-| Model           | LOCOMO      | LongMemEval | PrefEval-10 | PersonaMem  | 
+| Model           | LOCOMO      | LongMemEval | PrefEval-10 | PersonaMem  |
 |-----------------|-------------|-------------|-------------|-------------|
 | **GPT-4o-mini** | 52.75       | 55.4        | 2.8         | 43.46       |
 | **MemOS**       | **75.80**   | **77.80**   | **71.90**   | **61.17**   |

src/memos/api/config.py

@@ -427,7 +427,7 @@ def get_reader_config() -> dict[str, Any]:
             "config": {
                 "chunk_type": os.getenv("MEM_READER_CHAT_CHUNK_TYPE", "default"),
                 "chunk_length": int(os.getenv("MEM_READER_CHAT_CHUNK_TOKEN_SIZE", 1600)),
-                "chunk_session": int(os.getenv("MEM_READER_CHAT_CHUNK_SESS_SIZE", 20)),
+                "chunk_session": int(os.getenv("MEM_READER_CHAT_CHUNK_SESS_SIZE", 10)),
                 "chunk_overlap": int(os.getenv("MEM_READER_CHAT_CHUNK_OVERLAP", 2)),
             },
         }

src/memos/configs/memory.py

@@ -184,7 +184,7 @@ class TreeTextMemoryConfig(BaseTextMemoryConfig):
         ),
     )
 
-    search_strategy: dict[str, bool] | None = Field(
+    search_strategy: dict[str, Any] | None = Field(
         default=None,
         description=(
             'Set search strategy for this memory configuration.{"bm25": true, "cot": false}'

src/memos/mem_reader/strategy_struct.py

@@ -43,7 +43,7 @@ def _get_llm_response(self, mem_str: str) -> dict:
         template = STRATEGY_PROMPT_DICT["chat"][lang]
         examples = STRATEGY_PROMPT_DICT["chat"][f"{lang}_example"]
         prompt = template.replace("${conversation}", mem_str)
-        if self.config.remove_prompt_example:
+        if self.config.remove_prompt_example:  # TODO unused
             prompt = prompt.replace(examples, "")
         messages = [{"role": "user", "content": prompt}]
         try:
@@ -112,6 +112,19 @@ def get_scene_data_info(self, scene_data: list, type: str) -> list[str]:
 
                             results.append([overlap_item, item])
                             current_length = overlap_length + content_length
+            else:
+                cut_size, cut_overlap = (
+                    self.chat_chunker["chunk_session"],
+                    self.chat_chunker["chunk_overlap"],
+                )
+                for items in scene_data:
+                    step = cut_size - cut_overlap
+                    end = len(items) - cut_overlap
+                    if end <= 0:
+                        results.extend([items[:]])
+                    else:
+                        results.extend([items[i : i + cut_size] for i in range(0, end, step)])
+
         elif type == "doc":
             parser_config = ParserConfigFactory.model_validate(
                 {

src/memos/memories/textual/simple_tree.py

@@ -66,7 +66,9 @@ def __init__(
         time_start_bm = time.time()
         self.search_strategy = config.search_strategy
         self.bm25_retriever = (
-            EnhancedBM25() if self.search_strategy and self.search_strategy["bm25"] else None
+            EnhancedBM25()
+            if self.search_strategy and self.search_strategy.get("bm25", False)
+            else None
         )
         logger.info(f"time init: bm25_retriever time is: {time.time() - time_start_bm}")
 

src/memos/memories/textual/tree_text_memory/retrieve/retrieval_mid_structs.py

@@ -13,3 +13,4 @@ class ParsedTaskGoal:
     rephrased_query: str | None = None
     internet_search: bool = False
     goal_type: str | None = None  # e.g., 'default', 'explanation', etc.
+    context: str = ""

src/memos/memories/textual/tree_text_memory/retrieve/retrieve_utils.py

@@ -17,7 +17,7 @@ def find_project_root(marker=".git"):
         if (current / marker).exists():
             return current
         current = current.parent
-    logger.warn(f"The project root directory tag file was not found: {marker}")
+    return Path(".")
 
 
 PROJECT_ROOT = find_project_root()

src/memos/memories/textual/tree_text_memory/retrieve/searcher.py

@@ -30,8 +30,8 @@
 
 logger = get_logger(__name__)
 COT_DICT = {
-    "fast": {"en": COT_PROMPT, "zh": COT_PROMPT_ZH},
-    "fine": {"en": SIMPLE_COT_PROMPT, "zh": SIMPLE_COT_PROMPT_ZH},
+    "fine": {"en": COT_PROMPT, "zh": COT_PROMPT_ZH},
+    "fast": {"en": SIMPLE_COT_PROMPT, "zh": SIMPLE_COT_PROMPT_ZH},
 }
 
 
@@ -59,12 +59,8 @@ def __init__(
         # Create internet retriever from config if provided
         self.internet_retriever = internet_retriever
         self.moscube = moscube
-        self.vec_cot = (
-            search_strategy.get("vec_cot", "false") == "true" if search_strategy else False
-        )
-        self.use_fast_graph = (
-            search_strategy.get("fast_graph", "false") == "true" if search_strategy else False
-        )
+        self.vec_cot = search_strategy.get("cot", False) if search_strategy else False
+        self.use_fast_graph = search_strategy.get("fast_graph", False) if search_strategy else False
 
         self._usage_executor = ContextThreadPoolExecutor(max_workers=4, thread_name_prefix="usage")
 
@@ -287,6 +283,7 @@ def _retrieve_paths(
                     search_filter,
                     user_name,
                     id_filter,
+                    mode=mode,
                 )
             )
             tasks.append(
@@ -369,6 +366,7 @@ def _retrieve_from_long_term_and_user(
         search_filter: dict | None = None,
         user_name: str | None = None,
         id_filter: dict | None = None,
+        mode: str = "fast",
     ):
         """Retrieve and rerank from LongTermMemory and UserMemory"""
         results = []
@@ -377,7 +375,7 @@ def _retrieve_from_long_term_and_user(
         # chain of thinking
         cot_embeddings = []
         if self.vec_cot:
-            queries = self._cot_query(query)
+            queries = self._cot_query(query, mode=mode, context=parsed_goal.context)
             if len(queries) > 1:
                 cot_embeddings = self.embedder.embed(queries)
             cot_embeddings.extend(query_embedding)
@@ -566,7 +564,6 @@ def _cot_query(
             prompt = template.replace("${original_query}", query).replace(
                 "${split_num_threshold}", str(split_num)
             )
-        logger.info("COT process")
 
         messages = [{"role": "user", "content": prompt}]
         try:

src/memos/memories/textual/tree_text_memory/retrieve/task_goal_parser.py

@@ -39,7 +39,7 @@ def parse(
         - mode == 'fine': use LLM to parse structured topic/keys/tags
         """
         if mode == "fast":
-            return self._parse_fast(task_description, **kwargs)
+            return self._parse_fast(task_description, context=context, **kwargs)
         elif mode == "fine":
             if not self.llm:
                 raise ValueError("LLM not provided for slow mode.")
@@ -51,6 +51,7 @@ def _parse_fast(self, task_description: str, **kwargs) -> ParsedTaskGoal:
         """
         Fast mode: simple jieba word split.
         """
+        context = kwargs.get("context", "")
         use_fast_graph = kwargs.get("use_fast_graph", False)
         if use_fast_graph:
             desc_tokenized = self.tokenizer.tokenize_mixed(task_description)
@@ -61,6 +62,7 @@ def _parse_fast(self, task_description: str, **kwargs) -> ParsedTaskGoal:
                 goal_type="default",
                 rephrased_query=task_description,
                 internet_search=False,
+                context=context,
             )
         else:
             return ParsedTaskGoal(
@@ -70,6 +72,7 @@ def _parse_fast(self, task_description: str, **kwargs) -> ParsedTaskGoal:
                 goal_type="default",
                 rephrased_query=task_description,
                 internet_search=False,
+                context=context,
             )
 
     def _parse_fine(
@@ -91,16 +94,17 @@ def _parse_fine(
             logger.info(f"Parsing Goal... LLM input is {prompt}")
             response = self.llm.generate(messages=[{"role": "user", "content": prompt}])
             logger.info(f"Parsing Goal... LLM Response is {response}")
-            return self._parse_response(response)
+            return self._parse_response(response, context=context)
         except Exception:
             logger.warning(f"Fail to fine-parse query {query}: {traceback.format_exc()}")
-            return self._parse_fast(query)
+            return self._parse_fast(query, context=context)
 
-    def _parse_response(self, response: str) -> ParsedTaskGoal:
+    def _parse_response(self, response: str, **kwargs) -> ParsedTaskGoal:
         """
         Parse LLM JSON output safely.
         """
         try:
+            context = kwargs.get("context", "")
             response = response.replace("```", "").replace("json", "").strip()
             response_json = eval(response)
             return ParsedTaskGoal(
@@ -110,6 +114,7 @@ def _parse_response(self, response: str) -> ParsedTaskGoal:
                 rephrased_query=response_json.get("rephrased_instruction", None),
                 internet_search=response_json.get("internet_search", False),
                 goal_type=response_json.get("goal_type", "default"),
+                context=context,
             )
         except Exception as e:
             raise ValueError(f"Failed to parse LLM output: {e}\nRaw response:\n{response}") from e