.txt
```
-### Core Capabilities
+### Core capabilities
-[ReMeLight](reme/reme_light.py) is the core class of this memory system, providing complete memory management
-capabilities for AI Agents:
+[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 |
-|------------------------|------------------------------------|-------------------------------------------------------------------------------------------------------------------------------------------------------------|
-| `start` | 🚀 Start memory system | Initialize file store, file watcher, Embedding cache; clean up expired tool result files |
-| `close` | 📕 Close and clean up | Clean tool result files, stop file watcher, save Embedding cache |
-| `compact_memory` | 📦 Compact history to summary | [Compactor](reme/memory/file_based/compactor.py) — ReActAgent generates structured context checkpoint |
-| `summary_memory` | 📝 Write important memory to files | [Summarizer](reme/memory/file_based/summarizer.py) — ReActAgent + file tools (read / write / edit) |
-| `compact_tool_result` | ✂️ Compact oversized tool output | [ToolResultCompactor](reme/memory/file_based/tool_result_compactor.py) — Truncate and save to `tool_result/`, keep file reference in message |
-| `memory_search` | 🔍 Semantic memory search | [MemorySearch](reme/memory/tools/chunk/memory_search.py) — Vector + BM25 hybrid retrieval |
-| `get_in_memory_memory` | 🗂️ Create in-memory instance | [ReMeInMemoryMemory](reme/memory/file_based/reme_in_memory_memory.py) — Token-aware memory management, supports compression summary and state serialization |
+
+Category Method Function Key components Context Management check_context📊 Check context size ContextChecker — checks whether context exceeds thresholds and splits messagescompact_memory📦 Compact history into summary Compactor — ReActAgent that generates structured context summariescompact_tool_result✂️ Compact long tool outputs ToolResultCompactor — truncates long tool outputs and stores them in tool_result/ while keeping file references in messagespre_reasoning_hook🔄 Pre-reasoning hook compact_tool_result + check_context + compact_memory + summary_memory (async)Long-term Memory summary_memory📝 Persist important memory to files Summarizer — ReActAgent + file tools (read / write / edit)memory_search🔍 Semantic memory search MemorySearch — hybrid retrieval with vectors + BM25Session Memory get_in_memory_memory💾 Create in-session memory instance Returns ReMeInMemoryMemory with dialog_path configured for persistence await_summary_tasks⏳ Wait for async summary tasks Block until all background summary tasks complete - 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
---
-### 🚀 Quick Start
+### 🚀 Quick start
#### Installation
+**Install from source:**
+
```bash
+git clone https://github.com/agentscope-ai/ReMe.git
+cd ReMe
pip install -e ".[light]"
```
-#### Environment Variables
+**Update to the latest version:**
-`ReMeLight` environment variables configure Embedding and storage backend
+```bash
+git pull
+pip install -e ".[light]"
+```
-| Variable | Description | Example |
-|----------------------|--------------------------------|-----------------------------------------------------|
-| `LLM_API_KEY` | LLM API key | `sk-xxx` |
-| `LLM_BASE_URL` | LLM base URL | `https://dashscope.aliyuncs.com/compatible-mode/v1` |
-| `EMBEDDING_API_KEY` | Embedding API key (Optional) | `sk-xxx` |
-| `EMBEDDING_BASE_URL` | Embedding base URL (Optional) | `https://dashscope.aliyuncs.com/compatible-mode/v1` |
-| `LLM_MODEL_NAME` | LLM model name | `qwen3.5-plus` |
+#### Environment variables
-#### Python Usage
+`ReMeLight` uses environment variables to configure the embedding model and storage backends:
+
+| Variable | Description | Example |
+|----------------------|-------------------------------|-----------------------------------------------------|
+| `LLM_API_KEY` | LLM API key | `sk-xxx` |
+| `LLM_BASE_URL` | LLM base URL | `https://dashscope.aliyuncs.com/compatible-mode/v1` |
+| `EMBEDDING_API_KEY` | Embedding API key (optional) | `sk-xxx` |
+| `EMBEDDING_BASE_URL` | Embedding base URL (optional) | `https://dashscope.aliyuncs.com/compatible-mode/v1` |
+
+#### Python usage
```python
import asyncio
-from agentscope.message import Msg
from reme.reme_light import ReMeLight
async def main():
+ # Initialize ReMeLight
reme = ReMeLight(
- working_dir=".reme", # Memory file storage directory
- max_input_length=128000, # Model context window (tokens)
- memory_compact_ratio=0.7, # Trigger compaction when reaching max_input_length * 0.7
- language="zh", # Summary language (zh / "")
- tool_result_threshold=1000, # Auto-save tool outputs exceeding this character count
- retention_days=7, # tool_result/ file retention days
+ default_as_llm_config={"model_name": "qwen3.5-35b-a3b"},
+ # default_embedding_model_config={"model_name": "text-embedding-v4"},
+ default_file_store_config={"fts_enabled": True, "vector_enabled": False},
+ enable_load_env=True,
)
await reme.start()
- messages = [...]
+ messages = [...] # List of conversation messages
- # 1. Compact oversized tool outputs (prevent tool results from overflowing context)
+ # 1. Check context size (token counting, determine if compaction is needed)
+ messages_to_compact, messages_to_keep, is_valid = await reme.check_context(
+ messages=messages,
+ memory_compact_threshold=90000, # Threshold to trigger compaction (tokens)
+ memory_compact_reserve=10000, # Token count to reserve for recent messages
+ )
+
+ # 2. Compact conversation history into a structured summary
+ summary = await reme.compact_memory(
+ messages=messages,
+ previous_summary="",
+ max_input_length=128000, # Model context window (tokens)
+ compact_ratio=0.7, # Trigger compaction when exceeding max_input_length * 0.7
+ language="zh", # Summary language (e.g., "zh" / "")
+ )
+
+ # 3. Compact long tool outputs (prevent tool results from blowing up context)
messages = await reme.compact_tool_result(messages)
- # 2. Compact history to structured summary (trigger: context approaching limit), can pass previous summary for incremental update
- summary = await reme.compact_memory(messages=messages, previous_summary="")
+ # 4. Pre-reasoning hook (auto compact tool results + check context + generate summaries)
+ processed_messages, compressed_summary = await reme.pre_reasoning_hook(
+ messages=messages,
+ system_prompt="You are a helpful AI assistant.",
+ compressed_summary="",
+ max_input_length=128000,
+ compact_ratio=0.7,
+ memory_compact_reserve=10000,
+ enable_tool_result_compact=True,
+ tool_result_compact_keep_n=3,
+ )
- # 3. Submit async summary task in background (non-blocking, writes to memory/YYYY-MM-DD.md)
- reme.add_async_summary_task(messages=messages)
+ # 5. Persist important memory to files (writes to memory/YYYY-MM-DD.md)
+ summary_result = await reme.summary_memory(
+ messages=messages,
+ language="zh",
+ )
- # 4. Semantic memory search (Vector + BM25 hybrid retrieval)
+ # 6. Semantic memory search (vector + BM25 hybrid retrieval)
result = await reme.memory_search(query="Python version preference", max_results=5)
- # 5. Get in-memory instance (ReMeInMemoryMemory, manages single conversation context) AgentScope InMemoryMemory
- memory = reme.get_in_memory_memory()
- token_stats = await memory.estimate_tokens()
+ # 7. Create in-session memory instance (manages context for one conversation)
+ memory = reme.get_in_memory_memory() # Auto-configures dialog_path
+ for msg in messages:
+ await memory.add(msg)
+ token_stats = await memory.estimate_tokens(max_input_length=128000)
print(f"Current context usage: {token_stats['context_usage_ratio']:.1f}%")
- print(f"Message tokens: {token_stats['messages_tokens']}")
+ print(f"Message token count: {token_stats['messages_tokens']}")
print(f"Estimated total tokens: {token_stats['estimated_tokens']}")
- # 6. Wait for background tasks before closing
- summary_result = await reme.await_summary_tasks()
+ # 8. Mark messages as compressed (auto-persists to dialog/YYYY-MM-DD.jsonl)
+ # await memory.mark_messages_compressed(messages_to_compact)
- # Close ReMeLight
+ # Shutdown ReMeLight
await reme.close()
@@ -150,169 +225,266 @@ if __name__ == "__main__":
asyncio.run(main())
```
-### File-Based ReMeLight Memory System Architecture
+> 📂 Full example: [test_reme_light.py](tests/light/test_reme_light.py)
+> 📋 Sample run log: [test_reme_light_log.txt](tests/light/test_reme_light_log.txt) (223,838 tokens → 1,105 tokens, 99.5%
+> compression)
-[CoPaw MemoryManager](https://github.com/agentscope-ai/CoPaw/blob/main/src/copaw/agents/memory/memory_manager.py)
-inherits `ReMeLight` and integrates memory capabilities into the Agent reasoning flow:
+### Architecture of the file-based ReMeLight memory system
+
+#### Context data structure
```mermaid
-graph TB
- CoPaw["CoPaw MemoryManager\n(inherits ReMeLight)"] -->|pre_reasoning hook| Hook[MemoryCompactionHook]
- CoPaw --> ReMeLight[ReMeLight]
- Hook -->|exceeds threshold| ReMeLight
- ReMeLight --> CompactMemory[compact_memory\nHistory compaction]
- ReMeLight --> SummaryMemory[summary_memory\nWrite memory to files]
- ReMeLight --> CompactToolResult[compact_tool_result\nOversized tool output compaction]
- ReMeLight --> MemSearch[memory_search\nSemantic search]
- ReMeLight --> InMemory[get_in_memory_memory\nReMeInMemoryMemory]
- CompactMemory --> Compactor[Compactor\nReActAgent]
- SummaryMemory --> Summarizer[Summarizer\nReActAgent + file tools]
- CompactToolResult --> ToolResultCompactor[ToolResultCompactor\nTruncate + save to file]
- Summarizer --> FileIO[FileIO\nread / write / edit]
- FileIO --> MemoryFiles[memory/YYYY-MM-DD.md]
- ToolResultCompactor --> ToolResultFiles[tool_result/*.txt]
- MemoryFiles -.->|File change| FileWatcher[Async File Watcher]
- FileWatcher -->|Update index| FileStore[Local DB]
- MemSearch --> FileStore
-```
-
-### Context Compaction Mechanism
-
-#### Context Compaction
-
-[Compactor](reme/memory/file_based/compactor.py) uses ReActAgent to compact history into structured **context
-checkpoints**:
-
-| Field | Description |
-|-----------------------|-----------------------------------------------------|
-| `## Goal` | 🎯 User's objectives (can be multiple) |
-| `## Constraints` | ⚙️ Constraints and preferences mentioned by user |
-| `## Progress` | 📈 Completed / in progress / blocked tasks |
-| `## Key Decisions` | 🔑 Decisions made with brief reasons |
-| `## Next Steps` | 🗺️ Next action plan (ordered list) |
-| `## Critical Context` | 📌 File paths, function names, error messages, etc. |
-
-Supports **incremental updates**: when `previous_summary` is passed, automatically merges new conversation with old
-summary, preserving historical progress.
-
-#### Tool Result Compaction
-
-[ToolResultCompactor](reme/memory/file_based/tool_result_compactor.py) solves context overflow caused by oversized tool
-outputs (e.g., browser use):
-
-```mermaid
-graph LR
- A[tool_result message] --> B{Content length > threshold?}
- B -->|No| C[Keep as-is]
- B -->|Yes| D[Truncate to threshold characters]
- D --> E[Write full content to tool_result/uuid.txt]
- E --> F[Append file reference path to message]
-```
-
-Expired files (exceeding `retention_days`) are automatically cleaned up during `start` / `close` /
-`compact_tool_result`.
-
-### Memory Summary: ReAct + File Tools
-
-[Summarizer](reme/memory/file_based/summarizer.py) uses the **ReAct + file tools** pattern, letting AI autonomously
-decide what to write and where:
-
-```mermaid
-graph LR
- A[Receive conversation] --> B{Think: What's worth recording?}
- B --> C[Act: read memory/YYYY-MM-DD.md]
- C --> D{Think: How to merge with existing content?}
- D --> E[Act: edit to update file]
- E --> F{Think: Anything missing?}
- F -->|Yes| B
- F -->|No| G[Done]
-```
-
-[FileIO](reme/memory/file_based/file_io.py) provides file operation tools:
-
-| Tool | Function | Use case |
-|---------|--------------------------------|-----------------------------------------|
-| `read` | Read file content (line range) | View existing memory, avoid duplicates |
-| `write` | Overwrite file | Create new memory file or major rewrite |
-| `edit` | Replace after exact match | Append or modify specific sections |
-
-### In-Memory Session Management
-
-[ReMeInMemoryMemory](reme/memory/file_based/reme_in_memory_memory.py) extends AgentScope's `InMemoryMemory`:
-
-| Feature | Description |
-|----------------------------------|---------------------------------------------------------------------|
-| `get_memory` | Filter messages by mark, auto-prepend compression summary |
-| `estimate_tokens` | Precisely estimate current context token usage and ratio |
-| `get_history_str` | Generate human-readable conversation summary (with token stats) |
-| `state_dict` / `load_state_dict` | Support state serialization / deserialization (session persistence) |
-| `mark_messages_compressed` | Mark messages as compressed state |
-| `get_compressed_summary` | Get compressed summary content |
-
-### Memory Retrieval
-
-[MemorySearch](reme/memory/tools/chunk/memory_search.py) provides **vector + BM25 hybrid retrieval**:
-
-| Retrieval | Strength | Weakness |
-|---------------------|-------------------------------------------------|----------------------------------------|
-| **Vector semantic** | Captures similar meaning with different wording | Weaker on exact token match |
-| **BM25 full-text** | Strong exact token match | No synonym or paraphrase understanding |
-
-**Fusion**: Both retrieval paths are weighted and summed (vector 0.7 + BM25 0.3), so both natural-language queries and
-exact lookups get reliable results.
-
-```mermaid
-graph LR
- Q[Search query] --> V[Vector search × 0.7]
-Q --> B[BM25 × 0.3]
-V --> M[Dedupe + weighted merge]
-B --> M
-M --> R[Top-N results]
+flowchart TD
+ A[Context] --> B[compact_summary]
+ B --> C[dialog path guide + Goal/Constraints/Progress/KeyDecisions/NextSteps]
+ A --> E[messages: full dialogue history]
+ A --> F[File System Cache]
+ F --> G[dialog/YYYY-MM-DD.jsonl]
+ F --> H[tool_result/uuid.txt N-day TTL]
```
---
-## 🗃️ Vector-Based Memory System
+[CoPaw MemoryManager](https://github.com/agentscope-ai/CoPaw/blob/main/src/copaw/agents/memory/reme_light_memory_manager.py)
+inherits `ReMeLight` and integrates its memory capabilities into the agent reasoning loop:
-[ReMe Vector Based](reme/reme.py) is the core class for the vector-based memory system, supporting unified management of
-three memory types:
-
-| Memory Type | Purpose | Usage Context |
-|------------------------------|-----------------------------------------------------|---------------|
-| **Personal memory** | User preferences, habits | `user_name` |
-| **Task / procedural memory** | Task execution experience, success/failure patterns | `task_name` |
-| **Tool memory** | Tool usage experience, parameter tuning | `tool_name` |
-
-### Core Capabilities
-
-| Method | Function | Description |
-|--------------------|---------------------|-----------------------------------------------------------|
-| `summarize_memory` | 🧠 Summarize memory | Automatically extract and store memory from conversations |
-| `retrieve_memory` | 🔍 Retrieve memory | Retrieve relevant memory by query |
-| `add_memory` | ➕ Add memory | Manually add memory to vector store |
-| `get_memory` | 📖 Get memory | Fetch a single memory by ID |
-| `update_memory` | ✏️ Update memory | Update content or metadata of existing memory |
-| `delete_memory` | 🗑️ Delete memory | Delete specified memory |
-| `list_memory` | 📋 List memory | List memories with filtering and sorting |
-
-### Installation
-
-```bash
-pip install -U reme-ai
+```mermaid
+graph LR
+ Agent[Agent] -->|Before each reasoning step| Hook[pre_reasoning_hook]
+ Hook --> TC[compact_tool_result
+
---
@@ -454,10 +699,19 @@ graph TB
## ⚖️ License
-This project is open source under the Apache License 2.0. See the [LICENSE](./LICENSE) file for details.
+This project is open-sourced under the Apache License 2.0. See [LICENSE](./LICENSE) for details.
---
-## 📈 Star History
+## 🤔 Why ReMe?
+
+ReMe stands for **Remember Me** and **Refine Me**, symbolizing our goal to help AI agents "remember" users and "refine"
+themselves through interactions. We hope ReMe is not just a cold memory module, but a partner that truly helps agents
+understand users, accumulate experience, and continuously evolve.
+
+---
+
+## 📈 Star history
[](https://www.star-history.com/#agentscope-ai/ReMe&Date)
+
diff --git a/README_ZH.md b/README_ZH.md
index b45b261c..7e6ccd58 100644
--- a/README_ZH.md
+++ b/README_ZH.md
@@ -14,6 +14,11 @@
+
@@ -24,12 +29,36 @@
---
-🧠 ReMe 是一个专为 **AI 智能体** 打造的记忆管理框架，同时提供基于文件系统和基于向量库的记忆系统。
+## 📰 最新文章
+
+| 日期 | 标题 |
+|------------|----------------------------------------------------|
+| 2026-03-30 | [CoPaw 上下文管理设计解析](docs/copaw_context_design_zh.md) |
+
+---
+
+🧠 ReMe 是一个专为 **AI 智能体** 打造的记忆管理框架，同时提供基于[文件系统](#-基于文件的记忆系统-remelight)
+和基于[向量库](#-基于向量库的记忆系统)的记忆系统。
它解决智能体记忆的两类核心问题：**上下文窗口有限**（长对话时早期信息被截断或丢失）、**会话无状态**（新对话无法继承历史，每次从零开始）。
ReMe 让智能体拥有**真正的记忆力**——旧对话自动浓缩，重要信息持久保存，下次对话自动想起来。
+在 LoCoMo 与 HaluMem 基准测试中，ReMe 取得了领先结果，详见[实验效果](#实验效果)。
+
+
+你可以用 ReMe 做什么
---
@@ -38,7 +67,7 @@ ReMe 让智能体拥有**真正的记忆力**——旧对话自动浓缩，重
> 记忆即文件，文件即记忆
将**记忆视为文件**——可读、可编辑、可复制。
-[CoPaw](https://github.com/agentscope-ai/CoPaw)通过继承 `ReMeLight` 实现了长期记忆和上下文的管理。
+[CoPaw](https://github.com/agentscope-ai/CoPaw) 通过继承 `ReMeLight` 实现了长期记忆和上下文的管理。
| 传统记忆系统 | File Based ReMe |
|-----------|-----------------|
@@ -49,9 +78,11 @@ ReMe 让智能体拥有**真正的记忆力**——旧对话自动浓缩，重
```
working_dir/
-├── MEMORY.md # 长期记忆：用户偏好、项目配置等持久信息
+├── MEMORY.md # 长期记忆：用户偏好等持久信息
├── memory/
-│ └── YYYY-MM-DD.md # 每日摘要日志：对话结束后自动写入
+│ └── YYYY-MM-DD.md # 每日日记：对话结束后自动写入
+├── dialog/ # 原始对话记录：压缩前的完整对话
+│ └── YYYY-MM-DD.jsonl # 按日期存储的对话消息（JSONL 格式）
└── tool_result/ # 超长工具输出缓存（自动管理，超期自动清理）
└── .txt
```
@@ -60,15 +91,19 @@ working_dir/
[ReMeLight](reme/reme_light.py) 是该记忆系统的核心类，为 AI Agent 提供完整的记忆管理能力：
-| 方法 | 功能 | 关键组件 |
-|------------------------|--------------|----------------------------------------------------------------------------------------------------------|
-| `start` | 🚀 启动记忆系统 | 初始化文件存储、文件监控、Embedding 缓存；清理过期工具结果文件 |
-| `close` | 📕 关闭并清理 | 清理工具结果文件、停止文件监控、保存 Embedding 缓存 |
-| `compact_memory` | 📦 压缩历史对话为摘要 | [Compactor](reme/memory/file_based/compactor.py) — ReActAgent 生成结构化上下文检查点 |
-| `summary_memory` | 📝 将重要记忆写入文件 | [Summarizer](reme/memory/file_based/summarizer.py) — ReActAgent + 文件工具（read / write / edit） |
-| `compact_tool_result` | ✂️ 压缩超长工具输出 | [ToolResultCompactor](reme/memory/file_based/tool_result_compactor.py) — 截断并转存到 `tool_result/`，消息中保留文件引用 | |
-| `memory_search` | 🔍 语义搜索记忆 | [MemorySearch](reme/memory/tools/chunk/memory_search.py) — 向量 + BM25 混合检索 |
-| `get_in_memory_memory` | 🗂️ 创建会话内存实例 | [ReMeInMemoryMemory](reme/memory/file_based/reme_in_memory_memory.py) — Token 感知的内存管理，支持压缩摘要和状态序列化 |
+
+类别 方法 功能 关键组件 上下文管理 check_context📊 检查上下文大小 ContextChecker — 检查上下文是否超出阈值并拆分 Messagecompact_memory📦 压缩历史对话为摘要 Compactor — ReActAgent 生成结构化上下文摘要compact_tool_result✂️ 压缩超长工具输出 ToolResultCompactor — 截断超长的工具调用结果并转存到 tool_result/，消息中保留文件引用pre_reasoning_hook🔄 推理前预处理钩子 compact_tool_result + check_context + compact_memory + summary_memory(async) 长期记忆 summary_memory📝 将重要记忆写入文件 Summarizer — ReActAgent + 文件工具（read / write / edit）memory_search🔍 语义搜索记忆 MemorySearch — 向量 + BM25 混合检索会话内存 get_in_memory_memory💾 创建会话内存实例 返回 ReMeInMemoryMemory，自动配置 dialog_path 实现对话持久化 await_summary_tasks⏳ 等待异步摘要任务 阻塞等待所有后台摘要任务完成 - start🚀 启动记忆系统 初始化文件存储、文件监控、Embedding 缓存；清理过期工具结果文件 - close📕 关闭并清理 清理工具结果文件、停止文件监控、保存 Embedding 缓存
---
@@ -76,7 +111,18 @@ working_dir/
#### 安装
+**从源码安装：**
+
```bash
+git clone https://github.com/agentscope-ai/ReMe.git
+cd ReMe
+pip install -e ".[light]"
+```
+
+**更新到最新版本：**
+
+```bash
+git pull
pip install -e ".[light]"
```
@@ -84,57 +130,85 @@ pip install -e ".[light]"
`ReMeLight` 环境变量配置 Embedding 和存储后端
-| Variable | Description | Example |
-|----------------------|--------------------|-----------------------------------------------------|
-| `LLM_API_KEY` | LLM API key | `sk-xxx` |
-| `LLM_BASE_URL` | LLM base URL | `https://dashscope.aliyuncs.com/compatible-mode/v1` |
-| `EMBEDDING_API_KEY` | Embedding API key | `sk-xxx` |
-| `EMBEDDING_BASE_URL` | Embedding base URL | `https://dashscope.aliyuncs.com/compatible-mode/v1` |
-| `LLM_MODEL_NAME` | LLM model name | `qwen3.5-plus` |
+| Variable | Description | Example |
+|----------------------|-------------------------|-----------------------------------------------------|
+| `LLM_API_KEY` | LLM API key | `sk-xxx` |
+| `LLM_BASE_URL` | LLM base URL | `https://dashscope.aliyuncs.com/compatible-mode/v1` |
+| `EMBEDDING_API_KEY` | Embedding API key (可选) | `sk-xxx` |
+| `EMBEDDING_BASE_URL` | Embedding base URL (可选) | `https://dashscope.aliyuncs.com/compatible-mode/v1` |
-#### Python使用
+#### Python 使用
```python
import asyncio
-from agentscope.message import Msg
from reme.reme_light import ReMeLight
async def main():
+ # 初始化 ReMeLight
reme = ReMeLight(
- working_dir=".reme", # 记忆文件存储目录
- max_input_length=128000, # 模型上下文窗口（tokens）
- memory_compact_ratio=0.7, # 达到 max_input_length * 0.7 时触发压缩
- language="zh", # 摘要语言（zh / ""）
- tool_result_threshold=1000, # 超过此字符数的工具输出自动转存
- retention_days=7, # tool_result/ 文件保留天数
+ default_as_llm_config={"model_name": "qwen3.5-35b-a3b"},
+ # default_embedding_model_config={"model_name": "text-embedding-v4"},
+ default_file_store_config={"fts_enabled": True, "vector_enabled": False},
+ enable_load_env=True,
)
await reme.start()
- messages = [...]
+ messages = [...] # 对话消息列表
- # 1. 压缩超长工具输出（防止工具结果撑爆上下文）
+ # 1. 检查上下文大小（Token 计数，判断是否需要压缩）
+ messages_to_compact, messages_to_keep, is_valid = await reme.check_context(
+ messages=messages,
+ memory_compact_threshold=90000, # 触发压缩的阈值（tokens）
+ memory_compact_reserve=10000, # 保留的近期消息 token 数
+ )
+
+ # 2. 将历史对话压缩为结构化摘要（可传入上轮摘要，实现增量更新）
+ summary = await reme.compact_memory(
+ messages=messages,
+ previous_summary="",
+ max_input_length=128000, # 模型上下文窗口（tokens）
+ compact_ratio=0.7, # 达到 max_input_length * 0.7 时触发压缩
+ language="zh", # 摘要语言（zh / ""）
+ )
+
+ # 3. 压缩超长工具输出（防止工具结果撑爆上下文）
messages = await reme.compact_tool_result(messages)
- # 2. 将历史对话压缩为结构化摘要（触发时机：上下文接近上限），可传入上轮摘要，实现增量更新
- summary = await reme.compact_memory(messages=messages, previous_summary="")
+ # 4. 推理前预处理钩子（自动压缩工具结果 + 检查上下文 + 生成摘要）
+ processed_messages, compressed_summary = await reme.pre_reasoning_hook(
+ messages=messages,
+ system_prompt="你是一个有帮助的 AI 助手。",
+ compressed_summary="",
+ max_input_length=128000,
+ compact_ratio=0.7,
+ memory_compact_reserve=10000,
+ enable_tool_result_compact=True,
+ tool_result_compact_keep_n=3,
+ )
- # 3. 后台异步提交摘要任务（不阻塞对话，摘要写入 memory/YYYY-MM-DD.md）
- reme.add_async_summary_task(messages=messages)
+ # 5. 将重要记忆写入文件（摘要写入 memory/YYYY-MM-DD.md）
+ summary_result = await reme.summary_memory(
+ messages=messages,
+ language="zh",
+ )
- # 4. 语义搜索记忆（向量 + BM25 混合检索）
+ # 6. 语义搜索记忆（向量 + BM25 混合检索）
result = await reme.memory_search(query="Python 版本偏好", max_results=5)
- # 5. 获取会话内存实例（ReMeInMemoryMemory，管理单次对话的上下文）AgentScope InMemoryMemory
- memory = reme.get_in_memory_memory()
- token_stats = await memory.estimate_tokens()
+ # 7. 创建会话内存实例（管理单次对话的上下文）
+ from reme.memory.file_based.reme_in_memory_memory import ReMeInMemoryMemory
+ memory = reme.get_in_memory_memory() # 自动配置 dialog_path
+ for msg in messages:
+ await memory.add(msg)
+ token_stats = await memory.estimate_tokens(max_input_length=128000)
print(f"当前上下文使用率: {token_stats['context_usage_ratio']:.1f}%")
print(f"消息 Token 数: {token_stats['messages_tokens']}")
print(f"预估总 Token 数: {token_stats['estimated_tokens']}")
- # 6. 关闭前等待后台任务完成
- summary_result = await reme.await_summary_tasks()
+ # 8. 标记消息为压缩状态（自动持久化到 dialog/YYYY-MM-DD.jsonl）
+ # await memory.mark_messages_compressed(messages_to_compact)
# 关闭 ReMeLight
await reme.close()
@@ -144,131 +218,235 @@ if __name__ == "__main__":
asyncio.run(main())
```
+> 📂 完整示例代码：[test_reme_light.py](tests/light/test_reme_light.py)
+> 📋 运行结果示例：[test_reme_light_log.txt](tests/light/test_reme_light_log.txt)（223,838 tokens → 1,105 tokens，压缩率99.5%）
+
### 基于文件的 ReMeLight 记忆系统架构
-[CoPaw MemoryManager](https://github.com/agentscope-ai/CoPaw/blob/main/src/copaw/agents/memory/memory_manager.py) 继承
+#### 上下文数据结构
+
+```mermaid
+flowchart TD
+ A[Context] --> B[compact_summary]
+ B --> C[dialog 路径引导 + Goal/Constraints/Progress/KeyDecisions/NextSteps]
+ A --> E[messages: 完整对话历史]
+ A --> F[文件系统缓存]
+ F --> G[dialog/YYYY-MM-DD.jsonl]
+ F --> H[tool_result/uuid.txt N天TTL]
+```
+
+---
+
+[CoPaw MemoryManager](https://github.com/agentscope-ai/CoPaw/blob/main/src/copaw/agents/memory/reme_light_memory_manager.py)
+继承
`ReMeLight`，将记忆能力集成到 Agent 推理流程中：
```mermaid
-graph TB
- CoPaw["CoPaw MemoryManager\n(继承 ReMeLight)"] -->|pre_reasoning hook| Hook[MemoryCompactionHook]
- CoPaw --> ReMeLight[ReMeLight]
- Hook -->|超出阈值| ReMeLight
- ReMeLight --> CompactMemory[compact_memory\n历史对话压缩]
- ReMeLight --> SummaryMemory[summary_memory\n记忆写入文件]
- ReMeLight --> CompactToolResult[compact_tool_result\n超长工具输出压缩]
- ReMeLight --> MemSearch[memory_search\n语义搜索]
- ReMeLight --> InMemory[get_in_memory_memory\nReMeInMemoryMemory]
- CompactMemory --> Compactor[Compactor\nReActAgent]
- SummaryMemory --> Summarizer[Summarizer\nReActAgent + 文件工具]
- CompactToolResult --> ToolResultCompactor[ToolResultCompactor\n截断 + 转存文件]
- Summarizer --> FileIO[FileIO\nread / write / edit]
- FileIO --> MemoryFiles[memory/YYYY-MM-DD.md]
- ToolResultCompactor --> ToolResultFiles[tool_result/*.txt]
- MemoryFiles -.->|文件变更| FileWatcher[异步文件监控]
- FileWatcher -->|更新索引| FileStore[本地数据库]
- MemSearch --> FileStore
+graph LR
+ Agent[Agent] -->|每轮推理前| Hook[pre_reasoning_hook]
+ Hook --> TC[compact_tool_result
+
+
---
## 📄 引用
@@ -441,6 +679,13 @@ graph TB
---
+## 🤔 为什么叫 ReMe？
+
+ReMe 是 **Remember Me** 和 **Refine Me** 的缩写，寓意让 AI 智能体「记住我」并在交互中「精进自我」。我们希望 ReMe
+不只是一个冷冰冰的记忆模块，而是能让智能体真正理解用户、积累经验、持续进化的伙伴。
+
+---
+
## 📈 Star 历史
[](https://www.star-history.com/#agentscope-ai/ReMe&Date)
diff --git a/docs/cookbook/appworld/quickstart.md b/benchmark/appworld/quickstart.md
similarity index 80%
rename from docs/cookbook/appworld/quickstart.md
rename to benchmark/appworld/quickstart.md
index 78f6a1a8..45ea0d78 100644
--- a/docs/cookbook/appworld/quickstart.md
+++ b/benchmark/appworld/quickstart.md
@@ -9,7 +9,7 @@ This guide helps you quickly set up and run AppWorld experiments with ReMe integ
```bash
git clone https://github.com/agentscope-ai/ReMe.git
-cd ReMe/cookbook/appworld
+cd ReMe/benchmark/appworld
```
### 2. Appworld Environment Setup
@@ -56,26 +56,16 @@ pip install .
Launch the ReMe service to enable memory library functionality:
```bash
-reme \
+reme2 \
backend=http \
http.port=8002 \
- llm.default.model_name=qwen-max-latest \
- embedding_model.default.model_name=text-embedding-v4 \
- vector_store.default.backend=elasticsearch
+ llms.default.model_name=qwen3-8b \
+ embedding_models.default.model_name=text-embedding-v4 \
+ vector_stores.default.backend=es \
+ vector_stores.default.collection_name=appworld \
+ vector_stores.default.hosts=http://xx.yy.zz.mm:nn
```
-add memories for appworld:
-```bash
-curl -X POST "http://0.0.0.0:8002/vector_store" \
- -H "Content-Type: application/json" \
- -d '{
- "workspace_id": "appworld",
- "action": "load",
- "path": "./docs/library"
- }'
-```
-Now you have loaded the ReMe memory library to enable memory-based agent!
-
### 4. Common Issues
**AppWorld data not found**: Ensure `appworld download data` completed successfully
@@ -95,21 +85,21 @@ python run_appworld.py
```
**What this does:**
-- Runs AppWorld tasks on the development dataset
+- Runs AppWorld tasks on the test-normal set
- Compares agent performance with ReMe memory (`use_memory=True`) vs without memory
- Uses multiple workers for parallel processing
- Runs each task multiple times for statistical significance
- Results are automatically saved to `./exp_result/` directory
**Configuration options in `run_appworld.py`:**
-- `max_workers`: Number of parallel workers (default: 8)
-- `num_runs`: Number of times each task is repeated (default: 1)
+- `max_workers`: Number of parallel workers (default: 16)
+- `num_runs`: Number of times each task is repeated (default: 4)
- `batch_size`: Number of concurrent tasks per batch (default: 8)
- `num_trials`: Maximum number of self-reflections, failure-aware reflection mechanism is triggered when num_trials>1 (default: 1)
-- `model_name`: Task execution model
-- `use_memory`: Whether to use ReMe memory library
-- `use_memory_addition`: Whether to enable selective addition
-- `use_memory_deletion`: Whether to enable utility-based deletion
+- `model_name`: Task execution model (default: "qwen3-8b")
+- `use_memory`: Whether to use ReMe memory library (default: True)
+- `use_memory_addition`: Whether to enable selective addition (default: False)
+- `use_memory_deletion`: Whether to enable utility-based deletion (default: False)
### 2. View Experiment Results
diff --git a/benchmark/bfcl/default_ids.py b/benchmark/bfcl/default_ids.py
new file mode 100644
index 00000000..43f4065e
--- /dev/null
+++ b/benchmark/bfcl/default_ids.py
@@ -0,0 +1,206 @@
+# pylint: disable=C0114
+DEFAULT_TRAIN_IDS: set[str] = {
+ "multi_turn_base_102",
+ "multi_turn_base_107",
+ "multi_turn_base_110",
+ "multi_turn_base_114",
+ "multi_turn_base_115",
+ "multi_turn_base_118",
+ "multi_turn_base_122",
+ "multi_turn_base_123",
+ "multi_turn_base_128",
+ "multi_turn_base_13",
+ "multi_turn_base_130",
+ "multi_turn_base_132",
+ "multi_turn_base_133",
+ "multi_turn_base_143",
+ "multi_turn_base_144",
+ "multi_turn_base_146",
+ "multi_turn_base_15",
+ "multi_turn_base_158",
+ "multi_turn_base_169",
+ "multi_turn_base_17",
+ "multi_turn_base_172",
+ "multi_turn_base_176",
+ "multi_turn_base_182",
+ "multi_turn_base_187",
+ "multi_turn_base_197",
+ "multi_turn_base_199",
+ "multi_turn_base_22",
+ "multi_turn_base_23",
+ "multi_turn_base_24",
+ "multi_turn_base_36",
+ "multi_turn_base_40",
+ "multi_turn_base_44",
+ "multi_turn_base_47",
+ "multi_turn_base_48",
+ "multi_turn_base_5",
+ "multi_turn_base_51",
+ "multi_turn_base_59",
+ "multi_turn_base_63",
+ "multi_turn_base_65",
+ "multi_turn_base_66",
+ "multi_turn_base_67",
+ "multi_turn_base_68",
+ "multi_turn_base_70",
+ "multi_turn_base_75",
+ "multi_turn_base_77",
+ "multi_turn_base_78",
+ "multi_turn_base_79",
+ "multi_turn_base_81",
+ "multi_turn_base_83",
+ "multi_turn_base_93",
+}
+
+DEFAULT_VAL_IDS: set[str] = {
+ "multi_turn_base_0",
+ "multi_turn_base_1",
+ "multi_turn_base_10",
+ "multi_turn_base_100",
+ "multi_turn_base_101",
+ "multi_turn_base_103",
+ "multi_turn_base_104",
+ "multi_turn_base_105",
+ "multi_turn_base_106",
+ "multi_turn_base_108",
+ "multi_turn_base_109",
+ "multi_turn_base_11",
+ "multi_turn_base_111",
+ "multi_turn_base_112",
+ "multi_turn_base_113",
+ "multi_turn_base_116",
+ "multi_turn_base_117",
+ "multi_turn_base_119",
+ "multi_turn_base_12",
+ "multi_turn_base_120",
+ "multi_turn_base_121",
+ "multi_turn_base_124",
+ "multi_turn_base_125",
+ "multi_turn_base_126",
+ "multi_turn_base_127",
+ "multi_turn_base_129",
+ "multi_turn_base_131",
+ "multi_turn_base_134",
+ "multi_turn_base_135",
+ "multi_turn_base_136",
+ "multi_turn_base_137",
+ "multi_turn_base_138",
+ "multi_turn_base_139",
+ "multi_turn_base_14",
+ "multi_turn_base_140",
+ "multi_turn_base_141",
+ "multi_turn_base_142",
+ "multi_turn_base_145",
+ "multi_turn_base_147",
+ "multi_turn_base_148",
+ "multi_turn_base_149",
+ "multi_turn_base_150",
+ "multi_turn_base_151",
+ "multi_turn_base_152",
+ "multi_turn_base_153",
+ "multi_turn_base_154",
+ "multi_turn_base_155",
+ "multi_turn_base_156",
+ "multi_turn_base_157",
+ "multi_turn_base_159",
+ "multi_turn_base_16",
+ "multi_turn_base_160",
+ "multi_turn_base_161",
+ "multi_turn_base_162",
+ "multi_turn_base_163",
+ "multi_turn_base_164",
+ "multi_turn_base_165",
+ "multi_turn_base_166",
+ "multi_turn_base_167",
+ "multi_turn_base_168",
+ "multi_turn_base_170",
+ "multi_turn_base_171",
+ "multi_turn_base_173",
+ "multi_turn_base_174",
+ "multi_turn_base_175",
+ "multi_turn_base_177",
+ "multi_turn_base_178",
+ "multi_turn_base_179",
+ "multi_turn_base_18",
+ "multi_turn_base_180",
+ "multi_turn_base_181",
+ "multi_turn_base_183",
+ "multi_turn_base_184",
+ "multi_turn_base_185",
+ "multi_turn_base_186",
+ "multi_turn_base_188",
+ "multi_turn_base_189",
+ "multi_turn_base_19",
+ "multi_turn_base_190",
+ "multi_turn_base_191",
+ "multi_turn_base_192",
+ "multi_turn_base_193",
+ "multi_turn_base_194",
+ "multi_turn_base_195",
+ "multi_turn_base_196",
+ "multi_turn_base_198",
+ "multi_turn_base_2",
+ "multi_turn_base_20",
+ "multi_turn_base_21",
+ "multi_turn_base_25",
+ "multi_turn_base_26",
+ "multi_turn_base_27",
+ "multi_turn_base_28",
+ "multi_turn_base_29",
+ "multi_turn_base_3",
+ "multi_turn_base_30",
+ "multi_turn_base_31",
+ "multi_turn_base_32",
+ "multi_turn_base_33",
+ "multi_turn_base_34",
+ "multi_turn_base_35",
+ "multi_turn_base_37",
+ "multi_turn_base_38",
+ "multi_turn_base_39",
+ "multi_turn_base_4",
+ "multi_turn_base_41",
+ "multi_turn_base_42",
+ "multi_turn_base_43",
+ "multi_turn_base_45",
+ "multi_turn_base_46",
+ "multi_turn_base_49",
+ "multi_turn_base_50",
+ "multi_turn_base_52",
+ "multi_turn_base_53",
+ "multi_turn_base_54",
+ "multi_turn_base_55",
+ "multi_turn_base_56",
+ "multi_turn_base_57",
+ "multi_turn_base_58",
+ "multi_turn_base_6",
+ "multi_turn_base_60",
+ "multi_turn_base_61",
+ "multi_turn_base_62",
+ "multi_turn_base_64",
+ "multi_turn_base_69",
+ "multi_turn_base_7",
+ "multi_turn_base_71",
+ "multi_turn_base_72",
+ "multi_turn_base_73",
+ "multi_turn_base_74",
+ "multi_turn_base_76",
+ "multi_turn_base_8",
+ "multi_turn_base_80",
+ "multi_turn_base_82",
+ "multi_turn_base_84",
+ "multi_turn_base_85",
+ "multi_turn_base_86",
+ "multi_turn_base_87",
+ "multi_turn_base_88",
+ "multi_turn_base_89",
+ "multi_turn_base_9",
+ "multi_turn_base_90",
+ "multi_turn_base_91",
+ "multi_turn_base_92",
+ "multi_turn_base_94",
+ "multi_turn_base_95",
+ "multi_turn_base_96",
+ "multi_turn_base_97",
+ "multi_turn_base_98",
+ "multi_turn_base_99",
+}
diff --git a/benchmark/bfcl/init_task_memory_pool.py b/benchmark/bfcl/init_task_memory_pool.py
index bc38046f..2a4fcc87 100644
--- a/benchmark/bfcl/init_task_memory_pool.py
+++ b/benchmark/bfcl/init_task_memory_pool.py
@@ -114,6 +114,9 @@ def post_to_summarizer(trajectories: List[Any], service_url: str) -> Dict[str, A
request_data = {
"trajectories": trajectory_dicts,
+ "success_threshold": 1.0,
+ "enable_soft_comparison": True,
+ "validation_threshold": 0.5,
}
try:
@@ -156,6 +159,9 @@ def process_trajectories_with_threads(
results.append(result)
if "memory_list" in result["metadata"]:
print(f'✅ Group {group_index} processed: {result["metadata"].get("memory_list", 0)}')
+ memory_list = result["metadata"].get("memory_list", [])
+ response = requests.post(url=f"{service_url}/add_task_memory", json={"memory_list": memory_list})
+ response.raise_for_status()
else:
print(f"❌ Group {group_index} processed: error")
except Exception as e:
@@ -174,7 +180,7 @@ def main():
"""Main function to convert JSONL to memories using ReMe service."""
parser = argparse.ArgumentParser(description="Convert JSONL to memories using ReMe service")
parser.add_argument("--jsonl_file", type=str, required=True, help="Path to the JSONL file")
- parser.add_argument("--service_url", type=str, default="http://localhost:8001", help="ReMe service URL")
+ parser.add_argument("--service_url", type=str, default="http://localhost:8002", help="ReMe service URL")
parser.add_argument("--output_file", type=str, help="Output file to save results (optional)")
parser.add_argument("--n_threads", type=int, default=4, help="Number of threads for processing")
@@ -226,21 +232,4 @@ def main():
if __name__ == "__main__":
- import sys
-
- if len(sys.argv) > 1:
- main()
- else:
- print("Running in compatibility mode...")
- with open("exp_result/qwen3-8b/with_think/bfcl-multi-turn-base-train_wo-exp.jsonl", "r") as f:
- data = [json.loads(line) for line in f]
-
- grouped_trajectories = group_trajectories_by_task_id(data)
- print(f"Total groups: {len(grouped_trajectories)}")
-
- results = process_trajectories_with_threads(
- grouped_trajectories,
- "http://localhost:8001",
- n_threads=4,
- )
- print(f"Processed {len(results)} groups")
+ main()
diff --git a/benchmark/bfcl/local_file_to_library.py b/benchmark/bfcl/local_file_to_library.py
deleted file mode 100644
index a2d9ec15..00000000
--- a/benchmark/bfcl/local_file_to_library.py
+++ /dev/null
@@ -1,30 +0,0 @@
-"""Load the library data and convert them to the new format"""
-
-import json
-
-with open("../../file_vector_store/bfcl_test.jsonl", "r", encoding="utf-8") as f:
- bfcl = [json.loads(line) for line in f]
-
-new_bfcl = []
-for exp in bfcl:
- new_exp = {}
- new_exp["workspace_id"] = exp["workspace_id"]
- new_exp["memory_id"] = exp["unique_id"]
- new_exp["memory_type"] = exp["metadata"]["memory_type"]
-
- new_exp["when_to_use"] = exp["content"]
- new_exp["content"] = exp["metadata"]["content"]
- new_exp["score"] = exp["metadata"]["score"]
-
- new_exp["time_created"] = exp["metadata"]["time_created"]
- new_exp["time_modified"] = exp["metadata"]["time_modified"]
- new_exp["author"] = exp["metadata"]["author"]
-
- new_exp["metadata"] = exp["metadata"]["metadata"]
- new_exp["metadata"]["utility"] = 0
- new_exp["metadata"]["freq"] = 0
-
- new_bfcl.append(new_exp)
-
-with open("../../library/bfcl_test.jsonl", "w", encoding="utf-8") as f:
- f.writelines(json.dumps(item, ensure_ascii=False) + "\n" for item in new_bfcl)
diff --git a/benchmark/bfcl/quickstart.md b/benchmark/bfcl/quickstart.md
new file mode 100644
index 00000000..c80ef75f
--- /dev/null
+++ b/benchmark/bfcl/quickstart.md
@@ -0,0 +1,129 @@
+# BFCL
+Experiment Quick Start Guide
+
+This guide helps you quickly set up and run BFCL experiments with ReMe integration.
+
+## Env Setup
+
+### 1. BFCL installation
+
+#### Clone the repository
+```bash
+cd ReMe/benchmark/bfcl
+git clone https://github.com/ShishirPatil/gorilla.git
+cd gorilla
+git checkout ea13468
+```
+
+#### Change directory to the `berkeley-function-call-leaderboard`
+```bash
+cd berkeley-function-call-leaderboard
+```
+
+#### Install the package in editable mode
+```bash
+pip install -e .
+cd ../..
+pip install -r requirements.txt
+```
+
+#### Move the dataset to the data folder under bfcl
+```bash
+cp -r gorilla/berkeley-function-call-leaderboard/bfcl_eval/data ./
+```
+
+#### Preprocess the data to get the suitable data format
+```bash
+python preprocess.py
+```
+
+**Note**: The original BFCL data is designed as a benchmark dataset and does not have a train/validation split, you can use ``split_into_trainval.py`` to split data into train and validation sets.
+
+```bash
+python split_into_trainval.py --input ./data/multiturn_data_base.jsonl --train ./data/multiturn_data_base_train.jsonl --val ./data/multiturn_data_base_val.jsonl
+```
+
+### 2. Start ReMe Service
+
+After collecting trajectories, Launch the ReMe service (make sure you have installed ReMe environment, if not please follow the steps in the [ReMe Installation Guide](https://github.com/agentscope-ai/ReMe/blob/main/doc/README.md) to install):
+
+```bash
+reme2 \
+ backend=http \
+ http.port=8002 \
+ llms.default.model_name=qwen3-8b \
+ embedding_models.default.model_name=text-embedding-v4 \
+ vector_stores.default.backend=local \
+ vector_stores.default.collection_name=bfcl
+```
+
+
+Option: init the task memory pool from scratch
+
+- First, collect agent trajectories on training data set without task memory:
+
+ ```bash
+ # important: num_runs = 8, use_memory = False, experiment_suffix="wo-memory", data_path="data/multiturn_data_base_train.jsonl"
+ python run_bfcl.py
+ ```
+
+- Second, using ReMe to construct the initial task memory pool:
+ ```bash
+ python init_task_memory_pool.py --jsonl_file ./exp_result/qwen3-8b/with_think/bfcl-multi-turn-base_wo-memory.jsonl
+ ```
+
+ > Parameters:
+ > `jsonl_file`: Path to the collloaded trajectories
+ > `service_url`: ReMe service URL (default: `http://localhost:8002`)
+ > `n_threads`: Number of threads for processing
+ > `output_file`: Output file to save results (optional)
+
+ Now you have inited the task memory pool using `local` backend. Then, run the following `curl` command to dump the memory library:
+ ```bash
+ curl -X POST "http://0.0.0.0:8002/dump_memory" \
+ -H "Content-Type: application/json" \
+ -d '{
+ "dump_file_path": "./library/bfcl.jsonl",
+ }'
+ ```
+
+- Next time, you can import this previously exported task memory data to populate the new started workspace with existing knowledge:
+ ```bash
+ curl -X POST "http://0.0.0.0:8002/load_memory" \
+ -H "Content-Type: application/json" \
+ -d '{
+ "load_file_path": "./library/bfcl.jsonl",
+ "clear_existing": true
+ }'
+ ```
+
+
+### 3. Run Experiments on Validation Set
+
+Run you can compare agent performance on the validation set with task memory (`use_memory=True`) and without task memory:
+
+```bash
+# remember to change the configuration options, e.g., `data_path=./data/multiturn_data_base_val.jsonl`
+python run_bfcl.py
+```
+
+**Note**:
+- `max_workers`: Number of parallel workers
+- `num_runs`: Number of times each task is repeated
+- `model_name`: LLM model name
+- `enable_thinking`: Control the model's thinking mode
+- `data_path`: Path to the training dataset (default: `./data/multiturn_data_base_val.jsonl`)
+- `answer_path`: Path to the possible answer, which are used to evaluate the model's output function (default: `./data/possible_answer`)
+- Results are automatically saved to `./exp_result/{model_name}/{no_think/with_think}` directory
+
+After running experiments, analyze the statistical results:
+
+```bash
+python run_exp_statistic.py
+```
+
+**What this script does:**
+- Processes all result files in `./exp_result/`
+- Calculates best@k&pass@k metrics for different k values
+- Generates a summary table showing performance comparisons
+- Saves results to `experiment_summary.csv`
diff --git a/benchmark/bfcl/requirements.txt b/benchmark/bfcl/requirements.txt
index 86ebcb1f..445bf2b3 100644
--- a/benchmark/bfcl/requirements.txt
+++ b/benchmark/bfcl/requirements.txt
@@ -2,4 +2,5 @@ jinja2
loguru
openai
ray
-pandas
\ No newline at end of file
+pandas
+soundfile
\ No newline at end of file
diff --git a/benchmark/bfcl/run_bfcl.py b/benchmark/bfcl/run_bfcl.py
index c01071ce..6ea8c325 100644
--- a/benchmark/bfcl/run_bfcl.py
+++ b/benchmark/bfcl/run_bfcl.py
@@ -131,7 +131,7 @@ def main():
run_agent(
max_workers=max_workers,
model_name=model_name,
- dataset_name="bfcl-multi-turn-base",
+ dataset_name="bfcl-multi-turn-base-val",
experiment_suffix="w-fixed-memory",
data_path="data/multiturn_data_base_val.jsonl",
answer_path=Path("data/possible_answer"),
diff --git a/benchmark/bfcl/run_exp_statistic.py b/benchmark/bfcl/run_exp_statistic.py
index 9eb9b3c8..18efcc8d 100644
--- a/benchmark/bfcl/run_exp_statistic.py
+++ b/benchmark/bfcl/run_exp_statistic.py
@@ -141,7 +141,7 @@ def run_exp_statistic():
# Sort columns by the number in column name (best@8, best@4, best@2, best@1)
# best_columns = [col for col in df.columns if col.startswith('best@')]
- best_columns = df.columns
+ best_columns = list(df.columns)
best_columns.sort(key=lambda x: x, reverse=False)
df = df[best_columns]
diff --git a/benchmark/bfcl/split_into_trainval.py b/benchmark/bfcl/split_into_trainval.py
index 82155855..e217def7 100644
--- a/benchmark/bfcl/split_into_trainval.py
+++ b/benchmark/bfcl/split_into_trainval.py
@@ -4,16 +4,46 @@ import argparse
import json
import random
+from default_ids import DEFAULT_TRAIN_IDS, DEFAULT_VAL_IDS
-def split_jsonl(input_file, train_file, val_file, ratio=0.8):
+
+def split_jsonl(
+ input_file: str,
+ train_file: str,
+ val_file: str,
+ ratio: float = 0.75,
+ random_split: bool = False,
+) -> None:
"""Split the JSONL file into train and validation sets."""
with open(input_file, "r", encoding="utf-8") as f:
data = [json.loads(line) for line in f]
- random.shuffle(data)
- split_idx = int(len(data) * ratio)
- train_data = data[:split_idx]
- val_data = data[split_idx:]
+ if random_split:
+ random.shuffle(data)
+ split_idx = int(len(data) * ratio)
+ train_data = data[:split_idx]
+ val_data = data[split_idx:]
+ else:
+ train_data = []
+ val_data = []
+ unknown_ids: list[str] = []
+ for obj in data:
+ if "id" not in obj:
+ raise ValueError(f"Missing 'id' field in input file: {input_file}")
+ obj_id = str(obj["id"])
+ if obj_id in DEFAULT_TRAIN_IDS:
+ train_data.append(obj)
+ elif obj_id in DEFAULT_VAL_IDS:
+ val_data.append(obj)
+ else:
+ unknown_ids.append(obj_id)
+
+ if len(train_data) + len(val_data) != len(data):
+ missing = len(data) - (len(train_data) + len(val_data))
+ examples = ", ".join(unknown_ids) if unknown_ids else "(none)"
+ raise ValueError(
+ f"{missing} samples in {input_file} not found in train_ref/val_ref id sets. Examples: {examples}",
+ )
with open(train_file, "w", encoding="utf-8") as f:
for item in train_data:
@@ -29,6 +59,11 @@ if __name__ == "__main__":
parser.add_argument("--train", required=True, help="Path to output train file")
parser.add_argument("--val", required=True, help="Path to output validation file")
parser.add_argument("--ratio", type=float, default=0.5, help="Train ratio (default: 0.8)")
-
+ parser.add_argument(
+ "--random",
+ action="store_true",
+ help="Whether to randomly split input into train/val. "
+ "If false, split strictly by default train/val id sets (see default_ids.py).",
+ )
args = parser.parse_args()
- split_jsonl(args.input, args.train, args.val, args.ratio)
+ split_jsonl(args.input, args.train, args.val, args.ratio, args.random)
diff --git a/benchmark/halumem/scripts.sh b/benchmark/halumem/cat_correct_scripts.sh
similarity index 100%
rename from benchmark/halumem/scripts.sh
rename to benchmark/halumem/cat_correct_scripts.sh
diff --git a/benchmark/halumem/eval_scripts.sh b/benchmark/halumem/eval_scripts.sh
new file mode 100755
index 00000000..c8a06881
--- /dev/null
+++ b/benchmark/halumem/eval_scripts.sh
@@ -0,0 +1,5 @@
+clear && python benchmark/halumem/eval_reme.py \
+ --data_path /Users/yuli/workspace/HaluMem/data/HaluMem-Medium.jsonl \
+ --reme_model_name qwen3.5-plus \
+ --batch_size 10000 \
+ --algo_version default
\ No newline at end of file
diff --git a/benchmark/halumem/quickstart.md b/benchmark/halumem/quickstart.md
new file mode 100644
index 00000000..59dc68fa
--- /dev/null
+++ b/benchmark/halumem/quickstart.md
@@ -0,0 +1,33 @@
+# Halumem
+Experiment Quick Start Guide
+This guide helps you quickly set up and run Halumem experiments with ReMe integration.
+
+### 1. Start ReMe Service
+Install ReMe (if not already installed)
+If you haven't installed the ReMe environment yet, follow these steps:
+```bash
+# Create ReMe environment
+conda create -p ./reme-env python==3.12
+conda activate ./reme-env
+
+# Install ReMe
+pip install .
+```
+
+### 2. Clone the Repository
+```bash
+cd ./benchmark/halumem
+git clone https://github.com/MemTensor/HaluMem.git
+```
+
+### 3. Run Experiments
+Launch the ReMe service to enable memory library functionality:
+```bash
+clear && python benchmark/halumem/eval_reme.py \
+ --data_path benchmark/halumem/HaluMem/data/HaluMem-Medium.jsonl \
+ --reme_model_name gpt-4o-mini-2024-07-18 \
+ --eval_model_name gpt-4o-mini-2024-07-18 \
+ --batch_size 40 \
+ --algo_version default
+```
+
diff --git a/benchmark/locomo/eval_reme.py b/benchmark/locomo/eval_reme.py
new file mode 100644
index 00000000..82d91031
--- /dev/null
+++ b/benchmark/locomo/eval_reme.py
@@ -0,0 +1,1107 @@
+"""
+Simplified evaluation script for ReMe on Locomo benchmark.
+
+This script performs a simplified evaluation pipeline:
+1. Load Locomo data
+2. Process each user's sessions with ReMe (summary + retrieve)
+3. Evaluate question answering
+4. Generate metrics and statistics
+
+Usage:
+ python bench/halumem/eval_reme_simple.py --data_path locomo10.json \
+ --top_k 20 --user_num 100 --max_concurrency 20
+"""
+
+import asyncio
+import json
+import os
+import re
+import shutil
+import time
+from pathlib import Path
+from datetime import datetime, timezone, timedelta
+from dataclasses import dataclass
+from typing import Any
+import yaml
+from loguru import logger
+from reme.core.enumeration import Role
+from reme.core.schema import Message
+
+
+from reme.reme import ReMe
+
+
+# ==================== Configuration ====================
+@dataclass
+class EvalConfig:
+ """Evaluation configuration parameters."""
+
+ data_path: str
+ top_k: int = 20
+ user_num: int = 1
+ max_concurrency: int = 2
+ batch_size: int = 40
+ output_dir: str = "bench_results/reme"
+ reme_model_name: str = "qwen-flash"
+ eval_model_name: str = "qwen3-max"
+ algo_version: str = "locomo"
+ enable_thinking_params: bool = False
+
+
+# ==================== Utilities ====================
+
+
+class DataLoader:
+ """Handles loading and parsing of HaluMem data."""
+
+ @staticmethod
+ def load_jsonl(file_path: str) -> list[dict]:
+ """Load all entries from a JSONL file."""
+ with open(file_path, "r", encoding="utf-8") as f:
+ return [json.loads(line.strip()) for line in f if line.strip()]
+
+ @staticmethod
+ def load_json(file_path: str) -> dict:
+ """Load dict from a JSON file."""
+ with open(file_path, "r", encoding="utf-8") as f:
+ return json.load(f)
+
+ @staticmethod
+ def format_dialogue_messages(
+ dialogue: list[dict],
+ speaker_a: str,
+ base_timestamp: datetime,
+ time_interval: int,
+ ) -> list[dict]:
+ """Format dialogue into ReMe message format with conversation_time."""
+
+ return [
+ {
+ "role": "user" if turn["speaker"] == speaker_a else "assistant",
+ "name": turn["speaker"],
+ "content": turn["text"],
+ "time_created": (base_timestamp + timedelta(seconds=idx * time_interval)).strftime("%Y-%m-%d %H:%M:%S"),
+ }
+ for idx, turn in enumerate(dialogue)
+ ]
+
+ @staticmethod
+ def format_dialogue_for_eval(dialogue: list[dict], user_name: str = None) -> str:
+ """Format dialogue into string for evaluation."""
+ formatted_turns = []
+ for turn in dialogue:
+ timestamp = (
+ datetime.strptime(
+ turn["timestamp"],
+ "%b %d, %Y, %H:%M:%S",
+ )
+ .replace(tzinfo=timezone.utc)
+ .strftime("%Y-%m-%d %H:%M:%S")
+ )
+
+ # Use user_name if role is 'user' and user_name is provided
+ role = user_name if turn["role"] == "user" and user_name else turn["role"]
+
+ formatted_turns.append(
+ f"Role: {role}\n" f"Content: {turn['content']}\n" f"Time: {timestamp}",
+ )
+ return "\n\n".join(formatted_turns)
+
+
+class FileManager:
+ """Manages file I/O operations."""
+
+ def __init__(self, base_dir: str):
+ self.base_dir = Path(base_dir)
+ self.tmp_dir = self.base_dir
+ self.tmp_dir.mkdir(parents=True, exist_ok=True)
+
+ def get_user_dir(self, user_name: str) -> Path:
+ """Get the directory path for a user."""
+ user_dir = self.tmp_dir / user_name
+ user_dir.mkdir(parents=True, exist_ok=True)
+ return user_dir
+
+ def get_session_file(self, user_name: str, session_id: int) -> Path:
+ """Get the file path for a specific session."""
+ return self.get_user_dir(user_name) / f"session_{session_id}.json"
+
+ def get_question_file(self, user_name: str) -> Path:
+ """Get the file path for a specific question."""
+ return self.get_user_dir(user_name) / "questions.json"
+
+ def save_session(self, user_name: str, session_id: int, data: dict):
+ """Save session data to file."""
+ file_path = self.get_session_file(user_name, session_id)
+ with open(file_path, "w", encoding="utf-8") as f:
+ json.dump(data, f, ensure_ascii=False, indent=2)
+ logger.info(f"✅ Saved session {session_id} to {file_path}")
+
+ def save_question(self, user_name: str, data: dict):
+ """Save question data to file"""
+ file_path = self.get_question_file(user_name)
+ with open(file_path, "w", encoding="utf-8") as f:
+ json.dump(data, f, ensure_ascii=False, indent=2)
+ logger.info(f"✅ Saved question to {file_path}")
+
+ def load_session(self, user_name: str, session_id: int) -> dict | None:
+ """Load session data from file."""
+ file_path = self.get_session_file(user_name, session_id)
+ if not file_path.exists():
+ return None
+ with open(file_path, "r", encoding="utf-8") as f:
+ return json.load(f)
+
+ def user_has_cache(self, user_name: str) -> bool:
+ """Check if user has cached results."""
+ user_dir = self.get_user_dir(user_name)
+ return any(f.name.startswith("session_") and f.suffix == ".json" for f in user_dir.iterdir())
+
+ def combine_results(self, output_file: str):
+ """Combine all user session files into a single JSONL file."""
+ with open(output_file, "w", encoding="utf-8") as f_out:
+ for user_dir in self.tmp_dir.iterdir():
+ if not user_dir.is_dir():
+ continue
+
+ session_files = sorted(
+ [f for f in user_dir.iterdir() if f.name.startswith("session_") and f.suffix == ".json"],
+ )
+
+ if not session_files:
+ continue
+
+ # Load first session to get user metadata
+ with open(session_files[0], "r", encoding="utf-8") as f_in:
+ first_session = json.load(f_in)
+
+ user_data = {
+ "uuid": first_session["uuid"],
+ "user_name": first_session["user_name"],
+ "sessions": [],
+ }
+
+ # Load all sessions
+ for session_file in session_files:
+ with open(session_file, "r", encoding="utf-8") as f_in:
+ session_data = json.load(f_in)
+ # Remove redundant user metadata
+ session_data.pop("uuid", None)
+ session_data.pop("user_name", None)
+ user_data["sessions"].append(session_data)
+
+ question_file = user_dir / "questions.json"
+ if not question_file.exists():
+ continue
+ with open(question_file, "r", encoding="utf-8") as f_in:
+ question_data = json.load(f_in)
+ user_data["evaluation_results"] = {
+ "question_answering_records": question_data,
+ }
+
+ f_out.write(json.dumps(user_data, ensure_ascii=False) + "\n")
+
+
+# ==================== Memory Operations ====================
+
+
+class MemoryProcessor:
+ """Handles ReMe memory operations."""
+
+ def __init__(
+ self,
+ reme: ReMe,
+ eval_model_name: str = "qwen3-max",
+ algo_version: str = "locomo",
+ enable_thinking_params: bool = False,
+ ):
+ self.reme = reme
+ self.eval_model_name = eval_model_name
+ self.algo_version = algo_version
+ self.enable_thinking_params = enable_thinking_params
+
+ async def add_memories(
+ self,
+ user_id: str,
+ messages: list[dict],
+ batch_size: int = 10000,
+ ) -> tuple[list[str], list, float]:
+ """
+ Add memories in batches using ReMe and return extracted memory contents.
+
+ Returns:
+ tuple: (extracted_memories, agent_messages, total_duration_ms)
+ """
+ extracted_memories = []
+ summary_messages = []
+ total_duration_ms = 0
+
+ for i in range(0, len(messages), batch_size):
+ batch = messages[i : i + batch_size]
+ start = time.time()
+
+ # Use new summary API
+ result = await self.reme.summarize_memory(
+ messages=batch,
+ user_name=user_id,
+ version=self.algo_version,
+ return_dict=True,
+ enable_time_filter=True,
+ enable_thinking_params=self.enable_thinking_params,
+ )
+
+ duration_ms = (time.time() - start) * 1000
+ total_duration_ms += duration_ms
+
+ extracted_memories.extend([m.model_dump(exclude_none=True) for m in result["answer"]])
+ summary_messages.extend([m.simple_dump(enable_argument_dict=True) for m in result["messages"]])
+
+ return extracted_memories, summary_messages, total_duration_ms
+
+ async def search_memory(
+ self,
+ query: str,
+ user_id: str,
+ top_k: int = 20,
+ ) -> tuple[dict, list, float]:
+ """
+ Search memory using ReMe and return structured answer with reasoning.
+
+ Returns:
+ tuple: (answer_dict, agent_messages, duration_ms)
+ answer_dict contains: {"reasoning": str, "answer": str, "memories": str}
+ """
+ start = time.time()
+
+ # Retrieve memories from ReMe using new API
+ result = await self.reme.retrieve_memory(
+ query=query,
+ retrieve_top_k=top_k,
+ user_name=user_id,
+ version=self.algo_version,
+ return_dict=True,
+ enable_time_filter=True,
+ enable_thinking_params=self.enable_thinking_params,
+ )
+
+ # Extract memories from response
+ memories = result["answer"]
+ agent_messages = [x.simple_dump(enable_argument_dict=True) for x in result["messages"]]
+ retrieved_nodes = [x.model_dump(exclude_none=True) for x in result["retrieved_nodes"]]
+
+ # Use LLM to generate structured answer from memories
+ answer_result = await answer_question_with_memories(
+ reme=self.reme,
+ question=query,
+ memories=memories,
+ user_id=user_id,
+ model_name=self.eval_model_name,
+ )
+
+ # Add original memories to the result
+ answer_result["memories"] = memories
+ answer_result["retrieved_nodes"] = retrieved_nodes
+
+ duration_ms = (time.time() - start) * 1000
+ return answer_result, agent_messages, duration_ms
+
+
+# ==================== Evaluation Functions ====================
+
+
+async def answer_question_with_memories(
+ reme: ReMe,
+ question: str,
+ memories: str,
+ user_id: str = None,
+ model_name: str = "qwen3-30b-a3b-instruct-2507",
+):
+ """
+ Answer a question using retrieved memories with PROMPT_MEMZERO_JSON template.
+
+ Args:
+ reme: ReMe instance with default_llm and prompt_handler
+ question: The question to answer
+ memories: The retrieved memories (formatted as context)
+ user_id: Optional user ID for context formatting
+ model_name: Model name to use for LLM request
+
+ Returns:
+ dict with 'reasoning' and 'answer' fields
+ """
+ # Format context with memories
+ if user_id:
+ context = reme.prompt_handler.prompt_format(
+ "TEMPLATE_MEMOS",
+ user_id=user_id,
+ memories=memories,
+ )
+ else:
+ context = f"Memories:\n{memories}"
+
+ # Use PROMPT_MEMZERO_JSON template for structured JSON response
+ prompt = reme.prompt_handler.prompt_format(
+ "PROMPT_MEMZERO_JSON",
+ context=context,
+ question=question,
+ )
+
+ result = await reme.get_llm("qwen3_max_instruct").simple_request_for_json(
+ prompt=prompt,
+ model_name=model_name,
+ )
+
+ return result
+
+
+async def evaluation_for_question(
+ reme: ReMe,
+ question: str,
+ golden_answer: str,
+ generated_answer: str,
+ model_name: str = "qwen3-max",
+):
+ """
+ Question-Answering Evaluation with optional Dialogue Context.
+
+ Args:
+ reme: ReMe instance with default_llm and prompt_handler
+ question: The question string to be evaluated.
+ golden_answer: The reference (gold-standard) answer.
+ generated_answer: The answer produced by the memory system.
+ model_name: Model name to use for LLM request
+
+ Returns:
+ dict with 'reasoning' and 'evaluation_result' fields
+ """
+ await asyncio.sleep(10)
+ # Use configured prompts
+ system_prompt = reme.prompt_handler.prompt_format(
+ "SYSTEM_PROMPT",
+ )
+ user_prompt = reme.prompt_handler.prompt_format(
+ "USER_PROMPT",
+ question=question,
+ golden_answer=golden_answer,
+ generated_answer=generated_answer,
+ )
+
+ reme_result = await reme.get_llm("qwen3_max_instruct").chat(
+ messages=[
+ Message(role=Role.SYSTEM, content=system_prompt),
+ Message(role=Role.USER, content=user_prompt),
+ ],
+ model_name=model_name,
+ )
+
+ content = reme_result.content
+ match = re.search(r'"label"\s*:\s*"([^"]*?)"', content)
+ if match:
+ label = match.group(1)
+ else:
+ label = "WRONG"
+ result = {
+ "reasoning": content,
+ "evaluation_result": label.strip().upper() == "CORRECT",
+ }
+ return result
+
+
+# ==================== Evaluation ====================
+
+
+class QuestionAnsweringEvaluator:
+ """Evaluates question answering performance."""
+
+ def __init__(self, memory_processor: MemoryProcessor, reme: ReMe, top_k: int, eval_model_name: str = "qwen3-max"):
+ self.memory_processor = memory_processor
+ self.reme = reme
+ self.top_k = top_k
+ self.eval_model_name = eval_model_name
+
+ async def evaluate_questions(
+ self,
+ questions: list[dict],
+ user_name: str,
+ uuid: str,
+ ) -> list[dict]:
+ """Evaluate all questions for a conversation."""
+ results = []
+
+ for qa in questions:
+ if qa["category"] == 5:
+ continue
+ answer_dict, agent_messages, duration_ms = await self.memory_processor.search_memory(
+ query=qa["question"],
+ user_id=user_name,
+ top_k=self.top_k,
+ )
+
+ # Extract answer and reasoning from the structured response
+ system_answer = answer_dict.get("answer", "")
+ system_reasoning = answer_dict.get("reasoning", "")
+ retrieved_memories = answer_dict.get("memories", "")
+ retrieved_nodes = answer_dict.get("retrieved_nodes", "")
+
+ # Evaluate response
+ eval_result = await evaluation_for_question(
+ reme=self.reme,
+ question=qa["question"],
+ golden_answer=qa["answer"],
+ generated_answer=system_answer,
+ model_name=self.eval_model_name,
+ )
+
+ eval_result_original_answer = await evaluation_for_question(
+ reme=self.reme,
+ question=qa["question"],
+ golden_answer=qa["answer"],
+ generated_answer=retrieved_memories,
+ model_name=self.eval_model_name,
+ )
+
+ # Build result record
+ qa_result = {
+ **qa,
+ "uuid": uuid,
+ "system_response": system_answer,
+ "system_reasoning": system_reasoning,
+ "retrieved_memories": retrieved_memories,
+ "retrieved_nodes": retrieved_nodes,
+ "retrieve_messages": agent_messages,
+ "search_duration_ms": duration_ms,
+ "result_type": eval_result.get("evaluation_result"),
+ "question_answering_reasoning": eval_result.get("reasoning", ""),
+ "original_result_type": eval_result_original_answer.get("evaluation_result"),
+ "original_question_answering_reasoning": eval_result_original_answer.get("reasoning", ""),
+ }
+ results.append(qa_result)
+
+ return results
+
+
+class MetricsAggregator:
+ """Aggregates evaluation metrics."""
+
+ @staticmethod
+ def _compute_single_metric(qa_records: list[dict], result_key: str) -> dict[str, Any]:
+ """Compute metrics for a single result type key."""
+ total = len(qa_records)
+ if total == 0:
+ return {
+ "correct_qa_ratio(all)": 0,
+ "correct_qa_ratio(valid)": 0,
+ "qa_valid_num": 0,
+ "qa_num": 0,
+ "category_1_accuracy": 0.0,
+ "category_2_accuracy": 0.0,
+ "category_3_accuracy": 0.0,
+ "category_4_accuracy": 0.0,
+ }
+
+ correct = 0
+ valid = 0
+
+ category_1_correct = 0
+ category_1_num = 0
+ category_1_valid = 0
+ category_2_correct = 0
+ category_2_num = 0
+ category_2_valid = 0
+ category_3_correct = 0
+ category_3_num = 0
+ category_3_valid = 0
+ category_4_correct = 0
+ category_4_num = 0
+ category_4_valid = 0
+
+ for qa in qa_records:
+ result_type = qa.get(result_key, "")
+ category = qa.get("category", 0)
+ if category == 1:
+ category_1_num += 1
+ elif category == 2:
+ category_2_num += 1
+ elif category == 3:
+ category_3_num += 1
+ elif category == 4:
+ category_4_num += 1
+
+ if result_type is not None and category in [1, 2, 3, 4]:
+ valid += 1
+ if result_type is True:
+ correct += 1
+
+ if category == 1:
+ category_1_valid += 1
+ if result_type is True:
+ category_1_correct += 1
+ elif category == 2:
+ category_2_valid += 1
+ if result_type is True:
+ category_2_correct += 1
+ elif category == 3:
+ category_3_valid += 1
+ if result_type is True:
+ category_3_correct += 1
+ elif category == 4:
+ category_4_valid += 1
+ if result_type is True:
+ category_4_correct += 1
+
+ metrics = {
+ "correct_qa_ratio(all)": correct / total,
+ "qa_valid_num": valid,
+ "qa_num": total,
+ "category_1_accuracy": category_1_correct / category_1_num if category_1_num > 0 else 0,
+ "category_1_num": category_1_num,
+ "category_1_valid_num": category_1_valid,
+ "category_2_accuracy": category_2_correct / category_2_num if category_2_num > 0 else 0,
+ "category_2_num": category_2_num,
+ "category_2_valid_num": category_2_valid,
+ "category_3_accuracy": category_3_correct / category_3_num if category_3_num > 0 else 0,
+ "category_3_num": category_3_num,
+ "category_3_valid_num": category_3_valid,
+ "category_4_accuracy": category_4_correct / category_4_num if category_4_num > 0 else 0,
+ "category_4_num": category_4_num,
+ "category_4_valid_num": category_4_valid,
+ }
+
+ if valid > 0:
+ metrics.update(
+ {
+ "correct_qa_ratio(valid)": correct / valid,
+ },
+ )
+ else:
+ metrics.update(
+ {
+ "correct_qa_ratio(valid)": 0,
+ },
+ )
+
+ return metrics
+
+ @staticmethod
+ def compute_qa_metrics(qa_records: list[dict]) -> dict[str, Any]:
+ """Compute question answering metrics for both result_type and original_result_type."""
+ return {
+ "with_llm_answer": MetricsAggregator._compute_single_metric(qa_records, "result_type"),
+ "with_original_memories": MetricsAggregator._compute_single_metric(qa_records, "original_result_type"),
+ }
+
+ @staticmethod
+ def compute_time_metrics(eval_results_file: str) -> dict[str, float]:
+ """Compute timing metrics from evaluation results."""
+ add_duration = 0
+ search_duration = 0
+
+ with open(eval_results_file, "r", encoding="utf-8") as f:
+ for line in f:
+ if not line.strip():
+ continue
+ user_data = json.loads(line)
+
+ for session in user_data["sessions"]:
+ add_duration += session.get("add_dialogue_duration_ms", 0)
+
+ eval_results = user_data.get("evaluation_results", {})
+ for qa in eval_results.get("question_answering_records", []):
+ search_duration += qa.get("search_duration_ms", 0)
+
+ # Convert to minutes
+ return {
+ "add_dialogue_duration_time": add_duration / 1000 / 60,
+ "search_memory_duration_time": search_duration / 1000 / 60,
+ "total_duration_time": (add_duration + search_duration) / 1000 / 60,
+ }
+
+
+# ==================== Evaluator ====================
+
+
+class LocomoEvaluator:
+ """
+ LOCOMO 评估器核心类
+ 用于评估 MemAgent 的记忆完整性、记忆准确性和问答准确性
+ """
+
+ def __init__(self, config: EvalConfig):
+ self.config = config
+ with open("eval_reme.yaml", "r", encoding="utf-8") as file:
+ data = yaml.safe_load(file)
+ self.summary_prompt_1 = data["user_message_summary_1"]
+ self.summary_prompt_2 = data["user_message_summary_2"]
+ self.retriever_prompt = data["user_message_retrieve"]
+
+ ops_dict = {
+ "personal_summarizer": {
+ "prompt_dict": {
+ "user_message_s1": self.summary_prompt_1,
+ "user_message_s2": self.summary_prompt_2,
+ },
+ },
+ "personal_retriever": {
+ "prompt_dict": {
+ "user_message": self.retriever_prompt,
+ },
+ "params": {
+ "return_memory_nodes": True,
+ },
+ },
+ }
+
+ self.reme = ReMe(
+ default_llm_config={
+ "model_name": self.config.reme_model_name,
+ },
+ ops=ops_dict,
+ )
+
+ # Load evaluation prompts into ReMe's prompt handler
+ prompts_yaml_path = Path(__file__).parent / "eval_reme.yaml"
+ self.reme.prompt_handler.load_prompt_by_file(prompts_yaml_path)
+
+ self.file_manager = FileManager(config.output_dir)
+ self.memory_processor = MemoryProcessor(
+ self.reme,
+ config.eval_model_name,
+ config.algo_version,
+ config.enable_thinking_params,
+ )
+ self.qa_evaluator = QuestionAnsweringEvaluator(
+ self.memory_processor,
+ self.reme,
+ config.top_k,
+ config.eval_model_name,
+ )
+ self.data_loader = DataLoader()
+
+ # For real-time updates
+ self._update_lock: asyncio.Lock | None = None
+ self._output_file: str | None = None
+
+ async def __aenter__(self):
+ """Async context manager entry."""
+ await self.reme.start()
+ return self
+
+ async def __aexit__(self, exc_type, exc_val, exc_tb):
+ """Async context manager exit with cleanup."""
+ await self.reme.close()
+ return False
+
+ async def process_user(self, user_data: dict) -> dict:
+ """Process all sessions for a user."""
+ speaker_a = user_data["conversation"]["speaker_a"]
+ speaker_b = user_data["conversation"]["speaker_b"]
+ uuid = f"{speaker_a}_{speaker_b}"
+ user_name = [speaker_a, speaker_b]
+ user_file_name = f"{speaker_a}_{speaker_b}"
+
+ new_user_data = {
+ "uuid": f"{user_data['conversation']['speaker_a']}_{user_data['conversation']['speaker_b']}",
+ "user_name": user_name,
+ "sessions": [],
+ "qas": [],
+ "eval_results": {},
+ }
+ logger.info(f"Processing user: {speaker_a} and {speaker_b}")
+ session_num = 19 if uuid == "Caroline_Melanie" else int(len(user_data["conversation"]) / 2 - 1)
+ time_interval = 60
+
+ # Process conversation
+ for idx in range(session_num):
+ conversation = user_data["conversation"]
+ logger.info(f"Processing user {user_name}: session {idx+1}/{session_num}")
+ session_data = {
+ "uuid": uuid,
+ "user_name": user_file_name,
+ "timestamp": conversation[f"session_{idx+1}_date_time"],
+ "session": conversation[f"session_{idx+1}"],
+ }
+
+ # Format dialogue
+ dialogue = conversation[f"session_{idx+1}"]
+ base_timestamp = parse_locomo_timestamp(session_data["timestamp"])
+ formatted_messages = self.data_loader.format_dialogue_messages(
+ dialogue,
+ speaker_a,
+ base_timestamp,
+ time_interval,
+ )
+ extracted_memories, agent_messages, duration_ms = await self.memory_processor.add_memories(
+ user_id=user_name,
+ messages=formatted_messages,
+ batch_size=self.config.batch_size,
+ )
+ session_data.update(
+ {
+ "dialogue": dialogue,
+ "extracted_memories": extracted_memories,
+ "summary_messages": agent_messages,
+ "add_dialogue_duration_ms": duration_ms,
+ },
+ )
+
+ self.file_manager.save_session(user_file_name, idx, session_data)
+
+ # Process questions
+ qas = user_data["qa"]
+ qa_results = await self.qa_evaluator.evaluate_questions(
+ questions=qas,
+ user_name=user_name,
+ uuid=uuid,
+ )
+
+ new_user_data["evaluation_results"] = {
+ "question_answering_records": qa_results,
+ }
+ self.file_manager.save_question(user_file_name, qa_results)
+
+ # Update results file after each conversation completes
+ await self._trigger_update()
+
+ return {"uuid": uuid, "user_name": user_name, "status": "ok"}
+
+ async def _trigger_update(self):
+ """Trigger real-time update of results and statistics."""
+ if self._update_lock is None or self._output_file is None:
+ return
+
+ async with self._update_lock:
+ self.file_manager.combine_results(self._output_file)
+ self._update_statistics(self._output_file)
+
+ async def run_evaluation(self):
+ """Run the complete evaluation pipeline using ReMe."""
+ start_time = time.time()
+
+ # Load user data first to get user names
+ all_users = self.data_loader.load_json(self.config.data_path)
+ users_to_process = all_users[: self.config.user_num]
+
+ # Extract all user names and delete all profiles
+ all_user_names = [
+ f"{user_data['conversation']['speaker_a']}_&_{user_data['conversation']['speaker_b']}"
+ for user_data in all_users
+ ]
+ if all_user_names:
+ for user_name in all_user_names:
+ self.reme.get_profile_handler(user_name).delete_all()
+ logger.info(f"Deleted all profiles for {len(all_user_names)} users")
+
+ # Clear existing data
+ await self.reme.default_vector_store.delete_all()
+
+ # Clear meta_memory directory
+ meta_memory_path = Path(f"meta_memory/{self.reme.default_vector_store.collection_name}")
+ if meta_memory_path.exists():
+ shutil.rmtree(meta_memory_path)
+ logger.info(f"Cleared meta_memory directory: {meta_memory_path}")
+ meta_memory_path.mkdir(parents=True, exist_ok=True)
+
+ print("\n" + "=" * 80)
+ print("LOCOMO EVALUATION - REME - QUESTION ANSWERING")
+ print(f"Users: {len(users_to_process)} | Concurrency: {self.config.max_concurrency}")
+ print("=" * 80 + "\n")
+
+ # Output file path for real-time updates
+ self._output_file = os.path.join(self.config.output_dir, "eval_results.jsonl")
+
+ # Lock for thread-safe file updates
+ self._update_lock = asyncio.Lock()
+
+ # Process users with concurrency control
+ semaphore = asyncio.Semaphore(self.config.max_concurrency)
+
+ async def process_with_cache_check(idx: int, user_data: dict):
+ async with semaphore:
+ user_name = f"{user_data['conversation']['speaker_a']}_{user_data['conversation']['speaker_b']}"
+
+ # Check cache
+ if self.file_manager.user_has_cache(user_name):
+ print(f"⚡ [{idx}/{len(users_to_process)}] Skipping {user_name} (cached)")
+ result = {"user_name": user_name, "status": "cached"}
+ # Also trigger update for cached users
+ await self._trigger_update()
+ else:
+ print(f"🔄 [{idx}/{len(users_to_process)}] Processing {user_name}...")
+ result = await self.process_user(user_data)
+ print(f"✅ [{idx}/{len(users_to_process)}] Completed {user_name}")
+
+ return result
+
+ tasks = [process_with_cache_check(idx, user) for idx, user in enumerate(users_to_process, 1)]
+ await asyncio.gather(*tasks, return_exceptions=True)
+
+ elapsed = time.time() - start_time
+ print(f"\n✅ Processing completed in {elapsed:.2f}s")
+ print(f"📁 Results: {self._output_file}\n")
+
+ # Final aggregation and report
+ await self.aggregate_and_report(self._output_file)
+
+ def _update_statistics(self, results_file: str):
+ """Update statistics file based on current results (for real-time monitoring)."""
+ if not os.path.exists(results_file):
+ return
+
+ # Collect all QA records
+ qa_records = []
+ try:
+ with open(results_file, "r", encoding="utf-8") as f:
+ for line in f:
+ if not line.strip():
+ continue
+ user_data = json.loads(line)
+ eval_results = user_data.get("evaluation_results", {})
+ qa_records.extend(
+ eval_results.get("question_answering_records", []),
+ )
+ except (json.JSONDecodeError, KeyError):
+ return
+
+ if not qa_records:
+ return
+
+ # Compute metrics
+ qa_metrics = MetricsAggregator.compute_qa_metrics(qa_records)
+ time_metrics = MetricsAggregator.compute_time_metrics(results_file)
+
+ final_results = {
+ "overall_score": {
+ "question_answering": qa_metrics,
+ "time_consuming": time_metrics,
+ },
+ "question_answering_records": qa_records,
+ }
+
+ # Save statistics
+ report_file = os.path.join(self.config.output_dir, "eval_statistics.json")
+ with open(report_file, "w", encoding="utf-8") as f:
+ json.dump(final_results, f, ensure_ascii=False, indent=4)
+
+ async def aggregate_and_report(self, results_file: str):
+ """Aggregate results and generate final report."""
+ print("=" * 80)
+ print("AGGREGATING METRICS")
+ print("=" * 80 + "\n")
+
+ # Collect all QA records
+ qa_records = []
+ print(results_file)
+ with open(results_file, "r", encoding="utf-8") as f:
+ for line in f:
+ if not line.strip():
+ continue
+ user_data = json.loads(line)
+ print(user_data)
+ eval_results = user_data.get("evaluation_results", {})
+ qa_records.extend(
+ eval_results.get("question_answering_records", []),
+ )
+
+ # Compute metrics
+ qa_metrics = MetricsAggregator.compute_qa_metrics(qa_records)
+ time_metrics = MetricsAggregator.compute_time_metrics(results_file)
+
+ final_results = {
+ "overall_score": {
+ "question_answering": qa_metrics,
+ "time_consuming": time_metrics,
+ },
+ "question_answering_records": qa_records,
+ }
+
+ # Save final report
+ report_file = os.path.join(self.config.output_dir, "eval_statistics.json")
+ with open(report_file, "w", encoding="utf-8") as f:
+ json.dump(final_results, f, ensure_ascii=False, indent=4)
+
+ print(f"📊 Statistics saved to: {report_file}\n")
+
+ # Print summary
+ self._print_summary(qa_metrics, time_metrics)
+
+ def _print_summary(self, qa_metrics: dict, time_metrics: dict):
+ """Print evaluation summary."""
+ print("=" * 80)
+ print("EVALUATION SUMMARY - REME")
+ print("=" * 80 + "\n")
+
+ # Print metrics for LLM-generated answer (result_type)
+ llm_metrics = qa_metrics["with_llm_answer"]
+ print(llm_metrics)
+ print("📊 Question Answering (with LLM answer):")
+ print(f" Correct (all): {llm_metrics['correct_qa_ratio(all)']:.4f}")
+ print(f" Correct (valid): {llm_metrics['correct_qa_ratio(valid)']:.4f}")
+ print(f" Valid/Total: {llm_metrics['qa_valid_num']}/{llm_metrics['qa_num']}")
+ print(f" Category 1 Accuracy: {llm_metrics['category_1_accuracy']:.4f}")
+ print(f" Category 2 Accuracy: {llm_metrics['category_2_accuracy']:.4f}")
+ print(f" Category 3 Accuracy: {llm_metrics['category_3_accuracy']:.4f}")
+ print(f" Category 4 Accuracy: {llm_metrics['category_4_accuracy']:.4f}")
+
+ # Print metrics for original retrieved memories (original_result_type)
+ orig_metrics = qa_metrics["with_original_memories"]
+ print("\n📊 Question Answering (with original memories):")
+ print(f" Correct (all): {orig_metrics['correct_qa_ratio(all)']:.4f}")
+ print(f" Correct (valid): {orig_metrics['correct_qa_ratio(valid)']:.4f}")
+ print(f" Valid/Total: {orig_metrics['qa_valid_num']}/{orig_metrics['qa_num']}")
+ print(f" Category 1 Accuracy: {orig_metrics['category_1_accuracy']:.4f}")
+ print(f" Category 2 Accuracy: {orig_metrics['category_2_accuracy']:.4f}")
+ print(f" Category 3 Accuracy: {orig_metrics['category_3_accuracy']:.4f}")
+ print(f" Category 4 Accuracy: {orig_metrics['category_4_accuracy']:.4f}")
+
+ print("\n⏱️ Time Metrics:")
+ print(f" Memory Addition: {time_metrics['add_dialogue_duration_time']:.2f} min")
+ print(f" Memory Search: {time_metrics['search_memory_duration_time']:.2f} min")
+ print(f" Total: {time_metrics['total_duration_time']:.2f} min")
+ print("\n" + "=" * 80)
+
+
+def parse_locomo_timestamp(timestamp_str: str):
+ """
+ Parse LoCoMo timestamp format.
+
+ Input format: "6:07 pm on 13 January, 2023"
+ Special value: "Unknown" or unparseable returns None
+ Output: datetime object or None
+ """
+ # Clean string
+ timestamp_str = timestamp_str.replace("\\s+", " ").strip()
+
+ # Handle special cases: Unknown or empty string
+ if timestamp_str.lower() == "unknown" or not timestamp_str:
+ # No time information, return None
+ return None
+
+ try:
+ return datetime.strptime(timestamp_str, "%I:%M %p on %d %B, %Y")
+ except ValueError:
+ # If parse fails, return None and print warning
+ print(f"⚠️ Warning: Failed to parse timestamp '{timestamp_str}', no timestamp will be set")
+ return None
+
+
+# ==================== Main Pipeline ====================
+
+
+async def main_async(
+ data_path: str,
+ top_k: int,
+ user_num: int,
+ max_concurrency: int,
+ reme_model_name: str = "qwen-flash",
+ eval_model_name: str = "qwen3-max",
+ algo_version: str = "halumem",
+ enable_thinking_params: bool = False,
+):
+ """Main async entry point for ReMe evaluation with proper resource cleanup."""
+ config = EvalConfig(
+ data_path=data_path,
+ top_k=top_k,
+ user_num=user_num,
+ max_concurrency=max_concurrency,
+ reme_model_name=reme_model_name,
+ eval_model_name=eval_model_name,
+ algo_version=algo_version,
+ enable_thinking_params=enable_thinking_params,
+ )
+
+ # Use async context manager for automatic cleanup
+ async with LocomoEvaluator(config) as evaluator:
+ await evaluator.run_evaluation()
+
+
+def main(
+ data_path: str,
+ top_k: int = 20,
+ user_num: int = 1,
+ max_concurrency: int = 2,
+ reme_model_name: str = "qwen-flash",
+ eval_model_name: str = "qwen3-max",
+ algo_version: str = "halumem",
+ enable_thinking_params: bool = False,
+):
+ """Synchronous entry point."""
+ asyncio.run(
+ main_async(
+ data_path=data_path,
+ top_k=top_k,
+ user_num=user_num,
+ max_concurrency=max_concurrency,
+ reme_model_name=reme_model_name,
+ eval_model_name=eval_model_name,
+ algo_version=algo_version,
+ enable_thinking_params=enable_thinking_params,
+ ),
+ )
+
+
+if __name__ == "__main__":
+ import argparse
+
+ parser = argparse.ArgumentParser(description="Simplified evaluation for ReMe on Locomo benchmark")
+ parser.add_argument(
+ "--data_path",
+ type=str,
+ required=True,
+ help="Path to Locomo data file (e.g., locomo10.jsonl)",
+ )
+ parser.add_argument(
+ "--top_k",
+ type=int,
+ default=20,
+ help="Number of top memories to retrieve (default: 20)",
+ )
+ parser.add_argument(
+ "--user_num",
+ type=int,
+ default=1,
+ help="Number of users to evaluate (default: 1)",
+ )
+ parser.add_argument(
+ "--max_concurrency",
+ type=int,
+ default=2,
+ help="Maximum concurrency for processing (default: 2)",
+ )
+ parser.add_argument(
+ "--reme_model_name",
+ type=str,
+ default="qwen-flash",
+ help="Model name for ReMe (default: qwen-flash)",
+ )
+ parser.add_argument(
+ "--eval_model_name",
+ type=str,
+ default="qwen3-max",
+ help="Model name for evaluation (default: qwen3-max)",
+ )
+ parser.add_argument(
+ "--algo_version",
+ type=str,
+ default="default",
+ help="Algorithm version for summary and retrieval (default: halumem)",
+ )
+ parser.add_argument(
+ "--enable_thinking_params",
+ action="store_true",
+ default=True,
+ help="Enable thinking parameters for summary and retrieval (default: False)",
+ )
+
+ args = parser.parse_args()
+ print(f"args={args}!")
+
+ main(
+ data_path=args.data_path,
+ top_k=args.top_k,
+ user_num=args.user_num,
+ max_concurrency=args.max_concurrency,
+ reme_model_name=args.reme_model_name,
+ eval_model_name=args.eval_model_name,
+ algo_version=args.algo_version,
+ enable_thinking_params=args.enable_thinking_params,
+ )
diff --git a/benchmark/locomo/eval_reme.yaml b/benchmark/locomo/eval_reme.yaml
new file mode 100644
index 00000000..113431b3
--- /dev/null
+++ b/benchmark/locomo/eval_reme.yaml
@@ -0,0 +1,180 @@
+TEMPLATE_MEMOS: |
+ Memories for user {user_id}:
+ {memories}
+
+PROMPT_MEMZERO_JSON: |
+ # CONTEXT:
+ {context}
+
+ # CONTEXT PRIORITY:
+ When the context contains information from multiple sources, follow this strict priority order:
+ 1. **Historical Dialogue** (highest priority) - Direct conversation content
+ 2. **Extracted Memories** (medium priority) - Summarized memory points
+ 3. **User Profile** (lowest priority) - General user information
+
+ # Question:
+ {question}
+
+ # INSTRUCTIONS:
+ 1. Carefully analyze all provided memories (facts and entities)
+ 2. Pay special attention to the timestamps (event_time) to determine when events occurred
+ 3. If the question asks about a specific event or fact, look for direct evidence in the memories
+ 4. If the memories contain contradictory information, prioritize the most recent memory
+ 5. Always convert relative time references to specific dates, months, or years
+ 6. Be as specific as possible when talking about people, places, and events
+ 7. Timestamps in memories represent the time the event was mentioned in a message, not the actual time the event occurred
+
+
+ # OUTPUT FORMAT:
+ Please provide your response in the following JSON format:
+
+ ```json
+ {{
+ "reasoning": "reasoning content",
+ "answer": "Provide a detailed answer"
+ }}
+ ```
+
+SYSTEM_PROMPT: |
+ You are an expert grader that determines if answers to questions match a gold standard answer
+
+USER_PROMPT: |
+ Your task is to label an answer to a question as 'CORRECT' or 'WRONG'. You will be given the following data:
+ (1) a question (posed by one user to another user),
+ (2) a 'gold' (ground truth) answer,
+ (3) a generated answer
+ which you will score as CORRECT/WRONG.
+
+ The point of the question is to ask about something one user should know about the other user based on their prior conversations.
+ The gold answer will usually be a concise and short answer that includes the referenced topic, for example:
+ Question: Do you remember what I got the last time I went to Hawaii?
+ Gold answer: A shell necklace
+ The generated answer might be much longer, but you should be generous with your grading - as long as it touches on the same topic as the gold answer, it should be counted as CORRECT.
+
+ For time related questions, the gold answer will be a specific date, month, year, etc. The generated answer might be much longer or use relative time references (like "last Tuesday" or "next month"), but you should be generous with your grading - as long as it refers to the same date or time period as the gold answer, it should be counted as CORRECT. Even if the format differs (e.g., "May 7th" vs "7 May"), consider it CORRECT if it's the same date.
+
+ Now it's time for the real question:
+ Question: {question}
+ Gold answer: {golden_answer}
+ Generated answer: {generated_answer}
+
+ First, provide a short (one sentence) explanation of your reasoning, then finish with CORRECT or WRONG.
+ Do NOT include both CORRECT and WRONG in your response, or it will break the evaluation script.
+
+ Just return the label CORRECT or WRONG in a json format with the key as "label".
+
+user_message_summary_1: |
+ You are a Memory Agent responsible for managing {memory_type} memories about {memory_target}.
+
+ ## Latest Conversation
+ Format: round [] :
+ {context}
+
+ ## Task
+ ### Step 1: Create Memory Draft
+ Use `add_draft_and_retrieve_similar_memory` to create a memory draft list based on the latest conversation.
+ - For each memory draft, fill in the required parameters:
+ * `message_time`: timestamp from the conversation (e.g., '2020-01-01 00:00:00')
+ * `memory_content`: concise memory content extracted from the conversation
+ - Use actual names from the conversation (e.g., "Bob likes apples") instead of generic references (e.g., "user likes apples")
+ - Extract all important information comprehensively—do not miss critical details, but avoid any fabrications or unfounded assumptions
+ - The tool will retrieve similar historical memories via vector search to help you in Step 2
+
+ ### Step 2: Add Memories
+ Review each memory draft from Step 1 and compare it with the retrieved historical memories, then use `add_memory` to manage all memories in one call:
+
+ - For each new memory, fill in the required parameters:
+ * `message_time`: timestamp from the conversation (e.g., '2020-01-01 00:00:00')
+ * `memory_content`: memory content
+ - Add memories when:
+ * The draft contains new information not present in historical memories
+
+
+ **General Guidelines:**
+ - **Skip** drafts if their content is already fully covered by historical memories (avoid redundancy)
+ - You can add memories in a single `add_memory` tool call
+
+user_message_summary_2: |
+ You are a Profile Agent responsible for managing profiles about {memory_target}.
+
+ ## Latest Conversation
+ Format: round [] :
+ {context}
+
+ ## Current Profiles
+ {profiles}
+
+ ## Task
+ Analyze the Latest Conversation and use `update_profiles` to manage profiles (both updates and additions in one call):
+
+ **For profiles_to_update** (updating existing profiles):
+ - For each profile to update, fill in the required parameters:
+ * `profile_id`: ID of the profile to update (from Current Profiles)
+ * `message_time`: timestamp from the conversation (e.g., '2020-01-01 00:00:00')
+ * `profile_key`: profile key or category (e.g., 'name', 'age', 'occupation')
+ * `profile_value`: updated profile value, please be concise. (e.g., 'John Smith')
+
+ **For profiles_to_add** (adding new profiles):
+ - For each new profile, fill in the required parameters:
+ * `message_time`: timestamp from the conversation (e.g., '2020-01-01 00:00:00')
+ * `profile_key`: profile key or category (e.g., 'name', 'age', 'occupation')
+ * `profile_value`: profile value (e.g., 'John Smith')
+ - Add profiles when:
+ * The information represents a new distinct profile not present in Current Profiles
+ * The profile key doesn't exist in Current Profiles
+ * The information cannot be merged into existing profiles
+
+ **General Guidelines:**
+ - Extract all important information comprehensively—do not miss critical details, but avoid any fabrications or unfounded assumptions
+ - You can update and add profiles in a single tool call
+
+user_message_retrieve: |
+ You are a Memory Retrieval Agent specialized in retrieving {memory_type} memories about {memory_target}.
+
+ ## User Profile
+ {user_profile}
+
+ ## User Question
+ {context}
+
+ ## Multi-Phase Retrieval Strategy
+ Follow these phases sequentially to gather comprehensive information:
+
+ ### Phase 1: Semantic Search (No Time Filter)
+ **Tool**: `retrieve_memory` (without time constraints)
+ **Objective**: Cast a wide net to find potentially relevant memories
+ **Approach**:
+ - Execute 3-5 diverse search queries using different formulations:
+ * Original question verbatim
+ * Rephrased variations (different wording, synonyms)
+ * Entity-focused queries (extract and search specific names, places, events)
+ * Keyword-based searches (core concepts, topics)
+ * Related context queries (broader themes)
+ - Review all results before proceeding to next phase
+
+ ### Phase 2: Deep Dive into History
+ **Tool**: `read_history`
+ **When to use**: After exhausting retrieval attempts OR when specific conversation context is needed
+ **Important Constraints**:
+ - Each history is very long and resource-intensive to read
+ - **Maximum limit: Read no more than 3 histories total**
+ - Only use this phase when absolutely necessary for answering the question
+ **Approach**:
+ - Extract `history_id` from retrieved memory references
+ - Prioritize the most relevant or recent histories
+ - Can read multiple histories at once by passing multiple history_ids
+ - Be selective: choose only the top 1-3 most promising histories
+ - Use this to understand the full conversation surrounding a memory
+
+ ## Response Guidelines
+ - Base your answer EXCLUSIVELY on user profile, retrieved memories, and history data
+ - Never infer, assume, or hallucinate information
+ - Always cite sources with timestamps: `[timestamp] Memory content`
+ - Present conflicting information transparently with respective timestamps
+ - If you find sufficient information to answer the user's question, you may output directly without exhausting all search phases
+ - Exhaust all search strategies before concluding information doesn't exist
+
+ ### Output any tangentially related findings, Format:
+ [timestamp] [memory/profile/history] [relevant content1]
+ [timestamp] [memory/profile/history] [relevant content2]
+
diff --git a/benchmark/longmemeval/quickstart.md b/benchmark/longmemeval/quickstart.md
new file mode 100644
index 00000000..c27d5c65
--- /dev/null
+++ b/benchmark/longmemeval/quickstart.md
@@ -0,0 +1,48 @@
+# Longmemeval
+Experiment Quick Start Guide
+This guide helps you quickly set up and run Longmemeval experiments with ReMe integration.
+
+### 1. Start ReMe Service
+Install ReMe (if not already installed)
+If you haven't installed the ReMe environment yet, follow these steps:
+```bash
+# Create ReMe environment
+conda create -p ./reme-env python==3.12
+conda activate ./reme-env
+
+# Install ReMe
+pip install .
+```
+
+### 2. Clone the Repository
+```bash
+cd ./benchmark/longmemeval
+mkdir -p data/
+cd data/
+wget https://huggingface.co/datasets/xiaowu0162/longmemeval-cleaned/resolve/main/longmemeval_oracle.json
+wget https://huggingface.co/datasets/xiaowu0162/longmemeval-cleaned/resolve/main/longmemeval_s_cleaned.json
+wget https://huggingface.co/datasets/xiaowu0162/longmemeval-cleaned/resolve/main/longmemeval_m_cleaned.json
+cd ..
+```
+
+### 3. Run Experiments
+Launch the ReMe service to enable memory library functionality:
+```bash
+clear && python benchmark/longmemeval/eval_longmemeval_reme.py \
+ --data_path benchmark/longmemeval/data/longmemeval_s_cleaned.json \
+ --reme_model_name qwen-flash \
+ --reme_model_name retrieve_model_name \
+ --eval_model_name gpt-4o-mini-2024-07-18 \
+ --batch_size 20 \
+ --algo_version default
+```
+
+
+### 4. Evaluate Results
+Evaluate the results of the experiments:
+```bash
+python benchmark/longmememeval/compute_stats.py \
+ --results_dir bench_results/longmemeval_reme \
+ --output_file bench_results/longmemeval_reme/statistics.json
+```
+The `compute_stats.py` script computes various statistics from the evaluation results.
\ No newline at end of file
diff --git a/docs/REME2_README.md b/docs/REME2_README.md
deleted file mode 100644
index fd6174a6..00000000
--- a/docs/REME2_README.md
+++ /dev/null
@@ -1,13 +0,0 @@
-
-
-# TODO
-- [] halumem bench开发
-- [] default版本开发，for cli版本体验
-- [] cli开发
-- [] locomo bench开发
-- [] task memory迁移
-- [] mcp开发
-- [] reme外层接口完善
-- [] reme2 readme完善
-- [] 看日志，看要这个default版本怎么优化。
-- [] 学习Clawdbot记忆系统
\ No newline at end of file
diff --git a/docs/cli/quick_start_en.md b/docs/cli/quick_start_en.md
index 59c77603..f1701930 100644
--- a/docs/cli/quick_start_en.md
+++ b/docs/cli/quick_start_en.md
@@ -176,7 +176,6 @@ Controls how context space is allocated and how memory is searched:
| `watch_paths` | Directories/files to monitor |
| `suffix_filters` | Which file suffixes to watch (`.md`) |
| `recursive` | Whether to recurse into subdirectories |
-| `scan_on_start` | Whether to do a full scan on startup |
**token_counters — Token Counter**
diff --git a/docs/cli/quick_start_zh.md b/docs/cli/quick_start_zh.md
index cf7c8ace..4673e4c5 100644
--- a/docs/cli/quick_start_zh.md
+++ b/docs/cli/quick_start_zh.md
@@ -172,7 +172,6 @@ pip install -e .
| `watch_paths` | 要监控的目录/文件 |
| `suffix_filters` | 只关心哪些后缀（`.md`） |
| `recursive` | 是否递归子目录 |
-| `scan_on_start` | 启动时先全量扫一遍 |
**token_counters — Token 计数器**
diff --git a/docs/cookbook/bfcl/quickstart.md b/docs/cookbook/bfcl/quickstart.md
deleted file mode 100644
index c814b0ac..00000000
--- a/docs/cookbook/bfcl/quickstart.md
+++ /dev/null
@@ -1,121 +0,0 @@
-# BFCL
-Experiment Quick Start Guide
-
-This guide helps you quickly set up and run BFCL experiments with ReMe integration.
-
-## Env Setup
-
-### 1. BFCL installation
-
-#### clone the repository
-```bash
-git clone https://github.com/ShishirPatil/gorilla.git
-```
-
-#### Change directory to the `berkeley-function-call-leaderboard`
-```bash
-cd gorilla/berkeley-function-call-leaderboard
-```
-
-#### Install the package in editable mode
-```bash
-conda create -n bfcl-env python==3.12
-conda activate bfcl-env
-pip install -e .
-pip install -r requirements.txt
-```
-
-#### Move the dataset to the data folder under bfcl
-```bash
-cp -r bfcl_eval/data {/path/to/bfcl/data}
-```
-
-**Note**: The original BFCL data is designed as a benchmark dataset and does not have a train/validation split, you can use ``split_into_trainval.py`` to split JSONL file into train and validation sets.
-
-### 2. Collect agent trajectories on training data set
-
-Run the main experiment script to collect agent trajectories on training data set without task memory(`use_memory=False`):
-
-```bash
-python run_bfcl.py
-```
-
-**Note**:
-- `max_workers`: Number of parallel workers (default: `4`)
-- `num_runs`: Number of times each task is repeated (default: `1`)
-- `model_name`: LLM model name (default: `qwen3-8b`)
-- `enable_thinking`: Control the model's thinking mode (default: `False`)
-- `data_path`: Path to the training dataset (default: `./data/multiturn_data_base_train.jsonl`)
-- `answer_path`: Path to the possible answer, which are used to evaluate the model's output function (default: `./data/possible_answer`)
-- Results are automatically saved to `./exp_result/{model_name}/{no_think/with_think}` directory
-
-### 3. Start ReMe Service and Init the task memory pool
-
-After collecting trajectories, Launch the ReMe service (make sure you have installed ReMe environment, if not please follow the steps in the [ReMe Installation Guide](https://github.com/agentscope-ai/ReMe/blob/main/doc/README.md) to install):
-
-```bash
-reme \
- backend=http \
- http.port=8002 \
- llm.default.model_name=qwen-max-2025-01-25 \
- embedding_model.default.model_name=text-embedding-v4 \
- vector_store.default.backend=local
-```
-
-and then init the task memory pool:
-
-```bash
-python init_task_memory_pool.py
-```
-
-**Configuration options in `init_task_memory_pool.py`:**
-- `jsonl_file`: Path to the collloaded trajectories
-- `service_url`: ReMe service URL (default: `http://localhost:8002`)
-- `workspace_id`: Workspace ID for the task memory pool (default: `bfcl_test`)
-- `n_threads`: Number of threads for processing (default: `4`)
-- `output_file`: Output file to save results (optional)
-
-Now you have inited the task memory pool using `local` backend (start on `http://localhost:8002`). Then, use `local_file_to_library.py` script to convert the local file to the memory library or run the following `curl` command:
-```bash
-curl -X POST "http://0.0.0.0:8002/vector_store" \
- -H "Content-Type: application/json" \
- -d '{
- "workspace_id": "bfcl_test",
- "action": "dump",
- "path": "./library"
- }'
-```
-to dump the memory library (default in `./library/bfcl_test.jsonl`).
-
-Next time, you can import this previously exported task memory data to populate the new started workspace with existing knowledge:
-```bash
-curl -X POST "http://0.0.0.0:8002/vector_store" \
- -H "Content-Type: application/json" \
- -d '{
- "workspace_id": "bfcl_test",
- "action": "load",
- "path": "./library"
- }'
-```
-
-
-### 4. Run Experiments on Validation Set
-
-Run you can compare agent performance on the validation set with task memory (`use_memory=True`) and without task memory:
-
-```bash
-# remember to change the configuration options, e.g., `data_path=./data/multiturn_data_base_val.jsonl`
-python run_bfcl.py
-```
-
-After running experiments, analyze the statistical results:
-
-```bash
-python run_exp_statistic.py
-```
-
-**What this script does:**
-- Processes all result files in `./exp_result/`
-- Calculates best@k metrics for different k values
-- Generates a summary table showing performance comparisons
-- Saves results to `experiment_summary.csv`
diff --git a/docs/cookbook/faq.md b/docs/cookbook/faq.md
index 85ed37a3..66603e37 100644
--- a/docs/cookbook/faq.md
+++ b/docs/cookbook/faq.md
@@ -10,17 +10,18 @@ This document provides answers to frequently asked questions about our paper "[R
reme2 \
backend=http \
http.port=8002 \
- llm.default.model_name=qwen3-8b \
- embedding_model.default.model_name=text-embedding-v4 \
- vector_store.default.backend=es
+ llms.default.model_name=qwen3-8b \
+ embedding_models.default.model_name=text-embedding-v4 \
+ vector_stores.default.backend=es \
+ vector_stores.default.hosts=http://xx.yy.zz.mm:nn
```
**Evaluation Code:** [run_appworld.py](https://github.com/agentscope-ai/ReMe/blob/main/benchmark/appworld/run_appworld.py) with the following parameters
|Experimental Settings|No Memory |ReMe (fixed) |ReMe (dynamic)|
|---|---|---|---|
-|max_workers| 16|16|16|
+|max_workers|16|16|16|
|batch_size|8|8|8|
|num_runs|4|4|1|
-|num_trials|1 |1|3|
+|num_trials|1|1|3|
|model_name|"qwen3-8b"|"qwen3-8b"|"qwen3-8b"|
|use_memory| False| True|True|
|use_memory_addition|False|False|True|
diff --git a/docs/copaw_context_design.md b/docs/copaw_context_design.md
new file mode 100644
index 00000000..13fa5fc6
--- /dev/null
+++ b/docs/copaw_context_design.md
@@ -0,0 +1,278 @@
+# CoPaw Context Management Design
+
+> This article focuses on **short-term context management** and does not cover the long-term memory module.
+
+---
+
+Sooner or later, every AI Agent hits the same wall: **the context window fills up**.
+
+Tool calls return walls of HTML, thousands of lines of logs, or entire file contents — all of which rapidly consume precious token budget. As the conversation grows, early information either gets truncated or blows up the window entirely, and the Agent's performance begins to degrade.
+
+[CoPaw](https://github.com/agentscope-ai/CoPaw) addresses this problem with a systematic approach to context management. This article provides a complete breakdown of the data structures and runtime mechanics behind **CoPaw Context Management V2**.
+
+---
+
+## What does the context look like?
+
+Before discussing "how to manage it", let's first understand "what is being managed".
+
+CoPaw's context is split into two layers: the **in-memory layer** and the **file system layer**.
+
+### In-Memory Layer
+
+Two core fields are maintained in memory:
+
+- **`compact_summary`** (optional): After conversation history has been compacted, this field holds a structured summary covering five dimensions — `Goal`, `Constraints`, `Progress`, `KeyDecisions`, and `NextSteps` — essentially a refined "work memo". It also includes a **path guide to the raw historical dialog**, pointing the Agent to `dialog/YYYY-MM-DD.jsonl` and suggesting reading from the end backwards.
+- **`messages`**: The complete list of messages for the current conversation — the data actually consumed by the Agent during reasoning.
+
+### File System Layer (File Cache)
+
+For content that is too large or too volatile to reside in memory long-term, CoPaw offloads it to the file system:
+
+- **Raw conversation history**: `dialog/YYYY-MM-DD.jsonl`, stored per day
+- **Tool call results**: `tool_result/{uuid}.txt`, with an N-day TTL and automatic cleanup on expiry
+
+```mermaid
+flowchart TD
+ A[Context] --> B[compact_summary]
+ B --> C[dialog path guide + Goal/Constraints/Progress/KeyDecisions/NextSteps]
+ A --> E[messages: full dialogue history]
+ A --> F[File System Cache]
+ F --> G[dialog/YYYY-MM-DD.jsonl]
+ F --> H[tool_result/uuid.txt N-day TTL]
+```
+
+This design allows the Agent to quickly access recent conversations in memory while being able to look up historical context on demand — without stuffing all history into the context window.
+
+---
+
+## What happens before reasoning? The Pre-Reasoning Hook
+
+Before each reasoning step begins, CoPaw executes a **Pre-Reasoning Hook** that automatically tidies up the context. The process runs in four steps:
+
+1. **Tool result compaction** (`ToolCallResultCompact`): Process tool call results first — truncate oversized content and offload it to the file system.
+2. **Context checking** (`ContextChecker`): Compute the current token usage and determine whether it exceeds the threshold.
+3. **If the threshold is exceeded**:
+ - Keep the most recent **X%** of tokens (to preserve conversational continuity).
+ - Call the `Compactor` on the earlier history to generate a structured summary.
+4. **Dialog persistence** (`SaveDialog`): Save the compacted raw conversation to the file system.
+
+```mermaid
+flowchart LR
+ A[Pre-Reasoning Hook] --> B[ToolCallResultCompact]
+ B --> C[ContextChecker]
+ C --> D{Token > threshold?}
+ D -->|Yes| E[Keep recent X% tokens]
+ E --> F[Compact & generate summary]
+ F --> G[SaveDialog: persist to file]
+ D -->|No| H[Normal reasoning]
+```
+
+This flow ensures that the context is in a "clean" state at the start of every reasoning step.
+
+---
+
+## Tool Result Offload: Unified Two-Phase Truncation
+
+Tool call results are one of the main causes of context bloat. CoPaw uses a **two-phase truncation** strategy that separates the timing of truncation from its aggressiveness:
+
+- **First truncation**: Triggered immediately when a tool call result is **written into the context**, uniformly applied to all tools (including `read_file`). The full raw content is saved to `tool_result/{uuid}.txt`, and the message is annotated with the file path and a start-line hint.
+- **Second truncation**: Triggered by the Pre-Reasoning Hook on **messages that have slid out of the `recent_n` window**, applying a more aggressive truncation to further shrink context usage.
+
+```mermaid
+flowchart LR
+ A[Tool call completes] --> T[First truncation>>` marker:
+
+```python
+def truncate_text_output(text, start_line=1, total_lines=0,
+ max_bytes=DEFAULT_MAX_BYTES,
+ file_path=None, encoding="utf-8") -> str:
+ if TRUNCATION_NOTICE_MARKER in text:
+ return _retruncate(text, max_bytes=max_bytes, encoding=encoding)
+ else:
+ return _truncate_fresh(text, start_line=start_line,
+ total_lines=total_lines,
+ max_bytes=max_bytes,
+ file_path=file_path, encoding=encoding)
+```
+
+#### First Truncation (`_truncate_fresh`)
+
+**When it fires**: Immediately when the tool call completes and the result is written into the context — the text does not yet contain a truncation marker at this point.
+
+**Core logic**:
+
+1. If the text size in bytes does not exceed `max_bytes`, return the original text as-is.
+2. Otherwise, slice by bytes, keep the last complete line before the cut point, and compute the start line for the next read.
+3. Append a truncation notice (`<<>>`) at the end, prompting the reader to continue from `start_line=N`.
+
+**Example**: Suppose a tool returns 3,000 lines of HTML (200 KB in total), and `max_bytes = 50 KB`:
+
+```
+# Original tool output (200 KB, 3000 lines)
+
+ ...
+
+ ...(large content)
+
+
+
+# After first truncation, written to context (50 KB, ~750 lines)
+
+ ...
+
+ ...(first 750 lines)
+<<>>
+The output above was truncated.
+The full content is saved to the file and contains 3000 lines in total.
+This excerpt starts at line 1 and covers the next 51200 bytes.
+If the current content is not enough, call `read_file` with file_path=tool_result/abc123.txt start_line=751 to read more.
+```
+
+The full raw content is simultaneously written to `tool_result/abc123.txt`; only the truncated fragment and the continuation hint are kept in the context.
+
+#### Second Truncation (`_retruncate`)
+
+**When it fires**: During Pre-Reasoning Hook processing, applied to messages that have slid out of the `recent_n` window to further shrink context usage.
+
+**Core logic**:
+
+1. Split the text into the raw content before `<<>>` and the notice section after it.
+2. If the raw content still fits within the new `max_bytes` (with a 100-byte slack), return the original text as-is.
+3. Otherwise, re-slice according to the new, smaller byte limit and use regex to update the **byte count** and **continuation line number** in the notice; `file_path` remains unchanged.
+
+**Example**: The same tool message from above, after it slides out of `recent_n`. Second truncation reduces `max_bytes` from 50 KB to 10 KB:
+
+```
+# Before second truncation (first-truncation result already in context, 50 KB)
+
+ ...
+
+ ...(first 750 lines)
+<<>>
+...This excerpt starts at line 1 and covers the next 51200 bytes.
+...call `read_file` with file_path=tool_result/abc123.txt start_line=751 to read more.
+
+# After second truncation (further compressed to 10 KB, ~150 lines)
+
+ ...
+
+ ...(first 150 lines)
+<<>>
+...This excerpt starts at line 1 and covers the next 10240 bytes.
+...call `read_file` with file_path=tool_result/abc123.txt start_line=151 to read more.
+```
+
+Key point: `file_path` always points to `tool_result/abc123.txt`. The Agent can retrieve the full original content via the file reference at any time; truncation only affects the context fragment and meta info.
+
+---
+
+## Special Handling for the ReadFile Tool
+
+`read_file` shares the same two-phase truncation mechanism as Browser Use tools, with one key difference: **the file it reads already exists on the file system**, so there is no need to save a separate copy during first truncation.
+
+| Phase | Behavior |
+|----------------------------|-----------------------------------------------------------------------------------------------|
+| First truncation | Truncation happens at read time; result is written to context; original file path is already known — no need to save to `tool_result/` |
+| Within `recent_n` | Pre-Reasoning Hook makes no changes; keeps the read-time truncation result |
+| Outside `recent_n` (second truncation) | Same as other tools — more aggressive truncation is applied to the message content, meta info is updated |
+
+```mermaid
+flowchart LR
+ A[ReadFile call] -->|Truncate at read time| B[Context: truncated content>>` 标记来分发到两个不同的函数：
+
+```python
+def truncate_text_output(text, start_line=1, total_lines=0,
+ max_bytes=DEFAULT_MAX_BYTES,
+ file_path=None, encoding="utf-8") -> str:
+ if TRUNCATION_NOTICE_MARKER in text:
+ return _retruncate(text, max_bytes=max_bytes, encoding=encoding)
+ else:
+ return _truncate_fresh(text, start_line=start_line,
+ total_lines=total_lines,
+ max_bytes=max_bytes,
+ file_path=file_path, encoding=encoding)
+```
+
+#### 一次截断（`_truncate_fresh`）
+
+**触发时机**：工具调用完成、结果写入上下文时立即执行，此时文本中尚不含截断标记。
+
+**核心逻辑**：
+
+1. 若文本字节数未超过 `max_bytes`，直接返回原文；
+2. 否则按字节切片，保留截断点前最后一个完整行，计算下一段应从哪一行开始；
+3. 在末尾追加截断通知（`<<>>`），提示后续从 `start_line=N` 继续读取。
+
+**示例**：假设一个工具返回了 3 000 行的 HTML 内容（共 200 KB），而 `max_bytes = 50 KB`：
+
+```
+# 原始工具输出（200 KB，共 3000 行）
+
+ ...
+
+ ...（大量内容）
+
+
+
+# 一次截断后写入上下文（50 KB，约 750 行）
+
+ ...
+
+ ...（前 750 行）
+<<>>
+The output above was truncated.
+The full content is saved to the file and contains 3000 lines in total.
+This excerpt starts at line 1 and covers the next 51200 bytes.
+If the current content is not enough, call `read_file` with file_path=tool_result/abc123.txt start_line=751 to read more.
+```
+
+完整原始内容同时写入 `tool_result/abc123.txt`，上下文中仅保留截断片段与续读提示。
+
+#### 二次截断（`_retruncate`）
+
+**触发时机**：Pre-Reasoning Hook 处理时，对已滑出 `recent_n` 范围的历史消息执行，进一步压缩上下文占用。
+
+**核心逻辑**：
+
+1. 从文本中分离出 `<<>>` 前的原始内容与后面的通知部分；
+2. 若原始内容仍未超出新的 `max_bytes`（带 100 字节宽松量），直接返回原文；
+3. 否则按新的更小字节限制重新切片，并通过正则替换通知中的 **字节数** 与 **续读行号**，`file_path` 保持不变。
+
+**示例**：同样是上面那条工具消息，在它滑出 `recent_n` 之后，二次截断将 `max_bytes` 从 50 KB 压缩到 10 KB：
+
+```
+# 二次截断前（上下文中已有一次截断结果，50 KB）
+
+ ...
+
+ ...（前 750 行）
+<<>>
+...This excerpt starts at line 1 and covers the next 51200 bytes.
+...call `read_file` with file_path=tool_result/abc123.txt start_line=751 to read more.
+
+# 二次截断后（进一步压缩至 10 KB，约 150 行）
+
+ ...
+
+ ...（前 150 行）
+<<>>
+...This excerpt starts at line 1 and covers the next 10240 bytes.
+...call `read_file` with file_path=tool_result/abc123.txt start_line=151 to read more.
+```
+
+关键点：`file_path` 始终指向 `tool_result/abc123.txt`，Agent 随时可通过文件引用获取完整原始内容；截断只影响上下文片段与 meta 信息。
+
+---
+
+## ReadFile 工具的特殊处理
+
+`read_file` 与 Browser Use 类工具共享同一套两阶段截断机制，但有一个关键区别：**它读取的文件本身已存在于文件系统**
+，无需在一次截断时另行保存。
+
+| 阶段 | 行为 |
+|-------------------|---------------------------------------------------|
+| 一次截断 | 在读取时即完成截断，结果写入上下文；原始文件路径已知，无需额外保存到 `tool_result/` |
+| 在 recent_n 内 | Pre-Reasoning Hook 不做任何修改，保持读取时的截断结果 |
+| 超出 recent_n（二次截断） | 与其他工具相同，对消息内容做更激进的截断，更新 meta 信息 |
+
+```mermaid
+flowchart LR
+ A[ReadFile 调用] -->|读取时截断| B[上下文: 截断内容` (e.g. `root@test`).
+- **password**: Database password (seekdb Docker images commonly set this via `ROOT_PASSWORD`).
+- **database**: Logical database name (default: `test`).
+- **index_metric**: Distance metric for the vector index: `cosine` or `ip` (inner product); default `cosine`.
+- **index_ef_search**: HNSW `ef_search` parameter passed to pyobvector (default: `100`).
+- **collection_name**: Table name for the collection (from `VectorStoreConfig`, default `reme`). Use lowercase names if your deployment restricts identifiers.
+
+**Local seekdb via Docker**
+
+```text
+docker run -d --name reme_seekdb -p 2881:2881 -e ROOT_PASSWORD= quay.io/oceanbase/seekdb:latest
+```
+
+**Integration tests** (requires a running server, embedding API credentials in `.env`, and matching DB password):
+
+```shell
+OBVEC_PASSWORD= python tests/test_vector_store.py --obvec
+```
+
## Configuration File Examples
Configure Vector Store in `flowllm/config/default.yaml` under the `vector_store` section. The basic structure is as follows:
@@ -128,7 +151,7 @@ vector_store.default.params.=
### Configuration Field Descriptions
-- **`backend`** (required): Vector store backend type. Options: `local`, `memory`, `chroma`, `qdrant`, `elasticsearch`.
+- **`backend`** (required): Vector store backend type. Options: `local`, `memory`, `chroma`, `qdrant`, `elasticsearch`, `obvec`.
- **`embedding_model`** (required): Name of the embedding model configuration, referencing the `embedding_model` section.
- **`params`** (optional): Dictionary of backend-specific parameters passed to the vector store constructor.
@@ -295,6 +318,35 @@ vector_store.default.backend=elasticsearch
vector_store.default.params.hosts='["http://es-node1:9200", "http://es-node2:9200", "http://es-node3:9200"]'
```
+#### 6. ObVecVectorStore Configuration (OceanBase / seekdb)
+
+**Implementation**: [`reme/core/vector_store/obvec_vector_store.py`](https://github.com/agentscope-ai/ReMe/blob/main/reme/core/vector_store/obvec_vector_store.py)
+
+**Example (seekdb on localhost)**:
+
+```yaml
+vector_stores:
+ default:
+ backend: obvec
+ embedding_model: default
+ collection_name: reme
+ uri: "127.0.0.1:2881"
+ user: "root"
+ password: "your-root-password"
+ database: "test"
+ index_metric: "cosine"
+ index_ef_search: 100
+```
+
+```shell
+vector_stores.default.backend=obvec
+vector_stores.default.uri=127.0.0.1:2881
+vector_stores.default.user=root
+vector_stores.default.password=your-root-password
+```
+
+ReMe service YAML uses the key `vector_stores` (plural); CLI overrides use the same nested paths.
+
### Complete Configuration Example
Below is a complete `default.yaml` example including both embedding model and vector store configurations:
@@ -352,8 +404,9 @@ Two types of metadata filtering are supported:
- **Development & Testing**: Use MemoryVectorStore or LocalVectorStore—no additional services required.
- **Small-Scale Applications**: Use LocalVectorStore or ChromaVectorStore for simplicity and ease of use.
-- **Production Environments**: Use QdrantVectorStore or EsVectorStore for high performance and scalability.
+- **Production Environments**: Use QdrantVectorStore, EsVectorStore, or ObVecVectorStore (OceanBase/seekdb) for high performance and scalability, depending on your existing infrastructure.
- **Hybrid Search**: Use EsVectorStore to combine vector search with full-text search capabilities.
+- **OceanBase / seekdb**: Use ObVecVectorStore when you standardize on pyobvector and SQL-accessible vector tables.
## Important Notes
diff --git a/example.env b/example.env
index 5f39a48d..d1a27415 100644
--- a/example.env
+++ b/example.env
@@ -2,6 +2,4 @@ LLM_API_KEY=sk-xxxx
LLM_BASE_URL=https://xxxx/v1
#EMBEDDING_API_KEY=sk-xxxx
#EMBEDDING_BASE_URL=https://xxxx/v1
-LLM_MODEL_NAME=qwen3.5-plus
-
#TAVILY_API_KEY=xxxx
diff --git a/pyproject.toml b/pyproject.toml
index 309f03fa..c55004fe 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -33,7 +33,6 @@ classifiers = [
keywords = ["llm", "memory", "experience", "memoryscope", "ai", "mcp", "http", "reme", "personal"]
dependencies = [
- "flowllm[reme]>=0.2.0.10",
"sqlite-vec>=0.1.6",
"prompt_toolkit>=3.0.52",
"rich>=14.2.0",
@@ -45,13 +44,15 @@ dependencies = [
"fastapi>=0.121.3",
"fastmcp>=2.14.1",
"httpx>=0.28.1",
- "litellm>=1.80.0",
"loguru>=0.7.3",
"mcp>=1.25.0",
"numpy>=2.2.6",
"openai>=2.8.1",
"pandas>=2.3.3",
"pydantic>=2.12.4",
+ "pyobvector>=0.1.20",
+ # pyobvector imports Expression from sqlglot; removed from sqlglot 30+ top-level API
+ "sqlglot>=25,<30",
"qdrant-client>=1.16.0",
"tavily-python>=0.7.13",
"tiktoken>=0.12.0",
@@ -81,8 +82,13 @@ full = [
"reme_ai[dev,ray,light]",
]
+litellm = [
+ "litellm==1.80.0",
+]
+
light = [
- "agentscope==1.0.16.dev0",
+ "agentscope==1.0.18",
+ "flowllm[reme]>=0.2.0.10",
]
[tool.setuptools.packages.find]
diff --git a/reme/__init__.py b/reme/__init__.py
index 9842ce1a..c1dbeb42 100644
--- a/reme/__init__.py
+++ b/reme/__init__.py
@@ -6,7 +6,7 @@ from . import extension
from . import memory
from .reme import ReMe
-__version__ = "0.3.0.6b1"
+__version__ = "0.3.1.8"
__all__ = [
"config",
diff --git a/reme/config/cli.yaml b/reme/config/cli.yaml
index 7d348d45..e7cab227 100644
--- a/reme/config/cli.yaml
+++ b/reme/config/cli.yaml
@@ -8,12 +8,20 @@ metadata:
vector_weight: 0.7
candidate_multiplier: 2
-llms:
+as_llms:
default:
backend: openai
- # model_name: qwen3-235b-a22b-thinking-2507
model_name: qwen3.5-plus
- request_interval: 1
+
+as_llm_formatters:
+ default:
+ backend: openai
+
+as_token_counters:
+ default:
+ backend: hf
+ pretrained_model_name_or_path: Qwen/Qwen3-Coder-30B-A3B-Instruct
+ use_mirror: true
embedding_models:
default:
@@ -39,13 +47,4 @@ file_watchers:
watch_paths: [ ".reme", ".reme/memory" ]
suffix_filters: [ ".md" ]
recursive: false
- scan_on_start: true
-token_counters:
- default:
- backend: base
-
- hf:
- backend: hf
- model_name: Qwen/Qwen3-Coder-30B-A3B-Instruct
- use_mirror: true
diff --git a/reme/config/light.yaml b/reme/config/light.yaml
index 2a83371a..3500f9a5 100644
--- a/reme/config/light.yaml
+++ b/reme/config/light.yaml
@@ -1,11 +1,34 @@
+as_llms:
+ default:
+ backend: openai
+ model_name: qwen3.5-plus
+
+thread_pool_max_workers: -1
+
+as_llm_formatters:
+ default:
+ backend: openai
+
+as_token_counters:
+ default:
+ backend: rule
+ token_count_estimate_divisor: 3.75
+
embedding_models:
default:
backend: openai
+ dimensions: 1024
+ use_dimensions: false
+ enable_cache: true
+ max_batch_size: 10
+ max_cache_size: 2000
+ max_input_length: 8192
file_stores:
default:
backend: chroma
embedding_model: default
+ store_name: "reme"
file_watchers:
default:
@@ -13,4 +36,3 @@ file_watchers:
file_store: default
suffix_filters: [ ".md" ]
recursive: false
- scan_on_start: true
diff --git a/reme/config/service.yaml b/reme/config/service.yaml
index d736f215..31334099 100644
--- a/reme/config/service.yaml
+++ b/reme/config/service.yaml
@@ -66,7 +66,7 @@ flows:
description: "Whether to enable soft comparison between highest and lowest scoring trajectories (default: true)."
enable_similarity_comparison:
type: boolean
- description: "Whether to enable similarity-based comparison between success and failure trajectories (default: true)."
+ description: "Whether to enable similarity-based comparison between success and failure trajectories (default: false)."
max_similarity_sequences:
type: integer
description: "Maximum number of sequences to compare for similarity (default: 5)."
@@ -155,6 +155,7 @@ flows:
description: "The path to the memories file."
required:
- dump_file_path
+
test:
flow_content: TestOp()
description: "test"
diff --git a/reme/core/__init__.py b/reme/core/__init__.py
index 5872e2ad..726a3455 100644
--- a/reme/core/__init__.py
+++ b/reme/core/__init__.py
@@ -1,5 +1,8 @@
"""Core"""
+from . import as_llm
+from . import as_llm_formatter
+from . import as_token_counter
from . import embedding
from . import enumeration
from . import file_store
@@ -21,6 +24,9 @@ from .service_context import ServiceContext
__all__ = [
# Submodules
+ "as_llm",
+ "as_llm_formatter",
+ "as_token_counter",
"embedding",
"enumeration",
"file_watcher",
diff --git a/reme/core/application.py b/reme/core/application.py
index 49537807..8b5b1932 100644
--- a/reme/core/application.py
+++ b/reme/core/application.py
@@ -1,11 +1,10 @@
"""High-level entry point for configuring and running ReMe services and flows."""
import asyncio
+import os
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path
-from loguru import logger
-
from .embedding import BaseEmbeddingModel
from .file_store import BaseFileStore
from .file_watcher import BaseFileWatcher
@@ -13,12 +12,23 @@ from .flow import BaseFlow
from .llm import BaseLLM
from .prompt_handler import PromptHandler
from .registry_factory import R
-from .schema import Response, ServiceConfig
+from .schema import (
+ EmbeddingModelConfig,
+ Response,
+ ServiceConfig,
+ LLMConfig,
+ VectorStoreConfig,
+ FileStoreConfig,
+ FileWatcherConfig,
+ TokenCounterConfig,
+)
from .service_context import ServiceContext
from .token_counter import BaseTokenCounter
-from .utils import execute_stream_task, PydanticConfigParser, init_logger, MCPClient, print_logo
+from .utils import execute_stream_task, PydanticConfigParser, init_logger, MCPClient, print_logo, get_logger, load_env
from .vector_store import BaseVectorStore
+logger = get_logger()
+
class Application:
"""Application wrapper that wires together service context, flows, and runtimes."""
@@ -34,7 +44,10 @@ class Application:
config_path: str | None = None,
enable_logo: bool = True,
log_to_console: bool = True,
+ enable_load_env: bool = True,
parser: type[PydanticConfigParser] | None = None,
+ default_as_llm_config: dict | None = None,
+ default_as_llm_formatter_config: dict | None = None,
default_llm_config: dict | None = None,
default_embedding_model_config: dict | None = None,
default_vector_store_config: dict | None = None,
@@ -43,18 +56,25 @@ class Application:
default_file_watcher_config: dict | None = None,
**kwargs,
):
+
+ if enable_load_env:
+ load_env()
+
+ self.llm_api_key = llm_api_key or os.getenv("LLM_API_KEY", "")
+ self.llm_base_url = llm_base_url or os.getenv("LLM_BASE_URL", "")
+ self.embedding_api_key = embedding_api_key or os.getenv("EMBEDDING_API_KEY", "")
+ self.embedding_base_url = embedding_base_url or os.getenv("EMBEDDING_BASE_URL", "")
+
self.service_context = ServiceContext(
*args,
- llm_api_key=llm_api_key,
- llm_base_url=llm_base_url,
- embedding_api_key=embedding_api_key,
- embedding_base_url=embedding_base_url,
service_config=None,
parser=parser,
working_dir=working_dir,
config_path=config_path,
enable_logo=enable_logo,
log_to_console=log_to_console,
+ default_as_llm_config=default_as_llm_config,
+ default_as_llm_formatter_config=default_as_llm_formatter_config,
default_llm_config=default_llm_config,
default_embedding_model_config=default_embedding_model_config,
default_vector_store_config=default_vector_store_config,
@@ -136,22 +156,62 @@ class Application:
if not ray.is_initialized():
ray.init(num_cpus=self.service_config.ray_max_workers)
- if (
+ if self.service_config.thread_pool_max_workers > 0 and (
self.service_context.thread_pool is None
or self.service_context.thread_pool._shutdown # pylint: disable=protected-access
):
self.service_context.thread_pool = ThreadPoolExecutor(
max_workers=self.service_config.thread_pool_max_workers,
)
+ elif self.service_config.thread_pool_max_workers <= 0:
+ logger.info("Thread pool is disabled (thread_pool_max_workers <= 0)")
if self.service_context.service_config.enable_logo:
print_logo(service_config=self.service_config)
+ for name, config in self.service_config.as_llms.items():
+ if config.backend not in R.as_llms:
+ logger.warning(f"AS LLM backend {config.backend} is not supported.")
+ else:
+ try:
+ config_dict = config.model_dump(exclude={"backend"})
+ if not config_dict.get("api_key", ""):
+ config_dict["api_key"] = self.llm_api_key
+ if "client_kwargs" not in config_dict:
+ config_dict["client_kwargs"] = {}
+ if not config_dict["client_kwargs"].get("base_url", ""):
+ config_dict["client_kwargs"]["base_url"] = self.llm_base_url
+ self.service_context.as_llms[name] = R.as_llms[config.backend](**config_dict)
+ except Exception as e:
+ logger.error(f"Failed to initialize AS LLM '{name}': {e}")
+
+ for name, config in self.service_config.as_llm_formatters.items():
+ if config.backend not in R.as_llm_formatters:
+ logger.warning(f"AS LLM formatter backend {config.backend} is not supported.")
+ else:
+ try:
+ config_dict = config.model_dump(exclude={"backend"})
+ self.service_context.as_llm_formatters[name] = R.as_llm_formatters[config.backend](**config_dict)
+ except Exception as e:
+ logger.error(f"Failed to initialize AS LLM formatter '{name}': {e}")
+
+ for name, config in self.service_config.as_token_counters.items():
+ if config.backend not in R.as_token_counters:
+ logger.warning(f"Token counter backend {config.backend} is not supported.")
+ else:
+ try:
+ config_dict = config.model_dump(exclude={"backend"})
+ self.service_context.as_token_counters[name] = R.as_token_counters[config.backend](**config_dict)
+ except Exception as e:
+ logger.error(f"Failed to initialize AS token counter '{name}': {e}")
+
for name, config in self.service_config.llms.items():
if config.backend not in R.llms:
logger.warning(f"LLM backend {config.backend} is not supported.")
else:
config_dict = config.model_dump(exclude={"backend"})
+ config_dict.setdefault("api_key", self.llm_api_key)
+ config_dict.setdefault("base_url", self.llm_base_url)
self.service_context.llms[name] = R.llms[config.backend](**config_dict)
await self.service_context.llms[name].start()
@@ -160,7 +220,9 @@ class Application:
logger.warning(f"Embedding model backend {config.backend} is not supported.")
else:
config_dict = config.model_dump(exclude={"backend"})
- config_dict["cache_dir"] = working_path / "embedding_cache"
+ config_dict.setdefault("api_key", self.embedding_api_key)
+ config_dict.setdefault("base_url", self.embedding_base_url)
+ config_dict.setdefault("cache_dir", working_path / "embedding_cache")
self.service_context.embedding_models[name] = R.embedding_models[config.backend](**config_dict)
await self.service_context.embedding_models[name].start()
@@ -212,8 +274,221 @@ class Application:
await self.prepare_mcp_servers()
self._started = True
+ logger.info("ReMe Application started")
return self
+ # pylint: disable=too-many-statements
+ async def restart(self, restart_config: dict):
+ """Restart the application with new config."""
+
+ working_path = Path(self.service_config.working_dir)
+ working_path.mkdir(parents=True, exist_ok=True)
+
+ # as_llms
+ if "as_llms" in restart_config:
+ as_llms_config = restart_config["as_llms"]
+ assert isinstance(as_llms_config, dict)
+ for name, config in as_llms_config.items():
+ if name in self.service_context.as_llms:
+ del self.service_context.as_llms[name]
+
+ if config.get("backend") not in R.as_llms:
+ logger.warning(f"AS LLM backend {config.get('backend')} is not supported.")
+ continue
+
+ try:
+ config_dict = {k: v for k, v in config.items() if k != "backend"}
+ if not config_dict.get("api_key", ""):
+ config_dict["api_key"] = self.llm_api_key
+ if "client_kwargs" not in config_dict:
+ config_dict["client_kwargs"] = {}
+ if not config_dict["client_kwargs"].get("base_url", ""):
+ config_dict["client_kwargs"]["base_url"] = self.llm_base_url
+ self.service_context.as_llms[name] = R.as_llms[config["backend"]](**config_dict)
+ logger.info(f"Restarted AS LLM: {name}")
+ except Exception as e:
+ logger.error(f"Failed to restart AS LLM '{name}': {e}")
+
+ # as_llm_formatters
+ if "as_llm_formatters" in restart_config:
+ as_llm_formatters_config = restart_config["as_llm_formatters"]
+ assert isinstance(as_llm_formatters_config, dict)
+ for name, config in as_llm_formatters_config.items():
+ if name in self.service_context.as_llm_formatters:
+ del self.service_context.as_llm_formatters[name]
+
+ if config.get("backend") not in R.as_llm_formatters:
+ logger.warning(f"AS LLM formatter backend {config.get('backend')} is not supported.")
+ continue
+ try:
+ config_dict = {k: v for k, v in config.items() if k != "backend"}
+ self.service_context.as_llm_formatters[name] = R.as_llm_formatters[config["backend"]](**config_dict)
+ logger.info(f"Restarted AS LLM formatter: {name}")
+ except Exception as e:
+ logger.error(f"Failed to restart AS LLM formatter '{name}': {e}")
+
+ # as_token_counters
+ if "as_token_counters" in restart_config:
+ as_token_counters_config = restart_config["as_token_counters"]
+ assert isinstance(as_token_counters_config, dict)
+ for name, config in as_token_counters_config.items():
+ if name in self.service_context.as_token_counters:
+ del self.service_context.as_token_counters[name]
+
+ if config.get("backend") not in R.as_token_counters:
+ logger.warning(f"Token counter backend {config.get('backend')} is not supported.")
+ continue
+ try:
+ config_dict = {k: v for k, v in config.items() if k != "backend"}
+ self.service_context.as_token_counters[name] = R.as_token_counters[config["backend"]](**config_dict)
+ logger.info(f"Restarted AS token counter: {name}")
+ except Exception as e:
+ logger.error(f"Failed to restart AS token counter '{name}': {e}")
+
+ # llms
+ if "llms" in restart_config:
+ llms_config = restart_config["llms"]
+ assert isinstance(llms_config, dict)
+ for name, config in llms_config.items():
+ if name in self.service_context.llms:
+ llm = self.service_context.llms.pop(name)
+ await llm.close()
+
+ if isinstance(config, dict):
+ config = LLMConfig(**config)
+ if config.backend not in R.llms:
+ logger.warning(f"LLM backend {config.backend} is not supported.")
+ continue
+ config_dict = config.model_dump(exclude={"backend"})
+ config_dict.setdefault("api_key", self.llm_api_key)
+ config_dict.setdefault("base_url", self.llm_base_url)
+ self.service_context.llms[name] = R.llms[config.backend](**config_dict)
+ await self.service_context.llms[name].start()
+ logger.info(f"Restarted LLM: {name}")
+
+ # embedding_models
+ if "embedding_models" in restart_config:
+ embedding_models_config = restart_config["embedding_models"]
+ assert isinstance(embedding_models_config, dict)
+ updated_names = set()
+ for name, config in embedding_models_config.items():
+ if name in self.service_context.embedding_models:
+ embedding_model = self.service_context.embedding_models.pop(name)
+ await embedding_model.close()
+
+ if isinstance(config, dict):
+ config = EmbeddingModelConfig(**config)
+ if config.backend not in R.embedding_models:
+ logger.warning(f"Embedding model backend {config.backend} is not supported.")
+ continue
+ config_dict = config.model_dump(exclude={"backend"})
+ config_dict.setdefault("api_key", self.embedding_api_key)
+ config_dict.setdefault("base_url", self.embedding_base_url)
+ config_dict.setdefault("cache_dir", working_path / "embedding_cache")
+ self.service_context.embedding_models[name] = R.embedding_models[config.backend](**config_dict)
+ await self.service_context.embedding_models[name].start()
+ logger.info(f"Restarted embedding model: {name}")
+ updated_names.add(name)
+
+ # update embedding_model attribute for existing vector_stores and file_stores
+ for name in updated_names:
+ for vs_name, vs_config in self.service_config.vector_stores.items():
+ if vs_config.embedding_model == name and vs_name in self.service_context.vector_stores:
+ self.service_context.vector_stores[vs_name].embedding_model = (
+ self.service_context.embedding_models[name]
+ )
+ logger.info(f"Updated embedding model for vector store: {vs_name}")
+ for fs_name, fs_config in self.service_config.file_stores.items():
+ if fs_config.embedding_model == name and fs_name in self.service_context.file_stores:
+ self.service_context.file_stores[fs_name].embedding_model = (
+ self.service_context.embedding_models[name]
+ )
+ logger.info(f"Updated embedding model for file store: {fs_name}")
+
+ # token_counters
+ if "token_counters" in restart_config:
+ token_counters_config = restart_config["token_counters"]
+ assert isinstance(token_counters_config, dict)
+ for name, config in token_counters_config.items():
+ if name in self.service_context.token_counters:
+ del self.service_context.token_counters[name]
+
+ if isinstance(config, dict):
+ config = TokenCounterConfig(**config)
+ if config.backend not in R.token_counters:
+ logger.warning(f"Token counter backend {config.backend} is not supported.")
+ continue
+ config_dict = config.model_dump(exclude={"backend"})
+ self.service_context.token_counters[name] = R.token_counters[config.backend](**config_dict)
+ logger.info(f"Restarted token counter: {name}")
+
+ # vector_stores
+ if "vector_stores" in restart_config:
+ vector_stores_config = restart_config["vector_stores"]
+ assert isinstance(vector_stores_config, dict)
+ for name, config in vector_stores_config.items():
+ if name in self.service_context.vector_stores:
+ vector_store = self.service_context.vector_stores.pop(name)
+ await vector_store.close()
+ if isinstance(config, dict):
+ config = VectorStoreConfig(**config)
+ if config.backend not in R.vector_stores:
+ logger.warning(f"Vector store backend {config.backend} is not supported.")
+ continue
+ config_dict = config.model_dump(exclude={"backend", "embedding_model"})
+ config_dict.update(
+ {
+ "embedding_model": self.service_context.embedding_models[config.embedding_model],
+ "db_path": working_path / "vector_store",
+ },
+ )
+ self.service_context.vector_stores[name] = R.vector_stores[config.backend](**config_dict)
+ await self.service_context.vector_stores[name].start()
+ logger.info(f"Restarted vector store: {name}")
+
+ # file_stores
+ if "file_stores" in restart_config:
+ file_stores_config = restart_config["file_stores"]
+ assert isinstance(file_stores_config, dict)
+ for name, config in file_stores_config.items():
+ if name in self.service_context.file_stores:
+ file_store = self.service_context.file_stores.pop(name)
+ await file_store.close()
+ if isinstance(config, dict):
+ config = FileStoreConfig(**config)
+ if config.backend not in R.file_stores:
+ logger.warning(f"File store backend {config.backend} is not supported.")
+ continue
+ config_dict = config.model_dump(exclude={"backend", "embedding_model"})
+ config_dict.update(
+ {
+ "embedding_model": self.service_context.embedding_models[config.embedding_model],
+ "db_path": working_path / "file_store",
+ },
+ )
+ self.service_context.file_stores[name] = R.file_stores[config.backend](**config_dict)
+ await self.service_context.file_stores[name].start()
+ logger.info(f"Restarted file store: {name}")
+
+ # file_watchers
+ if "file_watchers" in restart_config:
+ file_watchers_config = restart_config["file_watchers"]
+ assert isinstance(file_watchers_config, dict)
+ for name, config in file_watchers_config.items():
+ if name in self.service_context.file_watchers:
+ file_watcher = self.service_context.file_watchers.pop(name)
+ await file_watcher.close()
+ if isinstance(config, dict):
+ config = FileWatcherConfig(**config)
+ if config.backend not in R.file_watchers:
+ logger.warning(f"File watcher backend {config.backend} is not supported.")
+ continue
+ config_dict = config.model_dump(exclude={"backend", "file_store"})
+ config_dict["file_store"] = self.service_context.file_stores[config.file_store]
+ self.service_context.file_watchers[name] = R.file_watchers[config.backend](**config_dict)
+ await self.service_context.file_watchers[name].start()
+ logger.info(f"Restarted file watcher: {name}")
+
async def prepare_mcp_servers(self):
"""Prepare and initialize MCP server connections."""
mcp_client = MCPClient(config={"mcpServers": self.service_config.mcp_servers})
@@ -258,11 +533,12 @@ class Application:
self.shutdown_ray()
self._started = False
+ logger.info("ReMe Application closed")
return False
def shutdown_thread_pool(self, wait: bool = True):
"""Shutdown the thread pool executor."""
- if self.service_context.thread_pool:
+ if self.service_context.thread_pool is not None:
self.service_context.thread_pool.shutdown(wait=wait)
def shutdown_ray(self, wait: bool = True):
diff --git a/reme/core/as_llm/__init__.py b/reme/core/as_llm/__init__.py
new file mode 100644
index 00000000..9cf527af
--- /dev/null
+++ b/reme/core/as_llm/__init__.py
@@ -0,0 +1,9 @@
+"""Module for registering AgentScope LLM models."""
+
+from agentscope.model import DashScopeChatModel
+from agentscope.model import OpenAIChatModel
+
+from ..registry_factory import R
+
+R.as_llms.register("openai")(OpenAIChatModel)
+R.as_llms.register("dashscope")(DashScopeChatModel)
diff --git a/reme/core/as_llm_formatter/__init__.py b/reme/core/as_llm_formatter/__init__.py
new file mode 100644
index 00000000..1c7eee46
--- /dev/null
+++ b/reme/core/as_llm_formatter/__init__.py
@@ -0,0 +1,9 @@
+"""Module for registering AgentScope LLM formatters."""
+
+from agentscope.formatter import DashScopeChatFormatter
+
+from .reme_openai_chat_formatter import ReMeOpenAIChatFormatter
+from ..registry_factory import R
+
+R.as_llm_formatters.register("openai")(ReMeOpenAIChatFormatter)
+R.as_llm_formatters.register("dashscope")(DashScopeChatFormatter)
diff --git a/reme/core/as_llm_formatter/reme_openai_chat_formatter.py b/reme/core/as_llm_formatter/reme_openai_chat_formatter.py
new file mode 100644
index 00000000..d6807dfb
--- /dev/null
+++ b/reme/core/as_llm_formatter/reme_openai_chat_formatter.py
@@ -0,0 +1,215 @@
+"""ReMeOpenAIChatFormatter"""
+
+import json
+from typing import Any
+
+from agentscope.formatter import OpenAIChatFormatter
+from agentscope.formatter._openai_formatter import (
+ _format_openai_image_block,
+ _to_openai_audio_data,
+)
+from agentscope.message import Msg, TextBlock, ImageBlock, URLSource
+from loguru import logger
+
+
+def _format_openai_video_block(video_block: dict) -> dict[str, Any]:
+ """Format a video block for OpenAI API.
+
+ Args:
+ video_block: The video block to format.
+
+ Returns:
+ A dictionary with video content in OpenAI format.
+ """
+ source = video_block["source"]
+ if source["type"] == "url":
+ url = source["url"]
+ elif source["type"] == "base64":
+ data = source["data"]
+ media_type = source["media_type"]
+ url = f"data:{media_type};base64,{data}"
+ else:
+ raise ValueError(f"Unsupported video source type: {source['type']}")
+
+ return {
+ "type": "video_url",
+ "video_url": {
+ "url": url,
+ },
+ }
+
+
+class ReMeOpenAIChatFormatter(OpenAIChatFormatter):
+ """ReMeOpenAIChatFormatter"""
+
+ async def _format(
+ self,
+ msgs: list[Msg],
+ ) -> list[dict[str, Any]]:
+ """Format message objects into OpenAI API required format.
+
+ Args:
+ msgs (`list[Msg]`):
+ The list of Msg objects to format.
+
+ Returns:
+ `list[dict[str, Any]]`:
+ A list of dictionaries, where each dictionary has "name",
+ "role", and "content" keys.
+ """
+ self.assert_list_of_msgs(msgs)
+
+ messages: list[dict] = []
+ i = 0
+ while i < len(msgs):
+ msg = msgs[i]
+ content_blocks = []
+ tool_calls = []
+ reasoning_content_blocks = []
+
+ for block in msg.get_content_blocks():
+ typ = block.get("type")
+ if typ == "text":
+ content_blocks.append({**block})
+
+ elif typ == "thinking":
+ # Collect thinking blocks for reasoning_content field
+ # This is compatible with models like DeepSeek that support
+ # extended thinking via reasoning_content field
+ reasoning_content_blocks.append({**block})
+
+ elif typ == "tool_use":
+ tool_calls.append(
+ {
+ "id": block.get("id"),
+ "type": "function",
+ "function": {
+ "name": block.get("name"),
+ "arguments": json.dumps(
+ block.get("input", {}),
+ ensure_ascii=False,
+ ),
+ },
+ },
+ )
+
+ elif typ == "tool_result":
+ (
+ textual_output,
+ multimodal_data,
+ ) = self.convert_tool_result_to_string(block["output"])
+
+ messages.append(
+ {
+ "role": "tool",
+ "tool_call_id": block.get("id"),
+ "content": (textual_output), # type: ignore[arg-type]
+ "name": block.get("name"),
+ },
+ )
+
+ # Then, handle the multimodal data if any
+ promoted_blocks: list = []
+ for url, multimodal_block in multimodal_data:
+ if multimodal_block["type"] == "image" and self.promote_tool_result_images:
+ promoted_blocks.extend(
+ [
+ TextBlock(
+ type="text",
+ text=f"\n- The image from '{url}': ",
+ ),
+ ImageBlock(
+ type="image",
+ source=URLSource(
+ type="url",
+ url=url,
+ ),
+ ),
+ ],
+ )
+
+ if promoted_blocks:
+ # Insert promoted blocks as new user message(s)
+ promoted_blocks = [
+ TextBlock(
+ type="text",
+ text="The following are "
+ "the image contents from the tool "
+ f"result of '{block['name']}':",
+ ),
+ *promoted_blocks,
+ TextBlock(
+ type="text",
+ text=" ",
+ ),
+ ]
+
+ msgs.insert(
+ i + 1,
+ Msg(
+ name="user",
+ content=promoted_blocks,
+ role="user",
+ ),
+ )
+
+ elif typ == "image":
+ content_blocks.append(
+ _format_openai_image_block(
+ block, # type: ignore[arg-type]
+ ),
+ )
+
+ elif typ == "audio":
+ # Filter out audio content when the multimodal model
+ # outputs both text and audio, to prevent errors in
+ # subsequent model calls
+ if msg.role == "assistant":
+ continue
+ input_audio = _to_openai_audio_data(block["source"])
+ content_blocks.append(
+ {
+ "type": "input_audio",
+ "input_audio": input_audio,
+ },
+ )
+
+ elif typ == "video":
+ # Filter out video content when the multimodal model
+ # outputs both text and video, to prevent errors in
+ # subsequent model calls
+ if msg.role == "assistant":
+ continue
+ content_blocks.append(
+ _format_openai_video_block(block),
+ )
+
+ else:
+ logger.warning(
+ "Unsupported block type %s in the message, skipped.",
+ typ,
+ )
+
+ msg_openai = {
+ "role": msg.role,
+ "name": msg.name,
+ "content": content_blocks or None,
+ }
+
+ if tool_calls:
+ msg_openai["tool_calls"] = tool_calls
+
+ # Add reasoning_content for thinking blocks (compatible with DeepSeek, etc.)
+ if reasoning_content_blocks:
+ reasoning_msg = "\n".join(reasoning.get("thinking", "") for reasoning in reasoning_content_blocks)
+ if reasoning_msg:
+ msg_openai["reasoning_content"] = reasoning_msg
+
+ # When both content and tool_calls are None, skipped
+ if msg_openai["content"] or msg_openai.get("tool_calls"):
+ messages.append(msg_openai)
+
+ # Move to next message
+ i += 1
+
+ return messages
diff --git a/reme/core/as_token_counter/__init__.py b/reme/core/as_token_counter/__init__.py
new file mode 100644
index 00000000..95910b8e
--- /dev/null
+++ b/reme/core/as_token_counter/__init__.py
@@ -0,0 +1,8 @@
+"""Module for registering AgentScope token counters."""
+
+from .reme_token_counter import ReMeTokenCounter
+from .rule_token_counter import RuleTokenCounter
+from ..registry_factory import R
+
+R.as_token_counters.register("hf")(ReMeTokenCounter)
+R.as_token_counters.register("rule")(RuleTokenCounter)
diff --git a/reme/core/as_token_counter/reme_token_counter.py b/reme/core/as_token_counter/reme_token_counter.py
new file mode 100644
index 00000000..6722e92f
--- /dev/null
+++ b/reme/core/as_token_counter/reme_token_counter.py
@@ -0,0 +1,123 @@
+"""Token counter for ReMe."""
+
+import os
+from typing import Any
+
+from agentscope.token import HuggingFaceTokenCounter
+
+from ..utils import get_logger
+
+logger = get_logger()
+
+
+class ReMeTokenCounter(HuggingFaceTokenCounter):
+ """Token counter for CoPaw with configurable tokenizer support.
+
+ This class extends HuggingFaceTokenCounter to provide token counting
+ functionality with support for both local and remote tokenizers,
+ as well as HuggingFace mirror for users in China.
+
+ Attributes:
+ pretrained_model_name_or_path: The tokenizer model path or "default" for local tokenizer.
+ use_mirror: Whether to use HuggingFace mirror.
+ token_count_estimate_divisor: Divisor for token estimation.
+ """
+
+ def __init__(
+ self,
+ pretrained_model_name_or_path: str,
+ use_mirror: bool = True,
+ token_count_estimate_divisor: float = 3.75,
+ **kwargs,
+ ):
+ """Initialize the token counter with the specified configuration.
+
+ Args:
+ pretrained_model_name_or_path: The tokenizer model path.
+ use_mirror: Whether to use the HuggingFace mirror
+ (https://hf-mirror.com) for downloading tokenizers. Useful for
+ users in China.
+ token_count_estimate_divisor: Divisor for estimating tokens when
+ tokenizer is unavailable. Defaults to 3.75.
+ **kwargs: Additional keyword arguments passed to HuggingFaceTokenCounter.
+ """
+ self.pretrained_model_name_or_path = pretrained_model_name_or_path
+ self.use_mirror = use_mirror
+ self.token_count_estimate_divisor = token_count_estimate_divisor
+
+ # Set HuggingFace endpoint for mirror support
+ if use_mirror:
+ mirror = "https://hf-mirror.com"
+ else:
+ mirror = "https://huggingface.co"
+
+ os.environ["HF_ENDPOINT"] = mirror
+
+ # if the huggingface is already imported in other dependencies,
+ # we need to set the endpoint manually
+ import huggingface_hub.constants
+
+ huggingface_hub.constants.ENDPOINT = mirror
+ huggingface_hub.constants.HUGGINGFACE_CO_URL_TEMPLATE = mirror + "/{repo_id}/resolve/{revision}/{filename}"
+
+ try:
+ super().__init__(
+ pretrained_model_name_or_path=self.pretrained_model_name_or_path,
+ use_mirror=use_mirror,
+ use_fast=True,
+ trust_remote_code=True,
+ **kwargs,
+ )
+ self._tokenizer_available = True
+
+ except Exception as e:
+ logger.error(f"Failed to initialize tokenizer {e}")
+ self._tokenizer_available = False
+
+ async def count(
+ self,
+ messages: list[dict],
+ tools: list[dict] | None = None,
+ text: str | None = None,
+ **kwargs: Any,
+ ) -> int:
+ """Count tokens in messages or text.
+
+ If text is provided, counts tokens directly in the text string.
+ Otherwise, counts tokens in the messages using the parent class method.
+
+ Args:
+ messages: List of message dictionaries in chat format.
+ tools: Optional list of tool definitions for token counting.
+ text: Optional text string to count tokens directly.
+ **kwargs: Additional keyword arguments passed to parent count method.
+
+ Returns:
+ The number of tokens, guaranteed to be at least the estimated minimum.
+ """
+ if text:
+ if self._tokenizer_available:
+ try:
+ token_ids = self.tokenizer.encode(text)
+ return max(len(token_ids), self.estimate_tokens(text))
+ except Exception as e:
+ logger.exception("Failed to encode text with tokenizer: %s", e)
+ return self.estimate_tokens(text)
+ else:
+ return self.estimate_tokens(text)
+ else:
+ return await super().count(messages, tools, **kwargs)
+
+ def estimate_tokens(self, text: str) -> int:
+ """Estimate the number of tokens in a text string.
+
+ Provides a fast character-based estimation as a fallback or lower bound.
+ Uses the configured divisor from instance settings.
+
+ Args:
+ text: The text string to estimate tokens for.
+
+ Returns:
+ The estimated number of tokens in the text string.
+ """
+ return int(len(text.encode("utf-8")) / self.token_count_estimate_divisor + 0.5)
diff --git a/reme/core/as_token_counter/rule_token_counter.py b/reme/core/as_token_counter/rule_token_counter.py
new file mode 100644
index 00000000..1201e185
--- /dev/null
+++ b/reme/core/as_token_counter/rule_token_counter.py
@@ -0,0 +1,78 @@
+"""Rule-based token counter for fast estimation without loading tokenizer."""
+
+from typing import Any
+
+from agentscope.token import HuggingFaceTokenCounter
+
+
+class RuleTokenCounter(HuggingFaceTokenCounter):
+ """Lightweight token counter using rule-based estimation only.
+
+ This class provides fast token estimation without loading any tokenizer,
+ useful when exact token counts are not critical or for quick approximations.
+
+ Attributes:
+ token_count_estimate_divisor: Divisor for token estimation.
+ """
+
+ def __init__(
+ self,
+ token_count_estimate_divisor: float = 3.75,
+ **_kwargs,
+ ):
+ """Initialize the rule-based token counter.
+
+ Args:
+ token_count_estimate_divisor: Divisor for estimating tokens.
+ Defaults to 3.75 (approximately 4 characters per token).
+ **kwargs: Additional keyword arguments (ignored).
+ """
+ self.token_count_estimate_divisor = token_count_estimate_divisor
+ # Skip tokenizer initialization from parent
+ self._tokenizer_available = False
+
+ async def count(
+ self,
+ messages: list[dict],
+ _tools: list[dict] | None = None,
+ text: str | None = None,
+ **_kwargs: Any,
+ ) -> int:
+ """Count tokens using rule-based estimation.
+
+ Args:
+ messages: List of message dictionaries in chat format.
+ _tools: Optional list of tool definitions (ignored).
+ text: Optional text string to count tokens directly.
+ **_kwargs: Additional keyword arguments (ignored).
+
+ Returns:
+ The estimated number of tokens.
+ """
+ if text:
+ return self.estimate_tokens(text)
+
+ # Estimate from messages
+ total_text = ""
+ for msg in messages:
+ content = msg.get("content", "")
+ if isinstance(content, str):
+ total_text += content
+ elif isinstance(content, list):
+ for part in content:
+ if isinstance(part, dict) and "text" in part:
+ total_text += part["text"]
+ return self.estimate_tokens(total_text)
+
+ def estimate_tokens(self, text: str) -> int:
+ """Estimate the number of tokens in a text string.
+
+ Uses character-based estimation with the configured divisor.
+
+ Args:
+ text: The text string to estimate tokens for.
+
+ Returns:
+ The estimated number of tokens in the text string.
+ """
+ return int(len(text.encode("utf-8")) / self.token_count_estimate_divisor + 0.5)
diff --git a/reme/core/embedding/base_embedding_model.py b/reme/core/embedding/base_embedding_model.py
index 3a02674f..b9871a34 100644
--- a/reme/core/embedding/base_embedding_model.py
+++ b/reme/core/embedding/base_embedding_model.py
@@ -29,7 +29,7 @@ class BaseEmbeddingModel(ABC):
api_key: str | None = None,
base_url: str | None = None,
model_name: str = "",
- dimensions: int | None = 1024,
+ dimensions: int = 1024,
use_dimensions: bool = False,
max_batch_size: int = 10,
max_retries: int = 3,
@@ -56,8 +56,8 @@ class BaseEmbeddingModel(ABC):
enable_cache: Whether to enable embedding cache
**kwargs: Additional model-specific parameters
"""
- self._api_key: str = api_key
- self._base_url: str = base_url
+ self._api_key: str | None = api_key
+ self._base_url: str | None = base_url
self.model_name = model_name
self.dimensions = dimensions
self.use_dimensions = use_dimensions
@@ -108,7 +108,34 @@ class BaseEmbeddingModel(ABC):
"""Truncate a list of texts to max_input_length."""
return [self._truncate_text(text) for text in texts]
- def _get_cache_key(self, text: str) -> str:
+ def _validate_and_adjust_embedding(self, embedding: list[float]) -> list[float]:
+ """Validate and adjust embedding dimensions to match expected dimensions.
+
+ Args:
+ embedding: The embedding vector to validate
+
+ Returns:
+ Embedding vector adjusted to match self.dimensions
+ """
+ actual_len = len(embedding)
+ if actual_len == self.dimensions:
+ return embedding
+
+ elif actual_len < self.dimensions:
+ logger.warning(
+ f"[ACTUAL_EMB_LENGTH]Embedding dimensions {actual_len} is less than expected {self.dimensions}, "
+ f"padding with zeros",
+ )
+ return embedding + [0.0] * (self.dimensions - actual_len)
+
+ else:
+ logger.warning(
+ f"[ACTUAL_EMB_LENGTH]Embedding dimensions {actual_len} is greater than expected {self.dimensions}, "
+ f"truncating to {self.dimensions}",
+ )
+ return embedding[: self.dimensions]
+
+ def _get_cache_key(self, text: str, dimensions: int) -> str:
"""Generate a cache key by hashing text + model_name + dimensions.
This ensures that the same text produces different cache keys when
@@ -116,12 +143,13 @@ class BaseEmbeddingModel(ABC):
Args:
text: Input text to hash
+ dimensions: Vector dimensions of the embeddings
Returns:
SHA256 hash combining text, model name, and dimensions
"""
# Combine text, model_name, and dimensions to create unique cache key
- cache_string = f"{text}|{self.model_name}|{self.dimensions}"
+ cache_string = f"{text}|{self.model_name}|{dimensions}"
return hashlib.sha256(cache_string.encode("utf-8")).hexdigest()
def _get_cache_file_path(self) -> Path:
@@ -151,6 +179,7 @@ class BaseEmbeddingModel(ABC):
return
try:
+ load_start = time.time()
# Read all lines first (to load in reverse order)
with open(cache_file, "r", encoding="utf-8") as f:
lines = f.readlines()
@@ -173,6 +202,13 @@ class BaseEmbeddingModel(ABC):
if cache_key in self._embedding_cache:
continue
+ if len(embedding) != self.dimensions:
+ logger.warning(
+ f"Embedding dimensions mismatch for cache key {cache_key}, "
+ f"expected {self.dimensions}, got {len(embedding)}",
+ )
+ continue
+
# Respect max_cache_size during loading
if len(self._embedding_cache) >= self.max_cache_size:
logger.info(
@@ -186,7 +222,9 @@ class BaseEmbeddingModel(ABC):
logger.warning(f"Failed to parse line in cache file: {e}")
continue
- logger.info(f"Loaded {loaded_count} embeddings from cache file: {cache_file}")
+ logger.info(
+ f"Loaded {loaded_count} embeddings from cache file: {cache_file} in {time.time() - load_start:.2f}s",
+ )
except Exception as e:
logger.error(f"Failed to load cache from {cache_file}: {e}, deleting cache file")
try:
@@ -213,6 +251,12 @@ class BaseEmbeddingModel(ABC):
try:
with open(cache_file, "w", encoding="utf-8") as f:
for cache_key, embedding in self._embedding_cache.items():
+ if len(embedding) != self.dimensions:
+ logger.warning(
+ f"Embedding dimensions mismatch for cache key {cache_key}, "
+ f"expected {self.dimensions}, got {len(embedding)}",
+ )
+ continue
cache_entry = {cache_key: embedding}
f.write(json.dumps(cache_entry, ensure_ascii=False) + "\n")
@@ -232,16 +276,27 @@ class BaseEmbeddingModel(ABC):
if not self.enable_cache:
return None
- cache_key = self._get_cache_key(text)
+ cache_key = self._get_cache_key(text, self.dimensions)
if cache_key in self._embedding_cache:
+ embeddings: list[float] = self._embedding_cache[cache_key]
+
+ # Validate embedding dimensions match expected dimensions
+ if len(embeddings) != self.dimensions:
+ logger.warning(
+ f"Cached embedding dimensions mismatch: expected {self.dimensions}, "
+ f"got {len(embeddings)}. Removing invalid cache entry.",
+ )
+ del self._embedding_cache[cache_key]
+ self._cache_misses += 1
+ return None
+
# Move to end (most recently used)
self._embedding_cache.move_to_end(cache_key)
self._cache_hits += 1
text_preview = text[:50] + "..." if len(text) > 50 else text
- logger.info(
- f"Cache hit for text: '{text_preview}' (hits: {self._cache_hits}, misses: {self._cache_misses})",
- )
- return self._embedding_cache[cache_key]
+ logger.info(f"Cache hit for text: {text_preview} (hits: {self._cache_hits}, misses: {self._cache_misses})")
+ return embeddings
+
self._cache_misses += 1
return None
@@ -258,9 +313,15 @@ class BaseEmbeddingModel(ABC):
if self.max_cache_size <= 0:
return
- cache_key = self._get_cache_key(text)
+ cache_key = self._get_cache_key(text, self.dimensions)
+ if len(embedding) != self.dimensions:
+ logger.warning(
+ f"[PUT_TO_CACHE] Embedding dimensions mismatch for cache key {cache_key}, "
+ f"expected {self.dimensions}, got real length {len(embedding)}",
+ )
+ return
- # Remove oldest entry if cache is full
+ # Remove the oldest entry if cache is full
if len(self._embedding_cache) >= self.max_cache_size and cache_key not in self._embedding_cache:
self._embedding_cache.popitem(last=False)
@@ -308,7 +369,7 @@ class BaseEmbeddingModel(ABC):
for i in range(self.max_retries):
try:
result = await self._get_embeddings([truncated_text], **kwargs)
- embedding = result[0]
+ embedding = self._validate_and_adjust_embedding(result[0])
# Store in cache
self._put_to_cache(truncated_text, embedding)
return embedding
@@ -354,8 +415,9 @@ class BaseEmbeddingModel(ABC):
if batch_embeddings:
# Store results and cache them
for orig_idx, text, embedding in zip(batch_indices, batch_texts, batch_embeddings):
- results[orig_idx] = embedding
- self._put_to_cache(text, embedding)
+ adjusted_embedding = self._validate_and_adjust_embedding(embedding)
+ results[orig_idx] = adjusted_embedding
+ self._put_to_cache(text, adjusted_embedding)
break
except Exception as e:
logger.error(f"Model {self.model_name} batch failed: {e}")
@@ -380,7 +442,7 @@ class BaseEmbeddingModel(ABC):
for i in range(self.max_retries):
try:
result = self._get_embeddings_sync([truncated_text], **kwargs)
- embedding = result[0]
+ embedding = self._validate_and_adjust_embedding(result[0])
# Store in cache
self._put_to_cache(truncated_text, embedding)
return embedding
@@ -426,8 +488,9 @@ class BaseEmbeddingModel(ABC):
if batch_embeddings:
# Store results and cache them
for orig_idx, text, embedding in zip(batch_indices, batch_texts, batch_embeddings):
- results[orig_idx] = embedding
- self._put_to_cache(text, embedding)
+ adjusted_embedding = self._validate_and_adjust_embedding(embedding)
+ results[orig_idx] = adjusted_embedding
+ self._put_to_cache(text, adjusted_embedding)
break
except Exception as exc:
logger.error(f"Model {self.model_name} batch failed: {exc}")
diff --git a/reme/core/embedding/openai_embedding_model.py b/reme/core/embedding/openai_embedding_model.py
index 3857f782..d686c83c 100644
--- a/reme/core/embedding/openai_embedding_model.py
+++ b/reme/core/embedding/openai_embedding_model.py
@@ -45,7 +45,9 @@ class OpenAIEmbeddingModel(BaseEmbeddingModel):
result_emb = [[] for _ in range(len(input_text))]
for emb in completion.data:
- result_emb[emb.index] = emb.embedding
+ # BGE-M3 returns dense_embedding instead of embedding; use as fallback
+ vec = getattr(emb, "embedding", None) or getattr(emb, "dense_embedding", None)
+ result_emb[emb.index] = list(vec) if vec is not None else []
return result_emb
async def start(self):
diff --git a/reme/core/file_store/base_file_store.py b/reme/core/file_store/base_file_store.py
index 9c79f81f..6a6b0901 100644
--- a/reme/core/file_store/base_file_store.py
+++ b/reme/core/file_store/base_file_store.py
@@ -7,6 +7,9 @@ from pathlib import Path
from ..embedding import BaseEmbeddingModel
from ..enumeration import MemorySource
from ..schema import FileMetadata, MemoryChunk, MemorySearchResult
+from ..utils import get_logger
+
+logger = get_logger()
class BaseFileStore(ABC):
@@ -54,24 +57,46 @@ class BaseFileStore(ABC):
"""Generate a zero vector based on embedding model dimensions."""
return [0.0] * self.embedding_dim
+ def _disable_vector_search(self, reason: str = "embedding API error") -> None:
+ """Disable vector search and log a warning."""
+ if self.vector_enabled:
+ logger.warning(
+ f"[{self.store_name}] Disabling vector search due to {reason}. "
+ "Falling back to full-text search only.",
+ )
+ self.vector_enabled = False
+
async def get_embedding(self, query: str, **kwargs) -> list[float]:
"""Get embedding for a single query string."""
if not self.vector_enabled:
return self._get_mock_embedding()
- return await self.embedding_model.get_embedding(query, **kwargs)
+ try:
+ return await self.embedding_model.get_embedding(query, **kwargs)
+ except Exception as e:
+ self._disable_vector_search(str(e))
+ return self._get_mock_embedding()
async def get_embeddings(self, queries: list[str], **kwargs) -> list[list[float]]:
"""Get embeddings for a batch of query strings."""
if not self.vector_enabled:
return [self._get_mock_embedding() for _ in queries]
- return await self.embedding_model.get_embeddings(queries, **kwargs)
+ try:
+ return await self.embedding_model.get_embeddings(queries, **kwargs)
+ except Exception as e:
+ self._disable_vector_search(str(e))
+ return [self._get_mock_embedding() for _ in queries]
async def get_chunk_embedding(self, chunk: MemoryChunk, **kwargs) -> MemoryChunk:
"""Generate and populate embedding field for a single MemoryChunk object."""
if not self.vector_enabled:
chunk.embedding = self._get_mock_embedding()
return chunk
- return await self.embedding_model.get_chunk_embedding(chunk, **kwargs)
+ try:
+ return await self.embedding_model.get_chunk_embedding(chunk, **kwargs)
+ except Exception as e:
+ self._disable_vector_search(str(e))
+ chunk.embedding = self._get_mock_embedding()
+ return chunk
async def get_chunk_embeddings(self, chunks: list[MemoryChunk], **kwargs) -> list[MemoryChunk]:
"""Generate and populate embedding fields for a batch of MemoryChunk objects."""
@@ -80,7 +105,14 @@ class BaseFileStore(ABC):
for chunk in chunks:
chunk.embedding = mock_embedding.copy()
return chunks
- return await self.embedding_model.get_chunk_embeddings(chunks, **kwargs)
+ try:
+ return await self.embedding_model.get_chunk_embeddings(chunks, **kwargs)
+ except Exception as e:
+ self._disable_vector_search(str(e))
+ mock_embedding = self._get_mock_embedding()
+ for chunk in chunks:
+ chunk.embedding = mock_embedding.copy()
+ return chunks
@abstractmethod
async def start(self):
diff --git a/reme/core/file_store/chroma_file_store.py b/reme/core/file_store/chroma_file_store.py
index 390a2bc4..6dda41b8 100644
--- a/reme/core/file_store/chroma_file_store.py
+++ b/reme/core/file_store/chroma_file_store.py
@@ -1,22 +1,24 @@
"""ChromaDB storage backend for file store."""
import json
+import random
import time
from pathlib import Path
-from loguru import logger
-
from .base_file_store import BaseFileStore
from ..enumeration import MemorySource
from ..schema import FileMetadata, MemoryChunk, MemorySearchResult
+from ..utils import get_logger
+
+logger = get_logger()
try:
import chromadb
from chromadb.config import Settings
- CHROMADB_AVAILABLE = True
-except ImportError:
- CHROMADB_AVAILABLE = False
+ _CHROMADB_IMPORT_ERROR: Exception | None = None
+except Exception as e:
+ _CHROMADB_IMPORT_ERROR = e
chromadb = None
Settings = None
@@ -38,10 +40,8 @@ class ChromaFileStore(BaseFileStore):
self,
**kwargs,
):
- if not CHROMADB_AVAILABLE:
- raise ImportError(
- "chromadb package is required for ChromaFileStore. Install it with: pip install chromadb",
- )
+ if _CHROMADB_IMPORT_ERROR is not None:
+ raise _CHROMADB_IMPORT_ERROR
super().__init__(**kwargs)
self.client: "chromadb.ClientAPI | None" = None
@@ -355,12 +355,41 @@ class ChromaFileStore(BaseFileStore):
where_filter = {"source": {"$in": [s.value for s in sources]}}
# Perform vector search
- results = self.chunks_collection.query(
- query_embeddings=[query_embedding],
- n_results=limit,
- where=where_filter,
- include=["documents", "metadatas", "distances"],
- )
+ try:
+ results = self.chunks_collection.query(
+ query_embeddings=[query_embedding],
+ n_results=limit,
+ where=where_filter,
+ include=["documents", "metadatas", "distances"],
+ )
+ except Exception as e:
+ logger.error(f"Vector search failed: {e}, falling back to random results")
+ # Fallback: get some documents without vector search and assign random scores
+ try:
+ fallback_results = self.chunks_collection.get(
+ where=where_filter,
+ limit=limit,
+ include=["documents", "metadatas"],
+ )
+ search_results = []
+ if fallback_results["ids"]:
+ for i, _ in enumerate(fallback_results["ids"]):
+ metadata = fallback_results["metadatas"][i]
+ search_results.append(
+ MemorySearchResult(
+ path=metadata["path"],
+ start_line=metadata["start_line"],
+ end_line=metadata["end_line"],
+ score=random.uniform(0.3, 0.7), # Random score in middle range
+ snippet=fallback_results["documents"][i],
+ source=MemorySource(metadata["source"]),
+ raw_metric=None,
+ ),
+ )
+ return search_results
+ except Exception as fallback_e:
+ logger.error(f"Fallback search also failed: {fallback_e}")
+ return []
search_results = []
if results["ids"] and results["ids"][0]:
@@ -430,7 +459,7 @@ class ChromaFileStore(BaseFileStore):
# ChromaDB where_document uses $contains for substring matching (case-sensitive)
# Use multiple case variants to improve recall
if len(word_variants_list) == 1:
- where_document = {"$contains": word_variants_list[0]}
+ where_document: dict = {"$contains": word_variants_list[0]}
else:
where_document = {"$or": [{"$contains": w} for w in word_variants_list]}
diff --git a/reme/core/file_store/local_file_store.py b/reme/core/file_store/local_file_store.py
index 37df17f4..38757d7a 100644
--- a/reme/core/file_store/local_file_store.py
+++ b/reme/core/file_store/local_file_store.py
@@ -259,6 +259,8 @@ class LocalFileStore(BaseFileStore):
if not query_embedding:
return []
+ expected_dim = self.embedding_dim
+
# Collect candidate chunks with embeddings
candidates = [
chunk for chunk in self._chunks.values() if (not sources or chunk.source in sources) and chunk.embedding
@@ -267,9 +269,29 @@ class LocalFileStore(BaseFileStore):
if not candidates:
return []
+ # Validate and fix chunk embedding dimensions
+ valid_embeddings = []
+ for chunk in candidates:
+ emb = chunk.embedding
+ emb_len = len(emb)
+ if emb_len != expected_dim:
+ if emb_len < expected_dim:
+ emb = emb + [0.0] * (expected_dim - emb_len)
+ logger.warning(
+ f"Chunk embedding dimension {emb_len} < expected {expected_dim}, "
+ f"padded with zeros (chunk_id={chunk.id})",
+ )
+ else:
+ emb = emb[:expected_dim]
+ logger.warning(
+ f"Chunk embedding dimension {emb_len} > expected {expected_dim}, "
+ f"truncated to {expected_dim} (chunk_id={chunk.id})",
+ )
+ valid_embeddings.append(emb)
+
# Build embedding matrix and compute similarities in batch
query_array = np.array([query_embedding]) # Shape: (1, emb_size)
- chunk_embeddings = np.array([chunk.embedding for chunk in candidates]) # Shape: (n, emb_size)
+ chunk_embeddings = np.array(valid_embeddings) # Shape: (n, emb_size)
similarities = batch_cosine_similarity(query_array, chunk_embeddings)[0] # Shape: (n,)
# Build results
diff --git a/reme/core/file_store/sqlite_file_store.py b/reme/core/file_store/sqlite_file_store.py
index 7fd49762..0a494d2c 100644
--- a/reme/core/file_store/sqlite_file_store.py
+++ b/reme/core/file_store/sqlite_file_store.py
@@ -1,7 +1,7 @@
"""SQLite storage backend for file store."""
import json
-import sqlite3
+
import struct
import time
@@ -29,6 +29,7 @@ class SqliteFileStore(BaseFileStore):
def __init__(self, vec_ext_path: str = "", **kwargs):
super().__init__(**kwargs)
self.vec_ext_path = vec_ext_path
+ import sqlite3
self.conn: sqlite3.Connection | None = None
@@ -61,6 +62,7 @@ class SqliteFileStore(BaseFileStore):
"""Initialize database and load extensions."""
if self.conn is not None:
return
+ import sqlite3
self.conn = sqlite3.connect(self.db_path / "reme.db", check_same_thread=False)
diff --git a/reme/core/file_watcher/base_file_watcher.py b/reme/core/file_watcher/base_file_watcher.py
index 5030a9cb..b4644dd7 100644
--- a/reme/core/file_watcher/base_file_watcher.py
+++ b/reme/core/file_watcher/base_file_watcher.py
@@ -29,12 +29,13 @@ class BaseFileWatcher:
watch_paths: list[str] | str,
suffix_filters: list[str] | None = None,
recursive: bool = False,
- debounce: int = 500, # Millisecond debounce
+ debounce: int = 2000,
chunk_tokens: int = 400,
chunk_overlap: int = 80,
file_store: BaseFileStore | None = None,
callback: Callable[[set[tuple[Change, str]]], None | Coroutine[Any, Any, None]] | None = None,
- scan_on_start: bool = False,
+ rebuild_index_on_start: bool = True,
+ poll_delay_ms: int = 2000,
**kwargs,
):
"""
@@ -49,7 +50,9 @@ class BaseFileWatcher:
chunk_overlap: Overlap size for chunks
file_store: File store instance
callback: Callback function for changes
- scan_on_start: If True, scan existing files on start and trigger on_changes with Change.added
+ rebuild_index_on_start: If True, clear all indexed data on start and rescan existing files.
+ If False, only monitor new changes without initialization.
+ poll_delay_ms: Polling delay in milliseconds. If > 300ms, force_polling will be enabled automatically.
**kwargs: Additional keyword arguments
"""
self.watch_paths: list[str] = [watch_paths] if isinstance(watch_paths, str) else watch_paths
@@ -60,7 +63,8 @@ class BaseFileWatcher:
self.chunk_overlap: int = chunk_overlap
self.file_store: BaseFileStore = file_store
self.callback = callback
- self.scan_on_start: bool = scan_on_start
+ self.rebuild_index_on_start: bool = rebuild_index_on_start
+ self.poll_delay_ms: int = poll_delay_ms
self.kwargs: dict = kwargs
self._stop_event = asyncio.Event()
@@ -74,11 +78,14 @@ class BaseFileWatcher:
self._running = True
- # Scan existing files if requested
- if self.scan_on_start:
- await self._scan_existing_files()
+ async def _initialize_and_watch():
+ if self.rebuild_index_on_start:
+ await self.file_store.clear_all()
+ logger.info("Cleared all indexed data on start")
+ await self._scan_existing_files()
+ await self._watch_loop()
- self._watch_task = asyncio.create_task(self._watch_loop())
+ self._watch_task = asyncio.create_task(_initialize_and_watch())
logger.info(f"Started watching: {self.watch_paths}")
async def close(self):
@@ -140,35 +147,64 @@ class BaseFileWatcher:
else:
logger.info("[SCAN_ON_START] No existing files found matching watch criteria")
- files: list[str] = await self.file_store.list_files(MemorySource.MEMORY)
- for file_path in files:
- chunks = await self.file_store.get_file_chunks(file_path, MemorySource.MEMORY)
- logger.info(f"Found existing file: {file_path}, {len(chunks)} chunks")
+ if self.file_store is not None:
+ files: list[str] = await self.file_store.list_files(MemorySource.MEMORY)
+ for file_path in files:
+ chunks = await self.file_store.get_file_chunks(file_path, MemorySource.MEMORY)
+ logger.info(f"Found existing file: {file_path}, {len(chunks)} chunks")
+
+ async def _interruptible_sleep(self, seconds: float):
+ """Sleep that can be interrupted by stop_event."""
+ try:
+ await asyncio.wait_for(self._stop_event.wait(), timeout=seconds)
+ except asyncio.TimeoutError:
+ pass # Normal timeout, continue
async def _watch_loop(self):
- """Core monitoring loop"""
+ """Core monitoring loop with auto-restart on failure"""
if not self.watch_paths:
logger.warning("No watch paths specified")
return
- try:
- async for changes in awatch(
- *self.watch_paths,
- watch_filter=self.watch_filter,
- recursive=self.recursive,
- debounce=self.debounce,
- stop_event=self._stop_event,
- ):
- if self._stop_event.is_set():
- break
+ while not self._stop_event.is_set():
+ # Filter out non-existent paths before each watch attempt
+ valid_paths = [p for p in self.watch_paths if Path(p).exists()]
- await self.on_changes(changes)
- except FileNotFoundError as e:
- # Watch path was deleted, this is expected during cleanup
- logger.debug(f"Watch path no longer exists: {e}")
- except Exception as e:
- # Log other exceptions but don't crash
- logger.error(f"Error in watch loop: {e}", exc_info=True)
+ if not valid_paths:
+ logger.warning("No valid watch paths exist, waiting 10 seconds before retry...")
+ await self._interruptible_sleep(10)
+ continue
+
+ invalid_paths = set(self.watch_paths) - set(valid_paths)
+ if invalid_paths:
+ logger.warning(f"Skipping non-existent paths: {invalid_paths}")
+
+ try:
+ logger.info(f"Starting watch on valid paths: {valid_paths}")
+ async for changes in awatch(
+ *valid_paths,
+ watch_filter=self.watch_filter,
+ recursive=self.recursive,
+ debounce=self.debounce,
+ poll_delay_ms=self.poll_delay_ms,
+ stop_event=self._stop_event,
+ ):
+ if self._stop_event.is_set():
+ break
+
+ await self.on_changes(changes)
+
+ except FileNotFoundError as e:
+ # Watch path was deleted during monitoring
+ logger.error(f"Watch path no longer exists: {e}, restarting in 10 seconds...")
+ if not self._stop_event.is_set():
+ await self._interruptible_sleep(10)
+
+ except Exception as e:
+ # Log other exceptions and restart
+ logger.error(f"Error in watch loop: {e}, restarting in 10 seconds...", exc_info=True)
+ if not self._stop_event.is_set():
+ await self._interruptible_sleep(10)
async def _on_changes(self, changes: set[tuple[Change, str]]):
"""Callback method to handle file changes"""
diff --git a/reme/core/file_watcher/full_file_watcher.py b/reme/core/file_watcher/full_file_watcher.py
index 2d365f23..5b1852b9 100644
--- a/reme/core/file_watcher/full_file_watcher.py
+++ b/reme/core/file_watcher/full_file_watcher.py
@@ -44,6 +44,7 @@ class FullFileWatcher(BaseFileWatcher):
async def _on_changes(self, changes: set[tuple[Change, str]]):
"""Handle file changes with full synchronization"""
self.dirty = True
+
for change_type, path in changes:
if change_type in [Change.added, Change.modified]:
file_meta = await self._build_file_metadata(path)
diff --git a/reme/core/llm/base_llm.py b/reme/core/llm/base_llm.py
index ba13ff10..c08a423c 100644
--- a/reme/core/llm/base_llm.py
+++ b/reme/core/llm/base_llm.py
@@ -2,7 +2,6 @@
import asyncio
import json
-import os
import time
from abc import ABC, abstractmethod
from typing import Callable, Generator, AsyncGenerator, Any
@@ -36,8 +35,8 @@ class BaseLLM(ABC):
request_interval: Minimum seconds between requests (default: 0.0)
**kwargs: Additional model-specific parameters
"""
- self._api_key: str = api_key
- self._base_url: str = base_url
+ self.api_key: str = api_key
+ self.base_url: str = base_url
self.model_name: str = model_name
self.max_retries: int = max_retries
self.raise_exception: bool = raise_exception
@@ -47,16 +46,6 @@ class BaseLLM(ABC):
self._last_request_time: float = 0.0
self._request_lock: asyncio.Lock = asyncio.Lock()
- @property
- def api_key(self) -> str | None:
- """Get API key from environment variable."""
- return os.getenv("LLM_API_KEY") or self._api_key
-
- @property
- def base_url(self) -> str | None:
- """Get base URL from environment variable."""
- return os.getenv("LLM_BASE_URL") or self._base_url
-
@staticmethod
def _accumulate_tool_call_chunk(tool_call, ret_tools: list[ToolCall]):
"""Assemble incremental tool call chunks into complete ToolCall objects."""
diff --git a/reme/core/op/base_op.py b/reme/core/op/base_op.py
index af5158c4..b016bb47 100644
--- a/reme/core/op/base_op.py
+++ b/reme/core/op/base_op.py
@@ -7,6 +7,9 @@ from abc import ABCMeta
from pathlib import Path
from typing import Callable, Optional, Any
+from agentscope.formatter import FormatterBase
+from agentscope.model import ChatModelBase
+from agentscope.token import HuggingFaceTokenCounter
from loguru import logger
from tqdm import tqdm
@@ -42,6 +45,9 @@ class BaseOp(metaclass=ABCMeta):
language: str = "",
prompt_name: str = "",
prompt_path: str = "",
+ as_llm: str | ChatModelBase = "default",
+ as_llm_formatter: str | FormatterBase = "default",
+ as_token_counter: str | HuggingFaceTokenCounter = "default",
llm: str | BaseLLM = "default",
embedding_model: str | BaseEmbeddingModel = "default",
vector_store: str | BaseVectorStore = "default",
@@ -64,6 +70,9 @@ class BaseOp(metaclass=ABCMeta):
self.language = language
self.prompt = self._get_prompt_handler(prompt_name, prompt_path)
+ self._as_llm = as_llm
+ self._as_llm_formatter = as_llm_formatter
+ self._as_token_counter = as_token_counter
self._llm = llm
self._embedding_model = embedding_model
self._vector_store = vector_store
@@ -129,6 +138,27 @@ class BaseOp(metaclass=ABCMeta):
"""Access the service configuration."""
return self.service_context.service_config
+ @property
+ def as_llm(self) -> ChatModelBase:
+ """Get the AgentScope LLM instance from ServiceContext."""
+ if isinstance(self._as_llm, str):
+ self._as_llm = self.service_context.as_llms[self._as_llm]
+ return self._as_llm
+
+ @property
+ def as_llm_formatter(self) -> FormatterBase:
+ """Get the AgentScope LLM formatter instance from ServiceContext."""
+ if isinstance(self._as_llm_formatter, str):
+ self._as_llm_formatter = self.service_context.as_llm_formatters[self._as_llm_formatter]
+ return self._as_llm_formatter
+
+ @property
+ def as_token_counter(self) -> HuggingFaceTokenCounter:
+ """Get the token counter instance from ServiceContext."""
+ if isinstance(self._as_token_counter, str):
+ self._as_token_counter = self.service_context.as_token_counters[self._as_token_counter]
+ return self._as_token_counter
+
@property
def llm(self) -> BaseLLM:
"""Get the LLM instance from ServiceContext."""
@@ -271,7 +301,7 @@ class BaseOp(metaclass=ABCMeta):
def submit_sync_task(self, fn: Callable, *args, **kwargs) -> "BaseOp":
"""Submit a task to the thread pool or local queue."""
- if self.enable_parallel:
+ if self.enable_parallel and self.service_context.thread_pool is not None:
task = self.service_context.thread_pool.submit(fn, *args, **kwargs)
else:
task = (fn, args, kwargs)
diff --git a/reme/core/op/base_ray_op.py b/reme/core/op/base_ray_op.py
index f63e7b83..2086d953 100644
--- a/reme/core/op/base_ray_op.py
+++ b/reme/core/op/base_ray_op.py
@@ -10,11 +10,11 @@ from tqdm import tqdm
from .base_op import BaseOp
from ..base_dict import BaseDict
-_RAY_IMPORT_ERROR = None
+_RAY_IMPORT_ERROR: Exception | None = None
try:
import ray
-except ImportError as _e:
+except Exception as _e:
_RAY_IMPORT_ERROR = _e
ray = None
diff --git a/reme/core/registry_factory.py b/reme/core/registry_factory.py
index b319b049..921f27a8 100644
--- a/reme/core/registry_factory.py
+++ b/reme/core/registry_factory.py
@@ -34,6 +34,9 @@ class RegistryFactory:
def __init__(self):
self.llms = Registry()
+ self.as_llms = Registry()
+ self.as_llm_formatters = Registry()
+ self.as_token_counters = Registry()
self.embedding_models = Registry()
self.vector_stores = Registry()
self.file_stores = Registry()
diff --git a/reme/core/schema/__init__.py b/reme/core/schema/__init__.py
index de167c69..b6445a28 100644
--- a/reme/core/schema/__init__.py
+++ b/reme/core/schema/__init__.py
@@ -1,5 +1,6 @@
"""schema"""
+from .as_msg_stat import AsBlockStat, AsMsgStat
from .cut_point_result import CutPointResult
from .file_metadata import FileMetadata
from .memory_chunk import MemoryChunk
@@ -27,6 +28,8 @@ from .truncation_result import TruncationResult
from .vector_node import VectorNode
__all__ = [
+ "AsBlockStat",
+ "AsMsgStat",
"CutPointResult",
"CmdConfig",
"ContentBlock",
diff --git a/reme/core/schema/as_msg_stat.py b/reme/core/schema/as_msg_stat.py
new file mode 100644
index 00000000..a8d2daaa
--- /dev/null
+++ b/reme/core/schema/as_msg_stat.py
@@ -0,0 +1,96 @@
+"""Schema definitions for AgentScope message statistics."""
+
+from pydantic import BaseModel, Field
+
+_TRUNCATION_NOTICE_MARKER = "<<>>"
+
+_DEFAULT_MAX_BLOCK_TEXT_PREVIEW_LENGTH = 100
+_DEFAULT_MAX_FORMATTER_TEXT_LENGTH = 1000
+
+
+class AsBlockStat(BaseModel):
+ """Statistics and metadata for a single content block in an AgentScope message."""
+
+ block_type: str = Field(default=...)
+ text: str = Field(default="", description="Text content of the block")
+ token_count: int = Field(default=0, description="Token count of the block, including base64 data")
+
+ # For tool_use and tool_result blocks
+ tool_name: str = Field(default="", description="Tool name for tool_use/tool_result blocks")
+ tool_input: str = Field(default="", description="Tool input arguments for tool_use blocks")
+ tool_output: str = Field(default="", description="Tool output for tool_result blocks")
+
+ # For media blocks
+ media_url: str = Field(default="", description="URL for image/audio/video blocks")
+
+ @property
+ def preview(self) -> str:
+ """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:
+ """Truncate text with ellipsis, replacing newlines with spaces."""
+ text = text.replace("\n", " ")
+ 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.
+
+ Args:
+ max_length: Maximum length of text content in the output.
+ include_thinking: Whether to include thinking block content.
+
+ Returns:
+ Formatted string representation of the block.
+ """
+ if self.block_type == "text":
+ if not self.text:
+ return ""
+ return f"[text]: {self._truncate(self.text, max_length)}"
+ if self.block_type == "thinking":
+ if not include_thinking or not self.text:
+ return ""
+ return f"[think]: {self._truncate(self.text, max_length)}"
+ if self.block_type in ("image", "audio", "video"):
+ 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"[tool_use]: {content}"
+ if self.block_type == "tool_result":
+ if not self.tool_output:
+ return ""
+ display_output = self.tool_output.split(_TRUNCATION_NOTICE_MARKER)[0]
+ content = f"{self.tool_name} output={self._truncate(display_output, max_length)}"
+ return f"[tool_result]: {content}"
+ return ""
+
+
+class AsMsgStat(BaseModel):
+ """Statistics and metadata for a complete AgentScope message."""
+
+ name: str = Field(default=...)
+ role: str = Field(default="")
+ content: list[AsBlockStat] = Field(default_factory=list)
+ timestamp: str = Field(default="")
+ metadata: dict = Field(default_factory=dict)
+
+ @property
+ def total_tokens(self) -> int:
+ """Return the total token count across all content blocks."""
+ return sum(block.token_count for block in self.content)
+
+ @property
+ def preview(self) -> str:
+ """Return a short preview of the message content."""
+ return self.format(_DEFAULT_MAX_BLOCK_TEXT_PREVIEW_LENGTH)
+
+ def format(self, max_length: int = _DEFAULT_MAX_FORMATTER_TEXT_LENGTH, include_thinking: bool = True) -> str:
+ """Format message to string representation."""
+ time_str = f"[{self.timestamp}] " if self.timestamp else ""
+ header = f"{time_str}{self.name or self.role}:"
+ blocks = [block.format(max_length, include_thinking) for block in self.content]
+ return "\n".join([header] + [b for b in blocks if b])
diff --git a/reme/core/schema/service_config.py b/reme/core/schema/service_config.py
index 5b89367a..48b97e0a 100644
--- a/reme/core/schema/service_config.py
+++ b/reme/core/schema/service_config.py
@@ -58,74 +58,68 @@ class FlowConfig(ToolCall):
cache_expire_hours: float = Field(default=0.1)
-class LLMConfig(BaseModel):
+class BasicConfig(BaseModel):
+ """Configuration for basic service settings and parameters."""
+
+ model_config = ConfigDict(extra="allow")
+
+ backend: str = Field(default="")
+
+
+class ModelConfig(BasicConfig):
+ """Configuration for model-based services with backend and model name."""
+
+ model_name: str = Field(default="")
+
+
+class LLMConfig(ModelConfig):
"""Configuration for Large Language Model backend and model identification."""
- model_config = ConfigDict(extra="allow")
- backend: str = Field(default="")
- model_name: str = Field(default="")
-
-
-class EmbeddingModelConfig(BaseModel):
+class EmbeddingModelConfig(ModelConfig):
"""Configuration for embedding model backends and identity."""
- model_config = ConfigDict(extra="allow")
- backend: str = Field(default="")
- model_name: str = Field(default="")
-
-
-class VectorStoreConfig(BaseModel):
- """Configuration for vector database storage and associated embeddings."""
-
- model_config = ConfigDict(extra="allow")
-
- backend: str = Field(default="local")
- collection_name: str = Field(default="reme")
- embedding_model: str = Field(default="default")
-
-
-class FileStoreConfig(BaseModel):
- """Configuration for file store database storage and associated embeddings."""
-
- model_config = ConfigDict(extra="allow")
-
- backend: str = Field(default="sqlite")
- store_name: str = Field(default="reme")
- embedding_model: str = Field(default="default")
-
-
-class TokenCounterConfig(BaseModel):
+class TokenCounterConfig(ModelConfig):
"""Configuration for token counting services and model mapping."""
- model_config = ConfigDict(extra="allow")
- backend: str = Field(default="base")
- model_name: str = Field(default="")
+class StoreConfig(BasicConfig):
+ """Configuration for storage services with embedding model support."""
+
+ embedding_model: str = Field(default="default")
-class FileWatcherConfig(BaseModel):
+class VectorStoreConfig(StoreConfig):
+ """Configuration for vector database storage and associated embeddings."""
+
+ collection_name: str = Field(default="reme")
+
+
+class FileStoreConfig(StoreConfig):
+ """Configuration for file store database storage and associated embeddings."""
+
+ store_name: str = Field(default="reme")
+
+
+class FileWatcherConfig(BasicConfig):
"""Configuration for file watcher service."""
- model_config = ConfigDict(extra="allow")
-
- backend: str = Field(default="")
file_store: str = Field(default="")
watch_paths: list[str] = Field(default_factory=list)
-class ServiceConfig(BaseModel):
+class ServiceConfig(BasicConfig):
"""Root configuration schema aggregating all service-level settings and components."""
- model_config = ConfigDict(extra="allow")
-
- backend: str = Field(default="")
app_name: str = Field(default=os.getenv("APP_NAME", "ReMe"))
working_dir: str = Field(default=".reme")
enable_logo: bool = Field(default=True)
language: str = Field(default="")
- thread_pool_max_workers: int = Field(default=16)
+ thread_pool_max_workers: int = Field(
+ default=16,
+ description="Number of thread pool workers. Set to -1 to disable thread pool.",
+ )
ray_max_workers: int = Field(default=-1)
log_to_console: bool = Field(default=True)
disabled_flows: list[str] = Field(default_factory=list)
@@ -137,6 +131,9 @@ class ServiceConfig(BaseModel):
cmd: CmdConfig = Field(default_factory=CmdConfig)
ops: dict[str, OpConfig] = Field(default_factory=dict)
flows: dict[str, FlowConfig] = Field(default_factory=dict)
+ as_llms: dict[str, BasicConfig] = Field(default_factory=dict)
+ as_llm_formatters: dict[str, BasicConfig] = Field(default_factory=dict)
+ as_token_counters: dict[str, BasicConfig] = Field(default_factory=dict)
llms: dict[str, LLMConfig] = Field(default_factory=dict)
embedding_models: dict[str, EmbeddingModelConfig] = Field(default_factory=dict)
vector_stores: dict[str, VectorStoreConfig] = Field(default_factory=dict)
diff --git a/reme/core/service_context.py b/reme/core/service_context.py
index ecb6c3a3..2c58fe7d 100644
--- a/reme/core/service_context.py
+++ b/reme/core/service_context.py
@@ -1,6 +1,5 @@
"""Service context."""
-import os
from concurrent.futures import ThreadPoolExecutor
from typing import TYPE_CHECKING
@@ -8,9 +7,12 @@ from loguru import logger
from .base_dict import BaseDict
from .schema import ServiceConfig
-from .utils import load_env, PydanticConfigParser
+from .utils import PydanticConfigParser
if TYPE_CHECKING:
+ from agentscope.model import ChatModelBase
+ from agentscope.formatter import FormatterBase
+ from agentscope.token import TokenCounterBase
from .llm import BaseLLM
from .embedding import BaseEmbeddingModel
from .vector_store import BaseVectorStore
@@ -26,16 +28,15 @@ class ServiceContext(BaseDict):
def __init__(
self,
*args,
- llm_api_key: str | None = None,
- llm_base_url: str | None = None,
- embedding_api_key: str | None = None,
- embedding_base_url: str | None = None,
service_config: ServiceConfig | None = None,
parser: type[PydanticConfigParser] | None = None,
working_dir: str | None = None,
config_path: str | None = None,
enable_logo: bool = True,
log_to_console: bool = True,
+ default_as_llm_config: dict | None = None,
+ default_as_llm_formatter_config: dict | None = None,
+ default_as_token_counter_config: dict | None = None,
default_llm_config: dict | None = None,
default_embedding_model_config: dict | None = None,
default_vector_store_config: dict | None = None,
@@ -46,15 +47,6 @@ class ServiceContext(BaseDict):
):
super().__init__()
- # Load environment variables
- load_env()
-
- # Update common environment variables for LLM and embedding services.
- self.update_env("LLM_API_KEY", llm_api_key)
- self.update_env("LLM_BASE_URL", llm_base_url)
- self.update_env("EMBEDDING_API_KEY", embedding_api_key)
- self.update_env("EMBEDDING_BASE_URL", embedding_base_url)
-
if service_config is None:
parser_class = parser if parser is not None else PydanticConfigParser
parser_instance = parser_class(ServiceConfig)
@@ -64,6 +56,12 @@ class ServiceContext(BaseDict):
if args:
input_args.extend(args)
+ if default_as_llm_config:
+ self._update_section_config(kwargs, "as_llms", **default_as_llm_config)
+ if default_as_llm_formatter_config:
+ self._update_section_config(kwargs, "as_llm_formatters", **default_as_llm_formatter_config)
+ if default_as_token_counter_config:
+ self._update_section_config(kwargs, "as_token_counters", **default_as_token_counter_config)
if default_llm_config:
self._update_section_config(kwargs, "llms", **default_llm_config)
if default_embedding_model_config:
@@ -90,6 +88,9 @@ class ServiceContext(BaseDict):
self.service_config: ServiceConfig = service_config
self.thread_pool: ThreadPoolExecutor | None = None
+ self.as_llms: dict[str, "ChatModelBase"] = {}
+ self.as_llm_formatters: dict[str, "FormatterBase"] = {}
+ self.as_token_counters: dict[str, "TokenCounterBase"] = {}
self.llms: dict[str, "BaseLLM"] = {}
self.embedding_models: dict[str, "BaseEmbeddingModel"] = {}
self.token_counters: dict[str, "BaseTokenCounter"] = {}
@@ -99,12 +100,6 @@ class ServiceContext(BaseDict):
self.flows: dict[str, "BaseFlow"] = {}
self.mcp_server_mapping: dict[str, dict] = {}
- @staticmethod
- def update_env(key: str, value: str | None):
- """Update environment variable if value is provided."""
- if value:
- os.environ[key] = value
-
@staticmethod
def _update_section_config(config: dict, section_name: str, **kwargs):
"""Update a specific section of the service config with new values."""
diff --git a/reme/core/utils/__init__.py b/reme/core/utils/__init__.py
index b43784c8..26ca2827 100644
--- a/reme/core/utils/__init__.py
+++ b/reme/core/utils/__init__.py
@@ -11,12 +11,14 @@ from .horse import play_horse_easter_egg
from .http_client import HttpClient
from .llm_utils import extract_content, format_messages, deduplicate_memories
from .logger_utils import init_logger
+from .std_logger import get_logger
from .logo_utils import print_logo
from .mcp_client import MCPClient
from .pydantic_config_parser import PydanticConfigParser
from .pydantic_utils import create_pydantic_model
from .singleton import singleton
from .time import timer, get_now_time
+from .hf_token_counter_utils import get_hf_token_counter
__all__ = [
"convert_dashscope_to_agentscope",
@@ -39,6 +41,7 @@ __all__ = [
"format_messages",
"deduplicate_memories",
"init_logger",
+ "get_logger",
"print_logo",
"MCPClient",
"PydanticConfigParser",
@@ -46,4 +49,5 @@ __all__ = [
"singleton",
"timer",
"get_now_time",
+ "get_hf_token_counter",
]
diff --git a/reme/core/utils/hf_token_counter_utils.py b/reme/core/utils/hf_token_counter_utils.py
new file mode 100644
index 00000000..a8ab348c
--- /dev/null
+++ b/reme/core/utils/hf_token_counter_utils.py
@@ -0,0 +1,23 @@
+"""Utility functions for working with text."""
+
+from agentscope.token import HuggingFaceTokenCounter
+
+_token_counter = None
+
+
+def get_hf_token_counter(
+ pretrained_model_name_or_path="Qwen/Qwen2.5-7B-Instruct",
+ use_mirror=True,
+ use_fast=True,
+ trust_remote_code=True,
+):
+ """Get or initialize the global token counter instance."""
+ global _token_counter
+ if _token_counter is None:
+ _token_counter = HuggingFaceTokenCounter(
+ pretrained_model_name_or_path=pretrained_model_name_or_path,
+ use_mirror=use_mirror,
+ use_fast=use_fast,
+ trust_remote_code=trust_remote_code,
+ )
+ return _token_counter
diff --git a/reme/core/utils/llm_utils.py b/reme/core/utils/llm_utils.py
index 6ec2cd23..e828c24e 100644
--- a/reme/core/utils/llm_utils.py
+++ b/reme/core/utils/llm_utils.py
@@ -3,10 +3,60 @@
import json
import re
+from agentscope.message import Msg
from loguru import logger
from ..enumeration import Role
-from ..schema import Message, Trajectory, MemoryNode
+from ..schema import Message, Trajectory, MemoryNode, ToolCall
+
+
+def convert_as_msg_to_message(msg) -> Message:
+ """Convert an agentscope Msg object to the project's Message type."""
+ role_str = getattr(msg, "role", "user")
+ role = (
+ Role(role_str.lower())
+ if isinstance(role_str, str) and role_str.lower() in [r.value for r in Role]
+ else Role.USER
+ )
+
+ content_blocks = msg.get_content_blocks()
+ content = ""
+ reasoning_content = ""
+ tool_calls = []
+ tool_call_id = ""
+
+ for block in content_blocks:
+ block_type = block["type"]
+ if block_type == "thinking":
+ reasoning_content = block["thinking"]
+ elif block_type == "tool_use":
+ try:
+ tool_calls.append(
+ ToolCall(
+ id=block["id"],
+ name=block["name"],
+ arguments=json.dumps(block["input"], ensure_ascii=False),
+ ),
+ )
+ except (json.JSONDecodeError, TypeError):
+ pass
+ elif block_type == "tool_result":
+ role = Role.TOOL
+ tool_call_id = block["id"]
+ content = block["output"][0]["text"]
+ else:
+ content = block[block_type]
+
+ return Message(
+ name=getattr(msg, "name", None),
+ role=role,
+ content=content,
+ reasoning_content=reasoning_content,
+ tool_calls=tool_calls,
+ tool_call_id=tool_call_id,
+ time_created=getattr(msg, "timestamp", "") or "",
+ metadata=getattr(msg, "metadata", {}) or {},
+ )
def format_messages(
@@ -24,6 +74,8 @@ def format_messages(
for i, message in enumerate(messages):
if isinstance(message, dict):
message = Message(**message)
+ if isinstance(message, Msg):
+ message = convert_as_msg_to_message(message)
if not enable_system and message.role is Role.SYSTEM:
continue
diff --git a/reme/core/utils/logger_utils.py b/reme/core/utils/logger_utils.py
index 22819db8..512c0a42 100644
--- a/reme/core/utils/logger_utils.py
+++ b/reme/core/utils/logger_utils.py
@@ -18,26 +18,6 @@ def init_logger(log_dir: str = "logs", level: str = "INFO", log_to_console: bool
# Remove default handler to avoid duplicate logs
logger.remove()
- # Ensure the logging directory exists
- os.makedirs(log_dir, exist_ok=True)
-
- # Generate filename based on the current timestamp
- # Use dashes instead of colons for Windows compatibility
- current_ts = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
- log_filename = f"{current_ts}.log"
- log_filepath = os.path.join(log_dir, log_filename)
-
- # Configure file-based logging with rotation and compression
- logger.add(
- log_filepath,
- level=level,
- rotation="00:00",
- retention="7 days",
- compression="zip",
- encoding="utf-8",
- format="{time:YYYY-MM-DD HH:mm:ss} | {level} | {file}:{line} | {function} | {message}",
- )
-
# Configure colorized standard output logging if enabled
if log_to_console:
logger.add(
@@ -46,3 +26,27 @@ def init_logger(log_dir: str = "logs", level: str = "INFO", log_to_console: bool
format="{time:YYYY-MM-DD HH:mm:ss} | {level} | {file}:{line} | {function} | {message}",
colorize=True,
)
+
+ # Try to configure file-based logging (skip if permission denied)
+ try:
+ # Ensure the logging directory exists
+ os.makedirs(log_dir, exist_ok=True)
+
+ # Generate filename based on the current timestamp
+ # Use dashes instead of colons for Windows compatibility
+ current_ts = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
+ log_filename = f"{current_ts}.log"
+ log_filepath = os.path.join(log_dir, log_filename)
+
+ # Configure file-based logging with rotation and compression
+ logger.add(
+ log_filepath,
+ level=level,
+ rotation="00:00",
+ retention="7 days",
+ compression="zip",
+ encoding="utf-8",
+ format="{time:YYYY-MM-DD HH:mm:ss} | {level} | {file}:{line} | {function} | {message}",
+ )
+ except Exception as e:
+ logger.error(f"Error configuring file logging: {e}")
diff --git a/reme/core/utils/std_logger.py b/reme/core/utils/std_logger.py
new file mode 100644
index 00000000..bfa6b32a
--- /dev/null
+++ b/reme/core/utils/std_logger.py
@@ -0,0 +1,122 @@
+"""Standard logging module configuration with loguru-like features."""
+
+import logging
+import os
+import sys
+from datetime import datetime
+from logging.handlers import TimedRotatingFileHandler
+
+# Store created logger instances
+_loggers: dict[str, logging.Logger] = {}
+
+
+class CustomFormatter(logging.Formatter):
+ """Custom formatter with colorized output support."""
+
+ # ANSI color codes
+ COLORS = {
+ logging.DEBUG: "\033[36m", # Cyan
+ logging.INFO: "\033[32m", # Green
+ logging.WARNING: "\033[33m", # Yellow
+ logging.ERROR: "\033[31m", # Red
+ logging.CRITICAL: "\033[35m", # Magenta
+ }
+ RESET = "\033[0m"
+
+ def __init__(self, fmt: str, colorize: bool = False):
+ super().__init__(fmt)
+ self.colorize = colorize
+
+ def format(self, record: logging.LogRecord) -> str:
+ # Add custom attribute: simplified filename and line number
+ record.file_line = f"{record.filename}:{record.lineno}"
+
+ if self.colorize:
+ color = self.COLORS.get(record.levelno, self.RESET)
+ record.levelname = f"{color}{record.levelname}{self.RESET}"
+
+ return super().format(record)
+
+
+def get_loggerv2(
+ name: str = "reme",
+ log_dir: str = "logs",
+ level: str = "INFO",
+ log_to_console: bool = True,
+ log_to_file: bool = True,
+ log_file_prefix: str = "reme",
+ rotation: str = "midnight",
+ retention_days: int = 7,
+ force_update: bool = False,
+) -> logging.Logger:
+ """Get a configured logger instance.
+
+ Args:
+ name: Logger name for distinguishing different loggers.
+ log_dir: Directory path for log files.
+ level: Logging level (DEBUG, INFO, WARNING, ERROR, CRITICAL).
+ log_to_console: Whether to output logs to console.
+ log_to_file: Whether to output logs to file.
+ log_file_prefix: Prefix for log file names (e.g., 'reme' -> 'reme_2024-01-01.log').
+ rotation: Log rotation time, defaults to midnight.
+ retention_days: Number of days to retain log files.
+ force_update: Whether to force update the logger configuration even if it already exists.
+
+ Returns:
+ Configured Logger instance.
+ """
+ # Return existing logger if already created and not force updating
+ if name in _loggers and not force_update:
+ return _loggers[name]
+
+ # Create new logger without using root logger
+ logger = logging.getLogger(name)
+ logger.setLevel(getattr(logging, level.upper(), logging.INFO))
+ logger.propagate = False # Do not propagate to root logger
+
+ # Clear existing handlers
+ logger.handlers.clear()
+
+ # Log format
+ log_format = "%(asctime)s | %(levelname)s | %(file_line)s | %(funcName)s | %(message)s"
+
+ # Configure file logging
+ if log_to_file:
+ try:
+ os.makedirs(log_dir, exist_ok=True)
+ current_ts = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
+ log_filename = f"{log_file_prefix}_{current_ts}.log"
+ log_filepath = os.path.join(log_dir, log_filename)
+
+ file_handler = TimedRotatingFileHandler(
+ log_filepath,
+ when=rotation,
+ interval=1,
+ backupCount=retention_days,
+ encoding="utf-8",
+ )
+ file_handler.setLevel(getattr(logging, level.upper(), logging.INFO))
+ file_handler.setFormatter(CustomFormatter(log_format, colorize=False))
+ file_handler.suffix = "%Y-%m-%d"
+ logger.addHandler(file_handler)
+
+ except Exception as e:
+ logger.error(f"Error configuring file logging: {e}")
+
+ # Configure console logging
+ if log_to_console:
+ console_handler = logging.StreamHandler(sys.stdout)
+ console_handler.setLevel(getattr(logging, level.upper(), logging.INFO))
+ console_handler.setFormatter(CustomFormatter(log_format, colorize=True))
+ logger.addHandler(console_handler)
+
+ # Cache logger
+ _loggers[name] = logger
+ return logger
+
+
+def get_logger():
+ """Get a configured logger instance using loguru."""
+ from loguru import logger
+
+ return logger
diff --git a/reme/core/vector_store/__init__.py b/reme/core/vector_store/__init__.py
index ccd28408..162229a2 100644
--- a/reme/core/vector_store/__init__.py
+++ b/reme/core/vector_store/__init__.py
@@ -4,6 +4,7 @@ from .base_vector_store import BaseVectorStore
from .chroma_vector_store import ChromaVectorStore
from .es_vector_store import ESVectorStore
from .local_vector_store import LocalVectorStore
+from .obvec_vector_store import ObVecVectorStore
from .pgvector_store import PGVectorStore
from .qdrant_vector_store import QdrantVectorStore
from ..registry_factory import R
@@ -13,6 +14,7 @@ __all__ = [
"ChromaVectorStore",
"ESVectorStore",
"LocalVectorStore",
+ "ObVecVectorStore",
"PGVectorStore",
"QdrantVectorStore",
]
@@ -20,6 +22,7 @@ __all__ = [
R.vector_stores.register("chroma")(ChromaVectorStore)
R.vector_stores.register("es")(ESVectorStore)
R.vector_stores.register("local")(LocalVectorStore)
+R.vector_stores.register("obvec")(ObVecVectorStore)
R.vector_stores.register("pgvector")(PGVectorStore)
R.vector_stores.register("qdrant")(QdrantVectorStore)
diff --git a/reme/core/vector_store/chroma_vector_store.py b/reme/core/vector_store/chroma_vector_store.py
index 2bd9f4ae..c6bfc1c9 100644
--- a/reme/core/vector_store/chroma_vector_store.py
+++ b/reme/core/vector_store/chroma_vector_store.py
@@ -9,12 +9,12 @@ from .base_vector_store import BaseVectorStore
from ..embedding import BaseEmbeddingModel
from ..schema import VectorNode
-_CHROMADB_IMPORT_ERROR = None
+_CHROMADB_IMPORT_ERROR: Exception | None = None
try:
import chromadb
from chromadb.config import Settings
-except ImportError as e:
+except Exception as e:
_CHROMADB_IMPORT_ERROR = e
chromadb = None
Settings = None
diff --git a/reme/core/vector_store/es_vector_store.py b/reme/core/vector_store/es_vector_store.py
index e0df83a0..332bfa68 100644
--- a/reme/core/vector_store/es_vector_store.py
+++ b/reme/core/vector_store/es_vector_store.py
@@ -13,12 +13,12 @@ from .base_vector_store import BaseVectorStore
from ..embedding import BaseEmbeddingModel
from ..schema import VectorNode
-_ELASTICSEARCH_IMPORT_ERROR = None
+_ELASTICSEARCH_IMPORT_ERROR: Exception | None = None
try:
from elasticsearch import AsyncElasticsearch
from elasticsearch.helpers import async_bulk
-except ImportError as e:
+except Exception as e:
_ELASTICSEARCH_IMPORT_ERROR = e
AsyncElasticsearch = None
async_bulk = None
diff --git a/reme/core/vector_store/obvec_vector_store.py b/reme/core/vector_store/obvec_vector_store.py
new file mode 100644
index 00000000..c8ad17c5
--- /dev/null
+++ b/reme/core/vector_store/obvec_vector_store.py
@@ -0,0 +1,453 @@
+"""OceanBase / seekdb vector store for ReMe (pyobvector)."""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+from typing import Any
+
+from loguru import logger
+from sqlalchemy import Column, JSON, String, text as sa_text
+from sqlalchemy.dialects.mysql import LONGTEXT
+
+from .base_vector_store import BaseVectorStore
+from ..embedding import BaseEmbeddingModel
+from ..schema import VectorNode
+
+_OBVECTOR_IMPORT_ERROR: Exception | None = None
+
+try:
+ from pyobvector import IndexParams, ObVecClient, VecIndexType, VECTOR
+ from pyobvector import cosine_distance, inner_product
+except Exception as e:
+ _OBVECTOR_IMPORT_ERROR = e
+ IndexParams = None # type: ignore[misc, assignment]
+ ObVecClient = None # type: ignore[misc, assignment]
+ VecIndexType = None # type: ignore[misc, assignment]
+ VECTOR = None # type: ignore[misc, assignment]
+
+_COL_SELECT = "id, content, vector, metadata"
+
+
+def _is_safe_metadata_key(key: str) -> bool:
+ return bool(key.replace("_", "").replace(".", "").isalnum())
+
+
+def _coerce_db_vector(raw: Any) -> list[float] | None:
+ if raw is None:
+ return None
+ if isinstance(raw, list):
+ return [float(x) for x in raw]
+ if isinstance(raw, str):
+ try:
+ parsed = json.loads(raw)
+ if isinstance(parsed, list):
+ return [float(x) for x in parsed]
+ except (json.JSONDecodeError, TypeError, ValueError):
+ pass
+ return None
+
+
+def _coerce_db_metadata(raw: Any) -> dict[str, Any]:
+ if raw is None:
+ return {}
+ if isinstance(raw, dict):
+ return raw
+ if isinstance(raw, str):
+ try:
+ parsed = json.loads(raw)
+ if isinstance(parsed, dict):
+ return parsed
+ except (json.JSONDecodeError, TypeError):
+ pass
+ return {}
+
+
+def _build_metadata_filter_sql(filters: dict[str, Any] | None) -> str:
+ if not filters:
+ return ""
+
+ parts: list[str] = []
+ for key, value in filters.items():
+ if not _is_safe_metadata_key(key):
+ continue
+
+ path = f"$.{key}"
+ if isinstance(value, list) and len(value) == 2:
+ lo, hi = value[0], value[1]
+ if isinstance(lo, (int, float)) and isinstance(hi, (int, float)):
+ parts.append(
+ f"(JSON_EXTRACT(metadata, '{path}') >= {lo} AND " f"JSON_EXTRACT(metadata, '{path}') <= {hi})",
+ )
+ else:
+ parts.append(
+ f"(JSON_EXTRACT(metadata, '{path}') >= '{lo}' AND " f"JSON_EXTRACT(metadata, '{path}') <= '{hi}')",
+ )
+ elif isinstance(value, (int, float)):
+ parts.append(f"JSON_EXTRACT(metadata, '{path}') = {value}")
+ else:
+ parts.append(f"JSON_EXTRACT(metadata, '{path}') = '{value}'")
+
+ return " AND ".join(parts)
+
+
+def _format_vector_sql_literal(vector: list[float]) -> str:
+ return "[" + ",".join(str(float(v)) for v in vector) + "]"
+
+
+def _normalize_embedding_for_ann(raw: Any) -> list[float]:
+ if hasattr(raw, "tolist"):
+ raw = raw.tolist()
+ return [float(x) for x in raw]
+
+
+def _vector_node_from_db_row(row: tuple[Any, ...]) -> VectorNode:
+ return VectorNode(
+ vector_id=row[0],
+ content=row[1] or "",
+ vector=_coerce_db_vector(row[2]),
+ metadata=_coerce_db_metadata(row[3]),
+ )
+
+
+def _normalize_nodes(nodes: VectorNode | list[VectorNode]) -> list[VectorNode]:
+ return [nodes] if isinstance(nodes, VectorNode) else list(nodes)
+
+
+def _sql_table(name: str) -> str:
+ return f"`{name}`"
+
+
+class ObVecVectorStore(BaseVectorStore):
+ """OceanBase or seekdb vector store for dense vectors and kNN search.
+
+ Args:
+ index_metric: ``cosine`` or ``ip`` (inner product). Invalid values raise
+ ``ValueError``; unsupported strings are not mapped to another metric.
+ """
+
+ def __init__(
+ self,
+ collection_name: str,
+ db_path: str | Path,
+ embedding_model: BaseEmbeddingModel,
+ uri: str = "127.0.0.1:2881",
+ user: str = "root",
+ password: str = "",
+ database: str = "test",
+ index_metric: str = "cosine",
+ index_ef_search: int = 100,
+ **kwargs,
+ ):
+ if _OBVECTOR_IMPORT_ERROR is not None:
+ raise ImportError(
+ "ObVecVectorStore requires pyobvector. Install with `pip install pyobvector`",
+ ) from _OBVECTOR_IMPORT_ERROR
+
+ super().__init__(
+ collection_name=collection_name,
+ db_path=db_path,
+ embedding_model=embedding_model,
+ **kwargs,
+ )
+
+ key = index_metric.strip().lower()
+ if key not in ("cosine", "ip"):
+ raise ValueError(
+ f"ObVecVectorStore index_metric must be 'cosine' or 'ip', got {index_metric!r}",
+ )
+ self.uri = uri
+ self.user = user
+ self.password = password
+ self.database = database
+ self.index_metric = key
+ self.index_ef_search = index_ef_search
+
+ self.client: ObVecClient | None = None
+ self.embedding_model_dims = embedding_model.dimensions
+
+ def _require_client(self) -> ObVecClient:
+ if self.client is None:
+ raise RuntimeError("ObVecVectorStore.start() must be called before this operation")
+ return self.client
+
+ async def list_collections(self) -> list[str]:
+ client = self._require_client()
+ result = client.perform_raw_text_sql(f"SHOW TABLES FROM {_sql_table(self.database)}")
+ rows = result.fetchall()
+ return [row[0] for row in rows if row]
+
+ def _table_columns_for_create(self, dimensions: int) -> list[Any]:
+ return [
+ Column("id", String(255), primary_key=True),
+ Column("content", LONGTEXT),
+ Column("vector", VECTOR(dimensions)),
+ Column("metadata", JSON),
+ ]
+
+ def _hnsw_index_params(self, collection_name: str) -> IndexParams:
+ metric = "cosine" if self.index_metric == "cosine" else "inner_product"
+ vidxs = IndexParams()
+ vidxs.add_index(
+ "vector",
+ VecIndexType.HNSW,
+ f"{collection_name}_vidx",
+ metric_type=metric,
+ params={"efSearch": self.index_ef_search},
+ )
+ return vidxs
+
+ async def create_collection(self, collection_name: str, **kwargs):
+ client = self._require_client()
+ dimensions = kwargs.get("dimensions", self.embedding_model_dims)
+
+ if client.check_table_exists(collection_name):
+ logger.info("Collection {} already exists", collection_name)
+ return
+
+ columns = self._table_columns_for_create(dimensions)
+ vidxs = self._hnsw_index_params(collection_name)
+
+ client.create_table_with_index_params(
+ table_name=collection_name,
+ columns=columns,
+ vidxs=vidxs,
+ )
+ logger.info("Created collection {} with dimensions={}", collection_name, dimensions)
+
+ async def delete_collection(self, collection_name: str, **kwargs):
+ client = self._require_client()
+ client.drop_table_if_exist(collection_name)
+ logger.info("Deleted collection {}", collection_name)
+
+ async def copy_collection(self, collection_name: str, **kwargs):
+ client = self._require_client()
+
+ if not client.check_table_exists(self.collection_name):
+ raise ValueError(f"Source collection {self.collection_name} does not exist")
+
+ await self.create_collection(collection_name)
+
+ try:
+ source_data = await self.list(limit=None)
+ if source_data:
+ await self.insert(source_data, collection_name=collection_name)
+ logger.info("Copied collection {} to {}", self.collection_name, collection_name)
+ except Exception:
+ try:
+ client.drop_table_if_exist(collection_name)
+ except Exception as cleanup_err:
+ logger.warning("Cleanup after failed copy failed: {}", cleanup_err)
+ raise
+
+ async def insert(self, nodes: VectorNode | list[VectorNode], **kwargs):
+ nodes = _normalize_nodes(nodes)
+ if not nodes:
+ return
+
+ client = self._require_client()
+
+ need_emb = [n for n in nodes if n.vector is None]
+ if need_emb:
+ filled = await self.get_node_embeddings(need_emb)
+ by_id = {n.vector_id: n for n in filled}
+ nodes_to_insert = [by_id.get(n.vector_id, n) for n in nodes]
+ else:
+ nodes_to_insert = nodes
+
+ data = [
+ {
+ "id": node.vector_id,
+ "content": node.content,
+ "vector": node.vector if node.vector is not None else [],
+ "metadata": node.metadata if node.metadata else {},
+ }
+ for node in nodes_to_insert
+ ]
+ target = kwargs.get("collection_name", self.collection_name)
+ client.insert(table_name=target, data=data)
+ logger.info("Inserted {} documents into {}", len(nodes_to_insert), target)
+
+ async def search(
+ self,
+ query: str,
+ limit: int = 5,
+ filters: dict | None = None,
+ **kwargs,
+ ) -> list[VectorNode]:
+ client = self._require_client()
+ raw_vec = await self.get_embedding(query)
+ query_vector = _normalize_embedding_for_ann(raw_vec)
+ dist_fn = cosine_distance if self.index_metric == "cosine" else inner_product
+
+ filter_sql = _build_metadata_filter_sql(filters)
+ where_parts = [sa_text(filter_sql)] if filter_sql else None
+
+ results = client.ann_search(
+ table_name=self.collection_name,
+ vec_data=query_vector,
+ vec_column_name="vector",
+ distance_func=dist_fn,
+ with_dist=True,
+ topk=limit,
+ output_column_names=["id", "content", "metadata"],
+ where_clause=where_parts,
+ )
+
+ score_threshold = kwargs.get("score_threshold")
+ out: list[VectorNode] = []
+ for row in results:
+ if len(row) < 4:
+ raise RuntimeError(
+ "ann_search row must have id, content, metadata, distance " f"(got {len(row)} columns)",
+ )
+ vid, content, metadata_raw, distance = row[0], row[1], row[2], row[3]
+ dist_f = float(distance)
+ if self.index_metric == "cosine":
+ score = max(0.0, 1.0 - dist_f / 2.0)
+ else:
+ score = max(0.0, dist_f)
+ if score_threshold is not None and score < score_threshold:
+ continue
+ meta = _coerce_db_metadata(metadata_raw) if metadata_raw is not None else {}
+ meta["score"] = score
+ meta["_distance"] = dist_f
+ out.append(
+ VectorNode(
+ vector_id=vid,
+ content=content or "",
+ vector=None,
+ metadata=meta,
+ ),
+ )
+ return out
+
+ async def delete(self, vector_ids: str | list[str], **kwargs):
+ if isinstance(vector_ids, str):
+ vector_ids = [vector_ids]
+ if not vector_ids:
+ return
+ client = self._require_client()
+ client.delete(self.collection_name, ids=vector_ids)
+ logger.info("Deleted {} documents from {}", len(vector_ids), self.collection_name)
+
+ async def delete_all(self, **kwargs):
+ client = self._require_client()
+ client.delete(self.collection_name)
+ logger.info("Deleted all documents from {}", self.collection_name)
+
+ async def update(self, nodes: VectorNode | list[VectorNode], **kwargs):
+ nodes = _normalize_nodes(nodes)
+ if not nodes:
+ return
+
+ client = self._require_client()
+ need_emb = [n for n in nodes if n.vector is None and bool(n.content)]
+ if need_emb:
+ filled = await self.get_node_embeddings(need_emb)
+ by_id = {n.vector_id: n for n in filled}
+ nodes_to_update = [
+ by_id.get(n.vector_id, n) if (n.vector is None and bool(n.content)) else n for n in nodes
+ ]
+ else:
+ nodes_to_update = nodes
+
+ for node in nodes_to_update:
+ updates: list[str] = []
+ params: dict[str, Any] = {}
+
+ if node.content is not None:
+ updates.append("content = :content")
+ params["content"] = node.content
+
+ if node.vector is not None:
+ updates.append("vector = :vector")
+ params["vector"] = _format_vector_sql_literal(node.vector)
+
+ if node.metadata is not None:
+ updates.append("metadata = :metadata")
+ params["metadata"] = json.dumps(node.metadata)
+
+ if not updates:
+ continue
+
+ params["vid"] = node.vector_id
+ update_sql = f"UPDATE {_sql_table(self.collection_name)} SET {', '.join(updates)} WHERE id = :vid"
+ with client.engine.connect() as conn:
+ with conn.begin():
+ conn.execute(sa_text(update_sql), params)
+
+ logger.info("Updated {} documents in {}", len(nodes_to_update), self.collection_name)
+
+ async def get(self, vector_ids: str | list[str]) -> VectorNode | list[VectorNode] | None:
+ single = isinstance(vector_ids, str)
+ if single:
+ vector_ids = [vector_ids]
+ if not vector_ids:
+ return [] if not single else None
+
+ client = self._require_client()
+ ids_str = "', '".join(vector_ids)
+ select_sql = f"SELECT {_COL_SELECT} FROM {_sql_table(self.collection_name)} " f"WHERE id IN ('{ids_str}')"
+ result = client.perform_raw_text_sql(select_sql)
+ rows = result.fetchall()
+ parsed = [_vector_node_from_db_row(row) for row in rows if row]
+ if single:
+ return parsed[0] if parsed else None
+ return parsed
+
+ async def list(
+ self,
+ filters: dict | None = None,
+ limit: int | None = None,
+ sort_key: str | None = None,
+ reverse: bool = False,
+ ) -> list[VectorNode]:
+ client = self._require_client()
+ select_sql = f"SELECT {_COL_SELECT} FROM {_sql_table(self.collection_name)}"
+ where_clause = _build_metadata_filter_sql(filters)
+ if where_clause:
+ select_sql += f" WHERE {where_clause}"
+ if sort_key and _is_safe_metadata_key(sort_key):
+ order = "DESC" if reverse else "ASC"
+ select_sql += f" ORDER BY JSON_EXTRACT(metadata, '$.{sort_key}') {order}"
+ if limit is not None:
+ select_sql += f" LIMIT {limit}"
+ result = client.perform_raw_text_sql(select_sql)
+ rows = result.fetchall()
+ return [_vector_node_from_db_row(row) for row in rows if row]
+
+ async def collection_info(self) -> dict[str, Any]:
+ """Return collection name and row count."""
+ client = self._require_client()
+ count_sql = f"SELECT COUNT(*) FROM {_sql_table(self.collection_name)}"
+ result = client.perform_raw_text_sql(count_sql)
+ row = result.fetchone()
+ count = row[0] if row else 0
+ return {"name": self.collection_name, "count": count}
+
+ async def reset(self):
+ """Drop and recreate the current collection table."""
+ logger.warning("Resetting collection {}...", self.collection_name)
+ await self.delete_collection(self.collection_name)
+ await self.create_collection(self.collection_name)
+
+ async def reset_collection(self, collection_name: str):
+ self.collection_name = collection_name
+ await self.create_collection(collection_name)
+ logger.info("Collection reset to {}", collection_name)
+
+ async def start(self) -> None:
+ self.client = ObVecClient(
+ uri=self.uri,
+ user=self.user,
+ password=self.password,
+ db_name=self.database,
+ )
+
+ await super().start()
+ logger.info("seekdb / OceanBase vector table {} ready", self.collection_name)
+
+ async def close(self):
+ self.client = None
+ logger.info("ObVec client connection closed")
diff --git a/reme/core/vector_store/pgvector_store.py b/reme/core/vector_store/pgvector_store.py
index bc665cd9..03805076 100644
--- a/reme/core/vector_store/pgvector_store.py
+++ b/reme/core/vector_store/pgvector_store.py
@@ -11,12 +11,12 @@ from .base_vector_store import BaseVectorStore
from ..embedding import BaseEmbeddingModel
from ..schema import VectorNode
-_ASYNCPG_IMPORT_ERROR = None
+_ASYNCPG_IMPORT_ERROR: Exception | None = None
try:
import asyncpg
from asyncpg import Pool
-except ImportError as e:
+except Exception as e:
_ASYNCPG_IMPORT_ERROR = e
asyncpg = None
Pool = None
diff --git a/reme/core/vector_store/qdrant_vector_store.py b/reme/core/vector_store/qdrant_vector_store.py
index b981a23f..3e69d667 100644
--- a/reme/core/vector_store/qdrant_vector_store.py
+++ b/reme/core/vector_store/qdrant_vector_store.py
@@ -9,7 +9,7 @@ from .base_vector_store import BaseVectorStore
from ..embedding import BaseEmbeddingModel
from ..schema import VectorNode
-_QDRANT_IMPORT_ERROR = None
+_QDRANT_IMPORT_ERROR: Exception | None = None
try:
from qdrant_client import AsyncQdrantClient
@@ -23,7 +23,7 @@ try:
Range,
VectorParams,
)
-except ImportError as e:
+except Exception as e:
_QDRANT_IMPORT_ERROR = e
AsyncQdrantClient = None
Distance = None
diff --git a/reme/extension/procedural_memory/summary/comparative_extraction.py b/reme/extension/procedural_memory/summary/comparative_extraction.py
index 9ca78784..4e728de3 100644
--- a/reme/extension/procedural_memory/summary/comparative_extraction.py
+++ b/reme/extension/procedural_memory/summary/comparative_extraction.py
@@ -49,8 +49,8 @@ class ComparativeExtraction(BaseOp):
comparative_task_memories.extend(soft_task_memories)
# Hard comparison: success vs failure (if similarity search is enabled)
- if self.context.get("enable_similarity_comparison", True) and success_trajectories and failure_trajectories:
- similar_pairs = self._find_similar_step_sequences(success_trajectories, failure_trajectories)
+ if self.context.get("enable_similarity_comparison", False) and success_trajectories and failure_trajectories:
+ similar_pairs = await self._find_similar_step_sequences(success_trajectories, failure_trajectories)
logger.info(f"Found {len(similar_pairs)} similar pairs for hard comparison")
for success_steps, failure_steps, similarity_score in similar_pairs:
@@ -182,7 +182,7 @@ class ComparativeExtraction(BaseOp):
else:
return trajectory.messages
- def _find_similar_step_sequences(
+ async def _find_similar_step_sequences(
self,
success_trajectories: List[Trajectory],
failure_trajectories: List[Trajectory],
@@ -227,8 +227,8 @@ class ComparativeExtraction(BaseOp):
"embedding_model",
)
):
- success_embeddings = self.vector_store.embedding_model.get_embeddings(success_texts)
- failure_embeddings = self.vector_store.embedding_model.get_embeddings(failure_texts)
+ success_embeddings = await self.vector_store.get_embeddings(success_texts)
+ failure_embeddings = await self.vector_store.get_embeddings(failure_texts)
# Calculate similarity and find most similar pairs
similarity_threshold = self.context.get("similarity_threshold", 0.5)
diff --git a/reme/memory/__init__.py b/reme/memory/__init__.py
index ac825bc6..e209a3fe 100644
--- a/reme/memory/__init__.py
+++ b/reme/memory/__init__.py
@@ -1,11 +1,11 @@
"""memory"""
from . import file_based
-from . import tools
+from . import vector_tools
from . import vector_based
__all__ = [
"file_based",
- "tools",
+ "vector_tools",
"vector_based",
]
diff --git a/reme/memory/file_based/__init__.py b/reme/memory/file_based/__init__.py
index 29f33749..e4c392e5 100644
--- a/reme/memory/file_based/__init__.py
+++ b/reme/memory/file_based/__init__.py
@@ -1,33 +1,13 @@
-"""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:
- - MemoryFormatter: Converts message lists to formatted strings with token limiting
- - ReMeInMemoryMemory: Extended InMemoryMemory with bugfixes and summary support
- - 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
- - FileIO: File I/O operations with configurable working directory
-"""
+"""File-based Memory Module."""
+from . import components
+from . import tools
from . import utils
-from .compactor import Compactor
-from .file_io import FileIO
-from .memory_formatter import MemoryFormatter
-from .reme_chat_formatter import ReMeChatFormatter
from .reme_in_memory_memory import ReMeInMemoryMemory
-from .summarizer import Summarizer
-from .tool_result_compactor import ToolResultCompactor
__all__ = [
- "MemoryFormatter",
- "ReMeInMemoryMemory",
- "Summarizer",
- "Compactor",
- "ToolResultCompactor",
- "FileIO",
+ "tools",
"utils",
- "ReMeChatFormatter",
+ "components",
+ "ReMeInMemoryMemory",
]
diff --git a/reme/memory/file_based/compactor.py b/reme/memory/file_based/compactor.py
deleted file mode 100644
index d8186d86..00000000
--- a/reme/memory/file_based/compactor.py
+++ /dev/null
@@ -1,84 +0,0 @@
-"""Compactor module for memory compaction operations."""
-
-import logging
-
-from agentscope.agent import ReActAgent
-from agentscope.formatter import FormatterBase
-from agentscope.message import Msg
-from agentscope.model import ChatModelBase
-from agentscope.token import HuggingFaceTokenCounter
-
-from .memory_formatter import MemoryFormatter
-from ...core.op import BaseOp
-
-logger = logging.getLogger(__name__)
-
-
-class Compactor(BaseOp):
- """Compactor class for compacting memory messages."""
-
- def __init__(
- self,
- memory_compact_threshold: int,
- chat_model: ChatModelBase,
- formatter: FormatterBase,
- token_counter: HuggingFaceTokenCounter,
- **kwargs,
- ):
- super().__init__(**kwargs)
- self.memory_compact_threshold: int = memory_compact_threshold
-
- self.chat_model: ChatModelBase = chat_model
- self.formatter: FormatterBase = formatter
- self.as_token_counter: HuggingFaceTokenCounter = token_counter
-
- async def execute(self):
- messages: list[Msg] = self.context.get("messages", [])
- previous_summary: str = self.context.get("previous_summary", "")
-
- if not messages:
- return ""
-
- formatter = MemoryFormatter(
- token_counter=self.as_token_counter,
- memory_compact_threshold=self.memory_compact_threshold,
- )
- history_formatted_str: str = formatter.format(messages)
-
- if not history_formatted_str:
- logger.warning(f"No history to compact. messages={messages}")
- return ""
-
- agent = ReActAgent(
- name="reme_compactor",
- model=self.chat_model,
- sys_prompt=self.get_prompt("system_prompt"),
- formatter=self.formatter,
- )
-
- if previous_summary:
- prefix: str = self.get_prompt("update_user_message_prefix")
- suffix: str = self.get_prompt("update_user_message_suffix")
- user_message: str = (
- f"\n{history_formatted_str}\n \n\n"
- f"{prefix}\n\n"
- f"\n{previous_summary}\n \n\n"
- f"{suffix}"
- )
- else:
- user_message: str = f"\n{history_formatted_str}\n \n\n" + self.get_prompt(
- "initial_user_message",
- )
- logger.info(f"Compactor sys_prompt={agent.sys_prompt} user_message={user_message}")
-
- compact_msg: Msg = await agent.reply(
- Msg(
- name="reme",
- role="user",
- content=user_message,
- ),
- )
-
- history_compact: str = compact_msg.get_text_content()
- logger.info(f"Compactor Result:\n{history_compact}")
- return history_compact
diff --git a/reme/memory/file_based/components/__init__.py b/reme/memory/file_based/components/__init__.py
new file mode 100644
index 00000000..86574790
--- /dev/null
+++ b/reme/memory/file_based/components/__init__.py
@@ -0,0 +1,15 @@
+"""components"""
+
+from .compactor import Compactor
+from .context_checker import ContextChecker
+from .summarizer import Summarizer
+from .tool_result_compactor import ToolResultCompactor
+from .cli import CliAgent
+
+__all__ = [
+ "Compactor",
+ "Summarizer",
+ "ContextChecker",
+ "ToolResultCompactor",
+ "CliAgent",
+]
diff --git a/reme/memory/file_based/components/cli.py b/reme/memory/file_based/components/cli.py
new file mode 100644
index 00000000..138a3aad
--- /dev/null
+++ b/reme/memory/file_based/components/cli.py
@@ -0,0 +1,326 @@
+"""CLI component for interactive chat using agentscope-based memory tools."""
+
+import asyncio
+from datetime import datetime
+from pathlib import Path
+import zoneinfo
+
+from agentscope.agent import ReActAgent
+from agentscope.message import Msg, TextBlock
+from agentscope.pipeline import stream_printing_messages
+from agentscope.tool import Toolkit, ToolResponse
+
+from .compactor import Compactor
+from .context_checker import ContextChecker
+from .summarizer import Summarizer
+from ..tools import FileIO, MemorySearch
+from ....core.op import BaseOp
+from ....core.utils import format_messages
+from ....core.utils import get_logger
+
+logger = get_logger()
+# name + desc + "{working_dir}/skills/{skill_name}/SKILL.md"
+
+_DEFAULT_AGENT_SKILL_INSTRUCTION = (
+ "# Agent Skills\n"
+ "The agent skills are a collection of folds of instructions, scripts, "
+ "and resources that you can load dynamically to improve performance "
+ "on specialized tasks. Each agent skill has a `SKILL.md` file in its "
+ "folder that describes how to use the skill. If you want to use a "
+ "skill, you MUST read its `SKILL.md` file carefully."
+)
+
+_DEFAULT_AGENT_SKILL_TEMPLATE = """## {name}
+{description}
+Check "{dir}/SKILL.md" for how to use this skill"""
+
+
+class CliAgent(BaseOp):
+ """CLI agent for interactive chat with memory management."""
+
+ def __init__(
+ self,
+ working_dir: str,
+ vector_weight: float = 0.7,
+ candidate_multiplier: float = 3.0,
+ context_window_tokens: int = 128000,
+ reserve_tokens: int = 36000,
+ keep_recent_tokens: int = 20000,
+ language: str = "zh",
+ timezone: str | None = None,
+ **kwargs,
+ ):
+ super().__init__(**kwargs)
+ self.working_dir: str = working_dir
+ Path(self.working_dir).mkdir(parents=True, exist_ok=True)
+ self.vector_weight: float = vector_weight
+ self.candidate_multiplier: float = candidate_multiplier
+ self.context_window_tokens: int = context_window_tokens
+ self.reserve_tokens: int = reserve_tokens
+ self.keep_recent_tokens: int = keep_recent_tokens
+ self.language: str = language
+ self.timezone: str | None = timezone
+
+ # Initialize message history
+ self.messages: list[Msg] = []
+ self.previous_summary: str = ""
+ self.summary_tasks: list[asyncio.Task] = []
+
+ def add_summary_task(self, messages: list[Msg]):
+ """Add summary task to queue."""
+ remaining_tasks = []
+ for task in self.summary_tasks:
+ if task.done():
+ exc = task.exception()
+ if exc is not None:
+ logger.exception(f"Summary task failed: {exc}")
+ else:
+ result = task.result()
+ logger.info(f"Summary task completed: {result}")
+ else:
+ remaining_tasks.append(task)
+ self.summary_tasks = remaining_tasks
+
+ # Create a toolkit for the summarizer
+ toolkit = self._create_file_toolkit()
+
+ # Create summarizer instance
+ memory_path = Path(self.working_dir) / "memory"
+ summarizer = Summarizer(
+ working_dir=self.working_dir,
+ memory_dir=str(memory_path),
+ memory_compact_threshold=int(self.context_window_tokens * 0.7),
+ token_counter=self.as_token_counter,
+ toolkit=toolkit,
+ as_llm=self.as_llm,
+ as_llm_formatter=self.as_llm_formatter,
+ language=self.language if self.language == "zh" else "",
+ console_enabled=False, # We disable the terminal printing to avoid messy outputs
+ timezone=self.timezone,
+ )
+
+ # Create summary task
+ summary_task = asyncio.create_task(
+ summarizer.call(
+ messages=messages,
+ service_context=self.service_context,
+ ),
+ )
+ self.summary_tasks.append(summary_task)
+
+ def _create_file_toolkit(self):
+ """Create a toolkit with file operations."""
+
+ toolkit = Toolkit()
+ file_io = FileIO(working_dir=self.working_dir)
+ toolkit.register_tool_function(file_io.read_file)
+ toolkit.register_tool_function(file_io.write_file)
+ toolkit.register_tool_function(file_io.edit_file)
+
+ return toolkit
+
+ async def new(self) -> str:
+ """Reset conversation history using summary."""
+ if not self.messages:
+ self.messages.clear()
+ self.previous_summary = ""
+ return "No history to reset."
+
+ self.add_summary_task(self.messages)
+
+ self.messages.clear()
+ self.previous_summary = ""
+ return "History saved to memory files and reset."
+
+ async def context_check(self) -> dict:
+ """Check if messages exceed token limits."""
+ # Create context checker
+ checker = ContextChecker(
+ memory_compact_threshold=self.context_window_tokens - self.reserve_tokens,
+ memory_compact_reserve=self.keep_recent_tokens,
+ token_counter=self.as_token_counter,
+ )
+
+ return await checker.call(
+ messages=self.messages,
+ service_context=self.service_context,
+ )
+
+ async def compact(self, force_compact: bool = False) -> str:
+ """Compact history then reset."""
+ if not self.messages:
+ return "No history to compact."
+
+ # Check and find cut point
+ messages_to_compact, messages_to_keep, _ = await self.context_check()
+ tokens_before = len(self.messages)
+
+ if force_compact:
+ messages_to_summarize = self.messages
+ left_messages = []
+ elif not messages_to_compact:
+ return "History is within token limits, no compaction needed."
+ else:
+ messages_to_summarize = messages_to_compact
+ left_messages = messages_to_keep
+
+ # Create compactor
+ compactor = Compactor(
+ memory_compact_threshold=self.context_window_tokens - self.reserve_tokens,
+ token_counter=self.as_token_counter,
+ as_llm=self.as_llm,
+ as_llm_formatter=self.as_llm_formatter,
+ language=self.language if self.language == "zh" else "",
+ console_enabled=False, # We disable the terminal printing to avoid messy outputs
+ timezone=self.timezone,
+ )
+
+ summary_content = await compactor.call(
+ messages=messages_to_summarize,
+ previous_summary=self.previous_summary,
+ service_context=self.service_context,
+ )
+
+ self.add_summary_task(messages=messages_to_summarize)
+
+ # Assemble final messages
+ self.messages = left_messages
+ self.previous_summary = summary_content
+
+ return f"History compacted from {tokens_before} messages."
+
+ def format_history(self) -> str:
+ """Format history messages."""
+ return format_messages(
+ messages=self.messages,
+ add_index=False,
+ add_reasoning=False,
+ strip_markdown_headers=False,
+ )
+
+ async def _build_messages(self, query: str) -> list[Msg]:
+ """Build system prompt message."""
+ tz = zoneinfo.ZoneInfo(self.timezone) if self.timezone else None
+ current_time = datetime.now(tz).strftime("%Y-%m-%d %H:%M:%S %A")
+
+ # Create system prompt
+ system_prompt = self.prompt_format(
+ "system_prompt",
+ workspace_dir=self.working_dir,
+ current_time=current_time,
+ has_previous_summary=bool(self.previous_summary),
+ previous_summary=self.previous_summary or "",
+ )
+
+ logger.info(f"[{self.__class__.__name__}] system_prompt: {system_prompt}")
+
+ # Build message list
+ messages = [Msg(name="system", role="system", content=system_prompt)]
+ messages.extend(self.messages)
+ messages.append(Msg(name="user", role="user", content=query))
+
+ return messages
+
+ async def memory_search(self, query: str, max_results: int = 5, min_score: float = 0.1) -> ToolResponse:
+ """
+ Mandatory recall step: semantically search MEMORY.md + memory/*.md (and optional session transcripts)
+ before answering questions about prior work, decisions, dates, people, preferences, or todos;
+ returns top snippets with path + lines.
+
+ Args:
+ query: The semantic search query to find relevant memory snippets
+ max_results: Maximum number of search results to return (optional), default is 5
+ min_score: Minimum similarity score threshold for results (optional), default is 0.1
+
+ Returns:
+ Search results as formatted string
+ """
+ search_tool = MemorySearch(
+ vector_weight=self.vector_weight,
+ candidate_multiplier=self.candidate_multiplier,
+ )
+ search_result = await search_tool.call(
+ query=query,
+ max_results=max_results,
+ min_score=min_score,
+ service_context=self.service_context,
+ )
+ return ToolResponse(
+ content=[
+ TextBlock(
+ type="text",
+ text=search_result,
+ ),
+ ],
+ )
+
+ async def execute(self):
+ """Execute the agent."""
+ _ = await self.compact(force_compact=False)
+
+ # Build messages for the agent
+ query = self.context.query
+ messages = await self._build_messages(query)
+
+ toolkit = self._create_file_toolkit()
+ # Register memory search tool
+ toolkit.register_tool_function(self.memory_search)
+
+ # Create the ReAct agent
+ agent = ReActAgent(
+ name="reme_cli_agent",
+ model=self.as_llm,
+ sys_prompt=messages[0].content, # System prompt
+ formatter=self.as_llm_formatter,
+ toolkit=toolkit,
+ )
+
+ # We disable the terminal printing to avoid messy outputs
+ agent.set_console_output_enabled(False)
+
+ self.messages = messages[1:] # remove the first SYSTEM message
+ agent.memory.content.clear()
+
+ # Stream processing state
+ in_thinking = False
+ in_answer = False
+
+ # obtain the printing messages from the agent in a streaming way
+ last_text_content = ""
+ last_think_content = ""
+ async for msg, last in stream_printing_messages(
+ agents=[agent],
+ coroutine_task=agent(self.messages),
+ ):
+ # print(msg, last)
+ content_blocks = msg.get_content_blocks()
+ for block in content_blocks:
+ if block["type"] == "thinking":
+ if not in_thinking and len(block["thinking"]) > len(last_think_content):
+ print("\033[90m\nThinking: ", end="", flush=True)
+ in_thinking = True
+ print(block["thinking"][len(last_think_content) :], end="", flush=True)
+ last_think_content = block["thinking"]
+ elif block["type"] == "text":
+ if in_thinking:
+ print("\033[0m") # reset color after thinking
+ in_thinking = False
+ if not in_answer:
+ print("\nRemy: ", end="", flush=True)
+ in_answer = True
+ print(block["text"][len(last_text_content) :], end="", flush=True)
+ last_text_content = block["text"]
+ elif block["type"] == "tool_use":
+ if in_thinking:
+ print("\033[0m") # reset color after thinking
+ in_thinking = False
+ if last:
+ print(f"\033[36m -> Executing Tool: name={block['name']}, input={block['input']}\033[0m")
+ elif block["type"] == "tool_result":
+ if last:
+ last_think_content = "" # reset for further thinking
+ print(f"\033[36m -> Tool Result for `{block['name']}`: {block['output'][0]['text']}\033[0m")
+ else:
+ print(f"Unknown block type: {block['type']}")
+ if last:
+ self.messages.append(msg)
diff --git a/reme/memory/file_based/components/cli.yaml b/reme/memory/file_based/components/cli.yaml
new file mode 100644
index 00000000..29dbfb25
--- /dev/null
+++ b/reme/memory/file_based/components/cli.yaml
@@ -0,0 +1,95 @@
+system_prompt: |
+ You are a personal assistant named Remy.
+
+ ## Working Directory
+ {workspace_dir}
+
+ ## Current Time
+ {current_time}
+
+ ## Tools
+ - `read` Read file contents
+ - `write` Write file contents
+ - `edit` Edit file contents
+ - `memory_search` Search your memories via vector store
+
+ **Don't give up easily** — if a tool doesn't return what you expect, try a different angle or approach.
+
+ ## Memory System
+ You are spun up fresh at the start of every session. These files are how you maintain continuity:
+ - **Long-term memory:** `MEMORY.md` — when you pick up a lesson or catch yourself making a mistake, feel free to **read, edit, and update** MEMORY.md
+ - **Daily notes:** `memory/YYYY-MM-DD.md` — jot things down often. When the user says "remember this," or whenever you feel something is worth noting or adding as a todo, feel free to **read, edit, and update** `memory/YYYY-MM-DD.md`
+ - **Read before you write** — always use `read` to check existing content before updating with `edit` or `write`
+
+ ### Memory Retrieval
+ 1. Start with `memory_search` — if nothing comes up, try rephrasing from a different angle
+ 2. To review a specific daily note (`memory/YYYY-MM-DD.md`), use `read`
+
+ ## Response Style 😊
+ - Keep it short and natural — talk like a friend, not a manual
+ - Use emoji sparingly for warmth — no more than 1–2 per reply
+ - For quick confirmations (yes/no, got it), an emoji is fine (👍, ✅, 🤔)
+ - When explaining or performing actions, lead with substance over flair
+
+ ## 🛡️ Safety
+ - Never run destructive commands without asking first
+ - Prefer `trash` over `rm` — recoverable beats permanent
+ - When in doubt, ask
+
+ ## Continuous Improvement
+ This is just a starting point. When you spot useful patterns or lessons during your conversations, note them in `MEMORY.md`. Do not modify system-level config files.
+
+ [has_previous_summary]## Previous Conversation Summary
+ [has_previous_summary]
+ [has_previous_summary]{previous_summary}
+ [has_previous_summary]
+ [has_previous_summary]
+ [has_previous_summary]The above is a summary of our earlier conversation. Use it as context to maintain continuity.
+
+system_prompt_zh: |
+ 你是一个名叫 Remy 的个人助手。
+
+ ## 工作目录
+ {workspace_dir}
+
+ ## 当前时间
+ {current_time}
+
+ ## 工具集合
+ - `read` 读取文件内容
+ - `write` 写入文件内容
+ - `edit` 编辑文件内容
+ - `memory_search` 通过向量库检索你的记忆
+
+ **不要轻易放弃**：如果工具执行结果不符合预期，可以从不同的维度进行不同的尝试。
+
+ ## 记忆系统
+ 每次新会话开始时，你都会被重新唤醒。以下文件是你保持连续性的关键：
+ - **长期记忆：** `MEMORY.md`：当你学到经验，或者当你犯了错误，可以**自由地阅读、编辑和更新** MEMORY.md
+ - **每日笔记：** `memory/YYYY-MM-DD.md`：要勤记笔记，当用户说"记住这个"，或者你觉得要记笔记/todo，可以**自由地阅读、编辑和更新** `memory/YYYY-MM-DD.md`
+ - **写入前先读取** — 务必先用 `read` 读取已有内容，再用 `edit` 或 `write` 更新文件
+
+ ### 记忆检索策略
+ 1. 优先使用`memory_search`检索记忆，没有搜索结果可以从不同角度多次尝试
+ 2. 如果你需要阅读每日笔记 `memory/YYYY-MM-DD.md`，可以使用`read`
+
+ ## 回应风格 😊
+ - 保持简洁自然，像朋友对话一样
+ - 适当使用 emoji 增加亲和力，但不要过度 — 每条回复最多 1-2 个
+ - 简单确认类场景（是/否、收到）可以用 emoji 快速回应（👍, ✅, 🤔）
+ - 涉及操作或解释时，优先给出有实质内容的文字回复
+
+ ## 🛡️ 安全规则
+ - 不要在没有询问的情况下运行破坏性命令
+ - 优先使用 `trash` 而不是 `rm`（可恢复比永久删除更好）
+ - 有疑问时，先询问
+
+ ## 持续改进
+ 这只是一个起点。当你在与用户的交互中发现有用的经验或模式，可以记录到 `MEMORY.md` 中。但不要修改系统级配置文件。
+
+ [has_previous_summary]## 之前的对话摘要
+ [has_previous_summary]
+ [has_previous_summary]{previous_summary}
+ [has_previous_summary]
+ [has_previous_summary]
+ [has_previous_summary]以上是我们之前对话的摘要。使用它作为上下文以保持连续性。
diff --git a/reme/memory/file_based/components/compactor.py b/reme/memory/file_based/components/compactor.py
new file mode 100644
index 00000000..e8fe9383
--- /dev/null
+++ b/reme/memory/file_based/components/compactor.py
@@ -0,0 +1,115 @@
+"""Compactor module for memory compaction operations."""
+
+from agentscope.agent import ReActAgent
+from agentscope.message import Msg
+
+from ..utils import AsMsgHandler
+from ....core.op import BaseOp
+from ....core.utils import get_logger
+
+logger = get_logger()
+
+
+def _is_valid_summary(content: str) -> bool:
+ """Check if the summary content is valid.
+
+ Args:
+ content: The summary content to validate.
+
+ Returns:
+ True if valid, False otherwise.
+ """
+ if not content or not content.strip():
+ return False
+ if "##" not in content:
+ return False
+ return True
+
+
+class Compactor(BaseOp):
+ """Compactor class for compacting memory messages."""
+
+ def __init__(
+ self,
+ memory_compact_threshold: int,
+ console_enabled: bool = False,
+ return_dict: bool = False,
+ add_thinking_block: bool = True,
+ extra_instruction: str = "",
+ **kwargs,
+ ):
+ super().__init__(**kwargs)
+ self.memory_compact_threshold: int = memory_compact_threshold
+ self.console_enabled: bool = console_enabled
+ self.return_dict: bool = return_dict
+ self.add_thinking_block: bool = add_thinking_block
+ self.extra_instruction: str = extra_instruction
+
+ # pylint: disable=too-many-return-statements
+ async def execute(self):
+ messages: list[Msg] = self.context.get("messages", [])
+ previous_summary: str = self.context.get("previous_summary", "")
+
+ if not messages:
+ if self.return_dict:
+ return {"user_message": "", "history_compact": "", "is_valid": False}
+ return ""
+
+ msg_handler = AsMsgHandler(self.as_token_counter)
+ before_token_count = await msg_handler.count_msgs_token(messages)
+ history_formatted_str: str = await msg_handler.format_msgs_to_str(
+ messages=messages,
+ memory_compact_threshold=self.memory_compact_threshold,
+ include_thinking=self.add_thinking_block,
+ )
+ after_token_count = await 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}")
+ if self.return_dict:
+ return {"user_message": "", "history_compact": "", "is_valid": False}
+ return ""
+
+ agent = ReActAgent(
+ name="reme_compactor",
+ model=self.as_llm,
+ sys_prompt=self.get_prompt("system_prompt"),
+ formatter=self.as_llm_formatter,
+ )
+ agent.set_console_output_enabled(self.console_enabled)
+
+ if previous_summary:
+ user_message: str = (
+ f"# conversation\n{history_formatted_str}\n\n"
+ f"# previous-summary\n{previous_summary}\n\n" + self.get_prompt("update_user_message")
+ )
+ else:
+ user_message: str = f"# conversation\n{history_formatted_str}\n\n" + self.get_prompt("initial_user_message")
+
+ if self.extra_instruction:
+ user_message += f"\n\n# extra-instruction\n{self.extra_instruction}"
+ logger.info(f"Compactor sys_prompt={agent.sys_prompt} user_message={user_message}")
+
+ compact_msg: Msg = await agent.reply(
+ Msg(
+ name="reme",
+ role="user",
+ content=user_message,
+ ),
+ )
+
+ history_compact: str = compact_msg.get_text_content()
+ is_valid: bool = _is_valid_summary(history_compact)
+
+ if not is_valid:
+ logger.warning(f"Invalid summary result: {history_compact[:200]}...")
+ if self.return_dict:
+ return {"user_message": user_message, "history_compact": history_compact, "is_valid": False}
+ return ""
+
+ logger.info(f"Compactor Result:\n{history_compact}")
+
+ if self.return_dict:
+ return {"user_message": user_message, "history_compact": history_compact, "is_valid": True}
+ return history_compact
diff --git a/reme/memory/file_based/compactor.yaml b/reme/memory/file_based/components/compactor.yaml
similarity index 74%
rename from reme/memory/file_based/compactor.yaml
rename to reme/memory/file_based/components/compactor.yaml
index 2e3b15a8..6fef5383 100644
--- a/reme/memory/file_based/compactor.yaml
+++ b/reme/memory/file_based/components/compactor.yaml
@@ -7,10 +7,14 @@ system_prompt_zh: |
这些摘要可以在未来会话中用于恢复上下文。专注于保留关键信息，同时减少token数量。
initial_user_message: |
- The messages above are a conversation to summarize. Create a structured context checkpoint summary
- that another LLM will use to continue the work.
+ # Task
+ Create a structured summary from the conversation above.
- Use this EXACT format:
+ # Rules:
+ - Keep each section concise
+ - Preserve exact file paths, function names, and error messages
+
+ # Output Format:
## Goal
[What is the user trying to accomplish? Can be multiple items if the session covers different tasks.]
@@ -39,13 +43,17 @@ initial_user_message: |
- [Any data, examples, or references needed to continue]
- [Or "(none)" if not applicable]
- Keep each section concise. Preserve exact file paths, function names, and error messages.
+ Output the structured summary following the format above.
initial_user_message_zh: |
- 上述消息是一场需要总结的对话。创建一个结构化的上下文检查点摘要，
- 以便另一个LLM可以用来继续工作。
+ # 任务
+ 根据上面的对话创建一个结构化摘要。
- 使用此确切格式：
+ # 规则：
+ - 保持每个部分简洁
+ - 保留确切的文件路径、函数名称和错误消息
+
+ # 输出示例：
## 目标
[用户试图完成什么？如果会话涵盖不同任务，可以有多个项目。]
@@ -74,14 +82,13 @@ initial_user_message_zh: |
- [任何继续工作所需的数据、示例或参考资料]
- [或者如果不适用则为"(none)"]
- 保持每个部分简洁。保留确切的文件路径、函数名称和错误消息。
+ 请按照上面示例的格式，输出结构化摘要。
-update_user_message_prefix: |
- The messages above are NEW conversation messages to incorporate into the existing summary provided in
- tags.
+update_user_message: |
+ # Task
+ Update the structured summary with new conversation messages.
-update_user_message_suffix: |
- Update the existing structured summary with new information. RULES:
+ # Rules:
- PRESERVE all existing information from the previous summary
- ADD new progress, decisions, and context from the new messages
- UPDATE the Progress section: move items from "In Progress" to "Done" when completed
@@ -89,7 +96,7 @@ update_user_message_suffix: |
- PRESERVE exact file paths, function names, and error messages
- If something is no longer relevant, you may remove it
- Use this EXACT format:
+ # Output Format:
## Goal
[Preserve existing goals, add new ones if the task expanded]
@@ -116,13 +123,13 @@ update_user_message_suffix: |
## Critical Context
- [Preserve important context, add new if needed]
- Keep each section concise. Preserve exact file paths, function names, and error messages.
+ Output the structured summary following the format above.
-update_user_message_prefix_zh: |
- 上述消息是要整合到现有摘要中的新对话消息，这些消息在标签中提供。
+update_user_message_zh: |
+ # 任务
+ 使用新的对话内容来更新结构化摘要。
-update_user_message_suffix_zh: |
- 用新信息更新现有的结构化摘要。规则：
+ # 规则：
- 保留来自先前摘要的所有现有信息
- 从新消息中添加新的进展、决策和上下文
- 更新进度部分：当完成时将项目从"进行中"移到"已完成"
@@ -130,7 +137,7 @@ update_user_message_suffix_zh: |
- 保留确切的文件路径、函数名称和错误消息
- 如果某些内容不再相关，您可以删除它
- 使用此确切格式：
+ # 输出示例：
## 目标
[保留现有目标，如果任务扩展则添加新目标]
@@ -157,4 +164,4 @@ update_user_message_suffix_zh: |
## 关键上下文
- [保留重要上下文，如需要则添加新的]
- 保持每个部分简洁。保留确切的文件路径、函数名称和错误消息。
+ 请按照上面示例的格式，输出结构化摘要。
diff --git a/reme/memory/file_based/components/context_checker.py b/reme/memory/file_based/components/context_checker.py
new file mode 100644
index 00000000..f5a49b2b
--- /dev/null
+++ b/reme/memory/file_based/components/context_checker.py
@@ -0,0 +1,92 @@
+"""ContextChecker module for checking context size and splitting messages."""
+
+from agentscope.message import Msg
+
+from ..utils import AsMsgHandler
+from ....core.op import BaseOp
+from ....core.utils import get_logger
+
+logger = get_logger()
+
+
+class ContextChecker(BaseOp):
+ """
+ ContextChecker class for checking context size and splitting messages.
+
+ This class analyzes conversation messages to determine if the context
+ exceeds the specified token threshold and splits messages into two groups:
+ those that should be compacted and those to keep in context.
+
+ Attributes:
+ memory_compact_threshold (int): Token count threshold for triggering compaction.
+ memory_compact_reserve (int): Token count to reserve for recent messages.
+ """
+
+ def __init__(
+ self,
+ memory_compact_threshold: int,
+ memory_compact_reserve: int = 10000,
+ **kwargs,
+ ):
+ """
+ Initialize the ContextChecker.
+
+ Args:
+ memory_compact_threshold (int): Token count threshold for triggering
+ compaction. Messages exceeding this threshold will be split.
+ memory_compact_reserve (int): Token count to reserve for recent messages
+ to keep in context. Defaults to 10000 tokens.
+ **kwargs: Additional keyword arguments passed to BaseOp.
+ """
+ super().__init__(**kwargs)
+ self.memory_compact_threshold: int = memory_compact_threshold
+ self.memory_compact_reserve: int = memory_compact_reserve
+ assert self.memory_compact_threshold > self.memory_compact_reserve
+
+ async def execute(self) -> tuple[list[Msg], list[Msg], bool]:
+ """
+ Execute context check and split messages.
+
+ Retrieves messages from context and checks if they exceed the token
+ threshold. If so, splits them into messages to compact and messages
+ to keep.
+
+ Context Parameters:
+ messages (list[Msg]): List of conversation messages to check.
+ Retrieved from self.context.get("messages", []).
+
+ Returns:
+ tuple[list[Msg], list[Msg], bool]: A tuple containing:
+ - messages_to_compact (list[Msg]): Older messages that should
+ be compacted/summarized.
+ - messages_to_keep (list[Msg]): Recent messages to keep in context.
+ - is_valid (bool): True if the split is valid (tool calls aligned),
+ False if splitting would break conversation integrity.
+
+ Note:
+ - Returns ([], messages, True) if no compaction is needed.
+ - Ensures conversation pairs (user-assistant) are not split.
+ - is_valid=False indicates tool_use and tool_result are misaligned.
+ """
+ messages: list[Msg] = self.context.get("messages", [])
+
+ if not messages:
+ logger.info("ContextChecker: No messages to check.")
+ return [], [], True
+
+ msg_handler = AsMsgHandler(self.as_token_counter)
+ messages_to_compact, messages_to_keep, is_valid = await msg_handler.context_check(
+ messages=messages,
+ memory_compact_threshold=self.memory_compact_threshold,
+ memory_compact_reserve=self.memory_compact_reserve,
+ )
+
+ 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/components/summarizer.py b/reme/memory/file_based/components/summarizer.py
new file mode 100644
index 00000000..9733eaa5
--- /dev/null
+++ b/reme/memory/file_based/components/summarizer.py
@@ -0,0 +1,97 @@
+"""Summarizer module for memory summarization operations."""
+
+import datetime
+import zoneinfo
+
+from agentscope.agent import ReActAgent
+from agentscope.message import Msg
+from agentscope.tool import Toolkit
+
+from ..utils import AsMsgHandler
+from ....core.op import BaseOp
+from ....core.utils import get_logger
+
+logger = get_logger()
+
+
+class Summarizer(BaseOp):
+ """Summarizer class for summarizing memory messages."""
+
+ def __init__(
+ self,
+ working_dir: str,
+ memory_dir: str,
+ memory_compact_threshold: int,
+ toolkit: Toolkit | None = None,
+ console_enabled: bool = False,
+ timezone: str | None = None,
+ add_thinking_block: bool = True,
+ **kwargs,
+ ):
+ super().__init__(**kwargs)
+ self.working_dir: str = working_dir
+ self.memory_dir: str = memory_dir
+ self.memory_compact_threshold: int = memory_compact_threshold
+ self.toolkit: Toolkit | None = toolkit
+ self.console_enabled: bool = console_enabled
+ self.timezone: str | None = timezone
+ self.add_thinking_block: bool = add_thinking_block
+
+ def _get_current_datetime(self) -> datetime.datetime:
+ """Get current datetime with timezone, fallback to local time if timezone is invalid."""
+ if self.timezone:
+ try:
+ return datetime.datetime.now(zoneinfo.ZoneInfo(self.timezone))
+ except Exception as e:
+ logger.error(f"Invalid timezone: {self.timezone}, falling back to local time error={e}")
+ return datetime.datetime.now()
+
+ async def execute(self):
+ messages: list[Msg] = self.context.get("messages", [])
+
+ if not messages:
+ return ""
+
+ msg_handler = AsMsgHandler(self.as_token_counter)
+ before_token_count = await msg_handler.count_msgs_token(messages)
+ history_formatted_str: str = await msg_handler.format_msgs_to_str(
+ messages=messages,
+ memory_compact_threshold=self.memory_compact_threshold,
+ include_thinking=self.add_thinking_block,
+ )
+ after_token_count = await 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}")
+ return ""
+
+ agent = ReActAgent(
+ name="reme_summarizer",
+ model=self.as_llm,
+ sys_prompt="You are a helpful assistant.",
+ formatter=self.as_llm_formatter,
+ toolkit=self.toolkit,
+ )
+ agent.set_console_output_enabled(self.console_enabled)
+
+ user_message: str = f"# conversation\n{history_formatted_str}\n\n" + self.prompt_format(
+ "user_message",
+ date=self._get_current_datetime().strftime("%Y-%m-%d"),
+ working_dir=self.working_dir,
+ memory_dir=self.memory_dir,
+ )
+
+ summary_msg: Msg = await agent.reply(
+ Msg(
+ name="reme",
+ role="user",
+ content=user_message,
+ ),
+ )
+ for i, (msg, _) in enumerate(agent.memory.content):
+ logger.info(f"Summarizer memory[{i}]: {msg.content}")
+
+ history_summary: str = summary_msg.get_text_content()
+ logger.info(f"Summarizer Result:\n{history_summary}")
+ return history_summary
diff --git a/reme/memory/file_based/components/summarizer.yaml b/reme/memory/file_based/components/summarizer.yaml
new file mode 100644
index 00000000..14599f71
--- /dev/null
+++ b/reme/memory/file_based/components/summarizer.yaml
@@ -0,0 +1,62 @@
+user_message: |
+ Memory Pre-compression Flush Cycle.
+
+ The current session is about to enter the automatic compression phase. Please capture persistent memory AND session reflections, then write them to disk.
+
+ Current date: {date}
+ Working directory: {working_dir}
+
+ # Task
+ Immediately store persistent memory and reflections to: {memory_dir}/YYYY-MM-DD.md
+
+ # Workflow
+ 1. Extract and synthesize content from the current session:
+ - Persistent Memory: Facts, user profile updates, project states, and important events.
+ - Experience Reflection: Reusable thinking logic derived from user feedback, successful problem-solving strategies, mistakes made/pitfalls to avoid, and actionable insights for future interactions.
+ 2. `read` {memory_dir}/YYYY-MM-DD.md (if the file doesn’t exist, an error message will be returned)
+ - If the file doesn’t exist, use `write` tool directly.
+ - If the file exists, intelligently merge new information with existing content, prefer using `edit` to update specific sections.
+ - Use `write` to overwrite the entire file only if substantial restructuring is required.
+
+ # Principles
+ - Intelligently merge new information with existing content:
+ - Categorize clearly (e.g., separate "Factual Memory" from "Reflections & Logic").
+ - Avoid duplicating already recorded information.
+ - Enrich existing entries with new details where relevant.
+ - Maintain chronological order wherever applicable.
+ - Always preserve timestamps and any date/time-related context.
+ - Add only genuinely new or meaningfully enriching information.
+ - Reflections MUST focus on forming reusable cognitive frameworks based on user feedback, aiming to improve future task execution.
+ - Keep entries concise yet complete.
+ - If there’s nothing to store or reflect on, respond with [SILENT].
+
+user_message_zh: |
+ 预压缩内存刷新轮次。
+
+ 当前会话即将进入自动压缩阶段；请将持久化记忆与经验反思捕获并写入磁盘。
+
+ 当前日期：{date}
+ 工作目录：{working_dir}
+
+ # 任务
+ 立即存储持久化记忆与反思（使用路径 {memory_dir}/YYYY-MM-DD.md）。
+
+ # 工作流程
+ 1. 从当前会话中提取并综合两类内容：
+ - 持久化记忆：客观事实、用户信息更新、项目状态及重要事件。
+ - 经验反思：基于用户反馈形成的可复用思考逻辑、成功的问题解决策略、犯下的错误/应避免的陷阱，以及对未来交互有帮助的行动指南。
+ 2. `read` {memory_dir}/YYYY-MM-DD.md（如文件不存在，会返回错误提示）
+ - 若文件不存在，直接使用 `write` 工具写入。
+ - 若文件已存在，智能合并新信息与现有内容，尽可能使用 `edit` 更新特定部分，仅在需要大幅重构时使用 `write` 覆盖整个文件。
+
+ # 原则
+ - 智能合并新信息与现有内容：
+ - 将内容进行清晰的分类（例如明确区分“事实记忆”与“反思与逻辑”）。
+ - 避免重复已记录的信息。
+ - 在相关时丰富现有条目的新细节。
+ - 在适用时保持时间顺序。
+ - 始终保留时间戳、日期和时间相关上下文。
+ - 仅添加真正新的或有丰富价值的信息。
+ - 反思内容必须侧重于根据用户反馈构建可复用的思维逻辑，以改善未来的任务执行。
+ - 保持条目简洁但完整。
+ - 若无任何新内容可存储或反思，请回复 [SILENT]。
\ No newline at end of file
diff --git a/reme/memory/file_based/components/tool_result_compactor.py b/reme/memory/file_based/components/tool_result_compactor.py
new file mode 100644
index 00000000..d80b2f61
--- /dev/null
+++ b/reme/memory/file_based/components/tool_result_compactor.py
@@ -0,0 +1,161 @@
+"""Tool Result Compactor: truncate large tool results and save full content to files."""
+
+import os
+import sys
+import uuid
+from datetime import datetime, timedelta
+from pathlib import Path
+
+from agentscope.message import Msg
+
+from ..utils import truncate_text_output, DEFAULT_MAX_BYTES, TRUNCATION_NOTICE_MARKER
+from ....core.op import BaseOp
+from ....core.utils import get_logger
+
+logger = get_logger()
+
+
+class ToolResultCompactor(BaseOp):
+ """Truncate large tool_result outputs and save full content to files."""
+
+ def __init__(
+ self,
+ tool_result_dir: str | Path,
+ retention_days: int = 3,
+ old_max_bytes: int = 3000,
+ recent_max_bytes: int = DEFAULT_MAX_BYTES,
+ recent_n: int = 1,
+ encoding: str = "utf-8",
+ **kwargs,
+ ):
+ super().__init__(**kwargs)
+ self.tool_result_dir = Path(tool_result_dir)
+ self.retention_days = retention_days
+ self.old_max_bytes = old_max_bytes
+ self.recent_max_bytes = recent_max_bytes
+ self.recent_n = recent_n
+ self.encoding = encoding
+ self.tool_result_dir.mkdir(parents=True, exist_ok=True)
+
+ def _truncate(self, content: str, max_bytes: int) -> str:
+ if not content:
+ return content
+
+ try:
+ if TRUNCATION_NOTICE_MARKER in content:
+ return truncate_text_output(content, max_bytes=max_bytes, encoding=self.encoding)
+
+ if len(content.encode(self.encoding)) <= max_bytes + 100:
+ return content
+
+ saved_path: str | None = None
+ fp = self.tool_result_dir / f"{uuid.uuid4().hex}.txt"
+ fp.write_text(content, encoding=self.encoding)
+ saved_path = str(fp)
+
+ return truncate_text_output(
+ content,
+ 1,
+ content.count("\n") + 1,
+ max_bytes,
+ file_path=saved_path,
+ encoding=self.encoding,
+ )
+ except Exception as e:
+ logger.warning("Failed to truncate content, returning original: %s", e)
+ return content
+
+ def _compact(self, output: str | list[dict], max_bytes: int) -> str | list[dict]:
+ """Truncate output to max_bytes, saving full content to file if needed."""
+
+ if isinstance(output, str):
+ return self._truncate(output, max_bytes)
+ if isinstance(output, list):
+ for b in output:
+ if isinstance(b, dict) and b.get("type") == "text":
+ b["text"] = self._truncate(b.get("text", ""), max_bytes)
+ return output
+
+ async def execute(self) -> list[Msg]:
+ """Process all messages, truncating large tool results."""
+ messages: list[Msg] = self.context.get("messages", [])
+ if not messages:
+ return messages
+
+ recent_n = 0
+ for msg in reversed(messages):
+ if not isinstance(msg.content, list) or not any(
+ isinstance(b, dict) and b.get("type") == "tool_result" for b in msg.content
+ ):
+ break
+ recent_n += 1
+ split_index = max(0, len(messages) - max(recent_n, self.recent_n))
+
+ md_file_tool_ids = set()
+ try:
+ for msg in messages:
+ if not isinstance(msg.content, list):
+ continue
+
+ for block in msg.content:
+ if isinstance(block, dict) and block.get("type") == "tool_use":
+ tool_id = block.get("id", "")
+ if not tool_id:
+ continue
+
+ if (
+ block.get("name", "").lower() == "read_file"
+ and ".md" in (block.get("raw_input") or "").lower()
+ ):
+ md_file_tool_ids.add(tool_id)
+ except Exception as e:
+ logger.warning("Failed to detect md file tool ids: %s", e)
+ logger.info(f"md_file_tool_ids: {md_file_tool_ids}")
+
+ for idx, msg in enumerate(messages):
+ if not isinstance(msg.content, list):
+ continue
+ is_recent = idx >= split_index
+ max_bytes = self.recent_max_bytes if is_recent else self.old_max_bytes
+ for block in msg.content:
+ if isinstance(block, dict) and block.get("type") == "tool_result" and block.get("output"):
+ tool_use_id = block.get("id", "")
+ if tool_use_id in md_file_tool_ids:
+ effective_max_bytes = self.recent_max_bytes
+ else:
+ effective_max_bytes = max_bytes
+ block["output"] = self._compact(block["output"], effective_max_bytes)
+
+ return messages
+
+ def cleanup_expired_files(self) -> int:
+ """Clean up files older than retention_days.
+
+ Returns:
+ Number of files successfully deleted.
+ """
+ if not self.tool_result_dir.exists():
+ return 0
+
+ cutoff = datetime.now() - timedelta(days=self.retention_days)
+ deleted = failed = 0
+
+ for fp in self.tool_result_dir.glob("*.txt"):
+ try:
+ stat = os.stat(fp)
+ if sys.platform == "win32":
+ ts = stat.st_ctime # creation time on Windows
+ else:
+ ts = getattr(stat, "st_birthtime", stat.st_mtime) # macOS/BSD; Linux fallback to mtime
+ if datetime.fromtimestamp(ts) < cutoff:
+ fp.unlink()
+ deleted += 1
+ except FileNotFoundError:
+ pass # deleted by another process between glob and stat/unlink
+ except Exception as e:
+ failed += 1
+ logger.warning("Failed to delete %s: %s", fp, e)
+
+ if deleted or failed:
+ logger.info("Cleaned up %d expired files (%d failed)", deleted, failed)
+ return deleted
diff --git a/reme/memory/file_based/file_io.py b/reme/memory/file_based/file_io.py
deleted file mode 100644
index 161c3210..00000000
--- a/reme/memory/file_based/file_io.py
+++ /dev/null
@@ -1,247 +0,0 @@
-"""File I/O operations with a configurable working directory."""
-
-import os
-from pathlib import Path
-from typing import Optional
-
-from agentscope.message import TextBlock
-from agentscope.tool import ToolResponse
-
-
-class FileIO:
- """File I/O operations with a configurable working directory."""
-
- def __init__(self, working_dir: str | Path):
- """Initialize FileIO with a working directory.
-
- Args:
- working_dir (`str`):
- The working directory for resolving relative paths.
- """
- self.working_dir = Path(working_dir)
-
- def _resolve_file_path(self, file_path: str) -> str:
- """Resolve file path: use absolute path as-is,
- resolve relative path from working_dir.
-
- Args:
- file_path: The input file path (absolute or relative).
-
- Returns:
- The resolved absolute file path as string.
- """
- path = Path(file_path)
- if path.is_absolute():
- return str(path)
- else:
- return str(self.working_dir / file_path)
-
- async def read( # pylint: disable=too-many-return-statements
- self,
- file_path: str,
- start_line: Optional[int] = None,
- end_line: Optional[int] = None,
- ) -> ToolResponse:
- """Read a file. Relative paths resolve from working_dir.
-
- Use start_line/end_line to read a specific line range (output includes
- line numbers). Omit both to read the full file.
-
- Args:
- file_path (`str`):
- Path to the file.
- start_line (`int`, optional):
- First line to read (1-based, inclusive).
- end_line (`int`, optional):
- Last line to read (1-based, inclusive).
- """
- file_path = self._resolve_file_path(file_path)
-
- if not os.path.exists(file_path):
- return ToolResponse(
- content=[
- TextBlock(
- type="text",
- text=f"Error: The file {file_path} does not exist.",
- ),
- ],
- )
-
- if not os.path.isfile(file_path):
- return ToolResponse(
- content=[
- TextBlock(
- type="text",
- text=f"Error: The path {file_path} is not a file.",
- ),
- ],
- )
-
- try:
- with open(file_path, "r", encoding="utf-8") as f:
- all_lines = f.readlines()
-
- range_requested = start_line is not None or end_line is not None
-
- if range_requested:
- total = len(all_lines)
- s = max(1, start_line if start_line is not None else 1)
- e = min(total, end_line if end_line is not None else total)
-
- if s > total:
- return ToolResponse(
- content=[
- TextBlock(
- type="text",
- text=(f"Error: start_line {s} exceeds file length " f"({total} lines) in {file_path}."),
- ),
- ],
- )
-
- if s > e:
- return ToolResponse(
- content=[
- TextBlock(
- type="text",
- text=(f"Error: start_line ({s}) is greater than " f"end_line ({e}) in {file_path}."),
- ),
- ],
- )
-
- selected = all_lines[s - 1 : e]
- content = "".join(selected)
- header = f"{file_path} (lines {s}-{e} of {total})\n"
- return ToolResponse(
- content=[
- TextBlock(
- type="text",
- text=header + content,
- ),
- ],
- )
- else:
- content = "".join(all_lines)
- return ToolResponse(
- content=[
- TextBlock(
- type="text",
- text=content,
- ),
- ],
- )
-
- except Exception as e:
- return ToolResponse(
- content=[
- TextBlock(
- type="text",
- text=f"Error: Read file failed due to \n{e}",
- ),
- ],
- )
-
- async def write(
- self,
- file_path: str,
- content: str,
- ) -> ToolResponse:
- """Create or overwrite a file. Relative paths resolve from working_dir.
-
- Args:
- file_path (`str`):
- Path to the file.
- content (`str`):
- Content to write.
- """
- if not file_path:
- return ToolResponse(
- content=[
- TextBlock(
- type="text",
- text="Error: No `file_path` provide.",
- ),
- ],
- )
-
- file_path = self._resolve_file_path(file_path)
-
- try:
- with open(file_path, "w", encoding="utf-8") as file:
- file.write(content)
- return ToolResponse(
- content=[
- TextBlock(
- type="text",
- text=f"Wrote {len(content)} bytes to {file_path}.",
- ),
- ],
- )
- except Exception as e:
- return ToolResponse(
- content=[
- TextBlock(
- type="text",
- text=f"Error: Write file failed due to \n{e}",
- ),
- ],
- )
-
- async def edit(
- self,
- file_path: str,
- old_text: str,
- new_text: str,
- ) -> ToolResponse:
- """Find-and-replace text in a file. All occurrences of old_text are
- replaced with new_text. Relative paths resolve from working_dir.
-
- Args:
- file_path (`str`):
- Path to the file.
- old_text (`str`):
- Exact text to find.
- new_text (`str`):
- Replacement text.
- """
- response = await self.read(file_path=file_path)
- if response.content and len(response.content) > 0:
- error_text = response.content[0].get("text", "")
- if error_text.startswith("Error:"):
- return response
- if not response.content or len(response.content) == 0:
- return ToolResponse(
- content=[
- TextBlock(
- type="text",
- text=f"Error: Failed to read file {file_path}.",
- ),
- ],
- )
-
- content = response.content[0].get("text", "")
- if old_text not in content:
- return ToolResponse(
- content=[
- TextBlock(
- type="text",
- text=f"Error: The text to replace was not found in {file_path}.",
- ),
- ],
- )
-
- new_content = content.replace(old_text, new_text)
- write_response = await self.write(file_path=file_path, content=new_content)
-
- if write_response.content and len(write_response.content) > 0:
- write_text = write_response.content[0].get("text", "")
- if write_text.startswith("Error:"):
- return write_response
-
- return ToolResponse(
- content=[
- TextBlock(
- type="text",
- text=f"Successfully replaced text in {file_path}.",
- ),
- ],
- )
diff --git a/reme/memory/file_based/memory_formatter.py b/reme/memory/file_based/memory_formatter.py
deleted file mode 100644
index 6c0e22c3..00000000
--- a/reme/memory/file_based/memory_formatter.py
+++ /dev/null
@@ -1,249 +0,0 @@
-"""Memory Formatter for CoPaw agents.
-
-Provides memory formatting capabilities including:
-- Converting list of Msg to formatted string
-- Memory compaction with token threshold
-- Support for various content block types (text, tool_use, tool_result, etc.)
-"""
-
-import json
-import logging
-import os
-
-from agentscope.message import Msg
-from agentscope.token import HuggingFaceTokenCounter
-
-from .utils import safe_count_str_tokens, truncate_text
-
-logger = logging.getLogger(__name__)
-
-_DEFAULT_MAX_FORMATTER_TEXT_LENGTH = 2000
-
-
-class MemoryFormatter:
- """Formatter that converts list of Msg to formatted string.
-
- Formats messages into human-readable string representation with:
- - Role and timestamp information
- - Text content and tool calls
- - Memory compact threshold to limit total token count
- """
-
- def __init__(
- self,
- token_counter: HuggingFaceTokenCounter,
- memory_compact_threshold: int,
- ):
- """Initialize MemoryFormatter.
-
- Args:
- token_counter: Token counter for estimating token counts.
- memory_compact_threshold: Maximum token count before skipping
- older messages.
- """
- self._token_counter = token_counter
- self._memory_compact_threshold = memory_compact_threshold
- self.max_length = int(
- os.getenv("MAX_FORMATTER_TEXT_LENGTH", str(_DEFAULT_MAX_FORMATTER_TEXT_LENGTH)),
- )
-
- @staticmethod
- def _format_tool_result_output(output: str | list[dict]) -> str:
- """Convert tool result output to string.
-
- Args:
- output: Tool result output, either string or list of content blocks.
-
- Returns:
- Formatted string representation of the tool result.
- """
- if isinstance(output, str):
- return output
-
- textual_parts = []
-
- for block in output:
- try:
- if not isinstance(block, dict) or "type" not in block:
- logger.warning(
- "Invalid block: %s, expected a dict with 'type' key, skipped.",
- block,
- )
- continue
-
- block_type = block["type"]
-
- if block_type == "text":
- textual_parts.append(block.get("text", ""))
-
- elif block_type in ["image", "audio", "video"]:
- source = block.get("source", {})
- url = source.get("url", "")
- if url:
- textual_parts.append(
- f"[{block_type}] {url}",
- )
- else:
- textual_parts.append(f"[{block_type}]")
-
- elif block_type == "file":
- file_path = block.get("path", "") or block.get("url", "")
- file_name = block.get("name", file_path)
- textual_parts.append(f"[file] {file_name}: {file_path}")
-
- else:
- # Unknown block type: log warning and skip
- logger.warning(
- "Unsupported block type '%s' in tool result, skipped.",
- block_type,
- )
-
- except Exception as e:
- logger.warning(
- "Failed to process block %s: %s, skipped.",
- block,
- e,
- )
-
- if not textual_parts:
- return ""
- if len(textual_parts) == 1:
- return textual_parts[0]
- return "\n".join(f"- {part}" for part in textual_parts)
-
- def _format_single_msg(
- self,
- msg: Msg,
- index: int | None = None,
- add_time: bool = True,
- ) -> tuple[str, int]:
- """Format a single Msg into string representation.
-
- Similar to Message.format_message style.
-
- Args:
- msg: The Msg object to format.
- index: Optional message index for round numbering.
- add_time: Whether to include timestamp.
-
- Returns:
- Tuple of (formatted_string, token_count).
- """
- lines = []
- token_count = 0
-
- # Build header: "round{index} [{timestamp}] {role}:"
- prefix = f"round{index} " if index is not None else ""
- time_str = f"[{msg.timestamp}] " if add_time and msg.timestamp else ""
- role_str = msg.name or msg.role
- header = f"{prefix}{time_str}{role_str}:"
- lines.append(header)
- token_count += safe_count_str_tokens(self._token_counter, header)
-
- # Process content blocks
- for block in msg.get_content_blocks():
- typ = block.get("type")
-
- if typ == "text":
- text_content = truncate_text(block.get("text", ""), self.max_length)
- if text_content:
- lines.append(text_content)
- token_count += safe_count_str_tokens(self._token_counter, text_content)
-
- elif typ == "thinking":
- # Skip thinking blocks to save tokens
- pass
-
- elif typ in ["image", "audio", "video"]:
- source = block.get("source", {})
- url = source.get("url", "")
- if url:
- lines.append(f"[{typ}] {url}")
- else:
- lines.append(f"[{typ}]")
- # Estimate fixed token cost for media reference
- token_count += 10
-
- elif typ == "tool_use":
- tool_name = block.get("name", "")
- tool_input = block.get("input", {})
- try:
- arguments_str = json.dumps(tool_input, ensure_ascii=False)
- except (TypeError, ValueError):
- arguments_str = str(tool_input)
- truncated_args = truncate_text(arguments_str, self.max_length)
- tool_line = f" - tool_call={tool_name} params={truncated_args}"
- lines.append(tool_line)
- token_count += safe_count_str_tokens(self._token_counter, tool_line)
-
- elif typ == "tool_result":
- tool_name = block.get("name", "")
- output = block.get("output", "")
- formatted_output = self._format_tool_result_output(output)
- truncated_output = truncate_text(formatted_output, self.max_length)
- if truncated_output:
- result_line = f" - tool_result={tool_name} output={truncated_output}"
- lines.append(result_line)
- token_count += safe_count_str_tokens(self._token_counter, result_line)
-
- else:
- logger.warning(
- "Unsupported block type %s in message, skipped.",
- typ,
- )
-
- return "\n".join(lines), token_count
-
- def format(
- self,
- msgs: list[Msg],
- add_time: bool = True,
- add_index: bool = True,
- ) -> str:
- """Format list of Msg into a single formatted string.
-
- Messages are processed in reverse order (newest first) and older
- messages are skipped when token count exceeds memory_compact_threshold.
-
- Args:
- msgs: List of Msg objects to format.
- add_time: Whether to include timestamp in each message.
- add_index: Whether to include round index in each message.
-
- Returns:
- Formatted string with all messages joined by newlines.
- """
- if not msgs:
- return ""
-
- formatted_parts: list[str] = []
- total_token_count = 0
-
- # Process messages in reverse order (newest first)
- for i in range(len(msgs) - 1, -1, -1):
- msg = msgs[i]
- index = i if add_index else None
-
- formatted_msg, msg_token_count = self._format_single_msg(
- msg,
- index=index,
- add_time=add_time,
- )
-
- # Always include current message first, then check threshold, at least one msg
- formatted_parts.append(formatted_msg)
- total_token_count += msg_token_count
-
- # Check if we should stop adding older messages
- if total_token_count >= self._memory_compact_threshold:
- logger.info(
- "Skipping older messages: token count %d >= %d",
- total_token_count,
- self._memory_compact_threshold,
- )
- break
-
- # Reverse to restore chronological order
- formatted_parts.reverse()
-
- return "\n\n".join(formatted_parts)
diff --git a/reme/memory/file_based/reme_chat_formatter.py b/reme/memory/file_based/reme_chat_formatter.py
deleted file mode 100644
index b70e700c..00000000
--- a/reme/memory/file_based/reme_chat_formatter.py
+++ /dev/null
@@ -1,29 +0,0 @@
-"""ReMe chat formatter."""
-
-from typing import Any
-
-from agentscope.formatter import OpenAIChatFormatter
-from agentscope.token import HuggingFaceTokenCounter
-
-from .utils import _extract_text_from_messages
-
-
-class ReMeChatFormatter(OpenAIChatFormatter):
- """ReMe chat formatter class."""
-
- async def _count(self, msgs: list[dict[str, Any]]) -> int | None:
- """Count the number of tokens in the input messages. If token counter
- is not provided, `None` will be returned.
-
- Args:
- msgs (`list[Msg]`):
- The input messages to count tokens for.
- """
- if self.token_counter is None:
- return None
-
- assert isinstance(self.token_counter, HuggingFaceTokenCounter)
- text = _extract_text_from_messages(msgs)
- token_ids = self.token_counter.tokenizer.encode(text)
- token_count = len(token_ids)
- return token_count
diff --git a/reme/memory/file_based/reme_in_memory_memory.py b/reme/memory/file_based/reme_in_memory_memory.py
index 0e915382..d8293e0e 100644
--- a/reme/memory/file_based/reme_in_memory_memory.py
+++ b/reme/memory/file_based/reme_in_memory_memory.py
@@ -1,16 +1,18 @@
"""Custom memory implementation with bugfixes and extensions."""
-import logging
+import json
+from datetime import datetime
+from pathlib import Path
-from agentscope.agent._react_agent import _MemoryMark
-from agentscope.formatter import FormatterBase
+from agentscope.agent._react_agent import _MemoryMark # noqa
from agentscope.memory import InMemoryMemory
from agentscope.message import Msg
from agentscope.token import HuggingFaceTokenCounter
-from .utils import safe_count_message_tokens, safe_count_str_tokens, _get_block_tokens
+from .utils import AsMsgHandler
+from ...core.utils import get_logger
-logger = logging.getLogger(__name__)
+logger = get_logger()
class ReMeInMemoryMemory(InMemoryMemory):
@@ -19,62 +21,124 @@ class ReMeInMemoryMemory(InMemoryMemory):
def __init__(
self,
token_counter: HuggingFaceTokenCounter,
- formatter: FormatterBase,
- max_input_length: int = 0,
+ dialog_path: str | Path | None = None,
):
+ """Initialize the ReMeInMemoryMemory.
+
+ Args:
+ token_counter: Token counter for measuring content length.
+ dialog_path: Path to the dialog storage directory. If provided,
+ messages will be persisted to jsonl files when cleared or compressed.
+ """
super().__init__()
self._token_counter: HuggingFaceTokenCounter = token_counter
- self._formatter: FormatterBase = formatter
- self._max_input_length: int = max_input_length
+ self._msg_handler: AsMsgHandler = AsMsgHandler(token_counter)
+ self._dialog_path: Path | None = Path(dialog_path) if dialog_path else None
+ self._long_term_memory: str = ""
+
+ def _append_messages_to_dialog(self, messages: list[Msg]) -> int:
+ """Append messages to dialog storage file.
+
+ Saves messages to jsonl files named by message date (YYYY-mm-dd.jsonl).
+ Each line is a JSON representation of a message.
+ Messages are grouped by their timestamp date.
+
+ Args:
+ messages: List of messages to append to the dialog file.
+
+ Returns:
+ Number of messages successfully appended.
+ """
+ if not messages:
+ return 0
+
+ if self._dialog_path is None:
+ logger.warning("dialog_path is not set, skipping dialog persistence")
+ return 0
+
+ # Ensure dialog directory exists
+ try:
+ self._dialog_path.mkdir(parents=True, exist_ok=True)
+ except Exception as e:
+ logger.exception(f"Failed to create dialog directory {self._dialog_path}: {e}")
+ return 0
+
+ # Group messages by date (extracted from timestamp)
+ # timestamp format: "YYYY-mm-dd HH:MM:SS.fff"
+ messages_by_date: dict[str, list[Msg]] = {}
+ for msg in messages:
+ try:
+ if msg.timestamp:
+ # Extract date part from timestamp
+ date_str = msg.timestamp.split()[0] # "YYYY-mm-dd"
+ else:
+ date_str = datetime.now().strftime("%Y-%m-%d")
+
+ if date_str not in messages_by_date:
+ messages_by_date[date_str] = []
+ messages_by_date[date_str].append(msg)
+ except Exception as e:
+ logger.warning(f"Failed to process message timestamp: {e}, using today's date")
+ date_str = datetime.now().strftime("%Y-%m-%d")
+ if date_str not in messages_by_date:
+ messages_by_date[date_str] = []
+ messages_by_date[date_str].append(msg)
+
+ # Append messages to corresponding date files (sorted by timestamp within each date)
+ total_count = 0
+ for date_str, msgs in messages_by_date.items():
+ # Sort messages by timestamp within the same date
+ try:
+ msgs_sorted = sorted(msgs, key=lambda m: m.timestamp or "")
+ except Exception as e:
+ logger.warning(f"Failed to sort messages by timestamp: {e}")
+ msgs_sorted = msgs
+
+ filename = f"{date_str}.jsonl"
+ filepath = self._dialog_path / filename
+
+ try:
+ with open(filepath, "a", encoding="utf-8") as f:
+ for msg in msgs_sorted:
+ msg_dict = msg.to_dict()
+ f.write(json.dumps(msg_dict, ensure_ascii=False) + "\n")
+ total_count += 1
+ logger.info(f"Appended {len(msgs_sorted)} messages to {filepath}")
+ except Exception as e:
+ logger.exception(f"Failed to append messages to dialog file {filepath}: {e}")
+
+ return total_count
async def get_memory(
self,
- mark: str | None = None,
- exclude_mark: str | None = _MemoryMark.COMPRESSED,
prepend_summary: bool = True,
**_kwargs,
) -> list[Msg]:
"""Get the messages from the memory by mark (if provided).
Args:
- mark: Optional mark to filter messages
- exclude_mark: Optional mark to exclude messages
prepend_summary: Whether to prepend compressed summary
**_kwargs: Additional keyword arguments (ignored)
Returns:
List of filtered messages
"""
- if not (mark is None or isinstance(mark, str)):
- raise TypeError(f"The mark should be a string or None, but got {type(mark)}.")
-
- if not (exclude_mark is None or isinstance(exclude_mark, str)):
- raise TypeError(f"The exclude_mark should be a string or None, but got {type(exclude_mark)}.")
-
- # Filter messages based on mark
- filtered_content = [(msg, marks) for msg, marks in self.content if mark is None or mark in marks]
-
- # Further filter messages based on exclude_mark
- if exclude_mark is not None:
- filtered_content = [(msg, marks) for msg, marks in filtered_content if exclude_mark not in marks]
+ filtered_content = [(msg, marks) for msg, marks in self.content if _MemoryMark.COMPRESSED not in marks]
+ parts = []
+ if self._long_term_memory:
+ parts.append(f"# Memories\n\n{self._long_term_memory}")
if prepend_summary and self._compressed_summary:
- previous_summary = f"""
-
-{self._compressed_summary}
-
-The above is a summary of our previous conversation.
-Use it as context to maintain continuity.
- """.strip()
+ parts.append(
+ f"# Summary of previous conversation\n\n"
+ f"Previous conversation logs are offloaded to dialog/YYYY-MM-DD.jsonl (or nearby date files). "
+ "Here is the summary:\n\n"
+ f"{self._compressed_summary}\n"
+ f"The above is a summary of previous conversation, use it as context to maintain continuity.",
+ )
- return [
- Msg(
- "user",
- previous_summary,
- "user",
- ),
- *[msg for msg, _ in filtered_content],
- ]
+ if parts:
+ return [Msg("user", "\n\n".join(parts), "user"), *[msg for msg, _ in filtered_content]]
return [msg for msg, _ in filtered_content]
@@ -113,23 +177,59 @@ Use it as context to maintain continuity.
self._compressed_summary = state_dict.get("_compressed_summary", "")
async def mark_messages_compressed(self, messages: list[Msg]) -> int:
- """Mark messages as compressed and return count."""
- return await self.update_messages_mark(
- new_mark=_MemoryMark.COMPRESSED,
- msg_ids=[msg.id for msg in messages],
- )
+ """Mark messages as compressed, persist them to dialog, and remove from memory.
+
+ This method:
+ 1. Persists the given messages to the dialog storage
+ 2. Removes them from memory
+
+ Args:
+ messages: List of messages to mark as compressed.
+
+ Returns:
+ Number of messages marked as compressed.
+ """
+ if not messages:
+ return 0
+
+ # Persist messages to dialog storage instead of compressed
+ self._append_messages_to_dialog(messages)
+
+ # Remove messages from memory
+ msg_ids = {msg.id for msg in messages}
+ initial_size = len(self.content)
+ self.content = [(msg, marks) for msg, marks in self.content if msg.id not in msg_ids]
+ removed_count = initial_size - len(self.content)
+
+ logger.info(f"Marked {removed_count} messages as compressed and removed from memory")
+ return removed_count
def clear_compressed_summary(self):
"""Clear the compressed summary."""
self._compressed_summary = "" # pylint: disable=attribute-defined-outside-init
def clear_content(self):
- """Clear the content."""
- self.content.clear()
+ """Persist all messages to dialog storage and clear the content.
- async def estimate_tokens(self) -> dict:
+ This method:
+ 1. Persists all messages in memory to the dialog storage
+ 2. Clears the in-memory content
+ """
+ # Persist all messages to dialog storage
+ if self.content:
+ messages = [msg for msg, _ in self.content]
+ self._append_messages_to_dialog(messages)
+
+ # Clear in-memory content
+ self.content.clear()
+ logger.info("Cleared all messages from memory")
+
+ async def estimate_tokens(self, max_input_length: int) -> dict:
"""Estimate token usage for current memory.
+ Args:
+ max_input_length: Max input length for context usage calculation.
+
Returns:
Dict containing detailed token statistics:
- total_messages: Number of messages
@@ -138,70 +238,23 @@ Use it as context to maintain continuity.
- estimated_tokens: Total estimated tokens
- max_input_length: Max input length from config
- context_usage_ratio: Usage percentage
- - messages_detail: List of per-message token details
+ - messages_detail: List of per-message AsMsgStat objects
"""
- messages = await self.get_memory(
- exclude_mark=_MemoryMark.COMPRESSED,
- prepend_summary=False,
- )
+ messages = await self.get_memory(prepend_summary=False)
compressed_summary = self.get_compressed_summary()
- compressed_summary_tokens = safe_count_str_tokens(self._token_counter, compressed_summary)
+ compressed_summary_tokens = await self._msg_handler.count_str_token(compressed_summary)
- # Calculate total token count using formatter
- prompt = await self._formatter.format(msgs=messages)
- messages_tokens = safe_count_message_tokens(self._token_counter, prompt)
+ # Build per-message token details using AsMsgHandler
+ messages_detail = [await self._msg_handler.stat_message(msg) for msg in messages]
+
+ # Calculate total message tokens from stats
+ messages_tokens = sum(stat.total_tokens for stat in messages_detail)
estimated_tokens = messages_tokens + compressed_summary_tokens
# Calculate context usage ratio
- max_input_length = self._max_input_length
context_usage_ratio = (estimated_tokens / max_input_length * 100) if max_input_length > 0 else 0
- # Build per-message token details
- messages_detail = []
- for i, msg in enumerate(messages, 1):
- msg_detail = {
- "index": i,
- "role": msg.role,
- "text_tokens": 0,
- "blocks": [],
- "preview": "",
- }
- try:
- content = msg.content
- if isinstance(content, str):
- text_tokens = safe_count_str_tokens(self._token_counter, content)
- msg_detail["text_tokens"] = text_tokens
- msg_detail["preview"] = f"{content[:100]}..." if len(content) > 100 else content
- else:
- total_tokens = 0
- text_parts = []
- for block in content:
- if not isinstance(block, dict):
- continue
- block_type = block.get("type", "unknown")
- block_tokens, block_str = _get_block_tokens(
- block,
- block_type,
- self._token_counter,
- )
- total_tokens += block_tokens
- text_parts.append(block_str)
- msg_detail["blocks"].append(
- {
- "type": block_type,
- "tokens": block_tokens,
- },
- )
- msg_detail["text_tokens"] = total_tokens
- text_preview = "".join(text_parts)
- msg_detail["preview"] = f"{text_preview[:100]}..." if len(text_preview) > 100 else text_preview
- except Exception as e:
- msg_detail["error"] = str(e)
- msg_detail["preview"] = f""
-
- messages_detail.append(msg_detail)
-
return {
"total_messages": len(messages),
"compressed_summary_tokens": compressed_summary_tokens,
@@ -212,25 +265,28 @@ Use it as context to maintain continuity.
"messages_detail": messages_detail,
}
- async def get_history_str(self) -> str:
+ async def get_history_str(self, max_input_length: int) -> str:
"""Get formatted history string similar to /history command output.
+ Args:
+ max_input_length: Max input length for context usage calculation.
+
Returns:
Formatted string containing conversation history details
"""
- stats = await self.estimate_tokens()
+ stats = await self.estimate_tokens(max_input_length)
lines = []
- for msg_detail in stats["messages_detail"]:
+ for i, msg_stat in enumerate(stats["messages_detail"], 1):
blocks_info = ""
- if msg_detail["blocks"]:
- block_strs = [f"{b['type']}(tokens={b['tokens']})" for b in msg_detail["blocks"]]
+ if msg_stat.content:
+ block_strs = [f"{b.block_type}(tokens={b.token_count})" for b in msg_stat.content]
blocks_info = f"\n content: [{', '.join(block_strs)}]"
lines.append(
- f"[{msg_detail['index']}] **{msg_detail['role']}** "
- f"(text_tokens={msg_detail['text_tokens']})"
- f"{blocks_info}\n preview: {msg_detail['preview']}",
+ f"[{i}] **{msg_stat.role}** "
+ f"(total_tokens={msg_stat.total_tokens})"
+ f"{blocks_info}\n preview: {msg_stat.preview}",
)
return (
diff --git a/reme/memory/file_based/summarizer.py b/reme/memory/file_based/summarizer.py
deleted file mode 100644
index 462ea9c5..00000000
--- a/reme/memory/file_based/summarizer.py
+++ /dev/null
@@ -1,92 +0,0 @@
-"""Summarizer module for memory summarization operations."""
-
-import datetime
-import logging
-
-from agentscope.agent import ReActAgent
-from agentscope.formatter import FormatterBase
-from agentscope.message import Msg
-from agentscope.model import ChatModelBase
-from agentscope.token import HuggingFaceTokenCounter
-from agentscope.tool import Toolkit
-
-from .memory_formatter import MemoryFormatter
-from .file_io import FileIO
-from ...core.op import BaseOp
-
-logger = logging.getLogger(__name__)
-
-
-class Summarizer(BaseOp):
- """Summarizer class for summarizing memory messages."""
-
- def __init__(
- self,
- working_dir: str,
- memory_dir: str,
- memory_compact_threshold: int,
- chat_model: ChatModelBase,
- formatter: FormatterBase,
- token_counter: HuggingFaceTokenCounter,
- toolkit: Toolkit | None = None,
- **kwargs,
- ):
- super().__init__(**kwargs)
- self.working_dir: str = working_dir
- self.memory_dir: str = memory_dir
- self.memory_compact_threshold: int = memory_compact_threshold
-
- self.chat_model: ChatModelBase = chat_model
- self.formatter: FormatterBase = formatter
- self.as_token_counter: HuggingFaceTokenCounter = token_counter
- if toolkit is not None:
- self.toolkit: Toolkit = toolkit
- else:
- self.toolkit = Toolkit()
- file_io = FileIO(working_dir=self.working_dir)
- self.toolkit.register_tool_function(file_io.read)
- self.toolkit.register_tool_function(file_io.write)
- self.toolkit.register_tool_function(file_io.edit)
-
- async def execute(self):
- messages: list[Msg] = self.context.get("messages", [])
-
- if not messages:
- return ""
-
- formatter = MemoryFormatter(
- token_counter=self.as_token_counter,
- memory_compact_threshold=self.memory_compact_threshold,
- )
- history_formatted_str: str = formatter.format(messages)
-
- if not history_formatted_str:
- logger.warning(f"No history to summarize. messages={messages}")
- return ""
-
- agent = ReActAgent(
- name="reme_summarizer",
- model=self.chat_model,
- sys_prompt="You are a helpful assistant.",
- formatter=self.formatter,
- toolkit=self.toolkit,
- )
-
- user_message: str = f"\n{history_formatted_str}\n \n" + self.prompt_format(
- "user_message",
- date=datetime.datetime.now().strftime("%Y-%m-%d"),
- working_dir=self.working_dir,
- memory_dir=self.memory_dir,
- )
-
- summary_msg: Msg = await agent.reply(
- Msg(
- name="reme",
- role="user",
- content=user_message,
- ),
- )
-
- history_summary: str = summary_msg.get_text_content()
- logger.info(f"Summarizer Result:\n{history_summary}")
- return history_summary
diff --git a/reme/memory/file_based/summarizer.yaml b/reme/memory/file_based/summarizer.yaml
deleted file mode 100644
index 9aa0892b..00000000
--- a/reme/memory/file_based/summarizer.yaml
+++ /dev/null
@@ -1,50 +0,0 @@
-user_message: |
- Memory Pre-compression Flush Cycle Initiated
- The current session is about to enter the automatic compression phase. Please capture persistent memory and write it to disk.
-
- Current date: {date}
- Working directory: {working_dir}
-
- Immediately store persistent memory to: {memory_dir}/YYYY-MM-DD.md
-
- Workflow:
- 1. First, `read` {memory_dir}/YYYY-MM-DD.md (if the file doesn’t exist, an error message will be returned).
- 2. Intelligently merge new information with existing content (skip merging if the file doesn’t exist):
- - Avoid duplicating already recorded information
- - Enrich existing entries with new details where relevant
- - Maintain chronological order wherever applicable
- 3. Write the updated content:
- - Prefer using `edit` to update specific sections when possible
- - Use `write` to overwrite the entire file only if substantial restructuring is required
-
- Principles:
- - Always preserve timestamps and any date/time-related context
- - Add only genuinely new or meaningfully enriching information
- - Keep entries concise yet complete
- - If there’s nothing to store, respond with [SILENT]
-
-
-user_message_zh: |
- 预压缩内存刷新轮次。
- 当前会话即将进入自动压缩阶段；请将持久化记忆捕获并写入磁盘。
-
- 当前日期：{date}
- 工作目录：{working_dir}
-
- 立即存储持久化记忆（使用路径 {memory_dir}/YYYY-MM-DD.md）。
-
- 工作流程：
- 1. 先 `read` {memory_dir}/YYYY-MM-DD.md（如文件不存在，会返回错误提示）
- 2. 智能合并新信息与现有内容（若文件不存在则跳过合并）：
- - 避免重复已记录的信息
- - 在相关时丰富现有条目的新细节
- - 在适用时保持时间顺序
- 3. 写入更新后的内容：
- - 尽可能使用 `edit` 更新特定部分
- - 如需大幅重构则使用 `write` 覆盖整个文件
-
- 原则：
- - 始终保留时间戳、日期和时间相关上下文
- - 仅添加真正新的或有丰富价值的信息
- - 保持条目简洁但完整
- - 若无内容可存储，请回复 [SILENT]
diff --git a/reme/memory/file_based/tool_result_compactor.py b/reme/memory/file_based/tool_result_compactor.py
deleted file mode 100644
index 5b1ef0a4..00000000
--- a/reme/memory/file_based/tool_result_compactor.py
+++ /dev/null
@@ -1,105 +0,0 @@
-"""Tool Result Compactor: truncate large tool results and save full content to files."""
-
-import logging
-import uuid
-from datetime import datetime, timedelta
-from pathlib import Path
-
-from agentscope.message import Msg
-
-from .utils import is_truncated, truncate_text
-from ...core.op import BaseOp
-
-logger = logging.getLogger(__name__)
-
-
-class ToolResultCompactor(BaseOp):
- """Truncate large tool_result outputs and save full content to files."""
-
- def __init__(
- self,
- tool_result_dir: str | Path,
- tool_result_threshold: int,
- retention_days: int = 7,
- **kwargs,
- ):
- super().__init__(**kwargs)
- self.tool_result_dir = Path(tool_result_dir)
- self.tool_result_threshold = tool_result_threshold
- self.retention_days = retention_days
-
- def _save_and_truncate(self, content: str, tool_name: str) -> str:
- """Save full content to file and return truncated version with file reference."""
- if not content or is_truncated(content) or len(content) <= self.tool_result_threshold:
- return content
-
- # Save full content
- self.tool_result_dir.mkdir(parents=True, exist_ok=True)
- file_path = self.tool_result_dir / f"{uuid.uuid4().hex}.txt"
- created_at = datetime.now().isoformat()
-
- file_path.write_text(
- f"# tool_name: {tool_name}\n# created_at: {created_at}\n# ---\n{content}",
- encoding="utf-8",
- )
- logger.debug("Saved tool result to %s (len=%d)", file_path, len(content))
-
- # Return truncated with file reference
- return f"{truncate_text(content, self.tool_result_threshold)}\n\n[Full content saved to: {file_path}]"
-
- def _process_output(self, output: str | list[dict], tool_name: str) -> str | list[dict]:
- """Process tool result output, truncating if necessary."""
- if isinstance(output, str):
- return self._save_and_truncate(output, tool_name)
-
- if isinstance(output, list):
- return [
- (
- {**b, "text": self._save_and_truncate(b.get("text", ""), tool_name)}
- if isinstance(b, dict) and b.get("type") == "text"
- else b
- )
- for b in output
- ]
- return output
-
- async def execute(self) -> list[Msg]:
- """Process all messages, truncating large tool results."""
- messages: list[Msg] = self.context.get("messages", [])
- if not messages:
- return messages
-
- for msg in messages:
- if not isinstance(msg.content, list):
- continue
-
- for block in msg.content:
- if isinstance(block, dict) and block.get("type") == "tool_result":
- output = block.get("output")
- if output:
- block["output"] = self._process_output(output, block.get("name", "unknown"))
-
- return messages
-
- def cleanup_expired_files(self) -> int:
- """Clean up files older than retention_days."""
- if not self.tool_result_dir.exists():
- return 0
-
- cutoff = datetime.now() - timedelta(days=self.retention_days)
- deleted = 0
-
- for fp in self.tool_result_dir.glob("*.txt"):
- try:
- for line in fp.read_text(encoding="utf-8").splitlines()[:3]:
- if line.startswith("# created_at:"):
- if datetime.fromisoformat(line.split(":", 1)[1].strip()) < cutoff:
- fp.unlink()
- deleted += 1
- break
- except Exception as e:
- logger.warning("Failed to process %s: %s", fp, e)
-
- if deleted:
- logger.info("Cleaned up %d expired files", deleted)
- return deleted
diff --git a/reme/memory/file_based/tools/__init__.py b/reme/memory/file_based/tools/__init__.py
new file mode 100644
index 00000000..0fb0d814
--- /dev/null
+++ b/reme/memory/file_based/tools/__init__.py
@@ -0,0 +1,13 @@
+"""File-based memory tool implementations."""
+
+from .file_io import FileIO
+from .memory_get import MemoryGet
+from .memory_search import MemorySearch
+from .shell import Shell
+
+__all__ = [
+ "FileIO",
+ "MemoryGet",
+ "MemorySearch",
+ "Shell",
+]
diff --git a/reme/memory/file_based/tools/browser_control.py b/reme/memory/file_based/tools/browser_control.py
new file mode 100644
index 00000000..8baf9f50
--- /dev/null
+++ b/reme/memory/file_based/tools/browser_control.py
@@ -0,0 +1,2624 @@
+# -*- coding: utf-8 -*-
+# flake8: noqa: E501
+# pylint: disable=too-many-lines
+"""Browser automation tool using Playwright.
+
+Single tool with action-based API matching browser MCP: start, stop, open,
+navigate, navigate_back, screenshot, snapshot, click, type, eval, evaluate,
+resize, console_messages, handle_dialog, file_upload, fill_form, install,
+press_key, network_requests, run_code, drag, hover, select_option, tabs,
+wait_for, pdf, close. Uses refs from snapshot for ref-based actions.
+"""
+
+import asyncio
+import atexit
+import json
+import logging
+import os
+import subprocess
+import sys
+import time
+from concurrent.futures import ThreadPoolExecutor
+from typing import Any, Optional
+
+from agentscope.message import TextBlock
+from agentscope.tool import ToolResponse
+
+from ...config import (
+ get_playwright_chromium_executable_path,
+ get_system_default_browser,
+ is_running_in_container,
+)
+
+from .browser_snapshot import build_role_snapshot_from_aria
+
+logger = logging.getLogger(__name__)
+
+# Hybrid mode detection: Windows + Uvicorn reload mode requires sync Playwright
+# to avoid NotImplementedError with asyncio.create_subprocess_exec.
+# On other platforms or without reload, use async Playwright for better performance.
+_USE_SYNC_PLAYWRIGHT = sys.platform == "win32" and os.environ.get("COPAW_RELOAD_MODE") == "1"
+
+if _USE_SYNC_PLAYWRIGHT:
+ _executor: Optional[ThreadPoolExecutor] = None
+
+ def _get_executor() -> ThreadPoolExecutor:
+ global _executor
+ if _executor is None:
+ _executor = ThreadPoolExecutor(
+ max_workers=1,
+ thread_name_prefix="playwright",
+ )
+ return _executor
+
+ async def _run_sync(func, *args, **kwargs):
+ """Run a sync function in the thread pool and await the result."""
+ loop = asyncio.get_event_loop()
+ return await loop.run_in_executor(
+ _get_executor(),
+ lambda: func(*args, **kwargs),
+ )
+
+else:
+
+ async def _run_sync(func, *args, **kwargs):
+ """Fallback: directly call async function (should not be used in async mode)."""
+ return await func(*args, **kwargs)
+
+
+# Process-global browser state (one browser, multiple pages by page_id)
+_state: dict[str, Any] = {
+ "playwright": None,
+ "browser": None,
+ "context": None,
+ "pages": {},
+ "refs": {}, # page_id -> ref -> {role, name?, nth?}
+ "refs_frame": {}, # page_id -> frame for last snapshot
+ "console_logs": {}, # page_id -> list of {level, text}
+ "network_requests": {}, # page_id -> list of request dicts
+ "pending_dialogs": {}, # page_id -> dialog handlers
+ "pending_file_choosers": {}, # page_id -> FileChooser list
+ "headless": True,
+ "current_page_id": None,
+ "page_counter": 0, # monotonic counter for page_N ids, avoids reuse after close
+ "last_activity_time": 0.0, # monotonic timestamp of last browser activity
+ "_idle_task": None, # background asyncio.Task for idle watchdog
+ "_last_browser_error": None, # message when launch failed (for user-facing error)
+ "_sync_browser": None, # sync browser handle for hybrid mode
+ "_sync_context": None, # sync context handle for hybrid mode
+ "_sync_playwright": None, # sync playwright handle for hybrid mode
+}
+
+# Stop the browser after this many seconds of inactivity (default 30 minutes).
+_BROWSER_IDLE_TIMEOUT = 1800.0
+
+
+def _touch_activity() -> None:
+ """Record the current time as the last browser activity timestamp."""
+ _state["last_activity_time"] = time.monotonic()
+
+
+def _is_browser_running() -> bool:
+ """Check if browser is currently running (sync or async mode)."""
+ if _USE_SYNC_PLAYWRIGHT:
+ return _state.get("_sync_browser") is not None
+ return _state.get("browser") is not None
+
+
+def _reset_browser_state() -> None:
+ """Reset all browser-related state variables."""
+ # Clear sync/async specific state
+ _state["playwright"] = None
+ _state["browser"] = None
+ _state["context"] = None
+ _state["_sync_playwright"] = None
+ _state["_sync_browser"] = None
+ _state["_sync_context"] = None
+ # Clear shared state
+ _state["pages"].clear()
+ _state["refs"].clear()
+ _state["refs_frame"].clear()
+ _state["console_logs"].clear()
+ _state["network_requests"].clear()
+ _state["pending_dialogs"].clear()
+ _state["pending_file_choosers"].clear()
+ _state["current_page_id"] = None
+ _state["page_counter"] = 0
+ _state["last_activity_time"] = 0.0
+ _state["headless"] = True
+
+
+async def _idle_watchdog(idle_seconds: float = _BROWSER_IDLE_TIMEOUT) -> None:
+ """Background task: stop the browser after it has been idle for *idle_seconds*.
+
+ This reclaims Chrome renderer processes that accumulate when pages are
+ opened during agent tasks but never explicitly closed.
+ """
+ try:
+ while True:
+ await asyncio.sleep(60) # check every minute
+ if not _is_browser_running():
+ return
+ idle = time.monotonic() - _state.get("last_activity_time", 0.0)
+ if idle >= idle_seconds:
+ logger.info(
+ "Browser idle for %.0fs (limit %.0fs), stopping to release resources",
+ idle,
+ idle_seconds,
+ )
+ await _action_stop()
+ return
+ except asyncio.CancelledError:
+ pass
+
+
+def _atexit_cleanup() -> None:
+ """Best-effort browser cleanup registered with :func:`atexit`.
+
+ Playwright child processes are cleaned up by the OS when the parent
+ exits, but this gives Playwright a chance to flush any pending I/O and
+ close Chrome gracefully before the process disappears.
+ """
+ if not _is_browser_running():
+ return
+
+ try:
+ loop = asyncio.get_event_loop()
+ if not loop.is_running() and not loop.is_closed():
+ loop.run_until_complete(_action_stop())
+ except Exception:
+ pass
+
+
+atexit.register(_atexit_cleanup)
+
+
+def _tool_response(text: str) -> ToolResponse:
+ """Wrap text for agentscope Toolkit (return ToolResponse)."""
+ return ToolResponse(
+ content=[TextBlock(type="text", text=text)],
+ )
+
+
+def _chromium_launch_args() -> list[str]:
+ """Extra args for Chromium when running in container."""
+ if is_running_in_container():
+ return ["--no-sandbox", "--disable-dev-shm-usage"]
+ return []
+
+
+def _chromium_executable_path() -> str | None:
+ """Chromium executable path when set (e.g. container); else None."""
+ return get_playwright_chromium_executable_path()
+
+
+def _use_webkit_fallback() -> bool:
+ """True only on macOS when no system Chrome/Edge/Chromium found.
+ Use WebKit (Safari) to avoid downloading Chromium. Windows has no system
+ WebKit, so we never use webkit there.
+ """
+ return sys.platform == "darwin" and _chromium_executable_path() is None
+
+
+def _ensure_playwright_async():
+ """Import async_playwright; raise ImportError with hint if missing."""
+ try:
+ from playwright.async_api import async_playwright
+
+ return async_playwright
+ except ImportError as exc:
+ raise ImportError(
+ "Playwright not installed. Use the same Python that runs CoPaw (e.g. "
+ "activate your venv or use 'uv run'): "
+ f"'{sys.executable}' -m pip install playwright && "
+ f"'{sys.executable}' -m playwright install",
+ ) from exc
+
+
+def _ensure_playwright_sync():
+ """Import sync_playwright; raise ImportError with hint if missing."""
+ try:
+ from playwright.sync_api import sync_playwright
+
+ return sync_playwright
+ except ImportError as exc:
+ raise ImportError(
+ "Playwright not installed. Use the same Python that runs CoPaw (e.g. "
+ "activate your venv or use 'uv run'): "
+ f"'{sys.executable}' -m pip install playwright && "
+ f"'{sys.executable}' -m playwright install",
+ ) from exc
+
+
+def _sync_browser_launch(headless: bool):
+ """Launch browser using sync Playwright (for hybrid mode)."""
+ sync_playwright = _ensure_playwright_sync()
+ pw = sync_playwright().start() # Start without context manager
+ use_default = not is_running_in_container() and os.environ.get(
+ "COPAW_BROWSER_USE_DEFAULT",
+ "1",
+ ).strip().lower() in ("1", "true", "yes")
+ default_kind, default_path = get_system_default_browser() if use_default else (None, None)
+ exe: Optional[str] = None
+ if default_kind == "chromium" and default_path:
+ exe = default_path
+ elif default_kind != "webkit":
+ exe = _chromium_executable_path()
+
+ if exe:
+ launch_kwargs = {"headless": headless}
+ extra_args = _chromium_launch_args()
+ if extra_args:
+ launch_kwargs["args"] = extra_args
+ launch_kwargs["executable_path"] = exe
+ browser = pw.chromium.launch(**launch_kwargs)
+ elif default_kind == "webkit" or sys.platform == "darwin":
+ browser = pw.webkit.launch(headless=headless)
+ else:
+ launch_kwargs = {"headless": headless}
+ extra_args = _chromium_launch_args()
+ if extra_args:
+ launch_kwargs["args"] = extra_args
+ browser = pw.chromium.launch(**launch_kwargs)
+
+ context = browser.new_context()
+ _attach_context_listeners(context)
+ return pw, browser, context
+
+
+def _sync_browser_close():
+ """Close browser using sync Playwright (for hybrid mode)."""
+ if _state["_sync_browser"] is not None:
+ try:
+ _state["_sync_browser"].close()
+ except Exception:
+ pass
+ if _state["_sync_playwright"] is not None:
+ try:
+ _state["_sync_playwright"].stop()
+ except Exception:
+ pass
+
+
+def _parse_json_param(value: str, default: Any = None):
+ """Parse optional JSON string param (e.g. fields, paths, values)."""
+ if not value or not isinstance(value, str):
+ return default
+ value = value.strip()
+ if not value:
+ return default
+ try:
+ return json.loads(value)
+ except json.JSONDecodeError:
+ if "," in value:
+ return [x.strip() for x in value.split(",")]
+ return default
+
+
+async def browser_use( # pylint: disable=R0911,R0912
+ action: str,
+ url: str = "",
+ page_id: str = "default",
+ selector: str = "",
+ text: str = "",
+ code: str = "",
+ path: str = "",
+ wait: int = 0,
+ full_page: bool = False,
+ width: int = 0,
+ height: int = 0,
+ level: str = "info",
+ filename: str = "",
+ accept: bool = True,
+ prompt_text: str = "",
+ ref: str = "",
+ element: str = "",
+ paths_json: str = "",
+ fields_json: str = "",
+ key: str = "",
+ submit: bool = False,
+ slowly: bool = False,
+ include_static: bool = False,
+ screenshot_type: str = "png",
+ snapshot_filename: str = "",
+ double_click: bool = False,
+ button: str = "left",
+ modifiers_json: str = "",
+ start_ref: str = "",
+ end_ref: str = "",
+ start_selector: str = "",
+ end_selector: str = "",
+ start_element: str = "",
+ end_element: str = "",
+ values_json: str = "",
+ tab_action: str = "",
+ index: int = -1,
+ wait_time: float = 0,
+ text_gone: str = "",
+ frame_selector: str = "",
+ headed: bool = False,
+) -> ToolResponse:
+ """Control browser (Playwright). Default is headless. Use headed=True with
+ action=start to open a visible browser window. Flow: start, open(url),
+ snapshot to get refs, then click/type etc. with ref or selector. Use
+ page_id for multiple tabs.
+
+ Args:
+ action (str):
+ Required. Action type. Values: start, stop, open, navigate,
+ navigate_back, snapshot, screenshot, click, type, eval, evaluate,
+ resize, console_messages, network_requests, handle_dialog,
+ file_upload, fill_form, install, press_key, run_code, drag, hover,
+ select_option, tabs, wait_for, pdf, close.
+ url (str):
+ URL to open. Required for action=open or navigate.
+ page_id (str):
+ Page/tab identifier, default "default". Use different page_id for
+ multiple tabs.
+ selector (str):
+ CSS selector to locate element for click/type/hover etc. Prefer
+ ref when available.
+ text (str):
+ Text to type. Required for action=type.
+ code (str):
+ JavaScript code. Required for action=eval, evaluate, or run_code.
+ path (str):
+ File path for screenshot save or PDF export.
+ wait (int):
+ Milliseconds to wait after click. Used with action=click.
+ full_page (bool):
+ Whether to capture full page. Used with action=screenshot.
+ width (int):
+ Viewport width in pixels. Used with action=resize.
+ height (int):
+ Viewport height in pixels. Used with action=resize.
+ level (str):
+ Console log level filter, e.g. "info" or "error". Used with
+ action=console_messages.
+ filename (str):
+ Filename for saving logs or screenshot. Used with
+ console_messages, network_requests, screenshot.
+ accept (bool):
+ Whether to accept dialog (true) or dismiss (false). Used with
+ action=handle_dialog.
+ prompt_text (str):
+ Input for prompt dialog. Used with action=handle_dialog when
+ dialog is prompt.
+ ref (str):
+ Element ref from snapshot output; use for stable targeting. Prefer
+ ref for click/type/hover/screenshot/evaluate/select_option.
+ element (str):
+ Element description for evaluate etc. Prefer ref when available.
+ paths_json (str):
+ JSON array string of file paths. Used with action=file_upload.
+ fields_json (str):
+ JSON object string of form field name to value. Used with
+ action=fill_form.
+ key (str):
+ Key name, e.g. "Enter", "Control+a". Required for
+ action=press_key.
+ submit (bool):
+ Whether to submit (press Enter) after typing. Used with
+ action=type.
+ slowly (bool):
+ Whether to type character by character. Used with action=type.
+ include_static (bool):
+ Whether to include static resource requests. Used with
+ action=network_requests.
+ screenshot_type (str):
+ Screenshot format, "png" or "jpeg". Used with action=screenshot.
+ snapshot_filename (str):
+ File path to save snapshot output. Used with action=snapshot.
+ double_click (bool):
+ Whether to double-click. Used with action=click.
+ button (str):
+ Mouse button: "left", "right", or "middle". Used with
+ action=click.
+ modifiers_json (str):
+ JSON array of modifier keys, e.g. ["Shift","Control"]. Used with
+ action=click.
+ start_ref (str):
+ Drag start element ref. Used with action=drag.
+ end_ref (str):
+ Drag end element ref. Used with action=drag.
+ start_selector (str):
+ Drag start CSS selector. Used with action=drag.
+ end_selector (str):
+ Drag end CSS selector. Used with action=drag.
+ start_element (str):
+ Drag start element description. Used with action=drag.
+ end_element (str):
+ Drag end element description. Used with action=drag.
+ values_json (str):
+ JSON of option value(s) for select. Used with
+ action=select_option.
+ tab_action (str):
+ Tab action: list, new, close, or select. Required for
+ action=tabs.
+ index (int):
+ Tab index for tabs select, zero-based. Used with action=tabs.
+ wait_time (float):
+ Seconds to wait. Used with action=wait_for.
+ text_gone (str):
+ Wait until this text disappears from page. Used with
+ action=wait_for.
+ frame_selector (str):
+ iframe selector, e.g. "iframe#main". Set when operating inside
+ that iframe in snapshot/click/type etc.
+ headed (bool):
+ When True with action=start, launch a visible browser window
+ (non-headless). User can see the real browser. Default False.
+ """
+ action = (action or "").strip().lower()
+ if not action:
+ return _tool_response(
+ json.dumps(
+ {"ok": False, "error": "action required"},
+ ensure_ascii=False,
+ indent=2,
+ ),
+ )
+
+ page_id = (page_id or "default").strip() or "default"
+ current = _state.get("current_page_id")
+ pages = _state.get("pages") or {}
+ if page_id == "default" and current and current in pages:
+ page_id = current
+
+ try:
+ if action == "start":
+ return await _action_start(headed=headed)
+ if action == "stop":
+ return await _action_stop()
+ if action == "open":
+ return await _action_open(url, page_id)
+ if action == "navigate":
+ return await _action_navigate(url, page_id)
+ if action == "navigate_back":
+ return await _action_navigate_back(page_id)
+ if action in ("screenshot", "take_screenshot"):
+ return await _action_screenshot(
+ page_id,
+ path or filename,
+ full_page,
+ screenshot_type,
+ ref,
+ element,
+ frame_selector,
+ )
+ if action == "snapshot":
+ return await _action_snapshot(
+ page_id,
+ snapshot_filename or filename,
+ frame_selector,
+ )
+ if action == "click":
+ return await _action_click(
+ page_id,
+ selector,
+ ref,
+ element,
+ wait,
+ double_click,
+ button,
+ modifiers_json,
+ frame_selector,
+ )
+ if action == "type":
+ return await _action_type(
+ page_id,
+ selector,
+ ref,
+ element,
+ text,
+ submit,
+ slowly,
+ frame_selector,
+ )
+ if action == "eval":
+ return await _action_eval(page_id, code)
+ if action == "evaluate":
+ return await _action_evaluate(
+ page_id,
+ code,
+ ref,
+ element,
+ frame_selector,
+ )
+ if action == "resize":
+ return await _action_resize(page_id, width, height)
+ if action == "console_messages":
+ return await _action_console_messages(
+ page_id,
+ level,
+ filename or path,
+ )
+ if action == "handle_dialog":
+ return await _action_handle_dialog(page_id, accept, prompt_text)
+ if action == "file_upload":
+ return await _action_file_upload(page_id, paths_json)
+ if action == "fill_form":
+ return await _action_fill_form(page_id, fields_json)
+ if action == "install":
+ return await _action_install()
+ if action == "press_key":
+ return await _action_press_key(page_id, key)
+ if action == "network_requests":
+ return await _action_network_requests(
+ page_id,
+ include_static,
+ filename or path,
+ )
+ if action == "run_code":
+ return await _action_run_code(page_id, code)
+ if action == "drag":
+ return await _action_drag(
+ page_id,
+ start_ref,
+ end_ref,
+ start_selector,
+ end_selector,
+ start_element,
+ end_element,
+ frame_selector,
+ )
+ if action == "hover":
+ return await _action_hover(
+ page_id,
+ ref,
+ element,
+ selector,
+ frame_selector,
+ )
+ if action == "select_option":
+ return await _action_select_option(
+ page_id,
+ ref,
+ element,
+ values_json,
+ frame_selector,
+ )
+ if action == "tabs":
+ return await _action_tabs(page_id, tab_action, index)
+ if action == "wait_for":
+ return await _action_wait_for(page_id, wait_time, text, text_gone)
+ if action == "pdf":
+ return await _action_pdf(page_id, path)
+ if action == "close":
+ return await _action_close(page_id)
+ return _tool_response(
+ json.dumps(
+ {"ok": False, "error": f"Unknown action: {action}"},
+ ensure_ascii=False,
+ indent=2,
+ ),
+ )
+ except Exception as e:
+ logger.exception("Browser tool error: %s", e, exc_info=True)
+ return _tool_response(
+ json.dumps(
+ {"ok": False, "error": str(e)},
+ ensure_ascii=False,
+ indent=2,
+ ),
+ )
+
+
+def _get_page(page_id: str):
+ """Return page for page_id or None if not found."""
+ return _state["pages"].get(page_id)
+
+
+def _get_refs(page_id: str) -> dict[str, dict]:
+ """Return refs map for page_id (ref -> {role, name?, nth?})."""
+ return _state["refs"].setdefault(page_id, {})
+
+
+def _get_root(page, _page_id: str, frame_selector: str = ""):
+ """Return page or frame for frame_selector (ref/selector)."""
+ if not (frame_selector and frame_selector.strip()):
+ return page
+ return page.frame_locator(frame_selector.strip())
+
+
+def _get_locator_by_ref(
+ page,
+ page_id: str,
+ ref: str,
+ frame_selector: str = "",
+):
+ """Resolve snapshot ref to locator; frame_selector for iframe."""
+ refs = _get_refs(page_id)
+ info = refs.get(ref)
+ if not info:
+ return None
+ role = info.get("role", "generic")
+ name = info.get("name")
+ nth = info.get("nth", 0)
+ root = _get_root(page, page_id, frame_selector)
+ locator = root.get_by_role(role, name=name or None)
+ if nth is not None and nth > 0:
+ locator = locator.nth(nth)
+ return locator
+
+
+def _attach_page_listeners(page, page_id: str) -> None:
+ """Attach console and request listeners for a page."""
+ logs = _state["console_logs"].setdefault(page_id, [])
+
+ def on_console(msg):
+ logs.append({"level": msg.type, "text": msg.text})
+
+ page.on("console", on_console)
+ requests_list = _state["network_requests"].setdefault(page_id, [])
+
+ def on_request(req):
+ requests_list.append(
+ {
+ "url": req.url,
+ "method": req.method,
+ "resourceType": getattr(req, "resource_type", None),
+ },
+ )
+
+ def on_response(res):
+ for r in requests_list:
+ if r.get("url") == res.url and "status" not in r:
+ r["status"] = res.status
+ break
+
+ page.on("request", on_request)
+ page.on("response", on_response)
+ dialogs = _state["pending_dialogs"].setdefault(page_id, [])
+
+ def on_dialog(dialog):
+ dialogs.append(dialog)
+
+ page.on("dialog", on_dialog)
+ choosers = _state["pending_file_choosers"].setdefault(page_id, [])
+
+ def on_filechooser(chooser):
+ choosers.append(chooser)
+
+ page.on("filechooser", on_filechooser)
+
+
+def _next_page_id() -> str:
+ """Return a unique page_id (page_N).
+ Uses monotonic counter so IDs are not reused after close."""
+ _state["page_counter"] = _state.get("page_counter", 0) + 1
+ return f"page_{_state['page_counter']}"
+
+
+def _attach_context_listeners(context) -> None:
+ """When the page opens a new tab (e.g. target=_blank, window.open),
+ register it and set as current."""
+
+ def on_page(page):
+ new_id = _next_page_id()
+ _state["refs"][new_id] = {}
+ _state["console_logs"][new_id] = []
+ _state["network_requests"][new_id] = []
+ _state["pending_dialogs"][new_id] = []
+ _state["pending_file_choosers"][new_id] = []
+ _attach_page_listeners(page, new_id)
+ _state["pages"][new_id] = page
+ _state["current_page_id"] = new_id
+ logger.debug(
+ "New tab opened by page, registered as page_id=%s",
+ new_id,
+ )
+
+ context.on("page", on_page)
+
+
+async def _ensure_browser() -> bool: # pylint: disable=too-many-branches
+ """Start browser if not running. Return True if ready, False on failure."""
+ # Check browser state based on mode
+ if _USE_SYNC_PLAYWRIGHT:
+ if _state["_sync_browser"] is not None and _state["_sync_context"] is not None:
+ _touch_activity()
+ return True
+ else:
+ if _state["browser"] is not None and _state["context"] is not None:
+ _touch_activity()
+ return True
+
+ try:
+ if _USE_SYNC_PLAYWRIGHT:
+ # Hybrid mode: use sync Playwright in thread pool
+ loop = asyncio.get_event_loop()
+ pw, browser, context = await loop.run_in_executor(
+ _get_executor(),
+ lambda: _sync_browser_launch(_state["headless"]),
+ )
+ _state["_sync_playwright"] = pw
+ _state["_sync_browser"] = browser
+ _state["_sync_context"] = context
+ else:
+ # Standard mode: use async Playwright
+ async_playwright = _ensure_playwright_async()
+ pw = await async_playwright().start()
+ # Prefer OS default browser when available (e.g. user's default Chrome/Safari).
+ use_default = not is_running_in_container() and os.environ.get(
+ "COPAW_BROWSER_USE_DEFAULT",
+ "1",
+ ).strip().lower() in ("1", "true", "yes")
+ default_kind, default_path = get_system_default_browser() if use_default else (None, None)
+ exe: Optional[str] = None
+ if default_kind == "chromium" and default_path:
+ exe = default_path
+ elif default_kind != "webkit":
+ exe = _chromium_executable_path()
+ if exe:
+ # System Chrome/Edge/Chromium (default or discovered)
+ launch_kwargs: dict[str, Any] = {
+ "headless": _state["headless"],
+ }
+ extra_args = _chromium_launch_args()
+ if extra_args:
+ launch_kwargs["args"] = extra_args
+ launch_kwargs["executable_path"] = exe
+ pw_browser = await pw.chromium.launch(**launch_kwargs)
+ elif default_kind == "webkit" or sys.platform == "darwin":
+ # macOS: default Safari or no Chromium → use WebKit (Safari)
+ pw_browser = await pw.webkit.launch(
+ headless=_state["headless"],
+ )
+ else:
+ # Windows/Linux without system Chromium → Playwright's Chromium
+ launch_kwargs = {"headless": _state["headless"]}
+ extra_args = _chromium_launch_args()
+ if extra_args:
+ launch_kwargs["args"] = extra_args
+ pw_browser = await pw.chromium.launch(**launch_kwargs)
+ context = await pw_browser.new_context()
+ _attach_context_listeners(context)
+ _state["playwright"] = pw
+ _state["browser"] = pw_browser
+ _state["context"] = context
+ _state["_last_browser_error"] = None
+ _touch_activity()
+ _start_idle_watchdog()
+ return True
+ except Exception as e:
+ _state["_last_browser_error"] = str(e)
+ return False
+
+
+def _start_idle_watchdog() -> None:
+ """Cancel any existing idle watchdog and start a fresh one."""
+ old_task = _state.get("_idle_task")
+ if old_task and not old_task.done():
+ old_task.cancel()
+ _state["_idle_task"] = asyncio.ensure_future(_idle_watchdog())
+
+
+def _cancel_idle_watchdog() -> None:
+ """Cancel the idle watchdog, if running."""
+ task = _state.get("_idle_task")
+ if task and not task.done():
+ task.cancel()
+ _state["_idle_task"] = None
+
+
+# pylint: disable=R0912,R0915
+async def _action_start(
+ headed: bool = False,
+) -> ToolResponse:
+ # Check browser state based on mode
+ if _USE_SYNC_PLAYWRIGHT:
+ browser_exists = _state["_sync_browser"] is not None
+ current_headless = not _state.get("_sync_headless", True)
+ else:
+ browser_exists = _state["browser"] is not None
+ current_headless = _state["headless"]
+
+ # If user asks for visible window (headed=True)
+ # but browser is already running headless, restart with headed
+ if browser_exists:
+ if headed and current_headless:
+ _cancel_idle_watchdog()
+ try:
+ await _action_stop()
+ except Exception:
+ pass
+ else:
+ return _tool_response(
+ json.dumps(
+ {"ok": True, "message": "Browser already running"},
+ ensure_ascii=False,
+ indent=2,
+ ),
+ )
+ # Default: headless (background). Only headed=True (e.g. browser_visible skill) shows window.
+ _state["headless"] = not headed
+
+ try:
+ if _USE_SYNC_PLAYWRIGHT:
+ loop = asyncio.get_event_loop()
+ pw, browser, context = await loop.run_in_executor(
+ _get_executor(),
+ lambda: _sync_browser_launch(_state["headless"]),
+ )
+ _state["_sync_playwright"] = pw
+ _state["_sync_browser"] = browser
+ _state["_sync_context"] = context
+ _state["_sync_headless"] = not headed
+ else:
+ async_playwright = _ensure_playwright_async()
+ pw = await async_playwright().start()
+ use_default = not is_running_in_container() and os.environ.get(
+ "COPAW_BROWSER_USE_DEFAULT",
+ "1",
+ ).strip().lower() in ("1", "true", "yes")
+ default_kind, default_path = get_system_default_browser() if use_default else (None, None)
+ exe: Optional[str] = None
+ if default_kind == "chromium" and default_path:
+ exe = default_path
+ elif default_kind != "webkit":
+ exe = _chromium_executable_path()
+ if exe:
+ launch_kwargs = {"headless": _state["headless"]}
+ extra_args = _chromium_launch_args()
+ if extra_args:
+ launch_kwargs["args"] = extra_args
+ launch_kwargs["executable_path"] = exe
+ pw_browser = await pw.chromium.launch(**launch_kwargs)
+ elif default_kind == "webkit" or sys.platform == "darwin":
+ pw_browser = await pw.webkit.launch(
+ headless=_state["headless"],
+ )
+ else:
+ launch_kwargs = {"headless": _state["headless"]}
+ extra_args = _chromium_launch_args()
+ if extra_args:
+ launch_kwargs["args"] = extra_args
+ pw_browser = await pw.chromium.launch(**launch_kwargs)
+ context = await pw_browser.new_context()
+ _attach_context_listeners(context)
+ _state["playwright"] = pw
+ _state["browser"] = pw_browser
+ _state["context"] = context
+ _touch_activity()
+ _start_idle_watchdog()
+ msg = "Browser started (visible window)" if not _state["headless"] else "Browser started"
+ return _tool_response(
+ json.dumps(
+ {"ok": True, "message": msg},
+ ensure_ascii=False,
+ indent=2,
+ ),
+ )
+ except Exception as e:
+ return _tool_response(
+ json.dumps(
+ {"ok": False, "error": f"Browser start failed: {e!s}"},
+ ensure_ascii=False,
+ indent=2,
+ ),
+ )
+
+
+async def _action_stop() -> ToolResponse:
+ _cancel_idle_watchdog()
+
+ # Check browser state based on mode
+ if not _is_browser_running():
+ return _tool_response(
+ json.dumps(
+ {"ok": True, "message": "Browser not running"},
+ ensure_ascii=False,
+ indent=2,
+ ),
+ )
+
+ if _USE_SYNC_PLAYWRIGHT:
+ # Close sync browser in thread pool
+ loop = asyncio.get_event_loop()
+ try:
+ await loop.run_in_executor(
+ _get_executor(),
+ _sync_browser_close,
+ )
+ except Exception as e:
+ return _tool_response(
+ json.dumps(
+ {"ok": False, "error": f"Browser stop failed: {e!s}"},
+ ensure_ascii=False,
+ indent=2,
+ ),
+ )
+ finally:
+ _reset_browser_state()
+ else:
+ # Standard async mode
+ try:
+ await _state["browser"].close()
+ if _state["playwright"] is not None:
+ await _state["playwright"].stop()
+ except Exception as e:
+ return _tool_response(
+ json.dumps(
+ {"ok": False, "error": f"Browser stop failed: {e!s}"},
+ ensure_ascii=False,
+ indent=2,
+ ),
+ )
+ finally:
+ _reset_browser_state()
+
+ return _tool_response(
+ json.dumps(
+ {"ok": True, "message": "Browser stopped"},
+ ensure_ascii=False,
+ indent=2,
+ ),
+ )
+
+
+async def _action_open(url: str, page_id: str) -> ToolResponse:
+ url = (url or "").strip()
+ if not url:
+ return _tool_response(
+ json.dumps(
+ {"ok": False, "error": "url required for open"},
+ ensure_ascii=False,
+ indent=2,
+ ),
+ )
+ if not await _ensure_browser():
+ err = _state.get("_last_browser_error") or "Browser not started"
+ return _tool_response(
+ json.dumps(
+ {"ok": False, "error": err},
+ ensure_ascii=False,
+ indent=2,
+ ),
+ )
+ try:
+ if _USE_SYNC_PLAYWRIGHT:
+ # Hybrid mode: create page in thread pool
+ loop = asyncio.get_event_loop()
+ # pylint: disable=unnecessary-lambda
+ page = await loop.run_in_executor(
+ _get_executor(),
+ lambda: _state["_sync_context"].new_page(),
+ )
+ else:
+ # Standard async mode
+ page = await _state["context"].new_page()
+
+ _state["refs"][page_id] = {}
+ _state["console_logs"][page_id] = []
+ _state["network_requests"][page_id] = []
+ _state["pending_dialogs"][page_id] = []
+ _state["pending_file_choosers"][page_id] = []
+ _attach_page_listeners(page, page_id)
+
+ if _USE_SYNC_PLAYWRIGHT:
+ loop = asyncio.get_event_loop()
+ await loop.run_in_executor(
+ _get_executor(),
+ lambda: page.goto(url),
+ )
+ else:
+ await page.goto(url)
+
+ _state["pages"][page_id] = page
+ _state["current_page_id"] = page_id
+ return _tool_response(
+ json.dumps(
+ {
+ "ok": True,
+ "message": f"Opened {url}",
+ "page_id": page_id,
+ "url": url,
+ },
+ ensure_ascii=False,
+ indent=2,
+ ),
+ )
+ except Exception as e:
+ return _tool_response(
+ json.dumps(
+ {"ok": False, "error": f"Open failed: {e!s}"},
+ ensure_ascii=False,
+ indent=2,
+ ),
+ )
+
+
+async def _action_navigate(url: str, page_id: str) -> ToolResponse:
+ url = (url or "").strip()
+ if not url:
+ return _tool_response(
+ json.dumps(
+ {"ok": False, "error": "url required for navigate"},
+ ensure_ascii=False,
+ indent=2,
+ ),
+ )
+ page = _get_page(page_id)
+ if not page:
+ return _tool_response(
+ json.dumps(
+ {"ok": False, "error": f"Page '{page_id}' not found"},
+ ensure_ascii=False,
+ indent=2,
+ ),
+ )
+ try:
+ if _USE_SYNC_PLAYWRIGHT:
+ loop = asyncio.get_event_loop()
+ await loop.run_in_executor(
+ _get_executor(),
+ lambda: page.goto(url),
+ )
+ else:
+ await page.goto(url)
+ _state["current_page_id"] = page_id
+ return _tool_response(
+ json.dumps(
+ {
+ "ok": True,
+ "message": f"Navigated to {url}",
+ "url": page.url,
+ },
+ ensure_ascii=False,
+ indent=2,
+ ),
+ )
+ except Exception as e:
+ return _tool_response(
+ json.dumps(
+ {"ok": False, "error": f"Navigate failed: {e!s}"},
+ ensure_ascii=False,
+ indent=2,
+ ),
+ )
+
+
+async def _action_screenshot(
+ page_id: str,
+ path: str,
+ full_page: bool,
+ screenshot_type: str = "png",
+ ref: str = "",
+ element: str = "", # pylint: disable=unused-argument
+ frame_selector: str = "",
+) -> ToolResponse:
+ path = (path or "").strip()
+ if not path:
+ ext = "jpeg" if screenshot_type == "jpeg" else "png"
+ path = f"page-{int(time.time())}.{ext}"
+ page = _get_page(page_id)
+ if not page:
+ return _tool_response(
+ json.dumps(
+ {"ok": False, "error": f"Page '{page_id}' not found"},
+ ensure_ascii=False,
+ indent=2,
+ ),
+ )
+ try:
+ if ref and ref.strip():
+ locator = _get_locator_by_ref(
+ page,
+ page_id,
+ ref.strip(),
+ frame_selector,
+ )
+ if locator is None:
+ return _tool_response(
+ json.dumps(
+ {"ok": False, "error": f"Unknown ref: {ref}"},
+ ensure_ascii=False,
+ indent=2,
+ ),
+ )
+ if _USE_SYNC_PLAYWRIGHT:
+ await _run_sync(
+ locator.screenshot,
+ path=path,
+ type=screenshot_type if screenshot_type == "jpeg" else "png",
+ )
+ else:
+ await locator.screenshot(
+ path=path,
+ type=screenshot_type if screenshot_type == "jpeg" else "png",
+ )
+ else:
+ if frame_selector and frame_selector.strip():
+ root = _get_root(page, page_id, frame_selector)
+ locator = root.locator("body").first
+ if _USE_SYNC_PLAYWRIGHT:
+ await _run_sync(
+ locator.screenshot,
+ path=path,
+ type=screenshot_type if screenshot_type == "jpeg" else "png",
+ )
+ else:
+ await locator.screenshot(
+ path=path,
+ type=screenshot_type if screenshot_type == "jpeg" else "png",
+ )
+ else:
+ if _USE_SYNC_PLAYWRIGHT:
+ await _run_sync(
+ page.screenshot,
+ path=path,
+ full_page=full_page,
+ type=screenshot_type if screenshot_type == "jpeg" else "png",
+ )
+ else:
+ await page.screenshot(
+ path=path,
+ full_page=full_page,
+ type=screenshot_type if screenshot_type == "jpeg" else "png",
+ )
+ return _tool_response(
+ json.dumps(
+ {
+ "ok": True,
+ "message": f"Screenshot saved to {path}",
+ "path": path,
+ },
+ ensure_ascii=False,
+ indent=2,
+ ),
+ )
+ except Exception as e:
+ return _tool_response(
+ json.dumps(
+ {"ok": False, "error": f"Screenshot failed: {e!s}"},
+ ensure_ascii=False,
+ indent=2,
+ ),
+ )
+
+
+async def _action_click( # pylint: disable=too-many-branches
+ page_id: str,
+ selector: str,
+ ref: str = "",
+ element: str = "", # pylint: disable=unused-argument
+ wait: int = 0,
+ double_click: bool = False,
+ button: str = "left",
+ modifiers_json: str = "",
+ frame_selector: str = "",
+) -> ToolResponse:
+ ref = (ref or "").strip()
+ selector = (selector or "").strip()
+ if not ref and not selector:
+ return _tool_response(
+ json.dumps(
+ {"ok": False, "error": "selector or ref required for click"},
+ ensure_ascii=False,
+ indent=2,
+ ),
+ )
+ page = _get_page(page_id)
+ if not page:
+ return _tool_response(
+ json.dumps(
+ {"ok": False, "error": f"Page '{page_id}' not found"},
+ ensure_ascii=False,
+ indent=2,
+ ),
+ )
+ try:
+ if wait > 0:
+ await asyncio.sleep(wait / 1000.0)
+ mods = _parse_json_param(modifiers_json, [])
+ if not isinstance(mods, list):
+ mods = []
+ kwargs = {
+ "button": button if button in ("left", "right", "middle") else "left",
+ }
+ if mods:
+ kwargs["modifiers"] = [m for m in mods if m in ("Alt", "Control", "ControlOrMeta", "Meta", "Shift")]
+
+ if _USE_SYNC_PLAYWRIGHT:
+ loop = asyncio.get_event_loop()
+ if ref:
+ locator = _get_locator_by_ref(
+ page,
+ page_id,
+ ref,
+ frame_selector,
+ )
+ if locator is None:
+ return _tool_response(
+ json.dumps(
+ {"ok": False, "error": f"Unknown ref: {ref}"},
+ ensure_ascii=False,
+ indent=2,
+ ),
+ )
+ if double_click:
+ await loop.run_in_executor(
+ _get_executor(),
+ lambda: locator.dblclick(**kwargs),
+ )
+ else:
+ await loop.run_in_executor(
+ _get_executor(),
+ lambda: locator.click(**kwargs),
+ )
+ else:
+ root = _get_root(page, page_id, frame_selector)
+ locator = root.locator(selector).first
+ if double_click:
+ await loop.run_in_executor(
+ _get_executor(),
+ lambda: locator.dblclick(**kwargs),
+ )
+ else:
+ await loop.run_in_executor(
+ _get_executor(),
+ lambda: locator.click(**kwargs),
+ )
+ else:
+ # Standard async mode
+ if ref:
+ locator = _get_locator_by_ref(
+ page,
+ page_id,
+ ref,
+ frame_selector,
+ )
+ if locator is None:
+ return _tool_response(
+ json.dumps(
+ {"ok": False, "error": f"Unknown ref: {ref}"},
+ ensure_ascii=False,
+ indent=2,
+ ),
+ )
+ if double_click:
+ await locator.dblclick(**kwargs)
+ else:
+ await locator.click(**kwargs)
+ else:
+ root = _get_root(page, page_id, frame_selector)
+ locator = root.locator(selector).first
+ if double_click:
+ await locator.dblclick(**kwargs)
+ else:
+ await locator.click(**kwargs)
+
+ return _tool_response(
+ json.dumps(
+ {"ok": True, "message": f"Clicked {ref or selector}"},
+ ensure_ascii=False,
+ indent=2,
+ ),
+ )
+ except Exception as e:
+ return _tool_response(
+ json.dumps(
+ {"ok": False, "error": f"Click failed: {e!s}"},
+ ensure_ascii=False,
+ indent=2,
+ ),
+ )
+
+
+async def _action_type(
+ page_id: str,
+ selector: str,
+ ref: str = "",
+ element: str = "", # pylint: disable=unused-argument
+ text: str = "",
+ submit: bool = False,
+ slowly: bool = False,
+ frame_selector: str = "",
+) -> ToolResponse:
+ ref = (ref or "").strip()
+ selector = (selector or "").strip()
+ if not ref and not selector:
+ return _tool_response(
+ json.dumps(
+ {"ok": False, "error": "selector or ref required for type"},
+ ensure_ascii=False,
+ indent=2,
+ ),
+ )
+ page = _get_page(page_id)
+ if not page:
+ return _tool_response(
+ json.dumps(
+ {"ok": False, "error": f"Page '{page_id}' not found"},
+ ensure_ascii=False,
+ indent=2,
+ ),
+ )
+ try:
+ if ref:
+ locator = _get_locator_by_ref(page, page_id, ref, frame_selector)
+ if locator is None:
+ return _tool_response(
+ json.dumps(
+ {"ok": False, "error": f"Unknown ref: {ref}"},
+ ensure_ascii=False,
+ indent=2,
+ ),
+ )
+ if _USE_SYNC_PLAYWRIGHT:
+ loop = asyncio.get_event_loop()
+ if slowly:
+ await loop.run_in_executor(
+ _get_executor(),
+ lambda: locator.press_sequentially(text or ""),
+ )
+ else:
+ await loop.run_in_executor(
+ _get_executor(),
+ lambda: locator.fill(text or ""),
+ )
+ if submit:
+ await loop.run_in_executor(
+ _get_executor(),
+ lambda: locator.press("Enter"),
+ )
+ else:
+ if slowly:
+ await locator.press_sequentially(text or "")
+ else:
+ await locator.fill(text or "")
+ if submit:
+ await locator.press("Enter")
+ else:
+ root = _get_root(page, page_id, frame_selector)
+ loc = root.locator(selector).first
+ if _USE_SYNC_PLAYWRIGHT:
+ loop = asyncio.get_event_loop()
+ if slowly:
+ await loop.run_in_executor(
+ _get_executor(),
+ lambda: loc.press_sequentially(text or ""),
+ )
+ else:
+ await loop.run_in_executor(
+ _get_executor(),
+ lambda: loc.fill(text or ""),
+ )
+ if submit:
+ await loop.run_in_executor(
+ _get_executor(),
+ lambda: loc.press("Enter"),
+ )
+ else:
+ if slowly:
+ await loc.press_sequentially(text or "")
+ else:
+ await loc.fill(text or "")
+ if submit:
+ await loc.press("Enter")
+ return _tool_response(
+ json.dumps(
+ {"ok": True, "message": f"Typed into {ref or selector}"},
+ ensure_ascii=False,
+ indent=2,
+ ),
+ )
+ except Exception as e:
+ return _tool_response(
+ json.dumps(
+ {"ok": False, "error": f"Type failed: {e!s}"},
+ ensure_ascii=False,
+ indent=2,
+ ),
+ )
+
+
+async def _action_eval(page_id: str, code: str) -> ToolResponse:
+ code = (code or "").strip()
+ if not code:
+ return _tool_response(
+ json.dumps(
+ {"ok": False, "error": "code required for eval"},
+ ensure_ascii=False,
+ indent=2,
+ ),
+ )
+ page = _get_page(page_id)
+ if not page:
+ return _tool_response(
+ json.dumps(
+ {"ok": False, "error": f"Page '{page_id}' not found"},
+ ensure_ascii=False,
+ indent=2,
+ ),
+ )
+ try:
+ if code.strip().startswith("(") or code.strip().startswith("function"):
+ if _USE_SYNC_PLAYWRIGHT:
+ result = await _run_sync(page.evaluate, code)
+ else:
+ result = await page.evaluate(code)
+ else:
+ if _USE_SYNC_PLAYWRIGHT:
+ result = await _run_sync(
+ page.evaluate,
+ f"() => {{ return ({code}); }}",
+ )
+ else:
+ result = await page.evaluate(f"() => {{ return ({code}); }}")
+ try:
+ out = json.dumps(
+ {"ok": True, "result": result},
+ ensure_ascii=False,
+ indent=2,
+ )
+ except TypeError:
+ out = json.dumps(
+ {"ok": True, "result": str(result)},
+ ensure_ascii=False,
+ indent=2,
+ )
+ return _tool_response(out)
+ except Exception as e:
+ return _tool_response(
+ json.dumps(
+ {"ok": False, "error": f"Eval failed: {e!s}"},
+ ensure_ascii=False,
+ indent=2,
+ ),
+ )
+
+
+async def _action_pdf(page_id: str, path: str) -> ToolResponse:
+ path = (path or "page.pdf").strip() or "page.pdf"
+ page = _get_page(page_id)
+ if not page:
+ return _tool_response(
+ json.dumps(
+ {"ok": False, "error": f"Page '{page_id}' not found"},
+ ensure_ascii=False,
+ indent=2,
+ ),
+ )
+ try:
+ if _USE_SYNC_PLAYWRIGHT:
+ await _run_sync(page.pdf, path=path)
+ else:
+ await page.pdf(path=path)
+ return _tool_response(
+ json.dumps(
+ {"ok": True, "message": f"PDF saved to {path}", "path": path},
+ ensure_ascii=False,
+ indent=2,
+ ),
+ )
+ except Exception as e:
+ return _tool_response(
+ json.dumps(
+ {"ok": False, "error": f"PDF failed: {e!s}"},
+ ensure_ascii=False,
+ indent=2,
+ ),
+ )
+
+
+async def _action_close(page_id: str) -> ToolResponse:
+ page = _get_page(page_id)
+ if not page:
+ return _tool_response(
+ json.dumps(
+ {"ok": False, "error": f"Page '{page_id}' not found"},
+ ensure_ascii=False,
+ indent=2,
+ ),
+ )
+ try:
+ if _USE_SYNC_PLAYWRIGHT:
+ await _run_sync(page.close)
+ else:
+ await page.close()
+ del _state["pages"][page_id]
+ for key in (
+ "refs",
+ "refs_frame",
+ "console_logs",
+ "network_requests",
+ "pending_dialogs",
+ "pending_file_choosers",
+ ):
+ _state[key].pop(page_id, None)
+ if _state.get("current_page_id") == page_id:
+ remaining = list(_state["pages"].keys())
+ _state["current_page_id"] = remaining[0] if remaining else None
+ return _tool_response(
+ json.dumps(
+ {"ok": True, "message": f"Closed page '{page_id}'"},
+ ensure_ascii=False,
+ indent=2,
+ ),
+ )
+ except Exception as e:
+ return _tool_response(
+ json.dumps(
+ {"ok": False, "error": f"Close failed: {e!s}"},
+ ensure_ascii=False,
+ indent=2,
+ ),
+ )
+
+
+async def _action_snapshot(
+ page_id: str,
+ filename: str,
+ frame_selector: str = "",
+) -> ToolResponse:
+ page = _get_page(page_id)
+ if not page:
+ return _tool_response(
+ json.dumps(
+ {"ok": False, "error": f"Page '{page_id}' not found"},
+ ensure_ascii=False,
+ indent=2,
+ ),
+ )
+ try:
+ if _USE_SYNC_PLAYWRIGHT:
+ # Hybrid mode: execute in thread pool
+ loop = asyncio.get_event_loop()
+ root = _get_root(page, page_id, frame_selector)
+ locator = root.locator(":root")
+ raw = await loop.run_in_executor(
+ _get_executor(),
+ lambda: locator.aria_snapshot(), # pylint: disable=unnecessary-lambda
+ )
+ else:
+ root = _get_root(page, page_id, frame_selector)
+ locator = root.locator(":root")
+ raw = await locator.aria_snapshot()
+
+ raw_str = str(raw) if raw is not None else ""
+ snapshot, refs = build_role_snapshot_from_aria(
+ raw_str,
+ interactive=False,
+ compact=False,
+ )
+ _state["refs"][page_id] = refs
+ _state["refs_frame"][page_id] = frame_selector.strip() if frame_selector else ""
+ out = {
+ "ok": True,
+ "snapshot": snapshot,
+ "refs": list(refs.keys()),
+ "url": page.url,
+ }
+ if frame_selector and frame_selector.strip():
+ out["frame_selector"] = frame_selector.strip()
+ if filename and filename.strip():
+ with open(filename.strip(), "w", encoding="utf-8") as f:
+ f.write(snapshot)
+ out["filename"] = filename.strip()
+ return _tool_response(json.dumps(out, ensure_ascii=False, indent=2))
+ except Exception as e:
+ return _tool_response(
+ json.dumps(
+ {"ok": False, "error": f"Snapshot failed: {e!s}"},
+ ensure_ascii=False,
+ indent=2,
+ ),
+ )
+
+
+async def _action_navigate_back(page_id: str) -> ToolResponse:
+ page = _get_page(page_id)
+ if not page:
+ return _tool_response(
+ json.dumps(
+ {"ok": False, "error": f"Page '{page_id}' not found"},
+ ensure_ascii=False,
+ indent=2,
+ ),
+ )
+ try:
+ if _USE_SYNC_PLAYWRIGHT:
+ await _run_sync(page.go_back)
+ else:
+ await page.go_back()
+ return _tool_response(
+ json.dumps(
+ {"ok": True, "message": "Navigated back", "url": page.url},
+ ensure_ascii=False,
+ indent=2,
+ ),
+ )
+ except Exception as e:
+ return _tool_response(
+ json.dumps(
+ {"ok": False, "error": f"Navigate back failed: {e!s}"},
+ ensure_ascii=False,
+ indent=2,
+ ),
+ )
+
+
+async def _action_evaluate(
+ page_id: str,
+ code: str,
+ ref: str = "",
+ element: str = "", # pylint: disable=unused-argument
+ frame_selector: str = "",
+) -> ToolResponse:
+ code = (code or "").strip()
+ if not code:
+ return _tool_response(
+ json.dumps(
+ {"ok": False, "error": "code required for evaluate"},
+ ensure_ascii=False,
+ indent=2,
+ ),
+ )
+ page = _get_page(page_id)
+ if not page:
+ return _tool_response(
+ json.dumps(
+ {"ok": False, "error": f"Page '{page_id}' not found"},
+ ensure_ascii=False,
+ indent=2,
+ ),
+ )
+ try:
+ if ref and ref.strip():
+ locator = _get_locator_by_ref(
+ page,
+ page_id,
+ ref.strip(),
+ frame_selector,
+ )
+ if locator is None:
+ return _tool_response(
+ json.dumps(
+ {"ok": False, "error": f"Unknown ref: {ref}"},
+ ensure_ascii=False,
+ indent=2,
+ ),
+ )
+ if _USE_SYNC_PLAYWRIGHT:
+ result = await _run_sync(locator.evaluate, code)
+ else:
+ result = await locator.evaluate(code)
+ else:
+ if code.strip().startswith("(") or code.strip().startswith(
+ "function",
+ ):
+ if _USE_SYNC_PLAYWRIGHT:
+ result = await _run_sync(page.evaluate, code)
+ else:
+ result = await page.evaluate(code)
+ else:
+ if _USE_SYNC_PLAYWRIGHT:
+ result = await _run_sync(
+ page.evaluate,
+ f"() => {{ return ({code}); }}",
+ )
+ else:
+ result = await page.evaluate(
+ f"() => {{ return ({code}); }}",
+ )
+ try:
+ out = json.dumps(
+ {"ok": True, "result": result},
+ ensure_ascii=False,
+ indent=2,
+ )
+ except TypeError:
+ out = json.dumps(
+ {"ok": True, "result": str(result)},
+ ensure_ascii=False,
+ indent=2,
+ )
+ return _tool_response(out)
+ except Exception as e:
+ return _tool_response(
+ json.dumps(
+ {"ok": False, "error": f"Evaluate failed: {e!s}"},
+ ensure_ascii=False,
+ indent=2,
+ ),
+ )
+
+
+async def _action_resize(
+ page_id: str,
+ width: int,
+ height: int,
+) -> ToolResponse:
+ if width <= 0 or height <= 0:
+ return _tool_response(
+ json.dumps(
+ {"ok": False, "error": "width and height must be positive"},
+ ensure_ascii=False,
+ indent=2,
+ ),
+ )
+ page = _get_page(page_id)
+ if not page:
+ return _tool_response(
+ json.dumps(
+ {"ok": False, "error": f"Page '{page_id}' not found"},
+ ensure_ascii=False,
+ indent=2,
+ ),
+ )
+ try:
+ if _USE_SYNC_PLAYWRIGHT:
+ await _run_sync(
+ page.set_viewport_size,
+ {"width": width, "height": height},
+ )
+ else:
+ await page.set_viewport_size({"width": width, "height": height})
+ return _tool_response(
+ json.dumps(
+ {"ok": True, "message": f"Resized to {width}x{height}"},
+ ensure_ascii=False,
+ indent=2,
+ ),
+ )
+ except Exception as e:
+ return _tool_response(
+ json.dumps(
+ {"ok": False, "error": f"Resize failed: {e!s}"},
+ ensure_ascii=False,
+ indent=2,
+ ),
+ )
+
+
+async def _action_console_messages(
+ page_id: str,
+ level: str,
+ filename: str,
+) -> ToolResponse:
+ level = (level or "info").strip().lower()
+ order = ("error", "warning", "info", "debug")
+ idx = order.index(level) if level in order else 2
+ page = _get_page(page_id)
+ if not page:
+ return _tool_response(
+ json.dumps(
+ {"ok": False, "error": f"Page '{page_id}' not found"},
+ ensure_ascii=False,
+ indent=2,
+ ),
+ )
+ logs = _state["console_logs"].get(page_id, [])
+ filtered = [m for m in logs if order.index(m["level"]) <= idx] if level in order else logs
+ lines = [f"[{m['level']}] {m['text']}" for m in filtered]
+ text = "\n".join(lines)
+ if filename and filename.strip():
+ with open(filename.strip(), "w", encoding="utf-8") as f:
+ f.write(text)
+ return _tool_response(
+ json.dumps(
+ {
+ "ok": True,
+ "message": f"Console messages saved to {filename}",
+ "filename": filename.strip(),
+ },
+ ensure_ascii=False,
+ indent=2,
+ ),
+ )
+ return _tool_response(
+ json.dumps(
+ {"ok": True, "messages": filtered, "text": text},
+ ensure_ascii=False,
+ indent=2,
+ ),
+ )
+
+
+async def _action_handle_dialog(
+ page_id: str,
+ accept: bool,
+ prompt_text: str,
+) -> ToolResponse:
+ page = _get_page(page_id)
+ if not page:
+ return _tool_response(
+ json.dumps(
+ {"ok": False, "error": f"Page '{page_id}' not found"},
+ ensure_ascii=False,
+ indent=2,
+ ),
+ )
+ dialogs = _state["pending_dialogs"].get(page_id, [])
+ if not dialogs:
+ return _tool_response(
+ json.dumps(
+ {"ok": False, "error": "No pending dialog"},
+ ensure_ascii=False,
+ indent=2,
+ ),
+ )
+ try:
+ dialog = dialogs.pop(0)
+ if accept:
+ if prompt_text and hasattr(dialog, "accept"):
+ if _USE_SYNC_PLAYWRIGHT:
+ await _run_sync(dialog.accept, prompt_text)
+ else:
+ await dialog.accept(prompt_text)
+ else:
+ if _USE_SYNC_PLAYWRIGHT:
+ await _run_sync(dialog.accept)
+ else:
+ await dialog.accept()
+ else:
+ if _USE_SYNC_PLAYWRIGHT:
+ await _run_sync(dialog.dismiss)
+ else:
+ await dialog.dismiss()
+ return _tool_response(
+ json.dumps(
+ {"ok": True, "message": "Dialog handled"},
+ ensure_ascii=False,
+ indent=2,
+ ),
+ )
+ except Exception as e:
+ return _tool_response(
+ json.dumps(
+ {"ok": False, "error": f"Handle dialog failed: {e!s}"},
+ ensure_ascii=False,
+ indent=2,
+ ),
+ )
+
+
+async def _action_file_upload(page_id: str, paths_json: str) -> ToolResponse:
+ page = _get_page(page_id)
+ if not page:
+ return _tool_response(
+ json.dumps(
+ {"ok": False, "error": f"Page '{page_id}' not found"},
+ ensure_ascii=False,
+ indent=2,
+ ),
+ )
+ paths = _parse_json_param(paths_json, [])
+ if not isinstance(paths, list):
+ paths = []
+ try:
+ choosers = _state["pending_file_choosers"].get(page_id, [])
+ if not choosers:
+ return _tool_response(
+ json.dumps(
+ {
+ "ok": False,
+ "error": "No chooser. Click upload then file_upload.",
+ },
+ ensure_ascii=False,
+ indent=2,
+ ),
+ )
+ chooser = choosers.pop(0)
+ if paths:
+ if _USE_SYNC_PLAYWRIGHT:
+ await _run_sync(chooser.set_files, paths)
+ else:
+ await chooser.set_files(paths)
+ return _tool_response(
+ json.dumps(
+ {"ok": True, "message": f"Uploaded {len(paths)} file(s)"},
+ ensure_ascii=False,
+ indent=2,
+ ),
+ )
+ if _USE_SYNC_PLAYWRIGHT:
+ await _run_sync(chooser.set_files, [])
+ else:
+ await chooser.set_files([])
+ return _tool_response(
+ json.dumps(
+ {"ok": True, "message": "File chooser cancelled"},
+ ensure_ascii=False,
+ indent=2,
+ ),
+ )
+ except Exception as e:
+ return _tool_response(
+ json.dumps(
+ {"ok": False, "error": f"File upload failed: {e!s}"},
+ ensure_ascii=False,
+ indent=2,
+ ),
+ )
+
+
+async def _action_fill_form(page_id: str, fields_json: str) -> ToolResponse:
+ page = _get_page(page_id)
+ if not page:
+ return _tool_response(
+ json.dumps(
+ {"ok": False, "error": f"Page '{page_id}' not found"},
+ ensure_ascii=False,
+ indent=2,
+ ),
+ )
+ fields = _parse_json_param(fields_json, [])
+ if not isinstance(fields, list) or not fields:
+ return _tool_response(
+ json.dumps(
+ {"ok": False, "error": "fields required (JSON array)"},
+ ensure_ascii=False,
+ indent=2,
+ ),
+ )
+ refs = _get_refs(page_id)
+ # Use last snapshot's frame so fill_form works after iframe snapshot
+ frame = _state["refs_frame"].get(page_id, "")
+ try:
+ for f in fields:
+ ref = (f.get("ref") or "").strip()
+ if not ref or ref not in refs:
+ continue
+ locator = _get_locator_by_ref(page, page_id, ref, frame)
+ if locator is None:
+ continue
+ field_type = (f.get("type") or "textbox").lower()
+ value = f.get("value")
+ if field_type == "checkbox":
+ if isinstance(value, str):
+ value = value.strip().lower() in ("true", "1", "yes")
+ if _USE_SYNC_PLAYWRIGHT:
+ await _run_sync(locator.set_checked, bool(value))
+ else:
+ await locator.set_checked(bool(value))
+ elif field_type == "radio":
+ if _USE_SYNC_PLAYWRIGHT:
+ await _run_sync(locator.set_checked, True)
+ else:
+ await locator.set_checked(True)
+ elif field_type == "combobox":
+ if _USE_SYNC_PLAYWRIGHT:
+ await _run_sync(
+ locator.select_option,
+ label=value if isinstance(value, str) else None,
+ value=value,
+ )
+ else:
+ await locator.select_option(
+ label=value if isinstance(value, str) else None,
+ value=value,
+ )
+ elif field_type == "slider":
+ if _USE_SYNC_PLAYWRIGHT:
+ await _run_sync(locator.fill, str(value))
+ else:
+ await locator.fill(str(value))
+ else:
+ if _USE_SYNC_PLAYWRIGHT:
+ await _run_sync(
+ locator.fill,
+ str(value) if value is not None else "",
+ )
+ else:
+ await locator.fill(str(value) if value is not None else "")
+ return _tool_response(
+ json.dumps(
+ {"ok": True, "message": f"Filled {len(fields)} field(s)"},
+ ensure_ascii=False,
+ indent=2,
+ ),
+ )
+ except Exception as e:
+ return _tool_response(
+ json.dumps(
+ {"ok": False, "error": f"Fill form failed: {e!s}"},
+ ensure_ascii=False,
+ indent=2,
+ ),
+ )
+
+
+def _run_playwright_install() -> None:
+ """Run playwright install in a blocking way (for use in thread)."""
+ subprocess.run(
+ [sys.executable, "-m", "playwright", "install"],
+ check=True,
+ capture_output=True,
+ text=True,
+ timeout=600, # 10 minutes max
+ )
+
+
+async def _action_install() -> ToolResponse:
+ """Install Playwright browsers. If a system Chrome/Chromium/Edge is found,
+ use it and skip download. On macOS with no Chromium, use Safari (WebKit)
+ so no download is needed. Only run playwright install when necessary.
+ """
+ exe = _chromium_executable_path()
+ if exe:
+ return _tool_response(
+ json.dumps(
+ {
+ "ok": True,
+ "message": f"Using system browser (no download): {exe}",
+ },
+ ensure_ascii=False,
+ indent=2,
+ ),
+ )
+ if _use_webkit_fallback():
+ return _tool_response(
+ json.dumps(
+ {
+ "ok": True,
+ "message": "On macOS using Safari (WebKit); no browser download needed.",
+ },
+ ensure_ascii=False,
+ indent=2,
+ ),
+ )
+ try:
+ await asyncio.to_thread(_run_playwright_install)
+ return _tool_response(
+ json.dumps(
+ {"ok": True, "message": "Browser installed"},
+ ensure_ascii=False,
+ indent=2,
+ ),
+ )
+ except subprocess.TimeoutExpired:
+ return _tool_response(
+ json.dumps(
+ {
+ "ok": False,
+ "error": "Browser install timed out (10 min). Run manually in terminal: "
+ f"{sys.executable!s} -m playwright install",
+ },
+ ensure_ascii=False,
+ indent=2,
+ ),
+ )
+ except Exception as e:
+ return _tool_response(
+ json.dumps(
+ {
+ "ok": False,
+ "error": f"Install failed: {e!s}. Install manually: "
+ f"{sys.executable!s} -m pip install playwright && "
+ f"{sys.executable!s} -m playwright install",
+ },
+ ensure_ascii=False,
+ indent=2,
+ ),
+ )
+
+
+async def _action_press_key(page_id: str, key: str) -> ToolResponse:
+ key = (key or "").strip()
+ if not key:
+ return _tool_response(
+ json.dumps(
+ {"ok": False, "error": "key required for press_key"},
+ ensure_ascii=False,
+ indent=2,
+ ),
+ )
+ page = _get_page(page_id)
+ if not page:
+ return _tool_response(
+ json.dumps(
+ {"ok": False, "error": f"Page '{page_id}' not found"},
+ ensure_ascii=False,
+ indent=2,
+ ),
+ )
+ try:
+ if _USE_SYNC_PLAYWRIGHT:
+ await _run_sync(page.keyboard.press, key)
+ else:
+ await page.keyboard.press(key)
+ return _tool_response(
+ json.dumps(
+ {"ok": True, "message": f"Pressed key {key}"},
+ ensure_ascii=False,
+ indent=2,
+ ),
+ )
+ except Exception as e:
+ return _tool_response(
+ json.dumps(
+ {"ok": False, "error": f"Press key failed: {e!s}"},
+ ensure_ascii=False,
+ indent=2,
+ ),
+ )
+
+
+async def _action_network_requests(
+ page_id: str,
+ include_static: bool,
+ filename: str,
+) -> ToolResponse:
+ page = _get_page(page_id)
+ if not page:
+ return _tool_response(
+ json.dumps(
+ {"ok": False, "error": f"Page '{page_id}' not found"},
+ ensure_ascii=False,
+ indent=2,
+ ),
+ )
+ requests = _state["network_requests"].get(page_id, [])
+ if not include_static:
+ static = ("image", "stylesheet", "font", "media")
+ requests = [r for r in requests if r.get("resourceType") not in static]
+ lines = [f"{r.get('method', '')} {r.get('url', '')} {r.get('status', '')}" for r in requests]
+ text = "\n".join(lines)
+ if filename and filename.strip():
+ with open(filename.strip(), "w", encoding="utf-8") as f:
+ f.write(text)
+ return _tool_response(
+ json.dumps(
+ {
+ "ok": True,
+ "message": f"Network requests saved to {filename}",
+ "filename": filename.strip(),
+ },
+ ensure_ascii=False,
+ indent=2,
+ ),
+ )
+ return _tool_response(
+ json.dumps(
+ {"ok": True, "requests": requests, "text": text},
+ ensure_ascii=False,
+ indent=2,
+ ),
+ )
+
+
+async def _action_run_code(page_id: str, code: str) -> ToolResponse:
+ """Run JS in page (like eval). Use evaluate for element (ref)."""
+ code = (code or "").strip()
+ if not code:
+ return _tool_response(
+ json.dumps(
+ {"ok": False, "error": "code required for run_code"},
+ ensure_ascii=False,
+ indent=2,
+ ),
+ )
+ page = _get_page(page_id)
+ if not page:
+ return _tool_response(
+ json.dumps(
+ {"ok": False, "error": f"Page '{page_id}' not found"},
+ ensure_ascii=False,
+ indent=2,
+ ),
+ )
+ try:
+ if code.strip().startswith("(") or code.strip().startswith("function"):
+ if _USE_SYNC_PLAYWRIGHT:
+ result = await _run_sync(page.evaluate, code)
+ else:
+ result = await page.evaluate(code)
+ else:
+ if _USE_SYNC_PLAYWRIGHT:
+ result = await _run_sync(
+ page.evaluate,
+ f"() => {{ return ({code}); }}",
+ )
+ else:
+ result = await page.evaluate(f"() => {{ return ({code}); }}")
+ try:
+ out = json.dumps(
+ {"ok": True, "result": result},
+ ensure_ascii=False,
+ indent=2,
+ )
+ except TypeError:
+ out = json.dumps(
+ {"ok": True, "result": str(result)},
+ ensure_ascii=False,
+ indent=2,
+ )
+ return _tool_response(out)
+ except Exception as e:
+ return _tool_response(
+ json.dumps(
+ {"ok": False, "error": f"Run code failed: {e!s}"},
+ ensure_ascii=False,
+ indent=2,
+ ),
+ )
+
+
+async def _action_drag(
+ page_id: str,
+ start_ref: str,
+ end_ref: str,
+ start_selector: str = "",
+ end_selector: str = "",
+ start_element: str = "", # pylint: disable=unused-argument
+ end_element: str = "", # pylint: disable=unused-argument
+ frame_selector: str = "",
+) -> ToolResponse:
+ start_ref = (start_ref or "").strip()
+ end_ref = (end_ref or "").strip()
+ start_selector = (start_selector or "").strip()
+ end_selector = (end_selector or "").strip()
+ use_refs = bool(start_ref and end_ref)
+ use_selectors = bool(start_selector and end_selector)
+ if not use_refs and not use_selectors:
+ return _tool_response(
+ json.dumps(
+ {
+ "ok": False,
+ "error": ("drag needs (start_ref,end_ref) or (start_sel,end_sel)"),
+ },
+ ensure_ascii=False,
+ indent=2,
+ ),
+ )
+ page = _get_page(page_id)
+ if not page:
+ return _tool_response(
+ json.dumps(
+ {"ok": False, "error": f"Page '{page_id}' not found"},
+ ensure_ascii=False,
+ indent=2,
+ ),
+ )
+ try:
+ root = _get_root(page, page_id, frame_selector)
+ if use_refs:
+ start_locator = _get_locator_by_ref(
+ page,
+ page_id,
+ start_ref,
+ frame_selector,
+ )
+ end_locator = _get_locator_by_ref(
+ page,
+ page_id,
+ end_ref,
+ frame_selector,
+ )
+ if start_locator is None or end_locator is None:
+ return _tool_response(
+ json.dumps(
+ {"ok": False, "error": "Unknown ref for drag"},
+ ensure_ascii=False,
+ indent=2,
+ ),
+ )
+ else:
+ start_locator = root.locator(start_selector).first
+ end_locator = root.locator(end_selector).first
+ if _USE_SYNC_PLAYWRIGHT:
+ await _run_sync(start_locator.drag_to, end_locator)
+ else:
+ await start_locator.drag_to(end_locator)
+ return _tool_response(
+ json.dumps(
+ {"ok": True, "message": "Drag completed"},
+ ensure_ascii=False,
+ indent=2,
+ ),
+ )
+ except Exception as e:
+ return _tool_response(
+ json.dumps(
+ {"ok": False, "error": f"Drag failed: {e!s}"},
+ ensure_ascii=False,
+ indent=2,
+ ),
+ )
+
+
+async def _action_hover(
+ page_id: str,
+ ref: str = "",
+ element: str = "", # pylint: disable=unused-argument
+ selector: str = "",
+ frame_selector: str = "",
+) -> ToolResponse:
+ ref = (ref or "").strip()
+ selector = (selector or "").strip()
+ if not ref and not selector:
+ return _tool_response(
+ json.dumps(
+ {"ok": False, "error": "hover requires ref or selector"},
+ ensure_ascii=False,
+ indent=2,
+ ),
+ )
+ page = _get_page(page_id)
+ if not page:
+ return _tool_response(
+ json.dumps(
+ {"ok": False, "error": f"Page '{page_id}' not found"},
+ ensure_ascii=False,
+ indent=2,
+ ),
+ )
+ try:
+ if ref:
+ locator = _get_locator_by_ref(page, page_id, ref, frame_selector)
+ if locator is None:
+ return _tool_response(
+ json.dumps(
+ {"ok": False, "error": f"Unknown ref: {ref}"},
+ ensure_ascii=False,
+ indent=2,
+ ),
+ )
+ else:
+ root = _get_root(page, page_id, frame_selector)
+ locator = root.locator(selector).first
+ if _USE_SYNC_PLAYWRIGHT:
+ await _run_sync(locator.hover)
+ else:
+ await locator.hover()
+ return _tool_response(
+ json.dumps(
+ {"ok": True, "message": f"Hovered {ref or selector}"},
+ ensure_ascii=False,
+ indent=2,
+ ),
+ )
+ except Exception as e:
+ return _tool_response(
+ json.dumps(
+ {"ok": False, "error": f"Hover failed: {e!s}"},
+ ensure_ascii=False,
+ indent=2,
+ ),
+ )
+
+
+async def _action_select_option(
+ page_id: str,
+ ref: str = "",
+ element: str = "", # pylint: disable=unused-argument
+ values_json: str = "",
+ frame_selector: str = "",
+) -> ToolResponse:
+ ref = (ref or "").strip()
+ values = _parse_json_param(values_json, [])
+ if not isinstance(values, list):
+ values = [values] if values is not None else []
+ if not ref:
+ return _tool_response(
+ json.dumps(
+ {"ok": False, "error": "ref required for select_option"},
+ ensure_ascii=False,
+ indent=2,
+ ),
+ )
+ if not values:
+ return _tool_response(
+ json.dumps(
+ {
+ "ok": False,
+ "error": "values required (JSON array or comma-separated)",
+ },
+ ensure_ascii=False,
+ indent=2,
+ ),
+ )
+ page = _get_page(page_id)
+ if not page:
+ return _tool_response(
+ json.dumps(
+ {"ok": False, "error": f"Page '{page_id}' not found"},
+ ensure_ascii=False,
+ indent=2,
+ ),
+ )
+ try:
+ locator = _get_locator_by_ref(page, page_id, ref, frame_selector)
+ if locator is None:
+ return _tool_response(
+ json.dumps(
+ {"ok": False, "error": f"Unknown ref: {ref}"},
+ ensure_ascii=False,
+ indent=2,
+ ),
+ )
+ if _USE_SYNC_PLAYWRIGHT:
+ await _run_sync(locator.select_option, value=values)
+ else:
+ await locator.select_option(value=values)
+ return _tool_response(
+ json.dumps(
+ {"ok": True, "message": f"Selected {values}"},
+ ensure_ascii=False,
+ indent=2,
+ ),
+ )
+ except Exception as e:
+ return _tool_response(
+ json.dumps(
+ {"ok": False, "error": f"Select option failed: {e!s}"},
+ ensure_ascii=False,
+ indent=2,
+ ),
+ )
+
+
+async def _action_tabs( # pylint: disable=too-many-return-statements
+ page_id: str,
+ tab_action: str,
+ index: int,
+) -> ToolResponse:
+ tab_action = (tab_action or "").strip().lower()
+ if not tab_action:
+ return _tool_response(
+ json.dumps(
+ {
+ "ok": False,
+ "error": "tab_action required (list, new, close, select)",
+ },
+ ensure_ascii=False,
+ indent=2,
+ ),
+ )
+ pages = _state["pages"]
+ page_ids = list(pages.keys())
+ if tab_action == "list":
+ return _tool_response(
+ json.dumps(
+ {"ok": True, "tabs": page_ids, "count": len(page_ids)},
+ ensure_ascii=False,
+ indent=2,
+ ),
+ )
+ if tab_action == "new":
+ if _USE_SYNC_PLAYWRIGHT:
+ if not _state["_sync_context"]:
+ ok = await _ensure_browser()
+ if not ok:
+ err = _state.get("_last_browser_error") or "Browser not started"
+ return _tool_response(
+ json.dumps(
+ {"ok": False, "error": err},
+ ensure_ascii=False,
+ indent=2,
+ ),
+ )
+ else:
+ if not _state["context"]:
+ ok = await _ensure_browser()
+ if not ok:
+ err = _state.get("_last_browser_error") or "Browser not started"
+ return _tool_response(
+ json.dumps(
+ {"ok": False, "error": err},
+ ensure_ascii=False,
+ indent=2,
+ ),
+ )
+ try:
+ if _USE_SYNC_PLAYWRIGHT:
+ page = await _run_sync(_state["_sync_context"].new_page)
+ else:
+ page = await _state["context"].new_page()
+ new_id = _next_page_id()
+ _state["refs"][new_id] = {}
+ _state["console_logs"][new_id] = []
+ _state["network_requests"][new_id] = []
+ _state["pending_dialogs"][new_id] = []
+ _attach_page_listeners(page, new_id)
+ _state["pages"][new_id] = page
+ _state["current_page_id"] = new_id
+ return _tool_response(
+ json.dumps(
+ {
+ "ok": True,
+ "page_id": new_id,
+ "tabs": list(_state["pages"].keys()),
+ },
+ ensure_ascii=False,
+ indent=2,
+ ),
+ )
+ except Exception as e:
+ return _tool_response(
+ json.dumps(
+ {"ok": False, "error": f"New tab failed: {e!s}"},
+ ensure_ascii=False,
+ indent=2,
+ ),
+ )
+ if tab_action == "close":
+ target_id = page_ids[index] if 0 <= index < len(page_ids) else page_id
+ return await _action_close(target_id)
+ if tab_action == "select":
+ target_id = page_ids[index] if 0 <= index < len(page_ids) else page_id
+ _state["current_page_id"] = target_id
+ return _tool_response(
+ json.dumps(
+ {
+ "ok": True,
+ "message": f"Use page_id={target_id} for later actions",
+ "page_id": target_id,
+ },
+ ensure_ascii=False,
+ indent=2,
+ ),
+ )
+ return _tool_response(
+ json.dumps(
+ {"ok": False, "error": f"Unknown tab_action: {tab_action}"},
+ ensure_ascii=False,
+ indent=2,
+ ),
+ )
+
+
+async def _action_wait_for(
+ page_id: str,
+ wait_time: float,
+ text: str,
+ text_gone: str,
+) -> ToolResponse:
+ page = _get_page(page_id)
+ if not page:
+ return _tool_response(
+ json.dumps(
+ {"ok": False, "error": f"Page '{page_id}' not found"},
+ ensure_ascii=False,
+ indent=2,
+ ),
+ )
+ try:
+ if wait_time and wait_time > 0:
+ await asyncio.sleep(wait_time)
+ text = (text or "").strip()
+ text_gone = (text_gone or "").strip()
+ if text:
+ locator = page.get_by_text(text)
+ if _USE_SYNC_PLAYWRIGHT:
+ await _run_sync(
+ locator.wait_for,
+ state="visible",
+ timeout=30000,
+ )
+ else:
+ await locator.wait_for(
+ state="visible",
+ timeout=30000,
+ )
+ if text_gone:
+ locator = page.get_by_text(text_gone)
+ if _USE_SYNC_PLAYWRIGHT:
+ await _run_sync(
+ locator.wait_for,
+ state="hidden",
+ timeout=30000,
+ )
+ else:
+ await locator.wait_for(
+ state="hidden",
+ timeout=30000,
+ )
+ return _tool_response(
+ json.dumps(
+ {"ok": True, "message": "Wait completed"},
+ ensure_ascii=False,
+ indent=2,
+ ),
+ )
+ except Exception as e:
+ return _tool_response(
+ json.dumps(
+ {"ok": False, "error": f"Wait failed: {e!s}"},
+ ensure_ascii=False,
+ indent=2,
+ ),
+ )
diff --git a/reme/memory/file_based/tools/browser_snapshot.py b/reme/memory/file_based/tools/browser_snapshot.py
new file mode 100644
index 00000000..11ab8885
--- /dev/null
+++ b/reme/memory/file_based/tools/browser_snapshot.py
@@ -0,0 +1,248 @@
+# -*- coding: utf-8 -*-
+"""Build role snapshot + refs from Playwright aria_snapshot."""
+
+import re
+from typing import Any
+
+INTERACTIVE_ROLES = frozenset(
+ {
+ "button",
+ "link",
+ "textbox",
+ "checkbox",
+ "radio",
+ "combobox",
+ "listbox",
+ "menuitem",
+ "menuitemcheckbox",
+ "menuitemradio",
+ "option",
+ "searchbox",
+ "slider",
+ "spinbutton",
+ "switch",
+ "tab",
+ "treeitem",
+ },
+)
+
+CONTENT_ROLES = frozenset(
+ {
+ "heading",
+ "cell",
+ "gridcell",
+ "columnheader",
+ "rowheader",
+ "listitem",
+ "article",
+ "region",
+ "main",
+ "navigation",
+ },
+)
+
+STRUCTURAL_ROLES = frozenset(
+ {
+ "generic",
+ "group",
+ "list",
+ "table",
+ "row",
+ "rowgroup",
+ "grid",
+ "treegrid",
+ "menu",
+ "menubar",
+ "toolbar",
+ "tablist",
+ "tree",
+ "directory",
+ "document",
+ "application",
+ "presentation",
+ "none",
+ },
+)
+
+
+def _get_indent_level(line: str) -> int:
+ m = re.match(r"^(\s*)", line)
+ return int(len(m.group(1)) / 2) if m else 0
+
+
+def _create_tracker() -> dict[str, Any]:
+ counts: dict[str, int] = {}
+ refs_by_key: dict[str, list[str]] = {}
+
+ def get_key(role: str, name: str | None) -> str:
+ return f"{role}:{name or ''}"
+
+ def get_next_index(role: str, name: str | None) -> int:
+ key = get_key(role, name)
+ current = counts.get(key, 0)
+ counts[key] = current + 1
+ return current
+
+ def track_ref(role: str, name: str | None, ref: str) -> None:
+ key = get_key(role, name)
+ refs_by_key.setdefault(key, []).append(ref)
+
+ def get_duplicate_keys() -> set[str]:
+ return {k for k, refs in refs_by_key.items() if len(refs) > 1}
+
+ return {
+ "get_next_index": get_next_index,
+ "track_ref": track_ref,
+ "get_duplicate_keys": get_duplicate_keys,
+ "get_key": get_key,
+ }
+
+
+def _remove_nth_from_non_duplicates(
+ refs: dict[str, dict],
+ tracker: dict,
+) -> None:
+ dup_keys = tracker["get_duplicate_keys"]()
+ for _, data in list(refs.items()):
+ key = tracker["get_key"](data["role"], data.get("name"))
+ if key not in dup_keys and "nth" in data:
+ del data["nth"]
+
+
+def _compact_tree(tree: str) -> str:
+ lines = tree.split("\n")
+ result = []
+ for i, line in enumerate(lines):
+ if "[ref=" in line:
+ result.append(line)
+ continue
+ if ":" in line and not line.rstrip().endswith(":"):
+ result.append(line)
+ continue
+ current_indent = _get_indent_level(line)
+ has_relevant = False
+ for j in range(i + 1, len(lines)):
+ if _get_indent_level(lines[j]) <= current_indent:
+ break
+ if "[ref=" in lines[j]:
+ has_relevant = True
+ break
+ if has_relevant:
+ result.append(line)
+ return "\n".join(result)
+
+
+def _process_line( # pylint: disable=too-many-return-statements
+ line: str,
+ refs: dict[str, dict],
+ options: dict[str, Any],
+ tracker: dict,
+ next_ref: Any,
+) -> str | None:
+ depth = _get_indent_level(line)
+ max_depth_val = options.get("maxDepth")
+ if max_depth_val is not None and depth > max_depth_val:
+ return None
+
+ m = re.match(r'^(\s*-\s*)(\w+)(?:\s+"([^"]*)")?(.*)$', line)
+ if not m:
+ return None if options.get("interactive") else line
+
+ prefix, role_raw, name, suffix = m.groups()
+ if role_raw.startswith("/"):
+ return None if options.get("interactive") else line
+
+ role = role_raw.lower()
+ is_interactive = role in INTERACTIVE_ROLES
+ is_content = role in CONTENT_ROLES
+ is_structural = role in STRUCTURAL_ROLES
+
+ if options.get("interactive") and not is_interactive:
+ return None
+ if options.get("compact") and is_structural and not name:
+ return None
+
+ should_have_ref = is_interactive or (is_content and name)
+ if not should_have_ref:
+ return line
+
+ ref = next_ref()
+ nth = tracker["get_next_index"](role, name)
+ tracker["track_ref"](role, name, ref)
+ refs[ref] = {"role": role, "name": name, "nth": nth}
+
+ enhanced = f"{prefix}{role_raw}"
+ if name:
+ enhanced += f' "{name}"'
+ enhanced += f" [ref={ref}]"
+ if nth is not None and nth > 0:
+ enhanced += f" [nth={nth}]"
+ if suffix:
+ enhanced += suffix
+ return enhanced
+
+
+def build_role_snapshot_from_aria(
+ aria_snapshot: str,
+ *,
+ interactive: bool = False,
+ compact: bool = False,
+ max_depth: int | None = None,
+) -> tuple[str, dict[str, dict]]:
+ """Build snapshot + refs from Playwright locator.aria_snapshot() output."""
+ options = {
+ "interactive": interactive,
+ "compact": compact,
+ "maxDepth": max_depth,
+ }
+ lines = aria_snapshot.split("\n")
+ refs: dict[str, dict] = {}
+ tracker = _create_tracker()
+ counter = [0]
+
+ def next_ref() -> str:
+ counter[0] += 1
+ return f"e{counter[0]}"
+
+ if options.get("interactive"):
+ result_lines = []
+ for line in lines:
+ depth = _get_indent_level(line)
+ max_d = options.get("maxDepth")
+ if max_d is not None and depth > max_d:
+ continue
+ m = re.match(r'^(\s*-\s*)(\w+)(?:\s+"([^"]*)")?(.*)$', line)
+ if not m:
+ continue
+ _, role_raw, name, suffix = m.groups()
+ if role_raw.startswith("/"):
+ continue
+ role = role_raw.lower()
+ if role not in INTERACTIVE_ROLES:
+ continue
+ ref = next_ref()
+ nth = tracker["get_next_index"](role, name)
+ tracker["track_ref"](role, name, ref)
+ refs[ref] = {"role": role, "name": name, "nth": nth}
+ enhanced = f"- {role_raw}"
+ if name:
+ enhanced += f' "{name}"'
+ enhanced += f" [ref={ref}]"
+ if nth is not None and nth > 0:
+ enhanced += f" [nth={nth}]"
+ if "[" in suffix:
+ enhanced += suffix
+ result_lines.append(enhanced)
+ _remove_nth_from_non_duplicates(refs, tracker)
+ snapshot = "\n".join(result_lines) or "(no interactive elements)"
+ return snapshot, refs
+
+ result_lines = []
+ for line in lines:
+ processed = _process_line(line, refs, options, tracker, next_ref)
+ if processed is not None:
+ result_lines.append(processed)
+ _remove_nth_from_non_duplicates(refs, tracker)
+ tree = "\n".join(result_lines) or "(empty)"
+ snapshot = _compact_tree(tree) if options.get("compact") else tree
+ return snapshot, refs
diff --git a/reme/memory/file_based/tools/file_io.py b/reme/memory/file_based/tools/file_io.py
new file mode 100644
index 00000000..3e28addc
--- /dev/null
+++ b/reme/memory/file_based/tools/file_io.py
@@ -0,0 +1,360 @@
+"""File I/O operations with a configurable working directory."""
+
+import os
+from pathlib import Path
+from typing import Optional
+
+from agentscope.message import TextBlock
+from agentscope.tool import ToolResponse
+
+from ..utils import read_file_safe, truncate_text_output, TRUNCATION_NOTICE_MARKER
+
+
+class FileIO:
+ """File I/O operations with a configurable working directory."""
+
+ def __init__(self, working_dir: str | Path):
+ """Initialize FileIO with a working directory.
+
+ Args:
+ working_dir (`str`):
+ The working directory for resolving relative paths.
+ """
+ self.working_dir = Path(working_dir)
+
+ def _resolve_file_path(self, file_path: str) -> str:
+ """Resolve file path: use absolute path as-is,
+ resolve relative path from working_dir.
+
+ Args:
+ file_path: The input file path (absolute or relative).
+
+ Returns:
+ The resolved absolute file path as string.
+ """
+ path = Path(file_path).expanduser()
+ if path.is_absolute():
+ return str(path)
+ else:
+ return str(self.working_dir / file_path)
+
+ async def read_file( # pylint: disable=too-many-return-statements
+ self,
+ file_path: str,
+ start_line: Optional[int] = None,
+ end_line: Optional[int] = None,
+ ) -> ToolResponse:
+ """Read a file. Relative paths resolve from WORKING_DIR.
+
+ Use start_line/end_line to read a specific line range (output includes
+ line numbers). Omit both to read the full file.
+
+ Args:
+ file_path (`str`):
+ Path to the file.
+ start_line (`int`, optional):
+ First line to read (1-based, inclusive).
+ end_line (`int`, optional):
+ Last line to read (1-based, inclusive).
+ """
+
+ # Convert start_line/end_line to int if they are strings
+ if start_line is not None:
+ try:
+ start_line = int(start_line)
+ except (ValueError, TypeError):
+ return ToolResponse(
+ content=[
+ TextBlock(
+ type="text",
+ text=f"Error: start_line must be an integer, got {start_line!r}.",
+ ),
+ ],
+ )
+
+ if end_line is not None:
+ try:
+ end_line = int(end_line)
+ except (ValueError, TypeError):
+ return ToolResponse(
+ content=[
+ TextBlock(
+ type="text",
+ text=f"Error: end_line must be an integer, got {end_line!r}.",
+ ),
+ ],
+ )
+
+ file_path = self._resolve_file_path(file_path)
+
+ if not os.path.exists(file_path):
+ return ToolResponse(
+ content=[
+ TextBlock(
+ type="text",
+ text=f"Error: The file {file_path} does not exist.",
+ ),
+ ],
+ )
+
+ if not os.path.isfile(file_path):
+ return ToolResponse(
+ content=[
+ TextBlock(
+ type="text",
+ text=f"Error: The path {file_path} is not a file.",
+ ),
+ ],
+ )
+
+ try:
+ content = read_file_safe(file_path)
+ all_lines = content.split("\n")
+ total = len(all_lines)
+
+ # Determine read range
+ s = max(1, start_line if start_line is not None else 1)
+ e = min(total, end_line if end_line is not None else total)
+
+ if s > total:
+ return ToolResponse(
+ content=[
+ TextBlock(
+ type="text",
+ text=f"Error: start_line {s} exceeds file length ({total} lines).",
+ ),
+ ],
+ )
+
+ if s > e:
+ return ToolResponse(
+ content=[
+ TextBlock(
+ type="text",
+ text=f"Error: start_line ({s}) > end_line ({e}).",
+ ),
+ ],
+ )
+
+ # Extract selected lines
+ selected_content = "\n".join(all_lines[s - 1 : e])
+
+ # Apply smart truncation (consistent with shell output format)
+ text = truncate_text_output(
+ selected_content,
+ start_line=s,
+ total_lines=total,
+ file_path=file_path,
+ )
+
+ # Add continuation hint if partial read without truncation.
+ # Use TRUNCATION_NOTICE_MARKER format so ToolResultCompactor can
+ # re-truncate with the correct start_line when compacting old messages.
+ if text == selected_content and e < total:
+ content_bytes = len(text.encode("utf-8"))
+ notice = (
+ TRUNCATION_NOTICE_MARKER + f"\nThe output above was truncated."
+ f"\nThe full content is saved to the file "
+ f"and contains {total} lines in total."
+ f"\nThis excerpt starts at line {s} and "
+ f"covers the next {content_bytes} bytes."
+ "\nIf the current content is not enough, "
+ f"call `read_file` with file_path={file_path} start_line={e + 1} to read more."
+ )
+ text = text + notice
+
+ return ToolResponse(
+ content=[TextBlock(type="text", text=text)],
+ )
+
+ except Exception as e:
+ return ToolResponse(
+ content=[
+ TextBlock(
+ type="text",
+ text=f"Error: Read file failed due to \n{e}",
+ ),
+ ],
+ )
+
+ async def write_file(
+ self,
+ file_path: str,
+ content: str,
+ ) -> ToolResponse:
+ """Create or overwrite a file. Relative paths resolve from working_dir.
+
+ Args:
+ file_path (`str`):
+ Path to the file.
+ content (`str`):
+ Content to write.
+ """
+ if not file_path:
+ return ToolResponse(
+ content=[
+ TextBlock(
+ type="text",
+ text="Error: No `file_path` provided.",
+ ),
+ ],
+ )
+
+ file_path = self._resolve_file_path(file_path)
+
+ try:
+ with open(file_path, "w", encoding="utf-8") as file:
+ file.write(content)
+ return ToolResponse(
+ content=[
+ TextBlock(
+ type="text",
+ text=f"Wrote {len(content)} bytes to {file_path}.",
+ ),
+ ],
+ )
+ except Exception as e:
+ return ToolResponse(
+ content=[
+ TextBlock(
+ type="text",
+ text=f"Error: Write file failed due to \n{e}",
+ ),
+ ],
+ )
+
+ # pylint: disable=too-many-return-statements
+ async def edit_file(
+ self,
+ file_path: str,
+ old_text: str,
+ new_text: str,
+ ) -> ToolResponse:
+ """Find-and-replace text in a file. All occurrences of old_text are
+ replaced with new_text. Relative paths resolve from working_dir.
+
+ Args:
+ file_path (`str`):
+ Path to the file.
+ old_text (`str`):
+ Exact text to find.
+ new_text (`str`):
+ Replacement text.
+ """
+ if not file_path:
+ return ToolResponse(
+ content=[
+ TextBlock(
+ type="text",
+ text="Error: No `file_path` provided.",
+ ),
+ ],
+ )
+
+ resolved_path = self._resolve_file_path(file_path)
+
+ if not os.path.exists(resolved_path):
+ return ToolResponse(
+ content=[
+ TextBlock(
+ type="text",
+ text=f"Error: The file {resolved_path} does not exist.",
+ ),
+ ],
+ )
+
+ if not os.path.isfile(resolved_path):
+ return ToolResponse(
+ content=[
+ TextBlock(
+ type="text",
+ text=f"Error: The path {resolved_path} is not a file.",
+ ),
+ ],
+ )
+
+ try:
+ content = read_file_safe(resolved_path)
+ except Exception as e:
+ return ToolResponse(
+ content=[
+ TextBlock(
+ type="text",
+ text=f"Error: Read file failed due to \n{e}",
+ ),
+ ],
+ )
+
+ if old_text not in content:
+ return ToolResponse(
+ content=[
+ TextBlock(
+ type="text",
+ text=f"Error: The text to replace was not found in {file_path}.",
+ ),
+ ],
+ )
+
+ new_content = content.replace(old_text, new_text)
+ write_response = await self.write_file(file_path=resolved_path, content=new_content)
+
+ if write_response.content and len(write_response.content) > 0:
+ write_text = write_response.content[0].get("text", "")
+ if write_text.startswith("Error:"):
+ return write_response
+
+ return ToolResponse(
+ content=[
+ TextBlock(
+ type="text",
+ text=f"Successfully replaced text in {file_path}.",
+ ),
+ ],
+ )
+
+ async def append_file(
+ self,
+ file_path: str,
+ content: str,
+ ) -> ToolResponse:
+ """Append content to the end of a file. Relative paths resolve from
+ working_dir.
+
+ Args:
+ file_path (`str`):
+ Path to the file.
+ content (`str`):
+ Content to append.
+ """
+ if not file_path:
+ return ToolResponse(
+ content=[
+ TextBlock(
+ type="text",
+ text="Error: No `file_path` provided.",
+ ),
+ ],
+ )
+
+ file_path = self._resolve_file_path(file_path)
+
+ try:
+ with open(file_path, "a", encoding="utf-8") as file:
+ file.write(content)
+ return ToolResponse(
+ content=[
+ TextBlock(
+ type="text",
+ text=f"Appended {len(content)} bytes to {file_path}.",
+ ),
+ ],
+ )
+ except Exception as e:
+ return ToolResponse(
+ content=[
+ TextBlock(
+ type="text",
+ text=f"Error: Append file failed due to \n{e}",
+ ),
+ ],
+ )
diff --git a/reme/memory/tools/chunk/memory_get.py b/reme/memory/file_based/tools/memory_get.py
similarity index 97%
rename from reme/memory/tools/chunk/memory_get.py
rename to reme/memory/file_based/tools/memory_get.py
index 572a4a26..9cc16968 100644
--- a/reme/memory/tools/chunk/memory_get.py
+++ b/reme/memory/file_based/tools/memory_get.py
@@ -3,12 +3,15 @@
import os
from pathlib import Path
-from loguru import logger
from ....core import RuntimeContext
from ....core.op import BaseTool
from ....core.schema import ToolCall
+from ....core.utils import get_logger
+
+logger = get_logger()
+
class MemoryGet(BaseTool):
"""Read specific snippets from memory files."""
@@ -109,5 +112,5 @@ class MemoryGet(BaseTool):
except Exception as e:
# Return error message to LLM instead of raising
error_msg = f"{self.__class__.__name__} failed: {str(e)}"
- logger.error(error_msg)
+ logger.exception(error_msg)
return await self.after_execute(error_msg)
diff --git a/reme/memory/tools/chunk/memory_search.py b/reme/memory/file_based/tools/memory_search.py
similarity index 97%
rename from reme/memory/tools/chunk/memory_search.py
rename to reme/memory/file_based/tools/memory_search.py
index 7423b331..927a6fc1 100644
--- a/reme/memory/tools/chunk/memory_search.py
+++ b/reme/memory/file_based/tools/memory_search.py
@@ -2,12 +2,14 @@
import json
-from loguru import logger
from ....core.enumeration import MemorySource
from ....core.op import BaseTool
from ....core.runtime_context import RuntimeContext
from ....core.schema import ToolCall
+from ....core.utils import get_logger
+
+logger = get_logger()
class MemorySearch(BaseTool):
@@ -108,5 +110,5 @@ class MemorySearch(BaseTool):
except Exception as e:
# Return error message to LLM instead of raising
error_msg = f"{self.__class__.__name__} failed: {str(e)}"
- logger.error(error_msg)
+ logger.exception(error_msg)
return await self.after_execute(error_msg)
diff --git a/reme/memory/file_based/tools/shell.py b/reme/memory/file_based/tools/shell.py
new file mode 100644
index 00000000..cce138f9
--- /dev/null
+++ b/reme/memory/file_based/tools/shell.py
@@ -0,0 +1,223 @@
+# -*- coding: utf-8 -*-
+# flake8: noqa: E501
+# pylint: disable=line-too-long
+"""The shell command tool."""
+
+import asyncio
+import locale
+import subprocess
+import sys
+from pathlib import Path
+
+from agentscope.message import TextBlock
+from agentscope.tool import ToolResponse
+
+
+def _execute_subprocess_sync(
+ cmd: str,
+ cwd: str,
+ timeout: int,
+) -> tuple[int, str, str]:
+ """Execute subprocess synchronously in a thread.
+
+ This function runs in a separate thread to avoid Windows asyncio
+ subprocess limitations.
+
+ Args:
+ cmd (`str`):
+ The shell command to execute.
+ cwd (`str`):
+ The working directory for the command execution.
+ timeout (`int`):
+ The maximum time (in seconds) allowed for the command to run.
+
+ Returns:
+ `tuple[int, str, str]`:
+ A tuple containing the return code, standard output, and
+ standard error of the executed command. If timeout occurs, the
+ return code will be -1 and stderr will contain timeout information.
+ """
+ try:
+ result = subprocess.run(
+ cmd,
+ shell=True,
+ capture_output=True,
+ text=True,
+ cwd=cwd,
+ timeout=timeout,
+ encoding=locale.getpreferredencoding(False) or "utf-8",
+ errors="replace",
+ check=True,
+ )
+ return (
+ result.returncode,
+ result.stdout.strip("\n"),
+ result.stderr.strip("\n"),
+ )
+ except subprocess.TimeoutExpired:
+ return (
+ -1,
+ "",
+ f"Command execution exceeded the timeout of {timeout} seconds.",
+ )
+ except Exception as e:
+ return -1, "", str(e)
+
+
+class Shell:
+ """Shell command execution with a configurable working directory."""
+
+ def __init__(self, working_dir: str | Path):
+ """Initialize Shell with a working directory.
+
+ Args:
+ working_dir (`str | Path`):
+ The working directory for command execution.
+ """
+ self.working_dir = Path(working_dir)
+
+ # pylint: disable=too-many-branches, too-many-statements
+ async def execute_shell_command(
+ self,
+ command: str,
+ timeout: int = 60,
+ ) -> ToolResponse:
+ """Execute given command and return the return code, standard output and
+ error within >>"
-TRUNCATION_MARKER_END = "<<>>"
-
-
-def truncate_text(text: str, max_length: int) -> str:
- """Truncate text to max length, keeping head and tail portions.
-
- Args:
- text: The text to truncate
- max_length: Maximum allowed length
-
- Returns:
- Truncated text with unique markers indicating truncation
- """
- text = str(text) if text else ""
- if not text:
- return text
-
- if len(text) <= max_length:
- return text
-
- half_length = max_length // 2
- truncated_chars = len(text) - max_length
- logger.debug(
- "Text truncated: original %d chars, kept head %d + tail %d, removed %d chars.",
- len(text),
- half_length,
- half_length,
- truncated_chars,
- )
- return (
- f"{text[:half_length]}\n\n{TRUNCATION_MARKER_START} "
- f"({truncated_chars} characters omitted) "
- f"{TRUNCATION_MARKER_END}\n\n{text[-half_length:]}"
- )
-
-
-def is_truncated(text: str) -> bool:
- """Check if the text has been truncated (contains truncation markers).
-
- Args:
- text: The text to check
-
- Returns:
- bool: True if text contains truncation markers, False otherwise
- """
- if not text:
- return False
- return TRUNCATION_MARKER_START in text and TRUNCATION_MARKER_END in text
-
-
-def _extract_text_from_messages(messages: list[dict]) -> str:
- """Extract text content from messages and concatenate into a string.
-
- Handles various message formats:
- - Simple string content: {"role": "user", "content": "hello"}
- - List content with text blocks:
- {"role": "user", "content": [{"type": "text", "text": "hello"}]}
- - List content with tool_result blocks:
- {"role": "user", "content": [{"type": "tool_result", "output": "..."}]}
-
- Args:
- messages: List of message dictionaries in chat format.
-
- Returns:
- str: Concatenated text content from all messages.
- """
- parts = []
- for msg in messages:
- content = msg.get("content", "")
- if isinstance(content, str):
- parts.append(content)
- elif isinstance(content, list):
- for block in content:
- if isinstance(block, dict):
- block_type = block.get("type", "")
- if block_type == "tool_result":
- output = block.get("output", "")
- if isinstance(output, str) and output:
- parts.append(output)
- elif isinstance(output, list):
- for sub in output:
- if isinstance(sub, dict):
- sub_text = sub.get("text") or sub.get("content", "")
- if sub_text:
- parts.append(str(sub_text))
- else:
- text = block.get("text") or block.get("content", "")
- if text:
- parts.append(str(text))
- elif isinstance(block, str):
- parts.append(block)
- return "\n".join(parts)
-
-
-def safe_count_message_tokens(
- token_counter: HuggingFaceTokenCounter,
- messages: list[dict],
-) -> int:
- """Safely count tokens in messages with fallback estimation.
-
- This is a wrapper around count_message_tokens that catches exceptions
- and falls back to a character-based estimation (len // 4) if the
- tokenizer fails.
-
- Args:
- token_counter: Token counter instance.
- messages: List of message dictionaries in chat format.
-
- Returns:
- int: The estimated number of tokens in the messages.
- """
- try:
- text = _extract_text_from_messages(messages)
- token_ids = token_counter.tokenizer.encode(text)
- token_count = len(token_ids)
- return token_count
-
- except Exception as e:
- # Fallback to character-based estimation
- text = _extract_text_from_messages(messages)
- estimated_tokens = len(text) // 4
- logger.warning(
- "Failed to count tokens: %s, using estimated_tokens=%d",
- e,
- estimated_tokens,
- )
- return estimated_tokens
-
-
-def safe_count_str_tokens(
- token_counter: HuggingFaceTokenCounter,
- text: str,
-) -> int:
- """Safely count tokens in a string with fallback estimation.
-
- Uses the tokenizer to count tokens in the given text. If the tokenizer
- fails, falls back to a character-based estimation (len // 4).
-
- Args:
- token_counter: Token counter instance.
- text: The string to count tokens for.
-
- Returns:
- int: The estimated number of tokens in the string.
- """
- try:
- token_ids = token_counter.tokenizer.encode(text)
- token_count = len(token_ids)
- return token_count
- except Exception as e:
- # Fallback to character-based estimation
- estimated_tokens = len(text) // 4
- logger.warning(
- "Failed to count string tokens: %s, using estimated_tokens=%d",
- e,
- estimated_tokens,
- )
- return estimated_tokens
-
-
-def _get_block_tokens( # pylint: disable=too-many-return-statements
- block: dict,
- block_type: str,
- token_counter: HuggingFaceTokenCounter,
-) -> tuple[int, str]:
- """Get token count and content string for different block types.
-
- Args:
- block: The content block dict
- block_type: The type of the block
-
- Returns:
- Tuple of (token count, content string)
- """
- if block_type == "text":
- text = block.get("text", "")
- return (safe_count_str_tokens(token_counter, text), text) if text else (0, "")
-
- if block_type == "thinking":
- thinking = block.get("thinking", "")
- return (safe_count_str_tokens(token_counter, thinking), thinking) if thinking else (0, "")
-
- if block_type == "tool_use":
- # Count input dict and raw_input string
- input_dict = block.get("input", {})
- raw_input = block.get("raw_input", "")
- input_str = str(input_dict) if input_dict else ""
- total = input_str + raw_input
- return (safe_count_str_tokens(token_counter, total), total) if total else (0, "")
-
- if block_type == "tool_result":
- output = block.get("output")
- if isinstance(output, str):
- return (safe_count_str_tokens(token_counter, output), output) if output else (0, "")
-
- if isinstance(output, list):
- # Recursively count tokens in nested blocks
- total_tokens = 0
- total_str = ""
- for item in output:
- if isinstance(item, dict):
- item_type = item.get("type", "unknown")
- item_tokens, item_str = _get_block_tokens(item, item_type, token_counter)
- total_tokens += item_tokens
- total_str += item_str
- return total_tokens, total_str
- return 0, ""
-
- if block_type in ("image", "audio", "video"):
- # For media blocks, count the URL or indicate base64 size
- source = block.get("source", {})
- if source.get("type") == "url":
- url = source.get("url", "")
- return safe_count_str_tokens(token_counter, url), url
- if source.get("type") == "base64":
- # Base64 data can be large, return approximate token count
- data = source.get("data", "")
- return (len(data) // 4, "[base64]") if data else (0, "")
- return 0, ""
-
- return 0, ""
-
-
-_token_counter = None
-
-
-def get_token_counter():
- """Get or initialize the global token counter instance.
-
- Returns:
- TokenCounterBase: The token counter instance for Qwen models.
-
- Raises:
- RuntimeError: If token counter initialization fails.
- """
- global _token_counter
- if _token_counter is None:
- # Use Qwen tokenizer for DashScope models
- # Qwen3 series uses the same tokenizer as Qwen2.5
-
- # Try local tokenizer first, fall back to online if not found
- local_tokenizer_path = Path(__file__).parent.parent.parent / "tokenizer"
-
- if local_tokenizer_path.exists() and (local_tokenizer_path / "tokenizer.json").exists():
- tokenizer_path = str(local_tokenizer_path)
- logger.info(f"Using local Qwen tokenizer from {tokenizer_path}")
- else:
- tokenizer_path = "Qwen/Qwen2.5-7B-Instruct"
- logger.info(
- "Local tokenizer not found, downloading from HuggingFace",
- )
-
- _token_counter = HuggingFaceTokenCounter(
- pretrained_model_name_or_path=tokenizer_path,
- use_mirror=True, # Use HF mirror for users in China
- use_fast=True,
- trust_remote_code=True,
- )
- logger.debug("Token counter initialized with Qwen tokenizer")
- return _token_counter
diff --git a/reme/memory/file_based/utils/__init__.py b/reme/memory/file_based/utils/__init__.py
new file mode 100644
index 00000000..769a8a66
--- /dev/null
+++ b/reme/memory/file_based/utils/__init__.py
@@ -0,0 +1,12 @@
+"""utils"""
+
+from .as_msg_handler import AsMsgHandler
+from .file_utils import truncate_text_output, read_file_safe, DEFAULT_MAX_BYTES, TRUNCATION_NOTICE_MARKER
+
+__all__ = [
+ "AsMsgHandler",
+ "truncate_text_output",
+ "read_file_safe",
+ "DEFAULT_MAX_BYTES",
+ "TRUNCATION_NOTICE_MARKER",
+]
diff --git a/reme/memory/file_based/utils/as_msg_handler.py b/reme/memory/file_based/utils/as_msg_handler.py
new file mode 100644
index 00000000..e175a402
--- /dev/null
+++ b/reme/memory/file_based/utils/as_msg_handler.py
@@ -0,0 +1,398 @@
+"""Handler for AgentScope message processing, token counting, and context management."""
+
+import json
+
+from agentscope.message import Msg
+from agentscope.token import HuggingFaceTokenCounter
+
+from ....core.schema import AsMsgStat, AsBlockStat
+from ....core.utils import get_logger
+
+logger = get_logger()
+
+
+class AsMsgHandler:
+ """Handles token counting, formatting, and context compaction for AgentScope messages."""
+
+ def __init__(self, token_counter: HuggingFaceTokenCounter):
+ self._token_counter = token_counter
+
+ async def count_str_token(self, text: str) -> int:
+ """Count tokens in a string.
+
+ Args:
+ text: The text to count tokens for.
+
+ Returns:
+ The number of tokens in the text.
+ """
+ if not text:
+ return 0
+
+ try:
+ token_count = await self._token_counter.count(messages=[], text=text)
+ assert token_count > 0, "Invalid token count"
+ return token_count
+
+ except Exception as e:
+ estimated_tokens = int(len(text.encode("utf-8")) / 3.75)
+ logger.warning(f"Failed to count string tokens: {text}, e={e}")
+ return estimated_tokens
+
+ async def _format_tool_result_output(self, output: str | list[dict]) -> tuple[str, int]:
+ """Convert tool result output to string."""
+ if isinstance(output, str):
+ return output, await self.count_str_token(output)
+
+ textual_parts = []
+ total_token_count = 0
+ for block in output:
+ try:
+ if not isinstance(block, dict) or "type" not in block:
+ logger.warning(
+ f"Invalid block: {block}, expected a dict with 'type' key, skipped.",
+ )
+ continue
+
+ block_type = block["type"]
+
+ if block_type == "text":
+ textual_parts.append(block.get("text", ""))
+ total_token_count += await self.count_str_token(textual_parts[-1])
+
+ elif block_type in ["image", "audio", "video"]:
+ source = block.get("source", {})
+ if source.get("type") == "base64":
+ data = source.get("data", "")
+ total_token_count += len(data) // 4 if data else 10
+ else:
+ url = source.get("url", "")
+ total_token_count += await self.count_str_token(url) if url else 10
+ textual_parts.append(f"[{block_type}] {url}")
+
+ elif block_type == "file":
+ file_path = block.get("path", "") or block.get("url", "")
+ file_name = block.get("name", file_path)
+ textual_parts.append(f"[file] {file_name}: {file_path}")
+ total_token_count += await self.count_str_token(file_path)
+
+ else:
+ logger.warning(
+ f"Unsupported block type '{block_type}' in tool result, skipped.",
+ )
+
+ except Exception as e:
+ logger.warning(
+ f"Failed to process block {block}: {e}, skipped.",
+ )
+
+ return "\n".join(textual_parts), total_token_count
+
+ async def stat_message(self, message: Msg) -> AsMsgStat:
+ """Analyze a message and generate block statistics."""
+ blocks = []
+ if isinstance(message.content, str):
+ blocks.append(
+ AsBlockStat(
+ block_type="text",
+ text=message.content,
+ token_count=await self.count_str_token(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")
+
+ if block_type == "text":
+ text = block.get("text", "")
+ token_count = await self.count_str_token(text)
+ blocks.append(
+ AsBlockStat(
+ block_type=block_type,
+ text=text,
+ token_count=token_count,
+ ),
+ )
+
+ elif block_type == "thinking":
+ thinking = block.get("thinking", "")
+ token_count = await self.count_str_token(thinking)
+ blocks.append(
+ AsBlockStat(
+ block_type=block_type,
+ text=thinking,
+ token_count=token_count,
+ ),
+ )
+
+ elif block_type in ("image", "audio", "video"):
+ source = block.get("source", {})
+ url = source.get("url", "")
+ if source.get("type") == "base64":
+ data = source.get("data", "")
+ token_count = len(data) // 4 if data else 10
+ else:
+ token_count = await self.count_str_token(url) if url else 10
+ blocks.append(
+ AsBlockStat(
+ block_type=block_type,
+ text="",
+ token_count=token_count,
+ media_url=url,
+ ),
+ )
+
+ elif block_type == "tool_use":
+ tool_name = block.get("name", "")
+ tool_input = block.get("input", "")
+ try:
+ input_str = json.dumps(tool_input, ensure_ascii=False)
+ except (TypeError, ValueError):
+ input_str = str(tool_input)
+ token_count = await self.count_str_token(tool_name + input_str)
+ blocks.append(
+ AsBlockStat(
+ block_type=block_type,
+ text="",
+ token_count=token_count,
+ tool_name=tool_name,
+ tool_input=input_str,
+ ),
+ )
+
+ elif block_type == "tool_result":
+ tool_name = block.get("name", "")
+ output = block.get("output", "")
+ formatted_output, token_count = await self._format_tool_result_output(output)
+ blocks.append(
+ AsBlockStat(
+ block_type=block_type,
+ text="",
+ token_count=token_count,
+ tool_name=tool_name,
+ tool_output=formatted_output,
+ ),
+ )
+
+ else:
+ logger.warning(f"Unsupported block type {block_type}, skipped.")
+
+ return AsMsgStat(
+ name=message.name or message.role,
+ role=message.role,
+ content=blocks,
+ timestamp=message.timestamp or "",
+ metadata=message.metadata or {},
+ )
+
+ async def count_msgs_token(self, messages: list[Msg]) -> int:
+ """Count total token count of a list of messages."""
+ total = 0
+ for msg in messages:
+ stat = await self.stat_message(msg)
+ total += stat.total_tokens
+ return total
+
+ async def format_msgs_to_str(
+ self,
+ messages: list[Msg],
+ memory_compact_threshold: int,
+ include_thinking: bool = True,
+ ) -> str:
+ """Format list of messages to a single formatted string.
+
+ Messages are processed in reverse order (newest first) and older
+ messages are skipped when token count exceeds memory_compact_threshold.
+
+ Args:
+ messages: List of Msg objects to format.
+ memory_compact_threshold: Maximum token count before skipping older messages.
+ include_thinking: Whether to include thinking blocks in output.
+ """
+ if not messages:
+ return ""
+
+ formatted_parts: list[str] = []
+ total_token_count = 0
+
+ for i in range(len(messages) - 1, -1, -1):
+ stat = await self.stat_message(messages[i])
+ formatted_content = stat.format(include_thinking=include_thinking)
+ content_token_count = await self.count_str_token(formatted_content)
+
+ is_latest = i == len(messages) - 1
+ if not is_latest and total_token_count + content_token_count > memory_compact_threshold:
+ logger.info(
+ f"Skipping older messages: adding {content_token_count} tokens would exceed threshold "
+ f"{memory_compact_threshold} (current: {total_token_count})",
+ )
+ break
+
+ if is_latest and content_token_count > memory_compact_threshold:
+ logger.warning(
+ f"Latest message alone ({content_token_count} tokens) exceeds threshold "
+ f"{memory_compact_threshold}, including it anyway.",
+ )
+
+ formatted_parts.append(formatted_content)
+ total_token_count += content_token_count
+
+ formatted_parts.reverse()
+ return "\n\n".join(formatted_parts)
+
+ @staticmethod
+ def validate_tool_ids_alignment(messages: list[Msg]) -> bool:
+ """Check if tool_use_ids and tool_result_ids are properly aligned.
+
+ Args:
+ messages: List of Msg objects to validate.
+
+ Returns:
+ True if all tool_use ids have corresponding tool_result ids and vice versa.
+ """
+ tool_use_ids: set[str] = set()
+ tool_result_ids: set[str] = set()
+
+ for msg in messages:
+ for block in msg.get_content_blocks("tool_use"):
+ if tool_id := block.get("id"):
+ tool_use_ids.add(tool_id)
+ for block in msg.get_content_blocks("tool_result"):
+ if tool_id := block.get("id"):
+ tool_result_ids.add(tool_id)
+
+ return tool_use_ids == tool_result_ids
+
+ async def context_check(
+ self,
+ messages: list[Msg],
+ memory_compact_threshold: int,
+ memory_compact_reserve: int,
+ ) -> tuple[list[Msg], list[Msg], bool]:
+ """Check if context exceeds threshold and split messages accordingly.
+
+ Only when total tokens exceed memory_compact_threshold, messages are split into
+ messages_to_keep (within reserve limit) and messages_to_compact (older messages).
+
+ Args:
+ messages: List of Msg objects to check.
+ memory_compact_threshold: Maximum token count threshold to trigger compaction.
+ memory_compact_reserve: Token limit for messages to keep.
+
+ Returns:
+ A tuple of (messages_to_compact, messages_to_keep, tools_aligned):
+ - messages_to_compact: Older messages that exceed reserve limit
+ - messages_to_keep: Recent messages within the reserve limit
+ - tools_aligned: Whether tool_use and tool_result ids are aligned in messages_to_keep
+ """
+ if not messages:
+ return [], [], True
+
+ # Calculate total tokens and stats for all messages
+ msg_stats: list[tuple[Msg, AsMsgStat]] = []
+ total_tokens = 0
+ for msg in messages:
+ stat = await self.stat_message(msg)
+ msg_stats.append((msg, stat))
+ total_tokens += stat.total_tokens
+
+ # If total tokens don't exceed threshold, no split needed
+ if total_tokens < memory_compact_threshold:
+ return [], messages, True
+
+ # Collect all tool_use ids and their message indices
+ # tool_use_id -> message index
+ tool_use_locations: dict[str, int] = {}
+ # tool_result_id -> message index
+ tool_result_locations: dict[str, int] = {}
+
+ for idx, (msg, _) in enumerate(msg_stats):
+ for block in msg.get_content_blocks("tool_use"):
+ tool_id = block.get("id", "")
+ if tool_id:
+ tool_use_locations[tool_id] = idx
+
+ for block in msg.get_content_blocks("tool_result"):
+ tool_id = block.get("id", "")
+ if tool_id:
+ tool_result_locations[tool_id] = idx
+
+ # Iterate from the end, accumulating messages to keep within reserve limit
+ keep_indices: set[int] = set()
+ 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
+ if accumulated_tokens + stat.total_tokens > memory_compact_reserve:
+ logger.info(
+ f"Context check: adding message {i} with {stat.total_tokens} tokens would exceed reserve "
+ f"{memory_compact_reserve} (current: {accumulated_tokens})",
+ )
+ break
+
+ # Check tool_result dependencies - if this message has tool_result,
+ # we need to ensure the corresponding tool_use is also included
+ tool_result_ids = [
+ block.get("id", "") for block in msg.get_content_blocks("tool_result") if block.get("id", "")
+ ]
+
+ # Calculate extra tokens needed for dependent tool_use messages
+ extra_tokens = 0
+ dependent_indices: set[int] = set()
+
+ for tool_id in tool_result_ids:
+ if tool_id in tool_use_locations:
+ tool_use_idx = tool_use_locations[tool_id]
+ if tool_use_idx not in keep_indices and tool_use_idx != i:
+ dependent_indices.add(tool_use_idx)
+ _, dep_stat = msg_stats[tool_use_idx]
+ extra_tokens += dep_stat.total_tokens
+
+ # Check if we can fit this message plus its dependencies within reserve
+ if accumulated_tokens + stat.total_tokens + extra_tokens > memory_compact_reserve:
+ logger.info(
+ f"Context check: message {i} requires {extra_tokens} extra tokens for tool_use dependencies, "
+ f"total would exceed reserve {memory_compact_reserve}",
+ )
+ break
+
+ # Add this message and its dependencies
+ keep_indices.add(i)
+ keep_indices.update(dependent_indices)
+ accumulated_tokens += stat.total_tokens + extra_tokens
+
+ # Build final lists based on keep_indices (preserve original order)
+ messages_to_compact = []
+ messages_to_keep = []
+
+ for idx, (msg, _) in enumerate(msg_stats):
+ if idx in keep_indices:
+ messages_to_keep.append(msg)
+ else:
+ messages_to_compact.append(msg)
+
+ # Validate tool ids alignment for messages_to_keep
+ tools_aligned = self.validate_tool_ids_alignment(messages_to_keep)
+
+ logger.info(
+ f"Context check result: {len(messages_to_compact)} messages to compact, "
+ f"{len(messages_to_keep)} messages to keep, "
+ f"total tokens: {total_tokens}, threshold: {memory_compact_threshold}, "
+ f"reserve: {memory_compact_reserve}, kept tokens: {accumulated_tokens}, "
+ f"tools_aligned: {tools_aligned}",
+ )
+
+ return messages_to_compact, messages_to_keep, tools_aligned
diff --git a/reme/memory/file_based/utils/file_utils.py b/reme/memory/file_based/utils/file_utils.py
new file mode 100644
index 00000000..50248554
--- /dev/null
+++ b/reme/memory/file_based/utils/file_utils.py
@@ -0,0 +1,217 @@
+# -*- coding: utf-8 -*-
+"""Shared utilities for file and shell tools."""
+
+import re
+
+from ....core.utils import get_logger
+
+logger = get_logger()
+
+# Default truncation limit
+DEFAULT_MAX_BYTES = 50 * 1024
+
+# Maximum file size to read into memory (1GB)
+MAX_FILE_READ_BYTES = 1024 * 1024 * 1024
+
+# Marker prepended to every truncation notice.
+# Format:
+# <<>>
+# The output above was truncated.
+# The full content is saved to the file and contains Z lines in total.
+# This excerpt starts at line X and covers the next N bytes.
+# If the current content is not enough, call `read_file` with file_path= start_line=Y to read more.
+#
+# Split output on this marker to recover the original (untruncated) portion:
+# original = output.split(TRUNCATION_NOTICE_MARKER)[0]
+TRUNCATION_NOTICE_MARKER = "<<>>"
+
+
+def _truncate_fresh(
+ text: str,
+ start_line: int,
+ total_lines: int,
+ max_bytes: int,
+ file_path: str | None,
+ encoding: str,
+) -> str:
+ """Truncate fresh text (no prior truncation marker) by bytes with line integrity.
+
+ Slices at the byte boundary and appends a truncation notice with a continuation
+ hint so callers know which line to read next.
+
+ Returns the original text unchanged when it fits within max_bytes, or when the
+ last line itself exceeds max_bytes (unhandled edge case).
+ """
+ text_bytes = text.encode(encoding)
+
+ # Under the byte limit — return as-is without any modification.
+ if len(text_bytes) <= max_bytes:
+ return text
+
+ # Slice at the byte boundary.
+ # Assuming every single line is shorter than DEFAULT_MAX_BYTES, this cut always
+ # lands mid-line, guaranteeing at least one complete line before the boundary.
+ # Lines that exceed DEFAULT_MAX_BYTES are not handled and may be skipped entirely.
+ truncated = text_bytes[:max_bytes]
+ # Decode back to str; errors="ignore" drops any split multi-byte character
+ # at the cut boundary without raising an exception.
+ result = truncated.decode(encoding, errors="ignore")
+
+ # Count '\n' characters to determine how many complete lines are included.
+ # The tail after the final '\n' is a partial line that will be covered by
+ # the next read starting at next_line.
+ newline_count = result.count("\n")
+
+ # Compute the first line number not yet fully included in this chunk.
+ # max(1, ...) prevents next_line from equaling start_line when a single line
+ # exceeds max_bytes (newline_count == 0), which would make the caller retry
+ # the same range indefinitely.
+ next_line = start_line + max(1, newline_count)
+
+ if next_line <= total_lines:
+ # Truncation fell before the last line — continue reading from next_line.
+ read_from = next_line
+ elif start_line < total_lines:
+ # next_line overshot total_lines, meaning the cut landed inside the last line.
+ # Re-read from the start of the last line so the caller gets it in full.
+ read_from = total_lines
+ else:
+ # start_line == total_lines: the last line itself exceeds DEFAULT_MAX_BYTES.
+ # This case is outside our handled range — return without a truncation notice.
+ return result
+
+ notice = (
+ TRUNCATION_NOTICE_MARKER + f"\nThe output above was truncated."
+ f"\nThe full content is saved to the file and contains {total_lines} lines in total."
+ f"\nThis excerpt starts at line {start_line} and covers the next {max_bytes} bytes."
+ f"\nIf the current content is not enough, call `read_file` with file_path={file_path or ''} "
+ f"start_line={read_from} to read more."
+ )
+
+ return result + notice
+
+
+def _retruncate(
+ text: str,
+ max_bytes: int,
+ encoding: str,
+) -> str:
+ """Re-truncate text that was previously truncated (contains TRUNCATION_NOTICE_MARKER).
+
+ Extracts the original content before the marker, applies the new byte limit, and
+ updates the embedded notice (byte count and continuation line number) via regex.
+
+ Returns the original text unchanged when:
+ - the content already fits within max_bytes (with a small slack);
+ - required metadata fields cannot be parsed from the existing notice.
+ """
+ parts = text.split(TRUNCATION_NOTICE_MARKER, 1)
+ original_content = parts[0]
+ old_notice = parts[1]
+
+ text_bytes = original_content.encode(encoding)
+
+ # Allow a small slack to avoid unnecessary re-truncation when content is just
+ # barely over the limit (e.g. due to minor encoding differences).
+ if len(text_bytes) <= max_bytes + 100:
+ return text
+
+ # Parse start_line from notice; return text unchanged if not found
+ start_match = re.search(r"starts at line (\d+)", old_notice)
+ if not start_match:
+ return text
+ start_line_parsed = int(start_match.group(1))
+
+ # Re-slice to the new byte limit.
+ # Because every line is assumed to be shorter than DEFAULT_MAX_BYTES, the cut
+ # always falls somewhere mid-line, so at least one complete line is preserved.
+ truncated_bytes = text_bytes[:max_bytes]
+ # errors="ignore" silently drops any incomplete multi-byte character at the cut boundary.
+ result = truncated_bytes.decode(encoding, errors="ignore")
+ # Each '\n' in result corresponds to one fully-included line;
+ # anything after the last '\n' is a partial line that was cut off.
+ newline_count = result.count("\n")
+
+ # The next read should start at the line immediately after all complete lines.
+ # max(1, ...) guards against the theoretical zero-newline case
+ # (impossible when every line is shorter than DEFAULT_MAX_BYTES).
+ next_line = start_line_parsed + max(1, newline_count)
+
+ if not re.search(r"covers the next \d+ bytes", old_notice):
+ return text
+ # _truncate_fresh always includes a continuation hint, so both fields are always present.
+ new_notice = re.sub(r"covers the next \d+ bytes", f"covers the next {max_bytes} bytes", old_notice)
+ new_notice = re.sub(r"start_line=\d+ to read more", f"start_line={next_line} to read more", new_notice)
+
+ return result + TRUNCATION_NOTICE_MARKER + new_notice
+
+
+def truncate_text_output(
+ text: str,
+ start_line: int = 1,
+ total_lines: int = 0,
+ max_bytes: int = DEFAULT_MAX_BYTES,
+ file_path: str | None = None,
+ encoding: str = "utf-8",
+) -> str:
+ """Truncate file output by bytes with line integrity.
+
+ If text is under byte limit, return as-is.
+ If over limit, truncate at the last complete line that fits,
+ allowing the next read to start from a fresh line.
+
+ Dispatches to :func:`_truncate_fresh` for text seen for the first time, or to
+ :func:`_retruncate` when the text already contains a TRUNCATION_NOTICE_MARKER
+ from a previous pass.
+
+ Args:
+ text: The output text to truncate.
+ start_line: The starting line number (1-based). Ignored when text already
+ contains a truncation notice (values are parsed from the notice instead).
+ total_lines: Total lines in the original file. Ignored when text already
+ contains a truncation notice (values are parsed from the notice instead).
+ max_bytes: Maximum size in bytes.
+ file_path: Optional file path to include in the truncation notice.
+ encoding: Character encoding used for byte-length calculation and decoding.
+
+ Returns:
+ Truncated text with notice if truncated.
+ """
+ if not text:
+ return text
+ if max_bytes <= 0:
+ return text
+
+ try:
+ if TRUNCATION_NOTICE_MARKER in text:
+ return _retruncate(text, max_bytes=max_bytes, encoding=encoding)
+ else:
+ return _truncate_fresh(
+ text,
+ start_line=start_line,
+ total_lines=total_lines,
+ max_bytes=max_bytes,
+ file_path=file_path,
+ encoding=encoding,
+ )
+ except Exception:
+ logger.warning("truncate_text_output failed, returning original text", exc_info=True)
+ return text
+
+
+def read_file_safe(file_path: str, max_bytes: int = MAX_FILE_READ_BYTES) -> str:
+ """Read file with Unicode error handling and memory protection.
+
+ Args:
+ file_path: Path to the file.
+ max_bytes: Maximum bytes to read into memory (default 1GB).
+
+ Returns:
+ File content as string (up to max_bytes).
+ """
+ try:
+ with open(file_path, "r", encoding="utf-8") as f:
+ return f.read(max_bytes)
+ except UnicodeDecodeError:
+ with open(file_path, "r", encoding="utf-8", errors="ignore") as f:
+ return f.read(max_bytes)
diff --git a/reme/memory/tools/record/__init__.py b/reme/memory/tools/record/__init__.py
deleted file mode 100644
index e69de29b..00000000
diff --git a/reme/memory/tools/__init__.py b/reme/memory/vector_tools/__init__.py
similarity index 92%
rename from reme/memory/tools/__init__.py
rename to reme/memory/vector_tools/__init__.py
index af9b851f..fd5e721f 100644
--- a/reme/memory/tools/__init__.py
+++ b/reme/memory/vector_tools/__init__.py
@@ -3,8 +3,6 @@
from .base_memory_tool import BaseMemoryTool
# chunk tools
-from .chunk.memory_get import MemoryGet
-from .chunk.memory_search import MemorySearch
from .delegate_task import DelegateTask
# history tools
@@ -36,9 +34,6 @@ __all__ = [
# base
"BaseMemoryTool",
"DelegateTask",
- # chunk tools
- "MemoryGet",
- "MemorySearch",
# history tools
"AddHistory",
"ReadHistory",
diff --git a/reme/memory/tools/base_memory_tool.py b/reme/memory/vector_tools/base_memory_tool.py
similarity index 100%
rename from reme/memory/tools/base_memory_tool.py
rename to reme/memory/vector_tools/base_memory_tool.py
diff --git a/reme/memory/tools/delegate_task.py b/reme/memory/vector_tools/delegate_task.py
similarity index 100%
rename from reme/memory/tools/delegate_task.py
rename to reme/memory/vector_tools/delegate_task.py
diff --git a/cookbook/__init__.py b/reme/memory/vector_tools/history/__init__.py
similarity index 100%
rename from cookbook/__init__.py
rename to reme/memory/vector_tools/history/__init__.py
diff --git a/reme/memory/tools/history/add_history.py b/reme/memory/vector_tools/history/add_history.py
similarity index 100%
rename from reme/memory/tools/history/add_history.py
rename to reme/memory/vector_tools/history/add_history.py
diff --git a/reme/memory/tools/history/read_history.py b/reme/memory/vector_tools/history/read_history.py
similarity index 100%
rename from reme/memory/tools/history/read_history.py
rename to reme/memory/vector_tools/history/read_history.py
diff --git a/reme/memory/tools/history/read_history_v2.py b/reme/memory/vector_tools/history/read_history_v2.py
similarity index 100%
rename from reme/memory/tools/history/read_history_v2.py
rename to reme/memory/vector_tools/history/read_history_v2.py
diff --git a/cookbook/appworld/__init__.py b/reme/memory/vector_tools/profiles/__init__.py
similarity index 100%
rename from cookbook/appworld/__init__.py
rename to reme/memory/vector_tools/profiles/__init__.py
diff --git a/reme/memory/tools/profiles/add_draft_and_read_all_profiles.py b/reme/memory/vector_tools/profiles/add_draft_and_read_all_profiles.py
similarity index 100%
rename from reme/memory/tools/profiles/add_draft_and_read_all_profiles.py
rename to reme/memory/vector_tools/profiles/add_draft_and_read_all_profiles.py
diff --git a/reme/memory/tools/profiles/add_profile.py b/reme/memory/vector_tools/profiles/add_profile.py
similarity index 100%
rename from reme/memory/tools/profiles/add_profile.py
rename to reme/memory/vector_tools/profiles/add_profile.py
diff --git a/reme/memory/tools/profiles/delete_profile.py b/reme/memory/vector_tools/profiles/delete_profile.py
similarity index 100%
rename from reme/memory/tools/profiles/delete_profile.py
rename to reme/memory/vector_tools/profiles/delete_profile.py
diff --git a/reme/memory/tools/profiles/profile_handler.py b/reme/memory/vector_tools/profiles/profile_handler.py
similarity index 100%
rename from reme/memory/tools/profiles/profile_handler.py
rename to reme/memory/vector_tools/profiles/profile_handler.py
diff --git a/reme/memory/tools/profiles/read_all_profiles.py b/reme/memory/vector_tools/profiles/read_all_profiles.py
similarity index 100%
rename from reme/memory/tools/profiles/read_all_profiles.py
rename to reme/memory/vector_tools/profiles/read_all_profiles.py
diff --git a/reme/memory/tools/profiles/update_profile.py b/reme/memory/vector_tools/profiles/update_profile.py
similarity index 100%
rename from reme/memory/tools/profiles/update_profile.py
rename to reme/memory/vector_tools/profiles/update_profile.py
diff --git a/reme/memory/tools/profiles/update_profiles_v1.py b/reme/memory/vector_tools/profiles/update_profiles_v1.py
similarity index 100%
rename from reme/memory/tools/profiles/update_profiles_v1.py
rename to reme/memory/vector_tools/profiles/update_profiles_v1.py
diff --git a/cookbook/bfcl/__init__.py b/reme/memory/vector_tools/record/__init__.py
similarity index 100%
rename from cookbook/bfcl/__init__.py
rename to reme/memory/vector_tools/record/__init__.py
diff --git a/reme/memory/tools/record/add_and_retrieve_similar_memory.py b/reme/memory/vector_tools/record/add_and_retrieve_similar_memory.py
similarity index 100%
rename from reme/memory/tools/record/add_and_retrieve_similar_memory.py
rename to reme/memory/vector_tools/record/add_and_retrieve_similar_memory.py
diff --git a/reme/memory/tools/record/add_draft_and_retrieve_similar_memory.py b/reme/memory/vector_tools/record/add_draft_and_retrieve_similar_memory.py
similarity index 100%
rename from reme/memory/tools/record/add_draft_and_retrieve_similar_memory.py
rename to reme/memory/vector_tools/record/add_draft_and_retrieve_similar_memory.py
diff --git a/reme/memory/tools/record/add_memory.py b/reme/memory/vector_tools/record/add_memory.py
similarity index 100%
rename from reme/memory/tools/record/add_memory.py
rename to reme/memory/vector_tools/record/add_memory.py
diff --git a/reme/memory/tools/record/delete_memory.py b/reme/memory/vector_tools/record/delete_memory.py
similarity index 100%
rename from reme/memory/tools/record/delete_memory.py
rename to reme/memory/vector_tools/record/delete_memory.py
diff --git a/reme/memory/tools/record/memory_handler.py b/reme/memory/vector_tools/record/memory_handler.py
similarity index 100%
rename from reme/memory/tools/record/memory_handler.py
rename to reme/memory/vector_tools/record/memory_handler.py
diff --git a/reme/memory/tools/record/retrieve_memory.py b/reme/memory/vector_tools/record/retrieve_memory.py
similarity index 100%
rename from reme/memory/tools/record/retrieve_memory.py
rename to reme/memory/vector_tools/record/retrieve_memory.py
diff --git a/reme/memory/tools/record/retrieve_recent_memory.py b/reme/memory/vector_tools/record/retrieve_recent_memory.py
similarity index 100%
rename from reme/memory/tools/record/retrieve_recent_memory.py
rename to reme/memory/vector_tools/record/retrieve_recent_memory.py
diff --git a/reme/memory/tools/record/update_memory.py b/reme/memory/vector_tools/record/update_memory.py
similarity index 100%
rename from reme/memory/tools/record/update_memory.py
rename to reme/memory/vector_tools/record/update_memory.py
diff --git a/reme/memory/tools/record/update_memory_v1.py b/reme/memory/vector_tools/record/update_memory_v1.py
similarity index 100%
rename from reme/memory/tools/record/update_memory_v1.py
rename to reme/memory/vector_tools/record/update_memory_v1.py
diff --git a/reme/memory/tools/record/update_memory_v2.py b/reme/memory/vector_tools/record/update_memory_v2.py
similarity index 100%
rename from reme/memory/tools/record/update_memory_v2.py
rename to reme/memory/vector_tools/record/update_memory_v2.py
diff --git a/reme/reme.py b/reme/reme.py
index 6c5525ee..5e2aa5b2 100644
--- a/reme/reme.py
+++ b/reme/reme.py
@@ -7,7 +7,7 @@ from .config import ReMeConfigParser
from .core import Application
from .core.enumeration import MemoryType, Role
from .core.schema import Message, MemoryNode
-from .memory.tools import (
+from .memory.vector_tools import (
AddDraftAndRetrieveSimilarMemory,
AddHistory,
AddMemory,
@@ -17,8 +17,8 @@ from .memory.tools import (
RetrieveMemory,
UpdateProfilesV1,
)
-from .memory.tools.profiles.profile_handler import ProfileHandler
-from .memory.tools.record.memory_handler import MemoryHandler
+from .memory.vector_tools.profiles.profile_handler import ProfileHandler
+from .memory.vector_tools.record.memory_handler import MemoryHandler
from .memory.vector_based import (
BaseMemoryAgent,
PersonalRetriever,
@@ -162,6 +162,28 @@ class ReMe(Application):
return memory_type, memory_target
+ def _ensure_started(self) -> None:
+ """Ensure memory operations run only after services are initialized."""
+ if not self._started:
+ raise RuntimeError("ReMe is not started. Call `await reme.start()` before using memory APIs.")
+
+ @staticmethod
+ def _unwrap_memory_result(
+ result: str | dict,
+ operation_name: str,
+ return_dict: bool,
+ ) -> str | dict:
+ """Normalize memory API results and fail loudly on swallowed inner errors."""
+ if not isinstance(result, dict):
+ raise RuntimeError(f"{operation_name} failed before producing a structured result: {result}")
+
+ if "answer" not in result:
+ raise RuntimeError(f"{operation_name} returned an invalid result payload: missing 'answer'")
+
+ if return_dict:
+ return result
+ return result["answer"]
+
async def summarize_memory(
self,
messages: list[Message | dict],
@@ -173,10 +195,12 @@ class ReMe(Application):
version: str = "default",
retrieve_top_k: int = 20,
return_dict: bool = False,
+ raise_exception: bool = False,
llm_config_name: str = "default",
**kwargs,
) -> str | dict:
"""Summarize personal, procedural and tool memories for the given context."""
+ self._ensure_started()
format_messages: list[Message] = []
for message in messages:
if isinstance(message, dict):
@@ -185,19 +209,21 @@ class ReMe(Application):
format_messages.append(message)
if version == "default":
- personal_summarizer_tools = [
+ personal_summarizer_tools: list = [
AddDraftAndRetrieveSimilarMemory(
enable_thinking_params=enable_thinking_params,
enable_memory_target=False,
enable_when_to_use=False,
enable_multiple=True,
top_k=retrieve_top_k,
+ raise_exception=raise_exception,
),
AddMemory(
enable_thinking_params=enable_thinking_params,
enable_memory_target=False,
enable_when_to_use=False,
enable_multiple=True,
+ raise_exception=raise_exception,
),
]
if self.enable_profile:
@@ -207,18 +233,21 @@ class ReMe(Application):
enable_thinking_params=False,
enable_memory_target=False,
profile_dir=self.profile_dir,
+ raise_exception=raise_exception,
),
UpdateProfilesV1(
enable_thinking_params=enable_thinking_params,
enable_memory_target=False,
enable_multiple=True,
profile_dir=self.profile_dir,
+ raise_exception=raise_exception,
),
],
)
personal_summarizer: BaseMemoryAgent = PersonalSummarizer(
llm=llm_config_name,
tools=personal_summarizer_tools,
+ raise_exception=raise_exception,
)
else:
@@ -233,14 +262,17 @@ class ReMe(Application):
enable_when_to_use=False,
enable_multiple=True,
top_k=retrieve_top_k,
+ raise_exception=raise_exception,
),
AddMemory(
enable_thinking_params=enable_thinking_params,
enable_memory_target=False,
enable_when_to_use=False,
enable_multiple=True,
+ raise_exception=raise_exception,
),
],
+ raise_exception=raise_exception,
)
tool_summarizer: BaseMemoryAgent = ToolSummarizer(
llm=llm_config_name,
@@ -251,14 +283,17 @@ class ReMe(Application):
enable_when_to_use=False,
enable_multiple=True,
top_k=retrieve_top_k,
+ raise_exception=raise_exception,
),
AddMemory(
enable_thinking_params=enable_thinking_params,
enable_memory_target=False,
enable_when_to_use=False,
enable_multiple=True,
+ raise_exception=raise_exception,
),
],
+ raise_exception=raise_exception,
)
memory_agents = []
@@ -306,7 +341,11 @@ class ReMe(Application):
memory_agents = [personal_summarizer, procedural_summarizer, tool_summarizer]
reme_summarizer: BaseMemoryAgent = ReMeSummarizer(
- tools=[AddHistory(), DelegateTask(memory_agents=memory_agents)],
+ tools=[
+ AddHistory(raise_exception=raise_exception),
+ DelegateTask(memory_agents=memory_agents, raise_exception=raise_exception),
+ ],
+ raise_exception=raise_exception,
)
result = await reme_summarizer.call(
@@ -317,10 +356,7 @@ class ReMe(Application):
**kwargs,
)
- if return_dict:
- return result
- else:
- return result["answer"]
+ return self._unwrap_memory_result(result, "summarize_memory", return_dict)
async def retrieve_memory(
self,
@@ -335,10 +371,12 @@ class ReMe(Application):
retrieve_top_k: int = 20,
enable_time_filter: bool = True,
return_dict: bool = False,
+ raise_exception: bool = False,
llm_config_name: str = "default",
**kwargs,
) -> str | dict:
"""Retrieve relevant personal, procedural and tool memories for a query."""
+ self._ensure_started()
if version == "default":
personal_retriever_tools = []
@@ -348,6 +386,7 @@ class ReMe(Application):
enable_thinking_params=False,
enable_memory_target=False,
profile_dir=self.profile_dir,
+ raise_exception=raise_exception,
),
)
personal_retriever_tools.extend(
@@ -357,16 +396,19 @@ class ReMe(Application):
enable_thinking_params=enable_thinking_params,
enable_time_filter=enable_time_filter,
enable_multiple=True,
+ raise_exception=raise_exception,
),
ReadHistory(
enable_thinking_params=enable_thinking_params,
enable_multiple=True,
+ raise_exception=raise_exception,
),
],
)
personal_retriever: BaseMemoryAgent = PersonalRetriever(
llm=llm_config_name,
tools=personal_retriever_tools,
+ raise_exception=raise_exception,
)
else:
raise NotImplementedError(f"version={version} is not supported")
@@ -379,12 +421,15 @@ class ReMe(Application):
enable_thinking_params=enable_thinking_params,
enable_time_filter=False,
enable_multiple=True,
+ raise_exception=raise_exception,
),
ReadHistory(
enable_thinking_params=enable_thinking_params,
enable_multiple=True,
+ raise_exception=raise_exception,
),
],
+ raise_exception=raise_exception,
)
tool_retriever: BaseMemoryAgent = ToolRetriever(
llm=llm_config_name,
@@ -394,12 +439,15 @@ class ReMe(Application):
enable_thinking_params=enable_thinking_params,
enable_time_filter=False,
enable_multiple=True,
+ raise_exception=raise_exception,
),
ReadHistory(
enable_thinking_params=enable_thinking_params,
enable_multiple=True,
+ raise_exception=raise_exception,
),
],
+ raise_exception=raise_exception,
)
memory_agents = []
@@ -444,7 +492,8 @@ class ReMe(Application):
memory_agents = [personal_retriever, procedural_retriever, tool_retriever]
reme_retriever: BaseMemoryAgent = ReMeRetriever(
- tools=[DelegateTask(memory_agents=memory_agents)],
+ tools=[DelegateTask(memory_agents=memory_agents, raise_exception=raise_exception)],
+ raise_exception=raise_exception,
)
result = await reme_retriever.call(
@@ -456,10 +505,7 @@ class ReMe(Application):
**kwargs,
)
- if return_dict:
- return result
- else:
- return result["answer"]
+ return self._unwrap_memory_result(result, "retrieve_memory", return_dict)
async def add_memory(
self,
diff --git a/reme/reme_cli.py b/reme/reme_cli.py
new file mode 100644
index 00000000..5d086532
--- /dev/null
+++ b/reme/reme_cli.py
@@ -0,0 +1,192 @@
+"""ReMe File System"""
+
+import asyncio
+import sys
+from pathlib import Path
+
+from prompt_toolkit import PromptSession
+
+from .config import ReMeConfigParser
+from .core import Application
+
+from .core.utils import play_horse_easter_egg
+from .memory.file_based.components import CliAgent
+
+
+class ReMeCli(Application):
+ """ReMe Cli"""
+
+ def __init__(
+ self,
+ *args,
+ working_dir: str = ".reme",
+ config_path: str = "cli",
+ enable_logo: bool = True,
+ log_to_console: bool = True,
+ llm_api_key: str | None = None,
+ llm_base_url: str | None = None,
+ embedding_api_key: str | None = None,
+ embedding_base_url: str | None = None,
+ default_as_llm_config: dict | None = None,
+ default_embedding_model_config: dict | None = None,
+ default_file_store_config: dict | None = None,
+ default_token_counter_config: dict | None = None,
+ default_file_watcher_config: dict | None = None,
+ context_window_tokens: int = 128000,
+ reserve_tokens: int = 36000,
+ keep_recent_tokens: int = 20000,
+ vector_weight: float = 0.7,
+ candidate_multiplier: float = 3.0,
+ **kwargs,
+ ):
+ """Initialize ReMe with config."""
+ working_path = Path(working_dir)
+ working_path.mkdir(parents=True, exist_ok=True)
+ memory_path = working_path / "memory"
+ memory_path.mkdir(parents=True, exist_ok=True)
+ self.working_dir: str = str(working_path.absolute())
+
+ default_file_watcher_config = default_file_watcher_config or {}
+ if not default_file_watcher_config.get("watch_paths", None):
+ default_file_watcher_config["watch_paths"] = [
+ str(working_path / "MEMORY.md"),
+ str(working_path / "memory.md"),
+ str(memory_path),
+ ]
+ super().__init__(
+ *args,
+ llm_api_key=llm_api_key,
+ llm_base_url=llm_base_url,
+ embedding_api_key=embedding_api_key,
+ embedding_base_url=embedding_base_url,
+ working_dir=working_dir,
+ config_path=config_path,
+ enable_logo=enable_logo,
+ log_to_console=log_to_console,
+ parser=ReMeConfigParser,
+ default_as_llm_config=default_as_llm_config,
+ default_embedding_model_config=default_embedding_model_config,
+ default_file_store_config=default_file_store_config,
+ default_token_counter_config=default_token_counter_config,
+ default_file_watcher_config=default_file_watcher_config,
+ **kwargs,
+ )
+
+ self.service_config.metadata.setdefault("context_window_tokens", context_window_tokens)
+ self.service_config.metadata.setdefault("reserve_tokens", reserve_tokens)
+ self.service_config.metadata.setdefault("keep_recent_tokens", keep_recent_tokens)
+ self.service_config.metadata.setdefault("vector_weight", vector_weight)
+ self.service_config.metadata.setdefault("candidate_multiplier", candidate_multiplier)
+
+ self.commands = {
+ "/new": "Create a new conversation.",
+ "/compact": "Compact messages into a summary.",
+ "/exit": "Exit the application.",
+ "/clear": "Clear the history.",
+ "/help": "Show help.",
+ "/horse": "A surprise.",
+ }
+
+ async def chat_with_remy(self, **kwargs):
+ """Interactive CLI chat with Remy using simple streaming output."""
+ language = self.service_config.language
+ print(f"ReMe language={language or 'default'}")
+
+ cli_agent = CliAgent(
+ vector_weight=self.service_config.metadata["vector_weight"],
+ candidate_multiplier=self.service_config.metadata["candidate_multiplier"],
+ context_window_tokens=self.service_config.metadata["context_window_tokens"],
+ reserve_tokens=self.service_config.metadata["reserve_tokens"],
+ keep_recent_tokens=self.service_config.metadata["keep_recent_tokens"],
+ working_dir=self.working_dir,
+ language=language,
+ **kwargs,
+ )
+ session = PromptSession()
+
+ # Print welcome banner
+ print("\n========================================")
+ print(" Welcome to Remy Chat!")
+ print("========================================\n")
+
+ while True:
+ try:
+ # Get user input (async)
+ user_input = await session.prompt_async("You: ")
+ user_input = user_input.strip()
+ if not user_input:
+ continue
+
+ # Handle commands
+ if user_input == "/exit":
+ break
+
+ if user_input == "/new":
+ result = await cli_agent.new()
+ print(f"{result}\nConversation reset\n")
+ continue
+
+ if user_input == "/compact":
+ result = await cli_agent.compact(force_compact=True)
+ print(f"{result}\nHistory compacted.\n")
+ continue
+
+ if user_input == "/history":
+ result = cli_agent.format_history()
+ print(f"Formated History:\n{result}\n")
+ continue
+
+ if user_input == "/clear":
+ cli_agent.messages.clear()
+ print("History cleared.\n")
+ continue
+
+ if user_input == "/help":
+ print("\nCommands:")
+ for command, description in self.commands.items():
+ print(f" {command}: {description}")
+ continue
+
+ if user_input == "/horse":
+ play_horse_easter_egg()
+ continue
+
+ try:
+ await cli_agent.call(
+ query=user_input,
+ service_context=self.service_context,
+ )
+ except Exception as e:
+ print(f"\nStream error: {e}")
+
+ # End current streaming line
+ print("\n")
+ print("----------------------------------------\n")
+
+ except EOFError:
+ break
+ except KeyboardInterrupt:
+ print("\nInterrupted.")
+ break
+ except Exception as e:
+ print(f"Error: {e}")
+ import traceback
+
+ traceback.print_exc()
+
+ print("\nGoodbye!\n")
+
+
+async def async_main():
+ """Main function for testing the ReMeFs CLI."""
+ async with ReMeCli(*sys.argv[1:], log_to_console=False) as reme:
+ await reme.chat_with_remy()
+
+
+def main():
+ """Main function for testing the ReMeFs CLI."""
+ asyncio.run(async_main())
+
+
+if __name__ == "__main__":
+ main()
diff --git a/reme/reme_light.py b/reme/reme_light.py
index 26500833..91e738ed 100644
--- a/reme/reme_light.py
+++ b/reme/reme_light.py
@@ -15,52 +15,53 @@ Key Features:
"""
import asyncio
-import logging
-import os
-import platform
from pathlib import Path
from agentscope.formatter import FormatterBase
from agentscope.message import Msg, TextBlock
-from agentscope.model import ChatModelBase, OpenAIChatModel
+from agentscope.model import ChatModelBase
from agentscope.token import HuggingFaceTokenCounter
from agentscope.tool import Toolkit, ToolResponse
from .config import ReMeConfigParser
from .core import Application
-from .memory.file_based import Compactor, Summarizer, ToolResultCompactor, ReMeInMemoryMemory, ReMeChatFormatter
-from .memory.file_based.utils import get_token_counter
-from .memory.tools import MemorySearch
-from .core.utils import load_env
+from .core.utils import get_logger
+from .memory.file_based import ReMeInMemoryMemory
+from .memory.file_based.components import (
+ Compactor,
+ ContextChecker,
+ Summarizer,
+ ToolResultCompactor,
+)
+from .memory.file_based.tools import FileIO, MemorySearch
+from .memory.file_based.utils import AsMsgHandler
-logger = logging.getLogger(__name__)
+logger = get_logger()
class ReMeLight(Application):
"""
- ReMe Light Application Class
+ ReMe Light Application Class.
- A specialized application class that extends ReMe's core Application framework
- with advanced memory management capabilities. This class is designed to handle
- long-running conversations by providing intelligent memory compaction,
- summarization, and semantic search features.
+ A lightweight memory-enabled application that provides semantic search,
+ memory compaction, summarization, and tool result management capabilities.
+ Built on top of the core Application framework with integrated vector store
+ and file-based memory management.
+
+ This class is designed for applications requiring:
+ - Long conversation memory management with automatic compaction
+ - Semantic search over stored memories using hybrid vector/text search
+ - Background summarization of conversation history
+ - Automatic cleanup of expired tool results
Attributes:
- working_path (Path): Absolute path to the working directory for storing data
- memory_path (Path): Path to the memory storage directory
- tool_result_path (Path): Path to store large tool result files
- chat_model (ChatModelBase): Language model for generating summaries and processing
- formatter (FormatterBase): Formatter for structuring model inputs/outputs
- token_counter (HuggingFaceTokenCounter): Token counting utility for length management
- toolkit (Toolkit): Collection of tools available to the application
- max_input_length (int): Maximum allowed input length in tokens
- memory_compact_threshold (int): Threshold at which memory compaction triggers
- language (str): Language code for localization ("zh" for Chinese, empty for English)
- vector_weight (float): Weight for vector search in hybrid search (0.0-1.0)
- candidate_multiplier (float): Multiplier for candidate retrieval in search
- tool_result_threshold (int): Size threshold for tool result compaction
- retention_days (int): Number of days to retain tool result files
- summary_tasks (list[asyncio.Task]): List of background summarization tasks
+ working_path (Path): Absolute path to the working directory.
+ memory_path (Path): Path to the memory storage directory.
+ tool_result_path (Path): Path to the tool result storage directory.
+ dialog_path (Path): Path to the dialog storage directory for raw conversation records.
+ vector_weight (float): Weight for vector search in hybrid search (0-1).
+ candidate_multiplier (float): Multiplier for candidate retrieval count.
+ summary_tasks (list[asyncio.Task]): List of active background summary tasks.
"""
def __init__(
@@ -70,81 +71,83 @@ class ReMeLight(Application):
llm_base_url: str | None = None,
embedding_api_key: str | None = None,
embedding_base_url: str | None = None,
- chat_model: ChatModelBase | None = None,
- formatter: FormatterBase | None = None,
- token_counter: HuggingFaceTokenCounter | None = None,
- toolkit: Toolkit | None = None,
- max_input_length: int = 128000,
- memory_compact_ratio: float = 0.7,
- language: str = "zh",
+ default_as_llm_config: dict | None = None,
+ default_embedding_model_config: dict | None = None,
+ default_file_store_config: dict | None = None,
+ default_file_watcher_config: dict | None = None,
vector_weight: float = 0.7,
candidate_multiplier: float = 3.0,
- tool_result_threshold: int = 1000,
- retention_days: int = 7,
+ enable_load_env: bool = False,
):
+ """
+ Initialize the ReMeLight application.
+
+ Sets up the working directory structure, configures API connections,
+ and initializes memory management components.
+
+ Args:
+ working_dir (str): Base directory for all application data storage.
+ Defaults to ".reme". Will be created if it doesn't exist.
+ llm_api_key (str | None): API key for the language model service.
+ If None, will attempt to use environment variables.
+ llm_base_url (str | None): Base URL for the language model API endpoint.
+ If None, will use the default endpoint.
+ embedding_api_key (str | None): API key for the embedding model service.
+ If None, will attempt to use environment variables.
+ embedding_base_url (str | None): Base URL for the embedding API endpoint.
+ If None, will use the default endpoint.
+ default_as_llm_config (dict | None): Default configuration dictionary
+ for AgentScope language model. Overrides default settings.
+ default_embedding_model_config (dict | None): Default configuration
+ dictionary for the embedding model.
+ default_file_store_config (dict | None): Default configuration
+ dictionary for the file storage backend.
+ default_file_watcher_config (dict | None): Default configuration
+ dictionary for the file watcher. If ``watch_paths`` is included,
+ it is used as-is. Otherwise the built-in watch paths (MEMORY.md,
+ memory.md, and the memory directory) are used.
+ vector_weight (float): Weight assigned to vector similarity search
+ in hybrid search operations. Range [0.0, 1.0], default 0.7.
+ Higher values prioritize semantic similarity over keyword matching.
+ candidate_multiplier (float): Multiplier applied to max_results when
+ retrieving candidates for re-ranking. Default 3.0 means 3x more
+ candidates are retrieved than the final result count.
+ enable_load_env (bool): Whether to load environment variables from
+ .env file. Defaults to False.
+
+ Note:
+ The following directory structure will be created:
+ - {working_dir}/ - Root working directory
+ - {working_dir}/memory/ - Memory storage files
+ - {working_dir}/tool_result/ - Compacted tool result files
+ - {working_dir}/dialog/ - Raw conversation records
+ """
# Initialize working directory structure
- # All application data will be stored under this path
self.working_path = Path(working_dir).absolute()
self.working_path.mkdir(parents=True, exist_ok=True)
-
- # Create memory storage directory for persistent memory files
self.memory_path = self.working_path / "memory"
self.memory_path.mkdir(parents=True, exist_ok=True)
-
- # Create tool result directory for storing large tool outputs
self.tool_result_path = self.working_path / "tool_result"
self.tool_result_path.mkdir(parents=True, exist_ok=True)
+ self.dialog_path = self.working_path / "dialog"
+ self.dialog_path.mkdir(parents=True, exist_ok=True)
- # Initialize runtime parameters (will be updated via update_params)
- self.max_input_length: int = 0
- self.memory_compact_threshold: int = 0
- self.language: str = ""
-
- # Apply initial parameter configuration
- self.update_params(
- max_input_length=max_input_length,
- memory_compact_ratio=memory_compact_ratio,
- language=language,
- )
-
- # Store configuration parameters
self.vector_weight: float = vector_weight
self.candidate_multiplier: float = candidate_multiplier
- self.tool_result_threshold: int = tool_result_threshold
- self.retention_days: int = retention_days
- load_env()
-
- llm_model_name = self._safe_str("LLM_MODEL_NAME", "")
- embedding_model_name = self._safe_str("EMBEDDING_MODEL_NAME", "")
- embedding_dimensions = self._safe_int("EMBEDDING_DIMENSIONS", 1024)
- embedding_cache_enabled = self._safe_str("EMBEDDING_CACHE_ENABLED", "true").lower() == "true"
- embedding_max_cache_size = self._safe_int("EMBEDDING_MAX_CACHE_SIZE", 2000)
- embedding_max_input_length = self._safe_int("EMBEDDING_MAX_INPUT_LENGTH", 8192)
- embedding_max_batch_size = self._safe_int("EMBEDDING_MAX_BATCH_SIZE", 10)
-
- # Determine if vector search should be enabled based on configuration
- # Vector search requires either an API key or a local model name
- vector_enabled = bool(embedding_api_key) or bool(embedding_model_name)
- if vector_enabled:
- logger.info("Vector search enabled.")
+ # Build the file watcher config: use provided watch_paths if given, otherwise use defaults
+ _default_watch_paths = [
+ str(self.working_path / "MEMORY.md"),
+ str(self.working_path / "memory.md"),
+ str(self.memory_path),
+ ]
+ if default_file_watcher_config and default_file_watcher_config.get("watch_paths"):
+ _merged_file_watcher_config = default_file_watcher_config
else:
- logger.warning(
- "Vector search disabled. Memory search functionality will be restricted. "
- "To enable, configure: EMBEDDING_API_KEY, EMBEDDING_BASE_URL, EMBEDDING_MODEL_NAME.",
- )
-
- # Check if full-text search (FTS) is enabled via environment variable
- fts_enabled = os.environ.get("FTS_ENABLED", "true").lower() == "true"
-
- # Determine the memory store backend to use
- # "auto" selects based on platform (local for Windows, chroma otherwise)
- # Supported: auto, local, chroma, sqlite, seekdb
- memory_store_backend = os.environ.get("MEMORY_STORE_BACKEND", "auto")
- if memory_store_backend == "auto":
- memory_backend = "local" if platform.system() == "Windows" else "chroma"
- else:
- memory_backend = memory_store_backend
+ _merged_file_watcher_config = {
+ **(default_file_watcher_config or {}),
+ "watch_paths": _default_watch_paths,
+ }
# Initialize the parent Application class with comprehensive configuration
super().__init__(
@@ -156,155 +159,44 @@ class ReMeLight(Application):
config_path="light",
enable_logo=False,
log_to_console=False,
+ enable_load_env=enable_load_env,
parser=ReMeConfigParser,
- default_embedding_model_config={
- "model_name": embedding_model_name,
- "dimensions": embedding_dimensions,
- "enable_cache": embedding_cache_enabled,
- "use_dimensions": False,
- "max_cache_size": embedding_max_cache_size,
- "max_input_length": embedding_max_input_length,
- "max_batch_size": embedding_max_batch_size,
- },
- default_file_store_config={
- "backend": memory_backend,
- "store_name": "copaw",
- "vector_enabled": vector_enabled,
- "fts_enabled": fts_enabled,
- },
- default_file_watcher_config={
- "watch_paths": [
- str(self.working_path / "MEMORY.md"),
- str(self.working_path / "memory.md"),
- str(self.memory_path),
- ],
- },
+ default_as_llm_config=default_as_llm_config,
+ default_embedding_model_config=default_embedding_model_config,
+ default_file_store_config=default_file_store_config,
+ default_file_watcher_config=_merged_file_watcher_config,
)
- if chat_model is not None:
- self.chat_model: ChatModelBase = chat_model
- else:
- # add more params later
- self.chat_model = OpenAIChatModel(
- api_key=os.environ["LLM_API_KEY"],
- client_kwargs={"base_url": os.environ["LLM_BASE_URL"]},
- model_name=llm_model_name,
- )
-
- if token_counter is not None:
- self.token_counter: HuggingFaceTokenCounter = token_counter
- else:
- self.token_counter = get_token_counter()
-
- if formatter is not None:
- self.formatter: FormatterBase = formatter
- else:
- self.formatter = ReMeChatFormatter(token_counter=self.token_counter)
- self.toolkit: Toolkit | None = toolkit
-
# Initialize list to track background summarization tasks
self.summary_tasks: list[asyncio.Task] = []
- def update_params(
- self,
- max_input_length: int,
- memory_compact_ratio: float,
- language: str,
- ):
- """
- Update runtime parameters for memory management.
-
- This method allows dynamic adjustment of memory-related parameters during
- runtime. It recalculates the memory compaction threshold based on the
- new input length and compaction ratio.
-
- Args:
- max_input_length (int): New maximum input length in tokens
- memory_compact_ratio (float): Ratio at which to trigger compaction (0.0-1.0)
- language (str): Language code for localization ("zh" or other)
-
- Note:
- The memory_compact_threshold is calculated as:
- max_input_length * memory_compact_ratio * 0.9
- The 0.9 factor provides a safety margin before reaching the absolute limit
- """
- # Update the maximum allowed input length
- self.max_input_length = max_input_length
-
- # Calculate compaction threshold with safety margin
- # This ensures compaction happens before hitting the hard limit
- self.memory_compact_threshold = int(max_input_length * memory_compact_ratio * 0.9)
-
- # Set language for localization
- if language == "zh":
- self.language = "zh"
- else:
- self.language = ""
-
@staticmethod
- def _safe_str(key: str, default: str) -> str:
- """
- Safely retrieve a string value from an environment variable.
+ def calculate_memory_compact_threshold(max_input_length: float, compact_ratio: float) -> int:
+ """Calculate the memory compaction threshold based on input length and ratio.
Args:
- key (str): The name of the environment variable to retrieve
- default (str): The default value to return if the variable is not set
+ max_input_length: Maximum input length in tokens.
+ compact_ratio: Ratio of the input length to use as the threshold.
Returns:
- str: The value of the environment variable, or the default if not set
+ Computed compaction threshold as an integer.
"""
- return os.environ.get(key, default)
-
- @staticmethod
- def _safe_int(key: str, default: int) -> int:
- """
- Safely retrieve an integer value from an environment variable.
-
- This method handles cases where the environment variable is not set
- or contains a non-integer value by returning the specified default.
-
- Args:
- key (str): The name of the environment variable to retrieve
- default (int): The default value to return on failure or if not set
-
- Returns:
- int: The integer value of the environment variable, or the default
-
- Note:
- Logs a warning if the value exists but cannot be parsed as an integer
- """
- value = os.environ.get(key)
- if value is None:
- return default
-
- try:
- return int(value)
- except ValueError:
- logger.warning(f"Invalid int value '{value}' for key '{key}', using default {default}")
- return default
+ return int(max_input_length * compact_ratio * 0.95)
def _cleanup_tool_results(self) -> int:
"""
Clean up expired tool result files from the tool result directory.
This method removes tool result files that have exceeded the retention
- period specified during initialization. It helps manage disk space by
- automatically removing old, unused tool outputs.
+ period. It helps manage disk space by automatically removing old, unused
+ tool outputs.
Returns:
int: The number of files that were successfully deleted
-
- Note:
- Exceptions during cleanup are logged but do not raise errors,
- ensuring the application continues to function even if cleanup fails
"""
try:
- # Create a compactor instance with current configuration
- compactor = ToolResultCompactor(
- tool_result_dir=self.tool_result_path,
- tool_result_threshold=self.tool_result_threshold,
- retention_days=self.retention_days,
- )
+ # Create a compactor instance with default configuration
+ compactor = ToolResultCompactor(tool_result_dir=self.tool_result_path)
# Execute cleanup and return count of deleted files
return compactor.cleanup_expired_files()
except Exception as e:
@@ -316,19 +208,18 @@ class ReMeLight(Application):
"""
Start the application lifecycle.
- This method initializes the application by calling the parent class's
- start method and performs initial cleanup of expired tool result files.
+ Initializes all application components by calling the parent class start
+ method, then performs initial cleanup of expired tool result files.
Returns:
- The result from the parent class's start method
+ The result from the parent Application.start() method.
Note:
- Tool result cleanup runs after successful startup to ensure
- the application is fully initialized before performing maintenance
+ This method should be called before using any other application
+ functionality. It ensures all services are properly initialized.
"""
- # Initialize parent application components
result = await super().start()
- # Perform initial cleanup of old tool result files
+ # Perform initial cleanup of any expired tool result files
self._cleanup_tool_results()
return result
@@ -336,50 +227,70 @@ class ReMeLight(Application):
"""
Close the application and perform cleanup.
- This method performs final cleanup of expired tool result files before
- shutting down the application through the parent class's close method.
+ Performs final cleanup of expired tool result files and then shuts down
+ all application components by calling the parent class close method.
Returns:
- bool: True if shutdown was successful, False otherwise
+ bool: True if the application was closed successfully, False otherwise.
Note:
- Cleanup is performed before calling parent close to ensure
- all resources are available during the cleanup process
+ This method should be called when the application is no longer needed
+ to ensure proper resource cleanup and data persistence.
"""
- # Clean up tool results before shutting down
+ # Final cleanup of expired tool result files before shutdown
self._cleanup_tool_results()
- # Shutdown parent application components
return await super().close()
async def compact_tool_result(
self,
messages: list[Msg],
+ old_max_bytes: int = 3000,
+ recent_max_bytes: int = 100 * 1024,
+ retention_days: int = 3,
+ recent_n: int = 1,
) -> list[Msg]:
"""
Compact tool results by truncating large outputs and saving full content to files.
- This method processes a list of messages and identifies tool results that exceed
- the configured size threshold. Large tool outputs are truncated in the message
- list while their full content is saved to files for later retrieval.
+ This method processes a list of messages containing tool results and compacts
+ any that exceed the configured threshold. Large tool outputs are truncated
+ in the message while the full content is saved to separate files for later
+ retrieval if needed.
Args:
- messages (list[Msg]): List of messages to process for tool result compaction
+ messages (list[Msg]): List of messages potentially containing tool results
+ that may need compaction.
+ old_max_bytes (int): Byte threshold for old (non-recent) messages. Default 3000.
+ recent_max_bytes (int): Byte threshold for recent messages (trailing consecutive
+ tool-result messages). Default 100KB (102400 bytes). Content exceeding this
+ limit is saved to disk; the message retains the first 100KB with a
+ read_file-style truncation notice and the saved file path.
+ retention_days (int): Number of days to retain tool result files.
+ Default 3.
+ recent_n (int): Minimum number of most-recent tool-result messages to treat
+ as "recent" (using recent_max_bytes). The actual recent window is the
+ larger of this value and the trailing consecutive tool-result run.
+ Default 1.
Returns:
- list[Msg]: The processed message list with large tool results compacted
+ list[Msg]: The processed list of messages with large tool results compacted.
+ If an error occurs, returns the original unmodified messages.
Note:
- - Tool results below the threshold remain unchanged in the messages
- - Large results are replaced with truncated versions and file references
- - Expired files are cleaned up as part of the compaction process
- - If compaction fails, the original messages are returned unchanged
+ - Recent tool results (trailing consecutive tool-result messages) are truncated
+ to recent_max_bytes using read_file-style output with a file path hint.
+ - Old tool results are truncated to old_max_bytes bytes.
+ - Full content of truncated results is saved to tool_result_path.
+ - Expired files are automatically cleaned up during this operation.
"""
try:
# Create compactor with instance configuration
compactor = ToolResultCompactor(
tool_result_dir=self.tool_result_path,
- tool_result_threshold=self.tool_result_threshold,
- retention_days=self.retention_days,
+ retention_days=retention_days,
+ old_max_bytes=old_max_bytes,
+ recent_max_bytes=recent_max_bytes,
+ recent_n=recent_n,
)
# Execute compaction and get processed messages
@@ -395,38 +306,120 @@ class ReMeLight(Application):
logger.exception(f"Error compacting tool results: {e}")
return messages
- async def compact_memory(self, messages: list[Msg], previous_summary: str = "") -> str:
+ async def check_context(
+ self,
+ messages: list[Msg],
+ memory_compact_threshold: int,
+ memory_compact_reserve: int = 10000,
+ as_token_counter: str | HuggingFaceTokenCounter = "default",
+ ) -> tuple[list[Msg], list[Msg], bool]:
+ """
+ Check context size and determine if compaction is needed.
+
+ Analyzes the provided messages to determine if they exceed the configured
+ token threshold and splits them into two groups: messages that should be
+ compacted and messages to keep in context.
+
+ Args:
+ messages (list[Msg]): List of messages to check for context overflow.
+ memory_compact_threshold (int): Token count threshold for triggering
+ compaction. Messages exceeding this threshold will be split.
+ memory_compact_reserve (int): Token count to reserve for recent messages
+ to keep in context. Defaults to 10000 tokens.
+ as_token_counter (str | HuggingFaceTokenCounter): The token counter to use.
+
+ Returns:
+ tuple[list[Msg], list[Msg], bool]: A tuple containing:
+ - messages_to_compact (list[Msg]): Older messages that should
+ be compacted/summarized.
+ - messages_to_keep (list[Msg]): Recent messages to keep in context.
+ - is_valid (bool): True if the split is valid (tool calls aligned),
+ False if splitting would break conversation integrity.
+
+ Note:
+ - Returns ([], messages, True) if no compaction is needed.
+ - Ensures conversation pairs (user-assistant) are not split.
+ - is_valid=False indicates tool_use and tool_result are misaligned.
+ """
+ try:
+ checker = ContextChecker(
+ memory_compact_threshold=memory_compact_threshold,
+ memory_compact_reserve=memory_compact_reserve,
+ as_token_counter=as_token_counter,
+ )
+
+ return await checker.call(
+ messages=messages,
+ service_context=self.service_context,
+ )
+
+ except Exception as e:
+ logger.exception(f"Error checking context: {e}")
+ return [], messages, False
+
+ async def compact_memory(
+ self,
+ messages: list[Msg],
+ as_llm: str | ChatModelBase = "default",
+ as_llm_formatter: str | FormatterBase = "default",
+ as_token_counter: str | HuggingFaceTokenCounter = "default",
+ language: str = "zh",
+ max_input_length: float = 128 * 1024,
+ compact_ratio: float = 0.7,
+ previous_summary: str = "",
+ return_dict: bool = False,
+ add_thinking_block: bool = True,
+ extra_instruction: str = "",
+ ) -> str | dict:
"""
Compact a list of messages into a condensed summary.
- This method uses the Compactor to reduce the length of message history
- while preserving essential information. It's useful when conversation
- history approaches the maximum input length limit.
+ Uses the configured language model to generate a concise summary of the
+ provided messages. This is useful for reducing context window usage while
+ preserving important information from the conversation history.
Args:
- messages (list[Msg]): The list of messages to compact
- previous_summary (str): Optional previous summary to incorporate
- into the compaction process for continuity
+ messages (list[Msg]): List of messages to be compacted into a summary.
+ as_llm (str | ChatModelBase): Language model identifier or instance
+ to use for summarization. Defaults to "default".
+ as_llm_formatter (str | FormatterBase): Formatter for the language model.
+ Defaults to "default".
+ as_token_counter (str | HuggingFaceTokenCounter): Token counter for
+ measuring message length. Defaults to "default".
+ language (str): Language for the summary output. "zh" for Chinese,
+ any other value for English. Defaults to "zh".
+ max_input_length (float): Maximum input length in tokens for the model.
+ Defaults to 128K tokens.
+ compact_ratio (float): Ratio used to calculate compaction threshold.
+ Defaults to 0.7.
+ previous_summary (str): Previous summary to incorporate into the new
+ summary for continuity. Defaults to empty string.
+ return_dict (bool): If True, returns a dict with user_message,
+ history_compact, and is_valid. Defaults to False.
+ add_thinking_block (bool): If True, adds a thinking block to the summary.
+ extra_instruction (str): Optional additional instruction appended to the
+ compaction prompt. Use this to guide what information to keep or
+ remove. For example: "Remove debug logs and tool-call details. Keep
+ requirements, decisions, and pending tasks." Defaults to empty string
+ (no extra instruction, preserving default behavior).
Returns:
- str: A compacted summary of the messages, or empty string on failure
-
- Note:
- - Compaction uses the configured language model to generate summaries
- - The compaction threshold determines when compaction is triggered
- - If compaction fails, an empty string is returned
+ str | dict: The condensed summary string, or a dict containing
+ user_message, history_compact, and is_valid if return_dict=True.
+ Returns empty string or dict with empty values if an error occurred.
"""
try:
- # Initialize compactor with current configuration
compactor = Compactor(
- memory_compact_threshold=self.memory_compact_threshold,
- chat_model=self.chat_model,
- formatter=self.formatter,
- token_counter=self.token_counter,
- language=self.language,
+ memory_compact_threshold=self.calculate_memory_compact_threshold(max_input_length, compact_ratio),
+ as_llm=as_llm,
+ as_llm_formatter=as_llm_formatter,
+ as_token_counter=as_token_counter,
+ language=language if language == "zh" else "",
+ return_dict=return_dict,
+ add_thinking_block=add_thinking_block,
+ extra_instruction=extra_instruction,
)
- # Execute compaction with optional previous summary context
return await compactor.call(
messages=messages,
previous_summary=previous_summary,
@@ -434,69 +427,267 @@ class ReMeLight(Application):
)
except Exception as e:
- # Log error and return empty string to indicate failure
+ # Log error and return appropriate empty result
logger.exception(f"Error compacting memory: {e}")
+ if return_dict:
+ return {"user_message": str(e), "history_compact": str(e), "is_valid": False}
return ""
- async def summary_memory(self, messages: list[Msg]) -> str:
+ async def summary_memory(
+ self,
+ messages: list[Msg],
+ as_llm: str | ChatModelBase = "default",
+ as_llm_formatter: str | FormatterBase = "default",
+ as_token_counter: str | HuggingFaceTokenCounter = "default",
+ toolkit: Toolkit | None = None,
+ language: str = "zh",
+ max_input_length: float = 128 * 1024,
+ compact_ratio: float = 0.7,
+ timezone: str | None = None,
+ add_thinking_block: bool = True,
+ ) -> str:
"""
Generate a comprehensive summary of the given messages.
- This method uses the Summarizer to create a detailed summary of the
- conversation history, which can be stored as persistent memory. Unlike
- compaction, summarization aims to capture key information in a format
- suitable for long-term storage and retrieval.
+ Creates a detailed summary of the conversation history and persists it
+ to the memory directory as structured files. Unlike compact_memory, this
+ method produces more detailed summaries suitable for long-term storage.
Args:
- messages (list[Msg]): The list of messages to summarize
+ messages (list[Msg]): List of messages to summarize.
+ as_llm (str | ChatModelBase): Language model identifier or instance
+ for summarization. Defaults to "default".
+ as_llm_formatter (str | FormatterBase): Formatter for the language model.
+ Defaults to "default".
+ as_token_counter (str | HuggingFaceTokenCounter): Token counter for
+ measuring message length. Defaults to "default".
+ toolkit (Toolkit | None): Toolkit with file operations for persisting
+ summaries. If None, creates a default toolkit with read/write/edit.
+ language (str): Language for the summary output. "zh" for Chinese,
+ any other value for English. Defaults to "zh".
+ max_input_length (float): Maximum input length in tokens.
+ Defaults to 128K tokens.
+ compact_ratio (float): Ratio used to calculate compaction threshold.
+ Defaults to 0.7.
+ timezone (str | None): Timezone string for date formatting
+ (e.g., "America/Chicago"). Defaults to system local time if None.
Returns:
- str: A generated summary of the messages, or empty string on failure
+ str: The generated summary text, or an empty string if an error occurred.
Note:
- - Summarization may use tools from the toolkit to enhance the summary
- - The summary is typically stored in the memory directory
- - If summarization fails, an empty string is returned
+ This method may write summary files to the memory_path directory
+ using the provided or default toolkit.
"""
try:
- # Initialize summarizer with working directories and configuration
- compactor = Summarizer(
+ if toolkit is None:
+ toolkit = Toolkit()
+ file_io = FileIO(working_dir=str(self.working_path))
+ toolkit.register_tool_function(file_io.read_file)
+ toolkit.register_tool_function(file_io.write_file)
+ toolkit.register_tool_function(file_io.edit_file)
+
+ summarizer = Summarizer(
working_dir=str(self.working_path),
memory_dir=str(self.memory_path),
- memory_compact_threshold=self.memory_compact_threshold,
- chat_model=self.chat_model,
- formatter=self.formatter,
- token_counter=self.token_counter,
- toolkit=self.toolkit,
- language=self.language,
+ memory_compact_threshold=self.calculate_memory_compact_threshold(max_input_length, compact_ratio),
+ toolkit=toolkit,
+ as_llm=as_llm,
+ as_llm_formatter=as_llm_formatter,
+ as_token_counter=as_token_counter,
+ language=language if language == "zh" else "",
+ timezone=timezone,
+ add_thinking_block=add_thinking_block,
)
- # Execute summarization on the provided messages
- return await compactor.call(messages=messages, service_context=self.service_context)
+ return await summarizer.call(messages=messages, service_context=self.service_context)
except Exception as e:
- # Log error and return empty string to indicate failure
logger.exception(f"Error summarizing memory: {e}")
return ""
+ def add_async_summary_task(self, messages: list[Msg], **kwargs):
+ """
+ Add an asynchronous summary task for the given messages.
+
+ Creates a background task to generate a summary of the provided messages
+ without blocking the main execution flow. Completed tasks are automatically
+ cleaned up from the task list.
+
+ Args:
+ messages (list[Msg]): List of messages to be summarized asynchronously.
+ **kwargs: Additional keyword arguments passed to summary_memory().
+ Supported arguments include:
+ - as_llm: Language model identifier or instance
+ - as_llm_formatter: Formatter for the language model
+ - as_token_counter: Token counter instance
+ - toolkit: Toolkit for file operations
+ - language: Output language ("zh" or other)
+ - max_input_length: Maximum input token length
+ - compact_ratio: Compaction threshold ratio
+
+ Note:
+ - 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
+ """
+ remaining_tasks = []
+ for task in self.summary_tasks:
+ if task.done():
+ if task.cancelled():
+ logger.warning("Summary task was cancelled.")
+ continue
+ exc = task.exception()
+ if exc is not None:
+ logger.error(f"Summary task failed: {exc}")
+ else:
+ result = task.result()
+ logger.info(f"Summary task completed: {result}")
+ else:
+ remaining_tasks.append(task)
+ self.summary_tasks = remaining_tasks
+
+ task = asyncio.create_task(self.summary_memory(messages=messages, **kwargs))
+ self.summary_tasks.append(task)
+
+ @property
+ def default_as_token_counter(self) -> HuggingFaceTokenCounter:
+ """
+ Get the default token counter for the memory.
+
+ Returns:
+ HuggingFaceTokenCounter: The default token counter instance.
+ """
+ return self.service_context.as_token_counters["default"]
+
+ async def pre_reasoning_hook(
+ self,
+ messages: list[Msg],
+ system_prompt: str = "",
+ compressed_summary: str = "",
+ as_llm: str | ChatModelBase = "default",
+ as_llm_formatter: str | FormatterBase = "default",
+ as_token_counter: str | HuggingFaceTokenCounter = "default",
+ toolkit: Toolkit | None = None,
+ language: str = "zh",
+ max_input_length: float = 128 * 1024,
+ compact_ratio: float = 0.7,
+ memory_compact_reserve: int = 10000,
+ enable_tool_result_compact: bool = True,
+ tool_result_compact_keep_n: int = 3,
+ ) -> tuple[list[Msg], str]:
+ """
+ Hook called before reasoning to manage memory and context.
+
+ This method is designed to be called before each reasoning step to ensure
+ the conversation context fits within model limits. It performs tool result
+ compaction, checks context size, and triggers memory compaction if needed.
+
+ Args:
+ messages (list[Msg]): Current conversation messages to be processed.
+ system_prompt (str): System prompt that will be included in the context.
+ Used to calculate available space. Defaults to empty string.
+ compressed_summary (str): Existing compressed summary from previous
+ compactions. Defaults to empty string.
+ as_llm (str | ChatModelBase): Language model for compaction operations.
+ Defaults to "default".
+ as_llm_formatter (str | FormatterBase): Formatter for the language model.
+ Defaults to "default".
+ as_token_counter (str | HuggingFaceTokenCounter): Token counter for
+ measuring content length. Defaults to "default".
+ toolkit (Toolkit | None): Toolkit for file operations in summarization.
+ Defaults to None.
+ language (str): Language for generated summaries. Defaults to "zh".
+ max_input_length (float): Maximum context window size in tokens.
+ Defaults to 128K tokens.
+ compact_ratio (float): Ratio for calculating compaction threshold.
+ Defaults to 0.7.
+ memory_compact_reserve (int): Token count to reserve for new responses.
+ Defaults to 10000 tokens.
+ enable_tool_result_compact (bool): Whether to compact tool results.
+ Defaults to True.
+ tool_result_compact_keep_n (int): Number of recent messages to exclude
+ from tool result compaction. Defaults to 3.
+
+ Returns:
+ tuple[list[Msg], str]: A tuple containing:
+ - list[Msg]: Messages to keep in context (maybe reduced)
+ - str: Updated compressed summary incorporating compacted messages
+
+ Note:
+ - Automatically triggers background summarization for compacted messages
+ - Tool results in recent messages (keep_n) are not compacted
+ - Returns original messages unchanged if no compaction is needed
+ """
+ msg_handler = AsMsgHandler(self.default_as_token_counter)
+
+ system_token_count = await msg_handler.count_str_token(system_prompt)
+ compressed_token_count = await msg_handler.count_str_token(compressed_summary)
+ memory_compact_threshold = self.calculate_memory_compact_threshold(max_input_length, compact_ratio)
+ left_compact_threshold = memory_compact_threshold - (system_token_count + compressed_token_count)
+ logger.info(f"Left compact threshold: {left_compact_threshold}")
+
+ if enable_tool_result_compact and tool_result_compact_keep_n > 0:
+ compact_msgs = messages[:-tool_result_compact_keep_n]
+ await self.compact_tool_result(compact_msgs)
+
+ 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,
+ as_token_counter=as_token_counter,
+ )
+
+ if not messages_to_compact:
+ return messages, compressed_summary
+
+ if not is_valid:
+ logger.warning("Invalid messages to compact, skipping.")
+ return messages, compressed_summary
+
+ self.add_async_summary_task(
+ messages=messages_to_compact,
+ as_llm=as_llm,
+ as_llm_formatter=as_llm_formatter,
+ as_token_counter=as_token_counter,
+ toolkit=toolkit,
+ language=language,
+ max_input_length=max_input_length,
+ compact_ratio=compact_ratio,
+ )
+
+ compressed_summary = await self.compact_memory(
+ messages=messages_to_compact,
+ as_llm=as_llm,
+ as_llm_formatter=as_llm_formatter,
+ as_token_counter=as_token_counter,
+ language=language,
+ max_input_length=max_input_length,
+ compact_ratio=compact_ratio,
+ previous_summary=compressed_summary,
+ )
+
+ return messages_to_keep, compressed_summary
+
async def await_summary_tasks(self) -> str:
"""
Wait for all background summary tasks to complete and collect results.
- This method iterates through all pending summary tasks, waits for their
- completion, and collects their results or error information. It's used
- to synchronize with background summarization operations before shutdown
- or when results are needed.
+ Blocks until all pending summary tasks in the task list have completed,
+ canceled, or failed. Collects status information from each task and
+ clears the task list after processing.
Returns:
- str: A concatenated string containing the status and results of
- all summary tasks, with each task on a new line
+ str: A concatenated string of status messages for all tasks, including:
+ - Completion confirmations with results
+ - Cancellation notices
+ - Error messages for failed tasks
Note:
- - Completed tasks are processed immediately without waiting
- - Incomplete tasks are awaited with a timeout
- - Cancelled tasks and exceptions are logged and included in results
- - The task list is cleared after processing all tasks
+ - This method will block if any tasks are still running
+ - All tasks are removed from summary_tasks after this call
+ - Task exceptions are logged but do not raise to the caller
+ - Use this before application shutdown to ensure all summaries complete
"""
result = ""
for task in self.summary_tasks:
@@ -509,7 +700,7 @@ class ReMeLight(Application):
# Check if the task raised an exception
exc = task.exception()
if exc is not None:
- logger.exception(f"Summary task failed: {exc}")
+ logger.error(f"Summary task failed: {exc}")
result += f"Summary task failed: {exc}\n"
else:
# Task completed successfully, collect result
@@ -527,6 +718,7 @@ class ReMeLight(Application):
except asyncio.CancelledError:
logger.warning("Summary task was cancelled while waiting.")
result += "Summary task was cancelled.\n"
+
except Exception as e:
logger.exception(f"Summary task failed: {e}")
result += f"Summary task failed: {e}\n"
@@ -535,70 +727,21 @@ class ReMeLight(Application):
self.summary_tasks.clear()
return result
- def add_async_summary_task(self, messages: list[Msg]):
- """
- Add an asynchronous summary task for the given messages.
-
- This method creates a background task to summarize the provided messages
- without blocking the main execution flow. Before adding a new task, it
- cleans up any completed tasks from the task list to prevent memory leaks.
-
- Args:
- messages (list[Msg]): The list of messages to be summarized in the
- background task
-
- Note:
- - Completed tasks are removed from the tracking list before adding
- - Task status (success, failure, cancellation) is logged for monitoring
- - The new task is created using asyncio.create_task for true async execution
- - Failed or cancelled tasks are logged but do not prevent new tasks
- """
- # Clean up completed summary tasks before adding a new one
- remaining_tasks = []
- for task in self.summary_tasks:
- if task.done():
- # Process completed task status
- if task.cancelled():
- logger.warning("Summary task was cancelled.")
- continue
- exc = task.exception()
- if exc is not None:
- logger.exception(f"Summary task failed: {exc}")
- else:
- # Log successful completion with result summary
- result = task.result()
- logger.info(f"Summary task completed: {result}")
- else:
- # Keep incomplete tasks in the tracking list
- remaining_tasks.append(task)
- self.summary_tasks = remaining_tasks
-
- # Create and track the new background summarization task
- task = asyncio.create_task(self.summary_memory(messages=messages))
- self.summary_tasks.append(task)
-
async def memory_search(self, query: str, max_results: int = 5, min_score: float = 0.1) -> ToolResponse:
"""
- Perform semantic memory search using vector and full-text search.
-
- This method searches the memory store for content relevant to the given query
- using a hybrid approach combining vector similarity search and full-text search.
- Results are ranked by relevance and filtered by the minimum score threshold.
+ Mandatory recall step: semantically search MEMORY.md + memory/*.md
+ (and optional session transcripts) before answering questions about
+ prior work, decisions, dates, people, preferences, or todos; returns
+ top snippets with path + lines.
Args:
- query (str): The search query string. Must not be empty.
- max_results (int): Maximum number of results to return (1-100, default: 5)
- min_score (float): Minimum relevance score threshold (0.001-0.999, default: 0.1)
+ query (str): The semantic search query to find relevant memory snippets.
+ max_results (int): Maximum number of search results to return (optional), default 5.
+ min_score (float): Minimum similarity score threshold for results (optional), default 0.1.
Returns:
ToolResponse: A ToolResponse containing the search results as text,
- or an error message if the query is empty
-
- Note:
- - Vector search weight is controlled by self.vector_weight
- - Candidate retrieval uses self.candidate_multiplier for broader search
- - Parameters are validated and clamped to valid ranges
- - Requires vector search to be enabled via embedding configuration
+ or an error message if the query is empty.
"""
# Validate query parameter
if not query:
@@ -614,12 +757,25 @@ class ReMeLight(Application):
# Validate and clamp max_results to valid range [1, 100]
if isinstance(max_results, int):
max_results = min(max(max_results, 1), 100)
+
+ elif isinstance(max_results, str):
+ try:
+ max_results = min(max(int(max_results), 1), 100)
+ except ValueError:
+ max_results = 5
else:
max_results = 5
# Validate and clamp min_score to valid range [0.001, 0.999]
if isinstance(min_score, (int, float)):
- min_score = min(max(min_score, 0.001), 0.999)
+ min_score = float(min(max(min_score, 0.001), 0.999))
+
+ elif isinstance(min_score, str):
+ try:
+ min_score = float(min(max(float(min_score), 0.001), 0.999))
+ except ValueError:
+ min_score = 0.1
+
else:
min_score = 0.1
@@ -647,26 +803,31 @@ class ReMeLight(Application):
],
)
- def get_in_memory_memory(self):
+ def get_in_memory_memory(self, as_token_counter: HuggingFaceTokenCounter | None = None):
"""
Create and return an in-memory memory instance.
- This method instantiates a ReMeInMemoryMemory object configured with
- the current application's token counter, formatter, and input length limits.
- The in-memory memory provides fast, temporary storage for conversation
- context without persistence.
+ Factory method to create a ReMeInMemoryMemory instance configured with
+ the specified token counter. This memory instance stores messages in RAM
+ during the session, and automatically persists them to dialog_path when
+ messages are compressed or cleared.
+
+ Args:
+ as_token_counter (HuggingFaceTokenCounter): Token counter for
+ measuring content length in the memory.
Returns:
- ReMeInMemoryMemory: A configured in-memory memory instance ready
- for storing and retrieving conversation messages
+ ReMeInMemoryMemory: A new in-memory memory instance ready for use.
+ The instance is configured with self.dialog_path for persistence.
Note:
- - In-memory memory is volatile and cleared when the instance is destroyed
- - Useful for managing conversation context within a single session
- - Shares the same token counter and formatter as the main application
+ - Messages are stored in RAM during active session
+ - When messages are compressed via mark_messages_compressed(), they
+ are persisted to {dialog_path}/{date}.jsonl files
+ - When clear_content() is called, all messages are persisted before
+ clearing from memory
"""
return ReMeInMemoryMemory(
- token_counter=self.token_counter,
- formatter=self.formatter,
- max_input_length=self.max_input_length,
+ token_counter=as_token_counter or self.default_as_token_counter,
+ dialog_path=str(self.dialog_path),
)
diff --git a/cookbook/frozenlake/__init__.py b/test/cookbook/__init__.py
similarity index 100%
rename from cookbook/frozenlake/__init__.py
rename to test/cookbook/__init__.py
diff --git a/cookbook/simple_demo/__init__.py b/test/cookbook/appworld/__init__.py
similarity index 100%
rename from cookbook/simple_demo/__init__.py
rename to test/cookbook/appworld/__init__.py
diff --git a/cookbook/appworld/appworld_react_agent.py b/test/cookbook/appworld/appworld_react_agent.py
similarity index 87%
rename from cookbook/appworld/appworld_react_agent.py
rename to test/cookbook/appworld/appworld_react_agent.py
index ba194462..ca8ee2f3 100644
--- a/cookbook/appworld/appworld_react_agent.py
+++ b/test/cookbook/appworld/appworld_react_agent.py
@@ -96,10 +96,7 @@ class AppworldReactAgent:
def prompt_messages(self, run_id, task_index, previous_memories: None, world: AppWorld):
app_descriptions = json.dumps(
- [
- {"name": k, "description": v}
- for (k, v) in world.task.app_descriptions.items()
- ],
+ [{"name": k, "description": v} for (k, v) in world.task.app_descriptions.items()],
indent=1,
)
dictionary = {"supervisor": world.task.supervisor, "app_descriptions": app_descriptions}
@@ -112,7 +109,12 @@ class AppworldReactAgent:
self.retrieved_memory_list[run_id][task_index] = response["metadata"]["memory_list"]
task_memory = response["answer"]
logger.info(f"loaded task_memory: {task_memory}")
- query = "Task:\n" + query + "\n\nSome Related Experience to help you to complete the task:\n" + re.sub(r'(?i)\bMemory\s*(\d+)\s*[:]', r'Experience \1:', task_memory)
+ query = (
+ "Task:\n"
+ + query
+ + "\n\nSome Related Experience to help you to complete the task:\n"
+ + re.sub(r"(?i)\bMemory\s*(\d+)\s*[:]", r"Experience \1:", task_memory)
+ )
else:
formatted_memories = []
for i, memory in enumerate(previous_memories, 1):
@@ -120,14 +122,18 @@ class AppworldReactAgent:
memory_content = memory["content"]
memory_text = f"Experience {i}:\n When to use: {condition}\n Content: {memory_content}\n"
formatted_memories.append(memory_text)
- query = "Task:\n" + query + "\n\nSome Related Experience to help you to complete the task:\n" + "\n".join(formatted_memories)
+ query = (
+ "Task:\n"
+ + query
+ + "\n\nSome Related Experience to help you to complete the task:\n"
+ + "\n".join(formatted_memories)
+ )
messages = [
{"role": "system", "content": sys_prompt},
- {"role": "user", "content": query}
+ {"role": "user", "content": query},
]
self.history[run_id][task_index] = messages
-
@staticmethod
def get_reward(world) -> float:
tracker = world.evaluate()
@@ -136,7 +142,9 @@ class AppworldReactAgent:
return num_passes / (num_passes + num_failures)
def extract_code_and_fix_content(
- self, text: str, ignore_multiple_calls=True
+ self,
+ text: str,
+ ignore_multiple_calls=True,
) -> tuple[str, str]:
full_code_regex = r"```python\n(.*?)```"
partial_code_regex = r".*```python\n(.*)"
@@ -154,7 +162,9 @@ class AppworldReactAgent:
match_end = re_match.end()
# check for partial code match at end (no terminating ```) following the last match
partial_match = re.match(
- partial_code_regex, original_text[match_end:], flags=re.DOTALL
+ partial_code_regex,
+ original_text[match_end:],
+ flags=re.DOTALL,
)
if partial_match:
output_code += partial_match.group(1).strip()
@@ -180,7 +190,12 @@ class AppworldReactAgent:
before_score = self.get_reward(world)
for i in range(self.max_interactions):
if i == 0:
- self.prompt_messages(run_id=run_id, task_index=task_index, previous_memories=previous_memories, world=world)
+ self.prompt_messages(
+ run_id=run_id,
+ task_index=task_index,
+ previous_memories=previous_memories,
+ world=world,
+ )
code_msg = self.call_llm(self.history[run_id][task_index])
code, text = self.extract_code_and_fix_content(code_msg)
self.history[run_id][task_index].append({"role": "assistant", "content": code})
@@ -189,7 +204,9 @@ class AppworldReactAgent:
# if len(output) > self.max_response_size:
# # logger.warning(f"output exceed max size={len(output)}")
# output = output[: self.max_response_size]
- self.history[run_id][task_index].append({"role": "user", "content": "Output:\n```\n" + output + "```\n\n"})
+ self.history[run_id][task_index].append(
+ {"role": "user", "content": "Output:\n```\n" + output + "```\n\n"},
+ )
if world.task_completed():
break
@@ -199,7 +216,9 @@ class AppworldReactAgent:
if self.use_memory:
if self.use_memory_addition:
- new_traj_list = [self.get_traj_from_task_history(task_id, self.history[run_id][task_index], after_score)]
+ new_traj_list = [
+ self.get_traj_from_task_history(task_id, self.history[run_id][task_index], after_score),
+ ]
previous_memories = self.add_memory(new_traj_list)
if after_score != 1:
self.delete_memory_by_ids([mem["memory_id"] for mem in previous_memories])
@@ -209,7 +228,7 @@ class AppworldReactAgent:
self.update_memory_information(self.retrieved_memory_list[run_id][task_index], update_utility)
counter += 1
- if self.use_memory_deletion: # and counter % self.delete_freq == 0:
+ if self.use_memory_deletion: # and counter % self.delete_freq == 0:
self.delete_memory()
t_result = {
@@ -261,7 +280,7 @@ class AppworldReactAgent:
return {
"task_id": task_id,
"messages": task_history,
- "score": reward
+ "score": reward,
}
def add_memory(self, trajectories):
@@ -290,8 +309,8 @@ class AppworldReactAgent:
json={
"workspace_id": self.memory_workspace_id,
"action": "delete_ids",
- "memory_ids": memory_ids
- }
+ "memory_ids": memory_ids,
+ },
)
response.raise_for_status()
@@ -318,6 +337,7 @@ class AppworldReactAgent:
)
response.raise_for_status()
+
def main():
dataset_name = "train"
task_ids = load_task_ids(dataset_name)
diff --git a/cookbook/appworld/prompt.py b/test/cookbook/appworld/prompt.py
similarity index 100%
rename from cookbook/appworld/prompt.py
rename to test/cookbook/appworld/prompt.py
diff --git a/cookbook/appworld/requirements.txt b/test/cookbook/appworld/requirements.txt
similarity index 100%
rename from cookbook/appworld/requirements.txt
rename to test/cookbook/appworld/requirements.txt
diff --git a/cookbook/appworld/run_appworld.py b/test/cookbook/appworld/run_appworld.py
similarity index 98%
rename from cookbook/appworld/run_appworld.py
rename to test/cookbook/appworld/run_appworld.py
index 3379a16a..d286c64a 100644
--- a/cookbook/appworld/run_appworld.py
+++ b/test/cookbook/appworld/run_appworld.py
@@ -90,7 +90,7 @@ def run_agent(
utility_threshold: float = 0.5,
workspace_id: str = "appworld_v1",
api_url: str = "http://0.0.0.0:8002/",
- batch_size: int = 4
+ batch_size: int = 4,
):
experiment_name = dataset_name + "_" + experiment_suffix
path: Path = Path(f"./exp_result/{model_name}")
@@ -125,7 +125,7 @@ def run_agent(
future_list: list = []
for i, task_id in enumerate(batch_task_ids):
actor = AppworldReactAgent.remote(
- index=start_idx+i,
+ index=start_idx + i,
model_name=model_name,
task_ids=[task_id],
experiment_name=experiment_name,
@@ -193,9 +193,10 @@ def run_agent(
result.append(task_results)
dump_file()
+
def main():
max_workers = 8
- num_runs = 1 # Number of runs
+ num_runs = 1 # Number of runs
batch_size = 8 # Number of concurrent tasks per batch
num_trials = 2
@@ -206,7 +207,6 @@ def main():
workspace_id = "appworld"
api_url = "http://0.0.0.0:8002/"
-
# Clean up workspace before starting
logger.info("Deleting workspace...")
delete_workspace(workspace_id=workspace_id, api_url=api_url)
@@ -216,7 +216,6 @@ def main():
logger.info("Start load experiments to build task memories")
load_memory(workspace_id=workspace_id, api_url=api_url)
-
for i in range(num_runs):
run_agent(
model_name=model_name,
@@ -232,8 +231,9 @@ def main():
utility_threshold=0.5,
workspace_id=workspace_id,
api_url=api_url,
- batch_size=batch_size
+ batch_size=batch_size,
)
+
if __name__ == "__main__":
main()
diff --git a/cookbook/appworld/run_exp_statistic.py b/test/cookbook/appworld/run_exp_statistic.py
similarity index 100%
rename from cookbook/appworld/run_exp_statistic.py
rename to test/cookbook/appworld/run_exp_statistic.py
diff --git a/cookbook/tool_memory/__init__.py b/test/cookbook/bfcl/__init__.py
similarity index 100%
rename from cookbook/tool_memory/__init__.py
rename to test/cookbook/bfcl/__init__.py
diff --git a/cookbook/bfcl/bfcl_agent.py b/test/cookbook/bfcl/bfcl_agent.py
similarity index 98%
rename from cookbook/bfcl/bfcl_agent.py
rename to test/cookbook/bfcl/bfcl_agent.py
index 64eb6e6b..2c779a6a 100644
--- a/cookbook/bfcl/bfcl_agent.py
+++ b/test/cookbook/bfcl/bfcl_agent.py
@@ -195,7 +195,7 @@ class BFCLAgent:
# Extract memory list from response
memory_list = result.get("metadata", {}).get("memory_list", [])
- logger.info(f'add new memories: {memory_list}')
+ logger.info(f"add new memories: {memory_list}")
return memory_list
def delete_memory_by_ids(self, memory_ids):
@@ -204,8 +204,8 @@ class BFCLAgent:
json={
"workspace_id": self.memory_workspace_id,
"action": "delete_ids",
- "memory_ids": memory_ids
- }
+ "memory_ids": memory_ids,
+ },
)
response.raise_for_status()
@@ -647,7 +647,9 @@ class BFCLAgent:
reward = self.get_reward(run_id, task_index)
if self.use_memory:
if self.use_memory_addition: # selectively add memories when succeed
- new_traj_list = [self.get_traj_from_task_history(task_id, self.history[run_id][task_index], reward)]
+ new_traj_list = [
+ self.get_traj_from_task_history(task_id, self.history[run_id][task_index], reward),
+ ]
previous_memories = self.add_memory(new_traj_list)
if reward != 1:
self.delete_memory_by_ids([mem["memory_id"] for mem in previous_memories])
diff --git a/cookbook/bfcl/bfcl_utils.py b/test/cookbook/bfcl/bfcl_utils.py
similarity index 100%
rename from cookbook/bfcl/bfcl_utils.py
rename to test/cookbook/bfcl/bfcl_utils.py
diff --git a/cookbook/bfcl/init_exp_pool.py b/test/cookbook/bfcl/init_exp_pool.py
similarity index 100%
rename from cookbook/bfcl/init_exp_pool.py
rename to test/cookbook/bfcl/init_exp_pool.py
diff --git a/cookbook/bfcl/init_task_memory_pool.py b/test/cookbook/bfcl/init_task_memory_pool.py
similarity index 100%
rename from cookbook/bfcl/init_task_memory_pool.py
rename to test/cookbook/bfcl/init_task_memory_pool.py
diff --git a/cookbook/bfcl/local_file_to_library.py b/test/cookbook/bfcl/local_file_to_library.py
similarity index 100%
rename from cookbook/bfcl/local_file_to_library.py
rename to test/cookbook/bfcl/local_file_to_library.py
diff --git a/cookbook/bfcl/requirements.txt b/test/cookbook/bfcl/requirements.txt
similarity index 100%
rename from cookbook/bfcl/requirements.txt
rename to test/cookbook/bfcl/requirements.txt
diff --git a/cookbook/bfcl/run_bfcl.py b/test/cookbook/bfcl/run_bfcl.py
similarity index 99%
rename from cookbook/bfcl/run_bfcl.py
rename to test/cookbook/bfcl/run_bfcl.py
index fd8cc2db..177d91e1 100644
--- a/cookbook/bfcl/run_bfcl.py
+++ b/test/cookbook/bfcl/run_bfcl.py
@@ -91,7 +91,7 @@ def main():
num_runs = 1
num_trials = 2
- model_name="qwen3-8b"
+ model_name = "qwen3-8b"
use_memory = False
use_memory_addition = False
use_memory_deletion = False
diff --git a/cookbook/bfcl/run_exp_statistic.py b/test/cookbook/bfcl/run_exp_statistic.py
similarity index 100%
rename from cookbook/bfcl/run_exp_statistic.py
rename to test/cookbook/bfcl/run_exp_statistic.py
diff --git a/cookbook/bfcl/split_into_trainval.py b/test/cookbook/bfcl/split_into_trainval.py
similarity index 100%
rename from cookbook/bfcl/split_into_trainval.py
rename to test/cookbook/bfcl/split_into_trainval.py
diff --git a/reme/memory/tools/chunk/__init__.py b/test/cookbook/frozenlake/__init__.py
similarity index 100%
rename from reme/memory/tools/chunk/__init__.py
rename to test/cookbook/frozenlake/__init__.py
diff --git a/cookbook/frozenlake/frozenlake_prompts.yaml b/test/cookbook/frozenlake/frozenlake_prompts.yaml
similarity index 100%
rename from cookbook/frozenlake/frozenlake_prompts.yaml
rename to test/cookbook/frozenlake/frozenlake_prompts.yaml
diff --git a/cookbook/frozenlake/frozenlake_react_agent.py b/test/cookbook/frozenlake/frozenlake_react_agent.py
similarity index 100%
rename from cookbook/frozenlake/frozenlake_react_agent.py
rename to test/cookbook/frozenlake/frozenlake_react_agent.py
diff --git a/cookbook/frozenlake/map_manager.py b/test/cookbook/frozenlake/map_manager.py
similarity index 100%
rename from cookbook/frozenlake/map_manager.py
rename to test/cookbook/frozenlake/map_manager.py
diff --git a/cookbook/frozenlake/run_exp_statistic.py b/test/cookbook/frozenlake/run_exp_statistic.py
similarity index 100%
rename from cookbook/frozenlake/run_exp_statistic.py
rename to test/cookbook/frozenlake/run_exp_statistic.py
diff --git a/cookbook/frozenlake/run_frozenlake.py b/test/cookbook/frozenlake/run_frozenlake.py
similarity index 100%
rename from cookbook/frozenlake/run_frozenlake.py
rename to test/cookbook/frozenlake/run_frozenlake.py
diff --git a/reme/memory/tools/history/__init__.py b/test/cookbook/simple_demo/__init__.py
similarity index 100%
rename from reme/memory/tools/history/__init__.py
rename to test/cookbook/simple_demo/__init__.py
diff --git a/cookbook/simple_demo/import_usage_demo.py b/test/cookbook/simple_demo/import_usage_demo.py
similarity index 100%
rename from cookbook/simple_demo/import_usage_demo.py
rename to test/cookbook/simple_demo/import_usage_demo.py
diff --git a/cookbook/simple_demo/mcp_task_memory.jsonl b/test/cookbook/simple_demo/mcp_task_memory.jsonl
similarity index 100%
rename from cookbook/simple_demo/mcp_task_memory.jsonl
rename to test/cookbook/simple_demo/mcp_task_memory.jsonl
diff --git a/cookbook/simple_demo/personal_memory.jsonl b/test/cookbook/simple_demo/personal_memory.jsonl
similarity index 100%
rename from cookbook/simple_demo/personal_memory.jsonl
rename to test/cookbook/simple_demo/personal_memory.jsonl
diff --git a/cookbook/simple_demo/task_memory.jsonl b/test/cookbook/simple_demo/task_memory.jsonl
similarity index 100%
rename from cookbook/simple_demo/task_memory.jsonl
rename to test/cookbook/simple_demo/task_memory.jsonl
diff --git a/cookbook/simple_demo/task_messages.jsonl b/test/cookbook/simple_demo/task_messages.jsonl
similarity index 100%
rename from cookbook/simple_demo/task_messages.jsonl
rename to test/cookbook/simple_demo/task_messages.jsonl
diff --git a/cookbook/simple_demo/use_personal_memory_demo.py b/test/cookbook/simple_demo/use_personal_memory_demo.py
similarity index 100%
rename from cookbook/simple_demo/use_personal_memory_demo.py
rename to test/cookbook/simple_demo/use_personal_memory_demo.py
diff --git a/cookbook/simple_demo/use_task_memory_demo.py b/test/cookbook/simple_demo/use_task_memory_demo.py
similarity index 100%
rename from cookbook/simple_demo/use_task_memory_demo.py
rename to test/cookbook/simple_demo/use_task_memory_demo.py
diff --git a/cookbook/simple_demo/use_task_memory_mcp_demo.py b/test/cookbook/simple_demo/use_task_memory_mcp_demo.py
similarity index 100%
rename from cookbook/simple_demo/use_task_memory_mcp_demo.py
rename to test/cookbook/simple_demo/use_task_memory_mcp_demo.py
diff --git a/cookbook/simple_demo/use_tool_memory_demo.py b/test/cookbook/simple_demo/use_tool_memory_demo.py
similarity index 100%
rename from cookbook/simple_demo/use_tool_memory_demo.py
rename to test/cookbook/simple_demo/use_tool_memory_demo.py
diff --git a/reme/memory/tools/profiles/__init__.py b/test/cookbook/tool_memory/__init__.py
similarity index 100%
rename from reme/memory/tools/profiles/__init__.py
rename to test/cookbook/tool_memory/__init__.py
diff --git a/cookbook/tool_memory/query.json b/test/cookbook/tool_memory/query.json
similarity index 100%
rename from cookbook/tool_memory/query.json
rename to test/cookbook/tool_memory/query.json
diff --git a/cookbook/tool_memory/run_reme_tool_bench.py b/test/cookbook/tool_memory/run_reme_tool_bench.py
similarity index 100%
rename from cookbook/tool_memory/run_reme_tool_bench.py
rename to test/cookbook/tool_memory/run_reme_tool_bench.py
diff --git a/cookbook/working_memory/react_agent_with_working_memory.py b/test/cookbook/working_memory/react_agent_with_working_memory.py
similarity index 70%
rename from cookbook/working_memory/react_agent_with_working_memory.py
rename to test/cookbook/working_memory/react_agent_with_working_memory.py
index 52e3fdc6..27f2275e 100644
--- a/cookbook/working_memory/react_agent_with_working_memory.py
+++ b/test/cookbook/working_memory/react_agent_with_working_memory.py
@@ -28,9 +28,11 @@ class ReactAgent:
rather than on complex agent logic.
"""
- def __init__(self,
- model_name="",
- max_steps: int = 50):
+ def __init__(
+ self,
+ model_name="",
+ max_steps: int = 50,
+ ):
# You can replace this with your own LLM wrapper if needed.
self.llm = OpenAICompatibleLLM(model_name=model_name)
@@ -65,10 +67,16 @@ class ReactAgent:
# Prepare all available tools from the MCP server.
tool_dict: Dict[str, ToolCall] = {}
- async with FastMcpClient("reme_mcp_server", {
- "type": "sse",
- "url": "http://0.0.0.0:8002/sse",
- }) as mcp_client, HttpClient(base_url="http://localhost:8003") as http_client:
+ async with (
+ FastMcpClient(
+ "reme_mcp_server",
+ {
+ "type": "sse",
+ "url": "http://0.0.0.0:8002/sse",
+ },
+ ) as mcp_client,
+ HttpClient(base_url="http://localhost:8003") as http_client,
+ ):
tool_calls = await mcp_client.list_tool_calls()
for tool_call in tool_calls:
@@ -87,25 +95,30 @@ class ReactAgent:
# - compress long histories,
# - offload detailed context into working memory storage,
# - keep the recent message(s) for short-term reasoning.
- result = await http_client.execute_flow("summary_working_memory",
- messages=[x.simple_dump() for x in messages],
- working_summary_mode="auto",
- compact_ratio_threshold=0.75,
- max_total_tokens=20000,
- max_tool_message_tokens=2000,
- group_token_threshold=None,
- keep_recent_count=1,
- store_dir="./test_working_memory")
+ result = await http_client.execute_flow(
+ "summary_working_memory",
+ messages=[x.simple_dump() for x in messages],
+ working_summary_mode="auto",
+ compact_ratio_threshold=0.75,
+ max_total_tokens=20000,
+ max_tool_message_tokens=2000,
+ group_token_threshold=None,
+ keep_recent_count=1,
+ store_dir="./test_working_memory",
+ )
# Convert the API result back into `Message` objects for the LLM.
messages = [Message(**x) for x in result.answer]
# Ask the LLM what to do next.
# You can plug in your own tool-calling strategy here.
- assistant_message: Message = await self.llm.achat(messages=messages, tools=[
- tool_dict["grep_working_memory"],
- tool_dict["read_working_memory"],
- ])
+ assistant_message: Message = await self.llm.achat(
+ messages=messages,
+ tools=[
+ tool_dict["grep_working_memory"],
+ tool_dict["read_working_memory"],
+ ],
+ )
messages.append(assistant_message)
@@ -118,20 +131,25 @@ class ReactAgent:
logger.exception(f"unknown tool_call.name={tool_call.name}")
continue
- logger.info(f"round{i + 1}.{j} submit tool_calls={tool_call.name} "
- f"argument={tool_call.argument_dict}")
+ logger.info(
+ f"round{i + 1}.{j} submit tool_calls={tool_call.name} " f"argument={tool_call.argument_dict}",
+ )
# Execute the tool via MCP and parse the result.
- result = await mcp_client.call_tool(tool_call.name,
- arguments=tool_call.argument_dict,
- parse_result=True)
+ result = await mcp_client.call_tool(
+ tool_call.name,
+ arguments=tool_call.argument_dict,
+ parse_result=True,
+ )
# Attach the tool result as a TOOL-role message so the LLM
# can see and reason about it in the next step.
- messages.append(Message(
- role=Role.TOOL,
- tool_call_id=tool_call.id,
- content=result,
- ))
+ messages.append(
+ Message(
+ role=Role.TOOL,
+ tool_call_id=tool_call.id,
+ content=result,
+ ),
+ )
return messages
diff --git a/cookbook/working_memory/work_memory_demo.py b/test/cookbook/working_memory/work_memory_demo.py
similarity index 99%
rename from cookbook/working_memory/work_memory_demo.py
rename to test/cookbook/working_memory/work_memory_demo.py
index 612a637a..826cb0f2 100644
--- a/cookbook/working_memory/work_memory_demo.py
+++ b/test/cookbook/working_memory/work_memory_demo.py
@@ -108,7 +108,7 @@ async def main():
logger.info(
f"origin_token_count: {origin_token_count} "
f"after_token_count: {after_token_count} "
- f"compress_ratio={after_token_count / origin_token_count:.2f}"
+ f"compress_ratio={after_token_count / origin_token_count:.2f}",
)
diff --git a/reme/extension/cli/__init__.py b/test/test/cli/__init__.py
similarity index 100%
rename from reme/extension/cli/__init__.py
rename to test/test/cli/__init__.py
diff --git a/reme/extension/cli/fb_cli.py b/test/test/cli/fb_cli.py
similarity index 100%
rename from reme/extension/cli/fb_cli.py
rename to test/test/cli/fb_cli.py
diff --git a/reme/extension/cli/fb_cli.yaml b/test/test/cli/fb_cli.yaml
similarity index 100%
rename from reme/extension/cli/fb_cli.yaml
rename to test/test/cli/fb_cli.yaml
diff --git a/reme/extension/cli/fb_compactor.py b/test/test/cli/fb_compactor.py
similarity index 100%
rename from reme/extension/cli/fb_compactor.py
rename to test/test/cli/fb_compactor.py
diff --git a/reme/extension/cli/fb_compactor.yaml b/test/test/cli/fb_compactor.yaml
similarity index 100%
rename from reme/extension/cli/fb_compactor.yaml
rename to test/test/cli/fb_compactor.yaml
diff --git a/reme/extension/cli/fb_context_checker.py b/test/test/cli/fb_context_checker.py
similarity index 100%
rename from reme/extension/cli/fb_context_checker.py
rename to test/test/cli/fb_context_checker.py
diff --git a/reme/extension/cli/fb_summarizer.py b/test/test/cli/fb_summarizer.py
similarity index 100%
rename from reme/extension/cli/fb_summarizer.py
rename to test/test/cli/fb_summarizer.py
diff --git a/reme/extension/cli/fb_summarizer.yaml b/test/test/cli/fb_summarizer.yaml
similarity index 100%
rename from reme/extension/cli/fb_summarizer.yaml
rename to test/test/cli/fb_summarizer.yaml
diff --git a/reme/extension/reme_cli.py b/test/test/reme_cli.py
similarity index 100%
rename from reme/extension/reme_cli.py
rename to test/test/reme_cli.py
diff --git a/test/test_fs_compactor.py b/test/test/test_fs_compactor.py
similarity index 100%
rename from test/test_fs_compactor.py
rename to test/test/test_fs_compactor.py
diff --git a/test/test_fs_context_checker.py b/test/test/test_fs_context_checker.py
similarity index 100%
rename from test/test_fs_context_checker.py
rename to test/test/test_fs_context_checker.py
diff --git a/test/test_fs_file_watch_integration.py b/test/test/test_fs_file_watch_integration.py
similarity index 99%
rename from test/test_fs_file_watch_integration.py
rename to test/test/test_fs_file_watch_integration.py
index 9f679d38..a3534337 100644
--- a/test/test_fs_file_watch_integration.py
+++ b/test/test/test_fs_file_watch_integration.py
@@ -295,7 +295,6 @@ async def test_file_watch_integration():
"watch_paths": [TestConfig.WORKING_DIR, f"{TestConfig.WORKING_DIR}/memory"],
"suffix_filters": [".md"],
"recursive": False,
- "scan_on_start": True,
},
)
diff --git a/test/test_fs_memory_get.py b/test/test/test_fs_memory_get.py
similarity index 100%
rename from test/test_fs_memory_get.py
rename to test/test/test_fs_memory_get.py
diff --git a/test/test_fs_memory_search.py b/test/test/test_fs_memory_search.py
similarity index 100%
rename from test/test_fs_memory_search.py
rename to test/test/test_fs_memory_search.py
diff --git a/test/test_fs_summary.py b/test/test/test_fs_summary.py
similarity index 100%
rename from test/test_fs_summary.py
rename to test/test/test_fs_summary.py
diff --git a/tests/light/test_compactor.py b/tests/light/test_compactor.py
index 32dfd9d3..74cba241 100644
--- a/tests/light/test_compactor.py
+++ b/tests/light/test_compactor.py
@@ -1,23 +1,18 @@
"""Tests for Compactor."""
import asyncio
-import logging
from agentscope.message import Msg
-
from test_utils import (
get_dash_chat_model,
get_formatter,
get_token_counter,
)
-from reme.memory.file_based import Compactor
-# 配置日志输出到控制台
-logging.basicConfig(
- level=logging.INFO,
- format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
-)
-logger = logging.getLogger(__name__)
+from reme.core.utils import get_logger
+from reme.memory.file_based.components import Compactor
+
+logger = get_logger()
# ANSI 颜色码
@@ -101,9 +96,10 @@ def create_compactor():
"""Create a Compactor instance for testing."""
return Compactor(
memory_compact_threshold=4000,
- chat_model=get_dash_chat_model(),
- formatter=get_formatter(),
- token_counter=get_token_counter(),
+ as_token_counter=get_token_counter(),
+ as_llm=get_dash_chat_model(),
+ as_llm_formatter=get_formatter(),
+ language="zh",
)
@@ -286,9 +282,9 @@ def test_low_threshold():
"""Test compaction with low memory threshold."""
compactor = Compactor(
memory_compact_threshold=500,
- chat_model=get_dash_chat_model(),
- formatter=get_formatter(),
- token_counter=get_token_counter(),
+ as_token_counter=get_token_counter(),
+ as_llm=get_dash_chat_model(),
+ as_llm_formatter=get_formatter(),
)
messages = [
@@ -309,9 +305,9 @@ def test_high_threshold():
"""Test compaction with high memory threshold."""
compactor = Compactor(
memory_compact_threshold=10000,
- chat_model=get_dash_chat_model(),
- formatter=get_formatter(),
- token_counter=get_token_counter(),
+ as_token_counter=get_token_counter(),
+ as_llm=get_dash_chat_model(),
+ as_llm_formatter=get_formatter(),
)
messages = [
diff --git a/tests/light/test_context_check.py b/tests/light/test_context_check.py
new file mode 100644
index 00000000..e31ea278
--- /dev/null
+++ b/tests/light/test_context_check.py
@@ -0,0 +1,1353 @@
+"""Tests for AsMsgHandler.context_check method."""
+
+import asyncio
+
+from agentscope.message import Msg
+from test_utils import get_token_counter
+
+from reme.core.utils import get_logger
+from reme.memory.file_based.utils import AsMsgHandler
+
+logger = get_logger()
+
+
+# ANSI color codes
+class Colors:
+ """ANSI color codes for terminal output."""
+
+ GREEN = "\033[92m"
+ RED = "\033[91m"
+ YELLOW = "\033[93m"
+ BLUE = "\033[94m"
+ CYAN = "\033[96m"
+ BOLD = "\033[1m"
+ RESET = "\033[0m"
+
+
+def print_pass(test_name: str):
+ """Print test passed message."""
+ print(f"{Colors.GREEN}{Colors.BOLD}✓ {test_name} PASSED{Colors.RESET}")
+
+
+def print_fail(test_name: str, error: str):
+ """Print test failed message."""
+ print(f"{Colors.RED}{Colors.BOLD}✗ {test_name} FAILED: {error}{Colors.RESET}")
+
+
+def print_error(test_name: str, error: str):
+ """Print test error message."""
+ print(f"{Colors.YELLOW}{Colors.BOLD}⚠ {test_name} ERROR: {error}{Colors.RESET}")
+
+
+def print_test_header(test_name: str):
+ """Print test header."""
+ print(f"\n{Colors.CYAN}{'=' * 60}{Colors.RESET}")
+ print(f"{Colors.BLUE}{Colors.BOLD}Running: {test_name}{Colors.RESET}")
+ print(f"{Colors.CYAN}{'=' * 60}{Colors.RESET}")
+
+
+def create_handler() -> AsMsgHandler:
+ """Create an AsMsgHandler instance for testing."""
+ return AsMsgHandler(token_counter=get_token_counter())
+
+
+def verify_context_check_invariants(
+ handler: AsMsgHandler,
+ messages: list[Msg],
+ to_compact: list[Msg],
+ to_keep: list[Msg],
+ memory_compact_threshold: int,
+ memory_compact_reserve: int,
+ test_name: str,
+):
+ """Verify that context_check results satisfy all invariants.
+
+ This function checks:
+ 1. Threshold requirement: If total tokens <= threshold, no compaction should occur
+ 2. Reserve requirement: Kept messages' total tokens should not exceed reserve
+ 3. Order requirement: Both to_compact and to_keep should preserve original order
+
+ Args:
+ handler: The AsMsgHandler instance
+ messages: Original messages list
+ to_compact: Messages to compact returned by context_check
+ to_keep: Messages to keep returned by context_check
+ memory_compact_threshold: The threshold parameter used
+ memory_compact_reserve: The reserve parameter used
+ test_name: Name of the test for error reporting
+
+ Raises:
+ AssertionError: If any invariant is violated
+ """
+ # Calculate total tokens of original messages
+ total_tokens = sum(asyncio.run(handler.stat_message(m)).total_tokens for m in messages)
+
+ # 1. Threshold requirement check
+ if total_tokens <= memory_compact_threshold:
+ assert len(to_compact) == 0, (
+ f"[{test_name}] Threshold violation: total_tokens ({total_tokens}) <= "
+ f"threshold ({memory_compact_threshold}), but to_compact is not empty "
+ f"(has {len(to_compact)} messages)"
+ )
+ assert to_keep == messages, (
+ f"[{test_name}] Threshold violation: total_tokens ({total_tokens}) <= "
+ f"threshold ({memory_compact_threshold}), but to_keep differs from original messages"
+ )
+
+ # 2. Reserve requirement check
+ kept_tokens = sum(asyncio.run(handler.stat_message(m)).total_tokens for m in to_keep)
+ assert kept_tokens <= memory_compact_reserve or len(to_keep) == 0, (
+ f"[{test_name}] Reserve violation: kept_tokens ({kept_tokens}) > " f"reserve ({memory_compact_reserve})"
+ )
+
+ # 3. Order requirement check - both lists should preserve original order
+ # Create a mapping of message id to original index
+ msg_to_idx = {id(m): i for i, m in enumerate(messages)}
+
+ # Check to_compact order
+ compact_indices = [msg_to_idx.get(id(m), -1) for m in to_compact]
+ for i in range(len(compact_indices) - 1):
+ assert compact_indices[i] < compact_indices[i + 1], (
+ f"[{test_name}] Order violation in to_compact: message at original index "
+ f"{compact_indices[i]} appears before message at index {compact_indices[i + 1]}"
+ )
+
+ # Check to_keep order
+ keep_indices = [msg_to_idx.get(id(m), -1) for m in to_keep]
+ for i in range(len(keep_indices) - 1):
+ assert keep_indices[i] < keep_indices[i + 1], (
+ f"[{test_name}] Order violation in to_keep: message at original index "
+ f"{keep_indices[i]} appears before message at index {keep_indices[i + 1]}"
+ )
+
+ # 4. Additional check: to_compact indices should all be less than to_keep indices
+ # (compact messages come from the beginning, keep messages come from the end)
+ if to_compact and to_keep:
+ max_compact_idx = max(compact_indices) if compact_indices else -1
+ min_keep_idx = min(keep_indices) if keep_indices else len(messages)
+ assert max_compact_idx < min_keep_idx, (
+ f"[{test_name}] Partition violation: max compact index ({max_compact_idx}) >= "
+ f"min keep index ({min_keep_idx}). Compact and keep should be a clean partition."
+ )
+
+ # 5. Check that all messages are accounted for (no duplicates, no missing)
+ assert len(to_compact) + len(to_keep) == len(messages), (
+ f"[{test_name}] Count mismatch: to_compact ({len(to_compact)}) + "
+ f"to_keep ({len(to_keep)}) != original ({len(messages)})"
+ )
+
+ all_returned = set(id(m) for m in to_compact) | set(id(m) for m in to_keep)
+ all_original = set(id(m) for m in messages)
+ assert all_returned == all_original, f"[{test_name}] Message set mismatch: returned messages differ from original"
+
+
+def create_user_msg(content: str) -> Msg:
+ """Create a user message."""
+ return Msg(name="user", role="user", content=content)
+
+
+def create_assistant_msg(content: str) -> Msg:
+ """Create an assistant message."""
+ return Msg(name="assistant", role="assistant", content=content)
+
+
+def create_tool_use_msg(tool_id: str, tool_name: str, tool_input: dict) -> Msg:
+ """Create a message with tool_use content block."""
+ return Msg(
+ name="assistant",
+ role="assistant",
+ content=[
+ {
+ "type": "tool_use",
+ "id": tool_id,
+ "name": tool_name,
+ "input": tool_input,
+ },
+ ],
+ )
+
+
+def create_tool_result_msg(tool_id: str, tool_name: str, output: str) -> Msg:
+ """Create a message with tool_result content block."""
+ return Msg(
+ name="tool",
+ role="user",
+ content=[
+ {
+ "type": "tool_result",
+ "id": tool_id,
+ "name": tool_name,
+ "output": output,
+ },
+ ],
+ )
+
+
+def create_mixed_tool_msg(
+ tool_use_id: str,
+ tool_use_name: str,
+ tool_use_input: dict,
+ tool_result_id: str,
+ tool_result_name: str,
+ tool_result_output: str,
+) -> Msg:
+ """Create a message with both tool_use and tool_result blocks."""
+ return Msg(
+ name="assistant",
+ role="assistant",
+ content=[
+ {
+ "type": "tool_use",
+ "id": tool_use_id,
+ "name": tool_use_name,
+ "input": tool_use_input,
+ },
+ {
+ "type": "tool_result",
+ "id": tool_result_id,
+ "name": tool_result_name,
+ "output": tool_result_output,
+ },
+ ],
+ )
+
+
+# =============================================================================
+# Normal Cases
+# =============================================================================
+
+
+def test_empty_messages():
+ """Test context_check with empty messages list."""
+ handler = create_handler()
+ messages = []
+ threshold, reserve = 1000, 500
+ to_compact, to_keep, _ = asyncio.run(
+ handler.context_check(
+ messages=messages,
+ memory_compact_threshold=threshold,
+ memory_compact_reserve=reserve,
+ ),
+ )
+ assert not to_compact, f"Expected empty compact list, got: {to_compact}"
+ assert to_keep == [], f"Expected empty keep list, got: {to_keep}"
+ verify_context_check_invariants(handler, messages, to_compact, to_keep, threshold, reserve, "test_empty_messages")
+ print_pass("test_empty_messages")
+
+
+def test_below_threshold_returns_all():
+ """Test that messages below threshold are all kept."""
+ handler = create_handler()
+ messages = [
+ create_user_msg("Hello"),
+ create_assistant_msg("Hi there!"),
+ create_user_msg("How are you?"),
+ ]
+ threshold, reserve = 10000, 5000
+ to_compact, to_keep, _ = asyncio.run(
+ handler.context_check(
+ messages=messages,
+ memory_compact_threshold=threshold, # Very high threshold
+ memory_compact_reserve=reserve,
+ ),
+ )
+ assert not to_compact, f"Expected empty compact list, got: {len(to_compact)}"
+ assert len(to_keep) == 3, f"Expected 3 messages to keep, got: {len(to_keep)}"
+ assert to_keep == messages, "Messages to keep should be the original messages"
+ verify_context_check_invariants(
+ handler,
+ messages,
+ to_compact,
+ to_keep,
+ threshold,
+ reserve,
+ "test_below_threshold_returns_all",
+ )
+ print_pass("test_below_threshold_returns_all")
+
+
+def test_above_threshold_triggers_compaction():
+ """Test that messages above threshold are split correctly."""
+ handler = create_handler()
+ # Create messages that will exceed threshold
+ messages = [
+ create_user_msg("First message " * 100),
+ create_assistant_msg("Second message " * 100),
+ create_user_msg("Third message " * 100),
+ create_assistant_msg("Fourth message " * 100),
+ ]
+ threshold, reserve = 100, 200
+ to_compact, to_keep, _ = asyncio.run(
+ handler.context_check(
+ messages=messages,
+ memory_compact_threshold=threshold, # Low threshold to trigger compaction
+ memory_compact_reserve=reserve,
+ ),
+ )
+ # Should have some messages compacted and some kept
+ assert len(to_compact) + len(to_keep) == len(messages), "Total messages should match"
+ assert len(to_compact) > 0, "Expected some messages to be compacted"
+ verify_context_check_invariants(
+ handler,
+ messages,
+ to_compact,
+ to_keep,
+ threshold,
+ reserve,
+ "test_above_threshold_triggers_compaction",
+ )
+ print_pass("test_above_threshold_triggers_compaction")
+
+
+def test_message_order_preserved():
+ """Test that message order is preserved in both lists."""
+ handler = create_handler()
+ messages = [
+ create_user_msg("First " * 50),
+ create_assistant_msg("Second " * 50),
+ create_user_msg("Third " * 50),
+ create_assistant_msg("Fourth " * 50),
+ create_user_msg("Fifth " * 10),
+ ]
+ threshold, reserve = 100, 150
+ to_compact, to_keep, _ = asyncio.run(
+ handler.context_check(
+ messages=messages,
+ memory_compact_threshold=threshold, # Low threshold
+ memory_compact_reserve=reserve,
+ ),
+ )
+ # Check order preservation - compact messages should appear first in original
+ all_messages = to_compact + to_keep
+ for i, msg in enumerate(all_messages):
+ assert msg in messages, f"Message {i} not found in original messages"
+ verify_context_check_invariants(
+ handler,
+ messages,
+ to_compact,
+ to_keep,
+ threshold,
+ reserve,
+ "test_message_order_preserved",
+ )
+ print_pass("test_message_order_preserved")
+
+
+# =============================================================================
+# Edge Cases - Threshold and Reserve Boundaries
+# =============================================================================
+
+
+def test_single_message_below_threshold():
+ """Test single message below threshold."""
+ handler = create_handler()
+ messages = [create_user_msg("Short message")]
+ threshold, reserve = 1000, 500
+ to_compact, to_keep, _ = asyncio.run(
+ handler.context_check(
+ messages=messages,
+ memory_compact_threshold=threshold,
+ memory_compact_reserve=reserve,
+ ),
+ )
+ assert not to_compact, "Should not compact single message below threshold"
+ assert len(to_keep) == 1, "Should keep the single message"
+ verify_context_check_invariants(
+ handler,
+ messages,
+ to_compact,
+ to_keep,
+ threshold,
+ reserve,
+ "test_single_message_below_threshold",
+ )
+ print_pass("test_single_message_below_threshold")
+
+
+def test_single_message_above_threshold():
+ """Test single message that exceeds threshold - nothing can be kept in reserve."""
+ handler = create_handler()
+ long_content = "Very long message " * 1000
+ messages = [create_user_msg(long_content)]
+ threshold, reserve = 10, 5
+ to_compact, to_keep, _ = asyncio.run(
+ handler.context_check(
+ messages=messages,
+ memory_compact_threshold=threshold, # Very low threshold
+ memory_compact_reserve=reserve, # Even lower reserve
+ ),
+ )
+ # Message exceeds both threshold and reserve, so it's compacted
+ assert len(to_compact) == 1, "Single large message should be compacted"
+ assert len(to_keep) == 0, "Nothing can fit in reserve"
+ verify_context_check_invariants(
+ handler,
+ messages,
+ to_compact,
+ to_keep,
+ threshold,
+ reserve,
+ "test_single_message_above_threshold",
+ )
+ print_pass("test_single_message_above_threshold")
+
+
+def test_reserve_zero():
+ """Test with reserve=0, no messages can be kept."""
+ handler = create_handler()
+ messages = [
+ create_user_msg("Hello"),
+ create_assistant_msg("Hi there!"),
+ ]
+ threshold, reserve = 1, 0
+ to_compact, to_keep, _ = asyncio.run(
+ handler.context_check(
+ messages=messages,
+ memory_compact_threshold=threshold, # Trigger compaction
+ memory_compact_reserve=reserve, # Zero reserve
+ ),
+ )
+ # All messages should be compacted since reserve is 0
+ assert len(to_compact) == 2, f"All messages should be compacted, got {len(to_compact)}"
+ assert len(to_keep) == 0, f"No messages should be kept, got {len(to_keep)}"
+ verify_context_check_invariants(handler, messages, to_compact, to_keep, threshold, reserve, "test_reserve_zero")
+ print_pass("test_reserve_zero")
+
+
+def test_threshold_zero():
+ """Test with threshold=0, always triggers compaction."""
+ handler = create_handler()
+ messages = [create_user_msg("A")] # Minimal message
+ threshold, reserve = 0, 1000
+ to_compact, to_keep, _ = asyncio.run(
+ handler.context_check(
+ messages=messages,
+ memory_compact_threshold=threshold, # Zero threshold - always triggers
+ memory_compact_reserve=reserve,
+ ),
+ )
+ # Even minimal message triggers compaction with threshold=0
+ # But reserve is high so it should be kept
+ assert len(to_compact) == 0 or len(to_keep) == 1, "Message should fit in reserve"
+ verify_context_check_invariants(handler, messages, to_compact, to_keep, threshold, reserve, "test_threshold_zero")
+ print_pass("test_threshold_zero")
+
+
+def test_exact_threshold_boundary():
+ """Test messages exactly at threshold boundary."""
+ handler = create_handler()
+ messages = [create_user_msg("Test message")]
+
+ # Get exact token count
+ stat = asyncio.run(handler.stat_message(messages[0]))
+ exact_tokens = stat.total_tokens
+ threshold, reserve = exact_tokens, exact_tokens
+
+ # Test at exact boundary
+ to_compact, to_keep, _ = asyncio.run(
+ handler.context_check(
+ messages=messages,
+ memory_compact_threshold=threshold, # Exactly at boundary
+ memory_compact_reserve=reserve,
+ ),
+ )
+ # At exact boundary (<=), should not trigger compaction
+ assert not to_compact, "Should not compact at exact boundary"
+ assert len(to_keep) == 1, "Should keep message at exact boundary"
+ verify_context_check_invariants(
+ handler,
+ messages,
+ to_compact,
+ to_keep,
+ threshold,
+ reserve,
+ "test_exact_threshold_boundary",
+ )
+ print_pass("test_exact_threshold_boundary")
+
+
+def test_reserve_larger_than_threshold():
+ """Test when reserve is larger than threshold (unusual but valid config)."""
+ handler = create_handler()
+ messages = [
+ create_user_msg("Message one " * 20),
+ create_assistant_msg("Message two " * 20),
+ ]
+ threshold, reserve = 50, 10000
+ to_compact, to_keep, _ = asyncio.run(
+ handler.context_check(
+ messages=messages,
+ memory_compact_threshold=threshold, # Low threshold
+ memory_compact_reserve=reserve, # High reserve
+ ),
+ )
+ # Compaction triggered but reserve can hold everything
+ # Total messages should be preserved
+ assert len(to_compact) + len(to_keep) == 2
+ verify_context_check_invariants(
+ handler,
+ messages,
+ to_compact,
+ to_keep,
+ threshold,
+ reserve,
+ "test_reserve_larger_than_threshold",
+ )
+ print_pass("test_reserve_larger_than_threshold")
+
+
+# =============================================================================
+# Edge Cases - Tool Use/Result Pairing
+# =============================================================================
+
+
+def test_tool_use_result_paired():
+ """Test that tool_use and tool_result pairs are kept together."""
+ handler = create_handler()
+ messages = [
+ create_user_msg("Please run the tool " * 50),
+ create_tool_use_msg("call_001", "test_tool", {"arg": "value"}),
+ create_tool_result_msg("call_001", "test_tool", "Tool output"),
+ create_assistant_msg("The tool returned results"),
+ ]
+ threshold, reserve = 50, 1000
+ to_compact, to_keep, _ = asyncio.run(
+ handler.context_check(
+ messages=messages,
+ memory_compact_threshold=threshold, # Trigger compaction
+ memory_compact_reserve=reserve, # Enough for tool pair
+ ),
+ )
+
+ # If tool_result is kept, tool_use should also be kept
+ tool_result_in_keep = any(any(b.get("type") == "tool_result" for b in m.get_content_blocks()) for m in to_keep)
+ tool_use_in_keep = any(any(b.get("type") == "tool_use" for b in m.get_content_blocks()) for m in to_keep)
+
+ if tool_result_in_keep:
+ assert tool_use_in_keep, "tool_use should be kept when tool_result is kept"
+ verify_context_check_invariants(
+ handler,
+ messages,
+ to_compact,
+ to_keep,
+ threshold,
+ reserve,
+ "test_tool_use_result_paired",
+ )
+ print_pass("test_tool_use_result_paired")
+
+
+def test_tool_use_without_result():
+ """Test tool_use message without corresponding tool_result."""
+ handler = create_handler()
+ messages = [
+ create_user_msg("Run the tool"),
+ create_tool_use_msg("call_orphan", "orphan_tool", {"arg": "value"}),
+ create_assistant_msg("Something happened"),
+ ]
+ threshold, reserve = 10, 1000
+ to_compact, to_keep, _ = asyncio.run(
+ handler.context_check(
+ messages=messages,
+ memory_compact_threshold=threshold,
+ memory_compact_reserve=reserve,
+ ),
+ )
+ # Should not crash, just process normally
+ assert len(to_compact) + len(to_keep) == 3
+ verify_context_check_invariants(
+ handler,
+ messages,
+ to_compact,
+ to_keep,
+ threshold,
+ reserve,
+ "test_tool_use_without_result",
+ )
+ print_pass("test_tool_use_without_result")
+
+
+def test_tool_result_without_use():
+ """Test tool_result message without corresponding tool_use."""
+ handler = create_handler()
+ messages = [
+ create_user_msg("Here's a result"),
+ create_tool_result_msg("call_orphan", "orphan_tool", "Some output"),
+ create_assistant_msg("Got it"),
+ ]
+ threshold, reserve = 10, 1000
+ to_compact, to_keep, _ = asyncio.run(
+ handler.context_check(
+ messages=messages,
+ memory_compact_threshold=threshold,
+ memory_compact_reserve=reserve,
+ ),
+ )
+ # Should not crash even with orphan tool_result
+ assert len(to_compact) + len(to_keep) == 3
+ verify_context_check_invariants(
+ handler,
+ messages,
+ to_compact,
+ to_keep,
+ threshold,
+ reserve,
+ "test_tool_result_without_use",
+ )
+ print_pass("test_tool_result_without_use")
+
+
+def test_multiple_tool_pairs():
+ """Test multiple tool_use/tool_result pairs."""
+ handler = create_handler()
+ messages = [
+ create_user_msg("Start task " * 50),
+ create_tool_use_msg("call_001", "tool_a", {"a": 1}),
+ create_tool_result_msg("call_001", "tool_a", "Result A"),
+ create_tool_use_msg("call_002", "tool_b", {"b": 2}),
+ create_tool_result_msg("call_002", "tool_b", "Result B"),
+ create_tool_use_msg("call_003", "tool_c", {"c": 3}),
+ create_tool_result_msg("call_003", "tool_c", "Result C"),
+ create_assistant_msg("All done"),
+ ]
+ threshold, reserve = 50, 500
+ to_compact, to_keep, _ = asyncio.run(
+ handler.context_check(
+ messages=messages,
+ memory_compact_threshold=threshold,
+ memory_compact_reserve=reserve,
+ ),
+ )
+
+ # Verify tool pairs integrity - for each kept tool_result, its tool_use should be kept
+ for msg in to_keep:
+ for block in msg.get_content_blocks("tool_result"):
+ tool_id = block.get("id", "")
+ if tool_id:
+ # Find corresponding tool_use
+ tool_use_found = False
+ for keep_msg in to_keep:
+ for use_block in keep_msg.get_content_blocks("tool_use"):
+ if use_block.get("id") == tool_id:
+ tool_use_found = True
+ break
+ assert tool_use_found, f"tool_use for {tool_id} should be kept with tool_result"
+ verify_context_check_invariants(
+ handler,
+ messages,
+ to_compact,
+ to_keep,
+ threshold,
+ reserve,
+ "test_multiple_tool_pairs",
+ )
+ print_pass("test_multiple_tool_pairs")
+
+
+def test_tool_dependency_causes_extra_inclusion():
+ """Test that tool_use is included even if it exceeds simple reserve calculation."""
+ handler = create_handler()
+ # Create a scenario where:
+ # - First message (tool_use) is large
+ # - Later message (tool_result) references it
+ # - Reserve alone wouldn't fit tool_use, but dependency requires it
+ large_tool_input = {"data": "x" * 200}
+ messages = [
+ create_user_msg("Start " * 100), # Large message
+ create_tool_use_msg("call_dep", "dep_tool", large_tool_input), # Medium
+ create_user_msg("Middle " * 100), # Large message
+ create_tool_result_msg("call_dep", "dep_tool", "Result"), # Small
+ create_assistant_msg("End"), # Small
+ ]
+ threshold, reserve = 100, 500
+ to_compact, to_keep, _ = asyncio.run(
+ handler.context_check(
+ messages=messages,
+ memory_compact_threshold=threshold, # Trigger compaction
+ memory_compact_reserve=reserve, # Medium reserve
+ ),
+ )
+
+ # Check pair integrity
+ result_kept = any(
+ any(b.get("id") == "call_dep" and b.get("type") == "tool_result" for b in m.get_content_blocks())
+ for m in to_keep
+ )
+ use_kept = any(
+ any(b.get("id") == "call_dep" and b.get("type") == "tool_use" for b in m.get_content_blocks()) for m in to_keep
+ )
+
+ if result_kept:
+ assert use_kept, "Dependent tool_use should be included with tool_result"
+ verify_context_check_invariants(
+ handler,
+ messages,
+ to_compact,
+ to_keep,
+ threshold,
+ reserve,
+ "test_tool_dependency_causes_extra_inclusion",
+ )
+ print_pass("test_tool_dependency_causes_extra_inclusion")
+
+
+def test_tool_dependency_exceeds_reserve():
+ """Test when tool_result + its tool_use dependency would exceed reserve."""
+ handler = create_handler()
+ # tool_use is very large, making the pair not fit in reserve
+ very_large_input = {"data": "x" * 2000}
+ messages = [
+ create_user_msg("First"),
+ create_tool_use_msg("call_big", "big_tool", very_large_input), # Very large
+ create_tool_result_msg("call_big", "big_tool", "Small result"),
+ create_assistant_msg("Last message"),
+ ]
+ threshold, reserve = 10, 100
+ to_compact, to_keep, _ = asyncio.run(
+ handler.context_check(
+ messages=messages,
+ memory_compact_threshold=threshold, # Trigger compaction
+ memory_compact_reserve=reserve, # Small reserve - can't fit the pair
+ ),
+ )
+
+ # The tool pair is too large, so it should be excluded or partially handled
+ # Either both are compacted (pair excluded) or neither is kept
+ result_kept = any(
+ any(b.get("id") == "call_big" and b.get("type") == "tool_result" for b in m.get_content_blocks())
+ for m in to_keep
+ )
+
+ if result_kept:
+ # If result is kept, use must also be kept (pair integrity)
+ use_kept = any(
+ any(b.get("id") == "call_big" and b.get("type") == "tool_use" for b in m.get_content_blocks())
+ for m in to_keep
+ )
+ assert use_kept, "Pair integrity violated"
+ verify_context_check_invariants(
+ handler,
+ messages,
+ to_compact,
+ to_keep,
+ threshold,
+ reserve,
+ "test_tool_dependency_exceeds_reserve",
+ )
+ print_pass("test_tool_dependency_exceeds_reserve")
+
+
+def test_interleaved_tool_pairs():
+ """Test interleaved tool_use/tool_result (not strictly sequential)."""
+ handler = create_handler()
+ messages = [
+ create_user_msg("Multi-tool task " * 30),
+ create_tool_use_msg("call_a", "tool_a", {"a": 1}),
+ create_tool_use_msg("call_b", "tool_b", {"b": 2}), # Two uses before results
+ create_tool_result_msg("call_a", "tool_a", "Result A"),
+ create_tool_result_msg("call_b", "tool_b", "Result B"),
+ create_assistant_msg("Both done"),
+ ]
+ threshold, reserve = 50, 1000
+ to_compact, to_keep, _ = asyncio.run(
+ handler.context_check(
+ messages=messages,
+ memory_compact_threshold=threshold,
+ memory_compact_reserve=reserve,
+ ),
+ )
+
+ # Verify pair integrity for interleaved pairs
+ for msg in to_keep:
+ for block in msg.get_content_blocks("tool_result"):
+ tool_id = block.get("id", "")
+ if tool_id:
+ use_found = any(
+ any(ub.get("id") == tool_id and ub.get("type") == "tool_use" for ub in km.get_content_blocks())
+ for km in to_keep
+ )
+ assert use_found, f"Interleaved tool_use {tool_id} should be kept"
+ verify_context_check_invariants(
+ handler,
+ messages,
+ to_compact,
+ to_keep,
+ threshold,
+ reserve,
+ "test_interleaved_tool_pairs",
+ )
+ print_pass("test_interleaved_tool_pairs")
+
+
+# =============================================================================
+# Edge Cases - Message Content Variations
+# =============================================================================
+
+
+def test_message_with_empty_content():
+ """Test message with empty string content."""
+ handler = create_handler()
+ messages = [
+ create_user_msg(""), # Empty content
+ create_assistant_msg("Response"),
+ ]
+ threshold, reserve = 1000, 500
+ to_compact, to_keep, _ = asyncio.run(
+ handler.context_check(
+ messages=messages,
+ memory_compact_threshold=threshold,
+ memory_compact_reserve=reserve,
+ ),
+ )
+ assert len(to_compact) + len(to_keep) == 2
+ verify_context_check_invariants(
+ handler,
+ messages,
+ to_compact,
+ to_keep,
+ threshold,
+ reserve,
+ "test_message_with_empty_content",
+ )
+ print_pass("test_message_with_empty_content")
+
+
+def test_message_with_whitespace_only():
+ """Test message with whitespace-only content."""
+ handler = create_handler()
+ messages = [
+ create_user_msg(" \n\t "), # Whitespace only
+ create_assistant_msg("Response"),
+ ]
+ threshold, reserve = 1000, 500
+ to_compact, to_keep, _ = asyncio.run(
+ handler.context_check(
+ messages=messages,
+ memory_compact_threshold=threshold,
+ memory_compact_reserve=reserve,
+ ),
+ )
+ assert len(to_compact) + len(to_keep) == 2
+ verify_context_check_invariants(
+ handler,
+ messages,
+ to_compact,
+ to_keep,
+ threshold,
+ reserve,
+ "test_message_with_whitespace_only",
+ )
+ print_pass("test_message_with_whitespace_only")
+
+
+def test_very_long_single_message():
+ """Test very long single message that exceeds any reasonable reserve."""
+ handler = create_handler()
+ huge_content = "x" * 100000 # Very long
+ messages = [create_user_msg(huge_content)]
+ threshold, reserve = 100, 1000
+ to_compact, to_keep, _ = asyncio.run(
+ handler.context_check(
+ messages=messages,
+ memory_compact_threshold=threshold,
+ memory_compact_reserve=reserve,
+ ),
+ )
+ # Single huge message - either kept alone or compacted
+ assert len(to_compact) + len(to_keep) == 1
+ verify_context_check_invariants(
+ handler,
+ messages,
+ to_compact,
+ to_keep,
+ threshold,
+ reserve,
+ "test_very_long_single_message",
+ )
+ print_pass("test_very_long_single_message")
+
+
+def test_many_small_messages():
+ """Test many small messages."""
+ handler = create_handler()
+ messages = [create_user_msg(f"Msg {i}") for i in range(100)]
+ threshold, reserve = 100, 200
+ to_compact, to_keep, _ = asyncio.run(
+ handler.context_check(
+ messages=messages,
+ memory_compact_threshold=threshold, # Low threshold
+ memory_compact_reserve=reserve,
+ ),
+ )
+ # Should compact older messages and keep recent ones
+ assert len(to_compact) + len(to_keep) == 100
+ assert len(to_keep) > 0, "Should keep some messages"
+ verify_context_check_invariants(
+ handler,
+ messages,
+ to_compact,
+ to_keep,
+ threshold,
+ reserve,
+ "test_many_small_messages",
+ )
+ print_pass("test_many_small_messages")
+
+
+def test_unicode_content():
+ """Test messages with unicode characters."""
+ handler = create_handler()
+ messages = [
+ create_user_msg("你好世界！🎉 Emoji and 中文"),
+ create_assistant_msg("مرحبا العالم 🌍 Arabic and more"),
+ create_user_msg("日本語テスト 🇯🇵"),
+ ]
+ threshold, reserve = 1000, 500
+ to_compact, to_keep, _ = asyncio.run(
+ handler.context_check(
+ messages=messages,
+ memory_compact_threshold=threshold,
+ memory_compact_reserve=reserve,
+ ),
+ )
+ assert len(to_compact) + len(to_keep) == 3
+ verify_context_check_invariants(handler, messages, to_compact, to_keep, threshold, reserve, "test_unicode_content")
+ print_pass("test_unicode_content")
+
+
+def test_special_characters_content():
+ """Test messages with special characters."""
+ handler = create_handler()
+ messages = [
+ create_user_msg("Special chars: <>&\"'`~!@#$%^&*()[]{}|\\"),
+ create_assistant_msg("More: \n\r\t\0 nulls and newlines"),
+ ]
+ threshold, reserve = 1000, 500
+ to_compact, to_keep, _ = asyncio.run(
+ handler.context_check(
+ messages=messages,
+ memory_compact_threshold=threshold,
+ memory_compact_reserve=reserve,
+ ),
+ )
+ assert len(to_compact) + len(to_keep) == 2
+ verify_context_check_invariants(
+ handler,
+ messages,
+ to_compact,
+ to_keep,
+ threshold,
+ reserve,
+ "test_special_characters_content",
+ )
+ print_pass("test_special_characters_content")
+
+
+# =============================================================================
+# Edge Cases - Boundary Conditions
+# =============================================================================
+
+
+def test_all_messages_fit_exactly_in_reserve():
+ """Test when all messages fit exactly in reserve after threshold exceeded."""
+ handler = create_handler()
+ messages = [
+ create_user_msg("Message 1"),
+ create_assistant_msg("Message 2"),
+ ]
+
+ # Calculate total tokens
+ total = sum(asyncio.run(handler.stat_message(m)).total_tokens for m in messages)
+ threshold, reserve = total - 1, total
+
+ to_compact, to_keep, _ = asyncio.run(
+ handler.context_check(
+ messages=messages,
+ memory_compact_threshold=threshold, # Just below total to trigger
+ memory_compact_reserve=reserve, # Exactly fits all
+ ),
+ )
+ # All should be kept since reserve can hold everything
+ assert len(to_keep) == 2, f"All messages should fit in reserve, got {len(to_keep)}"
+ verify_context_check_invariants(
+ handler,
+ messages,
+ to_compact,
+ to_keep,
+ threshold,
+ reserve,
+ "test_all_messages_fit_exactly_in_reserve",
+ )
+ print_pass("test_all_messages_fit_exactly_in_reserve")
+
+
+def test_first_message_only_compacted():
+ """Test when only the first message is compacted."""
+ handler = create_handler()
+ messages = [
+ create_user_msg("Large first message " * 100), # Large
+ create_assistant_msg("Small"), # Small
+ create_user_msg("Tiny"), # Tiny
+ ]
+
+ # Calculate tokens to set appropriate reserve
+ small_msg_tokens = asyncio.run(handler.stat_message(messages[1])).total_tokens
+ tiny_msg_tokens = asyncio.run(handler.stat_message(messages[2])).total_tokens
+ threshold, reserve = 50, small_msg_tokens + tiny_msg_tokens + 10
+
+ to_compact, to_keep, _ = asyncio.run(
+ handler.context_check(
+ messages=messages,
+ memory_compact_threshold=threshold, # Low to trigger
+ memory_compact_reserve=reserve, # Fits last 2
+ ),
+ )
+
+ assert len(to_compact) >= 1, "At least first message should be compacted"
+ verify_context_check_invariants(
+ handler,
+ messages,
+ to_compact,
+ to_keep,
+ threshold,
+ reserve,
+ "test_first_message_only_compacted",
+ )
+ print_pass("test_first_message_only_compacted")
+
+
+def test_last_message_only_kept():
+ """Test when only the last message can be kept."""
+ handler = create_handler()
+ messages = [
+ create_user_msg("Large " * 200),
+ create_assistant_msg("Large " * 200),
+ create_user_msg("Tiny"), # Only this fits
+ ]
+
+ tiny_tokens = asyncio.run(handler.stat_message(messages[2])).total_tokens
+ threshold, reserve = 10, tiny_tokens + 5
+
+ to_compact, to_keep, _ = asyncio.run(
+ handler.context_check(
+ messages=messages,
+ memory_compact_threshold=threshold,
+ memory_compact_reserve=reserve, # Only fits last message
+ ),
+ )
+
+ if len(to_keep) == 1:
+ # Last message should be the one kept
+ assert to_keep[0] == messages[2], "Only last message should be kept"
+ verify_context_check_invariants(
+ handler,
+ messages,
+ to_compact,
+ to_keep,
+ threshold,
+ reserve,
+ "test_last_message_only_kept",
+ )
+ print_pass("test_last_message_only_kept")
+
+
+def test_all_messages_compacted():
+ """Test when all messages need to be compacted (nothing fits in reserve)."""
+ handler = create_handler()
+ messages = [
+ create_user_msg("Large message " * 100),
+ create_assistant_msg("Large message " * 100),
+ ]
+ threshold, reserve = 10, 1
+ to_compact, to_keep, _ = asyncio.run(
+ handler.context_check(
+ messages=messages,
+ memory_compact_threshold=threshold, # Trigger compaction
+ memory_compact_reserve=reserve, # Too small for anything
+ ),
+ )
+ assert len(to_compact) == 2, "All messages should be compacted"
+ assert len(to_keep) == 0, "No messages should be kept"
+ verify_context_check_invariants(
+ handler,
+ messages,
+ to_compact,
+ to_keep,
+ threshold,
+ reserve,
+ "test_all_messages_compacted",
+ )
+ print_pass("test_all_messages_compacted")
+
+
+# =============================================================================
+# Edge Cases - Message Roles
+# =============================================================================
+
+
+def test_system_message():
+ """Test handling of system role messages."""
+ handler = create_handler()
+ system_msg = Msg(name="system", role="system", content="You are a helpful assistant.")
+ messages = [
+ system_msg,
+ create_user_msg("Hello"),
+ create_assistant_msg("Hi there!"),
+ ]
+ threshold, reserve = 1000, 500
+ to_compact, to_keep, _ = asyncio.run(
+ handler.context_check(
+ messages=messages,
+ memory_compact_threshold=threshold,
+ memory_compact_reserve=reserve,
+ ),
+ )
+ assert len(to_compact) + len(to_keep) == 3
+ verify_context_check_invariants(handler, messages, to_compact, to_keep, threshold, reserve, "test_system_message")
+ print_pass("test_system_message")
+
+
+def test_mixed_roles():
+ """Test messages with various roles (user, assistant, system)."""
+ handler = create_handler()
+ # agentscope.message.Msg only supports: user, assistant, system
+ messages = [
+ Msg(name="system", role="system", content="System prompt"),
+ Msg(name="user", role="user", content="User message"),
+ Msg(name="assistant", role="assistant", content="Assistant response"),
+ Msg(name="tool", role="user", content="Tool output as user role"),
+ Msg(name="helper", role="assistant", content="Another assistant message"),
+ ]
+ threshold, reserve = 1000, 500
+ to_compact, to_keep, _ = asyncio.run(
+ handler.context_check(
+ messages=messages,
+ memory_compact_threshold=threshold,
+ memory_compact_reserve=reserve,
+ ),
+ )
+ assert len(to_compact) + len(to_keep) == 5
+ verify_context_check_invariants(handler, messages, to_compact, to_keep, threshold, reserve, "test_mixed_roles")
+ print_pass("test_mixed_roles")
+
+
+# =============================================================================
+# Edge Cases - Tool Block Variations
+# =============================================================================
+
+
+def test_tool_use_with_empty_id():
+ """Test tool_use block with empty id."""
+ handler = create_handler()
+ messages = [
+ create_user_msg("Run tool"),
+ create_tool_use_msg("", "test_tool", {"arg": "value"}), # Empty ID
+ create_assistant_msg("Done"),
+ ]
+ threshold, reserve = 10, 1000
+ to_compact, to_keep, _ = asyncio.run(
+ handler.context_check(
+ messages=messages,
+ memory_compact_threshold=threshold,
+ memory_compact_reserve=reserve,
+ ),
+ )
+ # Should handle gracefully
+ assert len(to_compact) + len(to_keep) == 3
+ verify_context_check_invariants(
+ handler,
+ messages,
+ to_compact,
+ to_keep,
+ threshold,
+ reserve,
+ "test_tool_use_with_empty_id",
+ )
+ print_pass("test_tool_use_with_empty_id")
+
+
+def test_tool_result_with_empty_id():
+ """Test tool_result block with empty id."""
+ handler = create_handler()
+ messages = [
+ create_user_msg("Got result"),
+ create_tool_result_msg("", "test_tool", "Output"), # Empty ID
+ create_assistant_msg("Noted"),
+ ]
+ threshold, reserve = 10, 1000
+ to_compact, to_keep, _ = asyncio.run(
+ handler.context_check(
+ messages=messages,
+ memory_compact_threshold=threshold,
+ memory_compact_reserve=reserve,
+ ),
+ )
+ # Should handle gracefully
+ assert len(to_compact) + len(to_keep) == 3
+ verify_context_check_invariants(
+ handler,
+ messages,
+ to_compact,
+ to_keep,
+ threshold,
+ reserve,
+ "test_tool_result_with_empty_id",
+ )
+ print_pass("test_tool_result_with_empty_id")
+
+
+def test_duplicate_tool_ids():
+ """Test messages with duplicate tool IDs (unusual but possible)."""
+ handler = create_handler()
+ messages = [
+ create_tool_use_msg("call_dup", "tool_a", {"a": 1}),
+ create_tool_result_msg("call_dup", "tool_a", "Result A"),
+ create_tool_use_msg("call_dup", "tool_b", {"b": 2}), # Same ID, different tool
+ create_tool_result_msg("call_dup", "tool_b", "Result B"),
+ ]
+ threshold, reserve = 10, 1000
+ to_compact, to_keep, _ = asyncio.run(
+ handler.context_check(
+ messages=messages,
+ memory_compact_threshold=threshold,
+ memory_compact_reserve=reserve,
+ ),
+ )
+ # Should not crash with duplicate IDs
+ assert len(to_compact) + len(to_keep) == 4
+ verify_context_check_invariants(
+ handler,
+ messages,
+ to_compact,
+ to_keep,
+ threshold,
+ reserve,
+ "test_duplicate_tool_ids",
+ )
+ print_pass("test_duplicate_tool_ids")
+
+
+def test_message_with_multiple_tool_blocks():
+ """Test single message containing multiple tool blocks."""
+ handler = create_handler()
+ msg_with_multiple_tools = Msg(
+ name="assistant",
+ role="assistant",
+ content=[
+ {"type": "tool_use", "id": "call_1", "name": "tool1", "input": {}},
+ {"type": "tool_use", "id": "call_2", "name": "tool2", "input": {}},
+ {"type": "tool_use", "id": "call_3", "name": "tool3", "input": {}},
+ ],
+ )
+ messages = [
+ create_user_msg("Do multiple things"),
+ msg_with_multiple_tools,
+ create_tool_result_msg("call_1", "tool1", "Result 1"),
+ create_tool_result_msg("call_2", "tool2", "Result 2"),
+ create_tool_result_msg("call_3", "tool3", "Result 3"),
+ ]
+ threshold, reserve = 10, 2000
+ to_compact, to_keep, _ = asyncio.run(
+ handler.context_check(
+ messages=messages,
+ memory_compact_threshold=threshold,
+ memory_compact_reserve=reserve,
+ ),
+ )
+ assert len(to_compact) + len(to_keep) == 5
+ verify_context_check_invariants(
+ handler,
+ messages,
+ to_compact,
+ to_keep,
+ threshold,
+ reserve,
+ "test_message_with_multiple_tool_blocks",
+ )
+ print_pass("test_message_with_multiple_tool_blocks")
+
+
+# =============================================================================
+# Run All Tests
+# =============================================================================
+
+
+def run_all_tests():
+ """Run all tests."""
+ tests = [
+ # Normal cases
+ test_empty_messages,
+ test_below_threshold_returns_all,
+ test_above_threshold_triggers_compaction,
+ test_message_order_preserved,
+ # Edge cases - boundaries
+ test_single_message_below_threshold,
+ test_single_message_above_threshold,
+ test_reserve_zero,
+ test_threshold_zero,
+ test_exact_threshold_boundary,
+ test_reserve_larger_than_threshold,
+ # Edge cases - tool pairing
+ test_tool_use_result_paired,
+ test_tool_use_without_result,
+ test_tool_result_without_use,
+ test_multiple_tool_pairs,
+ test_tool_dependency_causes_extra_inclusion,
+ test_tool_dependency_exceeds_reserve,
+ test_interleaved_tool_pairs,
+ # Edge cases - content variations
+ test_message_with_empty_content,
+ test_message_with_whitespace_only,
+ test_very_long_single_message,
+ test_many_small_messages,
+ test_unicode_content,
+ test_special_characters_content,
+ # Edge cases - boundaries
+ test_all_messages_fit_exactly_in_reserve,
+ test_first_message_only_compacted,
+ test_last_message_only_kept,
+ test_all_messages_compacted,
+ # Edge cases - roles
+ test_system_message,
+ test_mixed_roles,
+ # Edge cases - tool blocks
+ test_tool_use_with_empty_id,
+ test_tool_result_with_empty_id,
+ test_duplicate_tool_ids,
+ test_message_with_multiple_tool_blocks,
+ ]
+
+ passed = 0
+ failed = 0
+
+ for test in tests:
+ try:
+ print_test_header(test.__name__)
+ test()
+ passed += 1
+ except AssertionError as e:
+ print_fail(test.__name__, str(e))
+ failed += 1
+ except Exception as e:
+ print_error(test.__name__, str(e))
+ failed += 1
+
+ # Print summary
+ print(f"\n{Colors.CYAN}{'=' * 60}{Colors.RESET}")
+ print(f"{Colors.BOLD}Test Results Summary{Colors.RESET}")
+ print(f"{Colors.CYAN}{'=' * 60}{Colors.RESET}")
+ print(f"{Colors.GREEN}{Colors.BOLD}✓ Passed: {passed}{Colors.RESET}")
+ if failed > 0:
+ print(f"{Colors.RED}{Colors.BOLD}✗ Failed: {failed}{Colors.RESET}")
+ else:
+ print(f"{Colors.GREEN}✗ Failed: {failed}{Colors.RESET}")
+ print(f"{Colors.CYAN}{'=' * 60}{Colors.RESET}")
+
+ if failed == 0:
+ print(f"\n{Colors.GREEN}{Colors.BOLD}🎉 All tests passed!{Colors.RESET}")
+ else:
+ print(f"\n{Colors.RED}{Colors.BOLD}💥 Some tests failed!{Colors.RESET}")
+
+
+if __name__ == "__main__":
+ run_all_tests()
diff --git a/tests/light/test_format_msgs_to_str.py b/tests/light/test_format_msgs_to_str.py
new file mode 100644
index 00000000..978a89f4
--- /dev/null
+++ b/tests/light/test_format_msgs_to_str.py
@@ -0,0 +1,908 @@
+"""Tests for AsMsgHandler.format_msgs_to_str method."""
+
+# pylint: disable=W0212
+
+import asyncio
+import sys
+
+from agentscope.message import Msg
+from test_utils import get_token_counter
+
+from reme.core.utils import get_logger
+from reme.memory.file_based.utils import AsMsgHandler
+
+logger = get_logger()
+
+
+# ANSI 颜色码
+class Colors:
+ """ANSI color codes for terminal output."""
+
+ GREEN = "\033[92m"
+ RED = "\033[91m"
+ YELLOW = "\033[93m"
+ BLUE = "\033[94m"
+ CYAN = "\033[96m"
+ BOLD = "\033[1m"
+ RESET = "\033[0m"
+
+
+def print_pass(test_name: str):
+ """打印测试通过信息"""
+ print(f"{Colors.GREEN}{Colors.BOLD}✓ {test_name} PASSED{Colors.RESET}")
+
+
+def print_fail(test_name: str, error: str):
+ """打印测试失败信息"""
+ print(f"{Colors.RED}{Colors.BOLD}✗ {test_name} FAILED: {error}{Colors.RESET}")
+
+
+def print_error(test_name: str, error: str):
+ """打印测试错误信息"""
+ print(f"{Colors.YELLOW}{Colors.BOLD}⚠ {test_name} ERROR: {error}{Colors.RESET}")
+
+
+def print_test_header(test_name: str):
+ """打印测试标题"""
+ print(f"\n{Colors.CYAN}{'=' * 60}{Colors.RESET}")
+ print(f"{Colors.BLUE}{Colors.BOLD}Running: {test_name}{Colors.RESET}")
+ print(f"{Colors.CYAN}{'=' * 60}{Colors.RESET}")
+
+
+# ==================== Helper Functions ====================
+
+
+def create_handler() -> AsMsgHandler:
+ """Create an AsMsgHandler instance for testing."""
+ return AsMsgHandler(token_counter=get_token_counter())
+
+
+def verify_result_within_threshold(
+ handler: AsMsgHandler,
+ result: str,
+ threshold: int,
+ test_name: str = "",
+ msgs: list[Msg] | None = None,
+) -> None:
+ """Verify that the included messages' original token count does not exceed threshold.
+
+ Note: The format_msgs_to_str method uses message token statistics (not formatted
+ string tokens) for threshold checking. The formatted result may have more tokens
+ than the threshold due to added metadata (timestamps, role prefixes, etc.).
+
+ This verification checks that included messages' original token sum <= threshold.
+
+ Args:
+ handler: The AsMsgHandler instance used for token counting.
+ result: The formatted string result from format_msgs_to_str.
+ threshold: The memory_compact_threshold value used.
+ test_name: Optional test name for better error messages.
+ msgs: Optional list of original messages to verify against.
+
+ Raises:
+ AssertionError: If included messages' token count exceeds threshold.
+ """
+ if not result or not msgs:
+ return # Empty result or no messages to verify
+
+ # Calculate tokens of messages that were included in the result
+ included_tokens = 0
+ for msg in msgs:
+ stat = asyncio.run(handler.stat_message(msg))
+ # Check if this message's content appears in the result
+ _ = stat.format(include_thinking=True) # Use True to check all content
+ # Simple heuristic: if the message content is in result, count its tokens
+ content_blocks = msg.get_content_blocks()
+ msg_included = False
+ for block in content_blocks:
+ block_type = block.get("type", "")
+ if block_type == "text" and block.get("text", "") in result:
+ msg_included = True
+ break
+ if block_type == "tool_use" and f"{block.get('name', '')}" in result:
+ msg_included = True
+ break
+ if block_type == "tool_result" and f"{block.get('name', '')}" in result:
+ msg_included = True
+ break
+
+ if msg_included:
+ included_tokens += stat.total_tokens
+
+ # Verify included messages' token sum doesn't exceed threshold
+ # Allow small tolerance for edge cases
+ assert (
+ included_tokens <= threshold + 1
+ ), f"{test_name}: Included messages token count ({included_tokens}) exceeds threshold ({threshold})."
+
+
+def create_user_msg(content: str) -> Msg:
+ """Create a user message."""
+ return Msg(name="user", role="user", content=content)
+
+
+def create_assistant_msg(content: str) -> Msg:
+ """Create an assistant message."""
+ return Msg(name="assistant", role="assistant", content=content)
+
+
+def create_tool_use_msg(tool_name: str, tool_input: dict, tool_id: str = "call_123") -> Msg:
+ """Create a message with tool_use content block."""
+ return Msg(
+ name="assistant",
+ role="assistant",
+ content=[
+ {
+ "type": "tool_use",
+ "id": tool_id,
+ "name": tool_name,
+ "input": tool_input,
+ },
+ ],
+ )
+
+
+def create_tool_result_msg(tool_name: str, output: str | list[dict], tool_id: str = "call_123") -> Msg:
+ """Create a message with tool_result content block."""
+ return Msg(
+ name="tool",
+ role="user",
+ content=[
+ {
+ "type": "tool_result",
+ "id": tool_id,
+ "name": tool_name,
+ "output": output,
+ },
+ ],
+ )
+
+
+def create_thinking_msg(thinking_content: str, text_content: str = "") -> Msg:
+ """Create a message with thinking content block."""
+ content = [
+ {
+ "type": "thinking",
+ "thinking": thinking_content,
+ },
+ ]
+ if text_content:
+ content.append({"type": "text", "text": text_content})
+ return Msg(name="assistant", role="assistant", content=content)
+
+
+def create_image_msg(url: str = "") -> Msg:
+ """Create a message with image content block."""
+ content = [
+ {
+ "type": "image",
+ "source": {"url": url} if url else {},
+ },
+ ]
+ return Msg(name="assistant", role="assistant", content=content)
+
+
+def create_mixed_content_msg(
+ text: str = "",
+ thinking: str = "",
+ tool_name: str = "",
+ tool_input: dict | None = None,
+ image_url: str = "",
+) -> Msg:
+ """Create a message with mixed content blocks."""
+ content = []
+ if thinking:
+ content.append({"type": "thinking", "thinking": thinking})
+ if text:
+ content.append({"type": "text", "text": text})
+ if tool_name:
+ content.append(
+ {
+ "type": "tool_use",
+ "id": "call_mixed",
+ "name": tool_name,
+ "input": tool_input or {},
+ },
+ )
+ if image_url:
+ content.append({"type": "image", "source": {"url": image_url}})
+ return Msg(name="assistant", role="assistant", content=content)
+
+
+# ==================== Normal Case Tests ====================
+
+
+def test_format_msgs_to_str_empty_list():
+ """Test format_msgs_to_str with empty message list."""
+ handler = create_handler()
+ threshold = 4000
+ msgs = []
+ result = asyncio.run(handler.format_msgs_to_str(msgs, memory_compact_threshold=threshold))
+ assert result == "", f"Expected empty string for empty list, got: {result}"
+ verify_result_within_threshold(handler, result, threshold, "empty_list", msgs)
+ print_pass("test_format_msgs_to_str_empty_list")
+
+
+def test_format_msgs_to_str_single_message():
+ """Test format_msgs_to_str with a single message."""
+ handler = create_handler()
+ threshold = 4000
+ msgs = [create_user_msg("Hello, how are you?")]
+ result = asyncio.run(handler.format_msgs_to_str(msgs, memory_compact_threshold=threshold))
+
+ assert "user:" in result, f"Expected 'user:' in result, got: {result}"
+ assert "Hello, how are you?" in result, f"Expected content in result, got: {result}"
+ verify_result_within_threshold(handler, result, threshold, "single_message", msgs)
+ print_pass("test_format_msgs_to_str_single_message")
+
+
+def test_format_msgs_to_str_multiple_messages():
+ """Test format_msgs_to_str with multiple messages."""
+ handler = create_handler()
+ threshold = 4000
+ msgs = [
+ create_user_msg("What is Python?"),
+ create_assistant_msg("Python is a programming language."),
+ create_user_msg("Tell me more."),
+ create_assistant_msg("Python is known for its readability."),
+ ]
+ result = asyncio.run(handler.format_msgs_to_str(msgs, memory_compact_threshold=threshold))
+
+ assert "What is Python?" in result
+ assert "Python is a programming language." in result
+ assert "Tell me more." in result
+ assert "Python is known for its readability." in result
+ verify_result_within_threshold(handler, result, threshold, "multiple_messages", msgs)
+ print_pass("test_format_msgs_to_str_multiple_messages")
+
+
+def test_format_msgs_to_str_message_order():
+ """Test that messages are returned in correct order (oldest to newest)."""
+ handler = create_handler()
+ threshold = 4000
+ msgs = [
+ create_user_msg("First message"),
+ create_assistant_msg("Second message"),
+ create_user_msg("Third message"),
+ ]
+ result = asyncio.run(handler.format_msgs_to_str(msgs, memory_compact_threshold=threshold))
+
+ # Find positions of each message
+ first_pos = result.find("First message")
+ second_pos = result.find("Second message")
+ third_pos = result.find("Third message")
+
+ assert first_pos < second_pos < third_pos, (
+ f"Messages not in correct order. Positions: first={first_pos}, " f"second={second_pos}, third={third_pos}"
+ )
+ verify_result_within_threshold(handler, result, threshold, "message_order", msgs)
+ print_pass("test_format_msgs_to_str_message_order")
+
+
+def test_format_msgs_to_str_with_tool_use():
+ """Test format_msgs_to_str with tool_use message."""
+ handler = create_handler()
+ threshold = 4000
+ msgs = [create_tool_use_msg("read_file", {"path": "/test.txt"})]
+ result = asyncio.run(handler.format_msgs_to_str(msgs, memory_compact_threshold=threshold))
+
+ assert "read_file" in result, f"Expected tool_use in result, got: {result}"
+ verify_result_within_threshold(handler, result, threshold, "with_tool_use", msgs)
+ print_pass("test_format_msgs_to_str_with_tool_use")
+
+
+def test_format_msgs_to_str_with_tool_result():
+ """Test format_msgs_to_str with tool_result message."""
+ handler = create_handler()
+ threshold = 4000
+ msgs = [create_tool_result_msg("read_file", "file content here")]
+ result = asyncio.run(handler.format_msgs_to_str(msgs, memory_compact_threshold=threshold))
+
+ assert "read_file" in result, f"Expected tool_result in result, got: {result}"
+ verify_result_within_threshold(handler, result, threshold, "with_tool_result", msgs)
+ print_pass("test_format_msgs_to_str_with_tool_result")
+
+
+def test_format_msgs_to_str_with_image():
+ """Test format_msgs_to_str with image message."""
+ handler = create_handler()
+ threshold = 4000
+ msgs = [create_image_msg("https://example.com/image.png")]
+ result = asyncio.run(handler.format_msgs_to_str(msgs, memory_compact_threshold=threshold))
+
+ assert "" in result, f"Expected '' in result, got: {result}"
+ verify_result_within_threshold(handler, result, threshold, "with_image", msgs)
+ print_pass("test_format_msgs_to_str_with_image")
+
+
+def test_format_msgs_to_str_conversation_flow():
+ """Test format_msgs_to_str with a complete conversation flow."""
+ handler = create_handler()
+ threshold = 4000
+ msgs = [
+ create_user_msg("Read the file."),
+ create_tool_use_msg("read_file", {"path": "/data.txt"}),
+ create_tool_result_msg("read_file", "File content here"),
+ create_assistant_msg("The file contains: File content here"),
+ ]
+ result = asyncio.run(handler.format_msgs_to_str(msgs, memory_compact_threshold=threshold))
+
+ assert "user:" in result
+ assert "read_file" in result
+ assert "read_file" in result
+ assert "assistant:" in result
+ verify_result_within_threshold(handler, result, threshold, "conversation_flow", msgs)
+ print_pass("test_format_msgs_to_str_conversation_flow")
+
+
+# ==================== Thinking Block Tests ====================
+
+
+def test_format_msgs_to_str_thinking_excluded_by_default():
+ """Test that thinking blocks are excluded when include_thinking=False (default)."""
+ handler = create_handler()
+ threshold = 4000
+ msgs = [create_thinking_msg("Let me think about this...", "Here is my response")]
+ result = asyncio.run(handler.format_msgs_to_str(msgs, memory_compact_threshold=threshold, include_thinking=False))
+
+ assert "Let me think about this" not in result, f"Thinking content should be excluded, got: {result}"
+ assert "Here is my response" in result, f"Text content should be included, got: {result}"
+ verify_result_within_threshold(handler, result, threshold, "thinking_excluded_by_default", msgs)
+ print_pass("test_format_msgs_to_str_thinking_excluded_by_default")
+
+
+def test_format_msgs_to_str_thinking_included():
+ """Test that thinking blocks are included when include_thinking=True."""
+ handler = create_handler()
+ threshold = 4000
+ msgs = [create_thinking_msg("Let me think about this...", "Here is my response")]
+ result = asyncio.run(handler.format_msgs_to_str(msgs, memory_compact_threshold=threshold, include_thinking=True))
+
+ assert "Let me think about this" in result, f"Thinking content should be included, got: {result}"
+ assert "" in result, f"Expected thinking tag in result, got: {result}"
+ verify_result_within_threshold(handler, result, threshold, "thinking_included", msgs)
+ print_pass("test_format_msgs_to_str_thinking_included")
+
+
+def test_format_msgs_to_str_thinking_only_message():
+ """Test message with only thinking block."""
+ handler = create_handler()
+ threshold = 4000
+ msgs = [create_thinking_msg("Deep thoughts here")]
+
+ # With include_thinking=False
+ result_no_thinking = asyncio.run(
+ handler.format_msgs_to_str(
+ msgs,
+ memory_compact_threshold=threshold,
+ include_thinking=False,
+ ),
+ )
+ # With include_thinking=True
+ result_with_thinking = asyncio.run(
+ handler.format_msgs_to_str(
+ msgs,
+ memory_compact_threshold=threshold,
+ include_thinking=True,
+ ),
+ )
+
+ assert "Deep thoughts here" not in result_no_thinking
+ assert "Deep thoughts here" in result_with_thinking
+ verify_result_within_threshold(handler, result_no_thinking, threshold, "thinking_only_no_thinking", msgs)
+ verify_result_within_threshold(handler, result_with_thinking, threshold, "thinking_only_with_thinking", msgs)
+ print_pass("test_format_msgs_to_str_thinking_only_message")
+
+
+# ==================== Token Threshold Tests ====================
+
+
+def test_format_msgs_to_str_all_within_threshold():
+ """Test all messages fit within threshold."""
+ handler = create_handler()
+ threshold = 10000
+ msgs = [
+ create_user_msg("Short message 1"),
+ create_assistant_msg("Short message 2"),
+ create_user_msg("Short message 3"),
+ ]
+ result = asyncio.run(handler.format_msgs_to_str(msgs, memory_compact_threshold=threshold))
+
+ assert "Short message 1" in result
+ assert "Short message 2" in result
+ assert "Short message 3" in result
+ verify_result_within_threshold(handler, result, threshold, "all_within_threshold", msgs)
+ print_pass("test_format_msgs_to_str_all_within_threshold")
+
+
+def test_format_msgs_to_str_exceeds_threshold_truncate_older():
+ """Test that older messages are truncated when exceeding threshold."""
+ handler = create_handler()
+ threshold = 500
+ msgs = []
+ for i in range(20):
+ msgs.append(create_user_msg(f"Question {i}: " + "x" * 100))
+ msgs.append(create_assistant_msg(f"Answer {i}: " + "y" * 100))
+
+ result = asyncio.run(handler.format_msgs_to_str(msgs, memory_compact_threshold=threshold))
+
+ # The newest messages should be present
+ assert (
+ "Answer 19" in result or "Question 19" in result
+ ), f"Expected recent message in result, got: {result[:500]}..."
+ # Older messages should be truncated
+ assert "Question 0" not in result, "Older messages should be truncated"
+ verify_result_within_threshold(handler, result, threshold, "exceeds_threshold_truncate_older", msgs)
+ print_pass("test_format_msgs_to_str_exceeds_threshold_truncate_older")
+
+
+def test_format_msgs_to_str_single_message_exceeds_threshold():
+ """Test when a single message exceeds the threshold."""
+ handler = create_handler()
+ threshold = 10
+ # Create a very long message
+ long_text = "x" * 10000
+ msgs = [create_user_msg(long_text)]
+
+ # With very low threshold, even a single message won't fit
+ result = asyncio.run(handler.format_msgs_to_str(msgs, memory_compact_threshold=threshold))
+
+ # The message should be skipped entirely since it exceeds threshold
+ assert result == "" or len(result) > 0, "Result should be empty or contain truncated content"
+ verify_result_within_threshold(handler, result, threshold, "single_message_exceeds_threshold", msgs)
+ print_pass("test_format_msgs_to_str_single_message_exceeds_threshold")
+
+
+def test_format_msgs_to_str_first_message_exceeds_threshold():
+ """Test when the first (oldest) message exceeds threshold but newer ones don't."""
+ handler = create_handler()
+ threshold = 100
+ msgs = [
+ create_user_msg("x" * 5000), # Old, long message
+ create_assistant_msg("Short response"), # New, short message
+ ]
+
+ result = asyncio.run(handler.format_msgs_to_str(msgs, memory_compact_threshold=threshold))
+
+ # Newer message should be present
+ assert "Short response" in result, f"Expected newer message in result, got: {result}"
+ verify_result_within_threshold(handler, result, threshold, "first_message_exceeds_threshold", msgs)
+ print_pass("test_format_msgs_to_str_first_message_exceeds_threshold")
+
+
+def test_format_msgs_to_str_threshold_zero():
+ """Test with threshold of zero - latest message is still included."""
+ handler = create_handler()
+ threshold = 0
+ msgs = [create_user_msg("Test message")]
+
+ result = asyncio.run(handler.format_msgs_to_str(msgs, memory_compact_threshold=threshold))
+
+ # Latest message is included even with zero threshold (implementation behavior)
+ assert "Test message" in result, f"Expected message in result, got: {result}"
+ # Skip verify_result_within_threshold since latest message is always included
+ print_pass("test_format_msgs_to_str_threshold_zero")
+
+
+def test_format_msgs_to_str_threshold_exact_fit():
+ """Test when messages exactly fit the threshold."""
+ handler = create_handler()
+ # Create a message and measure its formatted string tokens
+ msg = create_user_msg("Test")
+ stat = asyncio.run(handler.stat_message(msg))
+ formatted_content = stat.format(include_thinking=False)
+ exact_threshold = asyncio.run(handler.count_str_token(formatted_content))
+
+ msgs = [msg]
+ result = asyncio.run(handler.format_msgs_to_str(msgs, memory_compact_threshold=exact_threshold))
+
+ assert "Test" in result, f"Message should fit exactly, got: {result}"
+ verify_result_within_threshold(handler, result, exact_threshold, "threshold_exact_fit", msgs)
+ print_pass("test_format_msgs_to_str_threshold_exact_fit")
+
+
+def test_format_msgs_to_str_threshold_one_less():
+ """Test when threshold is one less than needed - latest message is still included."""
+ handler = create_handler()
+ msg = create_user_msg("Test message")
+ stat = asyncio.run(handler.stat_message(msg))
+ formatted_content = stat.format(include_thinking=False)
+ threshold_minus_one = asyncio.run(handler.count_str_token(formatted_content)) - 1
+
+ msgs = [msg]
+ result = asyncio.run(handler.format_msgs_to_str(msgs, memory_compact_threshold=threshold_minus_one))
+
+ # Latest message is included even when it exceeds threshold (implementation behavior)
+ assert "Test message" in result, f"Expected message in result, got: {result}"
+ # Skip verify_result_within_threshold since latest message is always included
+ print_pass("test_format_msgs_to_str_threshold_one_less")
+
+
+def test_format_msgs_to_str_large_threshold():
+ """Test with very large threshold - all messages should be included."""
+ handler = create_handler()
+ threshold = 1000000
+ msgs = [create_user_msg("Message " + str(i) + " " + "x" * 100) for i in range(50)]
+
+ result = asyncio.run(handler.format_msgs_to_str(msgs, memory_compact_threshold=threshold))
+
+ # All messages should be included
+ for i in range(50):
+ assert f"Message {i}" in result, f"Message {i} should be included"
+ verify_result_within_threshold(handler, result, threshold, "large_threshold", msgs)
+ print_pass("test_format_msgs_to_str_large_threshold")
+
+
+# ==================== Edge Cases Tests ====================
+
+
+def test_format_msgs_to_str_special_characters():
+ """Test with special characters in content."""
+ handler = create_handler()
+ threshold = 4000
+ msgs = [create_user_msg("Test with 中文, 日本語, émojis 🎉 and symbols @#$%")]
+ result = asyncio.run(handler.format_msgs_to_str(msgs, memory_compact_threshold=threshold))
+
+ assert "中文" in result
+ assert "日本語" in result
+ assert "🎉" in result
+ verify_result_within_threshold(handler, result, threshold, "special_characters", msgs)
+ print_pass("test_format_msgs_to_str_special_characters")
+
+
+def test_format_msgs_to_str_empty_content():
+ """Test with empty content message."""
+ handler = create_handler()
+ threshold = 4000
+ msgs = [create_user_msg("")]
+ result = asyncio.run(handler.format_msgs_to_str(msgs, memory_compact_threshold=threshold))
+
+ assert "user:" in result, f"Expected role in result even with empty content, got: {result}"
+ verify_result_within_threshold(handler, result, threshold, "empty_content", msgs)
+ print_pass("test_format_msgs_to_str_empty_content")
+
+
+def test_format_msgs_to_str_whitespace_only():
+ """Test with whitespace-only content."""
+ handler = create_handler()
+ threshold = 4000
+ msgs = [create_user_msg(" \n\t ")]
+ result = asyncio.run(handler.format_msgs_to_str(msgs, memory_compact_threshold=threshold))
+
+ assert "user:" in result
+ verify_result_within_threshold(handler, result, threshold, "whitespace_only", msgs)
+ print_pass("test_format_msgs_to_str_whitespace_only")
+
+
+def test_format_msgs_to_str_newlines_in_content():
+ """Test with newlines in message content."""
+ handler = create_handler()
+ threshold = 4000
+ msgs = [create_user_msg("Line 1\nLine 2\nLine 3")]
+ result = asyncio.run(handler.format_msgs_to_str(msgs, memory_compact_threshold=threshold))
+
+ assert "Line 1" in result
+ assert "Line 2" in result
+ assert "Line 3" in result
+ verify_result_within_threshold(handler, result, threshold, "newlines_in_content", msgs)
+ print_pass("test_format_msgs_to_str_newlines_in_content")
+
+
+def test_format_msgs_to_str_very_long_single_word():
+ """Test with very long single word (no spaces)."""
+ handler = create_handler()
+ threshold = 10000
+ long_word = "a" * 5000
+ msgs = [create_user_msg(long_word)]
+ result = asyncio.run(handler.format_msgs_to_str(msgs, memory_compact_threshold=threshold))
+
+ # Should contain at least part of the word (may be truncated by formatter)
+ assert "aaa" in result, f"Expected long word content in result, got: {result[:100]}..."
+ verify_result_within_threshold(handler, result, threshold, "very_long_single_word", msgs)
+ print_pass("test_format_msgs_to_str_very_long_single_word")
+
+
+def test_format_msgs_to_str_mixed_content_blocks():
+ """Test message with mixed content blocks."""
+ handler = create_handler()
+ threshold = 4000
+ msgs = [
+ create_mixed_content_msg(
+ text="Text content",
+ thinking="Thinking content",
+ tool_name="test_tool",
+ tool_input={"key": "value"},
+ image_url="https://example.com/img.png",
+ ),
+ ]
+
+ result_no_thinking = asyncio.run(
+ handler.format_msgs_to_str(
+ msgs,
+ memory_compact_threshold=threshold,
+ include_thinking=False,
+ ),
+ )
+ result_with_thinking = asyncio.run(
+ handler.format_msgs_to_str(
+ msgs,
+ memory_compact_threshold=threshold,
+ include_thinking=True,
+ ),
+ )
+
+ assert "Text content" in result_no_thinking
+ assert "test_tool" in result_no_thinking
+ assert "" in result_no_thinking
+ assert "Thinking content" not in result_no_thinking
+ assert "Thinking content" in result_with_thinking
+ verify_result_within_threshold(handler, result_no_thinking, threshold, "mixed_content_no_thinking", msgs)
+ verify_result_within_threshold(handler, result_with_thinking, threshold, "mixed_content_with_thinking", msgs)
+ print_pass("test_format_msgs_to_str_mixed_content_blocks")
+
+
+def test_format_msgs_to_str_multiple_separators():
+ """Test that messages are separated by double newlines."""
+ handler = create_handler()
+ threshold = 4000
+ msgs = [
+ create_user_msg("Message 1"),
+ create_assistant_msg("Message 2"),
+ ]
+ result = asyncio.run(handler.format_msgs_to_str(msgs, memory_compact_threshold=threshold))
+
+ assert "\n\n" in result, f"Expected double newline separator, got: {result}"
+ verify_result_within_threshold(handler, result, threshold, "multiple_separators", msgs)
+ print_pass("test_format_msgs_to_str_multiple_separators")
+
+
+def test_format_msgs_to_str_tool_result_complex_output():
+ """Test tool_result with complex output (list of blocks)."""
+ handler = create_handler()
+ threshold = 4000
+ complex_output = [
+ {"type": "text", "text": "Operation completed"},
+ {"type": "image", "source": {"url": "https://example.com/result.png"}},
+ ]
+ msgs = [create_tool_result_msg("process_data", complex_output)]
+ result = asyncio.run(handler.format_msgs_to_str(msgs, memory_compact_threshold=threshold))
+
+ assert "process_data" in result
+ verify_result_within_threshold(handler, result, threshold, "tool_result_complex_output", msgs)
+ print_pass("test_format_msgs_to_str_tool_result_complex_output")
+
+
+def test_format_msgs_to_str_different_roles():
+ """Test with different roles (user, assistant, system, tool)."""
+ handler = create_handler()
+ threshold = 4000
+ msgs = [
+ Msg(name="system", role="system", content="System instruction"),
+ create_user_msg("User message"),
+ create_assistant_msg("Assistant response"),
+ create_tool_result_msg("tool", "Tool output"),
+ ]
+ result = asyncio.run(handler.format_msgs_to_str(msgs, memory_compact_threshold=threshold))
+
+ assert "system:" in result
+ assert "user:" in result
+ assert "assistant:" in result
+ verify_result_within_threshold(handler, result, threshold, "different_roles", msgs)
+ print_pass("test_format_msgs_to_str_different_roles")
+
+
+def test_format_msgs_to_str_incremental_threshold_check():
+ """Test incremental addition of messages until threshold is exceeded."""
+ handler = create_handler()
+
+ # Create messages with known approximate sizes
+ msgs = []
+ for i in range(10):
+ msgs.append(create_user_msg(f"Message {i} with some padding text"))
+
+ # Calculate total tokens
+ async def get_total_tokens():
+ total = 0
+ for msg in msgs:
+ stat = await handler.stat_message(msg)
+ total += stat.total_tokens
+ return total
+
+ total_tokens = asyncio.run(get_total_tokens())
+
+ # Use threshold that allows about half the messages
+ half_threshold = total_tokens // 2
+ result = asyncio.run(handler.format_msgs_to_str(msgs, memory_compact_threshold=half_threshold))
+
+ # Should have some but not all messages
+ included_count = sum(1 for i in range(10) if f"Message {i}" in result)
+ assert 0 < included_count < 10, f"Expected partial messages, got {included_count} messages included"
+ # Newer messages should be included (messages are processed from end)
+ assert "Message 9" in result, "Newest message should be included"
+ verify_result_within_threshold(handler, result, half_threshold, "incremental_threshold_check", msgs)
+ print_pass("test_format_msgs_to_str_incremental_threshold_check")
+
+
+def test_format_msgs_to_str_negative_threshold():
+ """Test with negative threshold value - latest message is still included."""
+ handler = create_handler()
+ threshold = -1
+ msgs = [create_user_msg("Test message")]
+
+ result = asyncio.run(handler.format_msgs_to_str(msgs, memory_compact_threshold=threshold))
+
+ # Latest message is included even with negative threshold (implementation behavior)
+ assert "Test message" in result, f"Expected message in result, got: {result}"
+ # Skip verify_result_within_threshold since latest message is always included
+ print_pass("test_format_msgs_to_str_negative_threshold")
+
+
+def test_format_msgs_to_str_preserves_newest_first():
+ """Test that newest messages are preserved when threshold is exceeded."""
+ handler = create_handler()
+ threshold = 300
+ msgs = [
+ create_user_msg("OLD MESSAGE " + "x" * 200),
+ create_assistant_msg("MIDDLE MESSAGE " + "y" * 200),
+ create_user_msg("NEW MESSAGE " + "z" * 200),
+ ]
+
+ # Use threshold that only allows ~1-2 messages
+ result = asyncio.run(handler.format_msgs_to_str(msgs, memory_compact_threshold=threshold))
+
+ # Newest message should be present
+ assert "NEW MESSAGE" in result, f"Expected newest message, got: {result}"
+ verify_result_within_threshold(handler, result, threshold, "preserves_newest_first", msgs)
+ print_pass("test_format_msgs_to_str_preserves_newest_first")
+
+
+def test_format_msgs_to_str_base64_image():
+ """Test with base64 encoded image."""
+ handler = create_handler()
+ threshold = 10000
+ msgs = [
+ Msg(
+ name="assistant",
+ role="assistant",
+ content=[
+ {
+ "type": "image",
+ "source": {
+ "type": "base64",
+ "data": "SGVsbG8gV29ybGQ=" * 100, # Simulated base64 data
+ },
+ },
+ ],
+ ),
+ ]
+ result = asyncio.run(handler.format_msgs_to_str(msgs, memory_compact_threshold=threshold))
+
+ assert "" in result
+ verify_result_within_threshold(handler, result, threshold, "base64_image", msgs)
+ print_pass("test_format_msgs_to_str_base64_image")
+
+
+def test_format_msgs_to_str_audio_video_blocks():
+ """Test with audio and video content blocks."""
+ handler = create_handler()
+ threshold = 4000
+ msgs = [
+ Msg(
+ name="assistant",
+ role="assistant",
+ content=[
+ {"type": "audio", "source": {"url": "https://example.com/audio.mp3"}},
+ {"type": "video", "source": {"url": "https://example.com/video.mp4"}},
+ ],
+ ),
+ ]
+ result = asyncio.run(handler.format_msgs_to_str(msgs, memory_compact_threshold=threshold))
+
+ assert "" in result
+ assert "" in result
+ verify_result_within_threshold(handler, result, threshold, "audio_video_blocks", msgs)
+ print_pass("test_format_msgs_to_str_audio_video_blocks")
+
+
+def test_format_msgs_to_str_unknown_block_type():
+ """Test that unknown block types are skipped gracefully."""
+ handler = create_handler()
+ threshold = 4000
+ msgs = [
+ Msg(
+ name="assistant",
+ role="assistant",
+ content=[
+ {"type": "unknown_type", "data": "some data"},
+ {"type": "text", "text": "Valid text"},
+ ],
+ ),
+ ]
+ result = asyncio.run(handler.format_msgs_to_str(msgs, memory_compact_threshold=threshold))
+
+ # Should still include valid content
+ assert "Valid text" in result
+ verify_result_within_threshold(handler, result, threshold, "unknown_block_type", msgs)
+ print_pass("test_format_msgs_to_str_unknown_block_type")
+
+
+def run_all_tests():
+ """Run all tests."""
+ tests = [
+ # Normal case tests
+ test_format_msgs_to_str_empty_list,
+ test_format_msgs_to_str_single_message,
+ test_format_msgs_to_str_multiple_messages,
+ test_format_msgs_to_str_message_order,
+ test_format_msgs_to_str_with_tool_use,
+ test_format_msgs_to_str_with_tool_result,
+ test_format_msgs_to_str_with_image,
+ test_format_msgs_to_str_conversation_flow,
+ # Thinking block tests
+ test_format_msgs_to_str_thinking_excluded_by_default,
+ test_format_msgs_to_str_thinking_included,
+ test_format_msgs_to_str_thinking_only_message,
+ # Token threshold tests
+ test_format_msgs_to_str_all_within_threshold,
+ test_format_msgs_to_str_exceeds_threshold_truncate_older,
+ test_format_msgs_to_str_single_message_exceeds_threshold,
+ test_format_msgs_to_str_first_message_exceeds_threshold,
+ test_format_msgs_to_str_threshold_zero,
+ test_format_msgs_to_str_threshold_exact_fit,
+ test_format_msgs_to_str_threshold_one_less,
+ test_format_msgs_to_str_large_threshold,
+ # Edge cases tests
+ test_format_msgs_to_str_special_characters,
+ test_format_msgs_to_str_empty_content,
+ test_format_msgs_to_str_whitespace_only,
+ test_format_msgs_to_str_newlines_in_content,
+ test_format_msgs_to_str_very_long_single_word,
+ test_format_msgs_to_str_mixed_content_blocks,
+ test_format_msgs_to_str_multiple_separators,
+ test_format_msgs_to_str_tool_result_complex_output,
+ test_format_msgs_to_str_different_roles,
+ test_format_msgs_to_str_incremental_threshold_check,
+ test_format_msgs_to_str_negative_threshold,
+ test_format_msgs_to_str_preserves_newest_first,
+ test_format_msgs_to_str_base64_image,
+ test_format_msgs_to_str_audio_video_blocks,
+ test_format_msgs_to_str_unknown_block_type,
+ ]
+
+ passed = 0
+ failed = 0
+
+ for test in tests:
+ try:
+ print_test_header(test.__name__)
+ test()
+ passed += 1
+ except AssertionError as e:
+ print_fail(test.__name__, str(e))
+ failed += 1
+ except Exception as e:
+ print_error(test.__name__, str(e))
+ failed += 1
+
+ # 打印最终统计结果
+ print(f"\n{Colors.CYAN}{'=' * 60}{Colors.RESET}")
+ print(f"{Colors.BOLD}Test Results Summary{Colors.RESET}")
+ print(f"{Colors.CYAN}{'=' * 60}{Colors.RESET}")
+ print(f"{Colors.GREEN}{Colors.BOLD}✓ Passed: {passed}{Colors.RESET}")
+ if failed > 0:
+ print(f"{Colors.RED}{Colors.BOLD}✗ Failed: {failed}{Colors.RESET}")
+ else:
+ print(f"{Colors.GREEN}✗ Failed: {failed}{Colors.RESET}")
+ print(f"{Colors.CYAN}{'=' * 60}{Colors.RESET}")
+
+ if failed == 0:
+ print(f"\n{Colors.GREEN}{Colors.BOLD}🎉 All tests passed!{Colors.RESET}")
+ else:
+ print(f"\n{Colors.RED}{Colors.BOLD}💥 Some tests failed!{Colors.RESET}")
+
+ return failed == 0
+
+
+if __name__ == "__main__":
+ success = run_all_tests()
+ sys.exit(0 if success else 1)
diff --git a/tests/light/test_memory_formatter.py b/tests/light/test_memory_formatter.py
deleted file mode 100644
index 00f45bb1..00000000
--- a/tests/light/test_memory_formatter.py
+++ /dev/null
@@ -1,489 +0,0 @@
-"""Tests for MemoryFormatter."""
-
-# pylint: disable=W0212
-
-import logging
-
-from agentscope.message import Msg
-
-from test_utils import get_token_counter
-from reme.memory.file_based import MemoryFormatter
-
-# 配置日志输出到控制台
-logging.basicConfig(
- level=logging.INFO,
- format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
-)
-logger = logging.getLogger(__name__)
-
-
-# ANSI 颜色码
-class Colors:
- """ANSI color codes for terminal output."""
-
- GREEN = "\033[92m"
- RED = "\033[91m"
- YELLOW = "\033[93m"
- BLUE = "\033[94m"
- CYAN = "\033[96m"
- BOLD = "\033[1m"
- RESET = "\033[0m"
-
-
-def print_pass(test_name: str):
- """打印测试通过信息"""
- print(f"{Colors.GREEN}{Colors.BOLD}✓ {test_name} PASSED{Colors.RESET}")
-
-
-def print_fail(test_name: str, error: str):
- """打印测试失败信息"""
- print(f"{Colors.RED}{Colors.BOLD}✗ {test_name} FAILED: {error}{Colors.RESET}")
-
-
-def print_error(test_name: str, error: str):
- """打印测试错误信息"""
- print(f"{Colors.YELLOW}{Colors.BOLD}⚠ {test_name} ERROR: {error}{Colors.RESET}")
-
-
-def print_test_header(test_name: str):
- """打印测试标题"""
- print(f"\n{Colors.CYAN}{'=' * 60}{Colors.RESET}")
- print(f"{Colors.BLUE}{Colors.BOLD}Running: {test_name}{Colors.RESET}")
- print(f"{Colors.CYAN}{'=' * 60}{Colors.RESET}")
-
-
-def create_user_msg(content: str) -> Msg:
- """Create a user message."""
- return Msg(name="user", role="user", content=content)
-
-
-def create_assistant_msg(content: str) -> Msg:
- """Create an assistant message."""
- return Msg(name="assistant", role="assistant", content=content)
-
-
-def create_tool_use_msg(tool_name: str, tool_input: dict) -> Msg:
- """Create a message with tool_use content block."""
- return Msg(
- name="assistant",
- role="assistant",
- content=[
- {
- "type": "tool_use",
- "id": "call_123",
- "name": tool_name,
- "input": tool_input,
- },
- ],
- )
-
-
-def create_tool_result_msg(tool_name: str, output: str | list[dict]) -> Msg:
- """Create a message with tool_result content block."""
- return Msg(
- name="tool",
- role="user",
- content=[
- {
- "type": "tool_result",
- "id": "call_123",
- "name": tool_name,
- "output": output,
- },
- ],
- )
-
-
-def create_thinking_msg(thinking_content: str) -> Msg:
- """Create a message with thinking content block."""
- return Msg(
- name="assistant",
- role="assistant",
- content=[
- {
- "type": "thinking",
- "text": thinking_content,
- },
- ],
- )
-
-
-def create_image_msg(url: str = "") -> Msg:
- """Create a message with image content block."""
- content = [
- {
- "type": "image",
- "source": {"url": url} if url else {},
- },
- ]
- return Msg(name="assistant", role="assistant", content=content)
-
-
-def create_formatter(memory_compact_threshold: int = 4000) -> MemoryFormatter:
- """Create a MemoryFormatter instance for testing."""
- return MemoryFormatter(
- token_counter=get_token_counter(),
- memory_compact_threshold=memory_compact_threshold,
- )
-
-
-# ==================== _format_tool_result_output Tests ====================
-
-
-def test_format_tool_result_output_string():
- """Test _format_tool_result_output with string input."""
- result = MemoryFormatter._format_tool_result_output("Hello, world!")
- assert result == "Hello, world!", f"Expected 'Hello, world!', got: {result}"
- print_pass("test_format_tool_result_output_string")
-
-
-def test_format_tool_result_output_text_block():
- """Test _format_tool_result_output with text block."""
- output = [{"type": "text", "text": "This is text content"}]
- result = MemoryFormatter._format_tool_result_output(output)
- assert result == "This is text content", f"Expected 'This is text content', got: {result}"
- print_pass("test_format_tool_result_output_text_block")
-
-
-def test_format_tool_result_output_image_block():
- """Test _format_tool_result_output with image block."""
- output = [{"type": "image", "source": {"url": "https://example.com/image.png"}}]
- result = MemoryFormatter._format_tool_result_output(output)
- assert "[image]" in result, f"Expected '[image]' in result, got: {result}"
- assert "https://example.com/image.png" in result, f"Expected URL in result, got: {result}"
- print_pass("test_format_tool_result_output_image_block")
-
-
-def test_format_tool_result_output_file_block():
- """Test _format_tool_result_output with file block."""
- output = [{"type": "file", "path": "/path/to/file.txt", "name": "file.txt"}]
- result = MemoryFormatter._format_tool_result_output(output)
- assert "[file]" in result, f"Expected '[file]' in result, got: {result}"
- assert "file.txt" in result, f"Expected 'file.txt' in result, got: {result}"
- print_pass("test_format_tool_result_output_file_block")
-
-
-def test_format_tool_result_output_multiple_blocks():
- """Test _format_tool_result_output with multiple blocks."""
- output = [
- {"type": "text", "text": "First part"},
- {"type": "text", "text": "Second part"},
- ]
- result = MemoryFormatter._format_tool_result_output(output)
- assert "First part" in result, f"Expected 'First part' in result, got: {result}"
- assert "Second part" in result, f"Expected 'Second part' in result, got: {result}"
- # Multiple parts should be joined with newlines and bullets
- assert "- " in result, f"Expected bullet format in result, got: {result}"
- print_pass("test_format_tool_result_output_multiple_blocks")
-
-
-def test_format_tool_result_output_empty_list():
- """Test _format_tool_result_output with empty list."""
- result = MemoryFormatter._format_tool_result_output([])
- assert result == "", f"Expected empty string, got: {result}"
- print_pass("test_format_tool_result_output_empty_list")
-
-
-def test_format_tool_result_output_invalid_block():
- """Test _format_tool_result_output with invalid block (missing type)."""
- output = [{"text": "No type key"}]
- result = MemoryFormatter._format_tool_result_output(output)
- assert result == "", f"Expected empty string for invalid block, got: {result}"
- print_pass("test_format_tool_result_output_invalid_block")
-
-
-def test_format_tool_result_output_unknown_type():
- """Test _format_tool_result_output with unknown block type."""
- output = [{"type": "unknown_type", "data": "some data"}]
- result = MemoryFormatter._format_tool_result_output(output)
- assert result == "", f"Expected empty string for unknown type, got: {result}"
- print_pass("test_format_tool_result_output_unknown_type")
-
-
-# ==================== format (single message) Tests ====================
-
-
-def test_format_empty_messages():
- """Test format with empty message list."""
- formatter = create_formatter()
- result = formatter.format([])
- assert result == "", f"Expected empty string, got: {result}"
- print_pass("test_format_empty_messages")
-
-
-def test_format_single_user_message():
- """Test format with a single user message."""
- formatter = create_formatter()
- msgs = [create_user_msg("Hello, how are you?")]
- result = formatter.format(msgs)
-
- assert "user:" in result, f"Expected 'user:' in result, got: {result}"
- assert "Hello, how are you?" in result, f"Expected content in result, got: {result}"
- print_pass("test_format_single_user_message")
-
-
-def test_format_single_assistant_message():
- """Test format with a single assistant message."""
- formatter = create_formatter()
- msgs = [create_assistant_msg("I am fine, thank you!")]
- result = formatter.format(msgs)
-
- assert "assistant:" in result, f"Expected 'assistant:' in result, got: {result}"
- assert "I am fine, thank you!" in result, f"Expected content in result, got: {result}"
- print_pass("test_format_single_assistant_message")
-
-
-def test_format_with_tool_use():
- """Test format with tool_use message."""
- formatter = create_formatter()
- msgs = [create_tool_use_msg("read_file", {"path": "/test.txt"})]
- result = formatter.format(msgs)
-
- assert "tool_call=read_file" in result, f"Expected 'tool_call=read_file' in result, got: {result}"
- assert "params=" in result, f"Expected 'params=' in result, got: {result}"
- print_pass("test_format_with_tool_use")
-
-
-def test_format_with_tool_result():
- """Test format with tool_result message."""
- formatter = create_formatter()
- msgs = [create_tool_result_msg("read_file", "file content here")]
- result = formatter.format(msgs)
-
- assert "tool_result=read_file" in result, f"Expected 'tool_result=read_file' in result, got: {result}"
- assert "output=" in result, f"Expected 'output=' in result, got: {result}"
- print_pass("test_format_with_tool_result")
-
-
-def test_format_with_thinking_block():
- """Test that thinking blocks are skipped."""
- formatter = create_formatter()
- msgs = [create_thinking_msg("Let me think about this...")]
- result = formatter.format(msgs)
-
- # Thinking content should NOT appear in the result
- assert "Let me think about this" not in result, f"Thinking content should be skipped, got: {result}"
- print_pass("test_format_with_thinking_block")
-
-
-def test_format_with_image():
- """Test format with image content block."""
- formatter = create_formatter()
- msgs = [create_image_msg("https://example.com/image.png")]
- result = formatter.format(msgs)
-
- assert "[image]" in result, f"Expected '[image]' in result, got: {result}"
- print_pass("test_format_with_image")
-
-
-# ==================== format (multiple messages) Tests ====================
-
-
-def test_format_conversation():
- """Test format with a conversation."""
- formatter = create_formatter()
- msgs = [
- create_user_msg("What is Python?"),
- create_assistant_msg("Python is a programming language."),
- create_user_msg("Tell me more."),
- create_assistant_msg("Python is known for its readability and simplicity."),
- ]
- result = formatter.format(msgs)
-
- assert "round0" in result, f"Expected 'round0' in result, got: {result}"
- assert "round1" in result, f"Expected 'round1' in result, got: {result}"
- assert "round2" in result, f"Expected 'round2' in result, got: {result}"
- assert "round3" in result, f"Expected 'round3' in result, got: {result}"
- print_pass("test_format_conversation")
-
-
-def test_format_without_index():
- """Test format without round index."""
- formatter = create_formatter()
- msgs = [
- create_user_msg("Hello"),
- create_assistant_msg("Hi there!"),
- ]
- result = formatter.format(msgs, add_index=False)
-
- assert "round" not in result, f"Expected no 'round' prefix, got: {result}"
- print_pass("test_format_without_index")
-
-
-def test_format_without_time():
- """Test format without timestamp."""
- formatter = create_formatter()
- msgs = [create_user_msg("Test message")]
- result = formatter.format(msgs, add_time=False)
-
- # The result should not have timestamp brackets at the beginning
- # Note: this test may need adjustment based on actual timestamp format
- assert "user:" in result, f"Expected 'user:' in result, got: {result}"
- print_pass("test_format_without_time")
-
-
-def test_format_with_tool_conversation():
- """Test format with tool use and result in conversation."""
- formatter = create_formatter()
- msgs = [
- create_user_msg("Read the file."),
- create_tool_use_msg("read_file", {"path": "/data.txt"}),
- create_tool_result_msg("read_file", "File content here"),
- create_assistant_msg("The file contains: File content here"),
- ]
- result = formatter.format(msgs)
-
- assert "user:" in result
- assert "tool_call=read_file" in result
- assert "tool_result=read_file" in result
- assert "assistant:" in result
- print_pass("test_format_with_tool_conversation")
-
-
-# ==================== Token Threshold Tests ====================
-
-
-def test_format_low_threshold():
- """Test that older messages are skipped with low threshold."""
- formatter = create_formatter(memory_compact_threshold=100)
- msgs = []
- for i in range(20):
- msgs.append(create_user_msg(f"Question {i}: " + "x" * 50))
- msgs.append(create_assistant_msg(f"Answer {i}: " + "y" * 50))
-
- result = formatter.format(msgs)
-
- # With low threshold, not all messages should be included
- # The newest messages should be present
- assert "round39" in result or "round38" in result, f"Expected recent round in result, got: {result}"
- # Older messages might be truncated
- logger.info(f"Result length: {len(result)}")
- print_pass("test_format_low_threshold")
-
-
-def test_format_high_threshold():
- """Test that all messages are included with high threshold."""
- formatter = create_formatter(memory_compact_threshold=100000)
- msgs = [
- create_user_msg("Message 1"),
- create_assistant_msg("Response 1"),
- create_user_msg("Message 2"),
- create_assistant_msg("Response 2"),
- ]
- result = formatter.format(msgs)
-
- # All messages should be included
- assert "round0" in result
- assert "round1" in result
- assert "round2" in result
- assert "round3" in result
- print_pass("test_format_high_threshold")
-
-
-# ==================== Edge Cases Tests ====================
-
-
-def test_format_long_text_truncation():
- """Test that long text is truncated."""
- formatter = create_formatter()
- long_text = "x" * 5000 # Much longer than default max length
- msgs = [create_user_msg(long_text)]
- result = formatter.format(msgs)
-
- # The result should be shorter due to truncation
- assert len(result) < len(long_text), f"Expected truncated result, got length: {len(result)}"
- print_pass("test_format_long_text_truncation")
-
-
-def test_format_special_characters():
- """Test format with special characters in content."""
- formatter = create_formatter()
- msgs = [create_user_msg("Test with 中文, 日本語, émojis 🎉")]
- result = formatter.format(msgs)
-
- assert "中文" in result, f"Expected Chinese characters in result, got: {result}"
- print_pass("test_format_special_characters")
-
-
-def test_format_tool_result_with_complex_output():
- """Test format with complex tool result output."""
- formatter = create_formatter()
- complex_output = [
- {"type": "text", "text": "Operation completed"},
- {"type": "image", "source": {"url": "https://example.com/result.png"}},
- ]
- msgs = [create_tool_result_msg("process_data", complex_output)]
- result = formatter.format(msgs)
-
- assert "tool_result=process_data" in result, f"Expected tool result in result, got: {result}"
- print_pass("test_format_tool_result_with_complex_output")
-
-
-def run_all_tests():
- """Run all tests."""
- tests = [
- # _format_tool_result_output tests
- test_format_tool_result_output_string,
- test_format_tool_result_output_text_block,
- test_format_tool_result_output_image_block,
- test_format_tool_result_output_file_block,
- test_format_tool_result_output_multiple_blocks,
- test_format_tool_result_output_empty_list,
- test_format_tool_result_output_invalid_block,
- test_format_tool_result_output_unknown_type,
- # format tests (single message)
- test_format_empty_messages,
- test_format_single_user_message,
- test_format_single_assistant_message,
- test_format_with_tool_use,
- test_format_with_tool_result,
- test_format_with_thinking_block,
- test_format_with_image,
- # format tests (multiple messages)
- test_format_conversation,
- test_format_without_index,
- test_format_without_time,
- test_format_with_tool_conversation,
- # threshold tests
- test_format_low_threshold,
- test_format_high_threshold,
- # edge cases
- test_format_long_text_truncation,
- test_format_special_characters,
- test_format_tool_result_with_complex_output,
- ]
-
- passed = 0
- failed = 0
-
- for test in tests:
- try:
- print_test_header(test.__name__)
- test()
- passed += 1
- except AssertionError as e:
- print_fail(test.__name__, str(e))
- failed += 1
- except Exception as e:
- print_error(test.__name__, str(e))
- failed += 1
-
- # 打印最终统计结果
- print(f"\n{Colors.CYAN}{'=' * 60}{Colors.RESET}")
- print(f"{Colors.BOLD}Test Results Summary{Colors.RESET}")
- print(f"{Colors.CYAN}{'=' * 60}{Colors.RESET}")
- print(f"{Colors.GREEN}{Colors.BOLD}✓ Passed: {passed}{Colors.RESET}")
- if failed > 0:
- print(f"{Colors.RED}{Colors.BOLD}✗ Failed: {failed}{Colors.RESET}")
- else:
- print(f"{Colors.GREEN}✗ Failed: {failed}{Colors.RESET}")
- print(f"{Colors.CYAN}{'=' * 60}{Colors.RESET}")
-
- if failed == 0:
- print(f"\n{Colors.GREEN}{Colors.BOLD}🎉 All tests passed!{Colors.RESET}")
- else:
- print(f"\n{Colors.RED}{Colors.BOLD}💥 Some tests failed!{Colors.RESET}")
-
-
-if __name__ == "__main__":
- run_all_tests()
diff --git a/tests/light/test_reme_light.py b/tests/light/test_reme_light.py
index 9e2ac70a..4fc7452a 100644
--- a/tests/light/test_reme_light.py
+++ b/tests/light/test_reme_light.py
@@ -1,197 +1,195 @@
-"""测试 ReMeLight"""
+"""测试 ReMeLight
+
+演示 ReMeLight 的完整功能，并使用 AsMsgHandler 跟踪每步 Token 变化：
+1. compact_tool_result - 压缩超长工具输出
+2. compact_memory - 生成压缩摘要
+3. summary_memory - 生成完整摘要并写入文件
+4. pre_reasoning_hook - 推理前预处理钩子
+5. memory_search - 语义搜索记忆
+6. ReMeInMemoryMemory.estimate_tokens - 估算 Token 使用
+7. ReMeInMemoryMemory.get_history_str - 获取格式化历史记录
+"""
import asyncio
-
-from agentscope.message import Msg
+import logging
+from test_utils import build_sample_messages, get_msg_handler
from reme.reme_light import ReMeLight
-# ==================== 消息创建辅助函数 ====================
-def create_user_msg(content: str) -> Msg:
- """创建用户消息"""
- return Msg(name="user", role="user", content=content)
-
-
-def create_assistant_msg(content: str) -> Msg:
- """创建助手消息"""
- return Msg(name="assistant", role="assistant", content=content)
-
-
-def create_tool_use_msg(tool_id: str, tool_name: str, tool_input: dict) -> Msg:
- """创建工具调用消息"""
- return Msg(
- name="assistant",
- role="assistant",
- content=[
- {
- "type": "tool_use",
- "id": tool_id,
- "name": tool_name,
- "input": tool_input,
- },
- ],
- )
-
-
-def create_tool_result_msg(tool_id: str, tool_name: str, output: str) -> Msg:
- """创建工具结果消息"""
- return Msg(
- name="tool",
- role="user",
- content=[
- {
- "type": "tool_result",
- "id": tool_id,
- "name": tool_name,
- "output": output,
- },
- ],
- )
-
-
-def create_thinking_msg(thinking_content: str) -> Msg:
- """创建思考消息"""
- return Msg(
- name="assistant",
- role="assistant",
- content=[
- {
- "type": "thinking",
- "text": thinking_content,
- },
- ],
- )
-
-
-# ==================== 构建模拟对话历史 ====================
-def build_sample_messages() -> list[Msg]:
- """构建一段包含多种消息类型的模拟对话"""
- messages = [
- # 用户询问 Python 版本
- create_user_msg("我想设置一个 Python 开发环境，你有什么建议？"),
- # 助手思考
- create_thinking_msg("用户想要搭建 Python 开发环境，我需要了解他的需求和偏好..."),
- # 助手回复
- create_assistant_msg(
- "好的！我建议使用 Python 3.11 或 3.12 版本，它们性能更好且功能丰富。"
- "你希望用于什么类型的开发？Web、数据科学还是其他？",
- ),
- # 用户提供更多信息
- create_user_msg("主要是做 Web 开发，使用 FastAPI 框架。另外我喜欢用 pyenv 管理版本。"),
- # 助手调用工具查询
- create_tool_use_msg(
- tool_id="call_001",
- tool_name="search_web",
- tool_input={"query": "FastAPI Python version compatibility 2024"},
- ),
- # 工具返回结果（模拟较长的输出）
- create_tool_result_msg(
- tool_id="call_001",
- tool_name="search_web",
- output=(
- "FastAPI 官方推荐使用 Python 3.8+ 版本，但 3.11/3.12 性能最佳。\n"
- "主要依赖：\n"
- "- Starlette: ASGI 框架\n"
- "- Pydantic v2: 数据验证\n"
- "- Uvicorn: ASGI 服务器\n"
- "最新版本 FastAPI 0.109+ 完全支持 Python 3.12。\n"
- "建议搭配 uv 或 pip-tools 进行依赖管理。"
- ),
- ),
- # 助手总结建议
- create_assistant_msg(
- "根据查询结果，我的建议是：\n"
- "1. **Python 版本**: 使用 Python 3.11 或 3.12（通过 pyenv 安装）\n"
- "2. **框架**: FastAPI 0.109+ 完全兼容这些版本\n"
- "3. **依赖管理**: 推荐使用 uv（更快）或 pip-tools\n"
- "4. **ASGI 服务器**: Uvicorn 配合 gunicorn 用于生产环境\n\n"
- "需要我帮你生成一个项目模板吗？",
- ),
- # 用户确认偏好
- create_user_msg("好的，我决定用 Python 3.12 + FastAPI + uv。请记住我的这些偏好。"),
- # 助手确认
- create_assistant_msg(
- "已记录你的开发偏好：\n"
- "- Python 版本: 3.12 (通过 pyenv 管理)\n"
- "- Web 框架: FastAPI\n"
- "- 包管理器: uv\n"
- "以后有相关问题我会参考这些偏好给你建议！",
- ),
- ]
- return messages
+def print_token_change(_step_name: str, before: int, after: int):
+ """打印 Token 变化统计。"""
+ change = after - before
+ change_pct = (change / before * 100) if before > 0 else 0
+ print(f" 📊 Token 统计: {before:,} → {after:,} (变化: {change:+,}, {change_pct:+.1f}%)")
# ==================== 主测试流程 ====================
async def main():
- """ReMeLight 主测试流程，演示完整的记忆管理功能。"""
+ """测试 ReMeLight 的完整功能，并跟踪每步 Token 变化。"""
+ # 初始化 AsMsgHandler 用于 Token 统计
+ msg_handler = get_msg_handler()
+
# 初始化 ReMeLight
reme = ReMeLight(
- working_dir=".reme", # 记忆文件存储目录
- max_input_length=128000, # 模型上下文窗口（tokens）
- memory_compact_ratio=0.7, # 达到 max_input_length * 0.7 时触发压缩
- language="zh", # 摘要语言（zh / ""）
- tool_result_threshold=1000, # 超过此字符数的工具输出自动转存
- retention_days=7, # tool_result/ 文件保留天数
+ default_as_llm_config={"model_name": "qwen3.5-35b-a3b"},
+ # default_embedding_model_config={"model_name": "text-embedding-v4"},
+ default_file_store_config={"fts_enabled": True, "vector_enabled": False},
+ enable_load_env=True,
)
+ logging.getLogger("reme").setLevel(logging.WARNING)
await reme.start()
- print("=" * 60)
+ print("=" * 70)
print("ReMeLight 已启动")
- print("=" * 60)
+ print("=" * 70)
- # 构建模拟对话历史
- messages = build_sample_messages()
- print(f"\n[原始消息数量]: {len(messages)} 条")
+ # 构建模拟对话历史（包含超长 tool_result，确保超过 128K token）
+ original_messages = build_sample_messages(include_large_tool_result=True)
+ initial_tokens = await msg_handler.count_msgs_token(original_messages)
- # 1. 压缩超长工具输出（防止工具结果撑爆上下文）
- print("\n" + "-" * 40)
- print("[步骤 1] 压缩超长工具输出...")
- messages = await reme.compact_tool_result(messages)
- print(f"处理后消息数量: {len(messages)} 条")
+ print(f"\n[原始消息]: {len(original_messages)} 条, {initial_tokens:,} tokens")
+ print(f" 目标阈值: 128K = {128 * 1024:,} tokens")
+ print(f" 超出阈值: {initial_tokens > 128 * 1024}")
- # 2. 将历史对话压缩为结构化摘要（触发时机：上下文接近上限）
- print("\n" + "-" * 40)
- print("[步骤 2] 生成结构化压缩摘要...")
- summary = await reme.compact_memory(
+ # ==================== 1. compact_tool_result ====================
+ print("\n" + "=" * 70)
+ print("[步骤 1] compact_tool_result - 压缩超长工具输出")
+ print("=" * 70)
+
+ # 重新获取原始消息
+ messages = build_sample_messages(include_large_tool_result=True)
+ tokens_before = await msg_handler.count_msgs_token(messages)
+ messages_after_step1 = await reme.compact_tool_result(messages)
+ tokens_after = await msg_handler.count_msgs_token(messages_after_step1)
+
+ print(f" 消息数量: {len(messages)} → {len(messages_after_step1)}")
+ print_token_change("compact_tool_result", tokens_before, tokens_after)
+
+ # ==================== 2. compact_memory ====================
+ print("\n" + "=" * 70)
+ print("[步骤 2] compact_memory - 生成结构化压缩摘要")
+ print("=" * 70)
+
+ # 重新获取原始消息
+ messages = build_sample_messages(include_large_tool_result=True)
+ tokens_before = await msg_handler.count_msgs_token(messages)
+ compact_summary = await reme.compact_memory(
messages=messages,
- previous_summary="", # 可传入上轮摘要，实现增量更新
+ previous_summary="",
)
- print(f"压缩摘要:\n{summary[:500]}..." if len(summary) > 500 else f"压缩摘要:\n{summary}")
+ summary_tokens = await msg_handler.count_str_token(compact_summary)
- # 3. 后台异步提交摘要任务（不阻塞对话，摘要写入 memory/YYYY-MM-DD.md）
- print("\n" + "-" * 40)
- print("[步骤 3] 提交后台异步摘要任务...")
- reme.add_async_summary_task(messages=messages)
- print("异步任务已提交")
+ print(f" 输入消息 tokens: {tokens_before:,}")
+ print(f" 压缩摘要长度: {len(compact_summary)} 字符, {summary_tokens:,} tokens")
+ print(f" 压缩比: {summary_tokens / tokens_before * 100:.1f}%" if tokens_before > 0 else " 压缩比: N/A")
+ print(f" 摘要预览: {compact_summary[:200]}..." if len(compact_summary) > 200 else f" 摘要: {compact_summary}")
- # 4. 语义搜索记忆（向量 + BM25 混合检索）
- print("\n" + "-" * 40)
- print("[步骤 4] 语义搜索记忆...")
- result = await reme.memory_search(query="Python 版本偏好", max_results=5)
- print(f"搜索结果: {result}")
+ # ==================== 3. summary_memory ====================
+ print("\n" + "=" * 70)
+ print("[步骤 3] summary_memory - 生成完整摘要并写入文件")
+ print("=" * 70)
- # 5. 获取会话内存实例（ReMeInMemoryMemory，管理单次对话的上下文）
- print("\n" + "-" * 40)
- print("[步骤 5] 获取会话内存实例并估算 Token 使用...")
+ # 重新获取原始消息
+ messages = build_sample_messages(include_large_tool_result=True)
+ tokens_before = await msg_handler.count_msgs_token(messages)
+ summary_result = await reme.summary_memory(messages=messages)
+
+ print(f" 输入消息 tokens: {tokens_before:,}")
+ print(f" 摘要结果长度: {len(summary_result)} 字符")
+ print(f" 摘要预览: {summary_result[:200]}..." if len(summary_result) > 200 else f" 摘要: {summary_result}")
+
+ # ==================== 4. pre_reasoning_hook ====================
+ print("\n" + "=" * 70)
+ print("[步骤 4] pre_reasoning_hook - 推理前预处理")
+ print("=" * 70)
+
+ # 重新获取原始消息
+ messages = build_sample_messages(include_large_tool_result=True)
+ tokens_before = await msg_handler.count_msgs_token(messages)
+ processed_messages, compressed_summary = await reme.pre_reasoning_hook(
+ messages=messages,
+ system_prompt="你是一个有帮助的 AI 助手。",
+ compressed_summary="",
+ max_input_length=128000,
+ compact_ratio=0.7,
+ memory_compact_reserve=10000,
+ enable_tool_result_compact=True,
+ tool_result_compact_keep_n=3,
+ )
+ tokens_after = await msg_handler.count_msgs_token(processed_messages)
+ compressed_summary_tokens = await msg_handler.count_str_token(compressed_summary)
+
+ print(f" 消息数量: {len(messages)} → {len(processed_messages)}")
+ print_token_change("pre_reasoning_hook", tokens_before, tokens_after)
+ print(f" 压缩摘要: {len(compressed_summary)} 字符, {compressed_summary_tokens:,} tokens")
+ print(f" 总上下文: {tokens_after + compressed_summary_tokens:,} tokens")
+
+ # ==================== 5. memory_search ====================
+ print("\n" + "=" * 70)
+ print("[步骤 5] memory_search - 语义搜索记忆")
+ print("=" * 70)
+
+ search_result = await reme.memory_search(query="Python 版本偏好", max_results=5)
+ if search_result.content:
+ print(f" 搜索结果: {search_result.content}")
+ else:
+ print(" 未找到相关记忆")
+
+ # ==================== 6 & 7. ReMeInMemoryMemory ====================
+ print("\n" + "=" * 70)
+ print("[步骤 6] ReMeInMemoryMemory - 会话内存管理")
+ print("=" * 70)
+
+ # 重新获取原始消息
+ messages = build_sample_messages(include_large_tool_result=True)
memory = reme.get_in_memory_memory()
- # 将消息添加到内存中以便估算
for msg in messages:
await memory.add(msg)
- token_stats = await memory.estimate_tokens()
- print(f"当前上下文使用率: {token_stats['context_usage_ratio']:.1f}%")
- print(f"消息 Token 数: {token_stats['messages_tokens']}")
- print(f"预估总 Token 数: {token_stats['estimated_tokens']}")
+ print(f" 已添加 {len(messages)} 条原始消息到内存")
- # 6. 关闭前等待后台任务完成
- print("\n" + "-" * 40)
- print("[步骤 6] 等待后台任务完成...")
- summary_result = await reme.await_summary_tasks()
- print(f"后台摘要任务完成，结果长度: {len(summary_result)} 字符")
+ # 6.1 estimate_tokens
+ print("\n[6.1] estimate_tokens - 估算 Token 使用:")
+ token_stats = await memory.estimate_tokens(max_input_length=128000)
+ print(f" - 总消息数: {token_stats['total_messages']}")
+ print(f" - 消息 Token 数: {token_stats['messages_tokens']:,}")
+ print(f" - 压缩摘要 Token 数: {token_stats['compressed_summary_tokens']:,}")
+ print(f" - 预估总 Token 数: {token_stats['estimated_tokens']:,}")
+ print(f" - 最大输入长度: {token_stats['max_input_length']:,}")
+ print(f" - 上下文使用率: {token_stats['context_usage_ratio']:.2f}%")
+
+ # 6.2 get_history_str
+ print("\n[6.2] get_history_str - 格式化历史记录:")
+ history_str = await memory.get_history_str(max_input_length=128000)
+ print(history_str[:1000] + "..." if len(history_str) > 1000 else history_str)
+
+ # ==================== 等待后台任务完成 ====================
+ print("\n" + "=" * 70)
+ print("[步骤 7] 等待后台任务完成")
+ print("=" * 70)
+ await_result = await reme.await_summary_tasks()
+ print(f" 后台任务完成，结果长度: {len(await_result)} 字符")
+
+ # ==================== 总结 ====================
+ print("\n" + "=" * 70)
+ print("📊 Token 变化总结")
+ print("=" * 70)
+ print(f" 原始消息: {initial_tokens:,} tokens")
+ print(f" Step 1 compact_tool_result 后: {(await msg_handler.count_msgs_token(messages_after_step1)):,} tokens")
+ print(f" Step 2 compact_memory 摘要: {summary_tokens:,} tokens")
+ print(
+ f" Step 4 pre_reasoning_hook 后: {tokens_after:,} tokens + 摘要 {compressed_summary_tokens:,} "
+ f"tokens = {tokens_after + compressed_summary_tokens:,} tokens",
+ )
+ print(
+ f" 最大节省: {initial_tokens - tokens_after:,} "
+ f"tokens ({(initial_tokens - tokens_after) / initial_tokens * 100:.1f}%)",
+ )
+ print(f" 目标阈值: {128 * 1024:,} tokens")
# 关闭 ReMeLight
await reme.close()
- print("\n" + "=" * 60)
+ print("\n" + "=" * 70)
print("ReMeLight 已关闭")
- print("=" * 60)
+ print("=" * 70)
if __name__ == "__main__":
diff --git a/tests/light/test_reme_light_log.txt b/tests/light/test_reme_light_log.txt
new file mode 100644
index 00000000..5a0c71a5
--- /dev/null
+++ b/tests/light/test_reme_light_log.txt
@@ -0,0 +1,84 @@
+======================================================================
+ReMeLight 已启动
+======================================================================
+
+[原始消息]: 18 条, 223,838 tokens
+ 目标阈值: 128K = 131,072 tokens
+ 超出阈值: True
+
+======================================================================
+[步骤 1] compact_tool_result - 压缩超长工具输出
+======================================================================
+ 消息数量: 18 → 18
+ 📊 Token 统计: 223,838 → 1,107 (变化: -222,731, -99.5%)
+
+======================================================================
+[步骤 2] compact_memory - 生成结构化压缩摘要
+======================================================================
+ 输入消息 tokens: 223,838
+ 压缩摘要长度: 1032 字符, 493 tokens
+ 压缩比: 0.2%
+
+======================================================================
+[步骤 3] summary_memory - 生成完整摘要并写入文件
+======================================================================
+
+reme_summarizer: [SILENT]
+ 输入消息 tokens: 223,838
+ 摘要结果长度: 8 字符
+ 摘要: [SILENT]
+
+======================================================================
+[步骤 4] pre_reasoning_hook - 推理前预处理
+======================================================================
+ 消息数量: 18 → 18
+ 📊 Token 统计: 223,838 → 1,105 (变化: -222,733, -99.5%)
+ 压缩摘要: 0 字符, 0 tokens
+ 总上下文: 1,105 tokens
+
+======================================================================
+[步骤 5] memory_search - 语义搜索记忆
+======================================================================
+ 搜索结果: [{'type': 'text', 'text': '[\n {\n "path": "/Users...
+
+======================================================================
+[步骤 6] ReMeInMemoryMemory - 会话内存管理
+======================================================================
+ 已添加 18 条原始消息到内存
+
+[6.1] estimate_tokens - 估算 Token 使用:
+ - 总消息数: 18
+ - 消息 Token 数: 223,838
+ - 压缩摘要 Token 数: 0
+ - 预估总 Token 数: 223,838
+ - 最大输入长度: 128,000
+ - 上下文使用率: 174.87%
+
+[6.2] get_history_str - 格式化历史记录:
+**Conversation History**
+
+- Total messages: 18
+- Estimated tokens: 223838
+- Max input length: 128000
+- Context usage: 174.9%
+- Compressed summary tokens: 0
+...
+
+======================================================================
+[步骤 7] 等待后台任务完成
+======================================================================
+ 后台任务完成，结果长度: 0 字符
+
+======================================================================
+📊 Token 变化总结
+======================================================================
+ 原始消息: 223,838 tokens
+ Step 1 compact_tool_result 后: 1,107 tokens
+ Step 2 compact_memory 摘要: 493 tokens
+ Step 4 pre_reasoning_hook 后: 1,105 tokens + 摘要 0 tokens = 1,105 tokens
+ 最大节省: 222,733 tokens (99.5%)
+ 目标阈值: 131,072 tokens
+
+======================================================================
+ReMeLight 已关闭
+======================================================================
diff --git a/tests/light/test_summarizer.py b/tests/light/test_summarizer.py
index 560efabd..5f8a67a1 100644
--- a/tests/light/test_summarizer.py
+++ b/tests/light/test_summarizer.py
@@ -2,25 +2,22 @@
import asyncio
import datetime
-import logging
import tempfile
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.memory.file_based import Summarizer
+from reme.core.utils import get_logger
+from reme.memory.file_based.components import Summarizer
+from reme.memory.file_based.tools import FileIO
-# 配置日志输出到控制台
-logging.basicConfig(
- level=logging.INFO,
- format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
-)
-logger = logging.getLogger(__name__)
+
+logger = get_logger()
# ANSI 颜色码
@@ -100,6 +97,16 @@ def create_tool_result_msg(tool_name: str, output: str) -> Msg:
)
+def create_toolkit(working_dir: str) -> Toolkit:
+ """Create a default Toolkit with FileIO tools for testing."""
+ toolkit = Toolkit()
+ file_io = FileIO(working_dir=working_dir)
+ toolkit.register_tool_function(file_io.read_file)
+ toolkit.register_tool_function(file_io.write_file)
+ toolkit.register_tool_function(file_io.edit_file)
+ return toolkit
+
+
def create_summarizer(working_dir: str = None, memory_dir: str = "memory"):
"""Create a Summarizer instance for testing."""
if working_dir is None:
@@ -114,9 +121,10 @@ def create_summarizer(working_dir: str = None, memory_dir: str = "memory"):
working_dir=working_dir,
memory_dir=memory_dir,
memory_compact_threshold=4000,
- chat_model=get_dash_chat_model(),
- formatter=get_formatter(),
- token_counter=get_token_counter(),
+ as_token_counter=get_token_counter(),
+ toolkit=create_toolkit(working_dir),
+ as_llm=get_dash_chat_model(),
+ as_llm_formatter=get_formatter(),
),
working_dir,
)
@@ -195,9 +203,10 @@ def test_consecutive_summaries():
working_dir=working_dir,
memory_dir=memory_dir,
memory_compact_threshold=4000,
- chat_model=get_dash_chat_model(),
- formatter=get_formatter(),
- token_counter=get_token_counter(),
+ as_token_counter=get_token_counter(),
+ toolkit=create_toolkit(working_dir),
+ as_llm=get_dash_chat_model(),
+ as_llm_formatter=get_formatter(),
)
# 第一轮对话
diff --git a/tests/light/test_tool_result_compactor.py b/tests/light/test_tool_result_compactor.py
index 059e626e..b05dad29 100644
--- a/tests/light/test_tool_result_compactor.py
+++ b/tests/light/test_tool_result_compactor.py
@@ -7,8 +7,8 @@ from pathlib import Path
from agentscope.message import Msg
-from reme.memory.file_based.tool_result_compactor import ToolResultCompactor
-from reme.memory.file_based.utils import TRUNCATION_MARKER_START
+from reme.memory.file_based.components import ToolResultCompactor
+from reme.memory.file_based.utils import TRUNCATION_NOTICE_MARKER
def create_tool_result_msg(output: str | list, tool_name: str = "test_tool") -> Msg:
@@ -33,7 +33,7 @@ class TestToolResultCompactor:
def test_no_truncation_when_under_threshold(self):
"""Test that short content is not truncated."""
with tempfile.TemporaryDirectory() as tmpdir:
- op = ToolResultCompactor(tool_result_dir=tmpdir, tool_result_threshold=1000)
+ op = ToolResultCompactor(tool_result_dir=tmpdir, recent_max_bytes=1000)
messages = [create_tool_result_msg("short content")]
result = asyncio.run(op.call(messages=messages))
@@ -45,14 +45,14 @@ class TestToolResultCompactor:
def test_truncation_when_over_threshold(self):
"""Test that long content is truncated and saved to file."""
with tempfile.TemporaryDirectory() as tmpdir:
- op = ToolResultCompactor(tool_result_dir=tmpdir, tool_result_threshold=100)
+ op = ToolResultCompactor(tool_result_dir=tmpdir, recent_max_bytes=100)
long_content = "x" * 500
messages = [create_tool_result_msg(long_content)]
_ = asyncio.run(op.call(messages=messages))
output = messages[0].content[0]["output"]
- assert TRUNCATION_MARKER_START in output
+ assert TRUNCATION_NOTICE_MARKER in output
assert "[Full content saved to:" in output
# Verify file was created
@@ -68,8 +68,8 @@ class TestToolResultCompactor:
def test_skip_already_truncated(self):
"""Test that already truncated content is not re-truncated."""
with tempfile.TemporaryDirectory() as tmpdir:
- op = ToolResultCompactor(tool_result_dir=tmpdir, tool_result_threshold=100)
- truncated_content = f"head{TRUNCATION_MARKER_START}(100 chars omitted)<<>>tail"
+ op = ToolResultCompactor(tool_result_dir=tmpdir, recent_max_bytes=100)
+ truncated_content = "head<<>>(100 chars omitted)<<>>tail"
messages = [create_tool_result_msg(truncated_content)]
asyncio.run(op.call(messages=messages))
@@ -80,20 +80,20 @@ class TestToolResultCompactor:
def test_truncation_list_output(self):
"""Test truncation of list output with text blocks."""
with tempfile.TemporaryDirectory() as tmpdir:
- op = ToolResultCompactor(tool_result_dir=tmpdir, tool_result_threshold=100)
+ op = ToolResultCompactor(tool_result_dir=tmpdir, recent_max_bytes=100)
list_output = [{"type": "text", "text": "y" * 500}]
messages = [create_tool_result_msg(list_output)]
asyncio.run(op.call(messages=messages))
text_block = messages[0].content[0]["output"][0]
- assert TRUNCATION_MARKER_START in text_block["text"]
+ assert TRUNCATION_NOTICE_MARKER in text_block["text"]
assert len(list(Path(tmpdir).glob("*.txt"))) == 1
def test_list_output_no_truncation_when_short(self):
"""Test that short list output is not truncated."""
with tempfile.TemporaryDirectory() as tmpdir:
- op = ToolResultCompactor(tool_result_dir=tmpdir, tool_result_threshold=1000)
+ op = ToolResultCompactor(tool_result_dir=tmpdir, recent_max_bytes=1000)
list_output = [{"type": "text", "text": "short"}]
messages = [create_tool_result_msg(list_output)]
@@ -105,7 +105,7 @@ class TestToolResultCompactor:
def test_list_output_multiple_text_blocks(self):
"""Test truncation of multiple text blocks in list output."""
with tempfile.TemporaryDirectory() as tmpdir:
- op = ToolResultCompactor(tool_result_dir=tmpdir, tool_result_threshold=100)
+ op = ToolResultCompactor(tool_result_dir=tmpdir, recent_max_bytes=100)
list_output = [
{"type": "text", "text": "a" * 500},
{"type": "text", "text": "short"},
@@ -116,15 +116,15 @@ class TestToolResultCompactor:
asyncio.run(op.call(messages=messages))
output = messages[0].content[0]["output"]
- assert TRUNCATION_MARKER_START in output[0]["text"]
+ assert TRUNCATION_NOTICE_MARKER in output[0]["text"]
assert output[1]["text"] == "short" # unchanged
- assert TRUNCATION_MARKER_START in output[2]["text"]
+ assert TRUNCATION_NOTICE_MARKER in output[2]["text"]
assert len(list(Path(tmpdir).glob("*.txt"))) == 2
def test_list_output_mixed_block_types(self):
"""Test that non-text blocks in list output are unchanged."""
with tempfile.TemporaryDirectory() as tmpdir:
- op = ToolResultCompactor(tool_result_dir=tmpdir, tool_result_threshold=100)
+ op = ToolResultCompactor(tool_result_dir=tmpdir, recent_max_bytes=100)
list_output = [
{"type": "text", "text": "c" * 500},
{"type": "image", "source": {"type": "url", "url": "http://example.com/img.png"}},
@@ -134,14 +134,14 @@ class TestToolResultCompactor:
asyncio.run(op.call(messages=messages))
output = messages[0].content[0]["output"]
- assert TRUNCATION_MARKER_START in output[0]["text"]
+ assert TRUNCATION_NOTICE_MARKER in output[0]["text"]
assert output[1] == {"type": "image", "source": {"type": "url", "url": "http://example.com/img.png"}}
assert len(list(Path(tmpdir).glob("*.txt"))) == 1
def test_cleanup_expired_files(self):
"""Test cleanup of expired files."""
with tempfile.TemporaryDirectory() as tmpdir:
- op = ToolResultCompactor(tool_result_dir=tmpdir, tool_result_threshold=100, retention_days=1)
+ op = ToolResultCompactor(tool_result_dir=tmpdir, recent_max_bytes=100, retention_days=1)
# Create an old file
old_time = (datetime.now() - timedelta(days=2)).isoformat()
@@ -162,7 +162,7 @@ class TestToolResultCompactor:
def test_string_content_msg_unchanged(self):
"""Test that messages with string content are unchanged."""
with tempfile.TemporaryDirectory() as tmpdir:
- op = ToolResultCompactor(tool_result_dir=tmpdir, tool_result_threshold=100)
+ op = ToolResultCompactor(tool_result_dir=tmpdir, recent_max_bytes=100)
messages = [Msg(name="user", role="user", content="hello world")]
asyncio.run(op.call(messages=messages))
diff --git a/tests/light/test_tools.py b/tests/light/test_tools.py
new file mode 100644
index 00000000..779c59a2
--- /dev/null
+++ b/tests/light/test_tools.py
@@ -0,0 +1,442 @@
+# -*- coding: utf-8 -*-
+# pylint: disable=redefined-outer-name
+"""Unit tests for Shell and FileIO tools."""
+
+import asyncio
+import os
+import re
+import shutil
+import tempfile
+
+import pytest
+
+from reme.memory.file_based.tools.file_io import FileIO
+from reme.memory.file_based.tools.shell import Shell
+from reme.memory.file_based.utils import DEFAULT_MAX_BYTES
+
+
+# ============ Shell Tests ============
+
+
+@pytest.fixture(scope="module")
+def shell_env():
+ """Create temporary directory and Shell instance."""
+ test_dir = tempfile.mkdtemp(prefix="test_shell_")
+ shell = Shell(working_dir=test_dir)
+ yield {"dir": test_dir, "shell": shell}
+ shutil.rmtree(test_dir, ignore_errors=True)
+
+
+def test_shell_echo_success(shell_env):
+ """Test successful echo command execution."""
+ result = asyncio.run(shell_env["shell"].execute_shell_command("echo hello"))
+ assert result.content
+ text = result.content[0].get("text", "")
+ assert "hello" in text
+
+
+def test_shell_pwd_in_working_dir(shell_env):
+ """Test command executes in correct working directory."""
+ result = asyncio.run(shell_env["shell"].execute_shell_command("pwd"))
+ text = result.content[0].get("text", "")
+ assert shell_env["dir"] in text
+
+
+def test_shell_command_failure(shell_env):
+ """Test failed command returns error information."""
+ result = asyncio.run(shell_env["shell"].execute_shell_command("exit 1"))
+ text = result.content[0].get("text", "")
+ assert "failed" in text.lower()
+ assert "exit code" in text.lower()
+
+
+def test_shell_command_with_stderr(shell_env):
+ """Test command with stderr output."""
+ result = asyncio.run(
+ shell_env["shell"].execute_shell_command("echo error >&2 && exit 1"),
+ )
+ text = result.content[0].get("text", "")
+ assert "error" in text
+
+
+def test_shell_no_output(shell_env):
+ """Test successful command with no output."""
+ result = asyncio.run(shell_env["shell"].execute_shell_command("true"))
+ text = result.content[0].get("text", "")
+ assert "successfully" in text.lower()
+
+
+def test_shell_multiline_output(shell_env):
+ """Test command with multiline output."""
+ result = asyncio.run(
+ shell_env["shell"].execute_shell_command("echo -e 'line1\nline2\nline3'"),
+ )
+ text = result.content[0].get("text", "")
+ assert "line1" in text
+ assert "line2" in text
+ assert "line3" in text
+
+
+def test_shell_timeout(shell_env):
+ """Test command timeout handling."""
+ result = asyncio.run(
+ shell_env["shell"].execute_shell_command("sleep 10", timeout=1),
+ )
+ text = result.content[0].get("text", "")
+ assert "timeout" in text.lower()
+
+
+# ============ FileIO Read Tests ============
+
+
+@pytest.fixture(scope="module")
+def fileio_env():
+ """Create temporary directory with test files."""
+ test_dir = tempfile.mkdtemp(prefix="test_fileio_")
+ file_io = FileIO(working_dir=test_dir)
+
+ # Create simple test file
+ simple_file = os.path.join(test_dir, "simple.txt")
+ with open(simple_file, "w", encoding="utf-8") as f:
+ f.write("line1\nline2\nline3\nline4\nline5")
+
+ # Create large file (exceeds DEFAULT_MAX_BYTES)
+ large_file = os.path.join(test_dir, "large.txt")
+ with open(large_file, "w", encoding="utf-8") as f:
+ # Each line is ~7-10 bytes ("line N\n"); generate enough to exceed limit.
+ # Line 1 is literally "line 1" so the head-kept assertion can match it.
+ line_count = (DEFAULT_MAX_BYTES // 7) + 1000
+ for i in range(1, line_count + 1):
+ f.write(f"line {i}\n")
+
+ # Create large bytes file (exceeds DEFAULT_MAX_BYTES)
+ # Lines are 101 bytes each; at DEFAULT_MAX_BYTES the cut lands mid-line → else branch
+ large_bytes_file = os.path.join(test_dir, "large_bytes.txt")
+ with open(large_bytes_file, "w", encoding="utf-8") as f:
+ content = "x" * 100 + "\n"
+ lines_needed = (DEFAULT_MAX_BYTES // 101) + 100
+ for _ in range(lines_needed):
+ f.write(content)
+
+ # Single line larger than DEFAULT_MAX_BYTES → newline_count==0 branch in truncate
+ huge_line_file = os.path.join(test_dir, "huge_line.txt")
+ with open(huge_line_file, "w", encoding="utf-8") as f:
+ f.write("A" * (DEFAULT_MAX_BYTES + 1000) + "\nline2\n")
+
+ # Empty file
+ empty_file = os.path.join(test_dir, "empty.txt")
+ with open(empty_file, "w", encoding="utf-8") as f:
+ f.write("")
+
+ yield {
+ "dir": test_dir,
+ "file_io": file_io,
+ "simple_file": simple_file,
+ "large_file": large_file,
+ "large_bytes_file": large_bytes_file,
+ "huge_line_file": huge_line_file,
+ "empty_file": empty_file,
+ }
+ shutil.rmtree(test_dir, ignore_errors=True)
+
+
+def test_read_file_success(fileio_env):
+ """Test successful file reading."""
+ result = asyncio.run(fileio_env["file_io"].read_file(fileio_env["simple_file"]))
+ text = result.content[0].get("text", "")
+ assert "line1" in text
+ assert "line5" in text
+
+
+def test_read_file_relative_path(fileio_env):
+ """Test reading file with relative path."""
+ result = asyncio.run(fileio_env["file_io"].read_file("simple.txt"))
+ text = result.content[0].get("text", "")
+ assert "line1" in text
+
+
+def test_read_file_not_exists(fileio_env):
+ """Test reading non-existent file."""
+ result = asyncio.run(fileio_env["file_io"].read_file("nonexistent.txt"))
+ text = result.content[0].get("text", "")
+ assert "Error" in text
+ assert "does not exist" in text
+
+
+def test_read_file_with_line_range(fileio_env):
+ """Test reading specific line range."""
+ result = asyncio.run(
+ fileio_env["file_io"].read_file(fileio_env["simple_file"], start_line=2, end_line=4),
+ )
+ text = result.content[0].get("text", "")
+ assert "line2" in text
+ assert "line4" in text
+ assert "lines 2-4" in text.lower()
+
+
+def test_read_file_start_line_exceeds(fileio_env):
+ """Test start_line exceeding file length."""
+ result = asyncio.run(
+ fileio_env["file_io"].read_file(fileio_env["simple_file"], start_line=100),
+ )
+ text = result.content[0].get("text", "")
+ assert "Error" in text
+ assert "exceeds" in text
+
+
+def test_read_file_invalid_range(fileio_env):
+ """Test invalid line range (start > end)."""
+ result = asyncio.run(
+ fileio_env["file_io"].read_file(fileio_env["simple_file"], start_line=4, end_line=2),
+ )
+ text = result.content[0].get("text", "")
+ assert "Error" in text
+
+
+def test_read_file_truncated(fileio_env):
+ """Test file truncation by byte limit."""
+ result = asyncio.run(fileio_env["file_io"].read_file(fileio_env["large_file"]))
+ text = result.content[0].get("text", "")
+ assert "line 1" in text # Head is kept
+ assert "continue" in text.lower()
+
+
+def test_read_file_truncated_by_bytes(fileio_env):
+ """Test file truncation by byte limit."""
+ result = asyncio.run(fileio_env["file_io"].read_file(fileio_env["large_bytes_file"]))
+ text = result.content[0].get("text", "")
+ assert "continue" in text.lower()
+ assert "KB limit" in text
+
+
+def test_read_directory_error(fileio_env):
+ """Test reading a directory returns error."""
+ result = asyncio.run(fileio_env["file_io"].read_file(fileio_env["dir"]))
+ text = result.content[0].get("text", "")
+ assert "Error" in text
+ assert "not a file" in text
+
+
+def test_read_file_single_line_range(fileio_env):
+ """Test reading exactly one line (start_line == end_line)."""
+ result = asyncio.run(
+ fileio_env["file_io"].read_file(fileio_env["simple_file"], start_line=3, end_line=3),
+ )
+ text = result.content[0].get("text", "")
+ assert "line3" in text
+ assert "line2" not in text
+ assert "line4" not in text
+
+
+def test_read_file_only_start_line(fileio_env):
+ """Test reading from start_line to end of file (no end_line)."""
+ result = asyncio.run(
+ fileio_env["file_io"].read_file(fileio_env["simple_file"], start_line=4),
+ )
+ text = result.content[0].get("text", "")
+ assert "line4" in text
+ assert "line5" in text
+ assert "line1" not in text
+ assert "line3" not in text
+
+
+def test_read_file_only_end_line(fileio_env):
+ """Test reading from beginning to end_line (no start_line)."""
+ result = asyncio.run(
+ fileio_env["file_io"].read_file(fileio_env["simple_file"], end_line=2),
+ )
+ text = result.content[0].get("text", "")
+ assert "line1" in text
+ assert "line2" in text
+ assert "line4" not in text
+ assert "line5" not in text
+
+
+def test_read_file_end_line_clamped(fileio_env):
+ """Test end_line beyond total lines is silently clamped to file end."""
+ result = asyncio.run(
+ fileio_env["file_io"].read_file(fileio_env["simple_file"], start_line=1, end_line=999),
+ )
+ text = result.content[0].get("text", "")
+ assert "Error" not in text
+ assert "line1" in text
+ assert "line5" in text
+
+
+def test_read_file_continuation_hint(fileio_env):
+ """Partial range read without truncation shows remaining-lines continuation hint."""
+ # simple.txt has 5 lines; reading 1-3 leaves 2 more
+ result = asyncio.run(
+ fileio_env["file_io"].read_file(fileio_env["simple_file"], start_line=1, end_line=3),
+ )
+ text = result.content[0].get("text", "")
+ assert "more lines" in text
+ assert "start_line=4" in text
+
+
+def test_read_file_truncated_next_line_hint(fileio_env):
+ """Truncated large file provides a valid start_line > 1 to continue."""
+ result = asyncio.run(fileio_env["file_io"].read_file(fileio_env["large_file"]))
+ text = result.content[0].get("text", "")
+ match = re.search(r"start_line=(\d+)", text)
+ assert match is not None, "Expected start_line hint in truncated output"
+ assert int(match.group(1)) > 1
+
+
+def test_read_file_truncated_mid_line_message(fileio_env):
+ """Truncation mid-line reports which line is truncated (else branch)."""
+ # large_bytes_file lines are 101 bytes; truncation lands mid-line
+ result = asyncio.run(fileio_env["file_io"].read_file(fileio_env["large_bytes_file"]))
+ text = result.content[0].get("text", "")
+ assert "is truncated" in text.lower()
+
+
+def test_read_file_huge_single_line(fileio_env):
+ """Single line exceeding byte limit triggers 'partially shown' notice (newline_count==0 branch)."""
+ result = asyncio.run(fileio_env["file_io"].read_file(fileio_env["huge_line_file"]))
+ text = result.content[0].get("text", "")
+ assert "partially shown" in text.lower()
+ assert "start_line=2" in text
+
+
+def test_read_file_invalid_start_line_type(fileio_env):
+ """Non-integer start_line returns a descriptive error."""
+ result = asyncio.run(
+ fileio_env["file_io"].read_file(fileio_env["simple_file"], start_line="abc"),
+ )
+ text = result.content[0].get("text", "")
+ assert "Error" in text
+ assert "start_line" in text
+
+
+def test_read_file_invalid_end_line_type(fileio_env):
+ """Non-integer end_line returns a descriptive error."""
+ result = asyncio.run(
+ fileio_env["file_io"].read_file(fileio_env["simple_file"], end_line="xyz"),
+ )
+ text = result.content[0].get("text", "")
+ assert "Error" in text
+ assert "end_line" in text
+
+
+def test_read_file_start_line_as_string(fileio_env):
+ """Numeric-string start_line/end_line are coerced to int successfully."""
+ result = asyncio.run(
+ fileio_env["file_io"].read_file(fileio_env["simple_file"], start_line="2", end_line="4"),
+ )
+ text = result.content[0].get("text", "")
+ assert "Error" not in text
+ assert "line2" in text
+ assert "line4" in text
+
+
+def test_read_file_empty(fileio_env):
+ """Reading an empty file returns without error."""
+ result = asyncio.run(fileio_env["file_io"].read_file(fileio_env["empty_file"]))
+ text = result.content[0].get("text", "")
+ assert "Error" not in text
+
+
+# ============ FileIO Write Tests ============
+
+
+@pytest.fixture
+def write_env():
+ """Create temporary directory for write tests."""
+ test_dir = tempfile.mkdtemp(prefix="test_fileio_write_")
+ file_io = FileIO(working_dir=test_dir)
+ yield {"dir": test_dir, "file_io": file_io}
+ shutil.rmtree(test_dir, ignore_errors=True)
+
+
+def test_write_new_file(write_env):
+ """Test writing a new file."""
+ file_path = os.path.join(write_env["dir"], "new_file.txt")
+ result = asyncio.run(write_env["file_io"].write_file(file_path, "test content"))
+ text = result.content[0].get("text", "")
+ assert "Wrote" in text
+
+ with open(file_path, "r", encoding="utf-8") as f:
+ assert f.read() == "test content"
+
+
+def test_write_overwrite_file(write_env):
+ """Test overwriting existing file."""
+ file_path = os.path.join(write_env["dir"], "overwrite.txt")
+ with open(file_path, "w", encoding="utf-8") as f:
+ f.write("old content")
+
+ result = asyncio.run(write_env["file_io"].write_file(file_path, "new content"))
+ text = result.content[0].get("text", "")
+ assert "Wrote" in text
+
+ with open(file_path, "r", encoding="utf-8") as f:
+ assert f.read() == "new content"
+
+
+def test_write_empty_path(write_env):
+ """Test writing with empty path."""
+ result = asyncio.run(write_env["file_io"].write_file("", "content"))
+ text = result.content[0].get("text", "")
+ assert "Error" in text
+
+
+def test_write_relative_path(write_env):
+ """Test writing file with relative path."""
+ result = asyncio.run(write_env["file_io"].write_file("relative.txt", "relative content"))
+ text = result.content[0].get("text", "")
+ assert "Wrote" in text
+
+ file_path = os.path.join(write_env["dir"], "relative.txt")
+ assert os.path.exists(file_path)
+
+
+# ============ FileIO Edit Tests ============
+
+
+@pytest.fixture
+def edit_env():
+ """Create temporary directory with edit test file."""
+ test_dir = tempfile.mkdtemp(prefix="test_fileio_edit_")
+ file_io = FileIO(working_dir=test_dir)
+
+ edit_file = os.path.join(test_dir, "edit_test.txt")
+ with open(edit_file, "w", encoding="utf-8") as f:
+ f.write("Hello World\nThis is a test\nHello Again")
+
+ yield {"dir": test_dir, "file_io": file_io, "edit_file": edit_file}
+ shutil.rmtree(test_dir, ignore_errors=True)
+
+
+def test_edit_replace_text(edit_env):
+ """Test replacing text in file."""
+ result = asyncio.run(
+ edit_env["file_io"].edit_file(edit_env["edit_file"], "Hello", "Hi"),
+ )
+ text = result.content[0].get("text", "")
+ assert "Successfully" in text
+
+ with open(edit_env["edit_file"], "r", encoding="utf-8") as f:
+ content = f.read()
+ assert "Hello" not in content
+ assert "Hi World" in content
+ assert "Hi Again" in content
+
+
+def test_edit_text_not_found(edit_env):
+ """Test editing when text not found."""
+ result = asyncio.run(
+ edit_env["file_io"].edit_file(edit_env["edit_file"], "NotExists", "Replacement"),
+ )
+ text = result.content[0].get("text", "")
+ assert "Error" in text
+ assert "not found" in text
+
+
+def test_edit_nonexistent_file(edit_env):
+ """Test editing non-existent file."""
+ result = asyncio.run(
+ edit_env["file_io"].edit_file("nonexistent.txt", "old", "new"),
+ )
+ text = result.content[0].get("text", "")
+ assert "Error" in text
diff --git a/tests/light/test_truncate_text_output.py b/tests/light/test_truncate_text_output.py
new file mode 100644
index 00000000..ba25a842
--- /dev/null
+++ b/tests/light/test_truncate_text_output.py
@@ -0,0 +1,381 @@
+# -*- coding: utf-8 -*-
+# pylint: disable=missing-function-docstring
+"""Unit tests for _truncate_fresh, _retruncate, and truncate_text_output.
+
+Assumptions that mirror the production code:
+- Every single line is shorter than DEFAULT_MAX_BYTES.
+- Lines that exceed DEFAULT_MAX_BYTES are explicitly ignored / not fully read.
+- _truncate_fresh always includes a continuation hint in the notice, so
+ _retruncate can assume the hint is always present.
+
+Run:
+ cd tests/light && python test_truncate_text_output.py
+"""
+
+import re
+
+
+from reme.memory.file_based.utils.file_utils import (
+ TRUNCATION_NOTICE_MARKER,
+ _truncate_fresh,
+ _retruncate,
+ truncate_text_output,
+)
+
+# ── helpers ──────────────────────────────────────────────────────────────────
+
+LINE_BYTES = 20 # every test line is exactly this many bytes
+ENC = "utf-8"
+
+
+def make_line(i: int) -> str:
+ """Return a line that is exactly LINE_BYTES bytes in UTF-8.
+
+ Format: "L{i}" padded with underscores, terminated with "\\n".
+ Works for i up to 999.
+ """
+ prefix = f"L{i}"
+ return prefix + "_" * (LINE_BYTES - 1 - len(prefix)) + "\n"
+
+
+def make_text(n: int, start: int = 1) -> str:
+ """Build n lines starting from line number `start`."""
+ return "".join(make_line(i) for i in range(start, start + n))
+
+
+def content_of(result: str) -> str:
+ """Return the portion before the truncation marker."""
+ return result.split(TRUNCATION_NOTICE_MARKER)[0]
+
+
+def notice_of(result: str) -> str:
+ """Return the portion after the truncation marker, or '' if absent."""
+ parts = result.split(TRUNCATION_NOTICE_MARKER, 1)
+ return parts[1] if len(parts) > 1 else ""
+
+
+def parse_next_line(result: str) -> int:
+ """Extract 'start_line=X' from the notice, or -1 if absent."""
+ m = re.search(r"start_line=(\d+) to read more", notice_of(result))
+ return int(m.group(1)) if m else -1
+
+
+def parse_covers_bytes(result: str) -> int:
+ """Extract 'covers the next X bytes' from the notice."""
+ m = re.search(r"covers the next (\d+) bytes", notice_of(result))
+ return int(m.group(1)) if m else -1
+
+
+def assert_eq(a, b, msg=""):
+ assert a == b, f"{msg}: expected {b!r}, got {a!r}"
+
+
+def assert_in(needle, haystack, msg=""):
+ assert needle in haystack, f"{msg}: {needle!r} not found in {haystack!r}"
+
+
+def assert_not_in(needle, haystack, msg=""):
+ assert needle not in haystack, f"{msg}: {needle!r} unexpectedly found in {haystack!r}"
+
+
+# ── _truncate_fresh tests ─────────────────────────────────────────────────────
+
+
+def test_fresh_no_truncation_when_under_limit():
+ text = make_text(3) # 60 bytes
+ result = _truncate_fresh(text, start_line=1, total_lines=3, max_bytes=100, file_path=None, encoding=ENC)
+ assert_eq(result, text, "under limit: no change")
+ assert_not_in(TRUNCATION_NOTICE_MARKER, result)
+
+
+def test_fresh_no_truncation_at_exact_limit():
+ text = make_text(3) # 60 bytes
+ result = _truncate_fresh(text, start_line=1, total_lines=3, max_bytes=60, file_path=None, encoding=ENC)
+ assert_eq(result, text, "at exact limit: no change")
+ assert_not_in(TRUNCATION_NOTICE_MARKER, result)
+
+
+def test_fresh_mid_file_correct_next_line():
+ # 10 lines × 20 bytes = 200 bytes; max_bytes=95
+ # 95 bytes → 4 complete lines (80 bytes) + 15 bytes into line 5
+ # newline_count=4, next_line=1+4=5, 5<=10 → read_from=5
+ text = make_text(10)
+ result = _truncate_fresh(text, start_line=1, total_lines=10, max_bytes=95, file_path=None, encoding=ENC)
+
+ assert_in(TRUNCATION_NOTICE_MARKER, result)
+ assert_eq(parse_next_line(result), 5, "next_line after mid-file cut")
+ assert_eq(parse_covers_bytes(result), 95)
+ assert_in("L4", content_of(result), "line 4 fully included")
+ assert_not_in("L5\n", content_of(result), "line 5 not fully included")
+
+
+def test_fresh_notice_contains_total_lines_and_start_line():
+ text = make_text(10, start=3)
+ result = _truncate_fresh(text, start_line=3, total_lines=12, max_bytes=95, file_path="/tmp/foo.txt", encoding=ENC)
+ notice = notice_of(result)
+ assert_in("contains 12 lines in total", notice)
+ assert_in("starts at line 3", notice)
+ assert_in("file_path=/tmp/foo.txt", notice)
+
+
+def test_fresh_next_line_equals_total_lines_reads_from_last():
+ # 5 lines × 20 = 100 bytes; max_bytes=85
+ # 85 bytes → 4 complete lines (80 bytes) + 5 bytes into line 5
+ # newline_count=4, next_line=1+4=5 = total_lines → read_from=5
+ text = make_text(5)
+ result = _truncate_fresh(text, start_line=1, total_lines=5, max_bytes=85, file_path=None, encoding=ENC)
+
+ assert_in(TRUNCATION_NOTICE_MARKER, result)
+ assert_eq(parse_next_line(result), 5, "should re-read the last line")
+
+
+def test_fresh_next_line_overshoots_total_lines():
+ # max_bytes=115 → 5 complete lines (100 bytes) + 15 bytes into line 6
+ # newline_count=5, next_line=1+5=6 > total_lines(5), start_line(1) < 5
+ # → read_from = total_lines = 5
+ text = make_text(10)
+ result = _truncate_fresh(text, start_line=1, total_lines=5, max_bytes=115, file_path=None, encoding=ENC)
+
+ assert_in(TRUNCATION_NOTICE_MARKER, result)
+ assert_eq(parse_next_line(result), 5, "overshoot: fall back to total_lines")
+
+
+def test_fresh_long_first_line_skips_to_next_line():
+ # Line 1 is 200 bytes (> max_bytes=100), no '\n' in truncated result.
+ # newline_count=0, next_line=1+max(1,0)=2, 2<=5 → read_from=2
+ long_line = "A" * 199 + "\n" # 200 bytes
+ text = long_line + make_text(4, start=2)
+ result = _truncate_fresh(text, start_line=1, total_lines=5, max_bytes=100, file_path=None, encoding=ENC)
+
+ assert_in(TRUNCATION_NOTICE_MARKER, result)
+ assert_eq(parse_next_line(result), 2, "skip to line 2 after long line 1")
+ assert_not_in("A" * 101, content_of(result))
+
+
+def test_fresh_long_middle_line_skips_to_next_line():
+ # Lines 1-2 normal (40 bytes); line 3 is 200 bytes; lines 4-5 normal.
+ # max_bytes=100: fits lines 1-2 (40 bytes) + 60 bytes of line 3 (no '\n')
+ # newline_count=2, next_line=1+2=3, 3<=5 → read_from=3
+ text = make_text(2, start=1) + "B" * 199 + "\n" + make_text(2, start=4)
+ result = _truncate_fresh(text, start_line=1, total_lines=5, max_bytes=100, file_path=None, encoding=ENC)
+
+ assert_in(TRUNCATION_NOTICE_MARKER, result)
+ assert_eq(parse_next_line(result), 3, "restart from long line 3")
+
+
+def test_fresh_long_last_line_start_equals_total_no_notice():
+ # start_line == total_lines, single line too long → unhandled case, no notice.
+ # newline_count=0, next_line=5+1=6 > 5, start_line(5)==total_lines(5)
+ long_line = "C" * 199 + "\n" # 200 bytes
+ result = _truncate_fresh(long_line, start_line=5, total_lines=5, max_bytes=100, file_path=None, encoding=ENC)
+
+ assert_not_in(TRUNCATION_NOTICE_MARKER, result)
+ assert_eq(len(result), 100, "only max_bytes bytes returned")
+
+
+def test_fresh_long_last_line_arrived_from_previous_chunk_no_notice():
+ long_line = "D" * 199 + "\n"
+ result = _truncate_fresh(long_line, start_line=5, total_lines=5, max_bytes=80, file_path=None, encoding=ENC)
+
+ assert_not_in(TRUNCATION_NOTICE_MARKER, result)
+
+
+# ── _retruncate tests ─────────────────────────────────────────────────────────
+
+
+def _first_pass(n_lines: int = 50, max_bytes: int = 490) -> str:
+ """Produce a first-truncated text via _truncate_fresh."""
+ return _truncate_fresh(
+ make_text(n_lines),
+ start_line=1,
+ total_lines=n_lines,
+ max_bytes=max_bytes,
+ file_path="/file.txt",
+ encoding=ENC,
+ )
+
+
+def test_retruncate_within_slack_returns_unchanged():
+ # content ≈ 490 bytes; re-truncate with max_bytes=400.
+ # 490 <= 400+100=500 → slack hit, return unchanged.
+ pass1 = _first_pass(n_lines=50, max_bytes=490)
+ result = _retruncate(pass1, max_bytes=400, encoding=ENC)
+ assert_eq(result, pass1, "within slack: unchanged")
+
+
+def test_retruncate_updates_byte_count_in_notice():
+ pass1 = _first_pass(n_lines=50, max_bytes=490)
+ result = _retruncate(pass1, max_bytes=195, encoding=ENC)
+
+ assert_in(TRUNCATION_NOTICE_MARKER, result)
+ assert_eq(parse_covers_bytes(result), 195, "notice byte count updated")
+
+
+def test_retruncate_updates_next_line():
+ # pass1 content ≈ 490 bytes (24 complete lines), start_line=1.
+ # re-truncate at 195 bytes → 9 complete lines, next_line=1+9=10.
+ pass1 = _first_pass(n_lines=50, max_bytes=490)
+ result = _retruncate(pass1, max_bytes=195, encoding=ENC)
+
+ assert_eq(parse_next_line(result), 10, "next_line updated to 10")
+
+
+def test_retruncate_content_is_smaller():
+ pass1 = _first_pass(n_lines=50, max_bytes=490)
+ result = _retruncate(pass1, max_bytes=195, encoding=ENC)
+
+ assert len(content_of(result)) < len(content_of(pass1)), "content should shrink"
+
+
+def test_retruncate_missing_starts_at_line_returns_unchanged():
+ # Notice missing "starts at line X" → return unchanged.
+ # Use large content (600 bytes) to bypass the 100-byte slack.
+ large_content = make_text(30) # 600 bytes
+ broken = (
+ large_content + TRUNCATION_NOTICE_MARKER + "\nThe output above was truncated."
+ "\nThe full content is saved to the file and contains 30 lines in total."
+ "\nThis excerpt covers the next 600 bytes." # no "starts at line X"
+ "\nIf the current content is not enough, call `read_file` with file_path=/f.txt start_line=5 to read more."
+ )
+ result = _retruncate(broken, max_bytes=10, encoding=ENC)
+ assert_eq(result, broken, "missing 'starts at line': unchanged")
+
+
+def test_retruncate_missing_covers_next_bytes_returns_unchanged():
+ # Notice has malformed "covers ??? bytes" → return unchanged.
+ large_content = "X" * 500 + "\n"
+ broken = (
+ large_content + TRUNCATION_NOTICE_MARKER + "\nThe output above was truncated."
+ "\nThe full content is saved to the file and contains 10 lines in total."
+ "\nThis excerpt starts at line 1 and covers ??? bytes."
+ "\nIf the current content is not enough, call `read_file` with file_path=/f.txt start_line=5 to read more."
+ )
+ result = _retruncate(broken, max_bytes=10, encoding=ENC)
+ assert_eq(result, broken, "missing 'covers the next N bytes': unchanged")
+
+
+# ── truncate_text_output dispatch / guard tests ───────────────────────────────
+
+
+def test_dispatch_empty_string():
+ result = truncate_text_output("", start_line=1, total_lines=0, max_bytes=10)
+ assert_eq(result, "", "empty string bypassed")
+
+
+def test_dispatch_max_bytes_zero():
+ text = make_text(5)
+ result = truncate_text_output(text, max_bytes=0)
+ assert_eq(result, text, "max_bytes=0 bypassed")
+
+
+def test_dispatch_routes_to_fresh_when_no_marker():
+ text = make_text(10)
+ result = truncate_text_output(text, start_line=1, total_lines=10, max_bytes=95)
+ assert_in(TRUNCATION_NOTICE_MARKER, result)
+ assert_eq(parse_next_line(result), 5)
+
+
+def test_dispatch_routes_to_retruncate_when_marker_present():
+ pass1 = _first_pass(n_lines=50, max_bytes=490)
+ pass2 = truncate_text_output(pass1, max_bytes=195)
+ assert_in(TRUNCATION_NOTICE_MARKER, pass2)
+ assert_eq(parse_covers_bytes(pass2), 195)
+
+
+# ── multi-pass integration tests ──────────────────────────────────────────────
+
+
+def test_three_pass_decreasing_truncation():
+ """Three successive truncations with shrinking max_bytes."""
+ n_lines = 50
+ text = make_text(n_lines) # 1000 bytes
+
+ pass1 = truncate_text_output(text, start_line=1, total_lines=n_lines, max_bytes=490, file_path="/f.txt")
+ assert TRUNCATION_NOTICE_MARKER in pass1
+ assert parse_covers_bytes(pass1) == 490
+ assert parse_next_line(pass1) > 1
+
+ # 490 > 195+100=295 → re-truncation proceeds
+ pass2 = truncate_text_output(pass1, max_bytes=195)
+ assert TRUNCATION_NOTICE_MARKER in pass2
+ assert parse_covers_bytes(pass2) == 195
+ assert parse_next_line(pass2) < parse_next_line(pass1), "next_line regresses"
+ assert len(content_of(pass2)) < len(content_of(pass1))
+
+ # content_of(pass2) ≈ 195 bytes; 195 > 90+100=190 → re-truncation proceeds
+ pass3 = truncate_text_output(pass2, max_bytes=90)
+ assert TRUNCATION_NOTICE_MARKER in pass3
+ assert parse_covers_bytes(pass3) == 90
+ assert parse_next_line(pass3) < parse_next_line(pass2), "next_line regresses further"
+ assert len(content_of(pass3)) < len(content_of(pass2))
+
+
+def test_three_pass_next_lines_are_consistent():
+ """next_line values should monotonically decrease with each re-truncation."""
+ n_lines = 50
+ text = make_text(n_lines)
+
+ pass1 = truncate_text_output(text, start_line=1, total_lines=n_lines, max_bytes=490, file_path="/f.txt")
+ pass2 = truncate_text_output(pass1, max_bytes=195)
+ pass3 = truncate_text_output(pass2, max_bytes=90)
+
+ n1 = parse_next_line(pass1)
+ n2 = parse_next_line(pass2)
+ n3 = parse_next_line(pass3)
+
+ assert n1 > n2 > n3 > 1, f"Expected n1 > n2 > n3 > 1, got {n1} > {n2} > {n3}"
+
+
+# ── runner ────────────────────────────────────────────────────────────────────
+
+
+def run_all():
+ tests = [
+ # _truncate_fresh
+ test_fresh_no_truncation_when_under_limit,
+ test_fresh_no_truncation_at_exact_limit,
+ test_fresh_mid_file_correct_next_line,
+ test_fresh_notice_contains_total_lines_and_start_line,
+ test_fresh_next_line_equals_total_lines_reads_from_last,
+ test_fresh_next_line_overshoots_total_lines,
+ test_fresh_long_first_line_skips_to_next_line,
+ test_fresh_long_middle_line_skips_to_next_line,
+ test_fresh_long_last_line_start_equals_total_no_notice,
+ test_fresh_long_last_line_arrived_from_previous_chunk_no_notice,
+ # _retruncate
+ test_retruncate_within_slack_returns_unchanged,
+ test_retruncate_updates_byte_count_in_notice,
+ test_retruncate_updates_next_line,
+ test_retruncate_content_is_smaller,
+ test_retruncate_missing_starts_at_line_returns_unchanged,
+ test_retruncate_missing_covers_next_bytes_returns_unchanged,
+ # truncate_text_output dispatch / guard
+ test_dispatch_empty_string,
+ test_dispatch_max_bytes_zero,
+ test_dispatch_routes_to_fresh_when_no_marker,
+ test_dispatch_routes_to_retruncate_when_marker_present,
+ # multi-pass integration
+ test_three_pass_decreasing_truncation,
+ test_three_pass_next_lines_are_consistent,
+ ]
+
+ passed = 0
+ failed = 0
+ for t in tests:
+ try:
+ t()
+ print(f" PASS {t.__name__}")
+ passed += 1
+ except Exception as e:
+ print(f" FAIL {t.__name__}: {e}")
+ failed += 1
+
+ print(f"\n{passed} passed, {failed} failed out of {len(tests)} tests.")
+ return failed == 0
+
+
+if __name__ == "__main__":
+ import sys
+
+ ok = run_all()
+ sys.exit(0 if ok else 1)
diff --git a/tests/light/test_utils.py b/tests/light/test_utils.py
index f5cae021..27583948 100644
--- a/tests/light/test_utils.py
+++ b/tests/light/test_utils.py
@@ -1,50 +1,22 @@
"""Test utilities for copaw tests."""
import os
-from pathlib import Path
-from typing import Any
-from loguru import logger
+from agentscope.message import Msg, ThinkingBlock, TextBlock, ToolUseBlock, ToolResultBlock
-_token_counter = None
+from reme.memory.file_based.utils import AsMsgHandler
def get_token_counter():
- """Get or initialize the global token counter instance.
+ """Get HF token counter instance."""
+ from reme.core.utils import get_hf_token_counter
- Returns:
- TokenCounterBase: The token counter instance for Qwen models.
+ return get_hf_token_counter()
- Raises:
- RuntimeError: If token counter initialization fails.
- """
- global _token_counter
- if _token_counter is None:
- from agentscope.token import HuggingFaceTokenCounter
- # Use Qwen tokenizer for DashScope models
- # Qwen3 series uses the same tokenizer as Qwen2.5
-
- # Try local tokenizer first, fall back to online if not found
- local_tokenizer_path = Path(__file__).parent.parent.parent / "tokenizer"
-
- if local_tokenizer_path.exists() and (local_tokenizer_path / "tokenizer.json").exists():
- tokenizer_path = str(local_tokenizer_path)
- logger.info(f"Using local Qwen tokenizer from {tokenizer_path}")
- else:
- tokenizer_path = "Qwen/Qwen2.5-7B-Instruct"
- logger.info(
- "Local tokenizer not found, downloading from HuggingFace",
- )
-
- _token_counter = HuggingFaceTokenCounter(
- pretrained_model_name_or_path=tokenizer_path,
- use_mirror=True, # Use HF mirror for users in China
- use_fast=True,
- trust_remote_code=True,
- )
- logger.debug("Token counter initialized with Qwen tokenizer")
- return _token_counter
+def get_msg_handler() -> AsMsgHandler:
+ """Get AsMsgHandler instance."""
+ return AsMsgHandler(token_counter=get_token_counter())
def get_dash_chat_model(model_name: str = "qwen3.5-plus"):
@@ -54,8 +26,8 @@ def get_dash_chat_model(model_name: str = "qwen3.5-plus"):
load_env()
return OpenAIChatModel(
- api_key=os.environ["REME_LLM_API_KEY"],
- client_kwargs={"base_url": os.environ["REME_LLM_BASE_URL"]},
+ api_key=os.environ["LLM_API_KEY"],
+ client_kwargs={"base_url": os.environ["LLM_BASE_URL"]},
model_name=model_name,
)
@@ -63,27 +35,422 @@ def get_dash_chat_model(model_name: str = "qwen3.5-plus"):
def get_formatter():
"""Get formatter instance."""
from agentscope.formatter import OpenAIChatFormatter
- from agentscope.token import HuggingFaceTokenCounter
- from reme.memory.file_based.utils import _extract_text_from_messages
- class ReMeChatFormatter(OpenAIChatFormatter):
- """ReMe chat formatter class."""
+ return OpenAIChatFormatter()
- async def _count(self, msgs: list[dict[str, Any]]) -> int | None:
- """Count the number of tokens in the input messages. If token counter
- is not provided, `None` will be returned.
- Args:
- msgs (`list[Msg]`):
- The input messages to count tokens for.
- """
- if self.token_counter is None:
- return None
+def generate_large_code_content(target_tokens: int = 50000) -> str:
+ """生成大量代码内容，用于测试超长 tool_result。
- assert isinstance(self.token_counter, HuggingFaceTokenCounter)
- text = _extract_text_from_messages(msgs)
- token_ids = self.token_counter.tokenizer.encode(text)
- token_count = len(token_ids)
- return token_count
+ Args:
+ target_tokens: 目标 token 数（约 4 字符/token）
- return ReMeChatFormatter(token_counter=get_token_counter())
+ Returns:
+ 生成的代码内容字符串
+ """
+ code_template = '''
+# === File: src/module_{idx}/handlers.py ===
+"""Handler module {idx} for processing requests."""
+
+import asyncio
+import logging
+from typing import Any, Dict, List, Optional
+from dataclasses import dataclass, field
+from datetime import datetime
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class RequestContext_{idx}:
+ """Context for request processing in module {idx}."""
+ request_id: str
+ user_id: str
+ timestamp: datetime = field(default_factory=datetime.now)
+ metadata: Dict[str, Any] = field(default_factory=dict)
+ headers: Dict[str, str] = field(default_factory=dict)
+ query_params: Dict[str, str] = field(default_factory=dict)
+ body: Optional[bytes] = None
+ processed: bool = False
+ error_message: Optional[str] = None
+
+
+class Handler_{idx}:
+ """Main handler class for module {idx}."""
+
+ def __init__(self, config: Dict[str, Any]):
+ self.config = config
+ self.cache: Dict[str, Any] = {{}}
+ self.metrics: Dict[str, int] = {{
+ "requests_processed": 0,
+ "errors": 0,
+ "cache_hits": 0,
+ "cache_misses": 0,
+ }}
+ self._initialized = False
+ logger.info(f"Handler_{idx} initialized with config: {{config}}")
+
+ async def initialize(self) -> None:
+ """Initialize the handler with async resources."""
+ if self._initialized:
+ logger.warning("Handler_{idx} already initialized")
+ return
+
+ # Simulate async initialization
+ await asyncio.sleep(0.01)
+ self._initialized = True
+ logger.info("Handler_{idx} initialization complete")
+
+ async def process_request(self, context: RequestContext_{idx}) -> Dict[str, Any]:
+ """Process an incoming request.
+
+ Args:
+ context: The request context containing all request data
+
+ Returns:
+ Dict containing the response data
+ """
+ if not self._initialized:
+ raise RuntimeError("Handler not initialized")
+
+ self.metrics["requests_processed"] += 1
+
+ try:
+ # Check cache first
+ cache_key = f"{{context.request_id}}_{{context.user_id}}"
+ if cache_key in self.cache:
+ self.metrics["cache_hits"] += 1
+ return self.cache[cache_key]
+
+ self.metrics["cache_misses"] += 1
+
+ # Process the request
+ result = await self._do_process(context)
+
+ # Cache the result
+ self.cache[cache_key] = result
+ context.processed = True
+
+ return result
+
+ except Exception as e:
+ self.metrics["errors"] += 1
+ context.error_message = str(e)
+ logger.exception(f"Error processing request {{context.request_id}}: {{e}}")
+ raise
+
+ async def _do_process(self, context: RequestContext_{idx}) -> Dict[str, Any]:
+ """Internal processing logic."""
+ # Simulate some processing
+ await asyncio.sleep(0.001)
+
+ return {{
+ "status": "success",
+ "request_id": context.request_id,
+ "user_id": context.user_id,
+ "processed_at": datetime.now().isoformat(),
+ "module": "module_{idx}",
+ "data": {{
+ "result": f"Processed by handler_{idx}",
+ "metadata": context.metadata,
+ }}
+ }}
+
+ def get_metrics(self) -> Dict[str, int]:
+ """Return current metrics."""
+ return self.metrics.copy()
+
+ async def cleanup(self) -> None:
+ """Cleanup resources."""
+ self.cache.clear()
+ self._initialized = False
+ logger.info("Handler_{idx} cleaned up")
+
+'''
+
+ # 每个模块约 2000 字符 ≈ 500 tokens
+ # 目标 target_tokens，需要 target_tokens / 500 个模块
+ num_modules = max(1, target_tokens // 500)
+
+ parts = [f"# 大型项目代码检索结果\n# 共找到 {num_modules} 个相关模块\n"]
+ for i in range(num_modules):
+ parts.append(code_template.format(idx=i))
+
+ return "".join(parts)
+
+
+def build_sample_messages(include_large_tool_result: bool = True) -> list[Msg]:
+ """构建一段包含多种消息类型的模拟对话。
+
+ Args:
+ include_large_tool_result: 是否包含大型 tool_result，确保超过 128K token
+
+ Returns:
+ 消息列表
+ """
+ messages = [
+ Msg(
+ name="user",
+ role="user",
+ content="我想设置一个 Python 开发环境，你有什么建议？",
+ ),
+ Msg(
+ name="assistant",
+ role="assistant",
+ content=[
+ ThinkingBlock(type="thinking", thinking="用户想要搭建 Python 开发环境，我需要了解他的需求和偏好..."),
+ TextBlock(
+ type="text",
+ text="好的！我建议使用 Python 3.11 或 3.12 版本，它们性能更好且功能丰富。"
+ "你希望用于什么类型的开发？Web、数据科学还是其他？",
+ ),
+ ],
+ ),
+ Msg(name="user", role="user", content="主要是做 Web 开发，使用 FastAPI 框架。另外我喜欢用 pyenv 管理版本。"),
+ Msg(
+ name="assistant",
+ role="assistant",
+ content=[
+ ThinkingBlock(type="thinking", thinking="用户要做 FastAPI Web 开发，我需要搜索相关信息..."),
+ ToolUseBlock(
+ type="tool_use",
+ id="call_001",
+ name="search_web",
+ input={"query": "FastAPI Python version compatibility 2024"},
+ ),
+ ],
+ ),
+ Msg(
+ name="system",
+ role="system",
+ content=[
+ ToolResultBlock(
+ type="tool_result",
+ id="call_001",
+ name="search_web",
+ output=[
+ TextBlock(
+ type="text",
+ text="FastAPI 官方推荐使用 Python 3.8+ 版本，但 3.11/3.12 性能最佳。\n"
+ "主要依赖：\n"
+ "- Starlette: ASGI 框架\n"
+ "- Pydantic v2: 数据验证\n"
+ "- Uvicorn: ASGI 服务器\n"
+ "最新版本 FastAPI 0.109+ 完全支持 Python 3.12。\n"
+ "建议搭配 uv 或 pip-tools 进行依赖管理。",
+ ),
+ ],
+ ),
+ ],
+ ),
+ Msg(
+ name="assistant",
+ role="assistant",
+ content=[
+ ThinkingBlock(
+ type="thinking",
+ thinking="根据搜索结果，FastAPI 完全支持 Python 3.12，结合用户使用 pyenv 的习惯，我来给出完整的环境配置建议...",
+ ),
+ TextBlock(
+ type="text",
+ text="根据最新信息，我为你整理了 FastAPI 开发环境配置方案：\n\n"
+ "**1. Python 版本**\n"
+ "推荐 Python 3.12，性能最佳且完全兼容 FastAPI 0.109+。\n\n"
+ "**2. pyenv 安装步骤**\n"
+ "```bash\n"
+ "pyenv install 3.12.2\n"
+ "pyenv local 3.12.2\n"
+ "```\n\n"
+ "**3. 虚拟环境与依赖管理**\n"
+ "建议使用 `uv` 作为包管理器，它比 pip 快 10-100 倍。\n\n"
+ "需要我帮你生成项目初始化脚本吗？",
+ ),
+ ],
+ ),
+ Msg(name="user", role="user", content="好的，请帮我生成一个完整的项目初始化脚本，包括 pyproject.toml 配置。"),
+ Msg(
+ name="assistant",
+ role="assistant",
+ content=[
+ ThinkingBlock(
+ type="thinking",
+ thinking="用户需要完整的项目初始化脚本，我需要创建 pyproject.toml 并设置好 FastAPI 相关依赖...",
+ ),
+ ToolUseBlock(
+ type="tool_use",
+ id="call_002",
+ name="create_file",
+ input={
+ "path": "pyproject.toml",
+ "content": "[project]\nname = 'my-fastapi-app'\nversion = '0.1.0'\nrequires-python = "
+ "'>=3.11'\ndependencies = ['fastapi>=0.109.0', 'uvicorn[standard]>=0.27.0', "
+ "'pydantic>=2.0']",
+ },
+ ),
+ ],
+ ),
+ Msg(
+ name="system",
+ role="system",
+ content=[
+ ToolResultBlock(
+ type="tool_result",
+ id="call_002",
+ name="create_file",
+ output=[
+ TextBlock(type="text", text="文件 pyproject.toml 已创建成功。"),
+ ],
+ ),
+ ],
+ ),
+ Msg(
+ name="assistant",
+ role="assistant",
+ content=[
+ ToolUseBlock(
+ type="tool_use",
+ id="call_003",
+ name="create_file",
+ input={
+ "path": "init_project.sh",
+ "content": "#!/bin/bash\nset -e\n\n# 使用 pyenv 设置 Python 版本\npyenv install 3.12.2 "
+ "--skip-existing\npyenv local 3.12.2\n\n# 创建虚拟环境\npython -m venv "
+ ".venv\nsource .venv/bin/activate\n\n# 安装 uv 并使用它安装依赖\npip install uv"
+ "\nuv pip install -e .\n\necho '环境初始化完成！'",
+ },
+ ),
+ ],
+ ),
+ Msg(
+ name="system",
+ role="system",
+ content=[
+ ToolResultBlock(
+ type="tool_result",
+ id="call_003",
+ name="create_file",
+ output=[
+ TextBlock(type="text", text="文件 init_project.sh 已创建成功。"),
+ ],
+ ),
+ ],
+ ),
+ Msg(
+ name="assistant",
+ role="assistant",
+ content=[
+ TextBlock(
+ type="text",
+ text="我已为你创建了两个文件：\n\n"
+ "1. **pyproject.toml** - 项目配置文件，包含 FastAPI 核心依赖\n"
+ "2. **init_project.sh** - 一键初始化脚本\n\n"
+ "运行以下命令即可初始化项目：\n"
+ "```bash\n"
+ "chmod +x init_project.sh && ./init_project.sh\n"
+ "```\n\n"
+ "还有什么需要帮助的吗？",
+ ),
+ ],
+ ),
+ Msg(name="user", role="user", content="太棒了！请帮我搜索一下项目中所有的 handler 相关代码。"),
+ ]
+
+ # 添加大型代码搜索结果（确保超过 128K token）
+ if include_large_tool_result:
+ # 生成超大的代码搜索结果，目标 ~140K tokens
+ large_code_content = generate_large_code_content(target_tokens=140000)
+
+ messages.extend(
+ [
+ Msg(
+ name="assistant",
+ role="assistant",
+ content=[
+ ThinkingBlock(
+ type="thinking",
+ thinking="用户要我搜索项目中的 handler 代码，我需要使用代码搜索工具...",
+ ),
+ ToolUseBlock(
+ type="tool_use",
+ id="call_004",
+ name="search_codebase",
+ input={"query": "handler class implementation"},
+ ),
+ ],
+ ),
+ Msg(
+ name="system",
+ role="system",
+ content=[
+ ToolResultBlock(
+ type="tool_result",
+ id="call_004",
+ name="search_codebase",
+ output=[
+ TextBlock(type="text", text=large_code_content),
+ ],
+ ),
+ ],
+ ),
+ Msg(
+ name="assistant",
+ role="assistant",
+ content=[
+ ThinkingBlock(type="thinking", thinking="搜索返回了大量 handler 代码，我需要为用户整理一下..."),
+ TextBlock(
+ type="text",
+ text="我已经找到了项目中所有的 handler 相关代码。\n\n"
+ "这些 handler 类包含：\n"
+ "- 请求处理逻辑\n"
+ "- 缓存管理\n"
+ "- 指标统计\n"
+ "- 异步初始化\n\n"
+ "你需要我详细解释某个具体的 handler 吗？",
+ ),
+ ],
+ ),
+ ],
+ )
+
+ # 添加更多对话
+ messages.extend(
+ [
+ Msg(
+ name="user",
+ role="user",
+ content="还有一个问题，我应该如何配置 VS Code 来获得最佳的 FastAPI 开发体验？",
+ ),
+ Msg(
+ name="assistant",
+ role="assistant",
+ content=[
+ ThinkingBlock(
+ type="thinking",
+ thinking="用户询问 VS Code 配置，我需要推荐适合 FastAPI 开发的扩展和设置...",
+ ),
+ TextBlock(
+ type="text",
+ text="VS Code 的 FastAPI 开发配置建议：\n\n"
+ "**推荐扩展：**\n"
+ "- Python (Microsoft)\n"
+ "- Pylance - 类型检查和智能补全\n"
+ "- Ruff - 快速 linter 和 formatter\n"
+ "- REST Client - API 测试\n\n"
+ "**settings.json 配置：**\n"
+ "```json\n"
+ "{\n"
+ ' "python.defaultInterpreterPath": ".venv/bin/python",\n'
+ ' "[python]": {\n'
+ ' "editor.defaultFormatter": "charliermarsh.ruff",\n'
+ ' "editor.formatOnSave": true\n'
+ " }\n"
+ "}\n"
+ "```\n\n"
+ "这样配置后，你就能获得完整的类型提示和自动格式化支持了！",
+ ),
+ ],
+ ),
+ ],
+ )
+
+ return messages
diff --git a/tests/test_base_file_watcher.py b/tests/test_base_file_watcher.py
new file mode 100644
index 00000000..234d03c2
--- /dev/null
+++ b/tests/test_base_file_watcher.py
@@ -0,0 +1,791 @@
+"""
+Async unit tests for BaseFileWatcher covering:
+- Existing paths and files monitoring
+- Non-existent paths handling
+- File suffix filtering
+- Start/stop lifecycle
+- Callback functionality
+- rebuild_index_on_start feature
+
+Usage:
+ pytest tests/test_base_file_watcher.py -v
+ pytest tests/test_base_file_watcher.py -v -k "test_existing"
+"""
+
+# pylint: disable=redefined-outer-name,protected-access,unused-argument
+
+import asyncio
+import tempfile
+from pathlib import Path
+from unittest.mock import AsyncMock, MagicMock
+
+import pytest
+from watchfiles import Change
+
+from reme.core.file_watcher.base_file_watcher import BaseFileWatcher
+
+
+# ==================== Fixtures ====================
+
+
+@pytest.fixture
+def temp_dir():
+ """Create a temporary directory for testing."""
+ with tempfile.TemporaryDirectory() as tmpdir:
+ yield Path(tmpdir)
+
+
+@pytest.fixture
+def temp_files(temp_dir: Path):
+ """Create temporary test files."""
+ files = {}
+
+ # Create .txt files
+ for i in range(3):
+ file_path = temp_dir / f"test_file_{i}.txt"
+ file_path.write_text(f"Content of test file {i}")
+ files[f"txt_{i}"] = file_path
+
+ # Create .py files
+ for i in range(2):
+ file_path = temp_dir / f"test_script_{i}.py"
+ file_path.write_text(f"# Python script {i}\nprint('hello')")
+ files[f"py_{i}"] = file_path
+
+ # Create .md file
+ md_file = temp_dir / "readme.md"
+ md_file.write_text("# README")
+ files["md_0"] = md_file
+
+ yield files
+
+
+@pytest.fixture
+def temp_nested_dir(temp_dir: Path):
+ """Create nested directory structure."""
+ # Create subdirectories
+ sub1 = temp_dir / "subdir1"
+ sub1.mkdir()
+ sub2 = temp_dir / "subdir2"
+ sub2.mkdir()
+ nested = sub1 / "nested"
+ nested.mkdir()
+
+ # Create files in subdirectories
+ (sub1 / "file1.txt").write_text("subdir1 file")
+ (sub2 / "file2.txt").write_text("subdir2 file")
+ (nested / "nested_file.txt").write_text("nested file")
+
+ yield temp_dir
+
+
+# ==================== Test Existing Paths ====================
+
+
+class TestExistingPaths:
+ """Tests for existing paths and files."""
+
+ @pytest.mark.asyncio
+ async def test_init_with_single_existing_path(self, temp_dir: Path):
+ """Test initialization with a single existing path."""
+ watcher = BaseFileWatcher(watch_paths=str(temp_dir))
+
+ assert watcher.watch_paths == [str(temp_dir)]
+ assert watcher.recursive is False
+ assert watcher.is_running() is False
+
+ @pytest.mark.asyncio
+ async def test_init_with_multiple_existing_paths(self, temp_dir: Path):
+ """Test initialization with multiple existing paths."""
+ sub1 = temp_dir / "dir1"
+ sub2 = temp_dir / "dir2"
+ sub1.mkdir()
+ sub2.mkdir()
+
+ watcher = BaseFileWatcher(watch_paths=[str(sub1), str(sub2)])
+
+ assert len(watcher.watch_paths) == 2
+ assert str(sub1) in watcher.watch_paths
+ assert str(sub2) in watcher.watch_paths
+
+ @pytest.mark.asyncio
+ async def test_init_with_existing_file(self, temp_files):
+ """Test initialization with existing file path."""
+ file_path = temp_files["txt_0"]
+ watcher = BaseFileWatcher(watch_paths=str(file_path))
+
+ assert watcher.watch_paths == [str(file_path)]
+
+ @pytest.mark.asyncio
+ async def test_start_with_existing_path(self, temp_dir: Path):
+ """Test starting watcher with existing path."""
+ watcher = BaseFileWatcher(watch_paths=str(temp_dir))
+
+ await watcher.start()
+ assert watcher.is_running() is True
+
+ await watcher.close()
+ assert watcher.is_running() is False
+
+ @pytest.mark.asyncio
+ async def test_start_stop_lifecycle(self, temp_dir: Path):
+ """Test watcher start/stop lifecycle."""
+ watcher = BaseFileWatcher(watch_paths=str(temp_dir))
+
+ # Start
+ await watcher.start()
+ assert watcher.is_running() is True
+ assert watcher._watch_task is not None
+
+ # Stop
+ await watcher.close()
+ assert watcher.is_running() is False
+
+ # Restart
+ await watcher.start()
+ assert watcher.is_running() is True
+
+ await watcher.close()
+
+ @pytest.mark.asyncio
+ async def test_multiple_start_calls(self, temp_dir: Path):
+ """Test that multiple start calls don't create multiple tasks."""
+ watcher = BaseFileWatcher(watch_paths=str(temp_dir))
+
+ await watcher.start()
+ task1 = watcher._watch_task
+
+ await watcher.start() # Second call should be ignored
+ task2 = watcher._watch_task
+
+ assert task1 is task2
+ await watcher.close()
+
+ @pytest.mark.asyncio
+ async def test_multiple_close_calls(self, temp_dir: Path):
+ """Test that multiple close calls are safe."""
+ watcher = BaseFileWatcher(watch_paths=str(temp_dir))
+
+ await watcher.start()
+ await watcher.close()
+ await watcher.close() # Second call should be safe
+
+ assert watcher.is_running() is False
+
+
+# ==================== Test Non-Existent Paths ====================
+
+
+class TestNonExistentPaths:
+ """Tests for non-existent paths handling."""
+
+ @pytest.mark.asyncio
+ async def test_init_with_nonexistent_path(self):
+ """Test initialization with non-existent path."""
+ nonexistent = "/path/that/does/not/exist"
+ watcher = BaseFileWatcher(watch_paths=nonexistent)
+
+ assert watcher.watch_paths == [nonexistent]
+
+ @pytest.mark.asyncio
+ async def test_start_with_nonexistent_path(self):
+ """Test starting watcher with non-existent path (should handle gracefully)."""
+ nonexistent = "/path/that/does/not/exist"
+ watcher = BaseFileWatcher(watch_paths=nonexistent)
+
+ await watcher.start()
+ assert watcher.is_running() is True
+
+ # Give it a moment to enter the watch loop and detect the missing path
+ await asyncio.sleep(0.1)
+
+ await watcher.close()
+ assert watcher.is_running() is False
+
+ @pytest.mark.asyncio
+ async def test_mixed_existing_and_nonexistent_paths(self, temp_dir: Path):
+ """Test with mix of existing and non-existent paths."""
+ nonexistent = "/path/that/does/not/exist"
+ watcher = BaseFileWatcher(watch_paths=[str(temp_dir), nonexistent])
+
+ await watcher.start()
+ assert watcher.is_running() is True
+
+ # Give it time to filter paths
+ await asyncio.sleep(0.1)
+
+ await watcher.close()
+
+ @pytest.mark.asyncio
+ async def test_all_paths_nonexistent(self):
+ """Test when all paths are non-existent."""
+ watcher = BaseFileWatcher(
+ watch_paths=["/nonexistent1", "/nonexistent2"],
+ )
+
+ await watcher.start()
+ assert watcher.is_running() is True
+
+ # Wait for retry logic
+ await asyncio.sleep(0.2)
+
+ await watcher.close()
+
+ @pytest.mark.asyncio
+ async def test_empty_watch_paths(self):
+ """Test with empty watch paths list."""
+ watcher = BaseFileWatcher(watch_paths=[])
+
+ await watcher.start()
+ assert watcher.is_running() is True
+
+ await asyncio.sleep(0.1)
+ await watcher.close()
+
+
+# ==================== Test File Filtering ====================
+
+
+class TestFileFiltering:
+ """Tests for file suffix filtering."""
+
+ @pytest.mark.asyncio
+ async def test_watch_filter_no_filters(self, temp_files):
+ """Test watch_filter with no suffix filters (should match all)."""
+ watcher = BaseFileWatcher(watch_paths="/tmp")
+
+ assert watcher.watch_filter(Change.added, "test.txt") is True
+ assert watcher.watch_filter(Change.added, "test.py") is True
+ assert watcher.watch_filter(Change.added, "test.md") is True
+ assert watcher.watch_filter(Change.added, "noextension") is True
+
+ @pytest.mark.asyncio
+ async def test_watch_filter_with_txt_suffix(self):
+ """Test watch_filter with .txt suffix filter."""
+ watcher = BaseFileWatcher(watch_paths="/tmp", suffix_filters=[".txt"])
+
+ assert watcher.watch_filter(Change.added, "test.txt") is True
+ assert watcher.watch_filter(Change.added, "test.py") is False
+ assert watcher.watch_filter(Change.added, "file.txt.bak") is False
+
+ @pytest.mark.asyncio
+ async def test_watch_filter_with_multiple_suffixes(self):
+ """Test watch_filter with multiple suffix filters."""
+ watcher = BaseFileWatcher(
+ watch_paths="/tmp",
+ suffix_filters=[".txt", ".py", ".md"],
+ )
+
+ assert watcher.watch_filter(Change.added, "test.txt") is True
+ assert watcher.watch_filter(Change.added, "script.py") is True
+ assert watcher.watch_filter(Change.added, "readme.md") is True
+ assert watcher.watch_filter(Change.added, "config.json") is False
+
+ @pytest.mark.asyncio
+ async def test_watch_filter_suffix_without_dot(self):
+ """Test watch_filter handles suffixes without leading dot."""
+ watcher = BaseFileWatcher(
+ watch_paths="/tmp",
+ suffix_filters=["txt", "py"], # Without dots
+ )
+
+ assert watcher.watch_filter(Change.added, "test.txt") is True
+ assert watcher.watch_filter(Change.added, "script.py") is True
+
+ @pytest.mark.asyncio
+ async def test_watch_filter_all_change_types(self):
+ """Test watch_filter works with all Change types."""
+ watcher = BaseFileWatcher(
+ watch_paths="/tmp",
+ suffix_filters=[".txt"],
+ )
+
+ # All change types should work with filter
+ assert watcher.watch_filter(Change.added, "test.txt") is True
+ assert watcher.watch_filter(Change.modified, "test.txt") is True
+ assert watcher.watch_filter(Change.deleted, "test.txt") is True
+
+
+# ==================== Test Callback Functionality ====================
+
+
+class TestCallbackFunctionality:
+ """Tests for callback functionality."""
+
+ @pytest.mark.asyncio
+ async def test_sync_callback(self, temp_dir: Path):
+ """Test synchronous callback function."""
+ callback_called = []
+
+ def sync_callback(changes):
+ callback_called.append(changes)
+
+ watcher = BaseFileWatcher(
+ watch_paths=str(temp_dir),
+ callback=sync_callback,
+ )
+
+ # Simulate changes
+ test_changes = {(Change.added, str(temp_dir / "test.txt"))}
+ await watcher.on_changes(test_changes)
+
+ assert len(callback_called) == 1
+ assert callback_called[0] == test_changes
+
+ @pytest.mark.asyncio
+ async def test_async_callback(self, temp_dir: Path):
+ """Test asynchronous callback function."""
+ callback_called = []
+
+ async def async_callback(changes):
+ callback_called.append(changes)
+
+ watcher = BaseFileWatcher(
+ watch_paths=str(temp_dir),
+ callback=async_callback,
+ )
+
+ # Simulate changes
+ test_changes = {(Change.modified, str(temp_dir / "test.txt"))}
+ await watcher.on_changes(test_changes)
+
+ assert len(callback_called) == 1
+ assert callback_called[0] == test_changes
+
+ @pytest.mark.asyncio
+ async def test_no_callback_uses_internal_handler(self, temp_dir: Path):
+ """Test that without callback, internal _on_changes is used."""
+ watcher = BaseFileWatcher(watch_paths=str(temp_dir))
+
+ # Mock internal _on_changes
+ watcher._on_changes = AsyncMock()
+
+ test_changes = {(Change.added, str(temp_dir / "test.txt"))}
+ await watcher.on_changes(test_changes)
+
+ watcher._on_changes.assert_called_once_with(test_changes)
+
+
+# ==================== Test Scan on Start ====================
+
+
+class TestRebuildIndexOnStart:
+ """Tests for rebuild_index_on_start feature."""
+
+ @pytest.mark.asyncio
+ async def test_rebuild_index_on_start_false(self, temp_files, temp_dir: Path):
+ """Test that rebuild_index_on_start=False doesn't scan existing files."""
+ callback_called = []
+
+ async def callback(changes):
+ callback_called.append(changes)
+
+ # Create mock file_store
+ mock_file_store = MagicMock()
+ mock_file_store.list_files = AsyncMock(return_value=[])
+ mock_file_store.get_file_chunks = AsyncMock(return_value=[])
+
+ watcher = BaseFileWatcher(
+ watch_paths=str(temp_dir),
+ rebuild_index_on_start=False,
+ callback=callback,
+ file_store=mock_file_store,
+ )
+
+ await watcher.start()
+ await asyncio.sleep(0.1)
+ await watcher.close()
+
+ # No callback should be called for existing files
+ assert len(callback_called) == 0
+
+ @pytest.mark.asyncio
+ async def test_rebuild_index_on_start_true_with_files(self, temp_files, temp_dir: Path):
+ """Test that rebuild_index_on_start=True scans existing files."""
+ callback_called = []
+
+ async def callback(changes):
+ callback_called.append(changes)
+
+ # Create mock file_store
+ mock_file_store = MagicMock()
+ mock_file_store.list_files = AsyncMock(return_value=[])
+ mock_file_store.get_file_chunks = AsyncMock(return_value=[])
+
+ watcher = BaseFileWatcher(
+ watch_paths=str(temp_dir),
+ rebuild_index_on_start=True,
+ callback=callback,
+ file_store=mock_file_store,
+ )
+
+ await watcher.start()
+ await asyncio.sleep(0.1)
+ await watcher.close()
+
+ # Callback should be called with existing files
+ assert len(callback_called) >= 1
+
+ # Check that files were detected as Change.added
+ all_changes = set()
+ for change_set in callback_called:
+ all_changes.update(change_set)
+
+ assert all(change == Change.added for change, _ in all_changes)
+
+ @pytest.mark.asyncio
+ async def test_rebuild_index_on_start_with_suffix_filter(self, temp_files, temp_dir: Path):
+ """Test rebuild_index_on_start respects suffix filters."""
+ callback_called = []
+
+ async def callback(changes):
+ callback_called.append(changes)
+
+ mock_file_store = MagicMock()
+ mock_file_store.list_files = AsyncMock(return_value=[])
+ mock_file_store.get_file_chunks = AsyncMock(return_value=[])
+
+ watcher = BaseFileWatcher(
+ watch_paths=str(temp_dir),
+ rebuild_index_on_start=True,
+ suffix_filters=[".txt"],
+ callback=callback,
+ file_store=mock_file_store,
+ )
+
+ await watcher.start()
+ await asyncio.sleep(0.1)
+ await watcher.close()
+
+ # Check only .txt files were scanned
+ if callback_called:
+ all_changes = set()
+ for change_set in callback_called:
+ all_changes.update(change_set)
+
+ for _, path in all_changes:
+ assert path.endswith(".txt"), f"Expected .txt file, got {path}"
+
+ @pytest.mark.asyncio
+ async def test_rebuild_index_on_start_recursive(self, temp_nested_dir: Path):
+ """Test rebuild_index_on_start with recursive=True."""
+ callback_called = []
+
+ async def callback(changes):
+ callback_called.append(changes)
+
+ mock_file_store = MagicMock()
+ mock_file_store.list_files = AsyncMock(return_value=[])
+ mock_file_store.get_file_chunks = AsyncMock(return_value=[])
+
+ watcher = BaseFileWatcher(
+ watch_paths=str(temp_nested_dir),
+ rebuild_index_on_start=True,
+ recursive=True,
+ suffix_filters=[".txt"],
+ callback=callback,
+ file_store=mock_file_store,
+ )
+
+ await watcher.start()
+ await asyncio.sleep(0.1)
+ await watcher.close()
+
+ # Should find files in nested directories
+ if callback_called:
+ all_changes = set()
+ for change_set in callback_called:
+ all_changes.update(change_set)
+
+ paths = [path for _, path in all_changes]
+ # Should find nested_file.txt
+ nested_found = any("nested_file.txt" in p for p in paths)
+ assert nested_found, "Should find files in nested directories"
+
+ @pytest.mark.asyncio
+ async def test_rebuild_index_on_start_non_recursive(self, temp_nested_dir: Path):
+ """Test rebuild_index_on_start with recursive=False."""
+ callback_called = []
+
+ async def callback(changes):
+ callback_called.append(changes)
+
+ mock_file_store = MagicMock()
+ mock_file_store.list_files = AsyncMock(return_value=[])
+ mock_file_store.get_file_chunks = AsyncMock(return_value=[])
+
+ watcher = BaseFileWatcher(
+ watch_paths=str(temp_nested_dir),
+ rebuild_index_on_start=True,
+ recursive=False,
+ suffix_filters=[".txt"],
+ callback=callback,
+ file_store=mock_file_store,
+ )
+
+ await watcher.start()
+ await asyncio.sleep(0.1)
+ await watcher.close()
+
+ # Should NOT find files in nested directories
+ if callback_called:
+ all_changes = set()
+ for change_set in callback_called:
+ all_changes.update(change_set)
+
+ paths = [path for _, path in all_changes]
+ nested_found = any("nested_file.txt" in p for p in paths)
+ assert not nested_found, "Should not find files in nested directories"
+
+ @pytest.mark.asyncio
+ async def test_rebuild_index_on_start_nonexistent_path(self):
+ """Test rebuild_index_on_start with non-existent path."""
+ callback_called = []
+
+ async def callback(changes):
+ callback_called.append(changes)
+
+ mock_file_store = MagicMock()
+ mock_file_store.list_files = AsyncMock(return_value=[])
+ mock_file_store.get_file_chunks = AsyncMock(return_value=[])
+
+ watcher = BaseFileWatcher(
+ watch_paths="/nonexistent/path",
+ rebuild_index_on_start=True,
+ callback=callback,
+ file_store=mock_file_store,
+ )
+
+ await watcher.start()
+ await asyncio.sleep(0.1)
+ await watcher.close()
+
+ # No files should be found
+ assert len(callback_called) == 0
+
+
+# ==================== Test Dynamic Path Management ====================
+
+
+class TestDynamicPathManagement:
+ """Tests for dynamic path add/remove."""
+
+ @pytest.mark.asyncio
+ async def test_add_path_when_stopped(self, temp_dir: Path):
+ """Test adding path when watcher is stopped."""
+ watcher = BaseFileWatcher(watch_paths=str(temp_dir))
+
+ new_dir = temp_dir / "new_dir"
+ new_dir.mkdir()
+
+ await watcher.add_path(str(new_dir))
+
+ assert str(new_dir) in watcher.watch_paths
+
+ @pytest.mark.asyncio
+ async def test_add_path_when_running(self, temp_dir: Path):
+ """Test adding path when watcher is running (triggers restart)."""
+ watcher = BaseFileWatcher(watch_paths=str(temp_dir))
+
+ await watcher.start()
+ assert watcher.is_running()
+
+ new_dir = temp_dir / "new_dir"
+ new_dir.mkdir()
+
+ await watcher.add_path(str(new_dir))
+
+ assert str(new_dir) in watcher.watch_paths
+ assert watcher.is_running()
+
+ await watcher.close()
+
+ @pytest.mark.asyncio
+ async def test_add_duplicate_path(self, temp_dir: Path):
+ """Test adding duplicate path is ignored."""
+ watcher = BaseFileWatcher(watch_paths=str(temp_dir))
+
+ original_count = len(watcher.watch_paths)
+ await watcher.add_path(str(temp_dir))
+
+ assert len(watcher.watch_paths) == original_count
+
+ @pytest.mark.asyncio
+ async def test_remove_path(self, temp_dir: Path):
+ """Test removing path."""
+ sub1 = temp_dir / "sub1"
+ sub2 = temp_dir / "sub2"
+ sub1.mkdir()
+ sub2.mkdir()
+
+ watcher = BaseFileWatcher(watch_paths=[str(sub1), str(sub2)])
+
+ await watcher.remove_path(str(sub1))
+
+ assert str(sub1) not in watcher.watch_paths
+ assert str(sub2) in watcher.watch_paths
+
+ @pytest.mark.asyncio
+ async def test_remove_nonexistent_path(self, temp_dir: Path):
+ """Test removing path that's not in watch list."""
+ watcher = BaseFileWatcher(watch_paths=str(temp_dir))
+
+ original_paths = watcher.watch_paths.copy()
+ await watcher.remove_path("/some/other/path")
+
+ assert watcher.watch_paths == original_paths
+
+
+# ==================== Test Configuration Options ====================
+
+
+class TestConfigurationOptions:
+ """Tests for various configuration options."""
+
+ @pytest.mark.asyncio
+ async def test_debounce_setting(self, temp_dir: Path):
+ """Test debounce configuration."""
+ watcher = BaseFileWatcher(
+ watch_paths=str(temp_dir),
+ debounce=1000,
+ )
+
+ assert watcher.debounce == 1000
+
+ @pytest.mark.asyncio
+ async def test_chunk_settings(self, temp_dir: Path):
+ """Test chunk configuration."""
+ watcher = BaseFileWatcher(
+ watch_paths=str(temp_dir),
+ chunk_tokens=500,
+ chunk_overlap=100,
+ )
+
+ assert watcher.chunk_tokens == 500
+ assert watcher.chunk_overlap == 100
+
+ @pytest.mark.asyncio
+ async def test_recursive_setting(self, temp_dir: Path):
+ """Test recursive configuration."""
+ watcher = BaseFileWatcher(
+ watch_paths=str(temp_dir),
+ recursive=True,
+ )
+
+ assert watcher.recursive is True
+
+ @pytest.mark.asyncio
+ async def test_kwargs_preserved(self, temp_dir: Path):
+ """Test that extra kwargs are preserved."""
+ watcher = BaseFileWatcher(
+ watch_paths=str(temp_dir),
+ custom_arg1="value1",
+ custom_arg2=123,
+ )
+
+ assert watcher.kwargs.get("custom_arg1") == "value1"
+ assert watcher.kwargs.get("custom_arg2") == 123
+
+
+# ==================== Test Edge Cases ====================
+
+
+class TestEdgeCases:
+ """Tests for edge cases and boundary conditions."""
+
+ @pytest.mark.asyncio
+ async def test_watch_single_file(self, temp_files):
+ """Test watching a single file instead of directory."""
+ file_path = temp_files["txt_0"]
+
+ mock_file_store = MagicMock()
+ mock_file_store.list_files = AsyncMock(return_value=[])
+ mock_file_store.get_file_chunks = AsyncMock(return_value=[])
+
+ callback_called = []
+
+ async def callback(changes):
+ callback_called.append(changes)
+
+ watcher = BaseFileWatcher(
+ watch_paths=str(file_path),
+ rebuild_index_on_start=True,
+ callback=callback,
+ file_store=mock_file_store,
+ )
+
+ await watcher.start()
+ await asyncio.sleep(0.1)
+ await watcher.close()
+
+ # Single file should be detected
+ if callback_called:
+ all_changes = set()
+ for change_set in callback_called:
+ all_changes.update(change_set)
+ assert len(all_changes) == 1
+
+ @pytest.mark.asyncio
+ async def test_empty_directory(self, temp_dir: Path):
+ """Test watching empty directory."""
+ empty_dir = temp_dir / "empty"
+ empty_dir.mkdir()
+
+ mock_file_store = MagicMock()
+ mock_file_store.list_files = AsyncMock(return_value=[])
+ mock_file_store.get_file_chunks = AsyncMock(return_value=[])
+
+ callback_called = []
+
+ async def callback(changes):
+ callback_called.append(changes)
+
+ watcher = BaseFileWatcher(
+ watch_paths=str(empty_dir),
+ rebuild_index_on_start=True,
+ callback=callback,
+ file_store=mock_file_store,
+ )
+
+ await watcher.start()
+ await asyncio.sleep(0.1)
+ await watcher.close()
+
+ # No files should be detected
+ assert len(callback_called) == 0
+
+ @pytest.mark.asyncio
+ async def test_special_characters_in_path(self, temp_dir: Path):
+ """Test paths with special characters."""
+ special_dir = temp_dir / "test dir with spaces"
+ special_dir.mkdir()
+
+ file_path = special_dir / "file with spaces.txt"
+ file_path.write_text("content")
+
+ watcher = BaseFileWatcher(watch_paths=str(special_dir))
+
+ assert watcher.watch_filter(Change.added, str(file_path)) is True
+
+ @pytest.mark.asyncio
+ async def test_unicode_in_path(self, temp_dir: Path):
+ """Test paths with unicode characters."""
+ unicode_dir = temp_dir / "测试目录"
+ unicode_dir.mkdir()
+
+ file_path = unicode_dir / "文件.txt"
+ file_path.write_text("内容")
+
+ watcher = BaseFileWatcher(
+ watch_paths=str(unicode_dir),
+ suffix_filters=[".txt"],
+ )
+
+ assert watcher.watch_filter(Change.added, str(file_path)) is True
+
+
+# ==================== Main Entry Point ====================
+
+
+if __name__ == "__main__":
+ pytest.main([__file__, "-v"])
diff --git a/tests/test_reme_memory_error_handling.py b/tests/test_reme_memory_error_handling.py
new file mode 100644
index 00000000..99953efd
--- /dev/null
+++ b/tests/test_reme_memory_error_handling.py
@@ -0,0 +1,132 @@
+"""Tests for ReMe memory error handling and raise_exception propagation."""
+
+import pytest
+
+import reme.reme as reme_module
+from reme import ReMe
+
+
+class Recorder:
+ """Stub that records constructor args for later inspection."""
+
+ instances = []
+
+ def __init__(self, *args, **kwargs):
+ self.args = args
+ self.kwargs = kwargs
+ self.call_kwargs = None
+ Recorder.instances.append(self)
+
+
+class TopLevelAgent(Recorder):
+ """Stub agent that returns a successful structured result."""
+
+ async def call(self, **kwargs):
+ """Simulate a successful agent call."""
+ self.call_kwargs = kwargs
+ return {"answer": "ok", "success": True}
+
+
+def _make_reme() -> ReMe:
+ """Create a ReMe instance with startup bypassed for unit testing."""
+ reme = ReMe(enable_logo=False, log_to_console=False, enable_profile=False)
+ reme._started = True # pylint: disable=protected-access
+ return reme
+
+
+def _patch_summarize_dependencies(monkeypatch: pytest.MonkeyPatch) -> None:
+ """Patch all summarize-path dependencies with stubs."""
+ Recorder.instances = []
+ monkeypatch.setattr(reme_module, "AddDraftAndRetrieveSimilarMemory", Recorder)
+ monkeypatch.setattr(reme_module, "AddMemory", Recorder)
+ monkeypatch.setattr(reme_module, "AddHistory", Recorder)
+ monkeypatch.setattr(reme_module, "DelegateTask", Recorder)
+ monkeypatch.setattr(reme_module, "PersonalSummarizer", Recorder)
+ monkeypatch.setattr(reme_module, "ProceduralSummarizer", Recorder)
+ monkeypatch.setattr(reme_module, "ToolSummarizer", Recorder)
+ monkeypatch.setattr(reme_module, "ReMeSummarizer", TopLevelAgent)
+
+
+def _patch_retrieve_dependencies(monkeypatch: pytest.MonkeyPatch) -> None:
+ """Patch all retrieve-path dependencies with stubs."""
+ Recorder.instances = []
+ monkeypatch.setattr(reme_module, "RetrieveMemory", Recorder)
+ monkeypatch.setattr(reme_module, "ReadHistory", Recorder)
+ monkeypatch.setattr(reme_module, "DelegateTask", Recorder)
+ monkeypatch.setattr(reme_module, "PersonalRetriever", Recorder)
+ monkeypatch.setattr(reme_module, "ProceduralRetriever", Recorder)
+ monkeypatch.setattr(reme_module, "ToolRetriever", Recorder)
+ monkeypatch.setattr(reme_module, "ReMeRetriever", TopLevelAgent)
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("raise_exception", [False, True])
+async def test_summarize_memory_propagates_raise_exception(
+ monkeypatch: pytest.MonkeyPatch,
+ raise_exception: bool,
+):
+ """Verify raise_exception is forwarded to every sub-agent in summarize."""
+ _patch_summarize_dependencies(monkeypatch)
+ reme = _make_reme()
+
+ result = await reme.summarize_memory(
+ messages=[{"role": "user", "content": "hi", "time_created": "2026-03-20 10:00:00"}],
+ task_name="demo-task",
+ raise_exception=raise_exception,
+ )
+
+ assert result == "ok"
+ assert Recorder.instances
+ assert all(instance.kwargs.get("raise_exception") is raise_exception for instance in Recorder.instances)
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("raise_exception", [False, True])
+async def test_retrieve_memory_propagates_raise_exception(
+ monkeypatch: pytest.MonkeyPatch,
+ raise_exception: bool,
+):
+ """Verify raise_exception is forwarded to every sub-agent in retrieve."""
+ _patch_retrieve_dependencies(monkeypatch)
+ reme = _make_reme()
+
+ result = await reme.retrieve_memory(
+ query="hello",
+ task_name="demo-task",
+ raise_exception=raise_exception,
+ )
+
+ assert result == "ok"
+ assert Recorder.instances
+ assert all(instance.kwargs.get("raise_exception") is raise_exception for instance in Recorder.instances)
+
+
+@pytest.mark.asyncio
+async def test_summarize_memory_raises_runtime_error_for_unstructured_result(monkeypatch: pytest.MonkeyPatch):
+ """Verify RuntimeError is raised when the top-level summarizer returns a plain string."""
+ Recorder.instances = []
+ reme = _make_reme()
+
+ monkeypatch.setattr(reme_module, "PersonalSummarizer", Recorder)
+ monkeypatch.setattr(reme_module, "ProceduralSummarizer", Recorder)
+ monkeypatch.setattr(reme_module, "ToolSummarizer", Recorder)
+ monkeypatch.setattr(reme_module, "AddDraftAndRetrieveSimilarMemory", Recorder)
+ monkeypatch.setattr(reme_module, "AddMemory", Recorder)
+ monkeypatch.setattr(reme_module, "AddHistory", Recorder)
+ monkeypatch.setattr(reme_module, "DelegateTask", Recorder)
+
+ class FailingTopLevelAgent(Recorder):
+ """Stub agent that returns a failure string instead of a dict."""
+
+ async def call(self, **kwargs):
+ """Simulate a failed agent call returning a plain error string."""
+ self.call_kwargs = kwargs
+ return "[ReMeSummarizer] failed: boom"
+
+ monkeypatch.setattr(reme_module, "ReMeSummarizer", FailingTopLevelAgent)
+
+ with pytest.raises(RuntimeError, match="summarize_memory failed before producing a structured result"):
+ await reme.summarize_memory(
+ messages=[{"role": "user", "content": "hi", "time_created": "2026-03-20 10:00:00"}],
+ task_name="demo-task",
+ )
diff --git a/tests/test_vector_store.py b/tests/test_vector_store.py
index d666987d..4f3594d3 100644
--- a/tests/test_vector_store.py
+++ b/tests/test_vector_store.py
@@ -2,7 +2,8 @@
"""Unified test suite for vector store implementations.
This module provides comprehensive test coverage for LocalVectorStore, ESVectorStore,
-PGVectorStore, QdrantVectorStore, ChromaVectorStore, and SeekdbVectorStore implementations.
+PGVectorStore, QdrantVectorStore, ChromaVectorStore, ObVecVectorStore, and SeekdbVectorStore
+implementations.
Tests can be run for specific vector stores or all implementations.
Usage:
@@ -11,13 +12,14 @@ Usage:
python test_vector_store.py --pgvector # Test PGVectorStore only
python test_vector_store.py --qdrant # Test QdrantVectorStore only
python test_vector_store.py --chroma # Test ChromaVectorStore only
+ python test_vector_store.py --obvec # Test ObVecVectorStore only (needs seekdb / OceanBase)
python test_vector_store.py --seekdb # Test SeekdbVectorStore only (requires pyseekdb)
python test_vector_store.py --all # Test all vector stores
-
"""
import argparse
import asyncio
+import os
import shutil
import tempfile
from pathlib import Path
@@ -33,6 +35,7 @@ from reme.core.vector_store import (
ChromaVectorStore,
LocalVectorStore,
ESVectorStore,
+ ObVecVectorStore,
PGVectorStore,
QdrantVectorStore,
)
@@ -45,6 +48,12 @@ except ImportError:
load_env()
+
+def _search_score_for_log(metadata: dict) -> object:
+ """Similarity score for log lines (implementations use ``metadata['score']``)."""
+ return metadata.get("score", metadata.get("_score", "N/A"))
+
+
# ==================== Configuration ====================
@@ -83,6 +92,15 @@ class TestConfig:
CHROMA_TENANT = None # Set for ChromaDB Cloud tenant
CHROMA_DATABASE = None # Set for ChromaDB Cloud database
+ # ObVecVectorStore: seekdb docker often uses user `root` + ROOT_PASSWORD; OceanBase
+ # multi-tenant commonly uses `root@` (see pyobvector defaults).
+ # OBVEC_PASSWORD default `root` matches docker-compose.obvec.yml only—override if your
+ # seekdb uses another ROOT_PASSWORD (e.g. another compose stack on the same port).
+ OBVEC_URI = os.environ.get("OBVEC_URI", "127.0.0.1:2881")
+ OBVEC_USER = os.environ.get("OBVEC_USER", "root")
+ OBVEC_PASSWORD = os.environ.get("OBVEC_PASSWORD", "root")
+ OBVEC_DATABASE = os.environ.get("OBVEC_DATABASE", "test")
+
# Embedding model settings
EMBEDDING_MODEL_NAME = "text-embedding-v4"
EMBEDDING_DIMENSIONS = 64
@@ -192,7 +210,7 @@ def get_store_type(store: BaseVectorStore) -> str:
store: Vector store instance
Returns:
- str: Type identifier ("local", "es", "pgvector", "qdrant", "chroma", or "seekdb")
+ str: Type identifier ("local", "es", "pgvector", "qdrant", "chroma", "obvec", or "seekdb")
"""
if isinstance(store, LocalVectorStore):
return "local"
@@ -204,6 +222,8 @@ def get_store_type(store: BaseVectorStore) -> str:
return "pgvector"
elif isinstance(store, ChromaVectorStore):
return "chroma"
+ elif isinstance(store, ObVecVectorStore):
+ return "obvec"
elif SeekdbVectorStore is not None and isinstance(store, SeekdbVectorStore):
return "seekdb"
else:
@@ -214,7 +234,9 @@ def create_vector_store(store_type: str, collection_name: str) -> BaseVectorStor
"""Create a vector store instance based on type.
Args:
- store_type: Type of vector store ("local", "es", "pgvector", "qdrant", "chroma", or "seekdb")
+ store_type: Type of vector store (
+ "local", "es", "pgvector", "qdrant", "chroma", "obvec", or "seekdb"
+ )
collection_name: Name of the collection
Returns:
@@ -277,6 +299,18 @@ def create_vector_store(store_type: str, collection_name: str) -> BaseVectorStor
tenant=config.CHROMA_TENANT,
database=config.CHROMA_DATABASE,
)
+ elif store_type == "obvec":
+ return ObVecVectorStore(
+ collection_name=collection_name,
+ embedding_model=embedding_model,
+ db_path=tempfile.mkdtemp(prefix="test_obvec_"),
+ uri=config.OBVEC_URI,
+ user=config.OBVEC_USER,
+ password=config.OBVEC_PASSWORD,
+ database=config.OBVEC_DATABASE,
+ index_metric="cosine",
+ index_ef_search=100,
+ )
elif store_type == "seekdb":
if SeekdbVectorStore is None:
raise ImportError("SeekdbVectorStore not available; install pyseekdb")
@@ -350,7 +384,7 @@ async def test_search(store: BaseVectorStore, _store_name: str):
logger.info(f"Search returned {len(results)} results")
for i, r in enumerate(results, 1):
- score = r.metadata.get("score", r.metadata.get("_score", "N/A"))
+ score = _search_score_for_log(r.metadata)
logger.info(f" Result {i}: {r.content[:60]}... (score: {score})")
assert len(results) > 0, "Search should return results"
@@ -604,9 +638,9 @@ async def test_copy_collection(store: BaseVectorStore, store_name: str):
config = TestConfig()
copy_collection_name = f"{config.TEST_COLLECTION_PREFIX}_{store_name}_copy"
- # Elasticsearch and PostgreSQL require lowercase table/index names
+ # Elasticsearch, PostgreSQL and OceanBase require lowercase table/index names
store_type = get_store_type(store)
- if store_type in ("es", "pgvector"):
+ if store_type in ("es", "pgvector", "obvec"):
copy_collection_name = copy_collection_name.lower()
# seekdb uses collection names as-is
@@ -1034,7 +1068,7 @@ async def test_search_relevance_ranking(store: BaseVectorStore, _store_name: str
logger.info(f"Search results for: '{query}'")
for i, result in enumerate(results, 1):
- score = result.metadata.get("score", result.metadata.get("_score", "N/A"))
+ score = _search_score_for_log(result.metadata)
relevance = result.metadata.get("relevance", "unknown")
logger.info(f" {i}. [{relevance}] score={score}: {result.content[:60]}...")
@@ -1052,7 +1086,7 @@ async def test_search_relevance_ranking(store: BaseVectorStore, _store_name: str
results2 = await store.search(query=query2, limit=5)
logger.info(f"\nSearch results for: '{query2}'")
for i, result in enumerate(results2, 1):
- score = result.metadata.get("score", result.metadata.get("_score", "N/A"))
+ score = _search_score_for_log(result.metadata)
logger.info(f" {i}. score={score}: {result.content[:60]}...")
logger.info("✓ Search relevance ranking test passed")
@@ -1742,7 +1776,7 @@ async def cleanup_store(store: BaseVectorStore, store_type: str):
Args:
store: Vector store instance
- store_type: Type of vector store ("local" or "es")
+ store_type: Backend key (e.g. ``"local"``, ``"obvec"``)
"""
logger.info("=" * 20 + " CLEANUP " + "=" * 20)
@@ -1776,6 +1810,13 @@ async def cleanup_store(store: BaseVectorStore, store_type: str):
shutil.rmtree(test_dir)
logger.info(f"Cleaned up chroma directory: {config.CHROMA_PATH}")
+ # ObVecVectorStore uses a temp db_path per run (reserved for local sidecar files).
+ if store_type == "obvec":
+ obvec_dir = getattr(store, "db_path", None)
+ if obvec_dir and Path(obvec_dir).exists():
+ shutil.rmtree(obvec_dir, ignore_errors=True)
+ logger.info(f"Cleaned up obvec temp directory: {obvec_dir}")
+
# Clean up temp directory if SeekdbVectorStore used temp dir
if store_type == "seekdb" and config.SEEKDB_PATH is None and getattr(store, "db_path", None):
test_dir = Path(store.db_path)
@@ -1803,6 +1844,7 @@ Examples:
python test_vector_store.py --pgvector # Test PGVectorStore only
python test_vector_store.py --qdrant # Test QdrantVectorStore only
python test_vector_store.py --chroma # Test ChromaVectorStore only
+ python test_vector_store.py --obvec # Test ObVecVectorStore (seekdb / OceanBase)
python test_vector_store.py --seekdb # Test SeekdbVectorStore only
python test_vector_store.py --all # Test all vector stores
""",
@@ -1832,6 +1874,11 @@ Examples:
action="store_true",
help="Test ChromaVectorStore",
)
+ parser.add_argument(
+ "--obvec",
+ action="store_true",
+ help="Test ObVecVectorStore",
+ )
parser.add_argument(
"--seekdb",
action="store_true",
@@ -1855,6 +1902,7 @@ Examples:
("pgvector", "PGVectorStore"),
("qdrant", "QdrantVectorStore"),
("chroma", "ChromaVectorStore"),
+ ("obvec", "ObVecVectorStore"),
]
if SEEKDB_AVAILABLE:
stores_to_test.append(("seekdb", "SeekdbVectorStore"))
@@ -1870,6 +1918,8 @@ Examples:
stores_to_test.append(("qdrant", "QdrantVectorStore"))
if args.chroma:
stores_to_test.append(("chroma", "ChromaVectorStore"))
+ if args.obvec:
+ stores_to_test.append(("obvec", "ObVecVectorStore"))
if args.seekdb:
if not SEEKDB_AVAILABLE:
raise ImportError("seekdb tests require pyseekdb; install with: pip install pyseekdb")
@@ -1883,12 +1933,13 @@ Examples:
("pgvector", "PGVectorStore"),
("qdrant", "QdrantVectorStore"),
("chroma", "ChromaVectorStore"),
+ ("obvec", "ObVecVectorStore"),
]
if SEEKDB_AVAILABLE:
stores_to_test.append(("seekdb", "SeekdbVectorStore"))
print("No vector store specified, defaulting to test all vector stores")
print(
- "Use --local/--es/--pgvector/--qdrant/--chroma/--seekdb to test specific ones\n",
+ "Use --local/--es/--pgvector/--qdrant/--chroma/--obvec/--seekdb to test specific ones\n",
)
# Run tests for each vector store
diff --git a/tests/vector/test_reme_vector.py b/tests/vector/test_reme_vector.py
new file mode 100644
index 00000000..3829c841
--- /dev/null
+++ b/tests/vector/test_reme_vector.py
@@ -0,0 +1,89 @@
+"""测试 ReMe 的 vector 搜索功能"""
+
+import asyncio
+
+from reme import ReMe
+
+
+async def main():
+ """测试 ReMe 的 vector 搜索功能"""
+ # 初始化 ReMe
+ reme = ReMe(
+ working_dir=".reme",
+ default_llm_config={
+ "backend": "openai",
+ "model_name": "qwen3.5-plus",
+ },
+ default_embedding_model_config={
+ "backend": "openai",
+ "model_name": "text-embedding-v4",
+ "dimensions": 1024,
+ },
+ default_vector_store_config={
+ "backend": "local", # 支持 local/chroma/qdrant/elasticsearch
+ },
+ )
+ await reme.start()
+
+ messages = [
+ {"role": "user", "content": "帮我写一个 Python 脚本", "time_created": "2026-02-28 10:00:00"},
+ {"role": "assistant", "content": "好的，我来帮你写", "time_created": "2026-02-28 10:00:05"},
+ ]
+
+ # 1. 从对话中总结记忆（自动提取用户偏好、任务经验等）
+ result = await reme.summarize_memory(
+ messages=messages,
+ user_name="alice", # 个人记忆
+ # task_name="code_writing", # 任务记忆
+ )
+ print(f"总结结果: {result}")
+
+ # 2. 检索相关记忆
+ memories = await reme.retrieve_memory(
+ query="Python 编程",
+ user_name="alice",
+ # task_name="code_writing",
+ )
+ print(f"检索结果: {memories}")
+
+ # 3. 手动添加记忆
+ memory_node = await reme.add_memory(
+ memory_content="用户喜欢简洁的代码风格",
+ user_name="alice",
+ )
+ print(f"添加的记忆: {memory_node}")
+ memory_id = memory_node.memory_id
+
+ # 4. 通过 ID 获取单条记忆
+ fetched_memory = await reme.get_memory(memory_id=memory_id)
+ print(f"获取的记忆: {fetched_memory}")
+
+ # 5. 更新记忆内容
+ updated_memory = await reme.update_memory(
+ memory_id=memory_id,
+ user_name="alice",
+ memory_content="用户喜欢简洁且带注释的代码风格",
+ )
+ print(f"更新后的记忆: {updated_memory}")
+
+ # 6. 列出用户的所有记忆（支持过滤和排序）
+ all_memories = await reme.list_memory(
+ user_name="alice",
+ limit=10,
+ sort_key="time_created",
+ reverse=True,
+ )
+ print(f"用户记忆列表: {all_memories}")
+
+ # 7. 删除指定记忆
+ await reme.delete_memory(memory_id=memory_id)
+ print(f"已删除记忆: {memory_id}")
+
+ # 8. 删除所有记忆（谨慎使用）
+ # await reme.delete_all()
+
+ await reme.close()
+
+
+if __name__ == "__main__":
+ asyncio.run(main())
+ 
+
+
+