Feat: add recall strategy (#414)

* update reader and search strategy

* set strategy reader and search config

* fix install problem

* fix

* fix test

---------

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

Author: 3 people

poetry.lock

@@ -107,7 +107,9 @@ all = [
     "markitdown[docx,pdf,pptx,xls,xlsx] (>=0.1.1,<0.2.0)",
     "pymilvus (>=2.6.1,<3.0.0)",
     "datasketch (>=1.6.5,<2.0.0)",
-
+    "jieba (>=0.38.1,<0.42.1)",
+    "rank-bm25 (>=0.2.2)",
+    "cachetools (>=6.0.0)",
     # NOT exist in the above optional groups
     # Because they are either huge-size dependencies or infrequently used dependencies.
     # We kindof don't want users to install them.

pyproject.toml

@@ -419,17 +419,31 @@ def get_embedder_config() -> dict[str, Any]:
                 },
             }
 
+    @staticmethod
+    def get_reader_config() -> dict[str, Any]:
+        """Get reader configuration."""
+        return {
+            "backend": os.getenv("MEM_READER_BACKEND", "simple_struct"),
+            "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_overlap": int(os.getenv("MEM_READER_CHAT_CHUNK_OVERLAP", 2)),
+            },
+        }
+
     @staticmethod
     def get_internet_config() -> dict[str, Any]:
         """Get embedder configuration."""
+        reader_config = APIConfig.get_reader_config()
         return {
             "backend": "bocha",
             "config": {
                 "api_key": os.getenv("BOCHA_API_KEY"),
                 "max_results": 15,
                 "num_per_request": 10,
                 "reader": {
-                    "backend": "simple_struct",
+                    "backend": reader_config["backend"],
                     "config": {
                         "llm": {
                             "backend": "openai",
@@ -455,6 +469,7 @@ def get_internet_config() -> dict[str, Any]:
                                 "min_sentences_per_chunk": 1,
                             },
                         },
+                        "chat_chunker": reader_config,
                     },
                 },
             },
@@ -656,6 +671,8 @@ def get_product_default_config() -> dict[str, Any]:
         openai_config = APIConfig.get_openai_config()
         qwen_config = APIConfig.qwen_config()
         vllm_config = APIConfig.vllm_config()
+        reader_config = APIConfig.get_reader_config()
+
         backend_model = {
             "openai": openai_config,
             "huggingface": qwen_config,
@@ -667,7 +684,7 @@ def get_product_default_config() -> dict[str, Any]:
             "user_id": os.getenv("MOS_USER_ID", "root"),
             "chat_model": {"backend": backend, "config": backend_model[backend]},
             "mem_reader": {
-                "backend": "simple_struct",
+                "backend": reader_config["backend"],
                 "config": {
                     "llm": APIConfig.get_memreader_config(),
                     "embedder": APIConfig.get_embedder_config(),
@@ -680,6 +697,7 @@ def get_product_default_config() -> dict[str, Any]:
                             "min_sentences_per_chunk": 1,
                         },
                     },
+                    "chat_chunker": reader_config,
                 },
             },
             "enable_textual_memory": True,
@@ -750,6 +768,7 @@ def create_user_config(user_name: str, user_id: str) -> tuple[MOSConfig, General
         qwen_config = APIConfig.qwen_config()
         vllm_config = APIConfig.vllm_config()
         mysql_config = APIConfig.get_mysql_config()
+        reader_config = APIConfig.get_reader_config()
         backend = os.getenv("MOS_CHAT_MODEL_PROVIDER", "openai")
         backend_model = {
             "openai": openai_config,
@@ -764,7 +783,7 @@ def create_user_config(user_name: str, user_id: str) -> tuple[MOSConfig, General
                 "config": backend_model[backend],
             },
             "mem_reader": {
-                "backend": "simple_struct",
+                "backend": reader_config["backend"],
                 "config": {
                     "llm": APIConfig.get_memreader_config(),
                     "embedder": APIConfig.get_embedder_config(),
@@ -777,6 +796,7 @@ def create_user_config(user_name: str, user_id: str) -> tuple[MOSConfig, General
                             "min_sentences_per_chunk": 1,
                         },
                     },
+                    "chat_chunker": reader_config,
                 },
             },
             "enable_textual_memory": True,
@@ -845,6 +865,10 @@ def create_user_config(user_name: str, user_id: str) -> tuple[MOSConfig, General
                                 "LongTermMemory": os.getenv("NEBULAR_LONGTERM_MEMORY", 1e6),
                                 "UserMemory": os.getenv("NEBULAR_USER_MEMORY", 1e6),
                             },
+                            "search_strategy": {
+                                "bm25": bool(os.getenv("BM25_CALL", "false") == "true"),
+                                "cot": bool(os.getenv("VEC_COT_CALL", "false") == "true"),
+                            },
                         },
                     },
                     "act_mem": {}
