From 5e08aa48b81b586ccfd2f6edd5a4a9953e6cbc3a Mon Sep 17 00:00:00 2001 From: jinliyl <6469360+jinliyl@users.noreply.github.com> Date: Sat, 7 Mar 2026 15:22:56 +0800 Subject: [PATCH] refactor(memory): restructure file-based memory components and enhance message handling (#145) --- README.md | 35 +++++++++---------- README_ZH.md | 32 ++++++++--------- reme/core/schema/as_msg_stat.py | 35 ++++++++++++------- reme/memory/file_based/__init__.py | 30 ++++------------ reme/memory/file_based/component/__init__.py | 0 reme/memory/file_based/components/__init__.py | 13 +++++++ .../{component => components}/compactor.py | 5 ++- .../{component => components}/compactor.yaml | 0 .../context_checker.py | 15 ++++---- .../{component => components}/summarizer.py | 5 ++- .../{component => components}/summarizer.yaml | 0 .../tool_result_compactor.py | 0 .../file_based/reme_in_memory_memory.py | 2 +- reme/memory/file_based/utils/__init__.py | 7 ++++ .../file_based/{ => utils}/as_msg_handler.py | 33 ++++++++++++++--- reme/reme_light.py | 18 +++++----- tests/light/test_compactor.py | 4 +-- tests/light/test_context_check.py | 4 +-- tests/light/test_format_msgs_to_str.py | 4 +-- tests/light/test_summarizer.py | 4 +-- tests/light/test_tool_result_compactor.py | 3 +- tests/light/test_tools.py | 2 +- tests/light/test_utils.py | 2 +- 23 files changed, 148 insertions(+), 105 deletions(-) delete mode 100644 reme/memory/file_based/component/__init__.py create mode 100644 reme/memory/file_based/components/__init__.py rename reme/memory/file_based/{component => components}/compactor.py (89%) rename reme/memory/file_based/{component => components}/compactor.yaml (100%) rename reme/memory/file_based/{component => components}/context_checker.py (92%) rename reme/memory/file_based/{component => components}/summarizer.py (88%) rename reme/memory/file_based/{component => components}/summarizer.yaml (100%) rename reme/memory/file_based/{component => components}/tool_result_compactor.py (100%) create mode 100644 reme/memory/file_based/utils/__init__.py rename reme/memory/file_based/{ => utils}/as_msg_handler.py (92%) diff --git a/README.md b/README.md index 9d18e6ce..7bc1ef86 100644 --- a/README.md +++ b/README.md @@ -20,7 +20,7 @@ A memory management toolkit for AI agents — Remember Me, Refine Me.

-> For the older version, please refer to the [0.2.x documentation](docs/README_0_2_x_ZH.md). +> For the older version, please refer to the [0.2.x documentation](docs/README_0_2_x.md). --- @@ -64,17 +64,17 @@ working_dir/ [ReMeLight](reme/reme_light.py) is the core class of the file-based memory system. It provides full memory management capabilities for AI agents: -| Method | Function | Key components | -|------------------------|--------------------------------------|--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------| -| `check_context` | 📊 Check context size | [ContextChecker](reme/memory/file_based/component/context_checker.py) — checks whether context exceeds thresholds and splits messages | -| `compact_memory` | 📦 Compact history into summary | [Compactor](reme/memory/file_based/component/compactor.py) — ReActAgent that generates structured context summaries | -| `summary_memory` | 📝 Persist important memory to files | [Summarizer](reme/memory/file_based/component/summarizer.py) — ReActAgent + file tools (`read` / `write` / `edit`) | -| `compact_tool_result` | ✂️ Compact long tool outputs | [ToolResultCompactor](reme/memory/file_based/component/tool_result_compactor.py) — truncates long tool outputs and stores them in `tool_result/` while keeping file references in messages | -| `memory_search` | 🔍 Semantic memory search | [MemorySearch](reme/memory/file_based/tools/memory_search.py) — hybrid retrieval with vectors + BM25 | -| `ReMeInMemoryMemory` | 🗂️ In-session memory class | [ReMeInMemoryMemory](reme/memory/file_based/reme_in_memory_memory.py) — token-aware memory management with summary compression and state serialization | -| `pre_reasoning_hook` | 🔄 Pre-reasoning hook | `compact_tool_result` + `check_context` + `compact_memory` + `summary_memory` (async) | -| `start` | 🚀 Start memory system | Initialize file storage, file watcher, and embedding cache; clean up expired tool result files | -| `close` | 📕 Shutdown and cleanup | Clean up tool result files, stop file watcher, and persist embedding cache | +| Method | Function | Key components | +|-----------------------|--------------------------------------|---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------| +| `check_context` | 📊 Check context size | [ContextChecker](reme/memory/file_based/components/context_checker.py) — checks whether context exceeds thresholds and splits messages | +| `compact_memory` | 📦 Compact history into summary | [Compactor](reme/memory/file_based/components/compactor.py) — ReActAgent that generates structured context summaries | +| `summary_memory` | 📝 Persist important memory to files | [Summarizer](reme/memory/file_based/components/summarizer.py) — ReActAgent + file tools (`read` / `write` / `edit`) | +| `compact_tool_result` | ✂️ Compact long tool outputs | [ToolResultCompactor](reme/memory/file_based/components/tool_result_compactor.py) — truncates long tool outputs and stores them in `tool_result/` while keeping file references in messages | +| `memory_search` | 🔍 Semantic memory search | [MemorySearch](reme/memory/file_based/tools/memory_search.py) — hybrid retrieval with vectors + BM25 | +| `ReMeInMemoryMemory` | 🗂️ In-session memory class | [ReMeInMemoryMemory](reme/memory/file_based/reme_in_memory_memory.py) — token-aware memory management with summary compression and state serialization | +| `pre_reasoning_hook` | 🔄 Pre-reasoning hook | `compact_tool_result` + `check_context` + `compact_memory` + `summary_memory` (async) | +| `start` | 🚀 Start memory system | Initialize file storage, file watcher, and embedding cache; clean up expired tool result files | +| `close` | 📕 Shutdown and cleanup | Clean up tool result files, stop file watcher, and persist embedding cache | --- @@ -186,7 +186,7 @@ graph LR CC -->|Exceeds limit| SM[summary_memory
Async persistence] SM -->|ReAct + FileIO| Files[memory/*.md] Agent -->|Explicit call| Search[memory_search
Vector+BM25] - Agent -->|In-session| InMem[ReMeInMemoryMemory
Token-aware memory] + Agent -->|In - session| InMem[ReMeInMemoryMemory
Token-aware memory] Files -.->|FileWatcher| Store[(FileStore
Vector+FTS index)] Search --> Store ``` @@ -195,7 +195,7 @@ graph LR #### 1. `check_context` — context checking -[ContextChecker](reme/memory/file_based/component/context_checker.py) uses token counting to determine whether the +[ContextChecker](reme/memory/file_based/components/context_checker.py) uses token counting to determine whether the context exceeds thresholds and automatically splits messages into a "to compact" group and a "to keep" group. ```mermaid @@ -217,7 +217,7 @@ graph LR #### 2. `compact_memory` — conversation compaction -[Compactor](reme/memory/file_based/component/compactor.py) uses a ReActAgent to compact conversation history into a * +[Compactor](reme/memory/file_based/components/compactor.py) uses a ReActAgent to compact conversation history into a * *structured context summary**. ```mermaid @@ -245,7 +245,7 @@ graph LR #### 3. `summary_memory` — persistent memory -[Summarizer](reme/memory/file_based/component/summarizer.py) uses a **ReAct + file tools** pattern so that the AI can +[Summarizer](reme/memory/file_based/components/summarizer.py) uses a **ReAct + file tools** pattern so that the AI can decide what to write and where to write it. ```mermaid @@ -271,7 +271,7 @@ graph LR #### 4. `compact_tool_result` — tool result compaction -[ToolResultCompactor](reme/memory/file_based/component/tool_result_compactor.py) addresses the problem of long tool +[ToolResultCompactor](reme/memory/file_based/components/tool_result_compactor.py) addresses the problem of long tool outputs bloating the context. ```mermaid @@ -534,7 +534,6 @@ We evaluate ReMe on the BFCL-V3 multi-turn-base task (random split 50 train / 15 For more details on how to reproduce the experiments, see [quickstart.md](benchmark/bfcl/quickstart.md). - ## ⭐ Community & support - **Star & Watch**: Starring helps more agent developers discover ReMe; Watching keeps you up to date with new releases diff --git a/README_ZH.md b/README_ZH.md index 8723177c..bd038f67 100644 --- a/README_ZH.md +++ b/README_ZH.md @@ -60,17 +60,18 @@ working_dir/ [ReMeLight](reme/reme_light.py) 是该记忆系统的核心类，为 AI Agent 提供完整的记忆管理能力： -| 方法 | 功能 | 关键组件 | -|------------------------|--------------|-----------------------------------------------------------------------------------------------------------------------------| -| `check_context` | 📊 检查上下文大小 | [ContextChecker](reme/memory/file_based/component/context_checker.py) — 检查上下文是否超出阈值并拆分Message | -| `compact_memory` | 📦 压缩历史对话为摘要 | [Compactor](reme/memory/file_based/component/compactor.py) — ReActAgent 生成结构化上下文摘要 | -| `summary_memory` | 📝 将重要记忆写入文件 | [Summarizer](reme/memory/file_based/component/summarizer.py) — ReActAgent + 文件工具（read / write / edit） | -| `compact_tool_result` | ✂️ 压缩超长工具输出 | [ToolResultCompactor](reme/memory/file_based/component/tool_result_compactor.py) — 截断超长的工具调用结果并转存到 `tool_result/`，消息中保留文件引用 | -| `memory_search` | 🔍 语义搜索记忆 | [MemorySearch](reme/memory/file_based/tools/memory_search.py) — 向量 + BM25 混合检索 | -| `ReMeInMemoryMemory` | 🗂️ 会话内存类 | [ReMeInMemoryMemory](reme/memory/file_based/reme_in_memory_memory.py) — Token 感知的内存管理，支持压缩摘要和状态序列化 | -| `pre_reasoning_hook` | 🔄 推理前预处理钩子 | compact_tool_result + check_context + compact_memory + summary_memory(async) | -| `start` | 🚀 启动记忆系统 | 初始化文件存储、文件监控、Embedding 缓存；清理过期工具结果文件 | -| `close` | 📕 关闭并清理 | 清理工具结果文件、停止文件监控、保存 Embedding 缓存 |· +| 方法 | 功能 | 关键组件 | +|-----------------------|--------------|------------------------------------------------------------------------------------------------------------------------------| +| `check_context` | 📊 检查上下文大小 | [ContextChecker](reme/memory/file_based/components/context_checker.py) — 检查上下文是否超出阈值并拆分Message | +| `compact_memory` | 📦 压缩历史对话为摘要 | [Compactor](reme/memory/file_based/components/compactor.py) — ReActAgent 生成结构化上下文摘要 | +| `summary_memory` | 📝 将重要记忆写入文件 | [Summarizer](reme/memory/file_based/components/summarizer.py) — ReActAgent + 文件工具（read / write / edit） | +| `compact_tool_result` | ✂️ 压缩超长工具输出 | [ToolResultCompactor](reme/memory/file_based/components/tool_result_compactor.py) — 截断超长的工具调用结果并转存到 `tool_result/`，消息中保留文件引用 | +| `memory_search` | 🔍 语义搜索记忆 | [MemorySearch](reme/memory/file_based/tools/memory_search.py) — 向量 + BM25 混合检索 | +| `ReMeInMemoryMemory` | 🗂️ 会话内存类 | [ReMeInMemoryMemory](reme/memory/file_based/reme_in_memory_memory.py) — Token 感知的内存管理，支持压缩摘要和状态序列化 | +| `pre_reasoning_hook` | 🔄 推理前预处理钩子 | compact_tool_result + check_context + compact_memory + summary_memory(async) | +| `start` | 🚀 启动记忆系统 | 初始化文件存储、文件监控、Embedding 缓存；清理过期工具结果文件 | +| `close` | 📕 关闭并清理 | 清理工具结果文件、停止文件监控、保存 Embedding 缓存 |· + --- ### 🚀 快速开始 @@ -188,7 +189,7 @@ graph LR #### 1. check_context — 上下文检查 -[ContextChecker](reme/memory/file_based/component/context_checker.py) 基于 Token 计数判断上下文是否超限，自动拆分为「待压缩」和「保留」两组消息。 +[ContextChecker](reme/memory/file_based/components/context_checker.py) 基于 Token 计数判断上下文是否超限，自动拆分为「待压缩」和「保留」两组消息。 ```mermaid graph LR @@ -208,7 +209,7 @@ graph LR #### 2. compact_memory — 对话压缩 -[Compactor](reme/memory/file_based/component/compactor.py) 使用 ReActAgent 将历史对话压缩为**结构化上下文摘要**。 +[Compactor](reme/memory/file_based/components/compactor.py) 使用 ReActAgent 将历史对话压缩为**结构化上下文摘要**。 ```mermaid graph LR @@ -235,7 +236,7 @@ graph LR #### 3. summary_memory — 记忆持久化 -[Summarizer](reme/memory/file_based/component/summarizer.py) 采用 **ReAct + 文件工具** 模式，让 AI 自主决定写什么、写到哪。 +[Summarizer](reme/memory/file_based/components/summarizer.py) 采用 **ReAct + 文件工具** 模式，让 AI 自主决定写什么、写到哪。 ```mermaid graph LR @@ -260,7 +261,7 @@ graph LR #### 4. compact_tool_result — 工具结果压缩 -[ToolResultCompactor](reme/memory/file_based/component/tool_result_compactor.py) 解决工具输出过长导致上下文膨胀的问题。 +[ToolResultCompactor](reme/memory/file_based/components/tool_result_compactor.py) 解决工具输出过长导致上下文膨胀的问题。 ```mermaid graph LR @@ -518,7 +519,6 @@ Pass@K 衡量在生成 K 个候选中，至少一个成功完成任务（score=1 关于如何复现实验的更多细节，见 [quickstart.md](benchmark/bfcl/quickstart.md) - ## ⭐ 社区与支持 - **Star 与 Watch**：Star 可让更多智能体开发者发现 ReMe；Watch 可助你第一时间获知新版本与特性。 diff --git a/reme/core/schema/as_msg_stat.py b/reme/core/schema/as_msg_stat.py index 2e861863..b4ef02d4 100644 --- a/reme/core/schema/as_msg_stat.py +++ b/reme/core/schema/as_msg_stat.py @@ -26,6 +26,12 @@ class AsBlockStat(BaseModel): """Return a short preview of the block content.""" return self.format(_DEFAULT_MAX_BLOCK_TEXT_PREVIEW_LENGTH) + def _truncate(self, text: str, max_length: int) -> str: + """Simple truncation with ellipsis.""" + if len(text) <= max_length: + return text + return text[:max_length] + "..." + # pylint: disable=too-many-return-statements def format(self, max_length: int = _DEFAULT_MAX_FORMATTER_TEXT_LENGTH, include_thinking: bool = True) -> str: """Format block content to string representation. @@ -37,22 +43,25 @@ class AsBlockStat(BaseModel): Returns: Formatted string representation of the block. """ - from ..utils import truncate_text - if self.block_type == "text": - return truncate_text(self.text, max_length) if self.text else "" + if not self.text: + return "" + return f"{self._truncate(self.text, max_length)}" if self.block_type == "thinking": - if include_thinking and self.text: - return f"\n{truncate_text(self.text, max_length)}\n" - return "" + if not include_thinking or not self.text: + return "" + return f"{self._truncate(self.text, max_length)}" if self.block_type in ("image", "audio", "video"): - return f"[{self.block_type}] {self.media_url}" if self.media_url else f"[{self.block_type}]" - if self.block_type in ("tool_use", "tool_result"): - if self.block_type == "tool_use": - return f" - tool_call={self.tool_name} params={truncate_text(self.tool_input, max_length)}" - else: - output = truncate_text(self.tool_output, max_length) - return f" - tool_result={self.tool_name} output={output}" if output else "" + content = self.media_url if self.media_url else "" + return f"<{self.block_type}>{content}" + if self.block_type == "tool_use": + content = f"{self.tool_name} params={self._truncate(self.tool_input, max_length)}" + return f"{content}" + if self.block_type == "tool_result": + if not self.tool_output: + return "" + content = f"{self.tool_name} output={self._truncate(self.tool_output, max_length)}" + return f"{content}" return "" diff --git a/reme/memory/file_based/__init__.py b/reme/memory/file_based/__init__.py index 766f9b3e..e4c392e5 100644 --- a/reme/memory/file_based/__init__.py +++ b/reme/memory/file_based/__init__.py @@ -1,29 +1,13 @@ -"""File-based Memory Module. +"""File-based Memory Module.""" -This module provides memory management components for CoPaw (Cooperative Paw) agents, -including memory formatting, compaction, summarization, and file I/O operations. - -Components: - - ReMeInMemoryMemory: Extended InMemoryMemory with bugfixes and summary support - - AsMsgHandler: Handles AgentScope message statistics, formatting, and context checking - - Summarizer: Generates memory summaries using LLM - - Compactor: Compacts memory content to reduce token usage - - ToolResultCompactor: Truncates large tool results and saves full content to files - - ContextChecker: Checks context size and splits messages for compaction -""" - -from .as_msg_handler import AsMsgHandler -from .component.compactor import Compactor -from .component.context_checker import ContextChecker -from .component.summarizer import Summarizer -from .component.tool_result_compactor import ToolResultCompactor +from . import components +from . import tools +from . import utils from .reme_in_memory_memory import ReMeInMemoryMemory __all__ = [ - "AsMsgHandler", + "tools", + "utils", + "components", "ReMeInMemoryMemory", - "Summarizer", - "Compactor", - "ContextChecker", - "ToolResultCompactor", ] diff --git a/reme/memory/file_based/component/__init__.py b/reme/memory/file_based/component/__init__.py deleted file mode 100644 index e69de29b..00000000 diff --git a/reme/memory/file_based/components/__init__.py b/reme/memory/file_based/components/__init__.py new file mode 100644 index 00000000..42ab2f6b --- /dev/null +++ b/reme/memory/file_based/components/__init__.py @@ -0,0 +1,13 @@ +"""components""" + +from .compactor import Compactor +from .context_checker import ContextChecker +from .summarizer import Summarizer +from .tool_result_compactor import ToolResultCompactor + +__all__ = [ + "Compactor", + "Summarizer", + "ContextChecker", + "ToolResultCompactor", +] diff --git a/reme/memory/file_based/component/compactor.py b/reme/memory/file_based/components/compactor.py similarity index 89% rename from reme/memory/file_based/component/compactor.py rename to reme/memory/file_based/components/compactor.py index 3292c874..e7b37b30 100644 --- a/reme/memory/file_based/component/compactor.py +++ b/reme/memory/file_based/components/compactor.py @@ -4,7 +4,7 @@ from agentscope.agent import ReActAgent from agentscope.message import Msg from agentscope.token import HuggingFaceTokenCounter -from ..as_msg_handler import AsMsgHandler +from ..utils import AsMsgHandler from ....core.op import BaseOp from ....core.utils import get_std_logger @@ -32,10 +32,13 @@ class Compactor(BaseOp): if not messages: return "" + before_token_count = self.msg_handler.count_msgs_token(messages) history_formatted_str: str = self.msg_handler.format_msgs_to_str( messages=messages, memory_compact_threshold=self.memory_compact_threshold, ) + after_token_count = self.msg_handler.count_str_token(history_formatted_str) + logger.info(f"Compactor before_token_count={before_token_count} after_token_count={after_token_count}") if not history_formatted_str: logger.warning(f"No history to compact. messages={messages}") diff --git a/reme/memory/file_based/component/compactor.yaml b/reme/memory/file_based/components/compactor.yaml similarity index 100% rename from reme/memory/file_based/component/compactor.yaml rename to reme/memory/file_based/components/compactor.yaml diff --git a/reme/memory/file_based/component/context_checker.py b/reme/memory/file_based/components/context_checker.py similarity index 92% rename from reme/memory/file_based/component/context_checker.py rename to reme/memory/file_based/components/context_checker.py index 82bb381e..18ac4bf0 100644 --- a/reme/memory/file_based/component/context_checker.py +++ b/reme/memory/file_based/components/context_checker.py @@ -3,7 +3,7 @@ from agentscope.message import Msg from agentscope.token import HuggingFaceTokenCounter -from ..as_msg_handler import AsMsgHandler +from ..utils import AsMsgHandler from ....core.op import BaseOp from ....core.utils import get_std_logger @@ -87,11 +87,12 @@ class ContextChecker(BaseOp): memory_compact_reserve=self.memory_compact_reserve, ) - logger.info( - f"ContextChecker Result: " - f"to_compact={len(messages_to_compact)}, " - f"to_keep={len(messages_to_keep)}, " - f"is_valid={is_valid}", - ) + if messages_to_compact: + logger.info( + f"ContextChecker Result: " + f"to_compact={len(messages_to_compact)}, " + f"to_keep={len(messages_to_keep)}, " + f"is_valid={is_valid}", + ) return messages_to_compact, messages_to_keep, is_valid diff --git a/reme/memory/file_based/component/summarizer.py b/reme/memory/file_based/components/summarizer.py similarity index 88% rename from reme/memory/file_based/component/summarizer.py rename to reme/memory/file_based/components/summarizer.py index db3522da..d4e057be 100644 --- a/reme/memory/file_based/component/summarizer.py +++ b/reme/memory/file_based/components/summarizer.py @@ -7,7 +7,7 @@ from agentscope.message import Msg from agentscope.token import HuggingFaceTokenCounter from agentscope.tool import Toolkit -from ..as_msg_handler import AsMsgHandler +from ..utils import AsMsgHandler from ....core.op import BaseOp from ....core.utils import get_std_logger @@ -40,10 +40,13 @@ class Summarizer(BaseOp): if not messages: return "" + before_token_count = self.msg_handler.count_msgs_token(messages) history_formatted_str: str = self.msg_handler.format_msgs_to_str( messages=messages, memory_compact_threshold=self.memory_compact_threshold, ) + after_token_count = self.msg_handler.count_str_token(history_formatted_str) + logger.info(f"Summarizer before_token_count={before_token_count} after_token_count={after_token_count}") if not history_formatted_str: logger.warning(f"No history to summarize. messages={messages}") diff --git a/reme/memory/file_based/component/summarizer.yaml b/reme/memory/file_based/components/summarizer.yaml similarity index 100% rename from reme/memory/file_based/component/summarizer.yaml rename to reme/memory/file_based/components/summarizer.yaml diff --git a/reme/memory/file_based/component/tool_result_compactor.py b/reme/memory/file_based/components/tool_result_compactor.py similarity index 100% rename from reme/memory/file_based/component/tool_result_compactor.py rename to reme/memory/file_based/components/tool_result_compactor.py diff --git a/reme/memory/file_based/reme_in_memory_memory.py b/reme/memory/file_based/reme_in_memory_memory.py index 16f18726..da2118e8 100644 --- a/reme/memory/file_based/reme_in_memory_memory.py +++ b/reme/memory/file_based/reme_in_memory_memory.py @@ -5,7 +5,7 @@ from agentscope.memory import InMemoryMemory from agentscope.message import Msg from agentscope.token import HuggingFaceTokenCounter -from .as_msg_handler import AsMsgHandler +from .utils import AsMsgHandler from ...core.utils import get_std_logger logger = get_std_logger() diff --git a/reme/memory/file_based/utils/__init__.py b/reme/memory/file_based/utils/__init__.py new file mode 100644 index 00000000..6249ab15 --- /dev/null +++ b/reme/memory/file_based/utils/__init__.py @@ -0,0 +1,7 @@ +"""utils""" + +from .as_msg_handler import AsMsgHandler + +__all__ = [ + "AsMsgHandler", +] diff --git a/reme/memory/file_based/as_msg_handler.py b/reme/memory/file_based/utils/as_msg_handler.py similarity index 92% rename from reme/memory/file_based/as_msg_handler.py rename to reme/memory/file_based/utils/as_msg_handler.py index 9db6cac2..30facedc 100644 --- a/reme/memory/file_based/as_msg_handler.py +++ b/reme/memory/file_based/utils/as_msg_handler.py @@ -5,8 +5,8 @@ import json from agentscope.message import Msg from agentscope.token import HuggingFaceTokenCounter -from ...core.schema import AsMsgStat, AsBlockStat -from ...core.utils import get_std_logger +from ....core.schema import AsMsgStat, AsBlockStat +from ....core.utils import get_std_logger logger = get_std_logger() @@ -111,6 +111,19 @@ class AsMsgHandler: metadata=message.metadata or {}, ) + if not isinstance(message.content, list): + logger.warning( + "Unexpected message.content type %s, expected str or list, returning empty stat.", + type(message.content), + ) + return AsMsgStat( + name=message.name or message.role, + role=message.role, + content=blocks, + timestamp=message.timestamp or "", + metadata=message.metadata or {}, + ) + for block in message.content: block_type = block.get("type", "unknown") @@ -155,7 +168,7 @@ class AsMsgHandler: elif block_type == "tool_use": tool_name = block.get("name", "") - tool_input = block.get("raw_input", "") + tool_input = block.get("input", "") try: input_str = json.dumps(tool_input, ensure_ascii=False) except (TypeError, ValueError): @@ -227,7 +240,8 @@ class AsMsgHandler: formatted_content = stat.format(include_thinking=include_thinking) content_token_count = self.count_str_token(formatted_content) - if total_token_count + content_token_count > memory_compact_threshold: + is_latest = i == len(messages) - 1 + if not is_latest and total_token_count + content_token_count > memory_compact_threshold: logger.info( "Skipping older messages: adding %d tokens would exceed threshold %d (current: %d)", content_token_count, @@ -236,6 +250,13 @@ class AsMsgHandler: ) break + if is_latest and content_token_count > memory_compact_threshold: + logger.warning( + "Latest message alone (%d tokens) exceeds threshold %d, including it anyway.", + content_token_count, + memory_compact_threshold, + ) + formatted_parts.append(formatted_content) total_token_count += content_token_count @@ -324,6 +345,10 @@ class AsMsgHandler: accumulated_tokens = 0 for i in range(len(msg_stats) - 1, -1, -1): + # Skip messages already added as tool_use dependencies to avoid double-counting tokens + if i in keep_indices: + continue + msg, stat = msg_stats[i] # Check if adding this message would exceed reserve limit diff --git a/reme/reme_light.py b/reme/reme_light.py index 922b846b..94309eb0 100644 --- a/reme/reme_light.py +++ b/reme/reme_light.py @@ -26,16 +26,15 @@ from agentscope.tool import Toolkit, ToolResponse from .config import ReMeConfigParser from .core import Application from .core.utils import get_hf_token_counter, get_std_logger -from .memory.file_based import ( +from .memory.file_based import ReMeInMemoryMemory +from .memory.file_based.components import ( Compactor, ContextChecker, Summarizer, ToolResultCompactor, - ReMeInMemoryMemory, - AsMsgHandler, ) -from .memory.file_based import MemorySearch -from .memory.file_based.tools import FileIO +from .memory.file_based.tools import FileIO, MemorySearch +from .memory.file_based.utils import AsMsgHandler logger = get_std_logger() @@ -487,7 +486,7 @@ class ReMeLight(Application): - compact_ratio: Compaction threshold ratio Note: - - Completed/failed/cancelled tasks are cleaned up before adding new ones + - Completed/failed/canceled tasks are cleaned up before adding new ones - Task results and errors are logged automatically - Use await_summary_tasks() to wait for all pending tasks to complete """ @@ -561,7 +560,7 @@ class ReMeLight(Application): Returns: tuple[list[Msg], str]: A tuple containing: - - list[Msg]: Messages to keep in context (may be reduced) + - list[Msg]: Messages to keep in context (maybe reduced) - str: Updated compressed summary incorporating compacted messages Note: @@ -584,10 +583,11 @@ class ReMeLight(Application): compact_msgs = messages[:-tool_result_compact_keep_n] await self.compact_tool_result(compact_msgs) - messages_to_compact, messages_to_keep, is_valid = msg_handler.context_check( + messages_to_compact, messages_to_keep, is_valid = await self.check_context( messages=messages, memory_compact_threshold=left_compact_threshold, memory_compact_reserve=memory_compact_reserve, + token_counter=token_counter, ) if not messages_to_compact: @@ -626,7 +626,7 @@ class ReMeLight(Application): Wait for all background summary tasks to complete and collect results. Blocks until all pending summary tasks in the task list have completed, - cancelled, or failed. Collects status information from each task and + canceled, or failed. Collects status information from each task and clears the task list after processing. Returns: diff --git a/tests/light/test_compactor.py b/tests/light/test_compactor.py index 8ae2a051..cb3c9dee 100644 --- a/tests/light/test_compactor.py +++ b/tests/light/test_compactor.py @@ -3,7 +3,6 @@ import asyncio from agentscope.message import Msg - from test_utils import ( get_dash_chat_model, get_formatter, @@ -11,8 +10,7 @@ from test_utils import ( ) from reme.core.utils import get_std_logger -from reme.memory.file_based import Compactor - +from reme.memory.file_based.components import Compactor logger = get_std_logger() diff --git a/tests/light/test_context_check.py b/tests/light/test_context_check.py index 872d6e2d..934a43f2 100644 --- a/tests/light/test_context_check.py +++ b/tests/light/test_context_check.py @@ -1,10 +1,10 @@ """Tests for AsMsgHandler.context_check method.""" from agentscope.message import Msg - from test_utils import get_token_counter + from reme.core.utils import get_std_logger -from reme.memory.file_based.as_msg_handler import AsMsgHandler +from reme.memory.file_based.utils import AsMsgHandler logger = get_std_logger() diff --git a/tests/light/test_format_msgs_to_str.py b/tests/light/test_format_msgs_to_str.py index 93dd7a2f..64e97330 100644 --- a/tests/light/test_format_msgs_to_str.py +++ b/tests/light/test_format_msgs_to_str.py @@ -5,10 +5,10 @@ import sys from agentscope.message import Msg - from test_utils import get_token_counter + from reme.core.utils import get_std_logger -from reme.memory.file_based.as_msg_handler import AsMsgHandler +from reme.memory.file_based.utils import AsMsgHandler logger = get_std_logger() diff --git a/tests/light/test_summarizer.py b/tests/light/test_summarizer.py index 9718cd2a..2789bca5 100644 --- a/tests/light/test_summarizer.py +++ b/tests/light/test_summarizer.py @@ -7,16 +7,16 @@ from pathlib import Path from agentscope.message import Msg from agentscope.tool import Toolkit - from test_utils import ( get_dash_chat_model, get_formatter, get_token_counter, ) from reme.core.utils import get_std_logger -from reme.memory.file_based import Summarizer +from reme.memory.file_based.components import Summarizer from reme.memory.file_based.tools import FileIO + logger = get_std_logger() diff --git a/tests/light/test_tool_result_compactor.py b/tests/light/test_tool_result_compactor.py index b6cb7c69..1697911d 100644 --- a/tests/light/test_tool_result_compactor.py +++ b/tests/light/test_tool_result_compactor.py @@ -6,8 +6,9 @@ from datetime import datetime, timedelta from pathlib import Path from agentscope.message import Msg -from reme.memory.file_based import ToolResultCompactor + from reme.core.utils import is_truncated +from reme.memory.file_based.components import ToolResultCompactor def create_tool_result_msg(output: str | list, tool_name: str = "test_tool") -> Msg: diff --git a/tests/light/test_tools.py b/tests/light/test_tools.py index 933d506d..891fa534 100644 --- a/tests/light/test_tools.py +++ b/tests/light/test_tools.py @@ -9,8 +9,8 @@ import tempfile import pytest -from reme.memory.file_based.tools.shell import Shell from reme.memory.file_based.tools.file_io import FileIO +from reme.memory.file_based.tools.shell import Shell from reme.memory.file_based.tools.utils import DEFAULT_MAX_LINES, DEFAULT_MAX_BYTES diff --git a/tests/light/test_utils.py b/tests/light/test_utils.py index 19740fe6..27583948 100644 --- a/tests/light/test_utils.py +++ b/tests/light/test_utils.py @@ -4,7 +4,7 @@ import os from agentscope.message import Msg, ThinkingBlock, TextBlock, ToolUseBlock, ToolResultBlock -from reme.memory.file_based import AsMsgHandler +from reme.memory.file_based.utils import AsMsgHandler def get_token_counter():