feat: change and add config for multi-memreader

Author: fridayL

src/memos/api/config.py

@@ -321,7 +321,7 @@ def get_activation_config() -> dict[str, Any]:
 
     @staticmethod
     def get_memreader_config() -> dict[str, Any]:
-        """Get MemReader configuration."""
+        """Get MemReader configuration for chat/doc extraction (fine-tuned 0.6B model)."""
         return {
             "backend": "openai",
             "config": {
@@ -338,6 +338,107 @@ def get_memreader_config() -> dict[str, Any]:
             },
         }
 
+    @staticmethod
+    def get_memreader_general_llm_config() -> dict[str, Any]:
+        """Get general LLM configuration for non-chat/doc tasks.
+
+        Used for: hallucination filter, memory rewrite, memory merge,
+        tool trajectory extraction, skill memory extraction.
+
+        This is the fallback for image_parser_llm and preference_extractor_llm.
+        Fallback chain: MEMREADER_GENERAL_MODEL -> MEMRADER_MODEL (memreader config)
+
+        Note: If you have fine-tuned a custom model for chat/doc extraction only,
+        you should configure MEMREADER_GENERAL_MODEL to use a general-purpose LLM
+        for other tasks. Otherwise, all tasks will use the same MEMRADER_MODEL.
+        """
+        # Check if specific general model is configured
+        general_model = os.getenv("MEMREADER_GENERAL_MODEL")
+        if general_model:
+            return {
+                "backend": os.getenv("MEMREADER_GENERAL_BACKEND", "openai"),
+                "config": {
+                    "model_name_or_path": general_model,
+                    "temperature": 0.6,
+                    "max_tokens": int(os.getenv("MEMREADER_GENERAL_MAX_TOKENS", "8000")),
+                    "top_p": 0.95,
+                    "top_k": 20,
+                    "api_key": os.getenv(
+                        "MEMREADER_GENERAL_API_KEY", os.getenv("OPENAI_API_KEY", "EMPTY")
+                    ),
+                    "api_base": os.getenv(
+                        "MEMREADER_GENERAL_API_BASE",
+                        os.getenv("OPENAI_API_BASE", "https://api.openai.com/v1"),
+                    ),
+                    "remove_think_prefix": True,
+                },
+            }
+        # Fallback to memreader config (same behavior as before for users who don't customize)
+        return APIConfig.get_memreader_config()
+
+    @staticmethod
+    def get_image_parser_llm_config() -> dict[str, Any]:
+        """Get LLM configuration for image parsing (requires vision model).
+
+        Used for: image content extraction and analysis.
+        Requires a vision-capable model like GPT-4V, GPT-4o, etc.
+
+        Fallback chain: IMAGE_PARSER_MODEL -> general_llm -> OpenAI config
+        """
+        image_model = os.getenv("IMAGE_PARSER_MODEL")
+        if image_model:
+            return {
+                "backend": os.getenv("IMAGE_PARSER_BACKEND", "openai"),
+                "config": {
+                    "model_name_or_path": image_model,
+                    "temperature": 0.6,
+                    "max_tokens": int(os.getenv("IMAGE_PARSER_MAX_TOKENS", "4096")),
+                    "top_p": 0.95,
+                    "top_k": 20,
+                    "api_key": os.getenv(
+                        "IMAGE_PARSER_API_KEY", os.getenv("OPENAI_API_KEY", "EMPTY")
+                    ),
+                    "api_base": os.getenv(
+                        "IMAGE_PARSER_API_BASE",
+                        os.getenv("OPENAI_API_BASE", "https://api.openai.com/v1"),
+                    ),
+                    "remove_think_prefix": True,
+                },
+            }
+        # Fallback to general_llm config (which itself falls back to OpenAI)
+        return APIConfig.get_memreader_general_llm_config()
+
+    @staticmethod
+    def get_preference_extractor_llm_config() -> dict[str, Any]:
+        """Get LLM configuration for preference extraction.
+
+        Used for: extracting user preferences from conversations.
+
+        Fallback chain: PREFERENCE_EXTRACTOR_MODEL -> general_llm -> OpenAI config
+        """
+        pref_model = os.getenv("PREFERENCE_EXTRACTOR_MODEL")
+        if pref_model:
+            return {
+                "backend": os.getenv("PREFERENCE_EXTRACTOR_BACKEND", "openai"),
+                "config": {
+                    "model_name_or_path": pref_model,
+                    "temperature": 0.6,
+                    "max_tokens": int(os.getenv("PREFERENCE_EXTRACTOR_MAX_TOKENS", "8000")),
+                    "top_p": 0.95,
+                    "top_k": 20,
+                    "api_key": os.getenv(
+                        "PREFERENCE_EXTRACTOR_API_KEY", os.getenv("OPENAI_API_KEY", "EMPTY")
+                    ),
+                    "api_base": os.getenv(
+                        "PREFERENCE_EXTRACTOR_API_BASE",
+                        os.getenv("OPENAI_API_BASE", "https://api.openai.com/v1"),
+                    ),
+                    "remove_think_prefix": True,
+                },
+            }
+        # Fallback to general_llm config (which itself falls back to OpenAI)
+        return APIConfig.get_memreader_general_llm_config()
+
     @staticmethod
     def get_activation_vllm_config() -> dict[str, Any]:
         """Get Ollama configuration."""