@@ -912,6 +936,10 @@ def get_default_cube_config() -> GeneralMemCubeConfig | None:
                                 "LongTermMemory": os.getenv("NEBULAR_LONGTERM_MEMORY", 1e6),
                                 "UserMemory": os.getenv("NEBULAR_USER_MEMORY", 1e6),
                             },
+                            "search_strategy": {
+                                "bm25": bool(os.getenv("BM25_CALL", "false") == "true"),
+                                "cot": bool(os.getenv("VEC_COT_CALL", "false") == "true"),
+                            },
                             "mode": os.getenv("ASYNC_MODE", "sync"),
                         },
                     },

src/memos/api/config.py

@@ -36,11 +36,19 @@ def parse_datetime(cls, value):
         description="whether remove example in memory extraction prompt to save token",
     )
 
+    chat_chunker: dict[str, Any] = Field(
+        default=None, description="Configuration for the MemReader chat chunk strategy"
+    )
+
 
 class SimpleStructMemReaderConfig(BaseMemReaderConfig):
     """SimpleStruct MemReader configuration class."""
 
 
+class StrategyStructMemReaderConfig(BaseMemReaderConfig):
+    """StrategyStruct MemReader configuration class."""
+
+
 class MemReaderConfigFactory(BaseConfig):
     """Factory class for creating MemReader configurations."""
 
@@ -49,6 +57,7 @@ class MemReaderConfigFactory(BaseConfig):
 
     backend_to_class: ClassVar[dict[str, Any]] = {
         "simple_struct": SimpleStructMemReaderConfig,
+        "strategy_struct": StrategyStructMemReaderConfig,
     }
 
     @field_validator("backend")

src/memos/configs/mem_reader.py

@@ -184,6 +184,13 @@ class TreeTextMemoryConfig(BaseTextMemoryConfig):
         ),
     )
 
+    search_strategy: dict[str, bool] | None = Field(
+        default=None,
+        description=(
+            'Set search strategy for this memory configuration.{"bm25": true, "cot": false}'
+        ),
+    )
+
     mode: str | None = Field(
         default="sync",
         description=("whether use asynchronous mode in memory add"),

src/memos/configs/memory.py

@@ -58,15 +58,18 @@ def clear_cache(cls):
         logger.info("OpenAI LLM instance cache cleared")
 
     @timed(log=True, log_prefix="OpenAI LLM")
-    def generate(self, messages: MessageList) -> str:
-        """Generate a response from OpenAI LLM."""
+    def generate(self, messages: MessageList, **kwargs) -> str:
+        """Generate a response from OpenAI LLM, optionally overriding generation params."""
+        temperature = kwargs.get("temperature", self.config.temperature)
+        max_tokens = kwargs.get("max_tokens", self.config.max_tokens)
+        top_p = kwargs.get("top_p", self.config.top_p)
         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,
+            temperature=temperature,
+            max_tokens=max_tokens,
+            top_p=top_p,
         )
         logger.info(f"Response from OpenAI: {response.model_dump_json()}")
         response_content = response.choices[0].message.content

src/memos/llms/openai.py

@@ -3,6 +3,7 @@
 from memos.configs.mem_reader import MemReaderConfigFactory
 from memos.mem_reader.base import BaseMemReader
 from memos.mem_reader.simple_struct import SimpleStructMemReader
+from memos.mem_reader.strategy_struct import StrategyStructMemReader
 from memos.memos_tools.singleton import singleton_factory
 
 
@@ -11,6 +12,7 @@ class MemReaderFactory(BaseMemReader):
 
     backend_to_class: ClassVar[dict[str, Any]] = {
         "simple_struct": SimpleStructMemReader,
+        "strategy_struct": StrategyStructMemReader,
     }
 
     @classmethod