From 3b27826bddcbc5830e1f973fd0291e9b6952f2c1 Mon Sep 17 00:00:00 2001 From: "yuan.wang" Date: Wed, 11 Feb 2026 15:14:29 +0800 Subject: [PATCH 1/3] fix: llm extra body bug --- src/memos/api/config.py | 4 ++-- src/memos/api/handlers/config_builders.py | 1 + src/memos/configs/llm.py | 1 + src/memos/llms/vllm.py | 16 ++-------------- 4 files changed, 6 insertions(+), 16 deletions(-) diff --git a/src/memos/api/config.py b/src/memos/api/config.py index 70d9366e3..27f77b6ab 100644 --- a/src/memos/api/config.py +++ b/src/memos/api/config.py @@ -335,7 +335,7 @@ def get_memreader_config() -> dict[str, Any]: # validation requirements during tests/import. "api_base": os.getenv("MEMRADER_API_BASE", "https://api.openai.com/v1"), "remove_think_prefix": True, - "extra_body": {"chat_template_kwargs": {"enable_thinking": False}}, + "extra_body": {"enable_thinking": False}, }, } @@ -531,7 +531,7 @@ def get_internet_config() -> dict[str, Any]: "api_key": os.getenv("MEMRADER_API_KEY", "EMPTY"), "api_base": os.getenv("MEMRADER_API_BASE"), "remove_think_prefix": True, - "extra_body": {"chat_template_kwargs": {"enable_thinking": False}}, + "extra_body": {"enable_thinking": False}, }, }, "embedder": APIConfig.get_embedder_config(), diff --git a/src/memos/api/handlers/config_builders.py b/src/memos/api/handlers/config_builders.py index 2b3fbdd35..7426a60ba 100644 --- a/src/memos/api/handlers/config_builders.py +++ b/src/memos/api/handlers/config_builders.py @@ -105,6 +105,7 @@ def build_chat_llm_config() -> list[dict[str, Any]]: } ), "support_models": cfg.get("support_models", None), + "extra_body": cfg.get("extra_body", None), } for cfg in configs ] diff --git a/src/memos/configs/llm.py b/src/memos/configs/llm.py index 70217b896..5487d117c 100644 --- a/src/memos/configs/llm.py +++ b/src/memos/configs/llm.py @@ -116,6 +116,7 @@ class VLLMLLMConfig(BaseLLMConfig): default=False, description="Enable reasoning outputs from vLLM", ) + extra_body: Any = Field(default=None, description="Extra options for API") class LLMConfigFactory(BaseConfig): diff --git a/src/memos/llms/vllm.py b/src/memos/llms/vllm.py index 362112f11..0efead97b 100644 --- a/src/memos/llms/vllm.py +++ b/src/memos/llms/vllm.py @@ -111,13 +111,7 @@ def _generate_with_api_client(self, messages: list[MessageDict], **kwargs) -> st "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), - "extra_body": { - "chat_template_kwargs": { - "enable_thinking": kwargs.get( - "enable_thinking", self.config.enable_thinking - ) - } - }, + "extra_body": kwargs.get("extra_body", self.config.extra_body), } if kwargs.get("tools"): completion_kwargs["tools"] = kwargs.get("tools") @@ -175,13 +169,7 @@ def generate_stream(self, messages: list[MessageDict], **kwargs): "max_tokens": kwargs.get("max_tokens", self.config.max_tokens), "top_p": kwargs.get("top_p", self.config.top_p), "stream": True, - "extra_body": { - "chat_template_kwargs": { - "enable_thinking": kwargs.get( - "enable_thinking", self.config.enable_thinking - ) - } - }, + "extra_body": kwargs.get("extra_body", self.config.extra_body), } stream = self.client.chat.completions.create(**completion_kwargs) From cd5272b2fa948e8b37f4cede5a02aa1550297976 Mon Sep 17 00:00:00 2001 From: "yuan.wang" Date: Thu, 12 Feb 2026 10:36:24 +0800 Subject: [PATCH 2/3] fix: fix bug --- src/memos/api/config.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/src/memos/api/config.py b/src/memos/api/config.py index 27f77b6ab..d2dd19266 100644 --- a/src/memos/api/config.py +++ b/src/memos/api/config.py @@ -335,7 +335,6 @@ def get_memreader_config() -> dict[str, Any]: # validation requirements during tests/import. "api_base": os.getenv("MEMRADER_API_BASE", "https://api.openai.com/v1"), "remove_think_prefix": True, - "extra_body": {"enable_thinking": False}, }, } @@ -531,7 +530,6 @@ def get_internet_config() -> dict[str, Any]: "api_key": os.getenv("MEMRADER_API_KEY", "EMPTY"), "api_base": os.getenv("MEMRADER_API_BASE"), "remove_think_prefix": True, - "extra_body": {"enable_thinking": False}, }, }, "embedder": APIConfig.get_embedder_config(), From 7e5695963865e9023485151d63494b231062c323 Mon Sep 17 00:00:00 2001 From: "yuan.wang" Date: Thu, 12 Feb 2026 11:38:24 +0800 Subject: [PATCH 3/3] feat: add timed log to function --- src/memos/mem_reader/multi_modal_struct.py | 2 ++ .../mem_reader/read_multi_modal/multi_modal_parser.py | 4 ++++ .../mem_reader/read_skill_memory/process_skill_memory.py | 7 +++++++ 3 files changed, 13 insertions(+) diff --git a/src/memos/mem_reader/multi_modal_struct.py b/src/memos/mem_reader/multi_modal_struct.py index 8b0968ca1..a27d64758 100644 --- a/src/memos/mem_reader/multi_modal_struct.py +++ b/src/memos/mem_reader/multi_modal_struct.py @@ -646,6 +646,7 @@ def _merge_memories_with_llm( return None + @timed def _process_string_fine( self, fast_memory_items: list[TextualMemoryItem], @@ -883,6 +884,7 @@ def _get_llm_tool_trajectory_response(self, mem_str: str) -> dict: logger.error(f"[MultiModalFine] Error calling LLM for tool trajectory: {e}") return [] + @timed def _process_tool_trajectory_fine( self, fast_memory_items: list[TextualMemoryItem], info: dict[str, Any], **kwargs ) -> list[TextualMemoryItem]: diff --git a/src/memos/mem_reader/read_multi_modal/multi_modal_parser.py b/src/memos/mem_reader/read_multi_modal/multi_modal_parser.py index 808410e65..81bd25902 100644 --- a/src/memos/mem_reader/read_multi_modal/multi_modal_parser.py +++ b/src/memos/mem_reader/read_multi_modal/multi_modal_parser.py @@ -11,6 +11,7 @@ from memos.log import get_logger from memos.memories.textual.item import SourceMessage, TextualMemoryItem from memos.types import MessagesType +from memos.utils import timed from .assistant_parser import AssistantParser from .base import BaseMessageParser @@ -120,6 +121,7 @@ def _get_parser(self, message: Any) -> BaseMessageParser | None: logger.warning(f"[MultiModalParser] Could not determine parser for message: {message}") return None + @timed def parse( self, message: MessagesType, @@ -157,6 +159,7 @@ def parse( logger.error(f"[MultiModalParser] Error parsing message: {e}") return [] + @timed def parse_batch( self, messages: list[MessagesType], @@ -182,6 +185,7 @@ def parse_batch( results.append(items) return results + @timed def process_transfer( self, source: SourceMessage, diff --git a/src/memos/mem_reader/read_skill_memory/process_skill_memory.py b/src/memos/mem_reader/read_skill_memory/process_skill_memory.py index d691e53a3..d39955ac2 100644 --- a/src/memos/mem_reader/read_skill_memory/process_skill_memory.py +++ b/src/memos/mem_reader/read_skill_memory/process_skill_memory.py @@ -36,6 +36,7 @@ TOOL_GENERATION_PROMPT, ) from memos.types import MessageList +from memos.utils import timed load_dotenv() @@ -64,6 +65,7 @@ def _generate_content_by_llm(llm: BaseLLM, prompt_template: str, **kwargs) -> An return {} if "json" in prompt_template.lower() else "" +@timed def _batch_extract_skills( task_chunks: dict[str, MessageList], related_memories_map: dict[str, list[TextualMemoryItem]], @@ -97,6 +99,7 @@ def _batch_extract_skills( return results +@timed def _batch_generate_skill_details( raw_skills_data: list[tuple[dict[str, Any], str, MessageList]], related_skill_memories_map: dict[str, list[TextualMemoryItem]], @@ -756,6 +759,7 @@ def _delete_skills( logger.warning(f"Error deleting local file: {e}") +@timed def _write_skills_to_file( skill_memory: dict[str, Any], info: dict[str, Any], skills_dir_config: dict[str, Any] ) -> str: @@ -1000,6 +1004,7 @@ def _get_skill_file_storage_location() -> str: return "LOCAL" +@timed def process_skill_memory_fine( fast_memory_items: list[TextualMemoryItem], info: dict[str, Any], @@ -1064,6 +1069,7 @@ def process_skill_memory_fine( ) related_skill_memories_by_task[task_name] = [] + @timed def _simple_extract(): # simple extract skill memory, only one stage memories = [] @@ -1096,6 +1102,7 @@ def _simple_extract(): ) return memories + @timed def _full_extract(): # full extract skill memory, include two stage raw_extraction_results = _batch_extract_skills(