@@ -358,7 +459,7 @@ def get_preference_memory_config() -> dict[str, Any]:
         return {
             "backend": "pref_text",
             "config": {
-                "extractor_llm": APIConfig.get_memreader_config(),
+                "extractor_llm": APIConfig.get_preference_extractor_llm_config(),
                 "vector_db": {
                     "backend": "milvus",
                     "config": APIConfig.get_milvus_config(),
@@ -802,6 +903,10 @@ def get_product_default_config() -> dict[str, Any]:
                 "backend": reader_config["backend"],
                 "config": {
                     "llm": APIConfig.get_memreader_config(),
+                    # General LLM for non-chat/doc tasks (hallucination filter, rewrite, merge, etc.)
+                    "general_llm": APIConfig.get_memreader_general_llm_config(),
+                    # Image parser LLM (requires vision model)
+                    "image_parser_llm": APIConfig.get_image_parser_llm_config(),
                     "embedder": APIConfig.get_embedder_config(),
                     "chunker": {
                         "backend": "sentence",
@@ -924,6 +1029,10 @@ def create_user_config(user_name: str, user_id: str) -> tuple["MOSConfig", "Gene
                 "backend": reader_config["backend"],
                 "config": {
                     "llm": APIConfig.get_memreader_config(),
+                    # General LLM for non-chat/doc tasks (hallucination filter, rewrite, merge, etc.)
+                    "general_llm": APIConfig.get_memreader_general_llm_config(),
+                    # Image parser LLM (requires vision model)
+                    "image_parser_llm": APIConfig.get_image_parser_llm_config(),
                     "embedder": APIConfig.get_embedder_config(),
                     "chunker": {
                         "backend": "sentence",

src/memos/api/handlers/config_builders.py

@@ -201,3 +201,26 @@ def build_nli_client_config() -> dict[str, Any]:
         NLI client configuration dictionary
     """
     return APIConfig.get_nli_config()
+
+
+def build_general_llm_config() -> dict[str, Any]:
+    """
+    Build general LLM configuration for non-chat/doc tasks.
+
+    Used for: hallucination filter, memory rewrite, memory merge,
+    tool trajectory extraction, skill memory extraction.
+
+    Returns:
+        Validated general LLM configuration dictionary
+    """
+    return LLMConfigFactory.model_validate(APIConfig.get_memreader_general_llm_config())
+
+
+def build_image_parser_llm_config() -> dict[str, Any]:
+    """
+    Build image parser LLM configuration (requires vision model).
+
+    Returns:
+        Validated image parser LLM configuration dictionary
+    """
+    return LLMConfigFactory.model_validate(APIConfig.get_image_parser_llm_config())

src/memos/configs/mem_reader.py

@@ -24,7 +24,18 @@ def parse_datetime(cls, value):
             return datetime.fromisoformat(value.replace("Z", "+00:00"))
         return value
 
-    llm: LLMConfigFactory = Field(..., description="LLM configuration for the MemReader")
+    llm: LLMConfigFactory = Field(
+        ..., description="LLM configuration for chat/doc memory extraction (fine-tuned model)"
+    )
+    general_llm: LLMConfigFactory | None = Field(
+        default=None,
+        description="General LLM for non-chat/doc tasks: hallucination filter, memory rewrite, "
+        "memory merge, tool trajectory, skill memory. Falls back to main llm if not set.",
+    )
+    image_parser_llm: LLMConfigFactory | None = Field(
+        default=None,
+        description="Vision LLM for image parsing. Falls back to main llm if not set.",
+    )
     embedder: EmbedderConfigFactory = Field(
         ..., description="Embedder configuration for the MemReader"
     )

src/memos/mem_reader/multi_modal_struct.py

@@ -39,6 +39,7 @@ def __init__(self, config: MultiModalStructMemReaderConfig):
             config: Configuration object for the reader
         """
         from memos.configs.mem_reader import SimpleStructMemReaderConfig
+        from memos.llms.factory import LLMFactory
 
         # Extract direct_markdown_hostnames before converting to SimpleStructMemReaderConfig
         direct_markdown_hostnames = getattr(config, "direct_markdown_hostnames", None)
@@ -56,10 +57,20 @@ def __init__(self, config: MultiModalStructMemReaderConfig):
         simple_config = SimpleStructMemReaderConfig(**config_dict)
         super().__init__(simple_config)
 
+        # Image parser LLM (requires vision model)
+        # Falls back to main llm if not configured
+        self.image_parser_llm = (
+            LLMFactory.from_config(config.image_parser_llm)
+            if config.image_parser_llm is not None
+            else self.llm
+        )
+
         # Initialize MultiModalParser for routing to different parsers
+        # Pass image_parser_llm for image parsing
         self.multi_modal_parser = MultiModalParser(
             embedder=self.embedder,
             llm=self.llm,
+            image_parser_llm=self.image_parser_llm,
             parser=None,
             direct_markdown_hostnames=direct_markdown_hostnames,
         )
@@ -631,7 +642,8 @@ def _merge_memories_with_llm(
         )
 
         try:
-            response_text = self.llm.generate([{"role": "user", "content": merge_prompt}])
+            # Use general_llm for memory merge (not fine-tuned for this task)
+            response_text = self.general_llm.generate([{"role": "user", "content": merge_prompt}])
             merge_result = parse_json_result(response_text)
 
             if merge_result.get("should_merge", False):
@@ -873,12 +885,14 @@ def get_chunk_idx(item_with_pos) -> int:
     def _get_llm_tool_trajectory_response(self, mem_str: str) -> dict:
         """
         Generete tool trajectory experience item by llm.
+        Uses general_llm as this task is not fine-tuned for the main model.
         """
         try:
             lang = detect_lang(mem_str)
             template = TOOL_TRAJECTORY_PROMPT_ZH if lang == "zh" else TOOL_TRAJECTORY_PROMPT_EN
             prompt = template.replace("{messages}", mem_str)
-            rsp = self.llm.generate([{"role": "user", "content": prompt}])
+            # Use general_llm for tool trajectory (not fine-tuned for this task)
+            rsp = self.general_llm.generate([{"role": "user", "content": prompt}])
             rsp = rsp.replace("```json", "").replace("```", "")
             return json.loads(rsp)
         except Exception as e:
@@ -1000,13 +1014,14 @@ def _process_multi_modal_data(
                 future_tool = executor.submit(
                     self._process_tool_trajectory_fine, fast_memory_items, info, **kwargs
                 )
+                # Use general_llm for skill memory extraction (not fine-tuned for this task)
                 future_skill = executor.submit(
                     process_skill_memory_fine,
                     fast_memory_items=fast_memory_items,
                     info=info,
                     searcher=self.searcher,
                     graph_db=self.graph_db,
-                    llm=self.llm,
+                    llm=self.general_llm,
                     embedder=self.embedder,
                     oss_config=self.oss_config,
                     skills_dir_config=self.skills_dir_config,
@@ -1067,12 +1082,13 @@ def _process_transfer_multi_modal_data(
             future_tool = executor.submit(
                 self._process_tool_trajectory_fine, raw_nodes, info, **kwargs
             )
+            # Use general_llm for skill memory extraction (not fine-tuned for this task)
             future_skill = executor.submit(
                 process_skill_memory_fine,
                 raw_nodes,
                 info,
                 searcher=self.searcher,
-                llm=self.llm,
+                llm=self.general_llm,
                 embedder=self.embedder,
                 graph_db=self.graph_db,
                 oss_config=self.oss_config,

src/memos/mem_reader/read_multi_modal/multi_modal_parser.py

@@ -35,6 +35,7 @@ def __init__(
         self,
         embedder: BaseEmbedder,
         llm: BaseLLM | None = None,
+        image_parser_llm: BaseLLM | None = None,
         parser: Any | None = None,
         direct_markdown_hostnames: list[str] | None = None,
     ):
@@ -43,14 +44,18 @@ def __init__(
 
         Args:
             embedder: Embedder for generating embeddings
-            llm: Optional LLM for fine mode processing
+            llm: Optional LLM for fine mode processing (chat/doc extraction)
+            image_parser_llm: Optional vision LLM for image parsing.
+                Falls back to llm if not provided.
             parser: Optional parser for parsing file contents
             direct_markdown_hostnames: List of hostnames that should return markdown directly
                 without parsing. If None, reads from FILE_PARSER_DIRECT_MARKDOWN_HOSTNAMES
                 environment variable (comma-separated). Default: ["139.196.232.20"]
         """
         self.embedder = embedder
         self.llm = llm
+        # Image parser LLM (requires vision model), falls back to main llm
+        self.image_parser_llm = image_parser_llm if image_parser_llm is not None else llm
         self.parser = parser
 
         # Initialize parsers for different message types
@@ -63,7 +68,8 @@ def __init__(
         self.file_content_parser = FileContentParser(
             embedder, llm, parser, direct_markdown_hostnames=direct_markdown_hostnames
         )
-        self.image_parser = ImageParser(embedder, llm)
+        # Use dedicated image_parser_llm for image parsing (requires vision model)
+        self.image_parser = ImageParser(embedder, self.image_parser_llm)
         self.audio_parser = None  # future
 
         self.role_parsers = {

src/memos/mem_reader/simple_struct.py

@@ -173,7 +173,15 @@ def __init__(self, config: SimpleStructMemReaderConfig):
             config: Configuration object for the reader
         """
         self.config = config
+        # Main LLM for chat/doc memory extraction (fine-tuned model)
         self.llm = LLMFactory.from_config(config.llm)
+        # General LLM for non-chat/doc tasks (hallucination filter, rewrite, merge, etc.)
+        # Falls back to main llm if not configured
+        self.general_llm = (
+            LLMFactory.from_config(config.general_llm)
+            if config.general_llm is not None
+            else self.llm
+        )
         self.embedder = EmbedderFactory.from_config(config.embedder)
         self.chunker = ChunkerFactory.from_config(config.chunker)
         self.save_rawfile = self.chunker.config.save_rawfile
@@ -505,8 +513,9 @@ def rewrite_memories(
         prompt = template.format(**prompt_args)
 
         # Optionally run filter and parse the output
+        # Use general_llm for rewrite (not fine-tuned for this task)
         try:
-            raw = self.llm.generate([{"role": "user", "content": prompt}])
+            raw = self.general_llm.generate([{"role": "user", "content": prompt}])
             success, parsed = parse_rewritten_response(raw)
             logger.info(
                 f"[rewrite_memories] Hallucination filter parsed successfully: {success}；prompt: {prompt}"
@@ -565,8 +574,9 @@ def filter_hallucination_in_memories(
         prompt = template.format(**prompt_args)
 
         # Optionally run filter and parse the output
+        # Use general_llm for hallucination filter (not fine-tuned for this task)
         try:
-            raw = self.llm.generate([{"role": "user", "content": prompt}])
+            raw = self.general_llm.generate([{"role": "user", "content": prompt}])
             success, parsed = parse_keep_filter_response(raw)
             logger.info(
                 f"[filter_hallucination_in_memories] Hallucination filter parsed successfully: {success}；prompt: {prompt}"