From 3347506e229b3ba60a12f11146e95669789c53c9 Mon Sep 17 00:00:00 2001 From: jinliyl <6469360+jinliyl@users.noreply.github.com> Date: Thu, 5 Mar 2026 14:34:00 +0800 Subject: [PATCH 01/59] feat(file-watcher): add configurable retry for file watcher (#140) * feat(file-watcher): enhance file watcher with robust path validation and interruptible sleep * feat(file-watcher): enhance file watcher with robust path validation and interruptible sleep --- reme/core/file_watcher/base_file_watcher.py | 72 +- reme/reme_light.py | 8 +- tests/test_base_file_watcher.py | 791 ++++++++++++++++++++ 3 files changed, 845 insertions(+), 26 deletions(-) create mode 100644 tests/test_base_file_watcher.py diff --git a/reme/core/file_watcher/base_file_watcher.py b/reme/core/file_watcher/base_file_watcher.py index 5030a9cb..89c6e54a 100644 --- a/reme/core/file_watcher/base_file_watcher.py +++ b/reme/core/file_watcher/base_file_watcher.py @@ -140,35 +140,63 @@ 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, + 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/reme_light.py b/reme/reme_light.py index 4396c16f..7c13c72e 100644 --- a/reme/reme_light.py +++ b/reme/reme_light.py @@ -459,7 +459,7 @@ class ReMeLight(Application): """ try: # Initialize summarizer with working directories and configuration - compactor = Summarizer( + summarizer = Summarizer( working_dir=str(self.working_path), memory_dir=str(self.memory_path), memory_compact_threshold=self.memory_compact_threshold, @@ -471,7 +471,7 @@ class ReMeLight(Application): ) # 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 @@ -508,7 +508,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 @@ -562,7 +562,7 @@ class ReMeLight(Application): continue exc = task.exception() if exc is not None: - logger.exception(f"Summary task failed: {exc}") + logger.error(f"Summary task failed: {exc}") else: # Log successful completion with result summary result = task.result() diff --git a/tests/test_base_file_watcher.py b/tests/test_base_file_watcher.py new file mode 100644 index 00000000..ca310eb3 --- /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 +- scan_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 TestScanOnStart: + """Tests for scan_on_start feature.""" + + @pytest.mark.asyncio + async def test_scan_on_start_false(self, temp_files, temp_dir: Path): + """Test that scan_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), + scan_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_scan_on_start_true_with_files(self, temp_files, temp_dir: Path): + """Test that scan_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), + scan_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_scan_on_start_with_suffix_filter(self, temp_files, temp_dir: Path): + """Test scan_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), + scan_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_scan_on_start_recursive(self, temp_nested_dir: Path): + """Test scan_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), + scan_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_scan_on_start_non_recursive(self, temp_nested_dir: Path): + """Test scan_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), + scan_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_scan_on_start_nonexistent_path(self): + """Test scan_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", + scan_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), + scan_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), + scan_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"]) From 3dc3c4bf523c2112d111219587dacac70fce6618 Mon Sep 17 00:00:00 2001 From: "jinli.yl" Date: Thu, 5 Mar 2026 20:27:24 +0800 Subject: [PATCH 02/59] feat(memory): replace memory formatter with AsMsgHandler for enhanced message processing --- reme/core/schema/as_msg_stat.py | 80 ++ reme/memory/file_based/__init__.py | 11 +- reme/memory/file_based/as_msg_handler.py | 351 ++++++ reme/memory/file_based/compactor.py | 14 +- reme/memory/file_based/memory_formatter.py | 249 ---- reme/memory/file_based/reme_chat_formatter.py | 2 +- .../file_based/reme_in_memory_memory.py | 90 +- reme/memory/file_based/summarizer.py | 21 +- reme/reme_light.py | 25 +- tests/light/test_context_check.py | 1090 +++++++++++++++++ tests/light/test_format_msgs_to_str.py | 883 +++++++++++++ 11 files changed, 2461 insertions(+), 355 deletions(-) create mode 100644 reme/core/schema/as_msg_stat.py create mode 100644 reme/memory/file_based/as_msg_handler.py delete mode 100644 reme/memory/file_based/memory_formatter.py create mode 100644 tests/light/test_context_check.py create mode 100644 tests/light/test_format_msgs_to_str.py diff --git a/reme/core/schema/as_msg_stat.py b/reme/core/schema/as_msg_stat.py new file mode 100644 index 00000000..32cc9956 --- /dev/null +++ b/reme/core/schema/as_msg_stat.py @@ -0,0 +1,80 @@ +from pydantic import BaseModel, Field + +_DEFAULT_MAX_BLOCK_TEXT_PREVIEW_LENGTH = 100 +_DEFAULT_MAX_FORMATTER_TEXT_LENGTH = 2000 + +# Unique marker for truncated text +TRUNCATION_MARKER_START = "<<>>" +TRUNCATION_MARKER_END = "<<>>" + + +def _truncate_text(text: str, max_length: int) -> str: + """Truncate text to max length, keeping head and tail portions.""" + text = str(text) if text else "" + if not text or len(text) <= max_length: + return text + half_length = max_length // 2 + truncated_chars = len(text) - max_length + return ( + f"{text[:half_length]}\n\n{TRUNCATION_MARKER_START} " + f"({truncated_chars} characters omitted) " + f"{TRUNCATION_MARKER_END}\n\n{text[-half_length:]}" + ) + + +class AsBlockStat(BaseModel): + 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 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 block content to string representation.""" + if self.block_type == "text": + return _truncate_text(self.text, max_length) if self.text else "" + if self.block_type == "thinking": + if include_thinking and self.text: + return f"\n{_truncate_text(self.text, max_length)}\n" + return "" + if self.block_type in ("image", "audio", "video"): + return f"[{self.block_type}] {self.media_url}" if self.media_url else f"[{self.block_type}]" + if self.block_type == "tool_use": + return f" - tool_call={self.tool_name} params={_truncate_text(self.tool_input, max_length)}" + if self.block_type == "tool_result": + output = _truncate_text(self.tool_output, max_length) + return f" - tool_result={self.tool_name} output={output}" if output else "" + return "" + + +class AsMsgStat(BaseModel): + 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 sum(block.token_count for block in self.content) + + @property + def preview(self) -> str: + 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/memory/file_based/__init__.py b/reme/memory/file_based/__init__.py index 29f33749..d90cf4dd 100644 --- a/reme/memory/file_based/__init__.py +++ b/reme/memory/file_based/__init__.py @@ -4,8 +4,9 @@ This module provides memory management components for CoPaw (Cooperative Paw) ag 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 + - ReMeOpenAIChatFormatter: Converts message lists to formatted strings with token limiting + - AsMsgHandler: Handles AgentScope message statistics, formatting, and context checking - Summarizer: Generates memory summaries using LLM - Compactor: Compacts memory content to reduce token usage - ToolResultCompactor: Truncates large tool results and saves full content to files @@ -13,21 +14,21 @@ Components: """ from . import utils +from .as_msg_handler import AsMsgHandler from .compactor import Compactor from .file_io import FileIO -from .memory_formatter import MemoryFormatter -from .reme_chat_formatter import ReMeChatFormatter +from .reme_chat_formatter import ReMeOpenAIChatFormatter from .reme_in_memory_memory import ReMeInMemoryMemory from .summarizer import Summarizer from .tool_result_compactor import ToolResultCompactor __all__ = [ - "MemoryFormatter", + "AsMsgHandler", "ReMeInMemoryMemory", "Summarizer", "Compactor", "ToolResultCompactor", "FileIO", "utils", - "ReMeChatFormatter", + "ReMeOpenAIChatFormatter", ] diff --git a/reme/memory/file_based/as_msg_handler.py b/reme/memory/file_based/as_msg_handler.py new file mode 100644 index 00000000..26d3d433 --- /dev/null +++ b/reme/memory/file_based/as_msg_handler.py @@ -0,0 +1,351 @@ +import json +import logging + +from agentscope.message import Msg +from agentscope.token import HuggingFaceTokenCounter + +from ...core.schema.as_msg_stat import AsMsgStat, AsBlockStat + +logger = logging.getLogger(__name__) + + +class AsMsgHandler: + + def __init__(self, token_counter: HuggingFaceTokenCounter): + self._token_counter = token_counter + + 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_ids = self._token_counter.tokenizer.encode(text) + token_count = len(token_ids) + return token_count + + except Exception as e: + estimated_tokens = len(text.encode("utf-8")) // 4 + logger.warning(f"Failed to count string tokens: {text}, using estimated_tokens={estimated_tokens}") + return estimated_tokens + + @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: + 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 stat_message(self, message: Msg) -> AsMsgStat: + """Analyze a message and generate block statistics.""" + blocks = [] + + for block in message.get_content_blocks(): + block_type = block.get("type", "unknown") + + if block_type == "text": + text = block.get("text", "") + token_count = 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 = 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", "") + # For media, estimate fixed token cost or count URL + if source.get("type") == "base64": + data = source.get("data", "") + token_count = len(data) // 4 if data else 10 + else: + token_count = 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 = 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 = self._format_tool_result_output(output) + token_count = self.count_str_token(formatted_output) + blocks.append(AsBlockStat( + block_type=block_type, + text="", + token_count=token_count, + tool_name=tool_name, + tool_output=formatted_output, + )) + + else: + logger.warning("Unsupported block type %s, skipped.", block_type) + + return AsMsgStat( + name=message.name or message.role, + role=message.role, + content=blocks, + timestamp=message.timestamp or "", + metadata=message.metadata or {}, + ) + + def format_msgs_to_str( + self, + messages: list[Msg], + memory_compact_threshold: int, + include_thinking: bool = False, + ) -> 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 = self.stat_message(messages[i]) + + if total_token_count + stat.total_tokens > memory_compact_threshold: + logger.info( + "Skipping older messages: adding %d tokens would exceed threshold %d (current: %d)", + stat.total_tokens, + memory_compact_threshold, + total_token_count, + ) + break + + formatted_parts.append(stat.format(include_thinking=include_thinking)) + total_token_count += stat.total_tokens + + formatted_parts.reverse() + return "\n\n".join(formatted_parts) + + def context_check( + self, + messages: list[Msg], + memory_compact_threshold: int, + memory_compact_reserve: int, + ) -> tuple[list[Msg], list[Msg]]: + """Check if context exceeds threshold and split messages accordingly. + + This method checks if the total token count of messages exceeds the + memory_compact_threshold. If not, returns empty list and original messages. + If exceeded, uses memory_compact_reserve as the limit to keep messages + from the end, ensuring tool_use and tool_result blocks are properly paired. + + 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 after compaction. + + Returns: + A tuple of (messages_to_compact, messages_to_keep): + - messages_to_compact: Older messages that need to be compacted + - messages_to_keep: Recent messages within the reserve limit + """ + if not messages: + return [], [] + + # Calculate total tokens and stats for all messages + msg_stats: list[tuple[Msg, AsMsgStat]] = [] + total_tokens = 0 + for msg in messages: + stat = self.stat_message(msg) + msg_stats.append((msg, stat)) + total_tokens += stat.total_tokens + + # If total tokens don't exceed threshold, no compaction needed + if total_tokens <= memory_compact_threshold: + return [], messages + + # 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): + 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( + "Context check: adding message %d with %d tokens would exceed reserve %d (current: %d)", + i, + stat.total_tokens, + memory_compact_reserve, + 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( + "Context check: message %d requires %d extra tokens for tool_use dependencies, " + "total would exceed reserve %d", + i, + extra_tokens, + 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) + + logger.info( + "Context check result: %d messages to compact, %d messages to keep, " + "total tokens: %d, threshold: %d, reserve: %d, kept tokens: %d", + len(messages_to_compact), + len(messages_to_keep), + total_tokens, + memory_compact_threshold, + memory_compact_reserve, + accumulated_tokens, + ) + + return messages_to_compact, messages_to_keep \ No newline at end of file diff --git a/reme/memory/file_based/compactor.py b/reme/memory/file_based/compactor.py index d8186d86..c7dbb496 100644 --- a/reme/memory/file_based/compactor.py +++ b/reme/memory/file_based/compactor.py @@ -8,7 +8,7 @@ from agentscope.message import Msg from agentscope.model import ChatModelBase from agentscope.token import HuggingFaceTokenCounter -from .memory_formatter import MemoryFormatter +from .as_msg_handler import AsMsgHandler from ...core.op import BaseOp logger = logging.getLogger(__name__) @@ -30,7 +30,7 @@ class Compactor(BaseOp): self.chat_model: ChatModelBase = chat_model self.formatter: FormatterBase = formatter - self.as_token_counter: HuggingFaceTokenCounter = token_counter + self.msg_handler = AsMsgHandler(token_counter=token_counter) async def execute(self): messages: list[Msg] = self.context.get("messages", []) @@ -39,11 +39,10 @@ class Compactor(BaseOp): if not messages: return "" - formatter = MemoryFormatter( - token_counter=self.as_token_counter, + history_formatted_str: str = self.msg_handler.format_msgs_to_str( + messages=messages, 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}") @@ -66,9 +65,8 @@ class Compactor(BaseOp): f"{suffix}" ) else: - user_message: str = f"\n{history_formatted_str}\n\n\n" + self.get_prompt( - "initial_user_message", - ) + 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( 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 index b70e700c..f6205688 100644 --- a/reme/memory/file_based/reme_chat_formatter.py +++ b/reme/memory/file_based/reme_chat_formatter.py @@ -8,7 +8,7 @@ from agentscope.token import HuggingFaceTokenCounter from .utils import _extract_text_from_messages -class ReMeChatFormatter(OpenAIChatFormatter): +class ReMeOpenAIChatFormatter(OpenAIChatFormatter): """ReMe chat formatter class.""" async def _count(self, msgs: list[dict[str, Any]]) -> int | None: diff --git a/reme/memory/file_based/reme_in_memory_memory.py b/reme/memory/file_based/reme_in_memory_memory.py index 0e915382..61943a6c 100644 --- a/reme/memory/file_based/reme_in_memory_memory.py +++ b/reme/memory/file_based/reme_in_memory_memory.py @@ -3,12 +3,11 @@ import logging from agentscope.agent._react_agent import _MemoryMark -from agentscope.formatter import FormatterBase 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 .as_msg_handler import AsMsgHandler logger = logging.getLogger(__name__) @@ -19,13 +18,10 @@ class ReMeInMemoryMemory(InMemoryMemory): def __init__( self, token_counter: HuggingFaceTokenCounter, - formatter: FormatterBase, - max_input_length: int = 0, ): 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) async def get_memory( self, @@ -127,9 +123,12 @@ Use it as context to maintain continuity. """Clear the content.""" self.content.clear() - async def estimate_tokens(self) -> dict: + 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,7 +137,7 @@ 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, @@ -146,62 +145,18 @@ Use it as context to maintain continuity. ) compressed_summary = self.get_compressed_summary() - compressed_summary_tokens = safe_count_str_tokens(self._token_counter, compressed_summary) + compressed_summary_tokens = 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 = [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 +167,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 index 462ea9c5..6f9f0b02 100644 --- a/reme/memory/file_based/summarizer.py +++ b/reme/memory/file_based/summarizer.py @@ -10,8 +10,7 @@ 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 .as_msg_handler import AsMsgHandler from ...core.op import BaseOp logger = logging.getLogger(__name__) @@ -28,7 +27,7 @@ class Summarizer(BaseOp): chat_model: ChatModelBase, formatter: FormatterBase, token_counter: HuggingFaceTokenCounter, - toolkit: Toolkit | None = None, + toolkit: Toolkit, **kwargs, ): super().__init__(**kwargs) @@ -38,15 +37,8 @@ class Summarizer(BaseOp): 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) + self.msg_handler = AsMsgHandler(token_counter=token_counter) + self.toolkit: Toolkit = toolkit async def execute(self): messages: list[Msg] = self.context.get("messages", []) @@ -54,11 +46,10 @@ class Summarizer(BaseOp): if not messages: return "" - formatter = MemoryFormatter( - token_counter=self.as_token_counter, + history_formatted_str: str = self.msg_handler.format_msgs_to_str( + messages=messages, 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}") diff --git a/reme/reme_light.py b/reme/reme_light.py index 7c13c72e..8a0bdd9e 100644 --- a/reme/reme_light.py +++ b/reme/reme_light.py @@ -28,7 +28,7 @@ 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 import Compactor, Summarizer, ToolResultCompactor, ReMeInMemoryMemory, ReMeOpenAIChatFormatter, FileIO from .memory.file_based.utils import get_token_counter from .memory.tools import MemorySearch from .core.utils import load_env @@ -95,11 +95,6 @@ class ReMeLight(Application): self.tool_result_path = self.working_path / "tool_result" self.tool_result_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, @@ -198,7 +193,7 @@ class ReMeLight(Application): if formatter is not None: self.formatter: FormatterBase = formatter else: - self.formatter = ReMeChatFormatter(token_counter=self.token_counter) + self.formatter = ReMeOpenAIChatFormatter(token_counter=self.token_counter) self.toolkit: Toolkit | None = toolkit # Initialize list to track background summarization tasks @@ -458,6 +453,16 @@ class ReMeLight(Application): - If summarization fails, an empty string is returned """ try: + # Create toolkit if not provided + if self.toolkit is not None: + toolkit = self.toolkit + else: + toolkit = Toolkit() + file_io = FileIO(working_dir=str(self.working_path)) + toolkit.register_tool_function(file_io.read) + toolkit.register_tool_function(file_io.write) + toolkit.register_tool_function(file_io.edit) + # Initialize summarizer with working directories and configuration summarizer = Summarizer( working_dir=str(self.working_path), @@ -466,7 +471,7 @@ class ReMeLight(Application): chat_model=self.chat_model, formatter=self.formatter, token_counter=self.token_counter, - toolkit=self.toolkit, + toolkit=toolkit, language=self.language, ) @@ -662,10 +667,8 @@ class ReMeLight(Application): 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 + - Shares the same token counter as the main application """ return ReMeInMemoryMemory( token_counter=self.token_counter, - formatter=self.formatter, - max_input_length=self.max_input_length, ) diff --git a/tests/light/test_context_check.py b/tests/light/test_context_check.py new file mode 100644 index 00000000..300f0b61 --- /dev/null +++ b/tests/light/test_context_check.py @@ -0,0 +1,1090 @@ +"""Tests for AsMsgHandler.context_check method.""" + +import logging + +from agentscope.message import Msg + +from test_utils import get_token_counter +from reme.memory.file_based.as_msg_handler import AsMsgHandler + +# Configure logging +logging.basicConfig( + level=logging.INFO, + format="%(asctime)s - %(name)s - %(levelname)s - %(message)s", +) +logger = logging.getLogger(__name__) + + +# 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(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(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 = handler.context_check( + messages=messages, + memory_compact_threshold=threshold, + memory_compact_reserve=reserve, + ) + assert 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 = handler.context_check( + messages=messages, + memory_compact_threshold=threshold, # Very high threshold + memory_compact_reserve=reserve, + ) + assert 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 = 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 = 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 = handler.context_check( + messages=messages, + memory_compact_threshold=threshold, + memory_compact_reserve=reserve, + ) + assert 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 = 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 = 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 = 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 = handler.stat_message(messages[0]) + exact_tokens = stat.total_tokens + threshold, reserve = exact_tokens, exact_tokens + + # Test at exact boundary + to_compact, to_keep = handler.context_check( + messages=messages, + memory_compact_threshold=threshold, # Exactly at boundary + memory_compact_reserve=reserve, + ) + # At exact boundary (<=), should not trigger compaction + assert 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 = 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 = 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 = 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 = 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 = 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 = 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 = 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 = 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 = 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 = 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 = 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 = 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 = 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 = 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(handler.stat_message(m).total_tokens for m in messages) + threshold, reserve = total - 1, total + + to_compact, to_keep = 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 = handler.stat_message(messages[1]).total_tokens + tiny_msg_tokens = handler.stat_message(messages[2]).total_tokens + threshold, reserve = 50, small_msg_tokens + tiny_msg_tokens + 10 + + to_compact, to_keep = 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 = handler.stat_message(messages[2]).total_tokens + threshold, reserve = 10, tiny_tokens + 5 + + to_compact, to_keep = 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 = 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 = 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 = 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 = 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 = 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 = 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 = 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..29e1b2cb --- /dev/null +++ b/tests/light/test_format_msgs_to_str.py @@ -0,0 +1,883 @@ +"""Tests for AsMsgHandler.format_msgs_to_str method.""" + +# pylint: disable=W0212 + +import logging + +from agentscope.message import Msg + +from test_utils import get_token_counter +from reme.memory.file_based.as_msg_handler import AsMsgHandler + +# 配置日志输出到控制台 +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}") + + +# ==================== 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 = handler.stat_message(msg) + # Check if this message's content appears in the result + formatted = 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 + elif block_type == "tool_use" and f"tool_call={block.get('name', '')}" in result: + msg_included = True + break + elif block_type == "tool_result" and f"tool_result={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 = 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 = 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 = 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 = 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 = handler.format_msgs_to_str(msgs, memory_compact_threshold=threshold) + + assert "tool_call=read_file" in result, f"Expected tool_call 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 = handler.format_msgs_to_str(msgs, memory_compact_threshold=threshold) + + assert "tool_result=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 = handler.format_msgs_to_str(msgs, memory_compact_threshold=threshold) + + assert "[image]" in result, f"Expected '[image]' 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 = handler.format_msgs_to_str(msgs, memory_compact_threshold=threshold) + + assert "user:" in result + assert "tool_call=read_file" in result + assert "tool_result=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 = 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 = 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 = handler.format_msgs_to_str( + msgs, memory_compact_threshold=threshold, include_thinking=False + ) + # With include_thinking=True + result_with_thinking = 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 = 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 = 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 = 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 = 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 - no messages should be included.""" + handler = create_handler() + threshold = 0 + msgs = [create_user_msg("Test message")] + + result = handler.format_msgs_to_str(msgs, memory_compact_threshold=threshold) + + assert result == "", f"Expected empty string with zero threshold, got: {result}" + verify_result_within_threshold(handler, result, threshold, "threshold_zero", msgs) + 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 tokens + msg = create_user_msg("Test") + stat = handler.stat_message(msg) + exact_threshold = stat.total_tokens + + msgs = [msg] + result = 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.""" + handler = create_handler() + msg = create_user_msg("Test message") + stat = handler.stat_message(msg) + threshold_minus_one = stat.total_tokens - 1 + + msgs = [msg] + result = handler.format_msgs_to_str(msgs, memory_compact_threshold=threshold_minus_one) + + # Message should be skipped since it doesn't fit + assert result == "", f"Expected empty string when threshold is insufficient, got: {result}" + verify_result_within_threshold(handler, result, threshold_minus_one, "threshold_one_less", msgs) + 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 = 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 = 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 = 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 = 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 = 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 = 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 = handler.format_msgs_to_str( + msgs, memory_compact_threshold=threshold, include_thinking=False + ) + result_with_thinking = handler.format_msgs_to_str( + msgs, memory_compact_threshold=threshold, include_thinking=True + ) + + assert "Text content" in result_no_thinking + assert "tool_call=test_tool" in result_no_thinking + assert "[image]" 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 = 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 = handler.format_msgs_to_str(msgs, memory_compact_threshold=threshold) + + assert "tool_result=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 = 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 + total_tokens = sum(handler.stat_message(msg).total_tokens for msg in msgs) + + # Use threshold that allows about half the messages + half_threshold = total_tokens // 2 + result = 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.""" + handler = create_handler() + threshold = -1 + msgs = [create_user_msg("Test message")] + + result = handler.format_msgs_to_str(msgs, memory_compact_threshold=threshold) + + # Negative threshold should result in empty string (nothing fits) + assert result == "", f"Expected empty string with negative threshold, got: {result}" + verify_result_within_threshold(handler, result, max(0, threshold), "negative_threshold", msgs) + 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 = 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 = handler.format_msgs_to_str(msgs, memory_compact_threshold=threshold) + + assert "[image]" 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 = handler.format_msgs_to_str(msgs, memory_compact_threshold=threshold) + + assert "[audio]" in result + assert "[video]" 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 = 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() + exit(0 if success else 1) From 57f8a7b42c94b71564450d1291587f9f2a9f11fc Mon Sep 17 00:00:00 2001 From: "jinli.yl" Date: Fri, 6 Mar 2026 01:33:48 +0800 Subject: [PATCH 03/59] feat(core): integrate AgentScope LLM support with enhanced memory management --- reme/config/light.yaml | 17 + reme/core/__init__.py | 5 +- reme/core/application.py | 73 ++-- reme/core/as_llm/__init__.py | 7 + reme/core/as_llm_formatter/__init__.py | 7 + reme/core/op/base_op.py | 21 +- reme/core/registry_factory.py | 2 + reme/core/schema/__init__.py | 3 + reme/core/schema/as_msg_stat.py | 28 +- reme/core/schema/service_config.py | 81 ++-- reme/core/service_context.py | 10 + reme/core/utils/__init__.py | 7 + reme/core/utils/hf_token_counter_utils.py | 23 ++ reme/core/utils/std_logger.py | 109 +++++ reme/core/utils/truncate_text_utils.py | 53 +++ reme/memory/file_based/__init__.py | 14 +- reme/memory/file_based/as_msg_handler.py | 16 +- reme/memory/file_based/reme_chat_formatter.py | 29 -- .../file_based/reme_in_memory_memory.py | 34 +- reme/memory/file_based/sub_agent/__init__.py | 0 .../file_based/{ => sub_agent}/compactor.py | 29 +- .../file_based/{ => sub_agent}/compactor.yaml | 0 .../file_based/{ => sub_agent}/summarizer.py | 32 +- .../{ => sub_agent}/summarizer.yaml | 0 .../{ => sub_agent}/tool_result_compactor.py | 18 +- reme/memory/file_based/utils.py | 271 ------------- reme/memory/tools/__init__.py | 4 - reme/memory/tools/file/__init__.py | 0 .../{file_based => tools/file}/file_io.py | 0 reme/reme_light.py | 382 +++--------------- 30 files changed, 467 insertions(+), 808 deletions(-) create mode 100644 reme/core/as_llm/__init__.py create mode 100644 reme/core/as_llm_formatter/__init__.py create mode 100644 reme/core/utils/hf_token_counter_utils.py create mode 100644 reme/core/utils/std_logger.py create mode 100644 reme/core/utils/truncate_text_utils.py delete mode 100644 reme/memory/file_based/reme_chat_formatter.py create mode 100644 reme/memory/file_based/sub_agent/__init__.py rename reme/memory/file_based/{ => sub_agent}/compactor.py (77%) rename reme/memory/file_based/{ => sub_agent}/compactor.yaml (100%) rename reme/memory/file_based/{ => sub_agent}/summarizer.py (74%) rename reme/memory/file_based/{ => sub_agent}/summarizer.yaml (100%) rename reme/memory/file_based/{ => sub_agent}/tool_result_compactor.py (91%) delete mode 100644 reme/memory/file_based/utils.py create mode 100644 reme/memory/tools/file/__init__.py rename reme/memory/{file_based => tools/file}/file_io.py (100%) diff --git a/reme/config/light.yaml b/reme/config/light.yaml index 2a83371a..89df5f4b 100644 --- a/reme/config/light.yaml +++ b/reme/config/light.yaml @@ -1,11 +1,28 @@ +as_llms: + default: + backend: openai + model_name: qwen3.5-plus + +as_llm_formatters: + default: + backend: openai + 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: diff --git a/reme/core/__init__.py b/reme/core/__init__.py index 5872e2ad..88e42175 100644 --- a/reme/core/__init__.py +++ b/reme/core/__init__.py @@ -1,5 +1,6 @@ """Core""" - +from . import as_llm +from . import as_llm_formatter from . import embedding from . import enumeration from . import file_store @@ -21,6 +22,8 @@ from .service_context import ServiceContext __all__ = [ # Submodules + "as_llm", + "as_llm_formatter", "embedding", "enumeration", "file_watcher", diff --git a/reme/core/application.py b/reme/core/application.py index 49537807..bc59c7f3 100644 --- a/reme/core/application.py +++ b/reme/core/application.py @@ -1,6 +1,7 @@ """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 @@ -24,24 +25,26 @@ class Application: """Application wrapper that wires together service context, flows, and runtimes.""" 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, - working_dir: str | None = None, - config_path: str | None = None, - enable_logo: bool = True, - log_to_console: bool = True, - parser: type[PydanticConfigParser] | None = None, - default_llm_config: dict | None = None, - default_embedding_model_config: dict | None = None, - default_vector_store_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, - **kwargs, + 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, + working_dir: str | None = None, + config_path: str | None = None, + enable_logo: bool = True, + log_to_console: 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, + default_file_store_config: dict | None = None, + default_token_counter_config: dict | None = None, + default_file_watcher_config: dict | None = None, + **kwargs, ): self.service_context = ServiceContext( *args, @@ -55,6 +58,8 @@ class Application: 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, @@ -137,8 +142,8 @@ class Application: ray.init(num_cpus=self.service_config.ray_max_workers) if ( - self.service_context.thread_pool is None - or self.service_context.thread_pool._shutdown # pylint: disable=protected-access + 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, @@ -147,6 +152,26 @@ class Application: 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: + config_dict = config.model_dump(exclude={"backend"}) + if not config_dict.get("api_key", ""): + config_dict["api_key"] = os.getenv("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"] = os.getenv("LLM_BASE_URL", "") + self.service_context.as_llms[name] = R.as_llms[config.backend](**config_dict) + + 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: + config_dict = config.model_dump(exclude={"backend"}) + self.service_context.as_llm_formatters[name] = R.as_llm_formatters[config.backend](**config_dict) + 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.") @@ -294,10 +319,10 @@ class Application: stream_queue = asyncio.Queue() task = asyncio.create_task(flow.call(stream_queue=stream_queue, **kwargs)) async for chunk in execute_stream_task( - stream_queue=stream_queue, - task=task, - task_name=name, - output_format="str", + stream_queue=stream_queue, + task=task, + task_name=name, + output_format="str", ): yield chunk diff --git a/reme/core/as_llm/__init__.py b/reme/core/as_llm/__init__.py new file mode 100644 index 00000000..888048e6 --- /dev/null +++ b/reme/core/as_llm/__init__.py @@ -0,0 +1,7 @@ +from agentscope.model import DashScopeChatModel +from agentscope.model import OpenAIChatModel + +from ..registry_factory import R + +R.as_llms.register(OpenAIChatModel, "openai") +R.as_llms.register(DashScopeChatModel, "dashscope") diff --git a/reme/core/as_llm_formatter/__init__.py b/reme/core/as_llm_formatter/__init__.py new file mode 100644 index 00000000..9c3a52cf --- /dev/null +++ b/reme/core/as_llm_formatter/__init__.py @@ -0,0 +1,7 @@ +from agentscope.formatter import DashScopeChatFormatter +from agentscope.formatter import OpenAIChatFormatter + +from ..registry_factory import R + +R.as_llm_formatters.register(OpenAIChatFormatter, "openai") +R.as_llm_formatters.register(DashScopeChatFormatter, "dashscope") diff --git a/reme/core/op/base_op.py b/reme/core/op/base_op.py index af5158c4..0f0580a8 100644 --- a/reme/core/op/base_op.py +++ b/reme/core/op/base_op.py @@ -21,7 +21,8 @@ from ..service_context import ServiceContext from ..token_counter import BaseTokenCounter from ..utils import camel_to_snake, CacheHandler, timer from ..vector_store import BaseVectorStore - +from agentscope.model import ChatModelBase +from agentscope.formatter import FormatterBase class BaseOp(metaclass=ABCMeta): """Base operator class for LLM workflow execution and composition.""" @@ -42,6 +43,8 @@ class BaseOp(metaclass=ABCMeta): language: str = "", prompt_name: str = "", prompt_path: str = "", + as_llm: str | ChatModelBase = "default", + as_llm_formatter: str | FormatterBase = "default", llm: str | BaseLLM = "default", embedding_model: str | BaseEmbeddingModel = "default", vector_store: str | BaseVectorStore = "default", @@ -64,6 +67,8 @@ 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._llm = llm self._embedding_model = embedding_model self._vector_store = vector_store @@ -129,6 +134,20 @@ 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 llm(self) -> BaseLLM: """Get the LLM instance from ServiceContext.""" diff --git a/reme/core/registry_factory.py b/reme/core/registry_factory.py index b319b049..f54ad3c1 100644 --- a/reme/core/registry_factory.py +++ b/reme/core/registry_factory.py @@ -34,6 +34,8 @@ class RegistryFactory: def __init__(self): self.llms = Registry() + self.as_llms = Registry() + self.as_llm_formatters = 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 index 32cc9956..1acb9e8a 100644 --- a/reme/core/schema/as_msg_stat.py +++ b/reme/core/schema/as_msg_stat.py @@ -3,24 +3,6 @@ from pydantic import BaseModel, Field _DEFAULT_MAX_BLOCK_TEXT_PREVIEW_LENGTH = 100 _DEFAULT_MAX_FORMATTER_TEXT_LENGTH = 2000 -# Unique marker for truncated text -TRUNCATION_MARKER_START = "<<>>" -TRUNCATION_MARKER_END = "<<>>" - - -def _truncate_text(text: str, max_length: int) -> str: - """Truncate text to max length, keeping head and tail portions.""" - text = str(text) if text else "" - if not text or len(text) <= max_length: - return text - half_length = max_length // 2 - truncated_chars = len(text) - max_length - return ( - f"{text[:half_length]}\n\n{TRUNCATION_MARKER_START} " - f"({truncated_chars} characters omitted) " - f"{TRUNCATION_MARKER_END}\n\n{text[-half_length:]}" - ) - class AsBlockStat(BaseModel): block_type: str = Field(default=...) @@ -41,18 +23,20 @@ class AsBlockStat(BaseModel): def format(self, max_length: int = _DEFAULT_MAX_FORMATTER_TEXT_LENGTH, include_thinking: bool = True) -> str: """Format block content to string representation.""" + from ..utils import truncate_text + if self.block_type == "text": - return _truncate_text(self.text, max_length) if self.text else "" + return truncate_text(self.text, max_length) if self.text else "" if self.block_type == "thinking": if include_thinking and self.text: - return f"\n{_truncate_text(self.text, max_length)}\n" + return f"\n{truncate_text(self.text, max_length)}\n" return "" if self.block_type in ("image", "audio", "video"): return f"[{self.block_type}] {self.media_url}" if self.media_url else f"[{self.block_type}]" if self.block_type == "tool_use": - return f" - tool_call={self.tool_name} params={_truncate_text(self.tool_input, max_length)}" + return f" - tool_call={self.tool_name} params={truncate_text(self.tool_input, max_length)}" if self.block_type == "tool_result": - output = _truncate_text(self.tool_output, max_length) + output = truncate_text(self.tool_output, max_length) return f" - tool_result={self.tool_name} output={output}" if output else "" return "" diff --git a/reme/core/schema/service_config.py b/reme/core/schema/service_config.py index 5b89367a..5e4cd212 100644 --- a/reme/core/schema/service_config.py +++ b/reme/core/schema/service_config.py @@ -58,69 +58,60 @@ 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) @@ -137,6 +128,8 @@ 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) 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..92d0d566 100644 --- a/reme/core/service_context.py +++ b/reme/core/service_context.py @@ -11,6 +11,8 @@ from .schema import ServiceConfig from .utils import load_env, PydanticConfigParser if TYPE_CHECKING: + from agentscope.model import ChatModelBase + from agentscope.formatter import FormatterBase from .llm import BaseLLM from .embedding import BaseEmbeddingModel from .vector_store import BaseVectorStore @@ -36,6 +38,8 @@ class ServiceContext(BaseDict): 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_llm_config: dict | None = None, default_embedding_model_config: dict | None = None, default_vector_store_config: dict | None = None, @@ -64,6 +68,10 @@ 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_llm_config: self._update_section_config(kwargs, "llms", **default_llm_config) if default_embedding_model_config: @@ -90,6 +98,8 @@ 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.llms: dict[str, "BaseLLM"] = {} self.embedding_models: dict[str, "BaseEmbeddingModel"] = {} self.token_counters: dict[str, "BaseTokenCounter"] = {} diff --git a/reme/core/utils/__init__.py b/reme/core/utils/__init__.py index b43784c8..c1f35adb 100644 --- a/reme/core/utils/__init__.py +++ b/reme/core/utils/__init__.py @@ -11,12 +11,15 @@ 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 as get_std_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 +from .truncate_text_utils import truncate_text, is_truncated __all__ = [ "convert_dashscope_to_agentscope", @@ -39,6 +42,7 @@ __all__ = [ "format_messages", "deduplicate_memories", "init_logger", + "get_std_logger", "print_logo", "MCPClient", "PydanticConfigParser", @@ -46,4 +50,7 @@ __all__ = [ "singleton", "timer", "get_now_time", + "get_hf_token_counter", + "truncate_text", + "is_truncated", ] 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..dfbc3dc6 --- /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/std_logger.py b/reme/core/utils/std_logger.py new file mode 100644 index 00000000..e2de0e9c --- /dev/null +++ b/reme/core/utils/std_logger.py @@ -0,0 +1,109 @@ +"""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_logger( + 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, +) -> 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. + + Returns: + Configured Logger instance. + """ + # Return existing logger if already created + if name in _loggers: + 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: + 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) + + # 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 diff --git a/reme/core/utils/truncate_text_utils.py b/reme/core/utils/truncate_text_utils.py new file mode 100644 index 00000000..ec85ec61 --- /dev/null +++ b/reme/core/utils/truncate_text_utils.py @@ -0,0 +1,53 @@ +from .std_logger import get_logger + +logger = get_logger() + +TRUNCATION_MARKER_START = "<<>>" +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 diff --git a/reme/memory/file_based/__init__.py b/reme/memory/file_based/__init__.py index d90cf4dd..a1f5be73 100644 --- a/reme/memory/file_based/__init__.py +++ b/reme/memory/file_based/__init__.py @@ -5,22 +5,17 @@ including memory formatting, compaction, summarization, and file I/O operations. Components: - ReMeInMemoryMemory: Extended InMemoryMemory with bugfixes and summary support - - ReMeOpenAIChatFormatter: Converts message lists to formatted strings with token limiting - AsMsgHandler: Handles AgentScope message statistics, formatting, and context checking - Summarizer: Generates memory summaries using LLM - Compactor: Compacts memory content to reduce token usage - ToolResultCompactor: Truncates large tool results and saves full content to files - - FileIO: File I/O operations with configurable working directory """ -from . import utils from .as_msg_handler import AsMsgHandler -from .compactor import Compactor -from .file_io import FileIO -from .reme_chat_formatter import ReMeOpenAIChatFormatter from .reme_in_memory_memory import ReMeInMemoryMemory -from .summarizer import Summarizer -from .tool_result_compactor import ToolResultCompactor +from .sub_agent.compactor import Compactor +from .sub_agent.summarizer import Summarizer +from .sub_agent.tool_result_compactor import ToolResultCompactor __all__ = [ "AsMsgHandler", @@ -28,7 +23,4 @@ __all__ = [ "Summarizer", "Compactor", "ToolResultCompactor", - "FileIO", - "utils", - "ReMeOpenAIChatFormatter", ] diff --git a/reme/memory/file_based/as_msg_handler.py b/reme/memory/file_based/as_msg_handler.py index 26d3d433..e967f81e 100644 --- a/reme/memory/file_based/as_msg_handler.py +++ b/reme/memory/file_based/as_msg_handler.py @@ -1,12 +1,12 @@ import json -import logging from agentscope.message import Msg from agentscope.token import HuggingFaceTokenCounter -from ...core.schema.as_msg_stat import AsMsgStat, AsBlockStat +from ...core.schema import AsMsgStat, AsBlockStat +from ...core.utils import get_std_logger -logger = logging.getLogger(__name__) +logger = get_std_logger() class AsMsgHandler: @@ -179,10 +179,10 @@ class AsMsgHandler: ) def format_msgs_to_str( - self, - messages: list[Msg], - memory_compact_threshold: int, - include_thinking: bool = False, + self, + messages: list[Msg], + memory_compact_threshold: int, + include_thinking: bool = False, ) -> str: """Format list of messages to a single formatted string. @@ -348,4 +348,4 @@ class AsMsgHandler: accumulated_tokens, ) - return messages_to_compact, messages_to_keep \ No newline at end of file + return messages_to_compact, messages_to_keep 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 f6205688..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 ReMeOpenAIChatFormatter(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 61943a6c..f08ee98e 100644 --- a/reme/memory/file_based/reme_in_memory_memory.py +++ b/reme/memory/file_based/reme_in_memory_memory.py @@ -1,34 +1,30 @@ """Custom memory implementation with bugfixes and extensions.""" -import logging - -from agentscope.agent._react_agent import _MemoryMark +from agentscope.agent._react_agent import _MemoryMark # noqa from agentscope.memory import InMemoryMemory from agentscope.message import Msg from agentscope.token import HuggingFaceTokenCounter from .as_msg_handler import AsMsgHandler +from ...core.utils import get_std_logger -logger = logging.getLogger(__name__) +logger = get_std_logger() class ReMeInMemoryMemory(InMemoryMemory): """Extended InMemoryMemory with bugfixes and summary support.""" - def __init__( - self, - token_counter: HuggingFaceTokenCounter, - ): + def __init__(self, token_counter: HuggingFaceTokenCounter): super().__init__() self._token_counter: HuggingFaceTokenCounter = token_counter self._msg_handler: AsMsgHandler = AsMsgHandler(token_counter) async def get_memory( - self, - mark: str | None = None, - exclude_mark: str | None = _MemoryMark.COMPRESSED, - prepend_summary: bool = True, - **_kwargs, + 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). @@ -192,10 +188,10 @@ Use it as context to maintain continuity. ) return ( - f"**Conversation History**\n\n" - f"- Total messages: {stats['total_messages']}\n" - f"- Estimated tokens: {stats['estimated_tokens']}\n" - f"- Max input length: {stats['max_input_length']}\n" - f"- Context usage: {stats['context_usage_ratio']:.1f}%\n" - f"- Compressed summary tokens: {stats['compressed_summary_tokens']}\n\n" + "\n\n".join(lines) + f"**Conversation History**\n\n" + f"- Total messages: {stats['total_messages']}\n" + f"- Estimated tokens: {stats['estimated_tokens']}\n" + f"- Max input length: {stats['max_input_length']}\n" + f"- Context usage: {stats['context_usage_ratio']:.1f}%\n" + f"- Compressed summary tokens: {stats['compressed_summary_tokens']}\n\n" + "\n\n".join(lines) ) diff --git a/reme/memory/file_based/sub_agent/__init__.py b/reme/memory/file_based/sub_agent/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/reme/memory/file_based/compactor.py b/reme/memory/file_based/sub_agent/compactor.py similarity index 77% rename from reme/memory/file_based/compactor.py rename to reme/memory/file_based/sub_agent/compactor.py index c7dbb496..571db449 100644 --- a/reme/memory/file_based/compactor.py +++ b/reme/memory/file_based/sub_agent/compactor.py @@ -1,35 +1,28 @@ """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 .as_msg_handler import AsMsgHandler -from ...core.op import BaseOp +from ..as_msg_handler import AsMsgHandler +from ....core.op import BaseOp +from ....core.utils import get_std_logger -logger = logging.getLogger(__name__) +logger = get_std_logger() 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, + self, + memory_compact_threshold: int, + 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.msg_handler = AsMsgHandler(token_counter=token_counter) async def execute(self): @@ -50,9 +43,9 @@ class Compactor(BaseOp): agent = ReActAgent( name="reme_compactor", - model=self.chat_model, + model=self.as_llm, sys_prompt=self.get_prompt("system_prompt"), - formatter=self.formatter, + formatter=self.as_llm_formatter, ) if previous_summary: @@ -66,7 +59,7 @@ class Compactor(BaseOp): ) else: user_message: str = f"\n{history_formatted_str}\n\n\n" \ - + self.get_prompt("initial_user_message") + + 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( diff --git a/reme/memory/file_based/compactor.yaml b/reme/memory/file_based/sub_agent/compactor.yaml similarity index 100% rename from reme/memory/file_based/compactor.yaml rename to reme/memory/file_based/sub_agent/compactor.yaml diff --git a/reme/memory/file_based/summarizer.py b/reme/memory/file_based/sub_agent/summarizer.py similarity index 74% rename from reme/memory/file_based/summarizer.py rename to reme/memory/file_based/sub_agent/summarizer.py index 6f9f0b02..e063757e 100644 --- a/reme/memory/file_based/summarizer.py +++ b/reme/memory/file_based/sub_agent/summarizer.py @@ -1,42 +1,36 @@ """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 .as_msg_handler import AsMsgHandler -from ...core.op import BaseOp +from ..as_msg_handler import AsMsgHandler +from ....core.op import BaseOp +from ....core.utils import get_std_logger -logger = logging.getLogger(__name__) +logger = get_std_logger() 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, - **kwargs, + self, + working_dir: str, + memory_dir: str, + memory_compact_threshold: int, + token_counter: HuggingFaceTokenCounter, + toolkit: Toolkit, + **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.msg_handler = AsMsgHandler(token_counter=token_counter) self.toolkit: Toolkit = toolkit @@ -57,9 +51,9 @@ class Summarizer(BaseOp): agent = ReActAgent( name="reme_summarizer", - model=self.chat_model, + model=self.as_llm, sys_prompt="You are a helpful assistant.", - formatter=self.formatter, + formatter=self.as_llm_formatter, toolkit=self.toolkit, ) diff --git a/reme/memory/file_based/summarizer.yaml b/reme/memory/file_based/sub_agent/summarizer.yaml similarity index 100% rename from reme/memory/file_based/summarizer.yaml rename to reme/memory/file_based/sub_agent/summarizer.yaml diff --git a/reme/memory/file_based/tool_result_compactor.py b/reme/memory/file_based/sub_agent/tool_result_compactor.py similarity index 91% rename from reme/memory/file_based/tool_result_compactor.py rename to reme/memory/file_based/sub_agent/tool_result_compactor.py index 5b1ef0a4..3b49c504 100644 --- a/reme/memory/file_based/tool_result_compactor.py +++ b/reme/memory/file_based/sub_agent/tool_result_compactor.py @@ -1,27 +1,27 @@ """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 +from ....core.op import BaseOp +from ....core.utils import get_std_logger +from ....core.utils import truncate_text, is_truncated -logger = logging.getLogger(__name__) +logger = get_std_logger() 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, + 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) diff --git a/reme/memory/file_based/utils.py b/reme/memory/file_based/utils.py deleted file mode 100644 index f460f96f..00000000 --- a/reme/memory/file_based/utils.py +++ /dev/null @@ -1,271 +0,0 @@ -"""Utility functions for working with text.""" - -import logging -from pathlib import Path - -from agentscope.token import HuggingFaceTokenCounter - -logger = logging.getLogger(__name__) - -# Unique marker for truncated text -TRUNCATION_MARKER_START = "<<>>" -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/tools/__init__.py b/reme/memory/tools/__init__.py index af9b851f..2ad25ef0 100644 --- a/reme/memory/tools/__init__.py +++ b/reme/memory/tools/__init__.py @@ -1,17 +1,14 @@ """memory tools""" 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 from .history.add_history import AddHistory from .history.read_history import ReadHistory from .history.read_history_v2 import ReadHistoryV2 - # profiles tools from .profiles.add_draft_and_read_all_profiles import AddDraftAndReadAllProfiles from .profiles.add_profile import AddProfile @@ -19,7 +16,6 @@ from .profiles.delete_profile import DeleteProfile from .profiles.read_all_profiles import ReadAllProfiles from .profiles.update_profile import UpdateProfile from .profiles.update_profiles_v1 import UpdateProfilesV1 - # record tools from .record.add_and_retrieve_similar_memory import AddAndRetrieveSimilarMemory from .record.add_draft_and_retrieve_similar_memory import AddDraftAndRetrieveSimilarMemory diff --git a/reme/memory/tools/file/__init__.py b/reme/memory/tools/file/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/reme/memory/file_based/file_io.py b/reme/memory/tools/file/file_io.py similarity index 100% rename from reme/memory/file_based/file_io.py rename to reme/memory/tools/file/file_io.py diff --git a/reme/reme_light.py b/reme/reme_light.py index 8a0bdd9e..5c435410 100644 --- a/reme/reme_light.py +++ b/reme/reme_light.py @@ -16,130 +16,56 @@ 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, ReMeOpenAIChatFormatter, FileIO +from .core.utils import get_hf_token_counter +from .memory.file_based import Compactor, Summarizer, ToolResultCompactor, ReMeInMemoryMemory, ReMeOpenAIChatFormatter, \ + FileIO from .memory.file_based.utils import get_token_counter from .memory.tools import MemorySearch -from .core.utils import load_env logger = logging.getLogger(__name__) class ReMeLight(Application): - """ - 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. - - 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 - """ + """ReMe Light Application Class""" def __init__( - self, - working_dir: str = ".reme", - llm_api_key: str | None = None, - 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", - vector_weight: float = 0.7, - candidate_multiplier: float = 3.0, - tool_result_threshold: int = 1000, - retention_days: int = 7, + self, + working_dir: str = ".reme", + 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, + vector_weight: float = 0.7, + candidate_multiplier: float = 3.0, + tool_result_threshold: int = 1000, + retention_days: int = 7, ): # 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) - # 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.") - 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) - 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 - # Initialize the parent Application class with comprehensive configuration super().__init__( llm_api_key=llm_api_key, @@ -151,21 +77,9 @@ class ReMeLight(Application): enable_logo=False, log_to_console=False, 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_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={ "watch_paths": [ str(self.working_path / "MEMORY.md"), @@ -175,107 +89,12 @@ class ReMeLight(Application): }, ) - 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 = ReMeOpenAIChatFormatter(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. - - Args: - key (str): The name of the environment variable to retrieve - default (str): The default value to return if the variable is not set - - Returns: - str: The value of the environment variable, or the default if not set - """ - 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 + def calculate_memory_compact_threshold(max_input_length: float, compact_ratio: float) -> int: + return int(max_input_length * compact_ratio * 0.9) def _cleanup_tool_results(self) -> int: """ @@ -287,10 +106,6 @@ class ReMeLight(Application): 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 @@ -307,67 +122,18 @@ class ReMeLight(Application): return 0 async def start(self): - """ - 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. - - Returns: - The result from the parent class's start method - - Note: - Tool result cleanup runs after successful startup to ensure - the application is fully initialized before performing maintenance - """ - # Initialize parent application components + """Start the application lifecycle.""" result = await super().start() - # Perform initial cleanup of old tool result files self._cleanup_tool_results() return result async def close(self) -> bool: - """ - 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. - - Returns: - bool: True if shutdown was successful, False otherwise - - Note: - Cleanup is performed before calling parent close to ensure - all resources are available during the cleanup process - """ - # Clean up tool results before shutting down + """Close the application and perform cleanup.""" self._cleanup_tool_results() - # Shutdown parent application components return await super().close() - async def compact_tool_result( - self, - messages: list[Msg], - ) -> 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. - - Args: - messages (list[Msg]): List of messages to process for tool result compaction - - Returns: - list[Msg]: The processed message list with large tool results compacted - - 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 - """ + async def compact_tool_result(self, messages: list[Msg]) -> list[Msg]: + """Compact tool results by truncating large outputs and saving full content to files.""" try: # Create compactor with instance configuration compactor = ToolResultCompactor( @@ -389,38 +155,30 @@ 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: - """ - 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. - - Args: - messages (list[Msg]): The list of messages to compact - previous_summary (str): Optional previous summary to incorporate - into the compaction process for continuity - - 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 - """ + async def compact_memory( + self, + messages: list[Msg], + as_llm: str | ChatModelBase = "default", + as_llm_formatter: str | FormatterBase = "default", + token_counter: HuggingFaceTokenCounter | None = None, + language: str = "zh", + max_input_length: float = 128 * 1024, + compact_ratio: float = 0.7, + previous_summary: str = "", + ) -> str: + """Compact a list of messages into a condensed summary.""" try: - # Initialize compactor with current configuration + if token_counter is None: + token_counter = get_hf_token_counter() + 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, + token_counter=token_counter, + language=language if language == "zh" else "", ) - # Execute compaction with optional previous summary context return await compactor.call( messages=messages, previous_summary=previous_summary, @@ -433,25 +191,7 @@ class ReMeLight(Application): return "" async def summary_memory(self, messages: list[Msg]) -> 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. - - Args: - messages (list[Msg]): The list of messages to summarize - - Returns: - str: A generated summary of the messages, or empty string on failure - - 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 - """ + """Generate a comprehensive summary of the given messages.""" try: # Create toolkit if not provided if self.toolkit is not None: @@ -651,24 +391,10 @@ class ReMeLight(Application): ], ) - def get_in_memory_memory(self): - """ - Create and return an in-memory memory instance. + @staticmethod + def get_in_memory_memory(token_counter: HuggingFaceTokenCounter | None = None): + """Create and return an in-memory memory instance.""" + if token_counter is None: + token_counter = get_hf_token_counter() - 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. - - Returns: - ReMeInMemoryMemory: A configured in-memory memory instance ready - for storing and retrieving conversation messages - - 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 as the main application - """ - return ReMeInMemoryMemory( - token_counter=self.token_counter, - ) + return ReMeInMemoryMemory(token_counter=token_counter) From fc7b1cdba82735be0b66e87d9f5f4494eb2f4c86 Mon Sep 17 00:00:00 2001 From: "jinli.yl" Date: Fri, 6 Mar 2026 01:57:29 +0800 Subject: [PATCH 04/59] refactor(core): update registry registration syntax and improve code formatting --- reme/core/__init__.py | 1 + reme/core/application.py | 52 +-- reme/core/as_llm/__init__.py | 6 +- reme/core/as_llm_formatter/__init__.py | 6 +- reme/core/op/base_op.py | 5 +- reme/core/schema/as_msg_stat.py | 26 +- reme/core/utils/hf_token_counter_utils.py | 8 +- reme/core/utils/truncate_text_utils.py | 2 + reme/memory/file_based/as_msg_handler.py | 95 +++-- .../file_based/reme_in_memory_memory.py | 22 +- reme/memory/file_based/sub_agent/compactor.py | 13 +- .../memory/file_based/sub_agent/summarizer.py | 14 +- .../sub_agent/tool_result_compactor.py | 10 +- reme/memory/tools/__init__.py | 4 + reme/memory/tools/file/__init__.py | 7 + reme/reme_light.py | 180 ++++----- tests/light/test_compactor.py | 9 +- tests/light/test_context_check.py | 359 ++++++++++++++---- tests/light/test_format_msgs_to_str.py | 168 ++++---- tests/light/test_memory_formatter.py | 10 +- tests/light/test_reme_light.py | 6 +- tests/light/test_summarizer.py | 9 +- tests/light/test_tool_result_compactor.py | 1 - tests/light/test_utils.py | 71 +--- 24 files changed, 616 insertions(+), 468 deletions(-) diff --git a/reme/core/__init__.py b/reme/core/__init__.py index 88e42175..053755cc 100644 --- a/reme/core/__init__.py +++ b/reme/core/__init__.py @@ -1,4 +1,5 @@ """Core""" + from . import as_llm from . import as_llm_formatter from . import embedding diff --git a/reme/core/application.py b/reme/core/application.py index bc59c7f3..46f4a934 100644 --- a/reme/core/application.py +++ b/reme/core/application.py @@ -25,26 +25,26 @@ class Application: """Application wrapper that wires together service context, flows, and runtimes.""" 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, - working_dir: str | None = None, - config_path: str | None = None, - enable_logo: bool = True, - log_to_console: 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, - default_file_store_config: dict | None = None, - default_token_counter_config: dict | None = None, - default_file_watcher_config: dict | None = None, - **kwargs, + 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, + working_dir: str | None = None, + config_path: str | None = None, + enable_logo: bool = True, + log_to_console: 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, + default_file_store_config: dict | None = None, + default_token_counter_config: dict | None = None, + default_file_watcher_config: dict | None = None, + **kwargs, ): self.service_context = ServiceContext( *args, @@ -142,8 +142,8 @@ class Application: ray.init(num_cpus=self.service_config.ray_max_workers) if ( - self.service_context.thread_pool is None - or self.service_context.thread_pool._shutdown # pylint: disable=protected-access + 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, @@ -319,10 +319,10 @@ class Application: stream_queue = asyncio.Queue() task = asyncio.create_task(flow.call(stream_queue=stream_queue, **kwargs)) async for chunk in execute_stream_task( - stream_queue=stream_queue, - task=task, - task_name=name, - output_format="str", + stream_queue=stream_queue, + task=task, + task_name=name, + output_format="str", ): yield chunk diff --git a/reme/core/as_llm/__init__.py b/reme/core/as_llm/__init__.py index 888048e6..9cf527af 100644 --- a/reme/core/as_llm/__init__.py +++ b/reme/core/as_llm/__init__.py @@ -1,7 +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(OpenAIChatModel, "openai") -R.as_llms.register(DashScopeChatModel, "dashscope") +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 index 9c3a52cf..88b326a7 100644 --- a/reme/core/as_llm_formatter/__init__.py +++ b/reme/core/as_llm_formatter/__init__.py @@ -1,7 +1,9 @@ +"""Module for registering AgentScope LLM formatters.""" + from agentscope.formatter import DashScopeChatFormatter from agentscope.formatter import OpenAIChatFormatter from ..registry_factory import R -R.as_llm_formatters.register(OpenAIChatFormatter, "openai") -R.as_llm_formatters.register(DashScopeChatFormatter, "dashscope") +R.as_llm_formatters.register("openai")(OpenAIChatFormatter) +R.as_llm_formatters.register("dashscope")(DashScopeChatFormatter) diff --git a/reme/core/op/base_op.py b/reme/core/op/base_op.py index 0f0580a8..86cfa3be 100644 --- a/reme/core/op/base_op.py +++ b/reme/core/op/base_op.py @@ -7,6 +7,8 @@ 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 loguru import logger from tqdm import tqdm @@ -21,8 +23,7 @@ from ..service_context import ServiceContext from ..token_counter import BaseTokenCounter from ..utils import camel_to_snake, CacheHandler, timer from ..vector_store import BaseVectorStore -from agentscope.model import ChatModelBase -from agentscope.formatter import FormatterBase + class BaseOp(metaclass=ABCMeta): """Base operator class for LLM workflow execution and composition.""" diff --git a/reme/core/schema/as_msg_stat.py b/reme/core/schema/as_msg_stat.py index 1acb9e8a..4bb69f99 100644 --- a/reme/core/schema/as_msg_stat.py +++ b/reme/core/schema/as_msg_stat.py @@ -1,3 +1,5 @@ +"""Schema definitions for AgentScope message statistics.""" + from pydantic import BaseModel, Field _DEFAULT_MAX_BLOCK_TEXT_PREVIEW_LENGTH = 100 @@ -5,6 +7,8 @@ _DEFAULT_MAX_FORMATTER_TEXT_LENGTH = 2000 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") @@ -19,10 +23,20 @@ class AsBlockStat(BaseModel): @property def preview(self) -> str: + """Return a short preview of the block content.""" return self.format(_DEFAULT_MAX_BLOCK_TEXT_PREVIEW_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.""" + """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. + """ from ..utils import truncate_text if self.block_type == "text": @@ -33,15 +47,17 @@ class AsBlockStat(BaseModel): return "" if self.block_type in ("image", "audio", "video"): return f"[{self.block_type}] {self.media_url}" if self.media_url else f"[{self.block_type}]" - if self.block_type == "tool_use": - return f" - tool_call={self.tool_name} params={truncate_text(self.tool_input, max_length)}" - if self.block_type == "tool_result": + if self.block_type in ("tool_use", "tool_result"): + if self.block_type == "tool_use": + return f" - tool_call={self.tool_name} params={truncate_text(self.tool_input, max_length)}" output = truncate_text(self.tool_output, max_length) return f" - tool_result={self.tool_name} output={output}" if output else "" 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) @@ -50,10 +66,12 @@ class AsMsgStat(BaseModel): @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: diff --git a/reme/core/utils/hf_token_counter_utils.py b/reme/core/utils/hf_token_counter_utils.py index dfbc3dc6..a8ab348c 100644 --- a/reme/core/utils/hf_token_counter_utils.py +++ b/reme/core/utils/hf_token_counter_utils.py @@ -6,10 +6,10 @@ _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, + 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 diff --git a/reme/core/utils/truncate_text_utils.py b/reme/core/utils/truncate_text_utils.py index ec85ec61..da0c473a 100644 --- a/reme/core/utils/truncate_text_utils.py +++ b/reme/core/utils/truncate_text_utils.py @@ -1,3 +1,5 @@ +"""Utility functions for truncating long text strings.""" + from .std_logger import get_logger logger = get_logger() diff --git a/reme/memory/file_based/as_msg_handler.py b/reme/memory/file_based/as_msg_handler.py index e967f81e..802157ab 100644 --- a/reme/memory/file_based/as_msg_handler.py +++ b/reme/memory/file_based/as_msg_handler.py @@ -1,3 +1,5 @@ +"""Handler for AgentScope message processing, token counting, and context management.""" + import json from agentscope.message import Msg @@ -10,6 +12,7 @@ logger = get_std_logger() class AsMsgHandler: + """Handles token counting, formatting, and context compaction for AgentScope messages.""" def __init__(self, token_counter: HuggingFaceTokenCounter): self._token_counter = token_counter @@ -33,7 +36,7 @@ class AsMsgHandler: except Exception as e: estimated_tokens = len(text.encode("utf-8")) // 4 - logger.warning(f"Failed to count string tokens: {text}, using estimated_tokens={estimated_tokens}") + logger.warning(f"Failed to count string tokens: {text}, e={e}") return estimated_tokens @staticmethod @@ -107,20 +110,24 @@ class AsMsgHandler: if block_type == "text": text = block.get("text", "") token_count = self.count_str_token(text) - blocks.append(AsBlockStat( - block_type=block_type, - text=text, - token_count=token_count, - )) + blocks.append( + AsBlockStat( + block_type=block_type, + text=text, + token_count=token_count, + ), + ) elif block_type == "thinking": thinking = block.get("thinking", "") token_count = self.count_str_token(thinking) - blocks.append(AsBlockStat( - block_type=block_type, - text=thinking, - token_count=token_count, - )) + blocks.append( + AsBlockStat( + block_type=block_type, + text=thinking, + token_count=token_count, + ), + ) elif block_type in ("image", "audio", "video"): source = block.get("source", {}) @@ -131,12 +138,14 @@ class AsMsgHandler: token_count = len(data) // 4 if data else 10 else: token_count = self.count_str_token(url) if url else 10 - blocks.append(AsBlockStat( - block_type=block_type, - text="", - token_count=token_count, - media_url=url, - )) + 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", "") @@ -146,26 +155,30 @@ class AsMsgHandler: except (TypeError, ValueError): input_str = str(tool_input) token_count = 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, - )) + 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 = self._format_tool_result_output(output) token_count = self.count_str_token(formatted_output) - blocks.append(AsBlockStat( - block_type=block_type, - text="", - token_count=token_count, - tool_name=tool_name, - tool_output=formatted_output, - )) + blocks.append( + AsBlockStat( + block_type=block_type, + text="", + token_count=token_count, + tool_name=tool_name, + tool_output=formatted_output, + ), + ) else: logger.warning("Unsupported block type %s, skipped.", block_type) @@ -179,10 +192,10 @@ class AsMsgHandler: ) def format_msgs_to_str( - self, - messages: list[Msg], - memory_compact_threshold: int, - include_thinking: bool = False, + self, + messages: list[Msg], + memory_compact_threshold: int, + include_thinking: bool = False, ) -> str: """Format list of messages to a single formatted string. @@ -219,10 +232,10 @@ class AsMsgHandler: return "\n\n".join(formatted_parts) def context_check( - self, - messages: list[Msg], - memory_compact_threshold: int, - memory_compact_reserve: int, + self, + messages: list[Msg], + memory_compact_threshold: int, + memory_compact_reserve: int, ) -> tuple[list[Msg], list[Msg]]: """Check if context exceeds threshold and split messages accordingly. @@ -294,9 +307,7 @@ class AsMsgHandler: # 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", "") + 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 diff --git a/reme/memory/file_based/reme_in_memory_memory.py b/reme/memory/file_based/reme_in_memory_memory.py index f08ee98e..16f18726 100644 --- a/reme/memory/file_based/reme_in_memory_memory.py +++ b/reme/memory/file_based/reme_in_memory_memory.py @@ -20,11 +20,11 @@ class ReMeInMemoryMemory(InMemoryMemory): self._msg_handler: AsMsgHandler = AsMsgHandler(token_counter) async def get_memory( - self, - mark: str | None = None, - exclude_mark: str | None = _MemoryMark.COMPRESSED, - prepend_summary: bool = True, - **_kwargs, + 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). @@ -188,10 +188,10 @@ Use it as context to maintain continuity. ) return ( - f"**Conversation History**\n\n" - f"- Total messages: {stats['total_messages']}\n" - f"- Estimated tokens: {stats['estimated_tokens']}\n" - f"- Max input length: {stats['max_input_length']}\n" - f"- Context usage: {stats['context_usage_ratio']:.1f}%\n" - f"- Compressed summary tokens: {stats['compressed_summary_tokens']}\n\n" + "\n\n".join(lines) + f"**Conversation History**\n\n" + f"- Total messages: {stats['total_messages']}\n" + f"- Estimated tokens: {stats['estimated_tokens']}\n" + f"- Max input length: {stats['max_input_length']}\n" + f"- Context usage: {stats['context_usage_ratio']:.1f}%\n" + f"- Compressed summary tokens: {stats['compressed_summary_tokens']}\n\n" + "\n\n".join(lines) ) diff --git a/reme/memory/file_based/sub_agent/compactor.py b/reme/memory/file_based/sub_agent/compactor.py index 571db449..3292c874 100644 --- a/reme/memory/file_based/sub_agent/compactor.py +++ b/reme/memory/file_based/sub_agent/compactor.py @@ -15,10 +15,10 @@ class Compactor(BaseOp): """Compactor class for compacting memory messages.""" def __init__( - self, - memory_compact_threshold: int, - token_counter: HuggingFaceTokenCounter, - **kwargs, + self, + memory_compact_threshold: int, + token_counter: HuggingFaceTokenCounter, + **kwargs, ): super().__init__(**kwargs) self.memory_compact_threshold: int = memory_compact_threshold @@ -58,8 +58,9 @@ class Compactor(BaseOp): f"{suffix}" ) else: - user_message: str = f"\n{history_formatted_str}\n\n\n" \ - + self.get_prompt("initial_user_message") + 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( diff --git a/reme/memory/file_based/sub_agent/summarizer.py b/reme/memory/file_based/sub_agent/summarizer.py index e063757e..db3522da 100644 --- a/reme/memory/file_based/sub_agent/summarizer.py +++ b/reme/memory/file_based/sub_agent/summarizer.py @@ -18,13 +18,13 @@ class Summarizer(BaseOp): """Summarizer class for summarizing memory messages.""" def __init__( - self, - working_dir: str, - memory_dir: str, - memory_compact_threshold: int, - token_counter: HuggingFaceTokenCounter, - toolkit: Toolkit, - **kwargs, + self, + working_dir: str, + memory_dir: str, + memory_compact_threshold: int, + token_counter: HuggingFaceTokenCounter, + toolkit: Toolkit, + **kwargs, ): super().__init__(**kwargs) self.working_dir: str = working_dir diff --git a/reme/memory/file_based/sub_agent/tool_result_compactor.py b/reme/memory/file_based/sub_agent/tool_result_compactor.py index 3b49c504..412df6de 100644 --- a/reme/memory/file_based/sub_agent/tool_result_compactor.py +++ b/reme/memory/file_based/sub_agent/tool_result_compactor.py @@ -17,11 +17,11 @@ 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, + 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) diff --git a/reme/memory/tools/__init__.py b/reme/memory/tools/__init__.py index 2ad25ef0..af9b851f 100644 --- a/reme/memory/tools/__init__.py +++ b/reme/memory/tools/__init__.py @@ -1,14 +1,17 @@ """memory tools""" 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 from .history.add_history import AddHistory from .history.read_history import ReadHistory from .history.read_history_v2 import ReadHistoryV2 + # profiles tools from .profiles.add_draft_and_read_all_profiles import AddDraftAndReadAllProfiles from .profiles.add_profile import AddProfile @@ -16,6 +19,7 @@ from .profiles.delete_profile import DeleteProfile from .profiles.read_all_profiles import ReadAllProfiles from .profiles.update_profile import UpdateProfile from .profiles.update_profiles_v1 import UpdateProfilesV1 + # record tools from .record.add_and_retrieve_similar_memory import AddAndRetrieveSimilarMemory from .record.add_draft_and_retrieve_similar_memory import AddDraftAndRetrieveSimilarMemory diff --git a/reme/memory/tools/file/__init__.py b/reme/memory/tools/file/__init__.py index e69de29b..8234e60d 100644 --- a/reme/memory/tools/file/__init__.py +++ b/reme/memory/tools/file/__init__.py @@ -0,0 +1,7 @@ +"""File-based memory tool implementations.""" + +from .file_io import FileIO + +__all__ = [ + "FileIO", +] diff --git a/reme/reme_light.py b/reme/reme_light.py index 5c435410..82a11850 100644 --- a/reme/reme_light.py +++ b/reme/reme_light.py @@ -15,7 +15,6 @@ Key Features: """ import asyncio -import logging from pathlib import Path from agentscope.formatter import FormatterBase @@ -26,32 +25,31 @@ from agentscope.tool import Toolkit, ToolResponse from .config import ReMeConfigParser from .core import Application -from .core.utils import get_hf_token_counter -from .memory.file_based import Compactor, Summarizer, ToolResultCompactor, ReMeInMemoryMemory, ReMeOpenAIChatFormatter, \ - FileIO -from .memory.file_based.utils import get_token_counter +from .core.utils import get_hf_token_counter, get_std_logger +from .memory.file_based import Compactor, Summarizer, ToolResultCompactor, ReMeInMemoryMemory from .memory.tools import MemorySearch +from .memory.tools.file import FileIO -logger = logging.getLogger(__name__) +logger = get_std_logger() class ReMeLight(Application): """ReMe Light Application Class""" def __init__( - self, - working_dir: str = ".reme", - 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, - vector_weight: float = 0.7, - candidate_multiplier: float = 3.0, - tool_result_threshold: int = 1000, - retention_days: int = 7, + self, + working_dir: str = ".reme", + 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, + vector_weight: float = 0.7, + candidate_multiplier: float = 3.0, + tool_result_threshold: int = 1000, + retention_days: int = 7, ): # Initialize working directory structure self.working_path = Path(working_dir).absolute() @@ -94,6 +92,15 @@ class ReMeLight(Application): @staticmethod 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: + max_input_length: Maximum input length in tokens. + compact_ratio: Ratio of the input length to use as the threshold. + + Returns: + Computed compaction threshold as an integer. + """ return int(max_input_length * compact_ratio * 0.9) def _cleanup_tool_results(self) -> int: @@ -156,15 +163,15 @@ class ReMeLight(Application): return messages async def compact_memory( - self, - messages: list[Msg], - as_llm: str | ChatModelBase = "default", - as_llm_formatter: str | FormatterBase = "default", - token_counter: HuggingFaceTokenCounter | None = None, - language: str = "zh", - max_input_length: float = 128 * 1024, - compact_ratio: float = 0.7, - previous_summary: str = "", + self, + messages: list[Msg], + as_llm: str | ChatModelBase = "default", + as_llm_formatter: str | FormatterBase = "default", + token_counter: HuggingFaceTokenCounter | None = None, + language: str = "zh", + max_input_length: float = 128 * 1024, + compact_ratio: float = 0.7, + previous_summary: str = "", ) -> str: """Compact a list of messages into a condensed summary.""" try: @@ -173,9 +180,9 @@ class ReMeLight(Application): compactor = Compactor( memory_compact_threshold=self.calculate_memory_compact_threshold(max_input_length, compact_ratio), + token_counter=token_counter, as_llm=as_llm, as_llm_formatter=as_llm_formatter, - token_counter=token_counter, language=language if language == "zh" else "", ) @@ -190,58 +197,69 @@ class ReMeLight(Application): logger.exception(f"Error compacting memory: {e}") 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", + token_counter: HuggingFaceTokenCounter | None = None, + toolkit: Toolkit | None = None, + language: str = "zh", + max_input_length: float = 128 * 1024, + compact_ratio: float = 0.7, + ) -> str: """Generate a comprehensive summary of the given messages.""" try: - # Create toolkit if not provided - if self.toolkit is not None: - toolkit = self.toolkit - else: + if token_counter is None: + token_counter = get_hf_token_counter() + + if toolkit is None: toolkit = Toolkit() file_io = FileIO(working_dir=str(self.working_path)) toolkit.register_tool_function(file_io.read) toolkit.register_tool_function(file_io.write) toolkit.register_tool_function(file_io.edit) - # Initialize summarizer with working directories and configuration 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, + memory_compact_threshold=self.calculate_memory_compact_threshold(max_input_length, compact_ratio), + token_counter=token_counter, toolkit=toolkit, - language=self.language, + as_llm=as_llm, + as_llm_formatter=as_llm_formatter, + language=language if language == "zh" else "", ) - # Execute summarization on the provided messages 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.""" + 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) + 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. - - Returns: - str: A concatenated string containing the status and results of - all summary tasks, with each task on a new line - - 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 - """ + """Wait for all background summary tasks to complete and collect results.""" result = "" for task in self.summary_tasks: if task.done(): @@ -279,48 +297,6 @@ 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.error(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. diff --git a/tests/light/test_compactor.py b/tests/light/test_compactor.py index 32dfd9d3..19891e29 100644 --- a/tests/light/test_compactor.py +++ b/tests/light/test_compactor.py @@ -1,7 +1,6 @@ """Tests for Compactor.""" import asyncio -import logging from agentscope.message import Msg @@ -10,14 +9,10 @@ from test_utils import ( get_formatter, get_token_counter, ) +from reme.core.utils import get_std_logger 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__) +logger = get_std_logger() # ANSI 颜色码 diff --git a/tests/light/test_context_check.py b/tests/light/test_context_check.py index 300f0b61..f65f961e 100644 --- a/tests/light/test_context_check.py +++ b/tests/light/test_context_check.py @@ -1,18 +1,12 @@ """Tests for AsMsgHandler.context_check method.""" -import logging - from agentscope.message import Msg from test_utils import get_token_counter +from reme.core.utils import get_std_logger from reme.memory.file_based.as_msg_handler import AsMsgHandler -# Configure logging -logging.basicConfig( - level=logging.INFO, - format="%(asctime)s - %(name)s - %(levelname)s - %(message)s", -) -logger = logging.getLogger(__name__) +logger = get_std_logger() # ANSI color codes @@ -101,8 +95,7 @@ def verify_context_check_invariants( # 2. Reserve requirement check kept_tokens = sum(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})" + f"[{test_name}] Reserve violation: kept_tokens ({kept_tokens}) > " f"reserve ({memory_compact_reserve})" ) # 3. Order requirement check - both lists should preserve original order @@ -143,9 +136,7 @@ def verify_context_check_invariants( 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" - ) + assert all_returned == all_original, f"[{test_name}] Message set mismatch: returned messages differ from original" def create_user_msg(content: str) -> Msg: @@ -234,7 +225,7 @@ def test_empty_messages(): memory_compact_threshold=threshold, memory_compact_reserve=reserve, ) - assert to_compact == [], f"Expected empty compact list, got: {to_compact}" + 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") @@ -254,10 +245,18 @@ def test_below_threshold_returns_all(): memory_compact_threshold=threshold, # Very high threshold memory_compact_reserve=reserve, ) - assert to_compact == [], f"Expected empty compact list, got: {len(to_compact)}" + 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") + verify_context_check_invariants( + handler, + messages, + to_compact, + to_keep, + threshold, + reserve, + "test_below_threshold_returns_all", + ) print_pass("test_below_threshold_returns_all") @@ -280,7 +279,15 @@ def test_above_threshold_triggers_compaction(): # 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") + verify_context_check_invariants( + handler, + messages, + to_compact, + to_keep, + threshold, + reserve, + "test_above_threshold_triggers_compaction", + ) print_pass("test_above_threshold_triggers_compaction") @@ -304,7 +311,15 @@ def test_message_order_preserved(): 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") + verify_context_check_invariants( + handler, + messages, + to_compact, + to_keep, + threshold, + reserve, + "test_message_order_preserved", + ) print_pass("test_message_order_preserved") @@ -323,9 +338,17 @@ def test_single_message_below_threshold(): memory_compact_threshold=threshold, memory_compact_reserve=reserve, ) - assert to_compact == [], "Should not compact single message below threshold" + 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") + verify_context_check_invariants( + handler, + messages, + to_compact, + to_keep, + threshold, + reserve, + "test_single_message_below_threshold", + ) print_pass("test_single_message_below_threshold") @@ -343,7 +366,15 @@ def test_single_message_above_threshold(): # 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") + verify_context_check_invariants( + handler, + messages, + to_compact, + to_keep, + threshold, + reserve, + "test_single_message_above_threshold", + ) print_pass("test_single_message_above_threshold") @@ -388,12 +419,12 @@ 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 = handler.stat_message(messages[0]) exact_tokens = stat.total_tokens threshold, reserve = exact_tokens, exact_tokens - + # Test at exact boundary to_compact, to_keep = handler.context_check( messages=messages, @@ -401,9 +432,17 @@ def test_exact_threshold_boundary(): memory_compact_reserve=reserve, ) # At exact boundary (<=), should not trigger compaction - assert to_compact == [], "Should not compact at exact boundary" + 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") + verify_context_check_invariants( + handler, + messages, + to_compact, + to_keep, + threshold, + reserve, + "test_exact_threshold_boundary", + ) print_pass("test_exact_threshold_boundary") @@ -423,7 +462,15 @@ def test_reserve_larger_than_threshold(): # 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") + verify_context_check_invariants( + handler, + messages, + to_compact, + to_keep, + threshold, + reserve, + "test_reserve_larger_than_threshold", + ) print_pass("test_reserve_larger_than_threshold") @@ -447,20 +494,22 @@ def test_tool_use_result_paired(): 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 - ) - + 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") + verify_context_check_invariants( + handler, + messages, + to_compact, + to_keep, + threshold, + reserve, + "test_tool_use_result_paired", + ) print_pass("test_tool_use_result_paired") @@ -480,7 +529,15 @@ def test_tool_use_without_result(): ) # 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") + verify_context_check_invariants( + handler, + messages, + to_compact, + to_keep, + threshold, + reserve, + "test_tool_use_without_result", + ) print_pass("test_tool_use_without_result") @@ -500,7 +557,15 @@ def test_tool_result_without_use(): ) # 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") + verify_context_check_invariants( + handler, + messages, + to_compact, + to_keep, + threshold, + reserve, + "test_tool_result_without_use", + ) print_pass("test_tool_result_without_use") @@ -523,7 +588,7 @@ def test_multiple_tool_pairs(): 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"): @@ -537,7 +602,15 @@ def test_multiple_tool_pairs(): 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") + verify_context_check_invariants( + handler, + messages, + to_compact, + to_keep, + threshold, + reserve, + "test_multiple_tool_pairs", + ) print_pass("test_multiple_tool_pairs") @@ -552,7 +625,7 @@ def test_tool_dependency_causes_extra_inclusion(): 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_user_msg("Middle " * 100), # Large message create_tool_result_msg("call_dep", "dep_tool", "Result"), # Small create_assistant_msg("End"), # Small ] @@ -562,22 +635,27 @@ def test_tool_dependency_causes_extra_inclusion(): 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()) + 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 + 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") + 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") @@ -598,24 +676,30 @@ def test_tool_dependency_exceeds_reserve(): 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()) + 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()) + 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") + verify_context_check_invariants( + handler, + messages, + to_compact, + to_keep, + threshold, + reserve, + "test_tool_dependency_exceeds_reserve", + ) print_pass("test_tool_dependency_exceeds_reserve") @@ -636,19 +720,26 @@ def test_interleaved_tool_pairs(): 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()) + 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") + verify_context_check_invariants( + handler, + messages, + to_compact, + to_keep, + threshold, + reserve, + "test_interleaved_tool_pairs", + ) print_pass("test_interleaved_tool_pairs") @@ -671,7 +762,15 @@ def test_message_with_empty_content(): 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") + verify_context_check_invariants( + handler, + messages, + to_compact, + to_keep, + threshold, + reserve, + "test_message_with_empty_content", + ) print_pass("test_message_with_empty_content") @@ -689,7 +788,15 @@ def test_message_with_whitespace_only(): 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") + verify_context_check_invariants( + handler, + messages, + to_compact, + to_keep, + threshold, + reserve, + "test_message_with_whitespace_only", + ) print_pass("test_message_with_whitespace_only") @@ -706,7 +813,15 @@ def test_very_long_single_message(): ) # 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") + verify_context_check_invariants( + handler, + messages, + to_compact, + to_keep, + threshold, + reserve, + "test_very_long_single_message", + ) print_pass("test_very_long_single_message") @@ -723,7 +838,15 @@ def test_many_small_messages(): # 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") + verify_context_check_invariants( + handler, + messages, + to_compact, + to_keep, + threshold, + reserve, + "test_many_small_messages", + ) print_pass("test_many_small_messages") @@ -760,7 +883,15 @@ def test_special_characters_content(): 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") + verify_context_check_invariants( + handler, + messages, + to_compact, + to_keep, + threshold, + reserve, + "test_special_characters_content", + ) print_pass("test_special_characters_content") @@ -776,11 +907,11 @@ def test_all_messages_fit_exactly_in_reserve(): create_user_msg("Message 1"), create_assistant_msg("Message 2"), ] - + # Calculate total tokens total = sum(handler.stat_message(m).total_tokens for m in messages) threshold, reserve = total - 1, total - + to_compact, to_keep = handler.context_check( messages=messages, memory_compact_threshold=threshold, # Just below total to trigger @@ -788,7 +919,15 @@ def test_all_messages_fit_exactly_in_reserve(): ) # 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") + 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") @@ -800,20 +939,28 @@ def test_first_message_only_compacted(): create_assistant_msg("Small"), # Small create_user_msg("Tiny"), # Tiny ] - + # Calculate tokens to set appropriate reserve small_msg_tokens = handler.stat_message(messages[1]).total_tokens tiny_msg_tokens = handler.stat_message(messages[2]).total_tokens threshold, reserve = 50, small_msg_tokens + tiny_msg_tokens + 10 - + to_compact, to_keep = 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") + verify_context_check_invariants( + handler, + messages, + to_compact, + to_keep, + threshold, + reserve, + "test_first_message_only_compacted", + ) print_pass("test_first_message_only_compacted") @@ -825,20 +972,28 @@ def test_last_message_only_kept(): create_assistant_msg("Large " * 200), create_user_msg("Tiny"), # Only this fits ] - + tiny_tokens = handler.stat_message(messages[2]).total_tokens threshold, reserve = 10, tiny_tokens + 5 - + to_compact, to_keep = 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") + verify_context_check_invariants( + handler, + messages, + to_compact, + to_keep, + threshold, + reserve, + "test_last_message_only_kept", + ) print_pass("test_last_message_only_kept") @@ -857,7 +1012,15 @@ def test_all_messages_compacted(): ) 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") + verify_context_check_invariants( + handler, + messages, + to_compact, + to_keep, + threshold, + reserve, + "test_all_messages_compacted", + ) print_pass("test_all_messages_compacted") @@ -929,7 +1092,15 @@ def test_tool_use_with_empty_id(): ) # 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") + 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") @@ -949,7 +1120,15 @@ def test_tool_result_with_empty_id(): ) # 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") + 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") @@ -970,7 +1149,15 @@ def test_duplicate_tool_ids(): ) # 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") + verify_context_check_invariants( + handler, + messages, + to_compact, + to_keep, + threshold, + reserve, + "test_duplicate_tool_ids", + ) print_pass("test_duplicate_tool_ids") @@ -1000,7 +1187,15 @@ def test_message_with_multiple_tool_blocks(): 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") + 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") diff --git a/tests/light/test_format_msgs_to_str.py b/tests/light/test_format_msgs_to_str.py index 29e1b2cb..bd69751a 100644 --- a/tests/light/test_format_msgs_to_str.py +++ b/tests/light/test_format_msgs_to_str.py @@ -2,19 +2,15 @@ # pylint: disable=W0212 -import logging +import sys from agentscope.message import Msg from test_utils import get_token_counter +from reme.core.utils import get_std_logger from reme.memory.file_based.as_msg_handler import AsMsgHandler -# 配置日志输出到控制台 -logging.basicConfig( - level=logging.INFO, - format="%(asctime)s - %(name)s - %(levelname)s - %(message)s", -) -logger = logging.getLogger(__name__) +logger = get_std_logger() # ANSI 颜色码 @@ -72,7 +68,7 @@ def verify_result_within_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: @@ -93,7 +89,7 @@ def verify_result_within_threshold( for msg in msgs: stat = handler.stat_message(msg) # Check if this message's content appears in the result - formatted = stat.format(include_thinking=True) # Use True to check all content + _ = 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 @@ -102,21 +98,21 @@ def verify_result_within_threshold( if block_type == "text" and block.get("text", "") in result: msg_included = True break - elif block_type == "tool_use" and f"tool_call={block.get('name', '')}" in result: + if block_type == "tool_use" and f"tool_call={block.get('name', '')}" in result: msg_included = True break - elif block_type == "tool_result" and f"tool_result={block.get('name', '')}" in result: + if block_type == "tool_result" and f"tool_result={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})." - ) + assert ( + included_tokens <= threshold + 1 + ), f"{test_name}: Included messages token count ({included_tokens}) exceeds threshold ({threshold})." def create_user_msg(content: str) -> Msg: @@ -199,12 +195,14 @@ def create_mixed_content_msg( 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 {}, - }) + 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) @@ -274,8 +272,7 @@ def test_format_msgs_to_str_message_order(): 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}" + 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") @@ -347,9 +344,7 @@ def test_format_msgs_to_str_thinking_excluded_by_default(): msgs = [create_thinking_msg("Let me think about this...", "Here is my response")] result = 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 "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") @@ -362,9 +357,7 @@ def test_format_msgs_to_str_thinking_included(): msgs = [create_thinking_msg("Let me think about this...", "Here is my response")] result = 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 "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") @@ -375,14 +368,18 @@ def test_format_msgs_to_str_thinking_only_message(): handler = create_handler() threshold = 4000 msgs = [create_thinking_msg("Deep thoughts here")] - + # With include_thinking=False result_no_thinking = handler.format_msgs_to_str( - msgs, memory_compact_threshold=threshold, include_thinking=False + msgs, + memory_compact_threshold=threshold, + include_thinking=False, ) # With include_thinking=True result_with_thinking = handler.format_msgs_to_str( - msgs, memory_compact_threshold=threshold, include_thinking=True + msgs, + memory_compact_threshold=threshold, + include_thinking=True, ) assert "Deep thoughts here" not in result_no_thinking @@ -425,9 +422,9 @@ def test_format_msgs_to_str_exceeds_threshold_truncate_older(): result = 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]}..." - ) + 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) @@ -517,10 +514,7 @@ 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) - ] + msgs = [create_user_msg("Message " + str(i) + " " + "x" * 100) for i in range(50)] result = handler.format_msgs_to_str(msgs, memory_compact_threshold=threshold) @@ -604,19 +598,25 @@ 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", - )] + 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 = handler.format_msgs_to_str( - msgs, memory_compact_threshold=threshold, include_thinking=False + msgs, + memory_compact_threshold=threshold, + include_thinking=False, ) result_with_thinking = handler.format_msgs_to_str( - msgs, memory_compact_threshold=threshold, include_thinking=True + msgs, + memory_compact_threshold=threshold, + include_thinking=True, ) assert "Text content" in result_no_thinking @@ -682,7 +682,7 @@ def 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): @@ -690,16 +690,14 @@ def test_format_msgs_to_str_incremental_threshold_check(): # Calculate total tokens total_tokens = sum(handler.stat_message(msg).total_tokens for msg in msgs) - + # Use threshold that allows about half the messages half_threshold = total_tokens // 2 result = 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" - ) + 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) @@ -743,17 +741,21 @@ 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 - }, - }], - )] + msgs = [ + Msg( + name="assistant", + role="assistant", + content=[ + { + "type": "image", + "source": { + "type": "base64", + "data": "SGVsbG8gV29ybGQ=" * 100, # Simulated base64 data + }, + }, + ], + ), + ] result = handler.format_msgs_to_str(msgs, memory_compact_threshold=threshold) assert "[image]" in result @@ -765,14 +767,16 @@ 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"}}, - ], - )] + 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 = handler.format_msgs_to_str(msgs, memory_compact_threshold=threshold) assert "[audio]" in result @@ -785,14 +789,16 @@ 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"}, - ], - )] + msgs = [ + Msg( + name="assistant", + role="assistant", + content=[ + {"type": "unknown_type", "data": "some data"}, + {"type": "text", "text": "Valid text"}, + ], + ), + ] result = handler.format_msgs_to_str(msgs, memory_compact_threshold=threshold) # Should still include valid content @@ -880,4 +886,4 @@ def run_all_tests(): if __name__ == "__main__": success = run_all_tests() - exit(0 if success else 1) + sys.exit(0 if success else 1) diff --git a/tests/light/test_memory_formatter.py b/tests/light/test_memory_formatter.py index 00f45bb1..8b31718b 100644 --- a/tests/light/test_memory_formatter.py +++ b/tests/light/test_memory_formatter.py @@ -2,19 +2,13 @@ # pylint: disable=W0212 -import logging - from agentscope.message import Msg from test_utils import get_token_counter +from reme.core.utils import get_std_logger from reme.memory.file_based import MemoryFormatter -# 配置日志输出到控制台 -logging.basicConfig( - level=logging.INFO, - format="%(asctime)s - %(name)s - %(levelname)s - %(message)s", -) -logger = logging.getLogger(__name__) +logger = get_std_logger() # ANSI 颜色码 diff --git a/tests/light/test_reme_light.py b/tests/light/test_reme_light.py index 9e2ac70a..49102826 100644 --- a/tests/light/test_reme_light.py +++ b/tests/light/test_reme_light.py @@ -3,6 +3,7 @@ import asyncio from agentscope.message import Msg + from reme.reme_light import ReMeLight @@ -127,9 +128,6 @@ 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/ 文件保留天数 ) @@ -176,7 +174,7 @@ async def main(): # 将消息添加到内存中以便估算 for msg in messages: await memory.add(msg) - token_stats = await memory.estimate_tokens() + 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']}") diff --git a/tests/light/test_summarizer.py b/tests/light/test_summarizer.py index 560efabd..bf8e3a78 100644 --- a/tests/light/test_summarizer.py +++ b/tests/light/test_summarizer.py @@ -2,7 +2,6 @@ import asyncio import datetime -import logging import tempfile from pathlib import Path @@ -13,14 +12,10 @@ from test_utils import ( get_formatter, get_token_counter, ) +from reme.core.utils import get_std_logger from reme.memory.file_based import Summarizer -# 配置日志输出到控制台 -logging.basicConfig( - level=logging.INFO, - format="%(asctime)s - %(name)s - %(levelname)s - %(message)s", -) -logger = logging.getLogger(__name__) +logger = get_std_logger() # ANSI 颜色码 diff --git a/tests/light/test_tool_result_compactor.py b/tests/light/test_tool_result_compactor.py index 059e626e..ef97558a 100644 --- a/tests/light/test_tool_result_compactor.py +++ b/tests/light/test_tool_result_compactor.py @@ -6,7 +6,6 @@ from datetime import datetime, timedelta 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 diff --git a/tests/light/test_utils.py b/tests/light/test_utils.py index f5cae021..fc93009e 100644 --- a/tests/light/test_utils.py +++ b/tests/light/test_utils.py @@ -1,50 +1,13 @@ """Test utilities for copaw tests.""" import os -from pathlib import Path -from typing import Any - -from loguru import logger - -_token_counter = None 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. - - 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 + return get_hf_token_counter() def get_dash_chat_model(model_name: str = "qwen3.5-plus"): @@ -54,8 +17,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 +26,5 @@ 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.""" - - 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 - - return ReMeChatFormatter(token_counter=get_token_counter()) + return OpenAIChatFormatter() From 22331ea9634ba7dd82ba5407895c6f81c311f162 Mon Sep 17 00:00:00 2001 From: "jinli.yl" Date: Fri, 6 Mar 2026 02:09:04 +0800 Subject: [PATCH 05/59] refactor(tests): update test configurations and remove unused test file --- tests/light/test_compactor.py | 17 +- tests/light/test_memory_formatter.py | 483 ---------------------- tests/light/test_summarizer.py | 22 +- tests/light/test_tool_result_compactor.py | 16 +- 4 files changed, 35 insertions(+), 503 deletions(-) delete mode 100644 tests/light/test_memory_formatter.py diff --git a/tests/light/test_compactor.py b/tests/light/test_compactor.py index 19891e29..dbd9952c 100644 --- a/tests/light/test_compactor.py +++ b/tests/light/test_compactor.py @@ -4,13 +4,13 @@ import asyncio from agentscope.message import Msg +from reme.core.utils import get_std_logger +from reme.memory.file_based import Compactor from test_utils import ( get_dash_chat_model, get_formatter, get_token_counter, ) -from reme.core.utils import get_std_logger -from reme.memory.file_based import Compactor logger = get_std_logger() @@ -96,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_llm=get_dash_chat_model(), + as_llm_formatter=get_formatter(), + language="zh", ) @@ -281,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_llm=get_dash_chat_model(), + as_llm_formatter=get_formatter(), ) messages = [ @@ -304,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_llm=get_dash_chat_model(), + as_llm_formatter=get_formatter(), ) messages = [ diff --git a/tests/light/test_memory_formatter.py b/tests/light/test_memory_formatter.py deleted file mode 100644 index 8b31718b..00000000 --- a/tests/light/test_memory_formatter.py +++ /dev/null @@ -1,483 +0,0 @@ -"""Tests for MemoryFormatter.""" - -# pylint: disable=W0212 - -from agentscope.message import Msg - -from test_utils import get_token_counter -from reme.core.utils import get_std_logger -from reme.memory.file_based import MemoryFormatter - -logger = get_std_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}") - - -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_summarizer.py b/tests/light/test_summarizer.py index bf8e3a78..a2f2d975 100644 --- a/tests/light/test_summarizer.py +++ b/tests/light/test_summarizer.py @@ -6,6 +6,7 @@ import tempfile from pathlib import Path from agentscope.message import Msg +from agentscope.tool import Toolkit from test_utils import ( get_dash_chat_model, @@ -14,6 +15,7 @@ from test_utils import ( ) from reme.core.utils import get_std_logger from reme.memory.file_based import Summarizer +from reme.memory.tools.file import FileIO logger = get_std_logger() @@ -95,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) + toolkit.register_tool_function(file_io.write) + toolkit.register_tool_function(file_io.edit) + return toolkit + + def create_summarizer(working_dir: str = None, memory_dir: str = "memory"): """Create a Summarizer instance for testing.""" if working_dir is None: @@ -109,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(), + toolkit=create_toolkit(working_dir), + as_llm=get_dash_chat_model(), + as_llm_formatter=get_formatter(), ), working_dir, ) @@ -190,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(), + 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 ef97558a..b6cb7c69 100644 --- a/tests/light/test_tool_result_compactor.py +++ b/tests/light/test_tool_result_compactor.py @@ -6,8 +6,8 @@ from datetime import datetime, timedelta 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 import ToolResultCompactor +from reme.core.utils import is_truncated def create_tool_result_msg(output: str | list, tool_name: str = "test_tool") -> Msg: @@ -51,7 +51,7 @@ class TestToolResultCompactor: _ = asyncio.run(op.call(messages=messages)) output = messages[0].content[0]["output"] - assert TRUNCATION_MARKER_START in output + assert is_truncated(output) assert "[Full content saved to:" in output # Verify file was created @@ -68,7 +68,7 @@ class TestToolResultCompactor: """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" + truncated_content = "head<<>>(100 chars omitted)<<>>tail" messages = [create_tool_result_msg(truncated_content)] asyncio.run(op.call(messages=messages)) @@ -86,7 +86,7 @@ class TestToolResultCompactor: asyncio.run(op.call(messages=messages)) text_block = messages[0].content[0]["output"][0] - assert TRUNCATION_MARKER_START in text_block["text"] + assert is_truncated(text_block["text"]) assert len(list(Path(tmpdir).glob("*.txt"))) == 1 def test_list_output_no_truncation_when_short(self): @@ -115,9 +115,9 @@ class TestToolResultCompactor: asyncio.run(op.call(messages=messages)) output = messages[0].content[0]["output"] - assert TRUNCATION_MARKER_START in output[0]["text"] + assert is_truncated(output[0]["text"]) assert output[1]["text"] == "short" # unchanged - assert TRUNCATION_MARKER_START in output[2]["text"] + assert is_truncated(output[2]["text"]) assert len(list(Path(tmpdir).glob("*.txt"))) == 2 def test_list_output_mixed_block_types(self): @@ -133,7 +133,7 @@ class TestToolResultCompactor: asyncio.run(op.call(messages=messages)) output = messages[0].content[0]["output"] - assert TRUNCATION_MARKER_START in output[0]["text"] + assert is_truncated(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 From 30278b4a4d4082e9e00da12db5d25b33688b1f3a Mon Sep 17 00:00:00 2001 From: "jinli.yl" Date: Fri, 6 Mar 2026 02:09:55 +0800 Subject: [PATCH 06/59] style(tests): reorder imports in test_compactor.py --- tests/light/test_compactor.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/tests/light/test_compactor.py b/tests/light/test_compactor.py index dbd9952c..8ae2a051 100644 --- a/tests/light/test_compactor.py +++ b/tests/light/test_compactor.py @@ -4,14 +4,16 @@ import asyncio from agentscope.message import Msg -from reme.core.utils import get_std_logger -from reme.memory.file_based import Compactor from test_utils import ( get_dash_chat_model, get_formatter, get_token_counter, ) +from reme.core.utils import get_std_logger +from reme.memory.file_based import Compactor + + logger = get_std_logger() From 32f9074235c3415dacc337005630bab9ece05493 Mon Sep 17 00:00:00 2001 From: "jinli.yl" Date: Fri, 6 Mar 2026 15:28:33 +0800 Subject: [PATCH 07/59] refactor(memory): update import paths and enhance message token counting --- reme/core/schema/as_msg_stat.py | 5 +- reme/memory/file_based/__init__.py | 6 +- reme/memory/file_based/as_msg_handler.py | 126 ++++-- .../{sub_agent => component}/__init__.py | 0 .../{sub_agent => component}/compactor.py | 0 .../{sub_agent => component}/compactor.yaml | 0 .../{sub_agent => component}/summarizer.py | 0 .../{sub_agent => component}/summarizer.yaml | 0 .../tool_result_compactor.py | 0 reme/reme_light.py | 87 +++- tests/light/test_reme_light.py | 319 +++++++------ tests/light/test_utils.py | 426 ++++++++++++++++++ 12 files changed, 758 insertions(+), 211 deletions(-) rename reme/memory/file_based/{sub_agent => component}/__init__.py (100%) rename reme/memory/file_based/{sub_agent => component}/compactor.py (100%) rename reme/memory/file_based/{sub_agent => component}/compactor.yaml (100%) rename reme/memory/file_based/{sub_agent => component}/summarizer.py (100%) rename reme/memory/file_based/{sub_agent => component}/summarizer.yaml (100%) rename reme/memory/file_based/{sub_agent => component}/tool_result_compactor.py (100%) diff --git a/reme/core/schema/as_msg_stat.py b/reme/core/schema/as_msg_stat.py index 4bb69f99..2e861863 100644 --- a/reme/core/schema/as_msg_stat.py +++ b/reme/core/schema/as_msg_stat.py @@ -50,8 +50,9 @@ class AsBlockStat(BaseModel): if self.block_type in ("tool_use", "tool_result"): if self.block_type == "tool_use": return f" - tool_call={self.tool_name} params={truncate_text(self.tool_input, max_length)}" - output = truncate_text(self.tool_output, max_length) - return f" - tool_result={self.tool_name} output={output}" if output else "" + else: + output = truncate_text(self.tool_output, max_length) + return f" - tool_result={self.tool_name} output={output}" if output else "" return "" diff --git a/reme/memory/file_based/__init__.py b/reme/memory/file_based/__init__.py index a1f5be73..2e01cc41 100644 --- a/reme/memory/file_based/__init__.py +++ b/reme/memory/file_based/__init__.py @@ -13,9 +13,9 @@ Components: from .as_msg_handler import AsMsgHandler from .reme_in_memory_memory import ReMeInMemoryMemory -from .sub_agent.compactor import Compactor -from .sub_agent.summarizer import Summarizer -from .sub_agent.tool_result_compactor import ToolResultCompactor +from .component.compactor import Compactor +from .component.summarizer import Summarizer +from .component.tool_result_compactor import ToolResultCompactor __all__ = [ "AsMsgHandler", diff --git a/reme/memory/file_based/as_msg_handler.py b/reme/memory/file_based/as_msg_handler.py index 802157ab..9db6cac2 100644 --- a/reme/memory/file_based/as_msg_handler.py +++ b/reme/memory/file_based/as_msg_handler.py @@ -39,21 +39,13 @@ class AsMsgHandler: logger.warning(f"Failed to count string tokens: {text}, e={e}") return estimated_tokens - @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. - """ + 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 + return output, 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: @@ -67,19 +59,23 @@ class AsMsgHandler: if block_type == "text": textual_parts.append(block.get("text", "")) + total_token_count += self.count_str_token(textual_parts[-1]) elif block_type in ["image", "audio", "video"]: source = block.get("source", {}) - url = source.get("url", "") - if url: - textual_parts.append(f"[{block_type}] {url}") + if source.get("type") == "base64": + data = source.get("data", "") + total_token_count += len(data) // 4 if data else 10 else: - textual_parts.append(f"[{block_type}]") + url = source.get("url", "") + total_token_count += 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 += self.count_str_token(file_path) else: logger.warning( @@ -94,17 +90,28 @@ class AsMsgHandler: 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) + return "\n".join(textual_parts), total_token_count 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=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.get_content_blocks(): + for block in message.content: block_type = block.get("type", "unknown") if block_type == "text": @@ -132,7 +139,6 @@ class AsMsgHandler: elif block_type in ("image", "audio", "video"): source = block.get("source", {}) url = source.get("url", "") - # For media, estimate fixed token cost or count URL if source.get("type") == "base64": data = source.get("data", "") token_count = len(data) // 4 if data else 10 @@ -149,7 +155,7 @@ class AsMsgHandler: elif block_type == "tool_use": tool_name = block.get("name", "") - tool_input = block.get("input", {}) + tool_input = block.get("raw_input", "") try: input_str = json.dumps(tool_input, ensure_ascii=False) except (TypeError, ValueError): @@ -168,8 +174,7 @@ class AsMsgHandler: elif block_type == "tool_result": tool_name = block.get("name", "") output = block.get("output", "") - formatted_output = self._format_tool_result_output(output) - token_count = self.count_str_token(formatted_output) + formatted_output, token_count = self._format_tool_result_output(output) blocks.append( AsBlockStat( block_type=block_type, @@ -191,6 +196,10 @@ class AsMsgHandler: metadata=message.metadata or {}, ) + def count_msgs_token(self, messages: list[Msg]) -> int: + """Count total token count of a list of messages.""" + return sum(self.stat_message(msg).total_tokens for msg in messages) + def format_msgs_to_str( self, messages: list[Msg], @@ -215,47 +224,71 @@ class AsMsgHandler: for i in range(len(messages) - 1, -1, -1): stat = self.stat_message(messages[i]) + formatted_content = stat.format(include_thinking=include_thinking) + content_token_count = self.count_str_token(formatted_content) - if total_token_count + stat.total_tokens > memory_compact_threshold: + if total_token_count + content_token_count > memory_compact_threshold: logger.info( "Skipping older messages: adding %d tokens would exceed threshold %d (current: %d)", - stat.total_tokens, + content_token_count, memory_compact_threshold, total_token_count, ) break - formatted_parts.append(stat.format(include_thinking=include_thinking)) - total_token_count += stat.total_tokens + 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 + def context_check( self, messages: list[Msg], memory_compact_threshold: int, memory_compact_reserve: int, - ) -> tuple[list[Msg], list[Msg]]: + ) -> tuple[list[Msg], list[Msg], bool]: """Check if context exceeds threshold and split messages accordingly. - This method checks if the total token count of messages exceeds the - memory_compact_threshold. If not, returns empty list and original messages. - If exceeded, uses memory_compact_reserve as the limit to keep messages - from the end, ensuring tool_use and tool_result blocks are properly paired. + 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 after compaction. + memory_compact_reserve: Token limit for messages to keep. Returns: - A tuple of (messages_to_compact, messages_to_keep): - - messages_to_compact: Older messages that need to be compacted + 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 [], [] + return [], [], True # Calculate total tokens and stats for all messages msg_stats: list[tuple[Msg, AsMsgStat]] = [] @@ -265,9 +298,9 @@ class AsMsgHandler: msg_stats.append((msg, stat)) total_tokens += stat.total_tokens - # If total tokens don't exceed threshold, no compaction needed - if total_tokens <= memory_compact_threshold: - return [], messages + # 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 @@ -348,15 +381,20 @@ class AsMsgHandler: 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( "Context check result: %d messages to compact, %d messages to keep, " - "total tokens: %d, threshold: %d, reserve: %d, kept tokens: %d", + "total tokens: %d, threshold: %d, reserve: %d, kept tokens: %d, " + "tools_aligned: %s", len(messages_to_compact), len(messages_to_keep), total_tokens, memory_compact_threshold, memory_compact_reserve, accumulated_tokens, + tools_aligned, ) - return messages_to_compact, messages_to_keep + return messages_to_compact, messages_to_keep, tools_aligned diff --git a/reme/memory/file_based/sub_agent/__init__.py b/reme/memory/file_based/component/__init__.py similarity index 100% rename from reme/memory/file_based/sub_agent/__init__.py rename to reme/memory/file_based/component/__init__.py diff --git a/reme/memory/file_based/sub_agent/compactor.py b/reme/memory/file_based/component/compactor.py similarity index 100% rename from reme/memory/file_based/sub_agent/compactor.py rename to reme/memory/file_based/component/compactor.py diff --git a/reme/memory/file_based/sub_agent/compactor.yaml b/reme/memory/file_based/component/compactor.yaml similarity index 100% rename from reme/memory/file_based/sub_agent/compactor.yaml rename to reme/memory/file_based/component/compactor.yaml diff --git a/reme/memory/file_based/sub_agent/summarizer.py b/reme/memory/file_based/component/summarizer.py similarity index 100% rename from reme/memory/file_based/sub_agent/summarizer.py rename to reme/memory/file_based/component/summarizer.py diff --git a/reme/memory/file_based/sub_agent/summarizer.yaml b/reme/memory/file_based/component/summarizer.yaml similarity index 100% rename from reme/memory/file_based/sub_agent/summarizer.yaml rename to reme/memory/file_based/component/summarizer.yaml diff --git a/reme/memory/file_based/sub_agent/tool_result_compactor.py b/reme/memory/file_based/component/tool_result_compactor.py similarity index 100% rename from reme/memory/file_based/sub_agent/tool_result_compactor.py rename to reme/memory/file_based/component/tool_result_compactor.py diff --git a/reme/reme_light.py b/reme/reme_light.py index 82a11850..5266ae76 100644 --- a/reme/reme_light.py +++ b/reme/reme_light.py @@ -26,7 +26,7 @@ from agentscope.tool import Toolkit, ToolResponse from .config import ReMeConfigParser from .core import Application from .core.utils import get_hf_token_counter, get_std_logger -from .memory.file_based import Compactor, Summarizer, ToolResultCompactor, ReMeInMemoryMemory +from .memory.file_based import Compactor, Summarizer, ToolResultCompactor, ReMeInMemoryMemory, AsMsgHandler from .memory.tools import MemorySearch from .memory.tools.file import FileIO @@ -258,6 +258,75 @@ class ReMeLight(Application): task = asyncio.create_task(self.summary_memory(messages=messages, **kwargs)) self.summary_tasks.append(task) + 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", + token_counter: HuggingFaceTokenCounter | None = None, + 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.""" + if token_counter is None: + token_counter = get_hf_token_counter() + + msg_handler = AsMsgHandler(token_counter=token_counter) + + system_token_count = msg_handler.count_str_token(system_prompt) + compressed_token_count = 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 = msg_handler.context_check( + messages=messages, + memory_compact_threshold=left_compact_threshold, + memory_compact_reserve=memory_compact_reserve, + ) + + 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, + token_counter=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, + token_counter=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.""" result = "" @@ -289,6 +358,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" @@ -334,12 +404,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 diff --git a/tests/light/test_reme_light.py b/tests/light/test_reme_light.py index 49102826..e864cd74 100644 --- a/tests/light/test_reme_light.py +++ b/tests/light/test_reme_light.py @@ -1,195 +1,194 @@ -"""测试 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", # 记忆文件存储目录 tool_result_threshold=1000, # 超过此字符数的工具输出自动转存 retention_days=7, # tool_result/ 文件保留天数 ) + 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 = 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 = msg_handler.count_msgs_token(messages) + messages_after_step1 = await reme.compact_tool_result(messages) + tokens_after = 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 = 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 = 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 使用...") - memory = reme.get_in_memory_memory() - # 将消息添加到内存中以便估算 + # 重新获取原始消息 + messages = build_sample_messages(include_large_tool_result=True) + tokens_before = 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 = 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 = msg_handler.count_msgs_token(processed_messages) + compressed_summary_tokens = 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 = ReMeLight.get_in_memory_memory() 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']}") + 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 后: {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_utils.py b/tests/light/test_utils.py index fc93009e..19740fe6 100644 --- a/tests/light/test_utils.py +++ b/tests/light/test_utils.py @@ -2,6 +2,10 @@ import os +from agentscope.message import Msg, ThinkingBlock, TextBlock, ToolUseBlock, ToolResultBlock + +from reme.memory.file_based import AsMsgHandler + def get_token_counter(): """Get HF token counter instance.""" @@ -10,6 +14,11 @@ def get_token_counter(): return get_hf_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"): """Get DashScope chat model instance.""" from agentscope.model import OpenAIChatModel @@ -28,3 +37,420 @@ def get_formatter(): from agentscope.formatter import OpenAIChatFormatter return OpenAIChatFormatter() + + +def generate_large_code_content(target_tokens: int = 50000) -> str: + """生成大量代码内容，用于测试超长 tool_result。 + + Args: + target_tokens: 目标 token 数（约 4 字符/token） + + 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 From c1e9faaeb28129c93586f8d83b0cc7ed1a4e9fff Mon Sep 17 00:00:00 2001 From: "jinli.yl" Date: Fri, 6 Mar 2026 15:36:38 +0800 Subject: [PATCH 08/59] feat(docs): update README with new pre_reasoning_hook method and enhanced examples --- README.md | 54 ++++++++++++++++++++++++++++++++++------------------ README_ZH.md | 54 ++++++++++++++++++++++++++++++++++------------------ 2 files changed, 72 insertions(+), 36 deletions(-) diff --git a/README.md b/README.md index 9c69d232..10907911 100644 --- a/README.md +++ b/README.md @@ -67,14 +67,15 @@ working_dir/ 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 | +| `pre_reasoning_hook` | 🔄 Pre-reasoning hook | Auto compact tool results + generate summary + async trigger memory summarization task | | `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 | +| `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 (static method) | --- @@ -108,38 +109,52 @@ from reme.reme_light import ReMeLight async def main(): - 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 - ) + # Initialize ReMeLight + reme = ReMeLight() await reme.start() - messages = [...] + messages = [...] # Conversation message list # 1. Compact oversized tool outputs (prevent tool results from overflowing 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="") + # 2. Compact history to structured summary (can pass previous summary for incremental update) + summary = await reme.compact_memory( + messages=messages, + previous_summary="", + max_input_length=128000, # Model context window (tokens) + compact_ratio=0.7, # Trigger compaction when reaching max_input_length * 0.7 + language="zh", # Summary language (zh / "") + ) # 3. Submit async summary task in background (non-blocking, writes to memory/YYYY-MM-DD.md) reme.add_async_summary_task(messages=messages) - # 4. Semantic memory search (Vector + BM25 hybrid retrieval) + # 4. Pre-reasoning hook (auto compact tool results + generate summary) + 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, + ) + + # 5. 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() + # 6. Get in-memory instance (static method, manages single conversation context) + memory = ReMeLight.get_in_memory_memory() + 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"Estimated total tokens: {token_stats['estimated_tokens']}") - # 6. Wait for background tasks before closing + # 7. Wait for background tasks before closing summary_result = await reme.await_summary_tasks() # Close ReMeLight @@ -150,6 +165,9 @@ if __name__ == "__main__": asyncio.run(main()) ``` +> 📂 Full example code: [test_reme_light.py](tests/light/test_reme_light.py) +> 📋 Example output: [test_reme_light.log](tests/light/test_reme_light.log) (223,838 tokens → 1,105 tokens, 99.5% compression ratio) + ### File-Based ReMeLight Memory System Architecture [CoPaw MemoryManager](https://github.com/agentscope-ai/CoPaw/blob/main/src/copaw/agents/memory/memory_manager.py) diff --git a/README_ZH.md b/README_ZH.md index b45b261c..4757e59e 100644 --- a/README_ZH.md +++ b/README_ZH.md @@ -66,9 +66,10 @@ working_dir/ | `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/`，消息中保留文件引用 | | +| `compact_tool_result` | ✂️ 压缩超长工具输出 | [ToolResultCompactor](reme/memory/file_based/tool_result_compactor.py) — 截断并转存到 `tool_result/`，消息中保留文件引用 | +| `pre_reasoning_hook` | 🔄 推理前预处理钩子 | 自动压缩工具结果 + 生成摘要 + 异步触发记忆总结任务 | | `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 感知的内存管理，支持压缩摘要和状态序列化 | +| `get_in_memory_memory` | 🗂️ 创建会话内存实例 | [ReMeInMemoryMemory](reme/memory/file_based/reme_in_memory_memory.py) — Token 感知的内存管理，支持压缩摘要和状态序列化（静态方法） | --- @@ -102,38 +103,52 @@ from reme.reme_light import ReMeLight async def main(): - 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/ 文件保留天数 - ) + # 初始化 ReMeLight + reme = ReMeLight() await reme.start() - messages = [...] + messages = [...] # 对话消息列表 # 1. 压缩超长工具输出（防止工具结果撑爆上下文） messages = await reme.compact_tool_result(messages) - # 2. 将历史对话压缩为结构化摘要（触发时机：上下文接近上限），可传入上轮摘要，实现增量更新 - summary = await reme.compact_memory(messages=messages, previous_summary="") + # 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. 后台异步提交摘要任务（不阻塞对话，摘要写入 memory/YYYY-MM-DD.md） reme.add_async_summary_task(messages=messages) - # 4. 语义搜索记忆（向量 + BM25 混合检索） + # 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, + ) + + # 5. 语义搜索记忆（向量 + 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() + # 6. 获取会话内存实例（静态方法，管理单次对话的上下文） + memory = ReMeLight.get_in_memory_memory() + 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. 关闭前等待后台任务完成 + # 7. 关闭前等待后台任务完成 summary_result = await reme.await_summary_tasks() # 关闭 ReMeLight @@ -144,6 +159,9 @@ if __name__ == "__main__": asyncio.run(main()) ``` +> 📂 完整示例代码：[test_reme_light.py](tests/light/test_reme_light.py) +> 📋 运行结果示例：[test_reme_light.log](tests/light/test_reme_light.log)（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) 继承 From 2af9f329d263ce182683b8331a940b9588b95bb7 Mon Sep 17 00:00:00 2001 From: "jinli.yl" Date: Fri, 6 Mar 2026 15:43:24 +0800 Subject: [PATCH 09/59] docs(readme): update mermaid graph syntax in Chinese documentation --- README.md | 20 ++++++++++---------- README_ZH.md | 20 ++++++++++---------- 2 files changed, 20 insertions(+), 20 deletions(-) diff --git a/README.md b/README.md index 10907911..63bafdd0 100644 --- a/README.md +++ b/README.md @@ -175,18 +175,18 @@ inherits `ReMeLight` and integrates memory capabilities into the Agent reasoning ```mermaid graph TB - CoPaw["CoPaw MemoryManager\n(inherits ReMeLight)"] -->|pre_reasoning hook| Hook[MemoryCompactionHook] + CoPaw["CoPaw MemoryManager
(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] + ReMeLight --> CompactMemory[compact_memory
History compaction] + ReMeLight --> SummaryMemory[summary_memory
Write memory to files] + ReMeLight --> CompactToolResult[compact_tool_result
Oversized tool output compaction] + ReMeLight --> MemSearch[memory_search
Semantic search] + ReMeLight --> InMemory[get_in_memory_memory
ReMeInMemoryMemory] + CompactMemory --> Compactor[Compactor
ReActAgent] + SummaryMemory --> Summarizer[Summarizer
ReActAgent + file tools] + CompactToolResult --> ToolResultCompactor[ToolResultCompactor
Truncate + save to file] + Summarizer --> FileIO[FileIO
read / write / edit] FileIO --> MemoryFiles[memory/YYYY-MM-DD.md] ToolResultCompactor --> ToolResultFiles[tool_result/*.txt] MemoryFiles -.->|File change| FileWatcher[Async File Watcher] diff --git a/README_ZH.md b/README_ZH.md index 4757e59e..637eb215 100644 --- a/README_ZH.md +++ b/README_ZH.md @@ -169,18 +169,18 @@ if __name__ == "__main__": ```mermaid graph TB - CoPaw["CoPaw MemoryManager\n(继承 ReMeLight)"] -->|pre_reasoning hook| Hook[MemoryCompactionHook] + CoPaw["CoPaw MemoryManager
(继承 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] + ReMeLight --> CompactMemory[compact_memory
历史对话压缩] + ReMeLight --> SummaryMemory[summary_memory
记忆写入文件] + ReMeLight --> CompactToolResult[compact_tool_result
超长工具输出压缩] + ReMeLight --> MemSearch[memory_search
语义搜索] + ReMeLight --> InMemory[get_in_memory_memory
ReMeInMemoryMemory] + CompactMemory --> Compactor[Compactor
ReActAgent] + SummaryMemory --> Summarizer[Summarizer
ReActAgent + 文件工具] + CompactToolResult --> ToolResultCompactor[ToolResultCompactor
截断 + 转存文件] + Summarizer --> FileIO[FileIO
read / write / edit] FileIO --> MemoryFiles[memory/YYYY-MM-DD.md] ToolResultCompactor --> ToolResultFiles[tool_result/*.txt] MemoryFiles -.->|文件变更| FileWatcher[异步文件监控] From 46ffe42a409e884b65513645d4b7efba1deba551 Mon Sep 17 00:00:00 2001 From: "jinli.yl" Date: Fri, 6 Mar 2026 15:50:47 +0800 Subject: [PATCH 10/59] feat(config): update ReMeLight initialization with default configurations --- README.md | 7 +++++-- README_ZH.md | 7 +++++-- reme/config/light.yaml | 1 - tests/light/test_reme_light.py | 6 +++--- 4 files changed, 13 insertions(+), 8 deletions(-) diff --git a/README.md b/README.md index 63bafdd0..27acedf5 100644 --- a/README.md +++ b/README.md @@ -97,7 +97,6 @@ pip install -e ".[light]" | `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` | #### Python Usage @@ -110,7 +109,11 @@ from reme.reme_light import ReMeLight async def main(): # Initialize ReMeLight - reme = ReMeLight() + reme = ReMeLight( + 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}, + ) await reme.start() messages = [...] # Conversation message list diff --git a/README_ZH.md b/README_ZH.md index 637eb215..ded1d8d1 100644 --- a/README_ZH.md +++ b/README_ZH.md @@ -91,7 +91,6 @@ pip install -e ".[light]" | `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` | #### Python使用 @@ -104,7 +103,11 @@ from reme.reme_light import ReMeLight async def main(): # 初始化 ReMeLight - reme = ReMeLight() + reme = ReMeLight( + 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}, + ) await reme.start() messages = [...] # 对话消息列表 diff --git a/reme/config/light.yaml b/reme/config/light.yaml index 89df5f4b..bc85c10d 100644 --- a/reme/config/light.yaml +++ b/reme/config/light.yaml @@ -23,7 +23,6 @@ file_stores: embedding_model: default store_name: "reme" - file_watchers: default: backend: full diff --git a/tests/light/test_reme_light.py b/tests/light/test_reme_light.py index e864cd74..4ec8f52a 100644 --- a/tests/light/test_reme_light.py +++ b/tests/light/test_reme_light.py @@ -31,9 +31,9 @@ async def main(): # 初始化 ReMeLight reme = ReMeLight( - working_dir=".reme", # 记忆文件存储目录 - 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}, ) logging.getLogger("reme").setLevel(logging.WARNING) await reme.start() From ce53bc051a04ac2abef39abd223b44420765dd18 Mon Sep 17 00:00:00 2001 From: "jinli.yl" Date: Fri, 6 Mar 2026 16:04:01 +0800 Subject: [PATCH 11/59] refactor(embedding): update environment variable names for API key and base URL --- reme/core/embedding/base_embedding_model.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/reme/core/embedding/base_embedding_model.py b/reme/core/embedding/base_embedding_model.py index 36c6def7..78a91b8b 100644 --- a/reme/core/embedding/base_embedding_model.py +++ b/reme/core/embedding/base_embedding_model.py @@ -81,12 +81,12 @@ class BaseEmbeddingModel(ABC): @property def api_key(self) -> str | None: """Get API key from environment variable.""" - return os.getenv("REME_EMBEDDING_API_KEY") or self._api_key + return os.getenv("EMBEDDING_API_KEY") or self._api_key @property def base_url(self) -> str | None: """Get base URL from environment variable.""" - return os.getenv("REME_EMBEDDING_BASE_URL") or self._base_url + return os.getenv("EMBEDDING_BASE_URL") or self._base_url def _truncate_text(self, text: str) -> str: """Truncate text to max_input_length if it exceeds the limit.""" From 65971bafe3221ae66e75f410324841f0c5582af5 Mon Sep 17 00:00:00 2001 From: zouyingcao <57442064+zouyingcao@users.noreply.github.com> Date: Fri, 6 Mar 2026 16:11:39 +0800 Subject: [PATCH 12/59] Update: check the code&docs for evaluation on bfcl&appworld (#141) * fix: df.columns bug * fix: await for asynchronous method * update: docs for bfcl&appworld quickstart * update: benchmark/bfcl for new version quickstart * slightly revise bfcl cookbook * update for pre-commit * handle boolean flags in split_into_trainval.py * fix typo in faq.md --- benchmark/bfcl/default_ids.py | 206 ++++++++++++++++++ benchmark/bfcl/init_task_memory_pool.py | 27 +-- benchmark/bfcl/local_file_to_library.py | 30 --- benchmark/bfcl/requirements.txt | 3 +- benchmark/bfcl/run_bfcl.py | 2 +- benchmark/bfcl/run_exp_statistic.py | 2 +- benchmark/bfcl/split_into_trainval.py | 49 ++++- docs/cookbook/appworld/quickstart.md | 38 ++-- docs/cookbook/bfcl/quickstart.md | 122 ++++++----- docs/cookbook/faq.md | 11 +- reme/config/service.yaml | 3 +- .../summary/comparative_extraction.py | 10 +- 12 files changed, 352 insertions(+), 151 deletions(-) create mode 100644 benchmark/bfcl/default_ids.py delete mode 100644 benchmark/bfcl/local_file_to_library.py 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/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/docs/cookbook/appworld/quickstart.md b/docs/cookbook/appworld/quickstart.md index 78f6a1a8..45ea0d78 100644 --- a/docs/cookbook/appworld/quickstart.md +++ b/docs/cookbook/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/docs/cookbook/bfcl/quickstart.md b/docs/cookbook/bfcl/quickstart.md index c814b0ac..c80ef75f 100644 --- a/docs/cookbook/bfcl/quickstart.md +++ b/docs/cookbook/bfcl/quickstart.md @@ -7,99 +7,98 @@ This guide helps you quickly set up and run BFCL experiments with ReMe integrati ### 1. BFCL installation -#### clone the repository +#### 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 gorilla/berkeley-function-call-leaderboard +cd 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 . +cd ../.. pip install -r requirements.txt ``` #### Move the dataset to the data folder under bfcl ```bash -cp -r bfcl_eval/data {/path/to/bfcl/data} +cp -r gorilla/berkeley-function-call-leaderboard/bfcl_eval/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. +#### Preprocess the data to get the suitable data format +```bash +python preprocess.py +``` -### 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`): +**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 run_bfcl.py +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 ``` -**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 +### 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 -reme \ +reme2 \ 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 + 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 ``` -and then init the task memory pool: +
+Option: init the task memory pool from scratch -```bash -python init_task_memory_pool.py -``` +- First, collect agent trajectories on training data set without task memory: -**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) + ```bash + # important: num_runs = 8, use_memory = False, experiment_suffix="wo-memory", data_path="data/multiturn_data_base_train.jsonl" + python run_bfcl.py + ``` -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`). +- 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 + ``` -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" - }' -``` + > 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", + }' + ``` -### 4. Run Experiments on Validation Set +- 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: @@ -108,6 +107,15 @@ Run you can compare agent performance on the validation set with task memory (`u 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 @@ -116,6 +124,6 @@ python run_exp_statistic.py **What this script does:** - Processes all result files in `./exp_result/` -- Calculates best@k metrics for different k values +- 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/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/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/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) From a0d3120d53684b671aef3cec02b24efb120263a9 Mon Sep 17 00:00:00 2001 From: "jinli.yl" Date: Fri, 6 Mar 2026 16:18:50 +0800 Subject: [PATCH 13/59] fix(tests): update context check tests to handle additional return value --- reme_old_doc/__init__.py | 0 {reme/extension => test}/cli/__init__.py | 0 {reme/extension => test}/cli/fb_cli.py | 0 {reme/extension => test}/cli/fb_cli.yaml | 0 {reme/extension => test}/cli/fb_compactor.py | 0 .../extension => test}/cli/fb_compactor.yaml | 0 .../cli/fb_context_checker.py | 0 {reme/extension => test}/cli/fb_summarizer.py | 0 .../extension => test}/cli/fb_summarizer.yaml | 0 {reme/extension => test}/reme_cli.py | 0 tests/light/test_context_check.py | 66 +++++++++---------- tests/light/test_format_msgs_to_str.py | 8 ++- 12 files changed, 38 insertions(+), 36 deletions(-) create mode 100644 reme_old_doc/__init__.py rename {reme/extension => test}/cli/__init__.py (100%) rename {reme/extension => test}/cli/fb_cli.py (100%) rename {reme/extension => test}/cli/fb_cli.yaml (100%) rename {reme/extension => test}/cli/fb_compactor.py (100%) rename {reme/extension => test}/cli/fb_compactor.yaml (100%) rename {reme/extension => test}/cli/fb_context_checker.py (100%) rename {reme/extension => test}/cli/fb_summarizer.py (100%) rename {reme/extension => test}/cli/fb_summarizer.yaml (100%) rename {reme/extension => test}/reme_cli.py (100%) diff --git a/reme_old_doc/__init__.py b/reme_old_doc/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/reme/extension/cli/__init__.py b/test/cli/__init__.py similarity index 100% rename from reme/extension/cli/__init__.py rename to test/cli/__init__.py diff --git a/reme/extension/cli/fb_cli.py b/test/cli/fb_cli.py similarity index 100% rename from reme/extension/cli/fb_cli.py rename to test/cli/fb_cli.py diff --git a/reme/extension/cli/fb_cli.yaml b/test/cli/fb_cli.yaml similarity index 100% rename from reme/extension/cli/fb_cli.yaml rename to test/cli/fb_cli.yaml diff --git a/reme/extension/cli/fb_compactor.py b/test/cli/fb_compactor.py similarity index 100% rename from reme/extension/cli/fb_compactor.py rename to test/cli/fb_compactor.py diff --git a/reme/extension/cli/fb_compactor.yaml b/test/cli/fb_compactor.yaml similarity index 100% rename from reme/extension/cli/fb_compactor.yaml rename to test/cli/fb_compactor.yaml diff --git a/reme/extension/cli/fb_context_checker.py b/test/cli/fb_context_checker.py similarity index 100% rename from reme/extension/cli/fb_context_checker.py rename to test/cli/fb_context_checker.py diff --git a/reme/extension/cli/fb_summarizer.py b/test/cli/fb_summarizer.py similarity index 100% rename from reme/extension/cli/fb_summarizer.py rename to test/cli/fb_summarizer.py diff --git a/reme/extension/cli/fb_summarizer.yaml b/test/cli/fb_summarizer.yaml similarity index 100% rename from reme/extension/cli/fb_summarizer.yaml rename to test/cli/fb_summarizer.yaml diff --git a/reme/extension/reme_cli.py b/test/reme_cli.py similarity index 100% rename from reme/extension/reme_cli.py rename to test/reme_cli.py diff --git a/tests/light/test_context_check.py b/tests/light/test_context_check.py index f65f961e..872d6e2d 100644 --- a/tests/light/test_context_check.py +++ b/tests/light/test_context_check.py @@ -220,7 +220,7 @@ def test_empty_messages(): handler = create_handler() messages = [] threshold, reserve = 1000, 500 - to_compact, to_keep = handler.context_check( + to_compact, to_keep, _ = handler.context_check( messages=messages, memory_compact_threshold=threshold, memory_compact_reserve=reserve, @@ -240,7 +240,7 @@ def test_below_threshold_returns_all(): create_user_msg("How are you?"), ] threshold, reserve = 10000, 5000 - to_compact, to_keep = handler.context_check( + to_compact, to_keep, _ = handler.context_check( messages=messages, memory_compact_threshold=threshold, # Very high threshold memory_compact_reserve=reserve, @@ -271,7 +271,7 @@ def test_above_threshold_triggers_compaction(): create_assistant_msg("Fourth message " * 100), ] threshold, reserve = 100, 200 - to_compact, to_keep = handler.context_check( + to_compact, to_keep, _ = handler.context_check( messages=messages, memory_compact_threshold=threshold, # Low threshold to trigger compaction memory_compact_reserve=reserve, @@ -302,7 +302,7 @@ def test_message_order_preserved(): create_user_msg("Fifth " * 10), ] threshold, reserve = 100, 150 - to_compact, to_keep = handler.context_check( + to_compact, to_keep, _ = handler.context_check( messages=messages, memory_compact_threshold=threshold, # Low threshold memory_compact_reserve=reserve, @@ -333,7 +333,7 @@ def test_single_message_below_threshold(): handler = create_handler() messages = [create_user_msg("Short message")] threshold, reserve = 1000, 500 - to_compact, to_keep = handler.context_check( + to_compact, to_keep, _ = handler.context_check( messages=messages, memory_compact_threshold=threshold, memory_compact_reserve=reserve, @@ -358,7 +358,7 @@ def test_single_message_above_threshold(): long_content = "Very long message " * 1000 messages = [create_user_msg(long_content)] threshold, reserve = 10, 5 - to_compact, to_keep = handler.context_check( + to_compact, to_keep, _ = handler.context_check( messages=messages, memory_compact_threshold=threshold, # Very low threshold memory_compact_reserve=reserve, # Even lower reserve @@ -386,7 +386,7 @@ def test_reserve_zero(): create_assistant_msg("Hi there!"), ] threshold, reserve = 1, 0 - to_compact, to_keep = handler.context_check( + to_compact, to_keep, _ = handler.context_check( messages=messages, memory_compact_threshold=threshold, # Trigger compaction memory_compact_reserve=reserve, # Zero reserve @@ -403,7 +403,7 @@ def test_threshold_zero(): handler = create_handler() messages = [create_user_msg("A")] # Minimal message threshold, reserve = 0, 1000 - to_compact, to_keep = handler.context_check( + to_compact, to_keep, _ = handler.context_check( messages=messages, memory_compact_threshold=threshold, # Zero threshold - always triggers memory_compact_reserve=reserve, @@ -426,7 +426,7 @@ def test_exact_threshold_boundary(): threshold, reserve = exact_tokens, exact_tokens # Test at exact boundary - to_compact, to_keep = handler.context_check( + to_compact, to_keep, _ = handler.context_check( messages=messages, memory_compact_threshold=threshold, # Exactly at boundary memory_compact_reserve=reserve, @@ -454,7 +454,7 @@ def test_reserve_larger_than_threshold(): create_assistant_msg("Message two " * 20), ] threshold, reserve = 50, 10000 - to_compact, to_keep = handler.context_check( + to_compact, to_keep, _ = handler.context_check( messages=messages, memory_compact_threshold=threshold, # Low threshold memory_compact_reserve=reserve, # High reserve @@ -489,7 +489,7 @@ def test_tool_use_result_paired(): create_assistant_msg("The tool returned results"), ] threshold, reserve = 50, 1000 - to_compact, to_keep = handler.context_check( + to_compact, to_keep, _ = handler.context_check( messages=messages, memory_compact_threshold=threshold, # Trigger compaction memory_compact_reserve=reserve, # Enough for tool pair @@ -522,7 +522,7 @@ def test_tool_use_without_result(): create_assistant_msg("Something happened"), ] threshold, reserve = 10, 1000 - to_compact, to_keep = handler.context_check( + to_compact, to_keep, _ = handler.context_check( messages=messages, memory_compact_threshold=threshold, memory_compact_reserve=reserve, @@ -550,7 +550,7 @@ def test_tool_result_without_use(): create_assistant_msg("Got it"), ] threshold, reserve = 10, 1000 - to_compact, to_keep = handler.context_check( + to_compact, to_keep, _ = handler.context_check( messages=messages, memory_compact_threshold=threshold, memory_compact_reserve=reserve, @@ -583,7 +583,7 @@ def test_multiple_tool_pairs(): create_assistant_msg("All done"), ] threshold, reserve = 50, 500 - to_compact, to_keep = handler.context_check( + to_compact, to_keep, _ = handler.context_check( messages=messages, memory_compact_threshold=threshold, memory_compact_reserve=reserve, @@ -630,7 +630,7 @@ def test_tool_dependency_causes_extra_inclusion(): create_assistant_msg("End"), # Small ] threshold, reserve = 100, 500 - to_compact, to_keep = handler.context_check( + to_compact, to_keep, _ = handler.context_check( messages=messages, memory_compact_threshold=threshold, # Trigger compaction memory_compact_reserve=reserve, # Medium reserve @@ -671,7 +671,7 @@ def test_tool_dependency_exceeds_reserve(): create_assistant_msg("Last message"), ] threshold, reserve = 10, 100 - to_compact, to_keep = handler.context_check( + to_compact, to_keep, _ = handler.context_check( messages=messages, memory_compact_threshold=threshold, # Trigger compaction memory_compact_reserve=reserve, # Small reserve - can't fit the pair @@ -715,7 +715,7 @@ def test_interleaved_tool_pairs(): create_assistant_msg("Both done"), ] threshold, reserve = 50, 1000 - to_compact, to_keep = handler.context_check( + to_compact, to_keep, _ = handler.context_check( messages=messages, memory_compact_threshold=threshold, memory_compact_reserve=reserve, @@ -756,7 +756,7 @@ def test_message_with_empty_content(): create_assistant_msg("Response"), ] threshold, reserve = 1000, 500 - to_compact, to_keep = handler.context_check( + to_compact, to_keep, _ = handler.context_check( messages=messages, memory_compact_threshold=threshold, memory_compact_reserve=reserve, @@ -782,7 +782,7 @@ def test_message_with_whitespace_only(): create_assistant_msg("Response"), ] threshold, reserve = 1000, 500 - to_compact, to_keep = handler.context_check( + to_compact, to_keep, _ = handler.context_check( messages=messages, memory_compact_threshold=threshold, memory_compact_reserve=reserve, @@ -806,7 +806,7 @@ def test_very_long_single_message(): huge_content = "x" * 100000 # Very long messages = [create_user_msg(huge_content)] threshold, reserve = 100, 1000 - to_compact, to_keep = handler.context_check( + to_compact, to_keep, _ = handler.context_check( messages=messages, memory_compact_threshold=threshold, memory_compact_reserve=reserve, @@ -830,7 +830,7 @@ def 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 = handler.context_check( + to_compact, to_keep, _ = handler.context_check( messages=messages, memory_compact_threshold=threshold, # Low threshold memory_compact_reserve=reserve, @@ -859,7 +859,7 @@ def test_unicode_content(): create_user_msg("日本語テスト 🇯🇵"), ] threshold, reserve = 1000, 500 - to_compact, to_keep = handler.context_check( + to_compact, to_keep, _ = handler.context_check( messages=messages, memory_compact_threshold=threshold, memory_compact_reserve=reserve, @@ -877,7 +877,7 @@ def test_special_characters_content(): create_assistant_msg("More: \n\r\t\0 nulls and newlines"), ] threshold, reserve = 1000, 500 - to_compact, to_keep = handler.context_check( + to_compact, to_keep, _ = handler.context_check( messages=messages, memory_compact_threshold=threshold, memory_compact_reserve=reserve, @@ -912,7 +912,7 @@ def test_all_messages_fit_exactly_in_reserve(): total = sum(handler.stat_message(m).total_tokens for m in messages) threshold, reserve = total - 1, total - to_compact, to_keep = handler.context_check( + to_compact, to_keep, _ = handler.context_check( messages=messages, memory_compact_threshold=threshold, # Just below total to trigger memory_compact_reserve=reserve, # Exactly fits all @@ -945,7 +945,7 @@ def test_first_message_only_compacted(): tiny_msg_tokens = handler.stat_message(messages[2]).total_tokens threshold, reserve = 50, small_msg_tokens + tiny_msg_tokens + 10 - to_compact, to_keep = handler.context_check( + to_compact, to_keep, _ = handler.context_check( messages=messages, memory_compact_threshold=threshold, # Low to trigger memory_compact_reserve=reserve, # Fits last 2 @@ -976,7 +976,7 @@ def test_last_message_only_kept(): tiny_tokens = handler.stat_message(messages[2]).total_tokens threshold, reserve = 10, tiny_tokens + 5 - to_compact, to_keep = handler.context_check( + to_compact, to_keep, _ = handler.context_check( messages=messages, memory_compact_threshold=threshold, memory_compact_reserve=reserve, # Only fits last message @@ -1005,7 +1005,7 @@ def test_all_messages_compacted(): create_assistant_msg("Large message " * 100), ] threshold, reserve = 10, 1 - to_compact, to_keep = handler.context_check( + to_compact, to_keep, _ = handler.context_check( messages=messages, memory_compact_threshold=threshold, # Trigger compaction memory_compact_reserve=reserve, # Too small for anything @@ -1039,7 +1039,7 @@ def test_system_message(): create_assistant_msg("Hi there!"), ] threshold, reserve = 1000, 500 - to_compact, to_keep = handler.context_check( + to_compact, to_keep, _ = handler.context_check( messages=messages, memory_compact_threshold=threshold, memory_compact_reserve=reserve, @@ -1061,7 +1061,7 @@ def test_mixed_roles(): Msg(name="helper", role="assistant", content="Another assistant message"), ] threshold, reserve = 1000, 500 - to_compact, to_keep = handler.context_check( + to_compact, to_keep, _ = handler.context_check( messages=messages, memory_compact_threshold=threshold, memory_compact_reserve=reserve, @@ -1085,7 +1085,7 @@ def test_tool_use_with_empty_id(): create_assistant_msg("Done"), ] threshold, reserve = 10, 1000 - to_compact, to_keep = handler.context_check( + to_compact, to_keep, _ = handler.context_check( messages=messages, memory_compact_threshold=threshold, memory_compact_reserve=reserve, @@ -1113,7 +1113,7 @@ def test_tool_result_with_empty_id(): create_assistant_msg("Noted"), ] threshold, reserve = 10, 1000 - to_compact, to_keep = handler.context_check( + to_compact, to_keep, _ = handler.context_check( messages=messages, memory_compact_threshold=threshold, memory_compact_reserve=reserve, @@ -1142,7 +1142,7 @@ def test_duplicate_tool_ids(): create_tool_result_msg("call_dup", "tool_b", "Result B"), ] threshold, reserve = 10, 1000 - to_compact, to_keep = handler.context_check( + to_compact, to_keep, _ = handler.context_check( messages=messages, memory_compact_threshold=threshold, memory_compact_reserve=reserve, @@ -1181,7 +1181,7 @@ def test_message_with_multiple_tool_blocks(): create_tool_result_msg("call_3", "tool3", "Result 3"), ] threshold, reserve = 10, 2000 - to_compact, to_keep = handler.context_check( + to_compact, to_keep, _ = handler.context_check( messages=messages, memory_compact_threshold=threshold, memory_compact_reserve=reserve, diff --git a/tests/light/test_format_msgs_to_str.py b/tests/light/test_format_msgs_to_str.py index bd69751a..93dd7a2f 100644 --- a/tests/light/test_format_msgs_to_str.py +++ b/tests/light/test_format_msgs_to_str.py @@ -481,10 +481,11 @@ def 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 tokens + # Create a message and measure its formatted string tokens msg = create_user_msg("Test") stat = handler.stat_message(msg) - exact_threshold = stat.total_tokens + formatted_content = stat.format(include_thinking=False) + exact_threshold = handler.count_str_token(formatted_content) msgs = [msg] result = handler.format_msgs_to_str(msgs, memory_compact_threshold=exact_threshold) @@ -499,7 +500,8 @@ def test_format_msgs_to_str_threshold_one_less(): handler = create_handler() msg = create_user_msg("Test message") stat = handler.stat_message(msg) - threshold_minus_one = stat.total_tokens - 1 + formatted_content = stat.format(include_thinking=False) + threshold_minus_one = handler.count_str_token(formatted_content) - 1 msgs = [msg] result = handler.format_msgs_to_str(msgs, memory_compact_threshold=threshold_minus_one) From 7e750a5c8eb0939c755adec504555d9bb1d3bf89 Mon Sep 17 00:00:00 2001 From: "jinli.yl" Date: Fri, 6 Mar 2026 16:21:04 +0800 Subject: [PATCH 14/59] delete --- reme_old_doc/__init__.py | 0 1 file changed, 0 insertions(+), 0 deletions(-) delete mode 100644 reme_old_doc/__init__.py diff --git a/reme_old_doc/__init__.py b/reme_old_doc/__init__.py deleted file mode 100644 index e69de29b..00000000 From dcf97dc77f439fd648cb9185b7866b440132fd85 Mon Sep 17 00:00:00 2001 From: "jinli.yl" Date: Fri, 6 Mar 2026 16:34:04 +0800 Subject: [PATCH 15/59] docs(readme): update documentation and examples --- README.md | 27 +++++----- README_ZH.md | 9 ++-- tests/light/test_reme_light_log.txt | 84 +++++++++++++++++++++++++++++ 3 files changed, 102 insertions(+), 18 deletions(-) create mode 100644 tests/light/test_reme_light_log.txt diff --git a/README.md b/README.md index 27acedf5..7b5e4061 100644 --- a/README.md +++ b/README.md @@ -66,15 +66,15 @@ working_dir/ [ReMeLight](reme/reme_light.py) is the core class of this memory system, providing complete 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 | -| `pre_reasoning_hook` | 🔄 Pre-reasoning hook | Auto compact tool results + generate summary + async trigger memory summarization task | -| `memory_search` | 🔍 Semantic memory search | [MemorySearch](reme/memory/tools/chunk/memory_search.py) — Vector + BM25 hybrid retrieval | +| 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 | +| `pre_reasoning_hook` | 🔄 Pre-reasoning hook | Auto compact tool results + generate summary + async trigger memory summarization task | +| `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 (static method) | --- @@ -125,9 +125,9 @@ async def main(): summary = await reme.compact_memory( messages=messages, previous_summary="", - max_input_length=128000, # Model context window (tokens) - compact_ratio=0.7, # Trigger compaction when reaching max_input_length * 0.7 - language="zh", # Summary language (zh / "") + max_input_length=128000, # Model context window (tokens) + compact_ratio=0.7, # Trigger compaction when reaching max_input_length * 0.7 + language="zh", # Summary language (zh / "") ) # 3. Submit async summary task in background (non-blocking, writes to memory/YYYY-MM-DD.md) @@ -169,7 +169,8 @@ if __name__ == "__main__": ``` > 📂 Full example code: [test_reme_light.py](tests/light/test_reme_light.py) -> 📋 Example output: [test_reme_light.log](tests/light/test_reme_light.log) (223,838 tokens → 1,105 tokens, 99.5% compression ratio) +> 📋 Example output: [test_reme_light.log](tests/light/test_reme_light_log.txt) (223,838 tokens → 1,105 tokens, 99.5% +> compression ratio) ### File-Based ReMeLight Memory System Architecture diff --git a/README_ZH.md b/README_ZH.md index ded1d8d1..3c379a2b 100644 --- a/README_ZH.md +++ b/README_ZH.md @@ -97,7 +97,6 @@ pip install -e ".[light]" ```python import asyncio -from agentscope.message import Msg from reme.reme_light import ReMeLight @@ -119,9 +118,9 @@ async def main(): 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 / ""） + max_input_length=128000, # 模型上下文窗口（tokens） + compact_ratio=0.7, # 达到 max_input_length * 0.7 时触发压缩 + language="zh", # 摘要语言（zh / ""） ) # 3. 后台异步提交摘要任务（不阻塞对话，摘要写入 memory/YYYY-MM-DD.md） @@ -163,7 +162,7 @@ if __name__ == "__main__": ``` > 📂 完整示例代码：[test_reme_light.py](tests/light/test_reme_light.py) -> 📋 运行结果示例：[test_reme_light.log](tests/light/test_reme_light.log)（223,838 tokens → 1,105 tokens，压缩率 99.5%） +> 📋 运行结果示例：[test_reme_light.log](tests/light/test_reme_light_log.txt)（223,838 tokens → 1,105 tokens，压缩率 99.5%） ### 基于文件的 ReMeLight 记忆系统架构 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 已关闭 +====================================================================== From d0c9d890929fcce4215e67dea271479f81d15b2f Mon Sep 17 00:00:00 2001 From: jinliyl <6469360+jinliyl@users.noreply.github.com> Date: Fri, 6 Mar 2026 23:43:42 +0800 Subject: [PATCH 16/59] feat(memory): add ContextChecker component for context size management (#144) * feat(memory): add ContextChecker component for context size management * refactor(memory): restructure file-based memory tools and update imports * docs(readme): update documentation with detailed architecture and components * docs(readme): update Chinese documentation with enhanced memory management diagrams * refactor(cookbook): move cookbook files to test directory and clean up docs * docs(readme): update link path for old version documentation * docs(readme): update documentation with improved architecture diagrams and component details * docs(readme): update documentation with improved clarity and structure * refactor(docs): update in-memory memory documentation * docs(readme): add experiment reproduction link to quickstart guide --- README.md | 571 ++++++++------ README_ZH.md | 343 +++++--- .../appworld/quickstart.md | 0 .../cookbook => benchmark}/bfcl/quickstart.md | 0 .../{scripts.sh => cat_correct_scripts.sh} | 0 benchmark/halumem/eval_scripts.sh | 5 + docs/REME2_README.md | 13 - docs/deprecated.txt | 13 - docs/future_work.md | 18 - docs/reme_v2_design.md | 735 ------------------ docs/todo.md | 3 - example.env | 2 - reme/memory/__init__.py | 4 +- reme/memory/file_based/__init__.py | 5 +- .../file_based/component/context_checker.py | 97 +++ reme/memory/file_based/tools/__init__.py | 13 + .../file => file_based/tools}/file_io.py | 91 ++- .../chunk => file_based/tools}/memory_get.py | 0 .../tools}/memory_search.py | 0 reme/memory/file_based/tools/shell.py | 229 ++++++ reme/memory/file_based/tools/utils.py | 112 +++ reme/memory/tools/file/__init__.py | 7 - reme/memory/tools/record/__init__.py | 0 .../{tools => vector_tools}/__init__.py | 5 - .../base_memory_tool.py | 0 .../{tools => vector_tools}/delegate_task.py | 0 .../memory/vector_tools/history}/__init__.py | 0 .../history/add_history.py | 0 .../history/read_history.py | 0 .../history/read_history_v2.py | 0 .../memory/vector_tools/profiles}/__init__.py | 0 .../add_draft_and_read_all_profiles.py | 0 .../profiles/add_profile.py | 0 .../profiles/delete_profile.py | 0 .../profiles/profile_handler.py | 0 .../profiles/read_all_profiles.py | 0 .../profiles/update_profile.py | 0 .../profiles/update_profiles_v1.py | 0 .../memory/vector_tools/record}/__init__.py | 0 .../record/add_and_retrieve_similar_memory.py | 0 .../add_draft_and_retrieve_similar_memory.py | 0 .../record/add_memory.py | 0 .../record/delete_memory.py | 0 .../record/memory_handler.py | 0 .../record/retrieve_memory.py | 0 .../record/retrieve_recent_memory.py | 0 .../record/update_memory.py | 0 .../record/update_memory_v1.py | 0 .../record/update_memory_v2.py | 0 reme/reme.py | 8 +- reme/reme_light.py | 379 ++++++++- .../frozenlake => test/cookbook}/__init__.py | 0 .../cookbook/appworld}/__init__.py | 0 .../appworld/appworld_react_agent.py | 54 +- .../cookbook}/appworld/prompt.py | 0 .../cookbook}/appworld/requirements.txt | 0 .../cookbook}/appworld/run_appworld.py | 12 +- .../cookbook}/appworld/run_exp_statistic.py | 0 .../cookbook/bfcl}/__init__.py | 0 .../cookbook}/bfcl/bfcl_agent.py | 10 +- .../cookbook}/bfcl/bfcl_utils.py | 0 .../cookbook}/bfcl/init_exp_pool.py | 0 .../cookbook}/bfcl/init_task_memory_pool.py | 0 .../cookbook}/bfcl/local_file_to_library.py | 0 .../cookbook}/bfcl/requirements.txt | 0 {cookbook => test/cookbook}/bfcl/run_bfcl.py | 2 +- .../cookbook}/bfcl/run_exp_statistic.py | 0 .../cookbook}/bfcl/split_into_trainval.py | 0 .../cookbook/frozenlake}/__init__.py | 0 .../frozenlake/frozenlake_prompts.yaml | 0 .../frozenlake/frozenlake_react_agent.py | 0 .../cookbook}/frozenlake/map_manager.py | 0 .../cookbook}/frozenlake/run_exp_statistic.py | 0 .../cookbook}/frozenlake/run_frozenlake.py | 0 .../cookbook/simple_demo}/__init__.py | 0 .../simple_demo/import_usage_demo.py | 0 .../simple_demo/mcp_task_memory.jsonl | 0 .../simple_demo/personal_memory.jsonl | 0 .../cookbook}/simple_demo/task_memory.jsonl | 0 .../cookbook}/simple_demo/task_messages.jsonl | 0 .../simple_demo/use_personal_memory_demo.py | 0 .../simple_demo/use_task_memory_demo.py | 0 .../simple_demo/use_task_memory_mcp_demo.py | 0 .../simple_demo/use_tool_memory_demo.py | 0 .../cookbook/tool_memory}/__init__.py | 0 .../cookbook}/tool_memory/query.json | 0 .../tool_memory/run_reme_tool_bench.py | 0 .../react_agent_with_working_memory.py | 78 +- .../working_memory/work_memory_demo.py | 2 +- test/{ => test}/cli/__init__.py | 0 test/{ => test}/cli/fb_cli.py | 0 test/{ => test}/cli/fb_cli.yaml | 0 test/{ => test}/cli/fb_compactor.py | 0 test/{ => test}/cli/fb_compactor.yaml | 0 test/{ => test}/cli/fb_context_checker.py | 0 test/{ => test}/cli/fb_summarizer.py | 0 test/{ => test}/cli/fb_summarizer.yaml | 0 test/{ => test}/reme_cli.py | 0 test/{ => test}/test_fs_compactor.py | 0 test/{ => test}/test_fs_context_checker.py | 0 .../test_fs_file_watch_integration.py | 0 test/{ => test}/test_fs_memory_get.py | 0 test/{ => test}/test_fs_memory_search.py | 0 test/{ => test}/test_fs_summary.py | 0 tests/light/test_summarizer.py | 2 +- tests/light/test_tools.py | 321 ++++++++ tests/vector/test_reme_vector.py | 89 +++ 107 files changed, 1922 insertions(+), 1301 deletions(-) rename {docs/cookbook => benchmark}/appworld/quickstart.md (100%) rename {docs/cookbook => benchmark}/bfcl/quickstart.md (100%) rename benchmark/halumem/{scripts.sh => cat_correct_scripts.sh} (100%) create mode 100755 benchmark/halumem/eval_scripts.sh delete mode 100644 docs/REME2_README.md delete mode 100644 docs/deprecated.txt delete mode 100644 docs/future_work.md delete mode 100644 docs/reme_v2_design.md delete mode 100644 docs/todo.md create mode 100644 reme/memory/file_based/component/context_checker.py create mode 100644 reme/memory/file_based/tools/__init__.py rename reme/memory/{tools/file => file_based/tools}/file_io.py (77%) rename reme/memory/{tools/chunk => file_based/tools}/memory_get.py (100%) rename reme/memory/{tools/chunk => file_based/tools}/memory_search.py (100%) create mode 100644 reme/memory/file_based/tools/shell.py create mode 100644 reme/memory/file_based/tools/utils.py delete mode 100644 reme/memory/tools/file/__init__.py delete mode 100644 reme/memory/tools/record/__init__.py rename reme/memory/{tools => vector_tools}/__init__.py (92%) rename reme/memory/{tools => vector_tools}/base_memory_tool.py (100%) rename reme/memory/{tools => vector_tools}/delegate_task.py (100%) rename {cookbook => reme/memory/vector_tools/history}/__init__.py (100%) rename reme/memory/{tools => vector_tools}/history/add_history.py (100%) rename reme/memory/{tools => vector_tools}/history/read_history.py (100%) rename reme/memory/{tools => vector_tools}/history/read_history_v2.py (100%) rename {cookbook/appworld => reme/memory/vector_tools/profiles}/__init__.py (100%) rename reme/memory/{tools => vector_tools}/profiles/add_draft_and_read_all_profiles.py (100%) rename reme/memory/{tools => vector_tools}/profiles/add_profile.py (100%) rename reme/memory/{tools => vector_tools}/profiles/delete_profile.py (100%) rename reme/memory/{tools => vector_tools}/profiles/profile_handler.py (100%) rename reme/memory/{tools => vector_tools}/profiles/read_all_profiles.py (100%) rename reme/memory/{tools => vector_tools}/profiles/update_profile.py (100%) rename reme/memory/{tools => vector_tools}/profiles/update_profiles_v1.py (100%) rename {cookbook/bfcl => reme/memory/vector_tools/record}/__init__.py (100%) rename reme/memory/{tools => vector_tools}/record/add_and_retrieve_similar_memory.py (100%) rename reme/memory/{tools => vector_tools}/record/add_draft_and_retrieve_similar_memory.py (100%) rename reme/memory/{tools => vector_tools}/record/add_memory.py (100%) rename reme/memory/{tools => vector_tools}/record/delete_memory.py (100%) rename reme/memory/{tools => vector_tools}/record/memory_handler.py (100%) rename reme/memory/{tools => vector_tools}/record/retrieve_memory.py (100%) rename reme/memory/{tools => vector_tools}/record/retrieve_recent_memory.py (100%) rename reme/memory/{tools => vector_tools}/record/update_memory.py (100%) rename reme/memory/{tools => vector_tools}/record/update_memory_v1.py (100%) rename reme/memory/{tools => vector_tools}/record/update_memory_v2.py (100%) rename {cookbook/frozenlake => test/cookbook}/__init__.py (100%) rename {cookbook/simple_demo => test/cookbook/appworld}/__init__.py (100%) rename {cookbook => test/cookbook}/appworld/appworld_react_agent.py (87%) rename {cookbook => test/cookbook}/appworld/prompt.py (100%) rename {cookbook => test/cookbook}/appworld/requirements.txt (100%) rename {cookbook => test/cookbook}/appworld/run_appworld.py (98%) rename {cookbook => test/cookbook}/appworld/run_exp_statistic.py (100%) rename {cookbook/tool_memory => test/cookbook/bfcl}/__init__.py (100%) rename {cookbook => test/cookbook}/bfcl/bfcl_agent.py (98%) rename {cookbook => test/cookbook}/bfcl/bfcl_utils.py (100%) rename {cookbook => test/cookbook}/bfcl/init_exp_pool.py (100%) rename {cookbook => test/cookbook}/bfcl/init_task_memory_pool.py (100%) rename {cookbook => test/cookbook}/bfcl/local_file_to_library.py (100%) rename {cookbook => test/cookbook}/bfcl/requirements.txt (100%) rename {cookbook => test/cookbook}/bfcl/run_bfcl.py (99%) rename {cookbook => test/cookbook}/bfcl/run_exp_statistic.py (100%) rename {cookbook => test/cookbook}/bfcl/split_into_trainval.py (100%) rename {reme/memory/tools/chunk => test/cookbook/frozenlake}/__init__.py (100%) rename {cookbook => test/cookbook}/frozenlake/frozenlake_prompts.yaml (100%) rename {cookbook => test/cookbook}/frozenlake/frozenlake_react_agent.py (100%) rename {cookbook => test/cookbook}/frozenlake/map_manager.py (100%) rename {cookbook => test/cookbook}/frozenlake/run_exp_statistic.py (100%) rename {cookbook => test/cookbook}/frozenlake/run_frozenlake.py (100%) rename {reme/memory/tools/history => test/cookbook/simple_demo}/__init__.py (100%) rename {cookbook => test/cookbook}/simple_demo/import_usage_demo.py (100%) rename {cookbook => test/cookbook}/simple_demo/mcp_task_memory.jsonl (100%) rename {cookbook => test/cookbook}/simple_demo/personal_memory.jsonl (100%) rename {cookbook => test/cookbook}/simple_demo/task_memory.jsonl (100%) rename {cookbook => test/cookbook}/simple_demo/task_messages.jsonl (100%) rename {cookbook => test/cookbook}/simple_demo/use_personal_memory_demo.py (100%) rename {cookbook => test/cookbook}/simple_demo/use_task_memory_demo.py (100%) rename {cookbook => test/cookbook}/simple_demo/use_task_memory_mcp_demo.py (100%) rename {cookbook => test/cookbook}/simple_demo/use_tool_memory_demo.py (100%) rename {reme/memory/tools/profiles => test/cookbook/tool_memory}/__init__.py (100%) rename {cookbook => test/cookbook}/tool_memory/query.json (100%) rename {cookbook => test/cookbook}/tool_memory/run_reme_tool_bench.py (100%) rename {cookbook => test/cookbook}/working_memory/react_agent_with_working_memory.py (70%) rename {cookbook => test/cookbook}/working_memory/work_memory_demo.py (99%) rename test/{ => test}/cli/__init__.py (100%) rename test/{ => test}/cli/fb_cli.py (100%) rename test/{ => test}/cli/fb_cli.yaml (100%) rename test/{ => test}/cli/fb_compactor.py (100%) rename test/{ => test}/cli/fb_compactor.yaml (100%) rename test/{ => test}/cli/fb_context_checker.py (100%) rename test/{ => test}/cli/fb_summarizer.py (100%) rename test/{ => test}/cli/fb_summarizer.yaml (100%) rename test/{ => test}/reme_cli.py (100%) rename test/{ => test}/test_fs_compactor.py (100%) rename test/{ => test}/test_fs_context_checker.py (100%) rename test/{ => test}/test_fs_file_watch_integration.py (100%) rename test/{ => test}/test_fs_memory_get.py (100%) rename test/{ => test}/test_fs_memory_search.py (100%) rename test/{ => test}/test_fs_summary.py (100%) create mode 100644 tests/light/test_tools.py create mode 100644 tests/vector/test_reme_vector.py diff --git a/README.md b/README.md index 7b5e4061..9d18e6ce 100644 --- a/README.md +++ b/README.md @@ -12,7 +12,7 @@

License English - 简体中文 + Simplified Chinese GitHub Stars

@@ -20,66 +20,65 @@ A memory management toolkit for AI agents — Remember Me, Refine Me.

-> For legacy versions, see [0.2.x Documentation](docs/README_0_2_x.md) +> For the older version, please refer to the [0.2.x documentation](docs/README_0_2_x_ZH.md). --- -🧠 ReMe is a **memory management framework** built for **AI agents**, offering both **file-based** and **vector-based** -memory systems. +🧠 ReMe is a memory management framework designed for **AI agents**, providing both file-based and vector-based memory +systems. -It addresses two core problems of agent memory: **limited context windows** (early information gets truncated or lost -during long conversations) and **stateless sessions** (new conversations cannot inherit history and always start from -scratch). - -ReMe gives agents **real memory** — old conversations are automatically condensed, important information is persisted, -and the next conversation can recall it automatically. +It tackles two core problems of agent memory: **limited context window** (early information is truncated or lost in long +conversations) and **stateless sessions** (new sessions cannot inherit history and always start from scratch). +ReMe gives agents **real memory** — old conversations are automatically compacted, important information is persistently +stored, and relevant context is automatically recalled in future interactions. --- -## 📁 File-Based Memory System (ReMeLight) +## 📁 File-based memory system (ReMeLight) -> Memory as files, files as memory +> Memory as files, files as memory. -Treat **memory as files** — readable, editable, and portable. -[CoPaw](https://github.com/agentscope-ai/CoPaw) implements long-term memory and context management by inheriting +Treat **memory as files** — readable, editable, and copyable. +[CoPaw](https://github.com/agentscope-ai/CoPaw) integrates long-term memory and context management by inheriting from `ReMeLight`. -| Traditional Memory Systems | File-Based ReMe | -|----------------------------|--------------------| -| 🗄️ Database storage | 📝 Markdown files | -| 🔒 Opaque | 👀 Read anytime | -| ❌ Hard to modify | ✏️ Edit directly | -| 🚫 Hard to migrate | 📦 Copy to migrate | +| Traditional memory system | File-based ReMe | +|---------------------------|----------------------| +| 🗄️ Database storage | 📝 Markdown files | +| 🔒 Opaque | 👀 Always readable | +| ❌ Hard to modify | ✏️ Directly editable | +| 🚫 Hard to migrate | 📦 Copy to migrate | ``` working_dir/ -├── MEMORY.md # Long-term memory: user preferences, project config, etc. +├── MEMORY.md # Long-term memory: persistent info such as user preferences ├── memory/ -│ └── YYYY-MM-DD.md # Daily summary logs: written automatically after conversation ends -└── tool_result/ # Cache for oversized tool outputs (auto-managed, auto-cleaned when expired) +│ └── YYYY-MM-DD.md # Daily journal: automatically written after each conversation +└── tool_result/ # Cache for long tool outputs (auto-managed, expired entries auto-cleaned) └── .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 | -| `pre_reasoning_hook` | 🔄 Pre-reasoning hook | Auto compact tool results + generate summary + async trigger memory summarization task | -| `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 (static method) | +| Method | Function | Key components | +|------------------------|--------------------------------------|--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------| +| `check_context` | 📊 Check context size | [ContextChecker](reme/memory/file_based/component/context_checker.py) — checks whether context exceeds thresholds and splits messages | +| `compact_memory` | 📦 Compact history into summary | [Compactor](reme/memory/file_based/component/compactor.py) — ReActAgent that generates structured context summaries | +| `summary_memory` | 📝 Persist important memory to files | [Summarizer](reme/memory/file_based/component/summarizer.py) — ReActAgent + file tools (`read` / `write` / `edit`) | +| `compact_tool_result` | ✂️ Compact long tool outputs | [ToolResultCompactor](reme/memory/file_based/component/tool_result_compactor.py) — truncates long tool outputs and stores them in `tool_result/` while keeping file references in messages | +| `memory_search` | 🔍 Semantic memory search | [MemorySearch](reme/memory/file_based/tools/memory_search.py) — hybrid retrieval with vectors + BM25 | +| `ReMeInMemoryMemory` | 🗂️ In-session memory class | [ReMeInMemoryMemory](reme/memory/file_based/reme_in_memory_memory.py) — token-aware memory management with summary compression and state serialization | +| `pre_reasoning_hook` | 🔄 Pre-reasoning hook | `compact_tool_result` + `check_context` + `compact_memory` + `summary_memory` (async) | +| `start` | 🚀 Start memory system | Initialize file storage, file watcher, and embedding cache; clean up expired tool result files | +| `close` | 📕 Shutdown and cleanup | Clean up tool result files, stop file watcher, and persist embedding cache | --- -### 🚀 Quick Start +### 🚀 Quick start #### Installation @@ -87,23 +86,22 @@ capabilities for AI Agents: pip install -e ".[light]" ``` -#### Environment Variables +#### Environment variables -`ReMeLight` environment variables configure Embedding and storage backend +`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` | +| 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 usage ```python import asyncio -from agentscope.message import Msg from reme.reme_light import ReMeLight @@ -116,24 +114,24 @@ async def main(): ) await reme.start() - messages = [...] # Conversation message list + messages = [...] # List of conversation messages - # 1. Compact oversized tool outputs (prevent tool results from overflowing context) + # 1. Compact long tool outputs (prevent tool results from blowing up context) messages = await reme.compact_tool_result(messages) - # 2. Compact history to structured summary (can pass previous summary for incremental update) + # 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 reaching max_input_length * 0.7 - language="zh", # Summary language (zh / "") + compact_ratio=0.7, # Trigger compaction when exceeding max_input_length * 0.7 + language="zh", # Summary language (e.g., "zh" / "") ) - # 3. Submit async summary task in background (non-blocking, writes to memory/YYYY-MM-DD.md) + # 3. Submit summary task asynchronously (non-blocking, writes to memory/YYYY-MM-DD.md) reme.add_async_summary_task(messages=messages) - # 4. Pre-reasoning hook (auto compact tool results + generate summary) + # 4. Pre-reasoning hook (auto compact tool results + generate summaries) processed_messages, compressed_summary = await reme.pre_reasoning_hook( messages=messages, system_prompt="You are a helpful AI assistant.", @@ -145,22 +143,23 @@ async def main(): tool_result_compact_keep_n=3, ) - # 5. Semantic memory search (Vector + BM25 hybrid retrieval) + # 5. Semantic memory search (vector + BM25 hybrid retrieval) result = await reme.memory_search(query="Python version preference", max_results=5) - # 6. Get in-memory instance (static method, manages single conversation context) - memory = ReMeLight.get_in_memory_memory() + # 6. Create in-session memory instance (manages context for one conversation) + from reme.memory.file_based.reme_in_memory_memory import ReMeInMemoryMemory + memory = ReMeInMemoryMemory() 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']}") - # 7. Wait for background tasks before closing + # 7. Wait for background summary tasks to complete before shutdown summary_result = await reme.await_summary_tasks() - # Close ReMeLight + # Shutdown ReMeLight await reme.close() @@ -168,173 +167,226 @@ if __name__ == "__main__": asyncio.run(main()) ``` -> 📂 Full example code: [test_reme_light.py](tests/light/test_reme_light.py) -> 📋 Example output: [test_reme_light.log](tests/light/test_reme_light_log.txt) (223,838 tokens → 1,105 tokens, 99.5% -> compression ratio) +> 📂 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) -### File-Based ReMeLight Memory System Architecture +### Architecture of the file-based ReMeLight memory system [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: - -```mermaid -graph TB - CoPaw["CoPaw MemoryManager
(inherits ReMeLight)"] -->|pre_reasoning hook| Hook[MemoryCompactionHook] - CoPaw --> ReMeLight[ReMeLight] - Hook -->|exceeds threshold| ReMeLight - ReMeLight --> CompactMemory[compact_memory
History compaction] - ReMeLight --> SummaryMemory[summary_memory
Write memory to files] - ReMeLight --> CompactToolResult[compact_tool_result
Oversized tool output compaction] - ReMeLight --> MemSearch[memory_search
Semantic search] - ReMeLight --> InMemory[get_in_memory_memory
ReMeInMemoryMemory] - CompactMemory --> Compactor[Compactor
ReActAgent] - SummaryMemory --> Summarizer[Summarizer
ReActAgent + file tools] - CompactToolResult --> ToolResultCompactor[ToolResultCompactor
Truncate + save to file] - Summarizer --> FileIO[FileIO
read / 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): +inherits +`ReMeLight` and integrates its memory capabilities into the agent reasoning loop: ```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] + Agent[Agent] -->|Before each reasoning step| Hook[pre_reasoning_hook] + Hook --> TC[compact_tool_result
Compact tool outputs] + TC --> CC[check_context
Token counting] + CC -->|Exceeds limit| CM[compact_memory
Generate summary] + CC -->|Exceeds limit| SM[summary_memory
Async persistence] + SM -->|ReAct + FileIO| Files[memory/*.md] + Agent -->|Explicit call| Search[memory_search
Vector+BM25] + Agent -->|In-session| InMem[ReMeInMemoryMemory
Token-aware memory] + Files -.->|FileWatcher| Store[(FileStore
Vector+FTS index)] + Search --> Store ``` --- -## 🗃️ Vector-Based Memory System +#### 1. `check_context` — context checking -[ReMe Vector Based](reme/reme.py) is the core class for the vector-based memory system, supporting unified management of -three memory types: +[ContextChecker](reme/memory/file_based/component/context_checker.py) uses token counting to determine whether the +context exceeds thresholds and automatically splits messages into a "to compact" group and a "to keep" group. -| 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 + M[messages] --> H[AsMsgHandler
Token counting] + H --> C{total > threshold?} + C -->|No| K[Return all messages] + C -->|Yes| S[Keep from tail
reserve tokens] + S --> CP[messages_to_compact
Earlier messages] + S --> KP[messages_to_keep
Recent messages] + S --> V{is_valid
Tool calls aligned?} ``` -### Environment Variables +- **Core logic**: keep `reserve` tokens from the tail; mark the rest as messages to compact. +- **Integrity guarantee**: preserves complete user-assistant turns and tool_use/tool_result pairs without splitting + them. -API keys are set via environment variables; you can put them in a `.env` file in the project root: +--- -| 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` | +#### 2. `compact_memory` — conversation compaction -### Python Usage +[Compactor](reme/memory/file_based/component/compactor.py) uses a ReActAgent to compact conversation history into a * +*structured context summary**. + +```mermaid +graph LR + M[messages] --> H[AsMsgHandler
format_msgs_to_str] + H --> A[ReActAgent
reme_compactor] + P[previous_summary] -->|Incremental update| A + A --> S[Structured summary
Goal/Progress/Decisions...] +``` + +**Summary structure** (context checkpoints): + +| Field | Description | +|-----------------------|------------------------------------------------------------------------| +| `## Goal` | User goals | +| `## Constraints` | Constraints and preferences | +| `## Progress` | Task progress | +| `## Key Decisions` | Key decisions | +| `## Next Steps` | Next step plans | +| `## Critical Context` | Critical data such as file paths, function names, error messages, etc. | + +- **Incremental updates**: when `previous_summary` is provided, new conversations are merged into the existing summary. + +--- + +#### 3. `summary_memory` — persistent memory + +[Summarizer](reme/memory/file_based/component/summarizer.py) uses a **ReAct + file tools** pattern so that the AI can +decide what to write and where to write it. + +```mermaid +graph LR + M[messages] --> A[ReActAgent
reme_summarizer] + A -->|read| R[Read memory/YYYY-MM-DD.md] + R --> T{Reason: how to merge?} + T -->|write| W[Overwrite] + T -->|edit| E[Edit in place] + W --> F[memory/YYYY-MM-DD.md] + E --> F +``` + +**File tools** ([FileIO](reme/memory/file_based/tools/file_io.py)): + +| Tool | Function | +|---------|-----------------------| +| `read` | Read file content | +| `write` | Overwrite file | +| `edit` | Find-and-replace edit | + +--- + +#### 4. `compact_tool_result` — tool result compaction + +[ToolResultCompactor](reme/memory/file_based/component/tool_result_compactor.py) addresses the problem of long tool +outputs bloating the context. + +```mermaid +graph LR + M[messages] --> L{Iterate tool_result
len > threshold?} + L -->|No| K[Keep as-is] + L -->|Yes| T[truncate_text
Truncate to threshold] + T --> S[Write full content
tool_result/uuid.txt] + S --> R[Append file path reference
to message] + R --> C[cleanup_expired_files
Delete expired files] +``` + +- **Auto cleanup**: expired files (older than `retention_days`) are deleted automatically during `start` / `close` / + `compact_tool_result`. + +--- + +#### 5. `memory_search` — memory retrieval + +[MemorySearch](reme/memory/file_based/tools/memory_search.py) provides **vector + BM25 hybrid retrieval**. + +```mermaid +graph LR + Q[query] --> E[Embedding
Vectorization] + E --> V[vector_search
Semantic similarity] + Q --> B[BM25
Keyword matching] + V -->|" weight: 0.7 "| M[Deduplicate + weighted merge] + B -->|" weight: 0.3 "| M + M --> F[min_score filter] + F --> R[Top-N results] +``` + +- **Fusion mechanism**: vector weight 0.7 + BM25 weight 0.3 — balancing semantic similarity and exact matches. + +--- + +#### 6. `ReMeInMemoryMemory` — in-session memory + +[ReMeInMemoryMemory](reme/memory/file_based/reme_in_memory_memory.py) extends AgentScope's `InMemoryMemory` to provide +token-aware memory management. + +```mermaid +graph LR + C[content] --> G[get_memory
exclude_mark=COMPRESSED] + G --> F[Filter out compressed messages] + F --> P{prepend_summary?} + P -->|Yes| S[Prepend previous summary] + S --> O[Output messages] + P -->|No| O +``` + +| Function | Description | +|----------------------------------|---------------------------------------------------| +| `get_memory` | Filter messages by mark and auto-append summary | +| `estimate_tokens` | Estimate token usage of the context | +| `state_dict` / `load_state_dict` | Serialize/deserialize state (session persistence) | + +--- + +#### 7. `pre_reasoning_hook` — pre-reasoning processing + +This is a unified entry point that wires all the above components together and automatically manages context before each +reasoning step. + +```mermaid +graph LR + M[messages] --> TC[compact_tool_result
Compact long tool outputs] + TC --> CC[check_context
Compute remaining space] + CC --> D{messages_to_compact
Non-empty?} + D -->|No| K[Return original messages + summary] + D -->|Yes| V{is_valid?} + V -->|No| K + V -->|Yes| CM[compact_memory
Sync summary generation] + V -->|Yes| SM[add_async_summary_task
Async persistence] + CM --> R[Return messages_to_keep + new summary] +``` + +**Execution flow**: + +1. `compact_tool_result` — compact long tool outputs. +2. `check_context` — check whether the context exceeds limits. +3. `compact_memory` — generate compact summary (sync). +4. `summary_memory` — persist memory (async in the background). + +--- + +## 🗃️ Vector-based memory system + +[ReMe Vector Based](reme/reme.py) is the core class for the vector-based memory system. It manages three types of +memories: + +| Memory type | Use case | +|-----------------------|-------------------------------------------------------------------| +| **Personal memory** | Records user preferences and habits | +| **Procedural memory** | Records task execution experience and patterns of success/failure | +| **Tool memory** | Records tool usage experience and parameter tuning | + +### Core capabilities + +| Method | Function | Description | +|--------------------|--------------|-------------------------------------------------------------| +| `summarize_memory` | 🧠 Summarize | Automatically extract and store memories from conversations | +| `retrieve_memory` | 🔍 Retrieve | Retrieve related memories based on a query | +| `add_memory` | ➕ Add | Manually add memories into the vector store | +| `get_memory` | 📖 Get | Get a single memory by ID | +| `update_memory` | ✏️ Update | Update existing memory content or metadata | +| `delete_memory` | 🗑️ Delete | Delete a specific memory | +| `list_memory` | 📋 List | List memories with filtering and sorting | + +### Installation and environment variables + +Installation and environment configuration are the same as [ReMeLight](#installation). +API keys are configured via environment variables and can be stored in a `.env` file at the project root. + +### Python usage ```python import asyncio @@ -363,34 +415,34 @@ async def main(): messages = [ {"role": "user", "content": "Help me write a Python script", "time_created": "2026-02-28 10:00:00"}, - {"role": "assistant", "content": "Sure, I'll help you write it", "time_created": "2026-02-28 10:00:05"}, + {"role": "assistant", "content": "Sure, I'll help you with that.", "time_created": "2026-02-28 10:00:05"}, ] - # 1. Summarize memory from conversation (auto-extract user preferences, task experience, etc.) + # 1. Summarize memories from conversation (automatically extract user preferences, task experience, etc.) result = await reme.summarize_memory( messages=messages, user_name="alice", # Personal memory - # task_name="code_writing", # Task memory + # task_name="code_writing", # Procedural memory ) - print(f"Summarize result: {result}") + print(f"Summary result: {result}") - # 2. Retrieve relevant memory + # 2. Retrieve related memories memories = await reme.retrieve_memory( query="Python programming", user_name="alice", # task_name="code_writing", ) - print(f"Retrieve result: {memories}") + print(f"Retrieved memories: {memories}") - # 3. Manually add memory + # 3. Manually add a memory memory_node = await reme.add_memory( - memory_content="User prefers concise code style", + memory_content="The user prefers concise code style.", user_name="alice", ) print(f"Added memory: {memory_node}") memory_id = memory_node.memory_id - # 4. Get single memory by ID + # 4. Get a single memory by ID fetched_memory = await reme.get_memory(memory_id=memory_id) print(f"Fetched memory: {fetched_memory}") @@ -398,11 +450,11 @@ async def main(): updated_memory = await reme.update_memory( memory_id=memory_id, user_name="alice", - memory_content="User prefers concise, well-commented code style", + memory_content="The user prefers concise code with comments.", ) print(f"Updated memory: {updated_memory}") - # 6. List all memories for user (with filtering and sorting) + # 6. List all memories for the user (supports filtering and sorting) all_memories = await reme.list_memory( user_name="alice", limit=10, @@ -411,11 +463,11 @@ async def main(): ) print(f"User memory list: {all_memories}") - # 7. Delete specified memory + # 7. Delete a specific memory await reme.delete_memory(memory_id=memory_id) print(f"Deleted memory: {memory_id}") - # 8. Delete all memories (use with caution) + # 8. Delete all memories (use with care) # await reme.delete_all() await reme.close() @@ -425,21 +477,21 @@ if __name__ == "__main__": asyncio.run(main()) ``` -### Technical Architecture +### Technical architecture ```mermaid -graph TB +graph LR User[User / Agent] --> ReMe[Vector Based ReMe] - ReMe --> Summarize[Memory Summarize] - ReMe --> Retrieve[Memory Retrieve] - ReMe --> CRUD[CRUD] + ReMe --> Summarize[Summarize memories] + ReMe --> Retrieve[Retrieve memories] + ReMe --> CRUD[CRUD operations] Summarize --> PersonalSum[PersonalSummarizer] Summarize --> ProceduralSum[ProceduralSummarizer] Summarize --> ToolSum[ToolSummarizer] Retrieve --> PersonalRet[PersonalRetriever] Retrieve --> ProceduralRet[ProceduralRetriever] Retrieve --> ToolRet[ToolRetriever] - PersonalSum --> VectorStore[Vector DB] + PersonalSum --> VectorStore[Vector database] ProceduralSum --> VectorStore ToolSum --> VectorStore PersonalRet --> VectorStore @@ -447,17 +499,53 @@ graph TB ToolRet --> VectorStore ``` -## ⭐ Community & Support +### Experimental results -- **Star & Watch**: Star helps more agent developers discover ReMe; Watch keeps you updated on new releases and - features. -- **Share your work**: In Issues or Discussions, share what ReMe unlocks for your agents — we're happy to highlight - great community examples. -- **Need a new feature?** Open a Feature Request; we'll iterate with the community. -- **Code contributions**: All forms of code contribution are welcome. See - the [Contribution Guide](docs/contribution.md). -- **Acknowledgments**: Thanks to OpenClaw, Mem0, MemU, CoPaw, and other open-source projects for inspiration and - support. +Coming soon... + +--- + +## 🧪 Procedural memory paper + +> Our procedural (task) memory paper is available on [arXiv](https://arxiv.org/abs/2512.10696). + +### 🌍 [Appworld benchmark](benchmark/appworld/quickstart.md) + +We evaluate ReMe on the Appworld environment using Qwen3-8B (non-thinking mode): + +| Method | Avg@4 | Pass@4 | +|----------|---------------------|---------------------| +| w/o ReMe | 0.1497 | 0.3285 | +| w/ ReMe | 0.1706 **(+2.09%)** | 0.3631 **(+3.46%)** | + +Pass@K measures the probability that at least one of K generated candidates successfully completes the task (score=1). +The current experiments use an internal AppWorld environment, which may differ slightly from the public version. + +For more details on how to reproduce the experiments, see [quickstart.md](benchmark/appworld/quickstart.md). + +### 🔧 [BFCL-V3 benchmark](benchmark/bfcl/quickstart.md) + +We evaluate ReMe on the BFCL-V3 multi-turn-base task (random split 50 train / 150 val) using Qwen3-8B (thinking mode): + +| Method | Avg@4 | Pass@4 | +|----------|---------------------|---------------------| +| w/o ReMe | 0.4033 | 0.5955 | +| w/ ReMe | 0.4450 **(+4.17%)** | 0.6577 **(+6.22%)** | + +For more details on how to reproduce the experiments, see [quickstart.md](benchmark/bfcl/quickstart.md). + + +## ⭐ Community & support + +- **Star & Watch**: Starring helps more agent developers discover ReMe; Watching keeps you up to date with new releases + and features. +- **Share your results**: Share how ReMe empowers your agents in Issues or Discussions — we are happy to showcase great + community use cases. +- **Need a new feature?** Open a feature request; we’ll evolve ReMe together with the community. +- **Code contributions**: All forms of contributions are welcome. Please see + the [contribution guide](docs/contribution.md). +- **Acknowledgements**: We thank excellent open-source projects such as OpenClaw, Mem0, MemU, and CoPaw for their + inspiration and support. --- @@ -476,10 +564,11 @@ 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 +## 📈 Star history [![Star History Chart](https://api.star-history.com/svg?repos=agentscope-ai/ReMe&type=Date)](https://www.star-history.com/#agentscope-ai/ReMe&Date) + diff --git a/README_ZH.md b/README_ZH.md index 3c379a2b..8723177c 100644 --- a/README_ZH.md +++ b/README_ZH.md @@ -38,7 +38,7 @@ ReMe 让智能体拥有**真正的记忆力**——旧对话自动浓缩，重 > 记忆即文件，文件即记忆 将**记忆视为文件**——可读、可编辑、可复制。 -[CoPaw](https://github.com/agentscope-ai/CoPaw)通过继承 `ReMeLight` 实现了长期记忆和上下文的管理。 +[CoPaw](https://github.com/agentscope-ai/CoPaw) 通过继承 `ReMeLight` 实现了长期记忆和上下文的管理。 | 传统记忆系统 | File Based ReMe | |-----------|-----------------| @@ -49,9 +49,9 @@ ReMe 让智能体拥有**真正的记忆力**——旧对话自动浓缩，重 ``` working_dir/ -├── MEMORY.md # 长期记忆：用户偏好、项目配置等持久信息 +├── MEMORY.md # 长期记忆：用户偏好等持久信息 ├── memory/ -│ └── YYYY-MM-DD.md # 每日摘要日志：对话结束后自动写入 +│ └── YYYY-MM-DD.md # 每日日记：对话结束后自动写入 └── tool_result/ # 超长工具输出缓存（自动管理，超期自动清理） └── .txt ``` @@ -60,17 +60,17 @@ 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/`，消息中保留文件引用 | -| `pre_reasoning_hook` | 🔄 推理前预处理钩子 | 自动压缩工具结果 + 生成摘要 + 异步触发记忆总结任务 | -| `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](reme/memory/file_based/component/context_checker.py) — 检查上下文是否超出阈值并拆分Message | +| `compact_memory` | 📦 压缩历史对话为摘要 | [Compactor](reme/memory/file_based/component/compactor.py) — ReActAgent 生成结构化上下文摘要 | +| `summary_memory` | 📝 将重要记忆写入文件 | [Summarizer](reme/memory/file_based/component/summarizer.py) — ReActAgent + 文件工具（read / write / edit） | +| `compact_tool_result` | ✂️ 压缩超长工具输出 | [ToolResultCompactor](reme/memory/file_based/component/tool_result_compactor.py) — 截断超长的工具调用结果并转存到 `tool_result/`，消息中保留文件引用 | +| `memory_search` | 🔍 语义搜索记忆 | [MemorySearch](reme/memory/file_based/tools/memory_search.py) — 向量 + BM25 混合检索 | +| `ReMeInMemoryMemory` | 🗂️ 会话内存类 | [ReMeInMemoryMemory](reme/memory/file_based/reme_in_memory_memory.py) — Token 感知的内存管理，支持压缩摘要和状态序列化 | +| `pre_reasoning_hook` | 🔄 推理前预处理钩子 | compact_tool_result + check_context + compact_memory + summary_memory(async) | +| `start` | 🚀 启动记忆系统 | 初始化文件存储、文件监控、Embedding 缓存；清理过期工具结果文件 | +| `close` | 📕 关闭并清理 | 清理工具结果文件、停止文件监控、保存 Embedding 缓存 |· --- ### 🚀 快速开始 @@ -85,14 +85,14 @@ 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` | +| 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 @@ -141,8 +141,9 @@ async def main(): # 5. 语义搜索记忆（向量 + BM25 混合检索） result = await reme.memory_search(query="Python 版本偏好", max_results=5) - # 6. 获取会话内存实例（静态方法，管理单次对话的上下文） - memory = ReMeLight.get_in_memory_memory() + # 6. 创建会话内存实例（管理单次对话的上下文） + from reme.memory.file_based.reme_in_memory_memory import ReMeInMemoryMemory + memory = ReMeInMemoryMemory() for msg in messages: await memory.add(msg) token_stats = await memory.estimate_tokens(max_input_length=128000) @@ -162,7 +163,7 @@ if __name__ == "__main__": ``` > 📂 完整示例代码：[test_reme_light.py](tests/light/test_reme_light.py) -> 📋 运行结果示例：[test_reme_light.log](tests/light/test_reme_light_log.txt)（223,838 tokens → 1,105 tokens，压缩率 99.5%） +> 📋 运行结果示例：[test_reme_light_log.txt](tests/light/test_reme_light_log.txt)（223,838 tokens → 1,105 tokens，压缩率99.5%） ### 基于文件的 ReMeLight 记忆系统架构 @@ -170,125 +171,188 @@ if __name__ == "__main__": `ReMeLight`，将记忆能力集成到 Agent 推理流程中： ```mermaid -graph TB - CoPaw["CoPaw MemoryManager
(继承 ReMeLight)"] -->|pre_reasoning hook| Hook[MemoryCompactionHook] - CoPaw --> ReMeLight[ReMeLight] - Hook -->|超出阈值| ReMeLight - ReMeLight --> CompactMemory[compact_memory
历史对话压缩] - ReMeLight --> SummaryMemory[summary_memory
记忆写入文件] - ReMeLight --> CompactToolResult[compact_tool_result
超长工具输出压缩] - ReMeLight --> MemSearch[memory_search
语义搜索] - ReMeLight --> InMemory[get_in_memory_memory
ReMeInMemoryMemory] - CompactMemory --> Compactor[Compactor
ReActAgent] - SummaryMemory --> Summarizer[Summarizer
ReActAgent + 文件工具] - CompactToolResult --> ToolResultCompactor[ToolResultCompactor
截断 + 转存文件] - Summarizer --> FileIO[FileIO
read / 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
压缩工具输出] + TC --> CC[check_context
Token 计数] + CC -->|超限| CM[compact_memory
生成摘要] + CC -->|超限| SM[summary_memory
异步持久化] + SM -->|ReAct + FileIO| Files[memory/*.md] + Agent -->|主动调用| Search[memory_search
向量+BM25] + Agent -->|会话内存| InMem[ReMeInMemoryMemory
Token感知内存] + Files -.->|FileWatcher| Store[(FileStore
向量+FTS索引)] + Search --> Store ``` -### 上下文压缩机制 +--- -#### 上下文压缩 +#### 1. check_context — 上下文检查 -[Compactor](reme/memory/file_based/compactor.py) 使用 ReActAgent 将历史对话压缩为结构化的**上下文检查点**： - -| 字段 | 说明 | -|-----------------------|-----------------------| -| `## Goal` | 🎯 用户要完成的目标（可多项） | -| `## Constraints` | ⚙️ 用户提到的约束和偏好 | -| `## Progress` | 📈 已完成 / 进行中 / 阻塞的任务 | -| `## Key Decisions` | 🔑 做出的决策及简短理由 | -| `## Next Steps` | 🗺️ 下一步行动计划（有序列表） | -| `## Critical Context` | 📌 文件路径、函数名、错误信息等关键数据 | - -支持**增量更新**：传入 `previous_summary` 时，自动将新对话与旧摘要合并，保留历史进展。 - -#### 工具结果压缩 - -[ToolResultCompactor](reme/memory/file_based/tool_result_compactor.py) 解决工具输出过长（比如 browser use）导致上下文膨胀的问题： +[ContextChecker](reme/memory/file_based/component/context_checker.py) 基于 Token 计数判断上下文是否超限，自动拆分为「待压缩」和「保留」两组消息。 ```mermaid graph LR - A[tool_result 消息] --> B{内容长度 > threshold?} - B -->|否| C[保留原样] - B -->|是| D[截断到 threshold 字符] - D --> E[完整内容写入 tool_result/uuid.txt] - E --> F[消息中追加文件引用路径] + M[messages] --> H[AsMsgHandler
Token 计数] + H --> C{total > threshold?} + C -->|否| K[返回全部消息] + C -->|是| S[从尾部向前保留
reserve tokens] + S --> CP[messages_to_compact
早期消息] + S --> KP[messages_to_keep
近期消息] + S --> V{is_valid
工具调用对齐?} ``` -过期文件（超过 `retention_days`）在 `start` / `close` / `compact_tool_result` 时自动清理。 +- **核心逻辑**：从尾部向前保留 `reserve` tokens，超出部分标记为待压缩 +- **完整性保证**：不拆分 user-assistant 对话对，不拆分 tool_use/tool_result 配对 -### 记忆总结：ReAct + 文件工具 +--- -[Summarizer](reme/memory/file_based/summarizer.py) 采用 **ReAct + 文件工具** 模式，让 AI 自主决定写什么、写到哪： +#### 2. compact_memory — 对话压缩 + +[Compactor](reme/memory/file_based/component/compactor.py) 使用 ReActAgent 将历史对话压缩为**结构化上下文摘要**。 ```mermaid graph LR - A[接收对话] --> B{思考: 有什么值得记录?} - B --> C[行动: read memory/YYYY-MM-DD.md] - C --> D{思考: 如何与现有内容合并?} - D --> E[行动: edit 更新文件] - E --> F{思考: 还有遗漏吗?} - F -->|是| B - F -->|否| G[完成] + M[messages] --> H[AsMsgHandler
format_msgs_to_str] + H --> A[ReActAgent
reme_compactor] + P[previous_summary] -->|增量更新| A + A --> S[结构化摘要
Goal/Progress/Decisions...] ``` -[FileIO](reme/memory/file_based/file_io.py) 提供文件操作工具集： +**摘要结构**（上下文检查点）： -| 工具 | 功能 | 使用场景 | -|---------|---------------|---------------| -| `read` | 读取文件内容（支持行范围） | 查看现有记忆，避免重复写入 | -| `write` | 覆盖写入文件 | 创建新记忆文件或大幅重构 | -| `edit` | 精确匹配后替换 | 追加新内容或修改特定段落 | +| 字段 | 说明 | +|-----------------------|--------------------| +| `## Goal` | 用户目标 | +| `## Constraints` | 约束和偏好 | +| `## Progress` | 任务进展 | +| `## Key Decisions` | 关键决策 | +| `## Next Steps` | 下一步计划 | +| `## Critical Context` | 文件路径、函数名、错误信息等关键数据 | -### 会话内存管理 +- **增量更新**：传入 `previous_summary` 时，自动将新对话与旧摘要合并 -[ReMeInMemoryMemory](reme/memory/file_based/reme_in_memory_memory.py) 扩展了 AgentScope 的 `InMemoryMemory`： +--- -| 功能 | 说明 | -|----------------------------------|---------------------------| -| `get_memory` | 按标记过滤消息，自动在头部追加压缩摘要 | -| `estimate_tokens` | 精确估算当前上下文 Token 用量及使用率 | -| `get_history_str` | 生成人类可读的对话历史摘要（含 Token 统计） | -| `state_dict` / `load_state_dict` | 支持状态序列化 / 反序列化（会话持久化） | -| `mark_messages_compressed` | 标记消息为已压缩状态 | -| `get_compressed_summary` | 获取已压缩的摘要内容 | +#### 3. summary_memory — 记忆持久化 -### 记忆检索 - -[MemorySearch](reme/memory/tools/chunk/memory_search.py) 提供**向量 + BM25 混合检索**能力： - -| 检索方式 | 优势 | 劣势 | -|-------------|-----------------|----------------| -| **向量语义** | 捕捉意义相近但措辞不同的内容 | 对精确 token 匹配较弱 | -| **BM25 全文** | 精确 token 命中效果极佳 | 无法理解同义词和改写 | - -**融合机制**：两路召回后按权重加权求和（向量 0.7 + BM25 0.3），自然语言与精确查找均可命中。 +[Summarizer](reme/memory/file_based/component/summarizer.py) 采用 **ReAct + 文件工具** 模式，让 AI 自主决定写什么、写到哪。 ```mermaid graph LR - Q[搜索查询] --> V[向量搜索 × 0.7] -Q --> B[BM25 × 0.3] -V --> M[去重 + 加权融合] -B --> M -M --> R[Top-N 结果] + M[messages] --> A[ReActAgent
reme_summarizer] + A -->|read| R[读取 memory/YYYY-MM-DD.md] + R --> T{思考: 如何合并?} + T -->|write| W[覆盖写入] + T -->|edit| E[精确替换] + W --> F[memory/YYYY-MM-DD.md] + E --> F ``` +**文件工具**（[FileIO](reme/memory/file_based/tools/file_io.py)）： + +| 工具 | 功能 | +|---------|---------| +| `read` | 读取文件内容 | +| `write` | 覆盖写入文件 | +| `edit` | 精确匹配后替换 | + +--- + +#### 4. compact_tool_result — 工具结果压缩 + +[ToolResultCompactor](reme/memory/file_based/component/tool_result_compactor.py) 解决工具输出过长导致上下文膨胀的问题。 + +```mermaid +graph LR + M[messages] --> L{遍历 tool_result
len > threshold?} + L -->|否| K[保留原样] + L -->|是| T[truncate_text
截断到 threshold] + T --> S[完整内容写入
tool_result/uuid.txt] + S --> R[消息追加文件路径引用] + R --> C[cleanup_expired_files
清理过期文件] +``` + +- **自动清理**：过期文件（超过 `retention_days`）在 `start`/`close`/`compact_tool_result` 时自动删除 + +--- + +#### 5. memory_search — 记忆检索 + +[MemorySearch](reme/memory/file_based/tools/memory_search.py) 提供**向量 + BM25 混合检索**能力。 + +```mermaid +graph LR + Q[query] --> E[Embedding
向量化] + E --> V[vector_search
语义相似] + Q --> B[BM25
关键词匹配] + V -->|" weight: 0.7 "| M[去重 + 加权融合] + B -->|" weight: 0.3 "| M + M --> F[min_score 过滤] + F --> R[Top-N 结果] +``` + +- **融合机制**：向量权重 0.7 + BM25 权重 0.3，兼顾语义相似和精确匹配 + +--- + +#### 6. ReMeInMemoryMemory — 会话内存 + +[ReMeInMemoryMemory](reme/memory/file_based/reme_in_memory_memory.py) 扩展 AgentScope 的 `InMemoryMemory`，提供 Token +感知的内存管理。 + +```mermaid +graph LR + C[content] --> G[get_memory
exclude_mark=COMPRESSED] + G --> F[排除已压缩消息] + F --> P{prepend_summary?} + P -->|是| S[头部插入 previous-summary] + S --> O[输出 messages] + P -->|否| O +``` + +| 功能 | 说明 | +|----------------------------------|-------------------| +| `get_memory` | 按标记过滤，自动追加压缩摘要 | +| `estimate_tokens` | 估算上下文 Token 用量 | +| `state_dict` / `load_state_dict` | 状态序列化/反序列化（会话持久化） | + +--- + +#### 7. pre_reasoning_hook — 推理前预处理 + +整合上述组件的统一入口，在每轮推理前自动管理上下文。 + +```mermaid +graph LR + M[messages] --> TC[compact_tool_result
压缩超长工具输出] + TC --> CC[check_context
计算剩余空间] + CC --> D{messages_to_compact
非空?} + D -->|否| K[返回原消息 + 原摘要] + D -->|是| V{is_valid?} + V -->|否| K + V -->|是| CM[compact_memory
同步生成摘要] + V -->|是| SM[add_async_summary_task
异步持久化] + CM --> R[返回 messages_to_keep + 新摘要] +``` + +**执行流程**： + +1. `compact_tool_result` — 压缩超长工具输出 +2. `check_context` — 检查上下文是否超限 +3. `compact_memory` — 生成压缩摘要（同步） +4. `summary_memory` — 持久化记忆（异步后台） + --- ## 🗃️ 基于向量库的记忆系统 [ReMe Vector Based](reme/reme.py) 是基于向量库的记忆系统核心类，支持三种记忆类型的统一管理： -| 记忆类型 | 用途 | 使用场景 | -|--------------|------------------|-------------| -| **个人记忆** | 记录用户偏好、习惯 | `user_name` | -| **任务/程序性记忆** | 记录任务执行经验、成功/失败模式 | `task_name` | -| **工具记忆** | 记录工具使用经验、参数优化 | `tool_name` | +| 记忆类型 | 用途 | +|--------------|------------------| +| **个人记忆** | 记录用户偏好、习惯 | +| **任务/程序性记忆** | 记录任务执行经验、成功/失败模式 | +| **工具记忆** | 记录工具使用经验、参数优化 | ### 核心能力 @@ -302,24 +366,11 @@ M --> R[Top-N 结果] | `delete_memory` | 🗑️ 删除记忆 | 删除指定记忆 | | `list_memory` | 📋 列出记忆 | 列出某类记忆，支持过滤和排序 | -### 安装 +### 安装与环境变量 -```bash -pip install -U reme-ai -``` +安装和环境变量配置与 [ReMeLight 一致](#安装)，通过环境变量设置 API 密钥，可写在项目根目录的 `.env` 文件中。 -### 环境变量 - -API 密钥通过环境变量设置，可写在项目根目录的 `.env` 文件中： - -| 环境变量 | 说明 | 示例 | -|----------------------|--------------------------|-----------------------------------------------------| -| `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 @@ -413,7 +464,7 @@ if __name__ == "__main__": ### 技术架构 ```mermaid -graph TB +graph LR User[用户 / Agent] --> ReMe[Vector Based ReMe] ReMe --> Summarize[记忆总结] ReMe --> Retrieve[记忆检索] @@ -432,6 +483,42 @@ graph TB ToolRet --> VectorStore ``` +### 实验效果 + +Coming soon... + +--- + +## 🧪 程序化记忆论文 + +> 我们的程序性（任务）记忆论文已在 [arXiv](https://arxiv.org/abs/2512.10696) 发布 + +### 🌍 [Appworld 实验](benchmark/appworld/quickstart.md) + +我们在 Appworld 环境上使用 Qwen3-8B（非思考模式）进行评测： + +| 方法 | Avg@4 | Pass@4 | +|---------|---------------------|---------------------| +| 无 ReMe | 0.1497 | 0.3285 | +| 使用 ReMe | 0.1706 **(+2.09%)** | 0.3631 **(+3.46%)** | + +Pass@K 衡量在生成 K 个候选中，至少一个成功完成任务（score=1）的概率。 +当前实验使用的是内部 AppWorld 环境，可能与对外版本存在轻微差异。 + +关于如何复现实验的更多细节，见 [quickstart.md](benchmark/appworld/quickstart.md) + +### 🔧 [BFCL-V3 实验](benchmark/bfcl/quickstart.md) + +我们在 BFCL-V3 multi-turn-base 任务（随机划分 50 train / 150 val）上，使用 Qwen3-8B（思考模式）进行评测： + +| 方法 | Avg@4 | Pass@4 | +|---------|---------------------|---------------------| +| 无 ReMe | 0.4033 | 0.5955 | +| 使用 ReMe | 0.4450 **(+4.17%)** | 0.6577 **(+6.22%)** | + +关于如何复现实验的更多细节，见 [quickstart.md](benchmark/bfcl/quickstart.md) + + ## ⭐ 社区与支持 - **Star 与 Watch**：Star 可让更多智能体开发者发现 ReMe；Watch 可助你第一时间获知新版本与特性。 diff --git a/docs/cookbook/appworld/quickstart.md b/benchmark/appworld/quickstart.md similarity index 100% rename from docs/cookbook/appworld/quickstart.md rename to benchmark/appworld/quickstart.md diff --git a/docs/cookbook/bfcl/quickstart.md b/benchmark/bfcl/quickstart.md similarity index 100% rename from docs/cookbook/bfcl/quickstart.md rename to benchmark/bfcl/quickstart.md 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/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/deprecated.txt b/docs/deprecated.txt deleted file mode 100644 index 06f9648f..00000000 --- a/docs/deprecated.txt +++ /dev/null @@ -1,13 +0,0 @@ -from loguru import logger - -用英文注释，完善module/class/function docstring，要一句话简洁，不要变更代码逻辑，符合pep和pylint规范，使用list而不是typing.List/Dict，不使用typing.Union - -看看代码有什么问题 -用英文注释，完善module/class/function docstring，要一句话简洁，代码要简洁，符合pep和pylint规范，使用list而不是typing.List，不使用typing.Union -C0114: Missing module docstring (missing-module-docstring) -C0115: Missing class docstring (missing-class-docstring) -C0116: Missing function or method docstring (missing-function-docstring) -done: { for f in ./*.py; do [[ "$f" != "./__init__.py" ]] && grep -v '^[[:space:]]*#' "$f"; done; } | pbcopy - -然后是一个完整的tests，但是不要用其他的包，只是test开头的函数或者类，要求from loguru import logger -写一个测试文件，不要使用pytest，普通的test，要求英文注释 \ No newline at end of file diff --git a/docs/future_work.md b/docs/future_work.md deleted file mode 100644 index 807abfed..00000000 --- a/docs/future_work.md +++ /dev/null @@ -1,18 +0,0 @@ -# Future Work - -- [ ] P0 ReMe documentation style migration: Recommend using the same doc and jupyter structure as Agentscope Runtime @jiaji -- [ ] P0 ReMe integration with agentscope Personal/Task/Tool @jinli -- [ ] P0 ReMe sample library examples [show case](https://github.com/agentscope-ai/agentscope-samples/tree/main/functionality/long_term_memory_mem0) @jinli -- [ ] P0 Decouple flowllm dependencies @jinli -- [ ] P0 ReMe support for import, improve code documentation @jinli -- [ ] P1 ReMe integration with asio tool_memory @jinli -- [ ] P2 ReMe integration with agentscope-Runtime tool_memory @jinli - -- [ ] P0 Task Memory Research Paper @zhoyin - -- [ ] P1 Context interface definition @jinli - -- [ ] P2 Database layer interface unification @jinli -- [ ] P2 Automatic Tool Exploration Mode @wangcan -- [ ] P2 Mem-Agent Exploration @weikang -- [ ] P2 Desktop Pet Personal Assistant diff --git a/docs/reme_v2_design.md b/docs/reme_v2_design.md deleted file mode 100644 index 6fa78140..00000000 --- a/docs/reme_v2_design.md +++ /dev/null @@ -1,735 +0,0 @@ -# ReMeV2 深度设计文档：渐进式 Agentic Memory 方案 - -## 一、 背景与现状分析 - -### 1.1 当前面临的挑战 - -* **外功修炼（接口易用性）**：现有的 `server-client` 模式对新手开发者不够友好，集成成本高，需要更直观、纯 Pythonic 的调用方式。 -* **内功修炼（架构深度）**：受 `skills` 和 `agentic memory` 启发，现有的存储检索较为机械。我们需要一种基于**渐进式检索（Progressive Retrieval）**与**渐进式总结（Progressive Summarization）**的智能体记忆方案。 - -### 1.2 核心目标 - -1. **极简开发体验**：开发者友好，全异步接口，支持本地直接运行与 CLI 体验。 -2. **认知架构升级**：引入 渐进式检索 & 渐进式总结 的 Agentic 模式，融合多种记忆，让记忆的存取具备“思考”过程。 -3. **生态融合**：原生支持 AgentScope、LangChain 等主流框架。 - ---- - -## 二、 竞品调研与启示 - -### 2.1 主流竞品深度对比 - -| 产品 | 设计哲学 | 核心优势 | 局限性 | -|-------------|----------|---------------------------------------------|-------------------| -| **mem0** | 智能便签本 | 原子事实提取，极高 Token 效率。 | 缺乏对复杂逻辑链条的支持。 | -| **Letta** | 带硬盘的 CPU | 模拟计算机三级存储（Core/Recall/Archival），Agent 自主控存。 | 状态机管理相对复杂。 | -| **MIRIX** | 认知架构图谱 | 实体-关系双引擎，支持记忆“进化”与“固化”。 | 侧重研究，落地集成门槛较高。 | -| **LangMem** | 用户档案系统 | 异步 Compaction（压缩），Schema 驱动，强一致性。 | 偏向 SaaS 应用，灵活性略逊。 | - -### 2.2 mem0 -- https://github.com/mem0ai/mem0 -- https://docs.mem0.ai/core-concepts/memory-operations/add -- https://docs.mem0.ai/core-concepts/memory-operations/search -- https://docs.mem0.ai/core-concepts/memory-operations/update -- https://docs.mem0.ai/core-concepts/memory-operations/delete - -#### 2.2.1 API Reference -| 接口名称 | 核心输入参数 (Inputs) | 核心输出 (Outputs) | 背后逻辑 (Internal Logic) | -| --- | --- | --- | --- | -| **Add** | `messages` (文本/对话), `user_id`, `metadata` | `id`, `event` (ADD/UPDATE), `data` | **提取与合并**：LLM 提取事实，自动去重并更新已有记忆，而非简单堆叠。 | -| **Search** | `query` (自然语言), `filters`, `limit` | `id`, `memory` (事实文本), `score`, `metadata` | **语义检索**：基于向量相似度查找最相关的“原子事实”，支持多维过滤。 | -| **Update** | `memory_id` (必填), `data` (新内容) | 操作状态 (Success/Fail) | **手动干预**：允许开发者对特定的事实进行精确修正。 | -| **Delete** | `memory_id` 或 `user_id` (清空) | 操作状态 (Success/Fail) | **遗忘机制**：物理删除或逻辑移除不再需要的信息。 | - -#### 2.2.2 Tech Strategy & Benefits -| 维度 | 技术方案 (Technical Solution) | 核心优势 (Key Advantages) | -| --- | --- | --- | -| **存储架构** | **混合存储**：向量数据库 (Vector) + 图数据库 (Graph) + 关系型元数据。 | **多维关联**：不仅能搜到相似内容，还能理解实体间的逻辑关系（如“父子”、“因果”）。 | -| **数据处理** | **原子化事实提取**：利用 LLM 将长篇对话压缩为简短的 Fact。 | **极高 Token 效率**：注入 Prompt 的内容更精炼，减少 90% 以上的冗余信息，大幅降本。 | -| **管理层级** | **多级联动**：User (长期) Agent (专业) Session (短期)。 | **个性化定制**：实现跨会话的“长效记忆”，AI 能记住用户一个月前说过的偏好。 | -| **冲突处理** | **自适应更新算法**：新信息进入时自动比对旧记忆。 | **数据一致性**：自动处理矛盾信息（如用户更换了住址），确保记忆库始终是“最新真理”。 | -| **兼容性** | **解耦设计**：支持多种 Embedding 模型与向量数据库后端。 | **快速集成**：几行代码即可为现有 LLM 应用增加记忆层，适配各种生产环境。 | - - ---- - -### 2.3 Letta -- https://github.com/letta-ai/letta -- https://docs.letta.com/guides/agents/archival-memory/ -- https://docs.letta.com/guides/agents/archival-search/ - -#### 2.3.1 存储架构层级 (Memory Tiering) - -Letta 将记忆分为三个物理/逻辑层，模拟计算机的存储架构： - -| 记忆层级 | 存储介质 | 访问方式 | 核心作用 | -| --- | --- | --- | --- | -| **Core Memory** | **上下文窗口 (Prompt)** | 直接读写 | **即时意识**：包含 `Persona`（AI 设定）和 `Human`（用户信息）。Agent 随时可见，响应最快。 | -| **Recall Memory** | **关系型数据库 (SQL)** | 分页检索 | **短期/历史回顾**：存储完整的对话流（Messages）。用于回答“你刚才说了什么”。 | -| **Archival Memory** | **向量数据库 (Vector)** | 语义搜索 | **长期知识库**：存储海量事实或文档。Agent 通过工具自主检索或存入。 | - -#### 2.3.2 核心操作接口 (API & Tool Reference) - -在 Letta 中，记忆的操作通常封装为 **Tools**，由 Agent 根据推理需求主动调用。 - -| 接口/工具名称 | 输入参数 (Inputs) | 核心输出 (Outputs) | 背后逻辑 (Internal Logic) | -| --- | --- | --- | --- | -| **`core_memory_update`** | `section`, `new_content` | 更新后的段落内容 | **原子替换**：直接修改 System Prompt 中的特定块（如：更新用户的职业或 AI 的性格偏好）。 | -| **`archival_memory_insert`** | `content` (字符串) | 写入状态/ID | **知识沉淀**：将当前对话中的重要信息或外部文件片段“持久化”到向量数据库。 | -| **`archival_memory_search`** | `query`, `page` | 匹配的文本块列表 | **主动 RAG**：Agent 意识到知识不足时，自主发起向量检索，并将结果拉入临时上下文。 | -| **`conversation_search`** | `query`, `start_date` | 历史消息记录 | **全文检索**：在 Recall Memory 中根据关键词或时间戳查找历史对话详情。 | -| **`send_message`** | `message`, `agent_id` | 响应流/状态更新 | **状态循环**：这是主入口，触发 Agent 的“思考-行动-观察”循环，自动处理内存同步。 | - -#### 2.3.3 技术策略与核心优势 (Tech Strategy & Benefits) - -| 维度 | 技术方案 (Technical Solution) | 核心优势 (Key Advantages) | -| --- | --- | --- | -| **状态持久化** | **Agent State Snapshot**：将 Agent 的所有内存、工具定义和历史记录打包存入数据库。 | **无限存续**：Agent 不再是无状态的 API 调用。重启服务器后，Agent 依然记得所有细节。 | -| **自主演进** | **Self-Editing Loop**：Agent 拥有修改自己 Core Memory 的权限（通过函数调用）。 | **认知闭环**：AI 能在交流中发现矛盾并自我更正，例如发现用户搬家后自动更新 `Human` 模块。 | -| **算力调度** | **OOC (Out-of-Context) 管理**：当对话过长，系统自动将旧消息从 Core 移入 Recall。 | **突破 Context 限制**：在 8k 窗口的模型上也能处理相当于 1M 窗口的逻辑量，且成本更低。 | -| **多代理协同** | **Letta Server 中控**：统一管理多个 Agent 的状态机与资源访问权限。 | **企业级扩展**：支持创建 Agent 团队，每个 Agent 拥有独立的记忆空间但可共享 Archival 库。 | -| **解耦灵活性** | **Provider Agnostic**：后端支持 Postgres/Chroma，前端支持 OpenAI/Anthropic/Local LLMs。 | **无缝迁移**：不绑定特定模型，开发者可以根据成本或能力随时更换底座。 | - -#### 2.3.4 与 mem0 的深度对比 - -* **设计哲学**： -* **mem0** 像是一个**“智能记事本”**，它在后台默默地帮你总结事实。 -* **Letta** 像是一个**“带硬盘的 CPU”**，它把记忆管理完全交给了 Agent 自己的逻辑推理。 - - -* **交互模式**： -* **mem0** 通常是外部干预（Add/Search）。 -* **Letta** 强调 **Agentic Control**（Agent 意识到需要搜索时才去搜索），这种模式更接近人类的思维过程。 - ---- - -### 2.4 MIRIX -- https://github.com/Mirix-AI/MIRIX -- https://docs.mirix.io/ - -#### 2.4.1 API Reference - -| 接口名称 | 核心输入参数 (Inputs) | 核心输出 (Outputs) | 背后逻辑 (Internal Logic) | -| --- | --- | --- | --- | -| **Add** | `content` (观察/对话), `agent_id`, `context_type` (如任务/闲聊) | `memory_id`, `graph_nodes`, `status` | **实体建模**：不只是提取事实，而是将信息拆解为实体（Entities）与关系（Relations），并挂载到智能体的知识图谱中。 | -| **Query** | `query` (意图), `scope` (全局/局部), `top_k` | `retrieved_memories`, `relation_paths`, `score` | **混合检索**：结合向量（Vector）的语义相关性和图（Graph）的拓扑连接性，寻找具有逻辑深度背景的记忆。 | -| **Evolve** | `target_memories` (可选), `agent_id` | `optimized_structure`, `merged_nodes` | **记忆固化/压缩**：模仿人类大脑的“睡眠”机制，自动合并碎片化记忆，将短期经验转化为长期的结构化知识。 | -| **Observe** | `interaction_stream`, `feedback` | `insights`, `priority_update` | **实时学习**：根据用户反馈或环境变化，动态调整记忆的权重（Importance）和置信度。 | - -#### 2.4.2 Tech Strategy & Benefits - -| 维度 | 技术方案 (Technical Solution) | 核心优势 (Key Advantages) | -| --- | --- | --- | -| **存储架构** | **语义-关系双引擎**：向量索引（Vector Index）+ 属性图（Property Graph）。 | **深度上下文**：不仅知道“是什么”，还能通过图路径推理出“为什么”，有效解决 LLM 幻觉问题。 | -| **记忆层级** | **三层架构**：感知记忆 (Perception) -> 语义记忆 (Semantic) -> 经验记忆 (Episodic)。 | **任务适应性**：不同任务自动匹配不同的记忆深度，短期任务关注细节，长期任务关注模式。 | -| **演化机制** | **自主固化 (Self-Consolidation)**：通过 LLM 定期对冗余、矛盾信息进行清洗和逻辑抽象。 | **永久生命力**：解决随时间推移记忆库膨胀导致的检索噪声，确保记忆库“越用越聪明”。 | -| **推理增强** | **基于记忆的 RAG+**：在检索到的事实基础上，额外提供关联的逻辑链条（Logic Chains）。 | **辅助决策**：为 Agent 提供决策支撑，使其在处理复杂流程时具备类似“长期经验值”的直觉。 | -| **多代理协同** | **内存共享协议**：支持 Agent 之间的记忆交换与知识同步。 | **群体智能**：多个 Agent 可以共享同一套底层知识体系，同时保留各自的私有工作记忆。 | - -#### 2.4.3 与 mem0 的主要区别 - -* **Mem0** 侧重于**个性化偏好存储**（Personalization），核心是记住“用户喜欢什么”。 -* **MIRIX** 侧重于**智能体认知架构**（Agent Cognition），核心是让 Agent 具备类似人类的“知识归纳”和“逻辑推理”记忆能力。 - ---- - -### 2.5 LangMem -- https://github.com/langchain-ai/langmem -- https://langchain-ai.github.io/langmem/ - -#### 2.5.1 API Reference - -| 接口名称 | 核心输入参数 (Inputs) | 核心输出 (Outputs) | 背后逻辑 (Internal Logic) | -| --- | --- | --- | --- | -| **Add Messages** | `thread_id`, `messages` (List), `user_id` | 操作确认 / 任务 ID | **流式注入**：将原始对话追加到指定的 Thread。LangMem 会自动关联用户上下文，准备进行后续的异步处理。 | -| **Query Memory** | `user_id`, `query` (语义描述), `namespace` | 结构化记忆对象 (JSON / Text) | **多维检索**：不仅支持向量相似度搜索，还能根据定义的 Schema 返回结构化的用户画像或知识状态。 | -| **Trigger Logic** | `thread_id`, `memory_type` | 更新后的 Memory State | **异步固化**：后台启动 LLM 任务，将长篇对话“压缩”并“提取”到长期存储中。支持自定义提取逻辑（如更新用户信息）。 | -| **Manage State** | `user_id`, `patch_data` (增量更新) | 成功/失败 状态 | **精确受控**：开发者可以直接修改持久化的状态（State），支持类似于 Git 的状态管理。 | - -#### 2.5.2 Tech Strategy & Benefits - -| 维度 | 技术方案 (Technical Solution) | 核心优势 (Key Advantages) | -| --- | --- | --- | -| **存储架构** | **Stateful Persistence**：基于关系型数据库 (Postgres) + 向量索引。 | **强一致性**：利用数据库事务确保记忆更新的可靠性，支持复杂的结构化查询与过滤。 | -| **数据处理** | **异步化 Compaction (压缩)**：在对话间隙通过后台 Worker 提取知识。 | **无感延迟**：核心对话流程不被记忆提取阻塞，通过定时或事件驱动完成“记忆固化”，优化用户体验。 | -| **管理层级** | **Thread -> User -> Organization**：三层级联记忆。 | **上下文隔离**：完美适配 SaaS 应用场景，既能记住单次对话（Thread），也能沉淀用户习惯（User）。 | -| **逻辑引擎** | **Schema-Driven (模式驱动)**：允许定义 JSON Schema 来规范记忆内容。 | **高度可预测**：输出不再是散乱的句子，而是结构化的字段，方便下游程序直接调用逻辑（如自动填充表单）。 | -| **集成生态** | **LangGraph 原生集成**：作为 Checkpointer 或存储节点直接接入。 | **生态协同**：如果你已经在用 LangChain，LangMem 可以无缝接管状态流转，无需重写底层存储逻辑。 | - -#### 2.5.3 与 mem0 的核心差异 - -* **mem0** 像是一个**“便签本”**：它擅长从每一句话里抠出零散的事实（如“我喜欢吃苹果”），然后把它们存成一条条语义片段。 -* **LangMem** 像是一个**“用户档案系统”**：它更擅长分析一整段对话，然后更新一个复杂的 JSON 档案（如更新用户的偏好模型、性格标签、历史任务状态）。 - ---- - -## 三、 ReMeV2 API 接口设计 - -### 3.1 Long-Term Memory (长期记忆) - -#### 3.1.1 Basic Usage (基础用法) - -The most straightforward way to use ReMe for long-term memory management. Supports basic summary and retrieval operations. - -```python -import os -from reme_ai import ReMe - -os.environ["REME_LLM_API_KEY"] = "sk-..." -os.environ["REME_LLM_BASE_URL"] = "https://dashscope.aliyuncs.com/compatible-mode/v1" -os.environ["REME_EMBEDDING_API_KEY"] = "sk-..." -os.environ["REME_EMBEDDING_BASE_URL"] = "https://dashscope.aliyuncs.com/compatible-mode/v1" - -memory = ReMe( - memory_space="remy", # workspace identifier - llm={"backend": "openai", "model": "qwen-plus", "temperature": 0.6}, - embedding={"backend": "openai", "model": "text-embedding-v4", "dimension": 1024}, - vector_store={"backend": "local_file"}, # supported: local_file, chromadb, qdrant, etc. -) - -# Summarize conversation into memory -result = await memory.summary( - messages=[ - {"role": "user", "content": "I'm travelling to SF"}, - {"role": "assistant", "content": "That's great to hear!"} - ], - user_id="Alice", - # memory_type="auto" # default: auto (auto, personal, procedural, tool) -) - -# Retrieve relevant memories -memories = await memory.retrieve( - query="what is your travel plan?", - limit=3, - user_id="Alice", - # memory_type="auto" # default: auto -) -memories_str = "\n".join(f"- {m['memory']}" for m in memories["results"]) -print(memories_str) -``` - -#### 3.1.2 CLI Chat Application (命令行聊天应用) - -A complete example demonstrating how to build a memory-enhanced chatbot with CLI interface. - -```python -import os -from reme_ai import ReMe -from openai import OpenAI - -os.environ["REME_LLM_API_KEY"] = "sk-..." -os.environ["REME_LLM_BASE_URL"] = "https://dashscope.aliyuncs.com/compatible-mode/v1" -os.environ["REME_EMBEDDING_API_KEY"] = "sk-..." -os.environ["REME_EMBEDDING_BASE_URL"] = "https://dashscope.aliyuncs.com/compatible-mode/v1" - -memory = ReMe( - memory_space="remy", - llm={"backend": "openai", "model": "qwen-plus", "temperature": 0.6}, - embedding={"backend": "openai", "model": "text-embedding-v4", "dimension": 1024}, - vector_store={"backend": "local_file"}, -) - -os.environ["OPENAI_API_KEY"] = "sk-..." -os.environ["OPENAI_BASE_URL"] = "https://dashscope.aliyuncs.com/compatible-mode/v1" -openai_client = OpenAI() - -def chat_with_memories( - query: str, - history_messages: list[dict], - user_name: str = "", - start_summary_size: int = 2, - keep_size: int = 0 -) -> str: - # Retrieve relevant memories for the query - memories = memory.retrieve(query=query, user_id=user_name, limit=3) - - # Build system prompt with memories - system_prompt = ( - "You are a helpful AI named `Remy`. Use the user memories to answer the question. " - "If you don't know the answer, just say you don't know. Don't try to make up an answer.\n" - ) - if memories: - memories_str = "\n".join(f"- {m['memory']}" for m in memories["results"]) - system_prompt += f"User Memories:\n{memories_str}\n" - - # Generate response - system_message = {"role": "system", "content": system_prompt} - history_messages.append({"role": "user", "content": query}) - response = openai_client.chat.completions.create( - model="qwen-plus", - messages=[system_message] + history_messages - ) - history_messages.append({"role": "assistant", "content": response.choices[0].message.content}) - - # Summarize history when it gets too long - if len(history_messages) >= start_summary_size: - memory.summary(history_messages[:-keep_size], user_id=user_name) - print("Current memories: " + memory.list_memories(user_id=user_name)) - history_messages = history_messages[-keep_size:] - - return history_messages[-1]["content"] - -def main(): - user_name = input("Enter your name: ").strip() - print("Chat with Remy (type 'exit' to quit)") - - messages = [] - while True: - user_input = input(f"{user_name}: ").strip() - if user_input.lower() == 'exit': - print("Goodbye!") - break - - print(f"Remy: {chat_with_memories(user_input, messages, user_name)}") - - # Cleanup - memory.delete_all_memories(user_id=user_name) - print("All memories deleted") - -if __name__ == "__main__": - main() -``` - -#### 3.1.3 Advanced Usage (高级用法) - -For advanced users who want to customize retriever and summarizer behavior with Agentic mode. - -```python -import os -from reme_ai import ReMe -from reme_ai.retriever import AgenticRetriever -from reme_ai.summarizer import AgenticSummarizer -from reme_ai.tools import ATool, BTool, CTool - -os.environ["REME_LLM_API_KEY"] = "sk-..." -os.environ["REME_LLM_BASE_URL"] = "https://dashscope.aliyuncs.com/compatible-mode/v1" -os.environ["REME_EMBEDDING_API_KEY"] = "sk-..." -os.environ["REME_EMBEDDING_BASE_URL"] = "https://dashscope.aliyuncs.com/compatible-mode/v1" - -memory = ReMe( - memory_space="remy", - llm={"backend": "openai", "model": "qwen-plus", "temperature": 0.6}, - embedding={"backend": "openai", "model": "text-embedding-v4", "dimension": 1024}, - vector_store={"backend": "local_file"}, - use_agentic_mode=True, -) - -# Customize retriever and summarizer with custom tools and prompts -memory.set_retriever( - AgenticRetriever(tools=[ATool(), BTool(), CTool()]), - system_prompt="Custom retrieval instructions..." -) -memory.set_summarizer( - AgenticSummarizer(tools=[ATool(), BTool(), CTool()]) -) - -# Use the customized memory system -result = memory.summary( - messages=[ - {"role": "user", "content": "I'm travelling to SF"}, - {"role": "assistant", "content": "That's great to hear!"} - ], - user_id="Alice", - memory_type="auto", # auto, personal, procedural, tool -) - -memories = memory.retrieve( - query="what is your travel plan?", - limit=3, - user_id="Alice", - memory_type="auto", -) -memories_str = "\n".join(f"- {m['memory']}" for m in memories["results"]) -print(memories_str) -``` - -### 3.2 Short-Term Memory (短期记忆) - -#### 3.2.1 Basic Usage (基础用法) - -Context offload/reload API for managing short-term conversational memory within a session. - -```python -import os -from reme_ai import ReMe - -os.environ["REME_LLM_API_KEY"] = "sk-..." -os.environ["REME_LLM_BASE_URL"] = "https://dashscope.aliyuncs.com/compatible-mode/v1" -os.environ["REME_EMBEDDING_API_KEY"] = "sk-..." -os.environ["REME_EMBEDDING_BASE_URL"] = "https://dashscope.aliyuncs.com/compatible-mode/v1" - -memory = ReMe( - memory_space="remy", - llm={"backend": "openai", "model": "qwen-plus", "temperature": 0.6}, - embedding={"backend": "openai", "model": "text-embedding-v4", "dimension": 1024}, - vector_store={"backend": "local_file"}, -) - -# Offload context when conversation gets too long -result = memory.offload_context( - messages=[ - {"role": "user", "content": "I'm travelling to SF"}, - {"role": "assistant", "content": "That's great to hear!"} - ], -) - -# Reload relevant context when needed -memories = memory.reload_context( - query="what is your travel plan?", - limit=3, -) -memories_str = "\n".join(f"- {m['memory']}" for m in memories["results"]) -print(memories_str) -``` - -### 3.3 Framework Integration (框架集成) - -#### 3.3.1 Integration with AgentScope - -Integration example for AgentScope ReActAgent with long-term memory support. - -```python -# TODO: Provide AgentScope integration example -``` - -#### 3.3.2 Integration with LangChain - -Integration example for LangChain agents with ReMe memory layer. - -```python -# TODO: Provide LangChain integration example -``` - -### 3.4 OpenAI Compatible Interface - -OpenAI-compatible API interface for seamless integration with existing OpenAI-based applications. - -```python -# TODO: Research and implement OpenAI-compatible interface -# - Support for threads and assistants API -# - Compatible with OpenAI SDK -# - Support for streaming responses -``` - - - ---- - -## 四、核心方案设计 - -### 4.1 设计概述 - -ReMeV2 采用简洁的架构设计，核心理念为：**ReMeV2 = Tool(s) + Agent(s)** - -- **Tool层**：提供原子化的记忆操作能力，包括增删改查、检索、元数据管理等基础操作 -- **Agent层**：基于Tool层构建的智能代理，负责复杂的记忆管理逻辑，如分类总结、渐进式检索等 -- **Runtime层**：内部调度机制，协调Tool和Agent的交互流程 - -### 4.2 Tool层设计 - -Tool层提供装饰器形式的记忆操作工具，每个工具类通过 `@tool` 装饰器注册，明确定义初始化参数和调用参数。 - -#### 4.2.1 基类：BaseMemoryToolOp - -**初始化参数：** -- `enable_multiple` (bool): Enable multi-item operation mode. Default: `True` -- `enable_thinking_params` (bool): Include thinking parameter in tool schema for model reasoning. Default: `False` -- `memory_metadata_dir` (str): Directory path for storing memory metadata. Default: `"./memory_metadata"` - -#### 4.2.2 Tool操作列表 - -以下是所有Tool操作的完整定义，包括继承关系、初始化参数和调用参数： - -| Tool类 | 继承自 | 初始化参数（除基类外） | Tool Call参数（单项模式） | Tool Call参数（多项模式） | -|----------------------------|------------------|------------------------------------------------------------------------------------------------------------------|-----------------------------------------------------------------------------------------------------------------|-----------------------------------------------------------------------------------------------------------------------| -| **AddMemoryOp** | BaseMemoryToolOp | `add_when_to_use` (bool, 默认: False)
`add_metadata` (bool, 默认: True) | `when_to_use` (str, 可选)
`memory_content` (str, 必需)
`metadata` (dict, 可选) | `memories` (array, 必需):
- `when_to_use` (str, 可选)
- `memory_content` (str, 必需)
- `metadata` (dict, 可选) | -| **UpdateMemoryOp** | BaseMemoryToolOp | 无 | `memory_id` (str, 必需)
`memory_content` (str, 必需)
`metadata` (dict, 可选) | `memories` (array, 必需):
- `memory_id` (str, 必需)
- `memory_content` (str, 必需)
- `metadata` (dict, 可选) | -| **DeleteMemoryOp** | BaseMemoryToolOp | 无 | `memory_id` (str, 必需) | `memory_ids` (array[str], 必需) | -| **VectorRetrieveMemoryOp** | BaseMemoryToolOp | `enable_summary_memory` (bool, 默认: False)
`add_memory_type_target` (bool, 默认: False)
`top_k` (int, 默认: 20) | `query` (str, 必需)
`memory_type` (str, 可选, 枚举: [identity, personal, procedural])
`memory_target` (str, 可选) | `query_items` (array, 必需):
- `query` (str, 必需)
- `memory_type` (str, 可选)
- `memory_target` (str, 可选) | -| **AddMetaMemoryOp** | BaseMemoryToolOp | 无 | `memory_type` (str, 必需, 枚举: [personal, procedural])
`memory_target` (str, 必需) | `meta_memories` (array, 必需):
- `memory_type` (str, 必需)
- `memory_target` (str, 必需) | -| **ReadMetaMemoryOp** | BaseMemoryToolOp | `enable_tool_memory` (bool, 默认: False)
`enable_identity_memory` (bool, 默认: False) | 无（无输入schema） | N/A (enable_multiple=False) | -| **AddHistoryMemoryOp** | BaseMemoryToolOp | 无 | `messages` (array[object], 必需) | N/A (enable_multiple=False) | -| **ReadHistoryMemoryOp** | BaseMemoryToolOp | 无 | `memory_id` (str, 必需) | `memory_ids` (array[str], 必需) | -| **AddSummaryMemoryOp** | AddMemoryOp | 无（继承自AddMemoryOp） | `summary_memory` (str, 必需)
`metadata` (dict, 可选) | N/A (enable_multiple=False) | -| **ReadIdentityMemoryOp** | BaseMemoryToolOp | 无 | 无（无输入schema） | N/A (enable_multiple=False) | -| **UpdateIdentityMemoryOp** | BaseMemoryToolOp | 无 | `identity_memory` (str, 必需) | N/A (enable_multiple=False) | -| **ThinkToolOp** | BaseAsyncToolOp | `add_output_reflection` (bool, 默认: False) | `reflection` (str, 必需) | N/A | -| **HandsOffOp** | BaseMemoryToolOp | 无 | `memory_type` (str, 必需, 枚举: [identity, personal, procedural, tool])
`memory_target` (str, 必需) | `memory_tasks` (array, 必需):
- `memory_type` (str, 必需)
- `memory_target` (str, 必需) | - -### 4.3 Agent层设计 - -#### 4.3.1 基类：BaseMemoryAgentOp - -Agent层构建在Tool层之上，封装复杂的记忆管理逻辑。每个Agent通过组合多个Tool实现特定的记忆管理任务。 - -#### 4.3.2 Agent操作列表 - -以下是所有Agent操作的完整定义，包括初始化参数、调用参数和可用工具： - -| Agent类 | 继承自 | 初始化参数（基类外） | Tool Call参数 | 可用工具 | -|--------------------------------|-------------------|------------------------------------------------------------------------------------|-----------------------------------------------------------------------------------------------------------------------------------------------------------------|--------------------------------------------------------------------------------------------------------------------------| -| **PersonalSummaryAgentV1Op** | BaseMemoryAgentOp | None | `workspace_id` (str, required)
`memory_target` (str, required)
`query` (str, optional)
`messages` (array, optional)
`ref_memory_id` (str, required) | add_memory
update_memory
delete_memory
vector_retrieve_memory | -| **ProceduralSummaryAgentV1Op** | BaseMemoryAgentOp | None | `workspace_id` (str, required)
`memory_target` (str, required)
`query` (str, optional)
`messages` (array, optional)
`ref_memory_id` (str, required) | add_memory
update_memory
delete_memory
vector_retrieve_memory | -| **ToolSummaryAgentV1Op** | BaseMemoryAgentOp | None | `workspace_id` (str, required)
`memory_target` (str, required)
`query` (str, optional)
`messages` (array, optional)
`ref_memory_id` (str, required) | add_memory
update_memory
vector_retrieve_memory | -| **IdentitySummaryAgentV1Op** | BaseMemoryAgentOp | None | `workspace_id` (str, required)
`query` (str, optional)
`messages` (array, optional) | read_identity_memory
update_identity_memory | -| **ReMeSummaryAgentV1Op** | BaseMemoryAgentOp | `enable_tool_memory` (bool, 默认: True)
`enable_identity_memory` (bool, 默认: True) | `workspace_id` (str, required)
`query` (str, optional)
`messages` (array, optional) | add_meta_memory
add_summary_memory
hands_off
(内部调用: add_history_memory, read_identity_memory, read_meta_memory) | -| **ReMeRetrieveAgentV1Op** | BaseMemoryAgentOp | `enable_tool_memory` (bool, 默认: True) | `workspace_id` (str, required)
`query` (str, optional)
`messages` (array, optional) | vector_retrieve_memory
read_history_memory
(内部调用: read_meta_memory) | -| **ReMyAgentV1Op** | BaseMemoryAgentOp | `enable_tool_memory` (bool, 默认: True)
`enable_identity_memory` (bool, 默认: True) | `workspace_id` (str, required)
`query` (str, optional)
`messages` (array, optional) | vector_retrieve_memory
read_history_memory
(内部调用: read_identity_memory, read_meta_memory) | - -### 4.4 Runtime层设计（内部实现） - -Runtime层负责协调Tool和Agent的调用流程，实现记忆的渐进式处理。 - -#### 4.4.1 渐进式总结流程（Summary） - -总结流程采用分层处理策略，首先保存历史对话，读取元信息，然后由主Agent协调多个专用Agent完成分类总结。 - -**流程结构：** - -```python -# Step 1: Save conversation history -AddHistoryMemoryOp() - -# Step 2: Load meta information (memory types and targets) -ReadMetaMemoryOp() - -# Step 3: Progressive summarization with delegation -ReMeSummaryAgentV1Op(tools=[ - # Add meta memory entries for new memory types/targets - AddMetaMemoryOp(list(memory_type, memory_target)), - - # Add general summary memory as fallback - AddSummaryMemoryOp(summary_memory), - - # Delegate to specialized summary agents - HandsOffOp(list(memory_type, memory_target), agents=[ - PersonalSummaryAgentV1Op, # Summarize personal memories - ProceduralSummaryAgentV1Op, # Summarize procedural memories - ToolSummaryAgentV1Op, # Summarize tool-related memories - IdentitySummaryAgentV1Op # Update identity memory - ]), -]) - -# Specialized agents and their available tools -PersonalSummaryAgentV1Op(tools=[AddMemoryOp, UpdateMemoryOp, DeleteMemoryOp, VectorRetrieveMemoryOp]) -ProceduralSummaryAgentV1Op(tools=[AddMemoryOp, UpdateMemoryOp, DeleteMemoryOp, VectorRetrieveMemoryOp]) -ToolSummaryAgentV1Op(tools=[AddMemoryOp, UpdateMemoryOp, VectorRetrieveMemoryOp]) -IdentitySummaryAgentV1Op(tools=[ReadIdentityMemoryOp, UpdateIdentityMemoryOp]) -``` - -#### 4.4.2 渐进式检索流程（Retrieve） - -检索流程采用三层检索策略，类似于技能系统的加载机制，逐层加载和过滤记忆。 - -**流程结构：** - -```python -# Progressive retrieval with three layers -ReMeRetrieveAgentV1Op(tools=[ - # Layer 0: Load meta memory (all available memory types and targets) - ReadMetaMemoryOp(), - # Output format example: - # - personal(jinli): Information about Jinli's personal life and preferences - # - personal(jiaji): Information about Jiaji's background and interests - # - personal(jinli&jiaji): Shared memories between Jinli and Jiaji - # - procedural(appworld): Procedural knowledge for AppWorld tasks - # - procedural(bfcl-v3): Procedural knowledge for BFCL-v3 benchmark - # - tool(tool_guidelines): Guidelines for tool usage - # - identity(self): Agent's self-identity information - - # Layer 1+2: Vector-based retrieval on structured memories - VectorRetrieveMemoryOp(list(memory_type, memory_target, query)), - - # Layer 3: Load full conversation history for specific memory - ReadHistoryMemoryOp(ref_memory_id), -]) -``` - -**与技能系统的类比：** - -```python -# Skill system hierarchy (for reference) -load_meta_skills # Load skill metadata -load_skills # Load skill implementations -load_reference_skills # Load detailed skill documentation -execute_shell # Execute actual commands -``` - -## 五、扩展设计与实验方向 - -### 5.1 Summary Memory机制 - -Summary Memory作为通用维度的记忆类型，提供兜底的原始对话索引能力。 - -**工作流程示例：** - -```txt -Step 1: Progressive summarization across sessions - session1: List[Message] -> session2: List[Message] -> session3: List[Message] -> ... -summary ✓ (always) ✓ (always) ✓ (always) -personal ✗ ✗ ✓ (when applicable) -procedural ✗ ✓ (when applicable) ✗ - -Step 2: Retrieval with fallback strategy -vector_retrieve_memory(query, memory_type="personal", memory_target="jinli") - -> Search in memory_type: ["personal", "summary"] # Fallback to summary if personal not found -``` - -**设计优势：** -1. Provides a universal dimension for memory extraction across all memory types -2. Ensures fallback indexing of original conversations when specific meta memory is not available -3. Maintains conversation context even when specialized memory extraction fails - -### 5.2 Thinking参数实验 - -探索不同的模型推理能力增强方案，受AgentScope和Claude启发。 - -#### 5.2.1 Thinking参数设计 - -```python -async def record_to_memory( - self, - thinking: str, - content: list[str], - **kwargs: Any, -) -> ToolResponse: - """Use this function to record important information that you may - need later. The target content should be specific and concise, e.g. - who, when, where, do what, why, how, etc. - - Args: - thinking (`str`): - Your thinking and reasoning about what to record - content (`list[str]`): - The content to remember, which is a list of strings. - """ -``` - -#### 5.2.2 实验对比方案 - -| 方案类型 | 说明 | 灵感来源 | -|-------------------------------|--------------------------------------------|----------------| -| Thinking Model | Native reasoning-capable models (e.g., o1) | OpenAI | -| Instruct Model | Standard instruction-following models | Baseline | -| Instruct Model + Thinking Params | Add thinking parameter to tool schema | AgentScope | -| Instruct Model + Thinking Tool | Dedicated thinking tool for explicit reasoning | Claude | - -### 5.3 多项操作模式实验 - -对比单次调用和批量调用的性能与准确性差异。 - -**两种模式对比：** - -| 模式 | Tool调用方式 | Model调用次数 | 优势 | 劣势 | -|--------------|----------------------------|---------------|------------------------------|--------------------------| -| 单项模式 | Single-item per call | Multiple | Fine-grained control | Higher latency, more tokens | -| 多项模式 | Batch multiple items | Single | Lower latency, fewer tokens | Potential batch errors | - -**实验目标：** -- Evaluate accuracy: single vs. batch operations -- Measure latency and token efficiency -- Identify optimal use cases for each mode - -### 5.4 多版本与扩展性 - -支持从基类继承实现自定义Agent，便于团队协作和功能迭代。 - -**扩展示例：** - -```python -# Version 2 implementations by different team members -PersonalSummaryAgentV2Op / PersonalRetrieveAgentV2Op # @weikang -ProceduralSummaryAgentV2Op / ProceduralRetrieveAgentV2Op # @zouyin - -# Inherit from BaseMemoryAgentOp -class PersonalSummaryAgentV2Op(BaseMemoryAgentOp): - """Enhanced personal memory summarization with improved algorithms""" - pass -``` - -### 5.5 文件系统集成（未来方向） - -探索将文件操作能力集成到记忆系统中，支持基于文件的记忆管理。 - -**挑战与考虑：** - -1. **操作适配性**：Current operations (retrieve/add/update/delete) need adaptation for file-based storage -2. **工具选择**：Consider file operation tools: `grep`, `glob`, `ls`, `read_file`, `write_file`, `edit_file` -3. **模型能力**：Base models have limited file operation capabilities; `qwen3-code` shows better performance - -**潜在架构：** - -```python -# File-based memory operations -FileMemoryOp(tools=[ - grep, # Search within files - glob, # File pattern matching - ls, # List directory contents - read_file, # Read file contents - write_file, # Write new memory files - edit_file, # Update existing memory files -]) -``` - -### 5.6 自我修改上下文 - -支持Agent动态修改自身的上下文状态，实现自适应记忆管理。 - -**实现方式：** - -1. **Summary Agent 主动修改**： - - `add_meta_memory` directly modifies agent context - - Updates available memory types and targets during execution - -2. **ReMy Agent 被动修改**： - - Retrieves `identity_memory` at each interaction - - Dynamically updates self-state based on retrieved identity - - Enables adaptive behavior based on accumulated identity knowledge - -## ReMe V2 开发路线图与实施计划 - -### 技术改造阶段 -1. **代码整合与兼容**：合并flowllm中reme必要的代码，保留现在server-client的依赖，兼容现在各个仓库的依赖代码 -2. **核心接口重构**：新的ReMe接口设计，支持summary，retrieve，context_offload, context_reload 4个核心接口 -3. **Agentic算法升级**：新的agentic算法方案开发 - -### 评估验证阶段 -4. **Benchmark测试** - - halumem - - locomo - - longmemevel - - personal-v2 ? - - appworld/bfcl-v3 - -### 发布推广阶段 -5. **技术报告**撰写与发布 -6. **生态更新**：更新各个仓库的依赖代码 - - agentscope - - agentscope-runtime - - evotraders - - alias(tool-memory) - - agentscope-java - - AgentEvolver - - cookbook: reme procedural memory paper - - tool-memory-upgrade（将要合并） - -**里程碑目标**：春节前完成小版本发布 - ---- - -## ReMe V2 核心竞争优势 - -### 1. 渐进式 Agentic Memory 架构【核心创新】 -融合了多种记忆的渐进式agentic方案，实现从短期到长期记忆的智能化演进 - -### 2. 全生命周期记忆管理 -同时支持长期记忆（Long-term Memory）和短期记忆（Working Memory），完整覆盖Agent认知周期 - -### 3. 模型 -提供开源小模型 - -### 4. 开发者友好生态 - 1. **简洁接口**：提供简洁的接口设计，全异步接口 - 2. **即开即用**：提供CLI工具，开箱即用的体验 - 3. **生态融合**：提供和AgentScope、LangChain无缝集成的方案 - 4. **高度可扩展**：支持Agentic算法的二次开发与定制 \ No newline at end of file diff --git a/docs/todo.md b/docs/todo.md deleted file mode 100644 index 5d73e04d..00000000 --- a/docs/todo.md +++ /dev/null @@ -1,3 +0,0 @@ -1. 如何更好的注册class -2. op的返回，使用return 还是 self.output -3. 如何把agent的东西放出来 \ No newline at end of file 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/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 2e01cc41..766f9b3e 100644 --- a/reme/memory/file_based/__init__.py +++ b/reme/memory/file_based/__init__.py @@ -9,18 +9,21 @@ Components: - Summarizer: Generates memory summaries using LLM - Compactor: Compacts memory content to reduce token usage - ToolResultCompactor: Truncates large tool results and saves full content to files + - ContextChecker: Checks context size and splits messages for compaction """ from .as_msg_handler import AsMsgHandler -from .reme_in_memory_memory import ReMeInMemoryMemory from .component.compactor import Compactor +from .component.context_checker import ContextChecker from .component.summarizer import Summarizer from .component.tool_result_compactor import ToolResultCompactor +from .reme_in_memory_memory import ReMeInMemoryMemory __all__ = [ "AsMsgHandler", "ReMeInMemoryMemory", "Summarizer", "Compactor", + "ContextChecker", "ToolResultCompactor", ] diff --git a/reme/memory/file_based/component/context_checker.py b/reme/memory/file_based/component/context_checker.py new file mode 100644 index 00000000..82bb381e --- /dev/null +++ b/reme/memory/file_based/component/context_checker.py @@ -0,0 +1,97 @@ +"""ContextChecker module for checking context size and splitting messages.""" + +from agentscope.message import Msg +from agentscope.token import HuggingFaceTokenCounter + +from ..as_msg_handler import AsMsgHandler +from ....core.op import BaseOp +from ....core.utils import get_std_logger + +logger = get_std_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. + msg_handler (AsMsgHandler): Handler for message processing and token counting. + """ + + def __init__( + self, + memory_compact_threshold: int, + memory_compact_reserve: int = 10000, + token_counter: HuggingFaceTokenCounter | None = None, + **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. + token_counter (HuggingFaceTokenCounter | None): Token counter for + measuring content length. If None, a default counter will be used. + **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 + + self.msg_handler = AsMsgHandler(token_counter=token_counter) + + 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 + + messages_to_compact, messages_to_keep, is_valid = self.msg_handler.context_check( + messages=messages, + memory_compact_threshold=self.memory_compact_threshold, + memory_compact_reserve=self.memory_compact_reserve, + ) + + logger.info( + f"ContextChecker Result: " + f"to_compact={len(messages_to_compact)}, " + f"to_keep={len(messages_to_keep)}, " + f"is_valid={is_valid}", + ) + + return messages_to_compact, messages_to_keep, is_valid 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/tools/file/file_io.py b/reme/memory/file_based/tools/file_io.py similarity index 77% rename from reme/memory/tools/file/file_io.py rename to reme/memory/file_based/tools/file_io.py index 161c3210..58e9a5f7 100644 --- a/reme/memory/tools/file/file_io.py +++ b/reme/memory/file_based/tools/file_io.py @@ -7,6 +7,8 @@ from typing import Optional from agentscope.message import TextBlock from agentscope.tool import ToolResponse +from .utils import DEFAULT_MAX_BYTES, read_file_safe, truncate_output + class FileIO: """File I/O operations with a configurable working directory.""" @@ -78,57 +80,64 @@ class FileIO: ) try: - with open(file_path, "r", encoding="utf-8") as f: - all_lines = f.readlines() + content = read_file_safe(file_path) + all_lines = content.split("\n") + total = len(all_lines) - range_requested = start_line is not None or end_line is not None + # 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 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" + if s > total: return ToolResponse( content=[ TextBlock( type="text", - text=header + content, + 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 (keep head for file reading) + truncated, was_truncated, output_lines, reason = truncate_output(selected_content, keep="head") + + # Build response with truncation hints + if was_truncated: + end_display = s + output_lines - 1 + next_line = end_display + 1 + if reason == "lines": + hint = f"\n\n[Lines {s}-{end_display} of {total}. Use start_line={next_line} to continue.]" + else: + hint = ( + f"\n\n[Lines {s}-{end_display} of {total} ({DEFAULT_MAX_BYTES // 1024}KB limit). " + f"Use start_line={next_line} to continue.]" + ) + text = truncated + hint + elif e < total: + remaining = total - e + text = ( + f"{file_path} (lines {s}-{e} of {total})\n{truncated}\n\n[{remaining} more lines. " + f"Use start_line={e + 1} to continue.]" + ) else: - content = "".join(all_lines) - return ToolResponse( - content=[ - TextBlock( - type="text", - text=content, - ), - ], - ) + text = truncated + + return ToolResponse( + content=[TextBlock(type="text", text=text)], + ) except Exception as e: return ToolResponse( diff --git a/reme/memory/tools/chunk/memory_get.py b/reme/memory/file_based/tools/memory_get.py similarity index 100% rename from reme/memory/tools/chunk/memory_get.py rename to reme/memory/file_based/tools/memory_get.py diff --git a/reme/memory/tools/chunk/memory_search.py b/reme/memory/file_based/tools/memory_search.py similarity index 100% rename from reme/memory/tools/chunk/memory_search.py rename to reme/memory/file_based/tools/memory_search.py diff --git a/reme/memory/file_based/tools/shell.py b/reme/memory/file_based/tools/shell.py new file mode 100644 index 00000000..c2714b87 --- /dev/null +++ b/reme/memory/file_based/tools/shell.py @@ -0,0 +1,229 @@ +# -*- 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 + +from .utils import truncate_shell_output + + +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 , and + tags. + + Args: + command (`str`): + The shell command to execute. + timeout (`int`, defaults to `60`): + The maximum time (in seconds) allowed for the command to run. + Default is 60 seconds. + + Returns: + `ToolResponse`: + The tool response 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. + """ + + cmd = (command or "").strip() + + # Set working directory + working_dir = self.working_dir + + try: + if sys.platform == "win32": + # Windows: use thread pool to avoid asyncio subprocess limitations + returncode, stdout_str, stderr_str = await asyncio.to_thread( + _execute_subprocess_sync, + cmd, + str(working_dir), + timeout, + ) + else: + proc = await asyncio.create_subprocess_shell( + cmd, + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + bufsize=0, + cwd=str(working_dir), + ) + + try: + # Apply timeout to communicate directly; wait()+communicate() + # can hang if descendants keep stdout/stderr pipes open. + stdout, stderr = await asyncio.wait_for( + proc.communicate(), + timeout=timeout, + ) + encoding = locale.getpreferredencoding(False) or "utf-8" + stdout_str = stdout.decode(encoding, errors="replace").strip( + "\n", + ) + stderr_str = stderr.decode(encoding, errors="replace").strip( + "\n", + ) + returncode = proc.returncode + + except asyncio.TimeoutError: + # Handle timeout + stderr_suffix = ( + f"⚠️ TimeoutError: The command execution exceeded " + f"the timeout of {timeout} seconds. " + f"Please consider increasing the timeout value if this command " + f"requires more time to complete." + ) + returncode = -1 + try: + proc.terminate() + # Wait a bit for graceful termination + try: + await asyncio.wait_for(proc.wait(), timeout=1) + except asyncio.TimeoutError: + # Force kill if graceful termination fails + proc.kill() + await proc.wait() + + # Avoid hanging forever while draining pipes after timeout. + try: + stdout, stderr = await asyncio.wait_for( + proc.communicate(), + timeout=1, + ) + except asyncio.TimeoutError: + stdout, stderr = b"", b"" + encoding = locale.getpreferredencoding(False) or "utf-8" + stdout_str = stdout.decode( + encoding, + errors="replace", + ).strip( + "\n", + ) + stderr_str = stderr.decode( + encoding, + errors="replace", + ).strip( + "\n", + ) + if stderr_str: + stderr_str += f"\n{stderr_suffix}" + else: + stderr_str = stderr_suffix + except ProcessLookupError: + stdout_str = "" + stderr_str = stderr_suffix + + # Apply output truncation + stdout_str = truncate_shell_output(stdout_str) + stderr_str = truncate_shell_output(stderr_str) + + # Format the response in a human-friendly way + if returncode == 0: + # Success case: just show the output + if stdout_str: + response_text = stdout_str + else: + response_text = "Command executed successfully (no output)." + else: + # Error case: show detailed information + response_parts = [f"Command failed with exit code {returncode}."] + if stdout_str: + response_parts.append(f"\n[stdout]\n{stdout_str}") + if stderr_str: + response_parts.append(f"\n[stderr]\n{stderr_str}") + response_text = "".join(response_parts) + + return ToolResponse( + content=[ + TextBlock( + type="text", + text=response_text, + ), + ], + ) + + except Exception as e: + return ToolResponse( + content=[ + TextBlock( + type="text", + text=f"Error: Shell command execution failed due to \n{e}", + ), + ], + ) diff --git a/reme/memory/file_based/tools/utils.py b/reme/memory/file_based/tools/utils.py new file mode 100644 index 00000000..17f58877 --- /dev/null +++ b/reme/memory/file_based/tools/utils.py @@ -0,0 +1,112 @@ +"""Shared utilities for file and shell tools.""" + +# Default truncation limits +DEFAULT_MAX_LINES = 1000 +DEFAULT_MAX_BYTES = 30 * 1024 # 30KB + + +def truncate_output( + text: str, + max_lines: int = DEFAULT_MAX_LINES, + max_bytes: int = DEFAULT_MAX_BYTES, + keep: str = "head", +) -> tuple[str, bool, int, str]: + """Smart truncation for large content. + + Args: + text: Text content to truncate. + max_lines: Maximum number of lines. + max_bytes: Maximum size in bytes. + keep: Which part to keep - "head" (first lines) or "tail" (last lines). + + Returns: + (truncated_content, was_truncated, output_line_count, truncate_reason) + """ + if not text: + return text, False, 0, "" + + lines = text.split("\n") + total_lines = len(lines) + + # No truncation needed + if total_lines <= max_lines and len(text.encode("utf-8")) <= max_bytes: + return text, False, total_lines, "" + + # Apply line limit + if total_lines > max_lines: + if keep == "tail": + lines = lines[-max_lines:] + else: + lines = lines[:max_lines] + reason = "lines" + else: + reason = "" + + # Apply byte limit + if len("\n".join(lines).encode("utf-8")) > max_bytes: + if keep == "tail": + while lines and len("\n".join(lines).encode("utf-8")) > max_bytes: + lines.pop(0) + else: + truncated = [] + current_bytes = 0 + for line in lines: + line_bytes = len(line.encode("utf-8")) + 1 + if current_bytes + line_bytes > max_bytes: + break + truncated.append(line) + current_bytes += line_bytes + lines = truncated + reason = "bytes" + + return "\n".join(lines), True, len(lines), reason + + +def truncate_shell_output(text: str) -> str: + """Truncate shell output to last N lines or M bytes, with truncation notice. + + Args: + text: The output text to truncate. + + Returns: + Truncated text with notice if truncated. + """ + if not text: + return text + + try: + total_lines = len(text.split("\n")) + truncated, was_truncated, output_lines, reason = truncate_output(text, keep="tail") + + if not was_truncated: + return text + + start_line = total_lines - output_lines + 1 + if reason == "lines": + notice = f"\n\n[Output truncated: showing lines {start_line}-{total_lines} of {total_lines} total]" + else: + notice = ( + f"\n\n[Output truncated: showing lines {start_line}-{total_lines} of {total_lines} " + f"({DEFAULT_MAX_BYTES // 1024}KB limit)]" + ) + + return truncated + notice + except Exception: + return text + + +def read_file_safe(file_path: str) -> str: + """Read file with Unicode error handling. + + Args: + file_path: Path to the file. + + Returns: + File content as string. + """ + try: + with open(file_path, "r", encoding="utf-8") as f: + return f.read() + except UnicodeDecodeError: + with open(file_path, "r", encoding="utf-8", errors="ignore") as f: + return f.read() diff --git a/reme/memory/tools/file/__init__.py b/reme/memory/tools/file/__init__.py deleted file mode 100644 index 8234e60d..00000000 --- a/reme/memory/tools/file/__init__.py +++ /dev/null @@ -1,7 +0,0 @@ -"""File-based memory tool implementations.""" - -from .file_io import FileIO - -__all__ = [ - "FileIO", -] 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..685cdcdd 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, @@ -185,7 +185,7 @@ 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, diff --git a/reme/reme_light.py b/reme/reme_light.py index 5266ae76..922b846b 100644 --- a/reme/reme_light.py +++ b/reme/reme_light.py @@ -26,15 +26,45 @@ from agentscope.tool import Toolkit, ToolResponse from .config import ReMeConfigParser from .core import Application from .core.utils import get_hf_token_counter, get_std_logger -from .memory.file_based import Compactor, Summarizer, ToolResultCompactor, ReMeInMemoryMemory, AsMsgHandler -from .memory.tools import MemorySearch -from .memory.tools.file import FileIO +from .memory.file_based import ( + Compactor, + ContextChecker, + Summarizer, + ToolResultCompactor, + ReMeInMemoryMemory, + AsMsgHandler, +) +from .memory.file_based import MemorySearch +from .memory.file_based.tools import FileIO logger = get_std_logger() class ReMeLight(Application): - """ReMe Light Application Class""" + """ + ReMe Light Application Class. + + 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. + memory_path (Path): Path to the memory storage directory. + tool_result_path (Path): Path to the tool result storage directory. + vector_weight (float): Weight for vector search in hybrid search (0-1). + candidate_multiplier (float): Multiplier for candidate retrieval count. + tool_result_threshold (int): Character threshold for tool result compaction. + retention_days (int): Number of days to retain tool result files. + summary_tasks (list[asyncio.Task]): List of active background summary tasks. + """ def __init__( self, @@ -51,6 +81,47 @@ class ReMeLight(Application): tool_result_threshold: int = 1000, retention_days: int = 7, ): + """ + 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. + 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. + tool_result_threshold (int): Character count threshold for tool result + compaction. Results exceeding this length will be truncated and + saved to files. Default 1000 characters. + retention_days (int): Number of days to retain tool result files + before automatic cleanup. Default 7 days. + + 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 + """ # Initialize working directory structure self.working_path = Path(working_dir).absolute() self.working_path.mkdir(parents=True, exist_ok=True) @@ -129,18 +200,64 @@ class ReMeLight(Application): return 0 async def start(self): - """Start the application lifecycle.""" + """ + Start the application lifecycle. + + 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 Application.start() method. + + Note: + This method should be called before using any other application + functionality. It ensures all services are properly initialized. + """ result = await super().start() + # Perform initial cleanup of any expired tool result files self._cleanup_tool_results() return result async def close(self) -> bool: - """Close the application and perform cleanup.""" + """ + Close the application and perform cleanup. + + 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 the application was closed successfully, False otherwise. + + Note: + This method should be called when the application is no longer needed + to ensure proper resource cleanup and data persistence. + """ + # Final cleanup of expired tool result files before shutdown self._cleanup_tool_results() return await super().close() async def compact_tool_result(self, messages: list[Msg]) -> list[Msg]: - """Compact tool results by truncating large outputs and saving full content to files.""" + """ + Compact tool results by truncating large outputs and saving full content to files. + + 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 potentially containing tool results + that may need compaction. + + Returns: + list[Msg]: The processed list of messages with large tool results compacted. + If an error occurs, returns the original unmodified messages. + + Note: + - Tool results shorter than tool_result_threshold are left unchanged + - 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( @@ -162,6 +279,61 @@ class ReMeLight(Application): logger.exception(f"Error compacting tool results: {e}") return messages + async def check_context( + self, + messages: list[Msg], + memory_compact_threshold: int, + memory_compact_reserve: int = 10000, + token_counter: HuggingFaceTokenCounter | None = None, + ) -> 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. + token_counter (HuggingFaceTokenCounter | None): Token counter for + measuring message length. If None, uses default HuggingFace counter. + + 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: + if token_counter is None: + token_counter = get_hf_token_counter() + + checker = ContextChecker( + memory_compact_threshold=memory_compact_threshold, + memory_compact_reserve=memory_compact_reserve, + token_counter=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], @@ -173,7 +345,34 @@ class ReMeLight(Application): compact_ratio: float = 0.7, previous_summary: str = "", ) -> str: - """Compact a list of messages into a condensed summary.""" + """ + Compact a list of messages into a condensed summary. + + 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]): 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". + token_counter (HuggingFaceTokenCounter | None): Token counter for + measuring message length. If None, uses default HuggingFace counter. + 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. + + Returns: + str: The condensed summary of the messages, or an empty string if + an error occurred during compaction. + """ try: if token_counter is None: token_counter = get_hf_token_counter() @@ -208,7 +407,37 @@ class ReMeLight(Application): max_input_length: float = 128 * 1024, compact_ratio: float = 0.7, ) -> str: - """Generate a comprehensive summary of the given messages.""" + """ + Generate a comprehensive summary of the given messages. + + 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]): 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". + token_counter (HuggingFaceTokenCounter | None): Token counter for + measuring message length. If None, uses default HuggingFace counter. + 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. + + Returns: + str: The generated summary text, or an empty string if an error occurred. + + Note: + This method may write summary files to the memory_path directory + using the provided or default toolkit. + """ try: if token_counter is None: token_counter = get_hf_token_counter() @@ -238,7 +467,30 @@ class ReMeLight(Application): return "" def add_async_summary_task(self, messages: list[Msg], **kwargs): - """Add an asynchronous summary task for the given messages.""" + """ + 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 + - 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/cancelled 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(): @@ -274,7 +526,49 @@ class ReMeLight(Application): enable_tool_result_compact: bool = True, tool_result_compact_keep_n: int = 3, ) -> tuple[list[Msg], str]: - """Hook called before reasoning.""" + """ + 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". + token_counter (HuggingFaceTokenCounter | None): Token counter for + measuring content length. If None, uses default counter. + 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 (may be 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 + """ if token_counter is None: token_counter = get_hf_token_counter() @@ -328,7 +622,25 @@ class ReMeLight(Application): return messages_to_keep, compressed_summary async def await_summary_tasks(self) -> str: - """Wait for all background summary tasks to complete and collect results.""" + """ + Wait for all background summary tasks to complete and collect results. + + Blocks until all pending summary tasks in the task list have completed, + cancelled, or failed. Collects status information from each task and + clears the task list after processing. + + Returns: + str: A concatenated string of status messages for all tasks, including: + - Completion confirmations with results + - Cancellation notices + - Error messages for failed tasks + + Note: + - 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: if task.done(): @@ -369,26 +681,19 @@ class ReMeLight(Application): 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: @@ -452,7 +757,25 @@ class ReMeLight(Application): @staticmethod def get_in_memory_memory(token_counter: HuggingFaceTokenCounter | None = None): - """Create and return an in-memory memory instance.""" + """ + Create and return an in-memory memory instance. + + Factory method to create a ReMeInMemoryMemory instance configured with + the specified token counter. This memory instance stores data in RAM + without persistence, suitable for temporary or session-based storage. + + Args: + token_counter (HuggingFaceTokenCounter | None): Token counter for + measuring content length in the memory. If None, creates a + default HuggingFace token counter. + + Returns: + ReMeInMemoryMemory: A new in-memory memory instance ready for use. + + Example: + >>> memory = ReMeLight.get_in_memory_memory() + >>> # Use memory for temporary storage during a session + """ if token_counter is None: token_counter = get_hf_token_counter() 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/test/cli/__init__.py b/test/test/cli/__init__.py similarity index 100% rename from test/cli/__init__.py rename to test/test/cli/__init__.py diff --git a/test/cli/fb_cli.py b/test/test/cli/fb_cli.py similarity index 100% rename from test/cli/fb_cli.py rename to test/test/cli/fb_cli.py diff --git a/test/cli/fb_cli.yaml b/test/test/cli/fb_cli.yaml similarity index 100% rename from test/cli/fb_cli.yaml rename to test/test/cli/fb_cli.yaml diff --git a/test/cli/fb_compactor.py b/test/test/cli/fb_compactor.py similarity index 100% rename from test/cli/fb_compactor.py rename to test/test/cli/fb_compactor.py diff --git a/test/cli/fb_compactor.yaml b/test/test/cli/fb_compactor.yaml similarity index 100% rename from test/cli/fb_compactor.yaml rename to test/test/cli/fb_compactor.yaml diff --git a/test/cli/fb_context_checker.py b/test/test/cli/fb_context_checker.py similarity index 100% rename from test/cli/fb_context_checker.py rename to test/test/cli/fb_context_checker.py diff --git a/test/cli/fb_summarizer.py b/test/test/cli/fb_summarizer.py similarity index 100% rename from test/cli/fb_summarizer.py rename to test/test/cli/fb_summarizer.py diff --git a/test/cli/fb_summarizer.yaml b/test/test/cli/fb_summarizer.yaml similarity index 100% rename from test/cli/fb_summarizer.yaml rename to test/test/cli/fb_summarizer.yaml diff --git a/test/reme_cli.py b/test/test/reme_cli.py similarity index 100% rename from test/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 100% rename from test/test_fs_file_watch_integration.py rename to test/test/test_fs_file_watch_integration.py 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_summarizer.py b/tests/light/test_summarizer.py index a2f2d975..9718cd2a 100644 --- a/tests/light/test_summarizer.py +++ b/tests/light/test_summarizer.py @@ -15,7 +15,7 @@ from test_utils import ( ) from reme.core.utils import get_std_logger from reme.memory.file_based import Summarizer -from reme.memory.tools.file import FileIO +from reme.memory.file_based.tools import FileIO logger = get_std_logger() diff --git a/tests/light/test_tools.py b/tests/light/test_tools.py new file mode 100644 index 00000000..933d506d --- /dev/null +++ b/tests/light/test_tools.py @@ -0,0 +1,321 @@ +# -*- coding: utf-8 -*- +# pylint: disable=redefined-outer-name +"""Unit tests for Shell and FileIO tools.""" + +import asyncio +import os +import shutil +import tempfile + +import pytest + +from reme.memory.file_based.tools.shell import Shell +from reme.memory.file_based.tools.file_io import FileIO +from reme.memory.file_based.tools.utils import DEFAULT_MAX_LINES, 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_truncated_output(shell_env): + """Test output truncation for large output.""" + lines_to_generate = DEFAULT_MAX_LINES + 500 + cmd = f"seq 1 {lines_to_generate}" + result = asyncio.run(shell_env["shell"].execute_shell_command(cmd)) + text = result.content[0].get("text", "") + + # Should contain truncation notice + assert "truncated" in text.lower() + # Should contain the last line (tail is kept) + assert str(lines_to_generate) in text + # Verify first numeric line is > 1 (truncated from head) + numeric_lines = [tl for tl in text.strip().split("\n") if tl.isdigit()] + if numeric_lines: + assert int(numeric_lines[0]) > 1 + + +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_LINES) + large_file = os.path.join(test_dir, "large.txt") + with open(large_file, "w", encoding="utf-8") as f: + for i in range(1, DEFAULT_MAX_LINES + 500): + f.write(f"line {i}\n") + + # Create large bytes file (exceeds DEFAULT_MAX_BYTES) + 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) + + yield { + "dir": test_dir, + "file_io": file_io, + "simple_file": simple_file, + "large_file": large_file, + "large_bytes_file": large_bytes_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(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("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("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(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(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(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_by_lines(fileio_env): + """Test file truncation by line limit.""" + result = asyncio.run(fileio_env["file_io"].read(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(fileio_env["large_bytes_file"])) + text = result.content[0].get("text", "") + assert "continue" in text.lower() or "KB" in text + + +def test_read_directory_error(fileio_env): + """Test reading a directory returns error.""" + result = asyncio.run(fileio_env["file_io"].read(fileio_env["dir"])) + text = result.content[0].get("text", "") + assert "Error" in text + assert "not a file" 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_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_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("", "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("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(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(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("nonexistent.txt", "old", "new"), + ) + text = result.content[0].get("text", "") + assert "Error" in text 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()) From 5e08aa48b81b586ccfd2f6edd5a4a9953e6cbc3a Mon Sep 17 00:00:00 2001 From: jinliyl <6469360+jinliyl@users.noreply.github.com> Date: Sat, 7 Mar 2026 15:22:56 +0800 Subject: [PATCH 17/59] refactor(memory): restructure file-based memory components and enhance message handling (#145) --- README.md | 35 +++++++++---------- README_ZH.md | 32 ++++++++--------- reme/core/schema/as_msg_stat.py | 35 ++++++++++++------- reme/memory/file_based/__init__.py | 30 ++++------------ reme/memory/file_based/component/__init__.py | 0 reme/memory/file_based/components/__init__.py | 13 +++++++ .../{component => components}/compactor.py | 5 ++- .../{component => components}/compactor.yaml | 0 .../context_checker.py | 15 ++++---- .../{component => components}/summarizer.py | 5 ++- .../{component => components}/summarizer.yaml | 0 .../tool_result_compactor.py | 0 .../file_based/reme_in_memory_memory.py | 2 +- reme/memory/file_based/utils/__init__.py | 7 ++++ .../file_based/{ => utils}/as_msg_handler.py | 33 ++++++++++++++--- reme/reme_light.py | 18 +++++----- tests/light/test_compactor.py | 4 +-- tests/light/test_context_check.py | 4 +-- tests/light/test_format_msgs_to_str.py | 4 +-- tests/light/test_summarizer.py | 4 +-- tests/light/test_tool_result_compactor.py | 3 +- tests/light/test_tools.py | 2 +- tests/light/test_utils.py | 2 +- 23 files changed, 148 insertions(+), 105 deletions(-) delete mode 100644 reme/memory/file_based/component/__init__.py create mode 100644 reme/memory/file_based/components/__init__.py rename reme/memory/file_based/{component => components}/compactor.py (89%) rename reme/memory/file_based/{component => components}/compactor.yaml (100%) rename reme/memory/file_based/{component => components}/context_checker.py (92%) rename reme/memory/file_based/{component => components}/summarizer.py (88%) rename reme/memory/file_based/{component => components}/summarizer.yaml (100%) rename reme/memory/file_based/{component => components}/tool_result_compactor.py (100%) create mode 100644 reme/memory/file_based/utils/__init__.py rename reme/memory/file_based/{ => utils}/as_msg_handler.py (92%) diff --git a/README.md b/README.md index 9d18e6ce..7bc1ef86 100644 --- a/README.md +++ b/README.md @@ -20,7 +20,7 @@ A memory management toolkit for AI agents — Remember Me, Refine Me.

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

@@ -33,6 +34,24 @@ conversations) and **stateless sessions** (new sessions cannot inherit history a ReMe gives agents **real memory** — old conversations are automatically compacted, important information is persistently stored, and relevant context is automatically recalled in future interactions. +

+What you can do with ReMe + +
+ +- **Personal assistant**: Provide long-term memory for agents like [CoPaw](https://github.com/agentscope-ai/CoPaw), + remembering user preferences and conversation history. +- **Coding assistant**: Record code style preferences and project context, maintaining a consistent development + experience across sessions. +- **Customer service bot**: Track user issue history and preference settings for personalized service. +- **Task automation**: Learn success/failure patterns from historical tasks to continuously optimize execution + strategies. +- **Knowledge Q&A**: Build a searchable knowledge base with semantic search and exact matching support. +- **Multi-turn dialogue**: Automatically compress long conversations while retaining key information within limited + context windows. + +
+ --- ## 📁 File-based memory system (ReMeLight) @@ -82,7 +101,18 @@ capabilities for AI agents: #### Installation +**Install from source:** + ```bash +git clone https://github.com/agentscope-ai/ReMe.git +cd ReMe +pip install -e ".[light]" +``` + +**Update to the latest version:** + +```bash +git pull pip install -e ".[light]" ``` @@ -546,6 +576,14 @@ For more details on how to reproduce the experiments, see [quickstart.md](benchm - **Acknowledgements**: We thank excellent open-source projects such as OpenClaw, Mem0, MemU, and CoPaw for their inspiration and support. +### Contributors + +Thanks to all who have contributed to ReMe: + + + Contributors + + --- ## 📄 Citation @@ -567,6 +605,14 @@ This project is open-sourced under the Apache License 2.0. See [LICENSE](./LICEN --- +## 🤔 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 [![Star History Chart](https://api.star-history.com/svg?repos=agentscope-ai/ReMe&type=Date)](https://www.star-history.com/#agentscope-ai/ReMe&Date) diff --git a/README_ZH.md b/README_ZH.md index bd038f67..1cc60ed4 100644 --- a/README_ZH.md +++ b/README_ZH.md @@ -14,6 +14,7 @@ English 简体中文 GitHub Stars + DeepWiki

@@ -30,6 +31,19 @@ ReMe 让智能体拥有**真正的记忆力**——旧对话自动浓缩，重要信息持久保存，下次对话自动想起来。 +

+你可以用 ReMe 做什么 + +
+ +- **个人助理**：为 [CoPaw](https://github.com/agentscope-ai/CoPaw) 等智能体提供长期记忆，记住用户偏好和历史对话。 +- **编程助手**：记录代码风格偏好、项目上下文，跨会话保持一致的开发体验。 +- **客服机器人**：记录用户问题历史、偏好设置，提供个性化服务。 +- **任务自动化**：从历史任务中学习成功/失败模式，持续优化执行策略。 +- **知识问答**：构建可检索的知识库，支持语义搜索和精确匹配。 +- **多轮对话**：自动压缩长对话，在有限上下文窗口内保留关键信息。 + +
--- @@ -78,7 +92,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]" ``` @@ -527,6 +552,14 @@ Pass@K 衡量在生成 K 个候选中，至少一个成功完成任务（score=1 - **代码贡献**：欢迎任何形式的代码贡献，请参阅 [贡献指南](docs/contribution.md)。 - **致谢**：感谢 OpenClaw、Mem0、MemU、CoPaw 等优秀的开源项目，为项目带来诸多启发与帮助。 +### 贡献者 + +感谢所有为 ReMe 做出贡献的朋友们： + + + 贡献者 + + --- ## 📄 引用 @@ -548,6 +581,13 @@ Pass@K 衡量在生成 K 个候选中，至少一个成功完成任务（score=1 --- +## 🤔 为什么叫 ReMe？ + +ReMe 是 **Remember Me** 和 **Refine Me** 的缩写，寓意让 AI 智能体「记住我」并在交互中「精进自我」。我们希望 ReMe +不只是一个冷冰冰的记忆模块，而是能让智能体真正理解用户、积累经验、持续进化的伙伴。 + +--- + ## 📈 Star 历史 [![Star History Chart](https://api.star-history.com/svg?repos=agentscope-ai/ReMe&type=Date)](https://www.star-history.com/#agentscope-ai/ReMe&Date) From f408d6ec4a6141aeddd3a76f943bcec83714a503 Mon Sep 17 00:00:00 2001 From: "jinli.yl" Date: Sat, 7 Mar 2026 15:57:00 +0800 Subject: [PATCH 20/59] docs(readme): update chinese badge text to simplified chinese --- README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/README.md b/README.md index 2a85f0ad..7324675e 100644 --- a/README.md +++ b/README.md @@ -12,7 +12,7 @@

License English - Simplified Chinese + 简体中文 GitHub Stars DeepWiki

From 8b45493634001a61f09b248449463bd56d398117 Mon Sep 17 00:00:00 2001 From: hyp-001 <66149011+hyp-001@users.noreply.github.com> Date: Mon, 9 Mar 2026 16:09:06 +0800 Subject: [PATCH 21/59] =?UTF-8?q?=E5=A2=9E=E5=8A=A0locomo=E7=9A=84?= =?UTF-8?q?=E4=BB=A3=E7=A0=81=20(#148)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-authored-by: huwen.hyp --- benchmark/locomo/eval_reme.py | 1107 +++++++++++++++++++++++++++++++ benchmark/locomo/eval_reme.yaml | 180 +++++ 2 files changed, 1287 insertions(+) create mode 100644 benchmark/locomo/eval_reme.py create mode 100644 benchmark/locomo/eval_reme.yaml 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] + From 083ed6a137d48e5e07b690d3a51c98a51237efa1 Mon Sep 17 00:00:00 2001 From: jinliyl <6469360+jinliyl@users.noreply.github.com> Date: Mon, 9 Mar 2026 21:05:30 +0800 Subject: [PATCH 22/59] update readme (#149) --- README.md | 22 +++++++++++----------- README_ZH.md | 22 +++++++++++----------- 2 files changed, 22 insertions(+), 22 deletions(-) diff --git a/README.md b/README.md index 7324675e..ca7726bd 100644 --- a/README.md +++ b/README.md @@ -83,17 +83,17 @@ working_dir/ [ReMeLight](reme/reme_light.py) is the core class of the file-based memory system. It provides full memory management capabilities for AI agents: -| Method | Function | Key components | -|-----------------------|--------------------------------------|---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------| -| `check_context` | 📊 Check context size | [ContextChecker](reme/memory/file_based/components/context_checker.py) — checks whether context exceeds thresholds and splits messages | -| `compact_memory` | 📦 Compact history into summary | [Compactor](reme/memory/file_based/components/compactor.py) — ReActAgent that generates structured context summaries | -| `summary_memory` | 📝 Persist important memory to files | [Summarizer](reme/memory/file_based/components/summarizer.py) — ReActAgent + file tools (`read` / `write` / `edit`) | -| `compact_tool_result` | ✂️ Compact long tool outputs | [ToolResultCompactor](reme/memory/file_based/components/tool_result_compactor.py) — truncates long tool outputs and stores them in `tool_result/` while keeping file references in messages | -| `memory_search` | 🔍 Semantic memory search | [MemorySearch](reme/memory/file_based/tools/memory_search.py) — hybrid retrieval with vectors + BM25 | -| `ReMeInMemoryMemory` | 🗂️ In-session memory class | [ReMeInMemoryMemory](reme/memory/file_based/reme_in_memory_memory.py) — token-aware memory management with summary compression and state serialization | -| `pre_reasoning_hook` | 🔄 Pre-reasoning hook | `compact_tool_result` + `check_context` + `compact_memory` + `summary_memory` (async) | -| `start` | 🚀 Start memory system | Initialize file storage, file watcher, and embedding cache; clean up expired tool result files | -| `close` | 📕 Shutdown and cleanup | Clean up tool result files, stop file watcher, and persist embedding cache | + + + + + + + + + + +
CategoryMethodFunctionKey components
Context Managementcheck_context📊 Check context sizeContextChecker — checks whether context exceeds thresholds and splits messages
compact_memory📦 Compact history into summaryCompactor — ReActAgent that generates structured context summaries
compact_tool_result✂️ Compact long tool outputsToolResultCompactor — truncates long tool outputs and stores them in tool_result/ while keeping file references in messages
pre_reasoning_hook🔄 Pre-reasoning hookcompact_tool_result + check_context + compact_memory + summary_memory (async)
Long-term Memorysummary_memory📝 Persist important memory to filesSummarizer — ReActAgent + file tools (read / write / edit)
memory_search🔍 Semantic memory searchMemorySearch — hybrid retrieval with vectors + BM25
-start🚀 Start memory systemInitialize file storage, file watcher, and embedding cache; clean up expired tool result files
-close📕 Shutdown and cleanupClean up tool result files, stop file watcher, and persist embedding cache
--- diff --git a/README_ZH.md b/README_ZH.md index 1cc60ed4..c3de853f 100644 --- a/README_ZH.md +++ b/README_ZH.md @@ -74,17 +74,17 @@ working_dir/ [ReMeLight](reme/reme_light.py) 是该记忆系统的核心类，为 AI Agent 提供完整的记忆管理能力： -| 方法 | 功能 | 关键组件 | -|-----------------------|--------------|------------------------------------------------------------------------------------------------------------------------------| -| `check_context` | 📊 检查上下文大小 | [ContextChecker](reme/memory/file_based/components/context_checker.py) — 检查上下文是否超出阈值并拆分Message | -| `compact_memory` | 📦 压缩历史对话为摘要 | [Compactor](reme/memory/file_based/components/compactor.py) — ReActAgent 生成结构化上下文摘要 | -| `summary_memory` | 📝 将重要记忆写入文件 | [Summarizer](reme/memory/file_based/components/summarizer.py) — ReActAgent + 文件工具（read / write / edit） | -| `compact_tool_result` | ✂️ 压缩超长工具输出 | [ToolResultCompactor](reme/memory/file_based/components/tool_result_compactor.py) — 截断超长的工具调用结果并转存到 `tool_result/`，消息中保留文件引用 | -| `memory_search` | 🔍 语义搜索记忆 | [MemorySearch](reme/memory/file_based/tools/memory_search.py) — 向量 + BM25 混合检索 | -| `ReMeInMemoryMemory` | 🗂️ 会话内存类 | [ReMeInMemoryMemory](reme/memory/file_based/reme_in_memory_memory.py) — Token 感知的内存管理，支持压缩摘要和状态序列化 | -| `pre_reasoning_hook` | 🔄 推理前预处理钩子 | compact_tool_result + check_context + compact_memory + summary_memory(async) | -| `start` | 🚀 启动记忆系统 | 初始化文件存储、文件监控、Embedding 缓存；清理过期工具结果文件 | -| `close` | 📕 关闭并清理 | 清理工具结果文件、停止文件监控、保存 Embedding 缓存 |· + + + + + + + + + + +
类别方法功能关键组件
上下文管理check_context📊 检查上下文大小ContextChecker — 检查上下文是否超出阈值并拆分 Message
compact_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 混合检索
-start🚀 启动记忆系统初始化文件存储、文件监控、Embedding 缓存；清理过期工具结果文件
-close📕 关闭并清理清理工具结果文件、停止文件监控、保存 Embedding 缓存
--- From 8b1698451e96b99c71ffe157f02e1f2e64061d16 Mon Sep 17 00:00:00 2001 From: jinliyl <6469360+jinliyl@users.noreply.github.com> Date: Tue, 10 Mar 2026 18:19:00 +0800 Subject: [PATCH 23/59] fin(emb_dim) (#150) * fix(logger): handle file logging configuration errors gracefully * fix(embedding): handle embedding dimension mismatches and improve caching * fix(file-watcher): clear file store on changes to prevent stale data * refactor(logger): update logger implementation and fix message translation * refactor(logger): update logger configuration and add documentation --- reme/__init__.py | 2 +- reme/core/embedding/base_embedding_model.py | 92 +++++++++++++++---- reme/core/file_store/chroma_file_store.py | 44 +++++++-- reme/core/file_store/local_file_store.py | 24 ++++- reme/core/file_watcher/delta_file_watcher.py | 1 + reme/core/file_watcher/full_file_watcher.py | 2 + reme/core/utils/logger_utils.py | 44 +++++---- reme/core/utils/std_logger.py | 43 +++++---- .../file_based/components/compactor.yaml | 2 +- 9 files changed, 192 insertions(+), 62 deletions(-) diff --git a/reme/__init__.py b/reme/__init__.py index c75dfbcc..41537822 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.6b2" +__version__ = "0.3.0.6b3" __all__ = [ "config", diff --git a/reme/core/embedding/base_embedding_model.py b/reme/core/embedding/base_embedding_model.py index 78a91b8b..2b60e8e6 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, @@ -99,7 +99,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 @@ -107,12 +134,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: @@ -164,6 +192,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( @@ -204,6 +239,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") @@ -223,16 +264,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 @@ -249,9 +301,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) @@ -299,7 +357,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 @@ -345,8 +403,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}") @@ -371,7 +430,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 @@ -417,8 +476,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/file_store/chroma_file_store.py b/reme/core/file_store/chroma_file_store.py index 390a2bc4..d87f5fef 100644 --- a/reme/core/file_store/chroma_file_store.py +++ b/reme/core/file_store/chroma_file_store.py @@ -1,6 +1,7 @@ """ChromaDB storage backend for file store.""" import json +import random import time from pathlib import Path @@ -355,12 +356,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 +460,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_watcher/delta_file_watcher.py b/reme/core/file_watcher/delta_file_watcher.py index 6148bd07..f35c9b2f 100644 --- a/reme/core/file_watcher/delta_file_watcher.py +++ b/reme/core/file_watcher/delta_file_watcher.py @@ -141,6 +141,7 @@ class DeltaFileWatcher(BaseFileWatcher): async def _on_changes(self, changes: set[tuple[Change, str]]): """Handle file changes with incremental synchronization.""" self.dirty = True + await self.file_store.clear_all() for change_type, path in changes: if change_type == Change.added: diff --git a/reme/core/file_watcher/full_file_watcher.py b/reme/core/file_watcher/full_file_watcher.py index 2d365f23..c49a94fa 100644 --- a/reme/core/file_watcher/full_file_watcher.py +++ b/reme/core/file_watcher/full_file_watcher.py @@ -44,6 +44,8 @@ class FullFileWatcher(BaseFileWatcher): async def _on_changes(self, changes: set[tuple[Change, str]]): """Handle file changes with full synchronization""" self.dirty = True + await self.file_store.clear_all() + 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/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 index e2de0e9c..dbdf1908 100644 --- a/reme/core/utils/std_logger.py +++ b/reme/core/utils/std_logger.py @@ -38,7 +38,7 @@ class CustomFormatter(logging.Formatter): return super().format(record) -def get_logger( +def get_loggerv2( name: str = "reme", log_dir: str = "logs", level: str = "INFO", @@ -80,22 +80,26 @@ def get_logger( # Configure file logging if log_to_file: - 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) + 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) + 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: @@ -107,3 +111,10 @@ def get_logger( # 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/memory/file_based/components/compactor.yaml b/reme/memory/file_based/components/compactor.yaml index 2e3b15a8..83c4f952 100644 --- a/reme/memory/file_based/components/compactor.yaml +++ b/reme/memory/file_based/components/compactor.yaml @@ -119,7 +119,7 @@ update_user_message_suffix: | Keep each section concise. Preserve exact file paths, function names, and error messages. update_user_message_prefix_zh: | - 上述消息是要整合到现有摘要中的新对话消息，这些消息在标签中提供。 + 以上消息是需要整合到现有摘要中的新对话内容，现有摘要位于标签中。 update_user_message_suffix_zh: | 用新信息更新现有的结构化摘要。规则： From e5bb8451966a144d3e67bf685629877daac64e8b Mon Sep 17 00:00:00 2001 From: zouyingcao <57442064+zouyingcao@users.noreply.github.com> Date: Thu, 12 Mar 2026 11:23:36 +0800 Subject: [PATCH 24/59] refactor(cli): using AgentScope components to reimplement the reme_cli logic (#153) * add: as_token_counters config for reme_cli * add: reme_cli function * update: format the terminal printing for reme_cli * update: check for pre-commit * single quotes for the inner dictionary keys * update the usage of get_std_logger for pre-commit * update the usage of get_std_logger for pre-commit * add 'console_enabled' param in compactor&summarizer --- reme/config/cli.yaml | 22 +- reme/core/__init__.py | 2 + reme/core/application.py | 7 + reme/core/as_token_counter/__init__.py | 9 + reme/core/op/base_op.py | 10 + reme/core/registry_factory.py | 1 + reme/core/schema/service_config.py | 1 + reme/core/service_context.py | 5 + reme/core/utils/llm_utils.py | 54 +++- reme/core/utils/std_logger.py | 6 +- reme/memory/file_based/components/__init__.py | 2 + reme/memory/file_based/components/cli.py | 303 ++++++++++++++++++ reme/memory/file_based/components/cli.yaml | 95 ++++++ .../memory/file_based/components/compactor.py | 7 +- .../file_based/components/summarizer.py | 7 +- reme/reme_cli.py | 192 +++++++++++ 16 files changed, 703 insertions(+), 20 deletions(-) create mode 100644 reme/core/as_token_counter/__init__.py create mode 100644 reme/memory/file_based/components/cli.py create mode 100644 reme/memory/file_based/components/cli.yaml create mode 100644 reme/reme_cli.py diff --git a/reme/config/cli.yaml b/reme/config/cli.yaml index 7d348d45..a6914170 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: @@ -41,11 +49,3 @@ file_watchers: 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/core/__init__.py b/reme/core/__init__.py index 053755cc..726a3455 100644 --- a/reme/core/__init__.py +++ b/reme/core/__init__.py @@ -2,6 +2,7 @@ from . import as_llm from . import as_llm_formatter +from . import as_token_counter from . import embedding from . import enumeration from . import file_store @@ -25,6 +26,7 @@ __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 46f4a934..85a33c4d 100644 --- a/reme/core/application.py +++ b/reme/core/application.py @@ -172,6 +172,13 @@ class Application: config_dict = config.model_dump(exclude={"backend"}) self.service_context.as_llm_formatters[name] = R.as_llm_formatters[config.backend](**config_dict) + 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: + config_dict = config.model_dump(exclude={"backend"}) + self.service_context.as_token_counters[name] = R.as_token_counters[config.backend](**config_dict) + 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.") diff --git a/reme/core/as_token_counter/__init__.py b/reme/core/as_token_counter/__init__.py new file mode 100644 index 00000000..51b3bd29 --- /dev/null +++ b/reme/core/as_token_counter/__init__.py @@ -0,0 +1,9 @@ +"""Module for registering AgentScope token counters.""" + +from agentscope.token import OpenAITokenCounter +from agentscope.token import HuggingFaceTokenCounter + +from ..registry_factory import R + +R.as_token_counters.register("openai")(OpenAITokenCounter) +R.as_token_counters.register("hf")(HuggingFaceTokenCounter) diff --git a/reme/core/op/base_op.py b/reme/core/op/base_op.py index 86cfa3be..f25da289 100644 --- a/reme/core/op/base_op.py +++ b/reme/core/op/base_op.py @@ -9,6 +9,7 @@ from typing import Callable, Optional, Any from agentscope.formatter import FormatterBase from agentscope.model import ChatModelBase +from agentscope.token import TokenCounterBase from loguru import logger from tqdm import tqdm @@ -46,6 +47,7 @@ class BaseOp(metaclass=ABCMeta): prompt_path: str = "", as_llm: str | ChatModelBase = "default", as_llm_formatter: str | FormatterBase = "default", + as_token_counter: str | TokenCounterBase = "default", llm: str | BaseLLM = "default", embedding_model: str | BaseEmbeddingModel = "default", vector_store: str | BaseVectorStore = "default", @@ -70,6 +72,7 @@ class BaseOp(metaclass=ABCMeta): 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 @@ -149,6 +152,13 @@ class BaseOp(metaclass=ABCMeta): 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) -> TokenCounterBase: + """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.""" diff --git a/reme/core/registry_factory.py b/reme/core/registry_factory.py index f54ad3c1..921f27a8 100644 --- a/reme/core/registry_factory.py +++ b/reme/core/registry_factory.py @@ -36,6 +36,7 @@ class RegistryFactory: 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/service_config.py b/reme/core/schema/service_config.py index 5e4cd212..758944b2 100644 --- a/reme/core/schema/service_config.py +++ b/reme/core/schema/service_config.py @@ -130,6 +130,7 @@ class ServiceConfig(BasicConfig): 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 92d0d566..4ab83ff7 100644 --- a/reme/core/service_context.py +++ b/reme/core/service_context.py @@ -13,6 +13,7 @@ from .utils import load_env, 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 @@ -40,6 +41,7 @@ class ServiceContext(BaseDict): 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, @@ -72,6 +74,8 @@ class ServiceContext(BaseDict): 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: @@ -100,6 +104,7 @@ class ServiceContext(BaseDict): 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"] = {} 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/std_logger.py b/reme/core/utils/std_logger.py index dbdf1908..bfa6b32a 100644 --- a/reme/core/utils/std_logger.py +++ b/reme/core/utils/std_logger.py @@ -47,6 +47,7 @@ def get_loggerv2( log_file_prefix: str = "reme", rotation: str = "midnight", retention_days: int = 7, + force_update: bool = False, ) -> logging.Logger: """Get a configured logger instance. @@ -59,12 +60,13 @@ def get_loggerv2( 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 - if name in _loggers: + # 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 diff --git a/reme/memory/file_based/components/__init__.py b/reme/memory/file_based/components/__init__.py index 42ab2f6b..86574790 100644 --- a/reme/memory/file_based/components/__init__.py +++ b/reme/memory/file_based/components/__init__.py @@ -4,10 +4,12 @@ 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..bc88bd56 --- /dev/null +++ b/reme/memory/file_based/components/cli.py @@ -0,0 +1,303 @@ +"""CLI component for interactive chat using agentscope-based memory tools.""" + +import asyncio +from datetime import datetime +from pathlib import Path + +from agentscope.agent import ReActAgent +from agentscope.message import Msg, TextBlock +from agentscope.tool import Toolkit, ToolResponse +from agentscope.pipeline import stream_printing_messages +from loguru import logger + +from ....core.op import BaseOp +from ....core.utils import format_messages +from .compactor import Compactor +from .context_checker import ContextChecker +from .summarizer import Summarizer +from ..tools import FileIO, MemorySearch + + +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", + **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 + + # 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 + ) + + # 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) + toolkit.register_tool_function(file_io.write) + toolkit.register_tool_function(file_io.edit) + + 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.reserve_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 + ) + + 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.""" + current_time = datetime.now().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) -> str: + """ + 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 + + # 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 index e7b37b30..fbe505ad 100644 --- a/reme/memory/file_based/components/compactor.py +++ b/reme/memory/file_based/components/compactor.py @@ -3,12 +3,10 @@ from agentscope.agent import ReActAgent from agentscope.message import Msg from agentscope.token import HuggingFaceTokenCounter +from loguru import logger from ..utils import AsMsgHandler from ....core.op import BaseOp -from ....core.utils import get_std_logger - -logger = get_std_logger() class Compactor(BaseOp): @@ -18,12 +16,14 @@ class Compactor(BaseOp): self, memory_compact_threshold: int, token_counter: HuggingFaceTokenCounter, + console_enabled: bool = True, **kwargs, ): super().__init__(**kwargs) self.memory_compact_threshold: int = memory_compact_threshold self.msg_handler = AsMsgHandler(token_counter=token_counter) + self.console_enabled: bool = console_enabled async def execute(self): messages: list[Msg] = self.context.get("messages", []) @@ -50,6 +50,7 @@ class Compactor(BaseOp): sys_prompt=self.get_prompt("system_prompt"), formatter=self.as_llm_formatter, ) + agent.set_console_output_enabled(self.console_enabled) if previous_summary: prefix: str = self.get_prompt("update_user_message_prefix") diff --git a/reme/memory/file_based/components/summarizer.py b/reme/memory/file_based/components/summarizer.py index d4e057be..83441149 100644 --- a/reme/memory/file_based/components/summarizer.py +++ b/reme/memory/file_based/components/summarizer.py @@ -6,12 +6,10 @@ from agentscope.agent import ReActAgent from agentscope.message import Msg from agentscope.token import HuggingFaceTokenCounter from agentscope.tool import Toolkit +from loguru import logger from ..utils import AsMsgHandler from ....core.op import BaseOp -from ....core.utils import get_std_logger - -logger = get_std_logger() class Summarizer(BaseOp): @@ -24,6 +22,7 @@ class Summarizer(BaseOp): memory_compact_threshold: int, token_counter: HuggingFaceTokenCounter, toolkit: Toolkit, + console_enabled: bool = True, **kwargs, ): super().__init__(**kwargs) @@ -33,6 +32,7 @@ class Summarizer(BaseOp): self.msg_handler = AsMsgHandler(token_counter=token_counter) self.toolkit: Toolkit = toolkit + self.console_enabled: bool = console_enabled async def execute(self): messages: list[Msg] = self.context.get("messages", []) @@ -59,6 +59,7 @@ class Summarizer(BaseOp): formatter=self.as_llm_formatter, toolkit=self.toolkit, ) + agent.set_console_output_enabled(self.console_enabled) user_message: str = f"\n{history_formatted_str}\n\n" + self.prompt_format( "user_message", 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() From 67ad153a2a350d51db66f9aacfd1a5571e723dc8 Mon Sep 17 00:00:00 2001 From: Zhouwk <57825291+nitwtog@users.noreply.github.com> Date: Mon, 16 Mar 2026 11:58:55 +0800 Subject: [PATCH 25/59] =?UTF-8?q?=E6=B7=BB=E5=8A=A0Reme=E5=9C=A8Halumem?= =?UTF-8?q?=E5=92=8CLocomo=E7=9A=84=E5=AE=9E=E9=AA=8C=E7=BB=93=E6=9E=9C=20?= =?UTF-8?q?(#155)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * feat(reme): 添加配置选项以启用或禁用个人资料功能 - 在 ReMe 初始化方法中添加 enable_profile 参数，默认值为 True - 根据 enable_profile 设置决定是否创建 profile 目录和设置 profile_dir - 在 PersonalSummarizer 中根据 enable_profile 条件性地添加个人资料相关工具 - 在 PersonalRetriever 中根据 enable_profile 条件性地添加 ReadAllProfiles 工具 - 修改 profile_path 属性以在禁用个人资料时返回 None - 修改 get_profile_handler 方法以在禁用个人资料时返回 None - 为 enable_profile 参数添加文档说明其用于云向量存储场景 * refactor(benchmark): 重构LongMemEval基准测试中的ReMe实例管理 - 移除未使用的shutil导入 - 将固定的ReMe实例改为每个问题创建独立实例以实现隔离 - 更新LLM配置名称从qwen3-max-think到qwen-max-t - 修改模型调用逻辑使用正确的model_name参数 - 添加qwen-flash和GPT-4o-mini等新模型配置 - 统一使用"User"作为用户名，通过集合名实现隔离 - 调整并发处理数从4降至1，批处理大小从10增至30 - 每个问题类型采样数从2增至4 - 添加异步上下文管理确保资源正确释放 * reformat 2 files * refactor(benchmark): 重构长记忆评估中的模型配置 - 将原有的 eval_model_name 替换为专门的 retrieve_model_name 用于检索操作 - 添加对 qwen-max 模型配置的支持 - 更新参数解析器以支持新的检索模型参数 - 修改最大并发数默认值从 1 提升到 4 - 调整样本数量默认值从 4 减少到 1 - 统一模型参数命名规范，区分摘要、检索和评估模型 - 优化内存处理器初始化逻辑，支持独立的检索模型配置 * fix(benchmark): 移除数据路径默认值并设为必填参数 - 将LongMemEval评估脚本中的data_path参数改为必需参数 - 将HaluMem评估脚本中的data_path参数改为必需参数 - 删除了硬编码的默认文件路径配置 - 强制用户显式指定数据集文件路径以避免路径错误 * Update __init__.py * Update __init__.py * fix(benchmark): 修复ReMe评估中的模型配置和空值处理问题 - 移除了retrieve_memory调用中不需要的llm_config_name参数 - 修复了长字符串打印的换行格式问题 - 添加了eval_result为空时的初始化处理 - 在accuracy评估中加入了eval_model_name参数传递 * style(benchmark): 格式化模型名称打印输出 - 移除了多行字符串中的换行符和多余空格 - 将模型名称信息合并为单行连续显示 - 保持了原有的打印格式和信息完整性 * docs(readme): 更新文档添加实验结果表格 - 在英文版 README 中添加 🧪 Experiments 章节 - 添加 LoCoMo 和 HaluMem 两个基准测试的结果表格 - 在中文版 README_ZH 中添加 🧪 实验 章节 - 添加 LoCoMo 和 HaluMem 测试集的实验配置说明 - 添加完整的实验数据对比表格和评估协议说明 * docs(readme): 更新文档中的内存系统链接 - 为基于文件的记忆系统添加锚点链接 - 为基于向量库的记忆系统添加锚点链接 - 修复英文文档中的链接格式 - 修复中文文档中的链接格式和空行问题 --- README.md | 41 +++++++++++++++++++++++++++++++++++++++-- README_ZH.md | 41 ++++++++++++++++++++++++++++++++++++++++- 2 files changed, 79 insertions(+), 3 deletions(-) diff --git a/README.md b/README.md index ca7726bd..f379652f 100644 --- a/README.md +++ b/README.md @@ -25,8 +25,7 @@ --- -🧠 ReMe is a memory management framework designed for **AI agents**, providing both file-based and vector-based memory -systems. +🧠 ReMe is a memory management framework designed for **AI agents**, providing both [file-based](#-file-based-memory-system-remelight) and [vector-based](#-vector-based-memory-system) memory systems. It tackles two core problems of agent memory: **limited context window** (early information is truncated or lost in long conversations) and **stateless sessions** (new sessions cannot inherit history and always start from scratch). @@ -416,6 +415,44 @@ memories: Installation and environment configuration are the same as [ReMeLight](#installation). API keys are configured via environment variables and can be stored in a `.env` file at the project root. +## 🧪 Experiments + +Evaluations are conducted on three benchmarks: **LoCoMo** and **HaluMem**. Experimental settings: + +1. **ReMe backbone**: as specified in each table. +2. **Evaluation protocol**: LLM-as-a-Judge following MemOS — each answer is scored by GPT-4o-mini. + +Baseline results are reproduced from their respective papers under aligned settings where possible. + + + +### LoCoMo + +| Method | Single Hop | Multi Hop | Temporal | Open Domain | Overall | +|--------|------------|-----------|-----------|-------------|-----------| +| MemoryOS | 62.43 | 56.50 | 37.18 | 40.28 | 54.70 | +| Mem0 | 66.71 | 58.16 | 55.45 | 40.62 | 61.00 | +| MemU | 72.77 | 62.41 | 33.96 | 46.88 | 61.15 | +| MemOS | 81.45 | 69.15 | 72.27 | 60.42 | 75.87 | +| HiMem | 89.22 | 70.92 | 74.77 | 54.86 | 80.71 | +| Zep | 88.11 | 71.99 | 74.45 | 66.67 | 81.06 | +| TiMem | 81.43 | 62.20 | 77.63 | 52.08 | 75.30 | +| TSM | 84.30 | 66.67 | 71.03 | 58.33 | 76.69 | +| MemR3 | 89.44 | 71.39 | 76.22 | 61.11 | 81.55 | +| **ReMe** | **89.89** | **82.98** | **83.80** | **71.88** | **86.23** | + + +### HaluMem + +| Method | Memory Integrity | Memory Accuracy | QA Accuracy | +|-------------|------------------|---------------|-------------| +| MemoBase | 14.55 | 92.24 | 35.53 | +| Supermemory | 41.53 | 90.32 | 54.07 | +| Mem0 | 42.91 | 86.26 | 53.02 | +| ProMem | **73.80** | 89.47 | 62.26 | +| **ReMe** | 67.72 | **94.06** | **88.78** | + + ### Python usage ```python diff --git a/README_ZH.md b/README_ZH.md index c3de853f..3d37e30f 100644 --- a/README_ZH.md +++ b/README_ZH.md @@ -25,7 +25,8 @@ --- -🧠 ReMe 是一个专为 **AI 智能体** 打造的记忆管理框架，同时提供基于文件系统和基于向量库的记忆系统。 +🧠 ReMe 是一个专为 **AI 智能体** 打造的记忆管理框架，同时提供基于[文件系统](#-基于文件的记忆系统-remelight)和基于[向量库](#-基于向量库的记忆系统)的记忆系统。 + 它解决智能体记忆的两类核心问题：**上下文窗口有限**（长对话时早期信息被截断或丢失）、**会话无状态**（新对话无法继承历史，每次从零开始）。 @@ -396,6 +397,44 @@ graph LR 安装和环境变量配置与 [ReMeLight 一致](#安装)，通过环境变量设置 API 密钥，可写在项目根目录的 `.env` 文件中。 +## 🧪 实验 + +本实验部分在 LoCoMo、LongMemEval、HaluMem 三个数据集上进行评测，实验设置如下： + +1. **ReMe 使用模型**：如各表 backbone 列所示。 +2. **评估使用模型**：采用 LLM-as-a-Judge 协议（参照 MemOS）——每条回答由 GPT-4o-mini 裁判模型打分。 + +实验设置尽量与各基线论文保持一致，以复用其公开结果。 + + +### LoCoMo + +| Method | Single Hop | Multi Hop | Temporal | Open Domain | Overall | +|--------|------------|-----------|-----------|-------------|-----------| +| MemoryOS | 62.43 | 56.50 | 37.18 | 40.28 | 54.70 | +| Mem0 | 66.71 | 58.16 | 55.45 | 40.62 | 61.00 | +| MemU | 72.77 | 62.41 | 33.96 | 46.88 | 61.15 | +| MemOS | 81.45 | 69.15 | 72.27 | 60.42 | 75.87 | +| HiMem | 89.22 | 70.92 | 74.77 | 54.86 | 80.71 | +| Zep | 88.11 | 71.99 | 74.45 | 66.67 | 81.06 | +| TiMem | 81.43 | 62.20 | 77.63 | 52.08 | 75.30 | +| TSM | 84.30 | 66.67 | 71.03 | 58.33 | 76.69 | +| MemR3 | 89.44 | 71.39 | 76.22 | 61.11 | 81.55 | +| **ReMe** | **89.89** | **82.98** | **83.80** | **71.88** | **86.23** | + + +### HaluMem + +| Method | Memory Integrity | Memory Accuracy | QA Accuracy | +|-------------|------------------|---------------|-------------| +| MemoBase | 14.55 | 92.24 | 35.53 | +| Supermemory | 41.53 | 90.32 | 54.07 | +| Mem0 | 42.91 | 86.26 | 53.02 | +| ProMem | **73.80** | 89.47 | 62.26 | +| **ReMe** | 67.72 | **94.06** | **88.78** | + + + ### Python 使用 ```python From 9a6cf2b994e04315614f474d5f98474edecd15e5 Mon Sep 17 00:00:00 2001 From: jinliyl <6469360+jinliyl@users.noreply.github.com> Date: Tue, 17 Mar 2026 11:07:31 +0800 Subject: [PATCH 26/59] Dev/token (#159) * update * refactor(memory): remove unnecessary type check and update error logging * refactor(core): standardize logger import and update agentscope dependency * fix(memory): disable console output and add logging for summarizer component * feat(core): replace OpenAI token counter with custom ReMe token counter - Replace OpenAITokenCounter with ReMeTokenCounter implementation - Add support for HuggingFace mirror and configurable tokenizer - Register ReMeTokenCounter as default token counter in registry - Update config to use hf backend with Qwen2.5-7B-Instruct model refactor(memory): convert token counting methods to async in message handlers - Change count_str_token, stat_message, count_msgs_token to async methods - Update format_msgs_to_str and context_check to use async token counting - Modify _format_tool_result_output to support async token counting - Adjust all dependent methods to await async token counting calls feat(memory): add dialog persistence to in-memory storage - Implement _append_messages_to_dialog for saving messages to JSONL files - Add dialog_path parameter to ReMeInMemoryMemory constructor - Persist messages to daily JSONL files based on timestamp grouping - Update mark_messages_compressed to save and remove compressed messages - Modify clear_content to persist all messages before clearing memory refactor(ops): update token counter type hints and initialization - Change BaseOp to use HuggingFaceTokenCounter instead of TokenCounterBase - Update type annotations for as_token_counter property and parameters - Remove direct token counter injection from Compactor and ContextChecker - Pass as_token_counter parameter through service context mechanism style(logging): improve error logging with exception details - Replace logger.error with logger.exception in browser control tool - Change logger.error to logger.exception in memory get tool error handling - Add proper exception logging with stack trace information chore(config): add token counter configuration to light YAML - Add as_token_counters section with default hf backend configuration - Configure Qwen/Qwen2.5-7B-Instruct model with mirror support enabled - Set up pretrained_model_name_or_path and use_mirror parameters test(context): update context check tests to async implementation - Convert verify_context_check_invariants to async function - Update context check test methods to use async calls - Change stat_message calls to await async implementation - Modify test_empty_messages and test_below_threshold_returns_all to async * feat(core): implement context checking and memory management features * refactor(core): replace direct loguru import with logger utility function * refactor(reme): remove RuntimeContext dependency and simplify context checking * feat(docs): add raw conversation persistence to ReMe framework --- README.md | 96 +- README_ZH.md | 95 +- pyproject.toml | 2 +- reme/__init__.py | 2 +- reme/config/light.yaml | 6 + reme/core/application.py | 8 +- reme/core/as_llm_formatter/__init__.py | 4 +- .../reme_openai_chat_formatter.py | 215 ++ reme/core/as_token_counter/__init__.py | 7 +- .../as_token_counter/reme_token_counter.py | 114 + reme/core/op/base_op.py | 6 +- reme/core/utils/__init__.py | 4 +- reme/memory/file_based/components/cli.py | 29 +- .../memory/file_based/components/compactor.py | 17 +- .../file_based/components/context_checker.py | 14 +- .../file_based/components/summarizer.py | 23 +- .../components/tool_result_compactor.py | 4 +- .../file_based/reme_in_memory_memory.py | 144 +- .../file_based/tools/browser_control.py | 2624 +++++++++++++++++ .../file_based/tools/browser_snapshot.py | 248 ++ reme/memory/file_based/tools/memory_get.py | 7 +- reme/memory/file_based/tools/memory_search.py | 6 +- .../memory/file_based/utils/as_msg_handler.py | 67 +- reme/reme_light.py | 104 +- tests/light/test_compactor.py | 10 +- tests/light/test_context_check.py | 350 ++- tests/light/test_format_msgs_to_str.py | 177 +- tests/light/test_reme_light.py | 22 +- tests/light/test_summarizer.py | 8 +- tests/light/test_tools.py | 2 +- 30 files changed, 3939 insertions(+), 476 deletions(-) create mode 100644 reme/core/as_llm_formatter/reme_openai_chat_formatter.py create mode 100644 reme/core/as_token_counter/reme_token_counter.py create mode 100644 reme/memory/file_based/tools/browser_control.py create mode 100644 reme/memory/file_based/tools/browser_snapshot.py diff --git a/README.md b/README.md index f379652f..db94195d 100644 --- a/README.md +++ b/README.md @@ -25,7 +25,8 @@ --- -🧠 ReMe is a memory management framework designed for **AI agents**, providing both [file-based](#-file-based-memory-system-remelight) and [vector-based](#-vector-based-memory-system) memory systems. +🧠 ReMe is a memory management framework designed for **AI agents**, providing +both [file-based](#-file-based-memory-system-remelight) and [vector-based](#-vector-based-memory-system) memory systems. It tackles two core problems of agent memory: **limited context window** (early information is truncated or lost in long conversations) and **stateless sessions** (new sessions cannot inherit history and always start from scratch). @@ -73,6 +74,8 @@ working_dir/ ├── MEMORY.md # Long-term memory: persistent info such as user preferences ├── memory/ │ └── YYYY-MM-DD.md # Daily journal: automatically written after each conversation +├── dialog/ # Raw conversation records: full dialog before compression +│ └── YYYY-MM-DD.jsonl # Daily conversation messages in JSONL format └── tool_result/ # Cache for long tool outputs (auto-managed, expired entries auto-cleaned) └── .txt ``` @@ -90,6 +93,8 @@ capabilities for AI agents: pre_reasoning_hook🔄 Pre-reasoning hookcompact_tool_result + check_context + compact_memory + summary_memory (async) Long-term Memorysummary_memory📝 Persist important memory to filesSummarizer — ReActAgent + file tools (read / write / edit) memory_search🔍 Semantic memory searchMemorySearch — hybrid retrieval with vectors + BM25 +Session Memoryget_in_memory_memory💾 Create in-session memory instanceReturns ReMeInMemoryMemory with dialog_path configured for persistence +await_summary_tasks⏳ Wait for async summary tasksBlock until all background summary tasks complete -start🚀 Start memory systemInitialize file storage, file watcher, and embedding cache; clean up expired tool result files -close📕 Shutdown and cleanupClean up tool result files, stop file watcher, and persist embedding cache @@ -145,8 +150,12 @@ async def main(): messages = [...] # List of conversation messages - # 1. Compact long tool outputs (prevent tool results from blowing up context) - messages = await reme.compact_tool_result(messages) + # 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( @@ -157,10 +166,10 @@ async def main(): language="zh", # Summary language (e.g., "zh" / "") ) - # 3. Submit summary task asynchronously (non-blocking, writes to memory/YYYY-MM-DD.md) - reme.add_async_summary_task(messages=messages) + # 3. Compact long tool outputs (prevent tool results from blowing up context) + messages = await reme.compact_tool_result(messages) - # 4. Pre-reasoning hook (auto compact tool results + generate summaries) + # 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.", @@ -172,12 +181,17 @@ async def main(): tool_result_compact_keep_n=3, ) - # 5. Semantic memory search (vector + BM25 hybrid retrieval) + # 5. Persist important memory to files (writes to memory/YYYY-MM-DD.md) + summary_result = await reme.summary_memory( + messages=messages, + language="zh", + ) + + # 6. Semantic memory search (vector + BM25 hybrid retrieval) result = await reme.memory_search(query="Python version preference", max_results=5) - # 6. Create in-session memory instance (manages context for one conversation) - from reme.memory.file_based.reme_in_memory_memory import ReMeInMemoryMemory - memory = ReMeInMemoryMemory() + # 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) @@ -185,8 +199,8 @@ async def main(): print(f"Message token count: {token_stats['messages_tokens']}") print(f"Estimated total tokens: {token_stats['estimated_tokens']}") - # 7. Wait for background summary tasks to complete before shutdown - 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) # Shutdown ReMeLight await reme.close() @@ -214,8 +228,11 @@ graph LR CC -->|Exceeds limit| CM[compact_memory
Generate summary] CC -->|Exceeds limit| SM[summary_memory
Async persistence] SM -->|ReAct + FileIO| Files[memory/*.md] + CC -->|Exceeds limit| MMC[mark_messages_compressed
Persist raw dialog] + MMC --> Dialog[dialog/*.jsonl] Agent -->|Explicit call| Search[memory_search
Vector+BM25] Agent -->|In - session| InMem[ReMeInMemoryMemory
Token-aware memory] + InMem -->|Compress/Clear| Dialog Files -.->|FileWatcher| Store[(FileStore
Vector+FTS index)] Search --> Store ``` @@ -340,7 +357,7 @@ graph LR #### 6. `ReMeInMemoryMemory` — in-session memory [ReMeInMemoryMemory](reme/memory/file_based/reme_in_memory_memory.py) extends AgentScope's `InMemoryMemory` to provide -token-aware memory management. +token-aware memory management and raw conversation persistence. ```mermaid graph LR @@ -350,13 +367,20 @@ graph LR P -->|Yes| S[Prepend previous summary] S --> O[Output messages] P -->|No| O + M[mark_messages_compressed] --> D[Persist to dialog/YYYY-MM-DD.jsonl] + D --> R[Remove from memory] ``` -| Function | Description | -|----------------------------------|---------------------------------------------------| -| `get_memory` | Filter messages by mark and auto-append summary | -| `estimate_tokens` | Estimate token usage of the context | -| `state_dict` / `load_state_dict` | Serialize/deserialize state (session persistence) | +| Function | Description | +|----------------------------------|----------------------------------------------------------| +| `get_memory` | Filter messages by mark and auto-append summary | +| `estimate_tokens` | Estimate token usage of the context | +| `state_dict` / `load_state_dict` | Serialize/deserialize state (session persistence) | +| `mark_messages_compressed` | Mark messages compressed and persist to dialog directory | +| `clear_content` | Persist all messages before clearing memory | + +**Raw conversation persistence**: When messages are compressed or cleared, they are automatically saved to +`{dialog_path}/{date}.jsonl` with one JSON-formatted message per line. --- @@ -424,34 +448,30 @@ Evaluations are conducted on three benchmarks: **LoCoMo** and **HaluMem**. Exper Baseline results are reproduced from their respective papers under aligned settings where possible. - - ### LoCoMo -| Method | Single Hop | Multi Hop | Temporal | Open Domain | Overall | -|--------|------------|-----------|-----------|-------------|-----------| +| Method | Single Hop | Multi Hop | Temporal | Open Domain | Overall | +|----------|------------|-----------|-----------|-------------|-----------| | MemoryOS | 62.43 | 56.50 | 37.18 | 40.28 | 54.70 | -| Mem0 | 66.71 | 58.16 | 55.45 | 40.62 | 61.00 | -| MemU | 72.77 | 62.41 | 33.96 | 46.88 | 61.15 | -| MemOS | 81.45 | 69.15 | 72.27 | 60.42 | 75.87 | -| HiMem | 89.22 | 70.92 | 74.77 | 54.86 | 80.71 | -| Zep | 88.11 | 71.99 | 74.45 | 66.67 | 81.06 | -| TiMem | 81.43 | 62.20 | 77.63 | 52.08 | 75.30 | -| TSM | 84.30 | 66.67 | 71.03 | 58.33 | 76.69 | -| MemR3 | 89.44 | 71.39 | 76.22 | 61.11 | 81.55 | +| Mem0 | 66.71 | 58.16 | 55.45 | 40.62 | 61.00 | +| MemU | 72.77 | 62.41 | 33.96 | 46.88 | 61.15 | +| MemOS | 81.45 | 69.15 | 72.27 | 60.42 | 75.87 | +| HiMem | 89.22 | 70.92 | 74.77 | 54.86 | 80.71 | +| Zep | 88.11 | 71.99 | 74.45 | 66.67 | 81.06 | +| TiMem | 81.43 | 62.20 | 77.63 | 52.08 | 75.30 | +| TSM | 84.30 | 66.67 | 71.03 | 58.33 | 76.69 | +| MemR3 | 89.44 | 71.39 | 76.22 | 61.11 | 81.55 | | **ReMe** | **89.89** | **82.98** | **83.80** | **71.88** | **86.23** | - ### HaluMem | Method | Memory Integrity | Memory Accuracy | QA Accuracy | -|-------------|------------------|---------------|-------------| -| MemoBase | 14.55 | 92.24 | 35.53 | -| Supermemory | 41.53 | 90.32 | 54.07 | -| Mem0 | 42.91 | 86.26 | 53.02 | -| ProMem | **73.80** | 89.47 | 62.26 | -| **ReMe** | 67.72 | **94.06** | **88.78** | - +|-------------|------------------|-----------------|-------------| +| MemoBase | 14.55 | 92.24 | 35.53 | +| Supermemory | 41.53 | 90.32 | 54.07 | +| Mem0 | 42.91 | 86.26 | 53.02 | +| ProMem | **73.80** | 89.47 | 62.26 | +| **ReMe** | 67.72 | **94.06** | **88.78** | ### Python usage diff --git a/README_ZH.md b/README_ZH.md index 3d37e30f..6218b87c 100644 --- a/README_ZH.md +++ b/README_ZH.md @@ -25,8 +25,8 @@ --- -🧠 ReMe 是一个专为 **AI 智能体** 打造的记忆管理框架，同时提供基于[文件系统](#-基于文件的记忆系统-remelight)和基于[向量库](#-基于向量库的记忆系统)的记忆系统。 - +🧠 ReMe 是一个专为 **AI 智能体** 打造的记忆管理框架，同时提供基于[文件系统](#-基于文件的记忆系统-remelight) +和基于[向量库](#-基于向量库的记忆系统)的记忆系统。 它解决智能体记忆的两类核心问题：**上下文窗口有限**（长对话时早期信息被截断或丢失）、**会话无状态**（新对话无法继承历史，每次从零开始）。 @@ -67,6 +67,8 @@ working_dir/ ├── MEMORY.md # 长期记忆：用户偏好等持久信息 ├── memory/ │ └── YYYY-MM-DD.md # 每日日记：对话结束后自动写入 +├── dialog/ # 原始对话记录：压缩前的完整对话 +│ └── YYYY-MM-DD.jsonl # 按日期存储的对话消息（JSONL 格式） └── tool_result/ # 超长工具输出缓存（自动管理，超期自动清理） └── .txt ``` @@ -83,6 +85,8 @@ working_dir/ 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 缓存 @@ -138,8 +142,12 @@ async def main(): messages = [...] # 对话消息列表 - # 1. 压缩超长工具输出（防止工具结果撑爆上下文） - messages = await reme.compact_tool_result(messages) + # 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( @@ -150,10 +158,10 @@ async def main(): language="zh", # 摘要语言（zh / ""） ) - # 3. 后台异步提交摘要任务（不阻塞对话，摘要写入 memory/YYYY-MM-DD.md） - reme.add_async_summary_task(messages=messages) + # 3. 压缩超长工具输出（防止工具结果撑爆上下文） + messages = await reme.compact_tool_result(messages) - # 4. 推理前预处理钩子（自动压缩工具结果 + 生成摘要） + # 4. 推理前预处理钩子（自动压缩工具结果 + 检查上下文 + 生成摘要） processed_messages, compressed_summary = await reme.pre_reasoning_hook( messages=messages, system_prompt="你是一个有帮助的 AI 助手。", @@ -165,12 +173,18 @@ async def main(): tool_result_compact_keep_n=3, ) - # 5. 语义搜索记忆（向量 + BM25 混合检索） + # 5. 将重要记忆写入文件（摘要写入 memory/YYYY-MM-DD.md） + summary_result = await reme.summary_memory( + messages=messages, + language="zh", + ) + + # 6. 语义搜索记忆（向量 + BM25 混合检索） result = await reme.memory_search(query="Python 版本偏好", max_results=5) - # 6. 创建会话内存实例（管理单次对话的上下文） + # 7. 创建会话内存实例（管理单次对话的上下文） from reme.memory.file_based.reme_in_memory_memory import ReMeInMemoryMemory - memory = 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) @@ -178,8 +192,8 @@ async def main(): print(f"消息 Token 数: {token_stats['messages_tokens']}") print(f"预估总 Token 数: {token_stats['estimated_tokens']}") - # 7. 关闭前等待后台任务完成 - 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() @@ -205,8 +219,11 @@ graph LR CC -->|超限| CM[compact_memory
生成摘要] CC -->|超限| SM[summary_memory
异步持久化] SM -->|ReAct + FileIO| Files[memory/*.md] + CC -->|超限| MMC[mark_messages_compressed
持久化原始对话] + MMC --> Dialog[dialog/*.jsonl] Agent -->|主动调用| Search[memory_search
向量+BM25] Agent -->|会话内存| InMem[ReMeInMemoryMemory
Token感知内存] + InMem -->|压缩/清空| Dialog Files -.->|FileWatcher| Store[(FileStore
向量+FTS索引)] Search --> Store ``` @@ -325,7 +342,7 @@ graph LR #### 6. ReMeInMemoryMemory — 会话内存 [ReMeInMemoryMemory](reme/memory/file_based/reme_in_memory_memory.py) 扩展 AgentScope 的 `InMemoryMemory`，提供 Token -感知的内存管理。 +感知的内存管理和原始对话持久化能力。 ```mermaid graph LR @@ -335,13 +352,19 @@ graph LR P -->|是| S[头部插入 previous-summary] S --> O[输出 messages] P -->|否| O + M[mark_messages_compressed] --> D[持久化到 dialog/YYYY-MM-DD.jsonl] + D --> R[从内存移除] ``` -| 功能 | 说明 | -|----------------------------------|-------------------| -| `get_memory` | 按标记过滤，自动追加压缩摘要 | -| `estimate_tokens` | 估算上下文 Token 用量 | -| `state_dict` / `load_state_dict` | 状态序列化/反序列化（会话持久化） | +| 功能 | 说明 | +|----------------------------------|-----------------------| +| `get_memory` | 按标记过滤，自动追加压缩摘要 | +| `estimate_tokens` | 估算上下文 Token 用量 | +| `state_dict` / `load_state_dict` | 状态序列化/反序列化（会话持久化） | +| `mark_messages_compressed` | 标记消息压缩并持久化到 dialog 目录 | +| `clear_content` | 持久化所有消息后清空内存 | + +**原始对话持久化**：当消息被压缩或清空时，自动保存到 `{dialog_path}/{date}.jsonl`，每行一条 JSON 格式的消息记录。 --- @@ -406,34 +429,30 @@ graph LR 实验设置尽量与各基线论文保持一致，以复用其公开结果。 - ### LoCoMo -| Method | Single Hop | Multi Hop | Temporal | Open Domain | Overall | -|--------|------------|-----------|-----------|-------------|-----------| +| Method | Single Hop | Multi Hop | Temporal | Open Domain | Overall | +|----------|------------|-----------|-----------|-------------|-----------| | MemoryOS | 62.43 | 56.50 | 37.18 | 40.28 | 54.70 | -| Mem0 | 66.71 | 58.16 | 55.45 | 40.62 | 61.00 | -| MemU | 72.77 | 62.41 | 33.96 | 46.88 | 61.15 | -| MemOS | 81.45 | 69.15 | 72.27 | 60.42 | 75.87 | -| HiMem | 89.22 | 70.92 | 74.77 | 54.86 | 80.71 | -| Zep | 88.11 | 71.99 | 74.45 | 66.67 | 81.06 | -| TiMem | 81.43 | 62.20 | 77.63 | 52.08 | 75.30 | -| TSM | 84.30 | 66.67 | 71.03 | 58.33 | 76.69 | -| MemR3 | 89.44 | 71.39 | 76.22 | 61.11 | 81.55 | +| Mem0 | 66.71 | 58.16 | 55.45 | 40.62 | 61.00 | +| MemU | 72.77 | 62.41 | 33.96 | 46.88 | 61.15 | +| MemOS | 81.45 | 69.15 | 72.27 | 60.42 | 75.87 | +| HiMem | 89.22 | 70.92 | 74.77 | 54.86 | 80.71 | +| Zep | 88.11 | 71.99 | 74.45 | 66.67 | 81.06 | +| TiMem | 81.43 | 62.20 | 77.63 | 52.08 | 75.30 | +| TSM | 84.30 | 66.67 | 71.03 | 58.33 | 76.69 | +| MemR3 | 89.44 | 71.39 | 76.22 | 61.11 | 81.55 | | **ReMe** | **89.89** | **82.98** | **83.80** | **71.88** | **86.23** | - ### HaluMem | Method | Memory Integrity | Memory Accuracy | QA Accuracy | -|-------------|------------------|---------------|-------------| -| MemoBase | 14.55 | 92.24 | 35.53 | -| Supermemory | 41.53 | 90.32 | 54.07 | -| Mem0 | 42.91 | 86.26 | 53.02 | -| ProMem | **73.80** | 89.47 | 62.26 | -| **ReMe** | 67.72 | **94.06** | **88.78** | - - +|-------------|------------------|-----------------|-------------| +| MemoBase | 14.55 | 92.24 | 35.53 | +| Supermemory | 41.53 | 90.32 | 54.07 | +| Mem0 | 42.91 | 86.26 | 53.02 | +| ProMem | **73.80** | 89.47 | 62.26 | +| **ReMe** | 67.72 | **94.06** | **88.78** | ### Python 使用 diff --git a/pyproject.toml b/pyproject.toml index 43da64dd..4b91b329 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -81,7 +81,7 @@ full = [ ] light = [ - "agentscope==1.0.16.dev0", + "agentscope==1.0.17", ] [tool.setuptools.packages.find] diff --git a/reme/__init__.py b/reme/__init__.py index 41537822..32ed21d0 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.6b3" +__version__ = "0.3.0.6" __all__ = [ "config", diff --git a/reme/config/light.yaml b/reme/config/light.yaml index bc85c10d..10a22349 100644 --- a/reme/config/light.yaml +++ b/reme/config/light.yaml @@ -7,6 +7,12 @@ as_llm_formatters: default: backend: openai +as_token_counters: + default: + backend: hf + pretrained_model_name_or_path: Qwen/Qwen2.5-7B-Instruct + use_mirror: true + embedding_models: default: backend: openai diff --git a/reme/core/application.py b/reme/core/application.py index 85a33c4d..8778dfba 100644 --- a/reme/core/application.py +++ b/reme/core/application.py @@ -5,8 +5,6 @@ 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 @@ -17,9 +15,11 @@ from .registry_factory import R from .schema import Response, ServiceConfig 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 from .vector_store import BaseVectorStore +logger = get_logger() + class Application: """Application wrapper that wires together service context, flows, and runtimes.""" @@ -244,6 +244,7 @@ class Application: await self.prepare_mcp_servers() self._started = True + logger.info("ReMe Application started") return self async def prepare_mcp_servers(self): @@ -290,6 +291,7 @@ class Application: self.shutdown_ray() self._started = False + logger.info("ReMe Application closed") return False def shutdown_thread_pool(self, wait: bool = True): diff --git a/reme/core/as_llm_formatter/__init__.py b/reme/core/as_llm_formatter/__init__.py index 88b326a7..1c7eee46 100644 --- a/reme/core/as_llm_formatter/__init__.py +++ b/reme/core/as_llm_formatter/__init__.py @@ -1,9 +1,9 @@ """Module for registering AgentScope LLM formatters.""" from agentscope.formatter import DashScopeChatFormatter -from agentscope.formatter import OpenAIChatFormatter +from .reme_openai_chat_formatter import ReMeOpenAIChatFormatter from ..registry_factory import R -R.as_llm_formatters.register("openai")(OpenAIChatFormatter) +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 index 51b3bd29..c8b017d2 100644 --- a/reme/core/as_token_counter/__init__.py +++ b/reme/core/as_token_counter/__init__.py @@ -1,9 +1,6 @@ """Module for registering AgentScope token counters.""" -from agentscope.token import OpenAITokenCounter -from agentscope.token import HuggingFaceTokenCounter - +from .reme_token_counter import ReMeTokenCounter from ..registry_factory import R -R.as_token_counters.register("openai")(OpenAITokenCounter) -R.as_token_counters.register("hf")(HuggingFaceTokenCounter) +R.as_token_counters.register("hf")(ReMeTokenCounter) 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..721b98db --- /dev/null +++ b/reme/core/as_token_counter/reme_token_counter.py @@ -0,0 +1,114 @@ +"""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: + os.environ["HF_ENDPOINT"] = "https://hf-mirror.com" + else: + os.environ.pop("HF_ENDPOINT", None) + + 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.exception("Failed to initialize tokenizer: %s", 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/op/base_op.py b/reme/core/op/base_op.py index f25da289..a84ba18d 100644 --- a/reme/core/op/base_op.py +++ b/reme/core/op/base_op.py @@ -9,7 +9,7 @@ from typing import Callable, Optional, Any from agentscope.formatter import FormatterBase from agentscope.model import ChatModelBase -from agentscope.token import TokenCounterBase +from agentscope.token import HuggingFaceTokenCounter from loguru import logger from tqdm import tqdm @@ -47,7 +47,7 @@ class BaseOp(metaclass=ABCMeta): prompt_path: str = "", as_llm: str | ChatModelBase = "default", as_llm_formatter: str | FormatterBase = "default", - as_token_counter: str | TokenCounterBase = "default", + as_token_counter: str | HuggingFaceTokenCounter = "default", llm: str | BaseLLM = "default", embedding_model: str | BaseEmbeddingModel = "default", vector_store: str | BaseVectorStore = "default", @@ -153,7 +153,7 @@ class BaseOp(metaclass=ABCMeta): return self._as_llm_formatter @property - def as_token_counter(self) -> TokenCounterBase: + 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] diff --git a/reme/core/utils/__init__.py b/reme/core/utils/__init__.py index c1f35adb..9dbc59dc 100644 --- a/reme/core/utils/__init__.py +++ b/reme/core/utils/__init__.py @@ -11,7 +11,7 @@ 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 as get_std_logger +from .std_logger import get_logger from .logo_utils import print_logo from .mcp_client import MCPClient from .pydantic_config_parser import PydanticConfigParser @@ -42,7 +42,7 @@ __all__ = [ "format_messages", "deduplicate_memories", "init_logger", - "get_std_logger", + "get_logger", "print_logo", "MCPClient", "PydanticConfigParser", diff --git a/reme/memory/file_based/components/cli.py b/reme/memory/file_based/components/cli.py index bc88bd56..f4dddc21 100644 --- a/reme/memory/file_based/components/cli.py +++ b/reme/memory/file_based/components/cli.py @@ -6,16 +6,32 @@ from pathlib import Path from agentscope.agent import ReActAgent from agentscope.message import Msg, TextBlock -from agentscope.tool import Toolkit, ToolResponse from agentscope.pipeline import stream_printing_messages -from loguru import logger +from agentscope.tool import Toolkit, ToolResponse -from ....core.op import BaseOp -from ....core.utils import format_messages 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): @@ -117,7 +133,7 @@ class CliAgent(BaseOp): # Create context checker checker = ContextChecker( memory_compact_threshold=self.context_window_tokens - self.reserve_tokens, - memory_compact_reserve=self.reserve_tokens, + memory_compact_reserve=self.keep_recent_tokens, token_counter=self.as_token_counter, ) @@ -199,7 +215,7 @@ class CliAgent(BaseOp): return messages - async def memory_search(self, query: str, max_results: int = 5, min_score: float = 0.1) -> str: + 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; @@ -257,6 +273,7 @@ class CliAgent(BaseOp): 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 diff --git a/reme/memory/file_based/components/compactor.py b/reme/memory/file_based/components/compactor.py index fbe505ad..b6725b0b 100644 --- a/reme/memory/file_based/components/compactor.py +++ b/reme/memory/file_based/components/compactor.py @@ -2,11 +2,12 @@ from agentscope.agent import ReActAgent from agentscope.message import Msg -from agentscope.token import HuggingFaceTokenCounter -from loguru import logger from ..utils import AsMsgHandler from ....core.op import BaseOp +from ....core.utils import get_logger + +logger = get_logger() class Compactor(BaseOp): @@ -15,14 +16,11 @@ class Compactor(BaseOp): def __init__( self, memory_compact_threshold: int, - token_counter: HuggingFaceTokenCounter, - console_enabled: bool = True, + console_enabled: bool = False, **kwargs, ): super().__init__(**kwargs) self.memory_compact_threshold: int = memory_compact_threshold - - self.msg_handler = AsMsgHandler(token_counter=token_counter) self.console_enabled: bool = console_enabled async def execute(self): @@ -32,12 +30,13 @@ class Compactor(BaseOp): if not messages: return "" - before_token_count = self.msg_handler.count_msgs_token(messages) - history_formatted_str: str = self.msg_handler.format_msgs_to_str( + 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, ) - after_token_count = self.msg_handler.count_str_token(history_formatted_str) + 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: diff --git a/reme/memory/file_based/components/context_checker.py b/reme/memory/file_based/components/context_checker.py index 18ac4bf0..f5a49b2b 100644 --- a/reme/memory/file_based/components/context_checker.py +++ b/reme/memory/file_based/components/context_checker.py @@ -1,13 +1,12 @@ """ContextChecker module for checking context size and splitting messages.""" from agentscope.message import Msg -from agentscope.token import HuggingFaceTokenCounter from ..utils import AsMsgHandler from ....core.op import BaseOp -from ....core.utils import get_std_logger +from ....core.utils import get_logger -logger = get_std_logger() +logger = get_logger() class ContextChecker(BaseOp): @@ -21,14 +20,12 @@ class ContextChecker(BaseOp): Attributes: memory_compact_threshold (int): Token count threshold for triggering compaction. memory_compact_reserve (int): Token count to reserve for recent messages. - msg_handler (AsMsgHandler): Handler for message processing and token counting. """ def __init__( self, memory_compact_threshold: int, memory_compact_reserve: int = 10000, - token_counter: HuggingFaceTokenCounter | None = None, **kwargs, ): """ @@ -39,8 +36,6 @@ class ContextChecker(BaseOp): 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. - token_counter (HuggingFaceTokenCounter | None): Token counter for - measuring content length. If None, a default counter will be used. **kwargs: Additional keyword arguments passed to BaseOp. """ super().__init__(**kwargs) @@ -48,8 +43,6 @@ class ContextChecker(BaseOp): self.memory_compact_reserve: int = memory_compact_reserve assert self.memory_compact_threshold > self.memory_compact_reserve - self.msg_handler = AsMsgHandler(token_counter=token_counter) - async def execute(self) -> tuple[list[Msg], list[Msg], bool]: """ Execute context check and split messages. @@ -81,7 +74,8 @@ class ContextChecker(BaseOp): logger.info("ContextChecker: No messages to check.") return [], [], True - messages_to_compact, messages_to_keep, is_valid = self.msg_handler.context_check( + 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, diff --git a/reme/memory/file_based/components/summarizer.py b/reme/memory/file_based/components/summarizer.py index 83441149..d553a5f6 100644 --- a/reme/memory/file_based/components/summarizer.py +++ b/reme/memory/file_based/components/summarizer.py @@ -4,12 +4,13 @@ import datetime from agentscope.agent import ReActAgent from agentscope.message import Msg -from agentscope.token import HuggingFaceTokenCounter from agentscope.tool import Toolkit -from loguru import logger from ..utils import AsMsgHandler from ....core.op import BaseOp +from ....core.utils import get_logger + +logger = get_logger() class Summarizer(BaseOp): @@ -20,18 +21,15 @@ class Summarizer(BaseOp): working_dir: str, memory_dir: str, memory_compact_threshold: int, - token_counter: HuggingFaceTokenCounter, - toolkit: Toolkit, - console_enabled: bool = True, + toolkit: Toolkit | None = None, + console_enabled: bool = False, **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.msg_handler = AsMsgHandler(token_counter=token_counter) - self.toolkit: Toolkit = toolkit + self.toolkit: Toolkit | None = toolkit self.console_enabled: bool = console_enabled async def execute(self): @@ -40,12 +38,13 @@ class Summarizer(BaseOp): if not messages: return "" - before_token_count = self.msg_handler.count_msgs_token(messages) - history_formatted_str: str = self.msg_handler.format_msgs_to_str( + 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, ) - after_token_count = self.msg_handler.count_str_token(history_formatted_str) + 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: @@ -75,6 +74,8 @@ class Summarizer(BaseOp): 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}") diff --git a/reme/memory/file_based/components/tool_result_compactor.py b/reme/memory/file_based/components/tool_result_compactor.py index 412df6de..c1c6fe00 100644 --- a/reme/memory/file_based/components/tool_result_compactor.py +++ b/reme/memory/file_based/components/tool_result_compactor.py @@ -7,10 +7,10 @@ from pathlib import Path from agentscope.message import Msg from ....core.op import BaseOp -from ....core.utils import get_std_logger +from ....core.utils import get_logger from ....core.utils import truncate_text, is_truncated -logger = get_std_logger() +logger = get_logger() class ToolResultCompactor(BaseOp): diff --git a/reme/memory/file_based/reme_in_memory_memory.py b/reme/memory/file_based/reme_in_memory_memory.py index da2118e8..b1c8c16c 100644 --- a/reme/memory/file_based/reme_in_memory_memory.py +++ b/reme/memory/file_based/reme_in_memory_memory.py @@ -1,23 +1,112 @@ """Custom memory implementation with bugfixes and extensions.""" +import json +from datetime import datetime +from pathlib import Path + 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 AsMsgHandler -from ...core.utils import get_std_logger +from ...core.utils import get_logger -logger = get_std_logger() +logger = get_logger() class ReMeInMemoryMemory(InMemoryMemory): """Extended InMemoryMemory with bugfixes and summary support.""" - def __init__(self, token_counter: HuggingFaceTokenCounter): + def __init__( + self, + token_counter: HuggingFaceTokenCounter, + 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._msg_handler: AsMsgHandler = AsMsgHandler(token_counter) + self._dialog_path: Path | None = Path(dialog_path) if dialog_path else None + + 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, @@ -105,19 +194,52 @@ 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 + 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.""" + """Persist all messages to dialog storage and clear the content. + + 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. @@ -141,10 +263,10 @@ Use it as context to maintain continuity. ) compressed_summary = self.get_compressed_summary() - compressed_summary_tokens = self._msg_handler.count_str_token(compressed_summary) + compressed_summary_tokens = await self._msg_handler.count_str_token(compressed_summary) # Build per-message token details using AsMsgHandler - messages_detail = [self._msg_handler.stat_message(msg) for msg in messages] + 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) 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/memory_get.py b/reme/memory/file_based/tools/memory_get.py index 572a4a26..9cc16968 100644 --- a/reme/memory/file_based/tools/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/file_based/tools/memory_search.py b/reme/memory/file_based/tools/memory_search.py index 7423b331..927a6fc1 100644 --- a/reme/memory/file_based/tools/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/utils/as_msg_handler.py b/reme/memory/file_based/utils/as_msg_handler.py index 30facedc..46bf81b6 100644 --- a/reme/memory/file_based/utils/as_msg_handler.py +++ b/reme/memory/file_based/utils/as_msg_handler.py @@ -6,9 +6,9 @@ from agentscope.message import Msg from agentscope.token import HuggingFaceTokenCounter from ....core.schema import AsMsgStat, AsBlockStat -from ....core.utils import get_std_logger +from ....core.utils import get_logger -logger = get_std_logger() +logger = get_logger() class AsMsgHandler: @@ -17,7 +17,7 @@ class AsMsgHandler: def __init__(self, token_counter: HuggingFaceTokenCounter): self._token_counter = token_counter - def count_str_token(self, text: str) -> int: + async def count_str_token(self, text: str) -> int: """Count tokens in a string. Args: @@ -30,19 +30,19 @@ class AsMsgHandler: return 0 try: - token_ids = self._token_counter.tokenizer.encode(text) - token_count = len(token_ids) + 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 = len(text.encode("utf-8")) // 4 + estimated_tokens = int(len(text.encode("utf-8")) / 3.75) logger.warning(f"Failed to count string tokens: {text}, e={e}") return estimated_tokens - def _format_tool_result_output(self, output: str | list[dict]) -> tuple[str, int]: + 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, self.count_str_token(output) + return output, await self.count_str_token(output) textual_parts = [] total_token_count = 0 @@ -59,7 +59,7 @@ class AsMsgHandler: if block_type == "text": textual_parts.append(block.get("text", "")) - total_token_count += self.count_str_token(textual_parts[-1]) + total_token_count += await self.count_str_token(textual_parts[-1]) elif block_type in ["image", "audio", "video"]: source = block.get("source", {}) @@ -68,14 +68,14 @@ class AsMsgHandler: total_token_count += len(data) // 4 if data else 10 else: url = source.get("url", "") - total_token_count += self.count_str_token(url) if url else 10 + 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 += self.count_str_token(file_path) + total_token_count += await self.count_str_token(file_path) else: logger.warning( @@ -92,7 +92,7 @@ class AsMsgHandler: return "\n".join(textual_parts), total_token_count - def stat_message(self, message: Msg) -> AsMsgStat: + async def stat_message(self, message: Msg) -> AsMsgStat: """Analyze a message and generate block statistics.""" blocks = [] if isinstance(message.content, str): @@ -100,7 +100,7 @@ class AsMsgHandler: AsBlockStat( block_type="text", text=message.content, - token_count=self.count_str_token(message.content), + token_count=await self.count_str_token(message.content), ), ) return AsMsgStat( @@ -111,25 +111,12 @@ class AsMsgHandler: metadata=message.metadata or {}, ) - if not isinstance(message.content, list): - logger.warning( - "Unexpected message.content type %s, expected str or list, returning empty stat.", - type(message.content), - ) - return AsMsgStat( - name=message.name or message.role, - role=message.role, - content=blocks, - timestamp=message.timestamp or "", - metadata=message.metadata or {}, - ) - for block in message.content: block_type = block.get("type", "unknown") if block_type == "text": text = block.get("text", "") - token_count = self.count_str_token(text) + token_count = await self.count_str_token(text) blocks.append( AsBlockStat( block_type=block_type, @@ -140,7 +127,7 @@ class AsMsgHandler: elif block_type == "thinking": thinking = block.get("thinking", "") - token_count = self.count_str_token(thinking) + token_count = await self.count_str_token(thinking) blocks.append( AsBlockStat( block_type=block_type, @@ -156,7 +143,7 @@ class AsMsgHandler: data = source.get("data", "") token_count = len(data) // 4 if data else 10 else: - token_count = self.count_str_token(url) if url else 10 + token_count = await self.count_str_token(url) if url else 10 blocks.append( AsBlockStat( block_type=block_type, @@ -173,7 +160,7 @@ class AsMsgHandler: input_str = json.dumps(tool_input, ensure_ascii=False) except (TypeError, ValueError): input_str = str(tool_input) - token_count = self.count_str_token(tool_name + input_str) + token_count = await self.count_str_token(tool_name + input_str) blocks.append( AsBlockStat( block_type=block_type, @@ -187,7 +174,7 @@ class AsMsgHandler: elif block_type == "tool_result": tool_name = block.get("name", "") output = block.get("output", "") - formatted_output, token_count = self._format_tool_result_output(output) + formatted_output, token_count = await self._format_tool_result_output(output) blocks.append( AsBlockStat( block_type=block_type, @@ -209,11 +196,15 @@ class AsMsgHandler: metadata=message.metadata or {}, ) - def count_msgs_token(self, messages: list[Msg]) -> int: + async def count_msgs_token(self, messages: list[Msg]) -> int: """Count total token count of a list of messages.""" - return sum(self.stat_message(msg).total_tokens for msg in messages) + total = 0 + for msg in messages: + stat = await self.stat_message(msg) + total += stat.total_tokens + return total - def format_msgs_to_str( + async def format_msgs_to_str( self, messages: list[Msg], memory_compact_threshold: int, @@ -236,9 +227,9 @@ class AsMsgHandler: total_token_count = 0 for i in range(len(messages) - 1, -1, -1): - stat = self.stat_message(messages[i]) + stat = await self.stat_message(messages[i]) formatted_content = stat.format(include_thinking=include_thinking) - content_token_count = self.count_str_token(formatted_content) + 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: @@ -286,7 +277,7 @@ class AsMsgHandler: return tool_use_ids == tool_result_ids - def context_check( + async def context_check( self, messages: list[Msg], memory_compact_threshold: int, @@ -315,7 +306,7 @@ class AsMsgHandler: msg_stats: list[tuple[Msg, AsMsgStat]] = [] total_tokens = 0 for msg in messages: - stat = self.stat_message(msg) + stat = await self.stat_message(msg) msg_stats.append((msg, stat)) total_tokens += stat.total_tokens diff --git a/reme/reme_light.py b/reme/reme_light.py index 94309eb0..c41320c3 100644 --- a/reme/reme_light.py +++ b/reme/reme_light.py @@ -25,7 +25,7 @@ from agentscope.tool import Toolkit, ToolResponse from .config import ReMeConfigParser from .core import Application -from .core.utils import get_hf_token_counter, get_std_logger +from .core.utils import get_logger from .memory.file_based import ReMeInMemoryMemory from .memory.file_based.components import ( Compactor, @@ -36,7 +36,7 @@ from .memory.file_based.components import ( from .memory.file_based.tools import FileIO, MemorySearch from .memory.file_based.utils import AsMsgHandler -logger = get_std_logger() +logger = get_logger() class ReMeLight(Application): @@ -58,6 +58,7 @@ class ReMeLight(Application): 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. tool_result_threshold (int): Character threshold for tool result compaction. @@ -120,6 +121,7 @@ class ReMeLight(Application): - {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 self.working_path = Path(working_dir).absolute() @@ -128,6 +130,8 @@ class ReMeLight(Application): self.memory_path.mkdir(parents=True, exist_ok=True) 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) self.vector_weight: float = vector_weight self.candidate_multiplier: float = candidate_multiplier @@ -283,7 +287,7 @@ class ReMeLight(Application): messages: list[Msg], memory_compact_threshold: int, memory_compact_reserve: int = 10000, - token_counter: HuggingFaceTokenCounter | None = None, + as_token_counter: str | HuggingFaceTokenCounter = "default", ) -> tuple[list[Msg], list[Msg], bool]: """ Check context size and determine if compaction is needed. @@ -298,8 +302,7 @@ class ReMeLight(Application): 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. - token_counter (HuggingFaceTokenCounter | None): Token counter for - measuring message length. If None, uses default HuggingFace counter. + as_token_counter (str | HuggingFaceTokenCounter): The token counter to use. Returns: tuple[list[Msg], list[Msg], bool]: A tuple containing: @@ -315,13 +318,10 @@ class ReMeLight(Application): - is_valid=False indicates tool_use and tool_result are misaligned. """ try: - if token_counter is None: - token_counter = get_hf_token_counter() - checker = ContextChecker( memory_compact_threshold=memory_compact_threshold, memory_compact_reserve=memory_compact_reserve, - token_counter=token_counter, + as_token_counter=as_token_counter, ) return await checker.call( @@ -338,7 +338,7 @@ class ReMeLight(Application): messages: list[Msg], as_llm: str | ChatModelBase = "default", as_llm_formatter: str | FormatterBase = "default", - token_counter: HuggingFaceTokenCounter | None = None, + as_token_counter: str | HuggingFaceTokenCounter = "default", language: str = "zh", max_input_length: float = 128 * 1024, compact_ratio: float = 0.7, @@ -357,8 +357,8 @@ class ReMeLight(Application): to use for summarization. Defaults to "default". as_llm_formatter (str | FormatterBase): Formatter for the language model. Defaults to "default". - token_counter (HuggingFaceTokenCounter | None): Token counter for - measuring message length. If None, uses default HuggingFace counter. + 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. @@ -373,14 +373,11 @@ class ReMeLight(Application): an error occurred during compaction. """ try: - if token_counter is None: - token_counter = get_hf_token_counter() - compactor = Compactor( memory_compact_threshold=self.calculate_memory_compact_threshold(max_input_length, compact_ratio), - token_counter=token_counter, as_llm=as_llm, as_llm_formatter=as_llm_formatter, + as_token_counter=as_token_counter, language=language if language == "zh" else "", ) @@ -400,7 +397,7 @@ class ReMeLight(Application): messages: list[Msg], as_llm: str | ChatModelBase = "default", as_llm_formatter: str | FormatterBase = "default", - token_counter: HuggingFaceTokenCounter | None = None, + as_token_counter: str | HuggingFaceTokenCounter = "default", toolkit: Toolkit | None = None, language: str = "zh", max_input_length: float = 128 * 1024, @@ -419,8 +416,8 @@ class ReMeLight(Application): for summarization. Defaults to "default". as_llm_formatter (str | FormatterBase): Formatter for the language model. Defaults to "default". - token_counter (HuggingFaceTokenCounter | None): Token counter for - measuring message length. If None, uses default HuggingFace counter. + 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, @@ -438,9 +435,6 @@ class ReMeLight(Application): using the provided or default toolkit. """ try: - if token_counter is None: - token_counter = get_hf_token_counter() - if toolkit is None: toolkit = Toolkit() file_io = FileIO(working_dir=str(self.working_path)) @@ -452,10 +446,10 @@ class ReMeLight(Application): working_dir=str(self.working_path), memory_dir=str(self.memory_path), memory_compact_threshold=self.calculate_memory_compact_threshold(max_input_length, compact_ratio), - token_counter=token_counter, toolkit=toolkit, as_llm=as_llm, as_llm_formatter=as_llm_formatter, + as_token_counter=as_token_counter, language=language if language == "zh" else "", ) @@ -479,7 +473,7 @@ class ReMeLight(Application): Supported arguments include: - as_llm: Language model identifier or instance - as_llm_formatter: Formatter for the language model - - token_counter: Token counter instance + - as_token_counter: Token counter instance - toolkit: Toolkit for file operations - language: Output language ("zh" or other) - max_input_length: Maximum input token length @@ -509,6 +503,16 @@ class ReMeLight(Application): 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], @@ -516,7 +520,7 @@ class ReMeLight(Application): compressed_summary: str = "", as_llm: str | ChatModelBase = "default", as_llm_formatter: str | FormatterBase = "default", - token_counter: HuggingFaceTokenCounter | None = None, + as_token_counter: str | HuggingFaceTokenCounter = "default", toolkit: Toolkit | None = None, language: str = "zh", max_input_length: float = 128 * 1024, @@ -542,8 +546,8 @@ class ReMeLight(Application): Defaults to "default". as_llm_formatter (str | FormatterBase): Formatter for the language model. Defaults to "default". - token_counter (HuggingFaceTokenCounter | None): Token counter for - measuring content length. If None, uses default counter. + 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". @@ -568,13 +572,10 @@ class ReMeLight(Application): - Tool results in recent messages (keep_n) are not compacted - Returns original messages unchanged if no compaction is needed """ - if token_counter is None: - token_counter = get_hf_token_counter() + msg_handler = AsMsgHandler(self.default_as_token_counter) - msg_handler = AsMsgHandler(token_counter=token_counter) - - system_token_count = msg_handler.count_str_token(system_prompt) - compressed_token_count = msg_handler.count_str_token(compressed_summary) + 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}") @@ -587,7 +588,7 @@ class ReMeLight(Application): messages=messages, memory_compact_threshold=left_compact_threshold, memory_compact_reserve=memory_compact_reserve, - token_counter=token_counter, + as_token_counter=as_token_counter, ) if not messages_to_compact: @@ -601,7 +602,7 @@ class ReMeLight(Application): messages=messages_to_compact, as_llm=as_llm, as_llm_formatter=as_llm_formatter, - token_counter=token_counter, + as_token_counter=as_token_counter, toolkit=toolkit, language=language, max_input_length=max_input_length, @@ -612,7 +613,7 @@ class ReMeLight(Application): messages=messages_to_compact, as_llm=as_llm, as_llm_formatter=as_llm_formatter, - token_counter=token_counter, + as_token_counter=as_token_counter, language=language, max_input_length=max_input_length, compact_ratio=compact_ratio, @@ -755,28 +756,31 @@ class ReMeLight(Application): ], ) - @staticmethod - def get_in_memory_memory(token_counter: HuggingFaceTokenCounter | None = None): + def get_in_memory_memory(self, as_token_counter: HuggingFaceTokenCounter | None = None): """ Create and return an in-memory memory instance. Factory method to create a ReMeInMemoryMemory instance configured with - the specified token counter. This memory instance stores data in RAM - without persistence, suitable for temporary or session-based storage. + 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: - token_counter (HuggingFaceTokenCounter | None): Token counter for - measuring content length in the memory. If None, creates a - default HuggingFace token counter. + as_token_counter (HuggingFaceTokenCounter): Token counter for + measuring content length in the memory. Returns: ReMeInMemoryMemory: A new in-memory memory instance ready for use. + The instance is configured with self.dialog_path for persistence. - Example: - >>> memory = ReMeLight.get_in_memory_memory() - >>> # Use memory for temporary storage during a session + Note: + - 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 """ - if token_counter is None: - token_counter = get_hf_token_counter() - - return ReMeInMemoryMemory(token_counter=token_counter) + return ReMeInMemoryMemory( + token_counter=as_token_counter or self.default_as_token_counter, + dialog_path=str(self.dialog_path), + ) diff --git a/tests/light/test_compactor.py b/tests/light/test_compactor.py index cb3c9dee..74cba241 100644 --- a/tests/light/test_compactor.py +++ b/tests/light/test_compactor.py @@ -9,10 +9,10 @@ from test_utils import ( get_token_counter, ) -from reme.core.utils import get_std_logger +from reme.core.utils import get_logger from reme.memory.file_based.components import Compactor -logger = get_std_logger() +logger = get_logger() # ANSI 颜色码 @@ -96,7 +96,7 @@ def create_compactor(): """Create a Compactor instance for testing.""" return Compactor( memory_compact_threshold=4000, - token_counter=get_token_counter(), + as_token_counter=get_token_counter(), as_llm=get_dash_chat_model(), as_llm_formatter=get_formatter(), language="zh", @@ -282,7 +282,7 @@ def test_low_threshold(): """Test compaction with low memory threshold.""" compactor = Compactor( memory_compact_threshold=500, - token_counter=get_token_counter(), + as_token_counter=get_token_counter(), as_llm=get_dash_chat_model(), as_llm_formatter=get_formatter(), ) @@ -305,7 +305,7 @@ def test_high_threshold(): """Test compaction with high memory threshold.""" compactor = Compactor( memory_compact_threshold=10000, - token_counter=get_token_counter(), + as_token_counter=get_token_counter(), as_llm=get_dash_chat_model(), as_llm_formatter=get_formatter(), ) diff --git a/tests/light/test_context_check.py b/tests/light/test_context_check.py index 934a43f2..e31ea278 100644 --- a/tests/light/test_context_check.py +++ b/tests/light/test_context_check.py @@ -1,12 +1,14 @@ """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_std_logger +from reme.core.utils import get_logger from reme.memory.file_based.utils import AsMsgHandler -logger = get_std_logger() +logger = get_logger() # ANSI color codes @@ -78,7 +80,7 @@ def verify_context_check_invariants( AssertionError: If any invariant is violated """ # Calculate total tokens of original messages - total_tokens = sum(handler.stat_message(m).total_tokens for m in 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: @@ -93,7 +95,7 @@ def verify_context_check_invariants( ) # 2. Reserve requirement check - kept_tokens = sum(handler.stat_message(m).total_tokens for m in to_keep) + 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})" ) @@ -220,10 +222,12 @@ def test_empty_messages(): handler = create_handler() messages = [] threshold, reserve = 1000, 500 - to_compact, to_keep, _ = handler.context_check( - messages=messages, - memory_compact_threshold=threshold, - memory_compact_reserve=reserve, + 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}" @@ -240,10 +244,12 @@ def test_below_threshold_returns_all(): create_user_msg("How are you?"), ] threshold, reserve = 10000, 5000 - to_compact, to_keep, _ = handler.context_check( - messages=messages, - memory_compact_threshold=threshold, # Very high threshold - memory_compact_reserve=reserve, + 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)}" @@ -271,10 +277,12 @@ def test_above_threshold_triggers_compaction(): create_assistant_msg("Fourth message " * 100), ] threshold, reserve = 100, 200 - to_compact, to_keep, _ = handler.context_check( - messages=messages, - memory_compact_threshold=threshold, # Low threshold to trigger compaction - memory_compact_reserve=reserve, + 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" @@ -302,10 +310,12 @@ def test_message_order_preserved(): create_user_msg("Fifth " * 10), ] threshold, reserve = 100, 150 - to_compact, to_keep, _ = handler.context_check( - messages=messages, - memory_compact_threshold=threshold, # Low threshold - memory_compact_reserve=reserve, + 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 @@ -333,10 +343,12 @@ def test_single_message_below_threshold(): handler = create_handler() messages = [create_user_msg("Short message")] threshold, reserve = 1000, 500 - to_compact, to_keep, _ = handler.context_check( - messages=messages, - memory_compact_threshold=threshold, - memory_compact_reserve=reserve, + 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" @@ -358,10 +370,12 @@ def test_single_message_above_threshold(): long_content = "Very long message " * 1000 messages = [create_user_msg(long_content)] threshold, reserve = 10, 5 - to_compact, to_keep, _ = handler.context_check( - messages=messages, - memory_compact_threshold=threshold, # Very low threshold - memory_compact_reserve=reserve, # Even lower reserve + 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" @@ -386,10 +400,12 @@ def test_reserve_zero(): create_assistant_msg("Hi there!"), ] threshold, reserve = 1, 0 - to_compact, to_keep, _ = handler.context_check( - messages=messages, - memory_compact_threshold=threshold, # Trigger compaction - memory_compact_reserve=reserve, # Zero reserve + 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)}" @@ -403,10 +419,12 @@ def test_threshold_zero(): handler = create_handler() messages = [create_user_msg("A")] # Minimal message threshold, reserve = 0, 1000 - to_compact, to_keep, _ = handler.context_check( - messages=messages, - memory_compact_threshold=threshold, # Zero threshold - always triggers - memory_compact_reserve=reserve, + 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 @@ -421,15 +439,17 @@ def test_exact_threshold_boundary(): messages = [create_user_msg("Test message")] # Get exact token count - stat = handler.stat_message(messages[0]) + 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, _ = handler.context_check( - messages=messages, - memory_compact_threshold=threshold, # Exactly at boundary - memory_compact_reserve=reserve, + 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" @@ -454,10 +474,12 @@ def test_reserve_larger_than_threshold(): create_assistant_msg("Message two " * 20), ] threshold, reserve = 50, 10000 - to_compact, to_keep, _ = handler.context_check( - messages=messages, - memory_compact_threshold=threshold, # Low threshold - memory_compact_reserve=reserve, # High reserve + 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 @@ -489,10 +511,12 @@ def test_tool_use_result_paired(): create_assistant_msg("The tool returned results"), ] threshold, reserve = 50, 1000 - to_compact, to_keep, _ = handler.context_check( - messages=messages, - memory_compact_threshold=threshold, # Trigger compaction - memory_compact_reserve=reserve, # Enough for tool pair + 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 @@ -522,10 +546,12 @@ def test_tool_use_without_result(): create_assistant_msg("Something happened"), ] threshold, reserve = 10, 1000 - to_compact, to_keep, _ = handler.context_check( - messages=messages, - memory_compact_threshold=threshold, - memory_compact_reserve=reserve, + 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 @@ -550,10 +576,12 @@ def test_tool_result_without_use(): create_assistant_msg("Got it"), ] threshold, reserve = 10, 1000 - to_compact, to_keep, _ = handler.context_check( - messages=messages, - memory_compact_threshold=threshold, - memory_compact_reserve=reserve, + 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 @@ -583,10 +611,12 @@ def test_multiple_tool_pairs(): create_assistant_msg("All done"), ] threshold, reserve = 50, 500 - to_compact, to_keep, _ = handler.context_check( - messages=messages, - memory_compact_threshold=threshold, - memory_compact_reserve=reserve, + 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 @@ -630,10 +660,12 @@ def test_tool_dependency_causes_extra_inclusion(): create_assistant_msg("End"), # Small ] threshold, reserve = 100, 500 - to_compact, to_keep, _ = handler.context_check( - messages=messages, - memory_compact_threshold=threshold, # Trigger compaction - memory_compact_reserve=reserve, # Medium reserve + 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 @@ -671,10 +703,12 @@ def test_tool_dependency_exceeds_reserve(): create_assistant_msg("Last message"), ] threshold, reserve = 10, 100 - to_compact, to_keep, _ = handler.context_check( - messages=messages, - memory_compact_threshold=threshold, # Trigger compaction - memory_compact_reserve=reserve, # Small reserve - can't fit the pair + 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 @@ -715,10 +749,12 @@ def test_interleaved_tool_pairs(): create_assistant_msg("Both done"), ] threshold, reserve = 50, 1000 - to_compact, to_keep, _ = handler.context_check( - messages=messages, - memory_compact_threshold=threshold, - memory_compact_reserve=reserve, + 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 @@ -756,10 +792,12 @@ def test_message_with_empty_content(): create_assistant_msg("Response"), ] threshold, reserve = 1000, 500 - to_compact, to_keep, _ = handler.context_check( - messages=messages, - memory_compact_threshold=threshold, - memory_compact_reserve=reserve, + 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( @@ -782,10 +820,12 @@ def test_message_with_whitespace_only(): create_assistant_msg("Response"), ] threshold, reserve = 1000, 500 - to_compact, to_keep, _ = handler.context_check( - messages=messages, - memory_compact_threshold=threshold, - memory_compact_reserve=reserve, + 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( @@ -806,10 +846,12 @@ def test_very_long_single_message(): huge_content = "x" * 100000 # Very long messages = [create_user_msg(huge_content)] threshold, reserve = 100, 1000 - to_compact, to_keep, _ = handler.context_check( - messages=messages, - memory_compact_threshold=threshold, - memory_compact_reserve=reserve, + 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 @@ -830,10 +872,12 @@ def 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, _ = handler.context_check( - messages=messages, - memory_compact_threshold=threshold, # Low threshold - memory_compact_reserve=reserve, + 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 @@ -859,10 +903,12 @@ def test_unicode_content(): create_user_msg("日本語テスト 🇯🇵"), ] threshold, reserve = 1000, 500 - to_compact, to_keep, _ = handler.context_check( - messages=messages, - memory_compact_threshold=threshold, - memory_compact_reserve=reserve, + 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") @@ -877,10 +923,12 @@ def test_special_characters_content(): create_assistant_msg("More: \n\r\t\0 nulls and newlines"), ] threshold, reserve = 1000, 500 - to_compact, to_keep, _ = handler.context_check( - messages=messages, - memory_compact_threshold=threshold, - memory_compact_reserve=reserve, + 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( @@ -909,13 +957,15 @@ def test_all_messages_fit_exactly_in_reserve(): ] # Calculate total tokens - total = sum(handler.stat_message(m).total_tokens for m in messages) + total = sum(asyncio.run(handler.stat_message(m)).total_tokens for m in messages) threshold, reserve = total - 1, total - to_compact, to_keep, _ = handler.context_check( - messages=messages, - memory_compact_threshold=threshold, # Just below total to trigger - memory_compact_reserve=reserve, # Exactly fits all + 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)}" @@ -941,14 +991,16 @@ def test_first_message_only_compacted(): ] # Calculate tokens to set appropriate reserve - small_msg_tokens = handler.stat_message(messages[1]).total_tokens - tiny_msg_tokens = handler.stat_message(messages[2]).total_tokens + 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, _ = handler.context_check( - messages=messages, - memory_compact_threshold=threshold, # Low to trigger - memory_compact_reserve=reserve, # Fits last 2 + 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" @@ -973,13 +1025,15 @@ def test_last_message_only_kept(): create_user_msg("Tiny"), # Only this fits ] - tiny_tokens = handler.stat_message(messages[2]).total_tokens + tiny_tokens = asyncio.run(handler.stat_message(messages[2])).total_tokens threshold, reserve = 10, tiny_tokens + 5 - to_compact, to_keep, _ = handler.context_check( - messages=messages, - memory_compact_threshold=threshold, - memory_compact_reserve=reserve, # Only fits last message + 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: @@ -1005,10 +1059,12 @@ def test_all_messages_compacted(): create_assistant_msg("Large message " * 100), ] threshold, reserve = 10, 1 - to_compact, to_keep, _ = handler.context_check( - messages=messages, - memory_compact_threshold=threshold, # Trigger compaction - memory_compact_reserve=reserve, # Too small for anything + 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" @@ -1039,10 +1095,12 @@ def test_system_message(): create_assistant_msg("Hi there!"), ] threshold, reserve = 1000, 500 - to_compact, to_keep, _ = handler.context_check( - messages=messages, - memory_compact_threshold=threshold, - memory_compact_reserve=reserve, + 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") @@ -1061,10 +1119,12 @@ def test_mixed_roles(): Msg(name="helper", role="assistant", content="Another assistant message"), ] threshold, reserve = 1000, 500 - to_compact, to_keep, _ = handler.context_check( - messages=messages, - memory_compact_threshold=threshold, - memory_compact_reserve=reserve, + 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") @@ -1085,10 +1145,12 @@ def test_tool_use_with_empty_id(): create_assistant_msg("Done"), ] threshold, reserve = 10, 1000 - to_compact, to_keep, _ = handler.context_check( - messages=messages, - memory_compact_threshold=threshold, - memory_compact_reserve=reserve, + 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 @@ -1113,10 +1175,12 @@ def test_tool_result_with_empty_id(): create_assistant_msg("Noted"), ] threshold, reserve = 10, 1000 - to_compact, to_keep, _ = handler.context_check( - messages=messages, - memory_compact_threshold=threshold, - memory_compact_reserve=reserve, + 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 @@ -1142,10 +1206,12 @@ def test_duplicate_tool_ids(): create_tool_result_msg("call_dup", "tool_b", "Result B"), ] threshold, reserve = 10, 1000 - to_compact, to_keep, _ = handler.context_check( - messages=messages, - memory_compact_threshold=threshold, - memory_compact_reserve=reserve, + 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 @@ -1181,10 +1247,12 @@ def test_message_with_multiple_tool_blocks(): create_tool_result_msg("call_3", "tool3", "Result 3"), ] threshold, reserve = 10, 2000 - to_compact, to_keep, _ = handler.context_check( - messages=messages, - memory_compact_threshold=threshold, - memory_compact_reserve=reserve, + 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( diff --git a/tests/light/test_format_msgs_to_str.py b/tests/light/test_format_msgs_to_str.py index 64e97330..978a89f4 100644 --- a/tests/light/test_format_msgs_to_str.py +++ b/tests/light/test_format_msgs_to_str.py @@ -2,15 +2,16 @@ # 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_std_logger +from reme.core.utils import get_logger from reme.memory.file_based.utils import AsMsgHandler -logger = get_std_logger() +logger = get_logger() # ANSI 颜色码 @@ -87,7 +88,7 @@ def verify_result_within_threshold( # Calculate tokens of messages that were included in the result included_tokens = 0 for msg in msgs: - stat = handler.stat_message(msg) + 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 @@ -98,10 +99,10 @@ def verify_result_within_threshold( if block_type == "text" and block.get("text", "") in result: msg_included = True break - if block_type == "tool_use" and f"tool_call={block.get('name', '')}" in result: + if block_type == "tool_use" and f"{block.get('name', '')}" in result: msg_included = True break - if block_type == "tool_result" and f"tool_result={block.get('name', '')}" in result: + if block_type == "tool_result" and f"{block.get('name', '')}" in result: msg_included = True break @@ -216,7 +217,7 @@ def test_format_msgs_to_str_empty_list(): handler = create_handler() threshold = 4000 msgs = [] - result = handler.format_msgs_to_str(msgs, memory_compact_threshold=threshold) + 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") @@ -227,7 +228,7 @@ def test_format_msgs_to_str_single_message(): handler = create_handler() threshold = 4000 msgs = [create_user_msg("Hello, how are you?")] - result = handler.format_msgs_to_str(msgs, memory_compact_threshold=threshold) + 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}" @@ -245,7 +246,7 @@ def test_format_msgs_to_str_multiple_messages(): create_user_msg("Tell me more."), create_assistant_msg("Python is known for its readability."), ] - result = handler.format_msgs_to_str(msgs, memory_compact_threshold=threshold) + 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 @@ -264,7 +265,7 @@ def test_format_msgs_to_str_message_order(): create_assistant_msg("Second message"), create_user_msg("Third message"), ] - result = handler.format_msgs_to_str(msgs, memory_compact_threshold=threshold) + result = asyncio.run(handler.format_msgs_to_str(msgs, memory_compact_threshold=threshold)) # Find positions of each message first_pos = result.find("First message") @@ -283,9 +284,9 @@ def test_format_msgs_to_str_with_tool_use(): handler = create_handler() threshold = 4000 msgs = [create_tool_use_msg("read_file", {"path": "/test.txt"})] - result = handler.format_msgs_to_str(msgs, memory_compact_threshold=threshold) + result = asyncio.run(handler.format_msgs_to_str(msgs, memory_compact_threshold=threshold)) - assert "tool_call=read_file" in result, f"Expected tool_call in result, got: {result}" + 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") @@ -295,9 +296,9 @@ def test_format_msgs_to_str_with_tool_result(): handler = create_handler() threshold = 4000 msgs = [create_tool_result_msg("read_file", "file content here")] - result = handler.format_msgs_to_str(msgs, memory_compact_threshold=threshold) + result = asyncio.run(handler.format_msgs_to_str(msgs, memory_compact_threshold=threshold)) - assert "tool_result=read_file" in result, f"Expected tool_result in result, got: {result}" + 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") @@ -307,9 +308,9 @@ def test_format_msgs_to_str_with_image(): handler = create_handler() threshold = 4000 msgs = [create_image_msg("https://example.com/image.png")] - result = handler.format_msgs_to_str(msgs, memory_compact_threshold=threshold) + result = asyncio.run(handler.format_msgs_to_str(msgs, memory_compact_threshold=threshold)) - assert "[image]" in result, f"Expected '[image]' in result, got: {result}" + 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") @@ -324,11 +325,11 @@ def test_format_msgs_to_str_conversation_flow(): create_tool_result_msg("read_file", "File content here"), create_assistant_msg("The file contains: File content here"), ] - result = handler.format_msgs_to_str(msgs, memory_compact_threshold=threshold) + result = asyncio.run(handler.format_msgs_to_str(msgs, memory_compact_threshold=threshold)) assert "user:" in result - assert "tool_call=read_file" in result - assert "tool_result=read_file" 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") @@ -342,7 +343,7 @@ def test_format_msgs_to_str_thinking_excluded_by_default(): handler = create_handler() threshold = 4000 msgs = [create_thinking_msg("Let me think about this...", "Here is my response")] - result = handler.format_msgs_to_str(msgs, memory_compact_threshold=threshold, include_thinking=False) + 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}" @@ -355,7 +356,7 @@ def test_format_msgs_to_str_thinking_included(): handler = create_handler() threshold = 4000 msgs = [create_thinking_msg("Let me think about this...", "Here is my response")] - result = handler.format_msgs_to_str(msgs, memory_compact_threshold=threshold, include_thinking=True) + 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}" @@ -370,16 +371,20 @@ def test_format_msgs_to_str_thinking_only_message(): msgs = [create_thinking_msg("Deep thoughts here")] # With include_thinking=False - result_no_thinking = handler.format_msgs_to_str( - msgs, - memory_compact_threshold=threshold, - 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 = handler.format_msgs_to_str( - msgs, - memory_compact_threshold=threshold, - 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 @@ -401,7 +406,7 @@ def test_format_msgs_to_str_all_within_threshold(): create_assistant_msg("Short message 2"), create_user_msg("Short message 3"), ] - result = handler.format_msgs_to_str(msgs, memory_compact_threshold=threshold) + 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 @@ -419,7 +424,7 @@ def test_format_msgs_to_str_exceeds_threshold_truncate_older(): msgs.append(create_user_msg(f"Question {i}: " + "x" * 100)) msgs.append(create_assistant_msg(f"Answer {i}: " + "y" * 100)) - result = handler.format_msgs_to_str(msgs, memory_compact_threshold=threshold) + result = asyncio.run(handler.format_msgs_to_str(msgs, memory_compact_threshold=threshold)) # The newest messages should be present assert ( @@ -440,7 +445,7 @@ def test_format_msgs_to_str_single_message_exceeds_threshold(): msgs = [create_user_msg(long_text)] # With very low threshold, even a single message won't fit - result = handler.format_msgs_to_str(msgs, memory_compact_threshold=threshold) + 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" @@ -457,7 +462,7 @@ def test_format_msgs_to_str_first_message_exceeds_threshold(): create_assistant_msg("Short response"), # New, short message ] - result = handler.format_msgs_to_str(msgs, memory_compact_threshold=threshold) + 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}" @@ -466,15 +471,16 @@ def test_format_msgs_to_str_first_message_exceeds_threshold(): def test_format_msgs_to_str_threshold_zero(): - """Test with threshold of zero - no messages should be included.""" + """Test with threshold of zero - latest message is still included.""" handler = create_handler() threshold = 0 msgs = [create_user_msg("Test message")] - result = handler.format_msgs_to_str(msgs, memory_compact_threshold=threshold) + result = asyncio.run(handler.format_msgs_to_str(msgs, memory_compact_threshold=threshold)) - assert result == "", f"Expected empty string with zero threshold, got: {result}" - verify_result_within_threshold(handler, result, threshold, "threshold_zero", msgs) + # 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") @@ -483,12 +489,12 @@ def test_format_msgs_to_str_threshold_exact_fit(): handler = create_handler() # Create a message and measure its formatted string tokens msg = create_user_msg("Test") - stat = handler.stat_message(msg) + stat = asyncio.run(handler.stat_message(msg)) formatted_content = stat.format(include_thinking=False) - exact_threshold = handler.count_str_token(formatted_content) + exact_threshold = asyncio.run(handler.count_str_token(formatted_content)) msgs = [msg] - result = handler.format_msgs_to_str(msgs, memory_compact_threshold=exact_threshold) + 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) @@ -496,19 +502,19 @@ def 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.""" + """Test when threshold is one less than needed - latest message is still included.""" handler = create_handler() msg = create_user_msg("Test message") - stat = handler.stat_message(msg) + stat = asyncio.run(handler.stat_message(msg)) formatted_content = stat.format(include_thinking=False) - threshold_minus_one = handler.count_str_token(formatted_content) - 1 + threshold_minus_one = asyncio.run(handler.count_str_token(formatted_content)) - 1 msgs = [msg] - result = handler.format_msgs_to_str(msgs, memory_compact_threshold=threshold_minus_one) + result = asyncio.run(handler.format_msgs_to_str(msgs, memory_compact_threshold=threshold_minus_one)) - # Message should be skipped since it doesn't fit - assert result == "", f"Expected empty string when threshold is insufficient, got: {result}" - verify_result_within_threshold(handler, result, threshold_minus_one, "threshold_one_less", msgs) + # 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") @@ -518,7 +524,7 @@ def test_format_msgs_to_str_large_threshold(): threshold = 1000000 msgs = [create_user_msg("Message " + str(i) + " " + "x" * 100) for i in range(50)] - result = handler.format_msgs_to_str(msgs, memory_compact_threshold=threshold) + result = asyncio.run(handler.format_msgs_to_str(msgs, memory_compact_threshold=threshold)) # All messages should be included for i in range(50): @@ -535,7 +541,7 @@ def test_format_msgs_to_str_special_characters(): handler = create_handler() threshold = 4000 msgs = [create_user_msg("Test with 中文, 日本語, émojis 🎉 and symbols @#$%")] - result = handler.format_msgs_to_str(msgs, memory_compact_threshold=threshold) + result = asyncio.run(handler.format_msgs_to_str(msgs, memory_compact_threshold=threshold)) assert "中文" in result assert "日本語" in result @@ -549,7 +555,7 @@ def test_format_msgs_to_str_empty_content(): handler = create_handler() threshold = 4000 msgs = [create_user_msg("")] - result = handler.format_msgs_to_str(msgs, memory_compact_threshold=threshold) + 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) @@ -561,7 +567,7 @@ def test_format_msgs_to_str_whitespace_only(): handler = create_handler() threshold = 4000 msgs = [create_user_msg(" \n\t ")] - result = handler.format_msgs_to_str(msgs, memory_compact_threshold=threshold) + 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) @@ -573,7 +579,7 @@ def test_format_msgs_to_str_newlines_in_content(): handler = create_handler() threshold = 4000 msgs = [create_user_msg("Line 1\nLine 2\nLine 3")] - result = handler.format_msgs_to_str(msgs, memory_compact_threshold=threshold) + result = asyncio.run(handler.format_msgs_to_str(msgs, memory_compact_threshold=threshold)) assert "Line 1" in result assert "Line 2" in result @@ -588,7 +594,7 @@ def test_format_msgs_to_str_very_long_single_word(): threshold = 10000 long_word = "a" * 5000 msgs = [create_user_msg(long_word)] - result = handler.format_msgs_to_str(msgs, memory_compact_threshold=threshold) + 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]}..." @@ -610,20 +616,24 @@ def test_format_msgs_to_str_mixed_content_blocks(): ), ] - result_no_thinking = handler.format_msgs_to_str( - msgs, - memory_compact_threshold=threshold, - include_thinking=False, + result_no_thinking = asyncio.run( + handler.format_msgs_to_str( + msgs, + memory_compact_threshold=threshold, + include_thinking=False, + ), ) - result_with_thinking = handler.format_msgs_to_str( - msgs, - memory_compact_threshold=threshold, - include_thinking=True, + 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 "tool_call=test_tool" in result_no_thinking - assert "[image]" 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) @@ -639,7 +649,7 @@ def test_format_msgs_to_str_multiple_separators(): create_user_msg("Message 1"), create_assistant_msg("Message 2"), ] - result = handler.format_msgs_to_str(msgs, memory_compact_threshold=threshold) + 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) @@ -655,9 +665,9 @@ def test_format_msgs_to_str_tool_result_complex_output(): {"type": "image", "source": {"url": "https://example.com/result.png"}}, ] msgs = [create_tool_result_msg("process_data", complex_output)] - result = handler.format_msgs_to_str(msgs, memory_compact_threshold=threshold) + result = asyncio.run(handler.format_msgs_to_str(msgs, memory_compact_threshold=threshold)) - assert "tool_result=process_data" in result + 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") @@ -672,7 +682,7 @@ def test_format_msgs_to_str_different_roles(): create_assistant_msg("Assistant response"), create_tool_result_msg("tool", "Tool output"), ] - result = handler.format_msgs_to_str(msgs, memory_compact_threshold=threshold) + result = asyncio.run(handler.format_msgs_to_str(msgs, memory_compact_threshold=threshold)) assert "system:" in result assert "user:" in result @@ -691,11 +701,18 @@ def test_format_msgs_to_str_incremental_threshold_check(): msgs.append(create_user_msg(f"Message {i} with some padding text")) # Calculate total tokens - total_tokens = sum(handler.stat_message(msg).total_tokens for msg in msgs) + 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 = handler.format_msgs_to_str(msgs, memory_compact_threshold=half_threshold) + 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) @@ -707,16 +724,16 @@ def test_format_msgs_to_str_incremental_threshold_check(): def test_format_msgs_to_str_negative_threshold(): - """Test with negative threshold value.""" + """Test with negative threshold value - latest message is still included.""" handler = create_handler() threshold = -1 msgs = [create_user_msg("Test message")] - result = handler.format_msgs_to_str(msgs, memory_compact_threshold=threshold) + result = asyncio.run(handler.format_msgs_to_str(msgs, memory_compact_threshold=threshold)) - # Negative threshold should result in empty string (nothing fits) - assert result == "", f"Expected empty string with negative threshold, got: {result}" - verify_result_within_threshold(handler, result, max(0, threshold), "negative_threshold", msgs) + # 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") @@ -731,7 +748,7 @@ def test_format_msgs_to_str_preserves_newest_first(): ] # Use threshold that only allows ~1-2 messages - result = handler.format_msgs_to_str(msgs, memory_compact_threshold=threshold) + 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}" @@ -758,9 +775,9 @@ def test_format_msgs_to_str_base64_image(): ], ), ] - result = handler.format_msgs_to_str(msgs, memory_compact_threshold=threshold) + result = asyncio.run(handler.format_msgs_to_str(msgs, memory_compact_threshold=threshold)) - assert "[image]" in result + assert "" in result verify_result_within_threshold(handler, result, threshold, "base64_image", msgs) print_pass("test_format_msgs_to_str_base64_image") @@ -779,10 +796,10 @@ def test_format_msgs_to_str_audio_video_blocks(): ], ), ] - result = handler.format_msgs_to_str(msgs, memory_compact_threshold=threshold) + result = asyncio.run(handler.format_msgs_to_str(msgs, memory_compact_threshold=threshold)) - assert "[audio]" in result - assert "[video]" in result + assert " The above is a summary of our previous conversation. -Use it as context to maintain continuity. +If there is a new instruction from the user, do not continue executing the previous content; +only execute the user's new instruction. """.strip() return [ diff --git a/reme/memory/file_based/utils/as_msg_handler.py b/reme/memory/file_based/utils/as_msg_handler.py index 46bf81b6..0692ebce 100644 --- a/reme/memory/file_based/utils/as_msg_handler.py +++ b/reme/memory/file_based/utils/as_msg_handler.py @@ -50,8 +50,7 @@ class AsMsgHandler: try: if not isinstance(block, dict) or "type" not in block: logger.warning( - "Invalid block: %s, expected a dict with 'type' key, skipped.", - block, + f"Invalid block: {block}, expected a dict with 'type' key, skipped.", ) continue @@ -79,15 +78,12 @@ class AsMsgHandler: else: logger.warning( - "Unsupported block type '%s' in tool result, skipped.", - block_type, + f"Unsupported block type '{block_type}' in tool result, skipped.", ) except Exception as e: logger.warning( - "Failed to process block %s: %s, skipped.", - block, - e, + f"Failed to process block {block}: {e}, skipped.", ) return "\n".join(textual_parts), total_token_count @@ -186,7 +182,7 @@ class AsMsgHandler: ) else: - logger.warning("Unsupported block type %s, skipped.", block_type) + logger.warning(f"Unsupported block type {block_type}, skipped.") return AsMsgStat( name=message.name or message.role, @@ -234,18 +230,15 @@ class AsMsgHandler: is_latest = i == len(messages) - 1 if not is_latest and total_token_count + content_token_count > memory_compact_threshold: logger.info( - "Skipping older messages: adding %d tokens would exceed threshold %d (current: %d)", - content_token_count, - memory_compact_threshold, - total_token_count, + 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( - "Latest message alone (%d tokens) exceeds threshold %d, including it anyway.", - content_token_count, - memory_compact_threshold, + f"Latest message alone ({content_token_count} tokens) exceeds threshold " + f"{memory_compact_threshold}, including it anyway.", ) formatted_parts.append(formatted_content) @@ -345,11 +338,8 @@ class AsMsgHandler: # Check if adding this message would exceed reserve limit if accumulated_tokens + stat.total_tokens > memory_compact_reserve: logger.info( - "Context check: adding message %d with %d tokens would exceed reserve %d (current: %d)", - i, - stat.total_tokens, - memory_compact_reserve, - accumulated_tokens, + f"Context check: adding message {i} with {stat.total_tokens} tokens would exceed reserve " + f"{memory_compact_reserve} (current: {accumulated_tokens})", ) break @@ -374,11 +364,8 @@ class AsMsgHandler: # 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( - "Context check: message %d requires %d extra tokens for tool_use dependencies, " - "total would exceed reserve %d", - i, - extra_tokens, - memory_compact_reserve, + f"Context check: message {i} requires {extra_tokens} extra tokens for tool_use dependencies, " + f"total would exceed reserve {memory_compact_reserve}", ) break @@ -401,16 +388,11 @@ class AsMsgHandler: tools_aligned = self.validate_tool_ids_alignment(messages_to_keep) logger.info( - "Context check result: %d messages to compact, %d messages to keep, " - "total tokens: %d, threshold: %d, reserve: %d, kept tokens: %d, " - "tools_aligned: %s", - len(messages_to_compact), - len(messages_to_keep), - total_tokens, - memory_compact_threshold, - memory_compact_reserve, - accumulated_tokens, - tools_aligned, + 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 From ad392738ab20b25cfac518a95c2d2a83058086cc Mon Sep 17 00:00:00 2001 From: jinliyl <6469360+jinliyl@users.noreply.github.com> Date: Tue, 17 Mar 2026 20:13:19 +0800 Subject: [PATCH 28/59] feat(file-watcher): add clear-on-start option and remove redundant clears (#161) --- reme/__init__.py | 2 +- reme/core/file_watcher/base_file_watcher.py | 11 ++++++++++- reme/core/file_watcher/delta_file_watcher.py | 1 - reme/core/file_watcher/full_file_watcher.py | 1 - 4 files changed, 11 insertions(+), 4 deletions(-) diff --git a/reme/__init__.py b/reme/__init__.py index 32ed21d0..9bb1dd0d 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.6" +__version__ = "0.3.0.7" __all__ = [ "config", diff --git a/reme/core/file_watcher/base_file_watcher.py b/reme/core/file_watcher/base_file_watcher.py index 89c6e54a..61e8a147 100644 --- a/reme/core/file_watcher/base_file_watcher.py +++ b/reme/core/file_watcher/base_file_watcher.py @@ -34,7 +34,8 @@ class BaseFileWatcher: 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, + scan_on_start: bool = True, + clear_on_start: bool = True, **kwargs, ): """ @@ -50,6 +51,8 @@ class BaseFileWatcher: 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 + clear_on_start: If True, clear all indexed data on start before scanning. + Useful for full rebuild of the index. **kwargs: Additional keyword arguments """ self.watch_paths: list[str] = [watch_paths] if isinstance(watch_paths, str) else watch_paths @@ -61,6 +64,7 @@ class BaseFileWatcher: self.file_store: BaseFileStore = file_store self.callback = callback self.scan_on_start: bool = scan_on_start + self.clear_on_start: bool = clear_on_start self.kwargs: dict = kwargs self._stop_event = asyncio.Event() @@ -74,6 +78,11 @@ class BaseFileWatcher: self._running = True + # Clear all indexed data if requested + if self.clear_on_start and self.file_store is not None: + await self.file_store.clear_all() + logger.info("Cleared all indexed data on start") + # Scan existing files if requested if self.scan_on_start: await self._scan_existing_files() diff --git a/reme/core/file_watcher/delta_file_watcher.py b/reme/core/file_watcher/delta_file_watcher.py index f35c9b2f..6148bd07 100644 --- a/reme/core/file_watcher/delta_file_watcher.py +++ b/reme/core/file_watcher/delta_file_watcher.py @@ -141,7 +141,6 @@ class DeltaFileWatcher(BaseFileWatcher): async def _on_changes(self, changes: set[tuple[Change, str]]): """Handle file changes with incremental synchronization.""" self.dirty = True - await self.file_store.clear_all() for change_type, path in changes: if change_type == Change.added: diff --git a/reme/core/file_watcher/full_file_watcher.py b/reme/core/file_watcher/full_file_watcher.py index c49a94fa..5b1852b9 100644 --- a/reme/core/file_watcher/full_file_watcher.py +++ b/reme/core/file_watcher/full_file_watcher.py @@ -44,7 +44,6 @@ class FullFileWatcher(BaseFileWatcher): async def _on_changes(self, changes: set[tuple[Change, str]]): """Handle file changes with full synchronization""" self.dirty = True - await self.file_store.clear_all() for change_type, path in changes: if change_type in [Change.added, Change.modified]: From 8d7cc4bbd65b0278a0c48ba66578e85918303d66 Mon Sep 17 00:00:00 2001 From: "jinli.yl" Date: Wed, 18 Mar 2026 01:16:05 +0800 Subject: [PATCH 29/59] feat(core): add rule-based token counter and update default configuration --- reme/__init__.py | 2 +- reme/config/light.yaml | 5 +- reme/core/as_token_counter/__init__.py | 2 + .../as_token_counter/rule_token_counter.py | 78 +++++++++++++++++++ 4 files changed, 83 insertions(+), 4 deletions(-) create mode 100644 reme/core/as_token_counter/rule_token_counter.py diff --git a/reme/__init__.py b/reme/__init__.py index 9bb1dd0d..74463d84 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.7" +__version__ = "0.3.0.8" __all__ = [ "config", diff --git a/reme/config/light.yaml b/reme/config/light.yaml index 8b721596..e18a1302 100644 --- a/reme/config/light.yaml +++ b/reme/config/light.yaml @@ -9,9 +9,8 @@ as_llm_formatters: as_token_counters: default: - backend: hf - pretrained_model_name_or_path: "Qwen/Qwen2.5-7B-Instruct" - use_mirror: true + backend: rule + token_count_estimate_divisor: 3.75 embedding_models: default: diff --git a/reme/core/as_token_counter/__init__.py b/reme/core/as_token_counter/__init__.py index c8b017d2..95910b8e 100644 --- a/reme/core/as_token_counter/__init__.py +++ b/reme/core/as_token_counter/__init__.py @@ -1,6 +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/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) From d313ae6e52223b7b2ec33cf529a238eba6d4d85e Mon Sep 17 00:00:00 2001 From: jinliyl <6469360+jinliyl@users.noreply.github.com> Date: Wed, 18 Mar 2026 15:16:10 +0800 Subject: [PATCH 30/59] feat(core): update version and enhance configuration management (#164) --- reme/__init__.py | 2 +- reme/core/application.py | 24 ++++++++++++------- .../as_token_counter/reme_token_counter.py | 13 ++++++++-- reme/core/embedding/base_embedding_model.py | 15 ++---------- reme/core/llm/base_llm.py | 15 ++---------- reme/core/service_context.py | 22 +---------------- 6 files changed, 33 insertions(+), 58 deletions(-) diff --git a/reme/__init__.py b/reme/__init__.py index 74463d84..22374421 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.8" +__version__ = "0.3.0.9" __all__ = [ "config", diff --git a/reme/core/application.py b/reme/core/application.py index 8778dfba..f54a47e8 100644 --- a/reme/core/application.py +++ b/reme/core/application.py @@ -15,7 +15,7 @@ from .registry_factory import R from .schema import Response, ServiceConfig from .service_context import ServiceContext from .token_counter import BaseTokenCounter -from .utils import execute_stream_task, PydanticConfigParser, init_logger, MCPClient, print_logo, get_logger +from .utils import execute_stream_task, PydanticConfigParser, init_logger, MCPClient, print_logo, get_logger, load_env from .vector_store import BaseVectorStore logger = get_logger() @@ -46,12 +46,16 @@ class Application: default_file_watcher_config: dict | None = None, **kwargs, ): + + 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, @@ -158,11 +162,11 @@ class Application: else: config_dict = config.model_dump(exclude={"backend"}) if not config_dict.get("api_key", ""): - config_dict["api_key"] = os.getenv("LLM_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"] = os.getenv("LLM_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) for name, config in self.service_config.as_llm_formatters.items(): @@ -184,6 +188,8 @@ class Application: 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() @@ -192,7 +198,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() diff --git a/reme/core/as_token_counter/reme_token_counter.py b/reme/core/as_token_counter/reme_token_counter.py index 39ad8f6d..6722e92f 100644 --- a/reme/core/as_token_counter/reme_token_counter.py +++ b/reme/core/as_token_counter/reme_token_counter.py @@ -47,9 +47,18 @@ class ReMeTokenCounter(HuggingFaceTokenCounter): # Set HuggingFace endpoint for mirror support if use_mirror: - os.environ["HF_ENDPOINT"] = "https://hf-mirror.com" + mirror = "https://hf-mirror.com" else: - os.environ.pop("HF_ENDPOINT", None) + 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__( diff --git a/reme/core/embedding/base_embedding_model.py b/reme/core/embedding/base_embedding_model.py index 2b60e8e6..28efddce 100644 --- a/reme/core/embedding/base_embedding_model.py +++ b/reme/core/embedding/base_embedding_model.py @@ -6,7 +6,6 @@ Defines the abstract base class and standard API for all embedding model impleme import asyncio import hashlib import json -import os import time from abc import ABC from collections import OrderedDict @@ -56,8 +55,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 = api_key + self.base_url: str = base_url self.model_name = model_name self.dimensions = dimensions self.use_dimensions = use_dimensions @@ -78,16 +77,6 @@ class BaseEmbeddingModel(ABC): self.cache_path: Path = Path(self.cache_dir) self.cache_path.mkdir(parents=True, exist_ok=True) - @property - def api_key(self) -> str | None: - """Get API key from environment variable.""" - return os.getenv("EMBEDDING_API_KEY") or self._api_key - - @property - def base_url(self) -> str | None: - """Get base URL from environment variable.""" - return os.getenv("EMBEDDING_BASE_URL") or self._base_url - def _truncate_text(self, text: str) -> str: """Truncate text to max_input_length if it exceeds the limit.""" if len(text) > self.max_input_length: 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/service_context.py b/reme/core/service_context.py index 4ab83ff7..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,7 +7,7 @@ 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 @@ -29,10 +28,6 @@ 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, @@ -52,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) @@ -114,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.""" From 09ab707e98c0f794c87218ba42e73f8d596c38c4 Mon Sep 17 00:00:00 2001 From: jinliyl <6469360+jinliyl@users.noreply.github.com> Date: Thu, 19 Mar 2026 11:06:11 +0800 Subject: [PATCH 31/59] feat(core): add application restart capability with enhanced configuration options (#166) --- reme/core/application.py | 218 ++++++++++++++++++++++++++++++++++++++- reme/reme_light.py | 4 + 2 files changed, 220 insertions(+), 2 deletions(-) diff --git a/reme/core/application.py b/reme/core/application.py index f54a47e8..6f04e663 100644 --- a/reme/core/application.py +++ b/reme/core/application.py @@ -12,7 +12,16 @@ 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, get_logger, load_env @@ -35,6 +44,7 @@ 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, @@ -47,7 +57,8 @@ class Application: **kwargs, ): - load_env() + 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", "") @@ -255,6 +266,209 @@ class Application: 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 + + 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}") + + # 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 + 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}") + + # 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 + 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}") + + # 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}) diff --git a/reme/reme_light.py b/reme/reme_light.py index c41320c3..46b0ffe9 100644 --- a/reme/reme_light.py +++ b/reme/reme_light.py @@ -80,6 +80,7 @@ class ReMeLight(Application): candidate_multiplier: float = 3.0, tool_result_threshold: int = 1000, retention_days: int = 7, + enable_load_env: bool = False, ): """ Initialize the ReMeLight application. @@ -115,6 +116,8 @@ class ReMeLight(Application): saved to files. Default 1000 characters. retention_days (int): Number of days to retain tool result files before automatic cleanup. Default 7 days. + enable_load_env (bool): Whether to load environment variables from + .env file. Defaults to False. Note: The following directory structure will be created: @@ -148,6 +151,7 @@ class ReMeLight(Application): config_path="light", enable_logo=False, log_to_console=False, + enable_load_env=enable_load_env, parser=ReMeConfigParser, default_as_llm_config=default_as_llm_config, default_embedding_model_config=default_embedding_model_config, From 940a2f47a9dba2644f97a998769ae49dde2f55cc Mon Sep 17 00:00:00 2001 From: aquamarine Date: Thu, 19 Mar 2026 02:24:52 -0500 Subject: [PATCH 32/59] fix(reme): use timezone-aware datetime in memory summarization (#165) Use user-specified timezone instead of system local time when generating daily note filenames and timestamps in memory summarization and CLI. Changes: - Summarizer: accept 'timezone' param in __init__; use datetime.now(zoneinfo.ZoneInfo(tz)) instead of naive datetime.now() - ReMeLight.summary_memory(): accept 'timezone' param and pass to Summarizer - CliAgent: accept 'timezone' param in __init__; pass to Summarizer and use for current_time timestamp in system prompts - Fallback to system local time if timezone is None (preserves original behavior) Fixes timezone mismatch when system timezone differs from user's actual location (e.g., server in UTC+8 but user in America/Chicago). --- reme/memory/file_based/components/cli.py | 8 +++++++- reme/memory/file_based/components/summarizer.py | 11 ++++++++++- reme/reme_light.py | 4 ++++ 3 files changed, 21 insertions(+), 2 deletions(-) diff --git a/reme/memory/file_based/components/cli.py b/reme/memory/file_based/components/cli.py index f4dddc21..11d45035 100644 --- a/reme/memory/file_based/components/cli.py +++ b/reme/memory/file_based/components/cli.py @@ -3,6 +3,7 @@ import asyncio from datetime import datetime from pathlib import Path +import zoneinfo from agentscope.agent import ReActAgent from agentscope.message import Msg, TextBlock @@ -46,6 +47,7 @@ class CliAgent(BaseOp): reserve_tokens: int = 36000, keep_recent_tokens: int = 20000, language: str = "zh", + timezone: str | None = None, **kwargs, ): super().__init__(**kwargs) @@ -57,6 +59,7 @@ class CliAgent(BaseOp): 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] = [] @@ -93,6 +96,7 @@ class CliAgent(BaseOp): 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 @@ -168,6 +172,7 @@ class CliAgent(BaseOp): 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( @@ -195,7 +200,8 @@ class CliAgent(BaseOp): async def _build_messages(self, query: str) -> list[Msg]: """Build system prompt message.""" - current_time = datetime.now().strftime("%Y-%m-%d %H:%M:%S %A") + 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( diff --git a/reme/memory/file_based/components/summarizer.py b/reme/memory/file_based/components/summarizer.py index d553a5f6..855d64c1 100644 --- a/reme/memory/file_based/components/summarizer.py +++ b/reme/memory/file_based/components/summarizer.py @@ -1,6 +1,7 @@ """Summarizer module for memory summarization operations.""" import datetime +import zoneinfo from agentscope.agent import ReActAgent from agentscope.message import Msg @@ -23,6 +24,7 @@ class Summarizer(BaseOp): memory_compact_threshold: int, toolkit: Toolkit | None = None, console_enabled: bool = False, + timezone: str | None = None, **kwargs, ): super().__init__(**kwargs) @@ -31,6 +33,7 @@ class Summarizer(BaseOp): self.memory_compact_threshold: int = memory_compact_threshold self.toolkit: Toolkit | None = toolkit self.console_enabled: bool = console_enabled + self.timezone: str | None = timezone async def execute(self): messages: list[Msg] = self.context.get("messages", []) @@ -62,7 +65,13 @@ class Summarizer(BaseOp): user_message: str = f"\n{history_formatted_str}\n\n" + self.prompt_format( "user_message", - date=datetime.datetime.now().strftime("%Y-%m-%d"), + date=( + datetime.datetime.now( + zoneinfo.ZoneInfo(self.timezone), + ) + if self.timezone + else datetime.datetime.now() + ).strftime("%Y-%m-%d"), working_dir=self.working_dir, memory_dir=self.memory_dir, ) diff --git a/reme/reme_light.py b/reme/reme_light.py index 46b0ffe9..c8561550 100644 --- a/reme/reme_light.py +++ b/reme/reme_light.py @@ -406,6 +406,7 @@ class ReMeLight(Application): language: str = "zh", max_input_length: float = 128 * 1024, compact_ratio: float = 0.7, + timezone: str | None = None, ) -> str: """ Generate a comprehensive summary of the given messages. @@ -430,6 +431,8 @@ class ReMeLight(Application): 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: The generated summary text, or an empty string if an error occurred. @@ -455,6 +458,7 @@ class ReMeLight(Application): as_llm_formatter=as_llm_formatter, as_token_counter=as_token_counter, language=language if language == "zh" else "", + timezone=timezone, ) return await summarizer.call(messages=messages, service_context=self.service_context) From f86a3e1f573b62a6d7451af8e5f842f6997e6212 Mon Sep 17 00:00:00 2001 From: Sen Huang <48879559+ployts@users.noreply.github.com> Date: Thu, 19 Mar 2026 16:29:52 +0800 Subject: [PATCH 33/59] feat(memory): enhance summarizer to include experience reflections (#167) --- .../file_based/components/summarizer.yaml | 87 +++++++++++-------- 1 file changed, 49 insertions(+), 38 deletions(-) diff --git a/reme/memory/file_based/components/summarizer.yaml b/reme/memory/file_based/components/summarizer.yaml index 9aa0892b..00dbdc36 100644 --- a/reme/memory/file_based/components/summarizer.yaml +++ b/reme/memory/file_based/components/summarizer.yaml @@ -1,50 +1,61 @@ 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. + Memory Pre-compression Flush Cycle Initiated - Current date: {date} - Working directory: {working_dir} + The current session is about to enter the automatic compression phase. Please capture persistent memory AND session reflections, then write them to disk. - Immediately store persistent memory to: {memory_dir}/YYYY-MM-DD.md + Current date: {date} + Working directory: {working_dir} - 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 + Immediately store persistent memory and reflections to: {memory_dir}/YYYY-MM-DD.md - 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] + Workflow: + 1. First, `read` {memory_dir}/YYYY-MM-DD.md (if the file doesn’t exist, an error message will be returned). + 2. 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. + 3. Intelligently merge new information with existing content (skip merging if the file doesn’t exist): + - 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. + 4. 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. + - 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）。 + 当前日期：{date} + 工作目录：{working_dir} - 工作流程： - 1. 先 `read` {memory_dir}/YYYY-MM-DD.md（如文件不存在，会返回错误提示） - 2. 智能合并新信息与现有内容（若文件不存在则跳过合并）： - - 避免重复已记录的信息 - - 在相关时丰富现有条目的新细节 - - 在适用时保持时间顺序 - 3. 写入更新后的内容： - - 尽可能使用 `edit` 更新特定部分 - - 如需大幅重构则使用 `write` 覆盖整个文件 + 立即存储持久化记忆与反思（使用路径 {memory_dir}/YYYY-MM-DD.md）。 - 原则： - - 始终保留时间戳、日期和时间相关上下文 - - 仅添加真正新的或有丰富价值的信息 - - 保持条目简洁但完整 - - 若无内容可存储，请回复 [SILENT] + 工作流程： + 1. 先 `read` {memory_dir}/YYYY-MM-DD.md（如文件不存在，会返回错误提示） + 2. 从当前会话中提取并综合两类内容： + - 持久化记忆：客观事实、用户信息更新、项目状态及重要事件。 + - 经验反思：基于用户反馈形成的可复用思考逻辑、成功的问题解决策略、犯下的错误/应避免的陷阱，以及对未来交互有帮助的行动指南。 + 3. 智能合并新信息与现有内容（若文件不存在则跳过合并）： + - 将内容进行清晰的分类（例如明确区分“事实记忆”与“反思与逻辑”）。 + - 避免重复已记录的信息。 + - 在相关时丰富现有条目的新细节。 + - 在适用时保持时间顺序。 + 4. 写入更新后的内容： + - 尽可能使用 `edit` 更新特定部分。 + - 如需大幅重构则使用 `write` 覆盖整个文件。 + + 原则： + - 始终保留时间戳、日期和时间相关上下文。 + - 仅添加真正新的或有丰富价值的信息。 + - 反思内容必须侧重于根据用户反馈构建可复用的思维逻辑，以改善未来的任务执行。 + - 保持条目简洁但完整。 + - 若无任何新内容可存储或反思，请回复 [SILENT]。 \ No newline at end of file From 6dd987a1d24c860816d9b61516d90fc0bd0a9640 Mon Sep 17 00:00:00 2001 From: jinliyl <6469360+jinliyl@users.noreply.github.com> Date: Thu, 19 Mar 2026 19:53:32 +0800 Subject: [PATCH 34/59] refactor(core): update text truncation utilities and tool result handling (#168) * refactor(memory): remove mark filtering parameters and simplify get_memory logic * feat(core): bump version to 0.3.1.0 * refactor(memory): update comment to clarify dialog storage persistence * refactor(core): update text truncation utilities and tool result handling - Add new truncate_text_head function for head-based truncation - Introduce TRUNCATION_MARKER_START constant for truncation detection - Replace tail-based truncation with head-based truncation in tool result compaction - Remove tool_result_threshold and retention_days parameters from RemeLight initialization - Update ToolResultCompactor to use configurable thresholds for recent vs old messages - Modify compact_tool_result method to accept multiple threshold parameters - Adjust cleanup logic to use default compactor configuration - Simplify is_truncated function to check only start marker * ``` feat(memory): add long line splitting in tool result compaction - Added _split_long_lines function to break oversized lines at 10000 characters - Implemented line splitting before saving tool results to files - Prevents extremely long lines from breaking file-based storage - Maintains compatibility with existing tool result format - Preserves original content integrity through chunked processing ``` --- reme/__init__.py | 2 +- reme/core/utils/__init__.py | 4 +- reme/core/utils/truncate_text_utils.py | 33 +++++++++- .../components/tool_result_compactor.py | 64 +++++++++++++++---- .../file_based/reme_in_memory_memory.py | 24 +------ reme/reme_light.py | 46 ++++++------- 6 files changed, 111 insertions(+), 62 deletions(-) diff --git a/reme/__init__.py b/reme/__init__.py index 22374421..874836a7 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.9" +__version__ = "0.3.1.0" __all__ = [ "config", diff --git a/reme/core/utils/__init__.py b/reme/core/utils/__init__.py index 9dbc59dc..b51936c7 100644 --- a/reme/core/utils/__init__.py +++ b/reme/core/utils/__init__.py @@ -19,7 +19,7 @@ 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 -from .truncate_text_utils import truncate_text, is_truncated +from .truncate_text_utils import truncate_text, truncate_text_head, is_truncated, TRUNCATION_MARKER_START __all__ = [ "convert_dashscope_to_agentscope", @@ -52,5 +52,7 @@ __all__ = [ "get_now_time", "get_hf_token_counter", "truncate_text", + "truncate_text_head", "is_truncated", + "TRUNCATION_MARKER_START", ] diff --git a/reme/core/utils/truncate_text_utils.py b/reme/core/utils/truncate_text_utils.py index da0c473a..1f60d19f 100644 --- a/reme/core/utils/truncate_text_utils.py +++ b/reme/core/utils/truncate_text_utils.py @@ -41,15 +41,42 @@ def truncate_text(text: str, max_length: int) -> str: ) +def truncate_text_head(text: str, max_length: int) -> str: + """Truncate text from the beginning, keeping only the head portion. + + Args: + text: The text to truncate + max_length: Maximum allowed length + + Returns: + Truncated text with marker indicating truncation at the end + """ + text = str(text) if text else "" + if not text: + return text + + if len(text) <= max_length: + return text + + truncated_chars = len(text) - max_length + logger.debug( + "Text truncated from head: original %d chars, kept %d, removed %d chars from tail.", + len(text), + max_length, + truncated_chars, + ) + return f"{text[:max_length]}{TRUNCATION_MARKER_START}" + + def is_truncated(text: str) -> bool: - """Check if the text has been truncated (contains truncation markers). + """Check if the text has been truncated (contains truncation marker). Args: text: The text to check Returns: - bool: True if text contains truncation markers, False otherwise + bool: True if text contains truncation marker, False otherwise """ if not text: return False - return TRUNCATION_MARKER_START in text and TRUNCATION_MARKER_END in text + return TRUNCATION_MARKER_START in text diff --git a/reme/memory/file_based/components/tool_result_compactor.py b/reme/memory/file_based/components/tool_result_compactor.py index c1c6fe00..002bb8ae 100644 --- a/reme/memory/file_based/components/tool_result_compactor.py +++ b/reme/memory/file_based/components/tool_result_compactor.py @@ -8,10 +8,26 @@ from agentscope.message import Msg from ....core.op import BaseOp from ....core.utils import get_logger -from ....core.utils import truncate_text, is_truncated +from ....core.utils import truncate_text_head, TRUNCATION_MARKER_START logger = get_logger() +MAX_LINE_LENGTH = 10000 + + +def _split_long_lines(text: str, max_len: int = MAX_LINE_LENGTH) -> str: + """Split lines that exceed max_len by inserting newlines.""" + lines = text.split("\n") + result = [] + for line in lines: + if len(line) <= max_len: + result.append(line) + else: + # Split line into chunks of max_len + for i in range(0, len(line), max_len): + result.append(line[i : i + max_len]) + return "\n".join(result) + class ToolResultCompactor(BaseOp): """Truncate large tool_result outputs and save full content to files.""" @@ -19,43 +35,59 @@ class ToolResultCompactor(BaseOp): def __init__( self, tool_result_dir: str | Path, - tool_result_threshold: int, retention_days: int = 7, + recent_n: int = 1, + old_threshold: int = 500, + recent_threshold: int = 30000, **kwargs, ): super().__init__(**kwargs) self.tool_result_dir = Path(tool_result_dir) - self.tool_result_threshold = tool_result_threshold self.retention_days = retention_days + self.recent_n = recent_n + self.old_threshold = old_threshold + self.recent_threshold = recent_threshold - def _save_and_truncate(self, content: str, tool_name: str) -> str: + def _save_and_truncate(self, content: str, tool_name: str, threshold: int) -> 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: + if not content: return content - # Save full content + # Check if content was previously truncated + if TRUNCATION_MARKER_START in content: + parts = content.split(TRUNCATION_MARKER_START, 1) + if len(parts[0]) <= threshold: + return content + return f"{truncate_text_head(parts[0], threshold)}{parts[1]}" + + # Not truncated before + if len(content) <= threshold: + return content + + # Save full content with long lines split 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() + processed_content = _split_long_lines(content) file_path.write_text( - f"# tool_name: {tool_name}\n# created_at: {created_at}\n# ---\n{content}", + f"# tool_name: {tool_name}\n# created_at: {created_at}\n# ---\n{processed_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}]" + return f"{truncate_text_head(content, threshold)}\n\n[Full content saved to: {file_path}]" - def _process_output(self, output: str | list[dict], tool_name: str) -> str | list[dict]: + def _process_output(self, output: str | list[dict], tool_name: str, threshold: int) -> str | list[dict]: """Process tool result output, truncating if necessary.""" if isinstance(output, str): - return self._save_and_truncate(output, tool_name) + return self._save_and_truncate(output, tool_name, threshold) if isinstance(output, list): return [ ( - {**b, "text": self._save_and_truncate(b.get("text", ""), tool_name)} + {**b, "text": self._save_and_truncate(b.get("text", ""), tool_name, threshold)} if isinstance(b, dict) and b.get("type") == "text" else b ) @@ -69,15 +101,21 @@ class ToolResultCompactor(BaseOp): if not messages: return messages - for msg in messages: + # Split messages into old and recent parts + split_index = max(0, len(messages) - self.recent_n) + + for idx, msg in enumerate(messages): if not isinstance(msg.content, list): continue + # Determine threshold based on message position + threshold = self.recent_threshold if idx >= split_index else self.old_threshold + 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")) + block["output"] = self._process_output(output, block.get("name", "unknown"), threshold) return messages diff --git a/reme/memory/file_based/reme_in_memory_memory.py b/reme/memory/file_based/reme_in_memory_memory.py index 4fcec627..1c53c314 100644 --- a/reme/memory/file_based/reme_in_memory_memory.py +++ b/reme/memory/file_based/reme_in_memory_memory.py @@ -110,34 +110,19 @@ class ReMeInMemoryMemory(InMemoryMemory): 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] if prepend_summary and self._compressed_summary: previous_summary = f""" @@ -210,7 +195,7 @@ only execute the user's new instruction. if not messages: return 0 - # Persist messages to dialog storage + # Persist messages to dialog storage instead of compressed self._append_messages_to_dialog(messages) # Remove messages from memory @@ -258,10 +243,7 @@ only execute the user's new instruction. - context_usage_ratio: Usage percentage - 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 = await self._msg_handler.count_str_token(compressed_summary) diff --git a/reme/reme_light.py b/reme/reme_light.py index c8561550..58008120 100644 --- a/reme/reme_light.py +++ b/reme/reme_light.py @@ -61,8 +61,6 @@ class ReMeLight(Application): 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. - tool_result_threshold (int): Character threshold for tool result compaction. - retention_days (int): Number of days to retain tool result files. summary_tasks (list[asyncio.Task]): List of active background summary tasks. """ @@ -78,8 +76,6 @@ class ReMeLight(Application): default_file_store_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, ): """ @@ -111,11 +107,6 @@ class ReMeLight(Application): 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. - tool_result_threshold (int): Character count threshold for tool result - compaction. Results exceeding this length will be truncated and - saved to files. Default 1000 characters. - retention_days (int): Number of days to retain tool result files - before automatic cleanup. Default 7 days. enable_load_env (bool): Whether to load environment variables from .env file. Defaults to False. @@ -138,8 +129,6 @@ class ReMeLight(Application): 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 # Initialize the parent Application class with comprehensive configuration super().__init__( @@ -186,19 +175,15 @@ class ReMeLight(Application): 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 """ 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: @@ -243,7 +228,14 @@ class ReMeLight(Application): self._cleanup_tool_results() return await super().close() - async def compact_tool_result(self, messages: list[Msg]) -> list[Msg]: + async def compact_tool_result( + self, + messages: list[Msg], + recent_n: int = 1, + old_threshold: int = 500, + recent_threshold: int = 30000, + retention_days: int = 7, + ) -> list[Msg]: """ Compact tool results by truncating large outputs and saving full content to files. @@ -255,13 +247,19 @@ class ReMeLight(Application): Args: messages (list[Msg]): List of messages potentially containing tool results that may need compaction. + recent_n (int): Number of recent messages to use recent_threshold for. + Default 1. + old_threshold (int): Character threshold for old messages. Default 500. + recent_threshold (int): Character threshold for recent messages. Default 30000. + retention_days (int): Number of days to retain tool result files. + Default 7. Returns: list[Msg]: The processed list of messages with large tool results compacted. If an error occurs, returns the original unmodified messages. Note: - - Tool results shorter than tool_result_threshold are left unchanged + - Tool results are truncated based on old_threshold/recent_threshold - Full content of truncated results is saved to tool_result_path - Expired files are automatically cleaned up during this operation """ @@ -269,8 +267,10 @@ class ReMeLight(Application): # 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, + recent_n=recent_n, + old_threshold=old_threshold, + recent_threshold=recent_threshold, ) # Execute compaction and get processed messages From 4f63fbf197d64762dfe7bbb45528262cdb3d342f Mon Sep 17 00:00:00 2001 From: jinliyl <6469360+jinliyl@users.noreply.github.com> Date: Thu, 19 Mar 2026 23:40:56 +0800 Subject: [PATCH 35/59] refactor(memory): update conversation continuity context handling (#170) * refactor(memory): update conversation continuity context handling * chore(version): bump version to 0.3.1.1 --- reme/__init__.py | 2 +- reme/memory/file_based/reme_in_memory_memory.py | 3 +-- 2 files changed, 2 insertions(+), 3 deletions(-) diff --git a/reme/__init__.py b/reme/__init__.py index 874836a7..97905514 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.1.0" +__version__ = "0.3.1.1" __all__ = [ "config", diff --git a/reme/memory/file_based/reme_in_memory_memory.py b/reme/memory/file_based/reme_in_memory_memory.py index 1c53c314..5882b900 100644 --- a/reme/memory/file_based/reme_in_memory_memory.py +++ b/reme/memory/file_based/reme_in_memory_memory.py @@ -130,8 +130,7 @@ class ReMeInMemoryMemory(InMemoryMemory): {self._compressed_summary} The above is a summary of our previous conversation. -If there is a new instruction from the user, do not continue executing the previous content; -only execute the user's new instruction. +Use it as context to maintain continuity. """.strip() return [ From 33f6822792b595d227044c42dc9e7e754aa0a9fb Mon Sep 17 00:00:00 2001 From: Aleksandr Mordvinov Date: Fri, 20 Mar 2026 13:05:04 +0700 Subject: [PATCH 36/59] fix: support BGE-M3 embedding (dense_embedding fallback) (#169) BGE-M3 returns dense_embedding instead of embedding; use dense_embedding as fallback when embedding is None to avoid TypeError. Made-with: Cursor Co-authored-by: AleksandrMordvinov --- reme/core/embedding/openai_embedding_model.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) 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): From b8619aaabcf4fd6bbfe6b0daa9fd2bbf812b8f25 Mon Sep 17 00:00:00 2001 From: "jinli.yl" Date: Fri, 20 Mar 2026 15:22:28 +0800 Subject: [PATCH 37/59] test(config): enable environment loading in test configuration --- tests/light/test_reme_light.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/light/test_reme_light.py b/tests/light/test_reme_light.py index fd884189..4fc7452a 100644 --- a/tests/light/test_reme_light.py +++ b/tests/light/test_reme_light.py @@ -34,6 +34,7 @@ async def main(): 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() From 8f48f91a43a735ed9ca36974b51155c02dbc90f8 Mon Sep 17 00:00:00 2001 From: jinliyl <6469360+jinliyl@users.noreply.github.com> Date: Fri, 20 Mar 2026 15:23:56 +0800 Subject: [PATCH 38/59] Enable environment loading in ReMeLight configuration --- README.md | 1 + 1 file changed, 1 insertion(+) diff --git a/README.md b/README.md index db94195d..47196a58 100644 --- a/README.md +++ b/README.md @@ -145,6 +145,7 @@ async def main(): 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() From e7993a469a6b48ba751878b3968f370108240239 Mon Sep 17 00:00:00 2001 From: jinliyl <6469360+jinliyl@users.noreply.github.com> Date: Fri, 20 Mar 2026 15:24:20 +0800 Subject: [PATCH 39/59] Enable environment loading in ReMeLight configuration --- README_ZH.md | 1 + 1 file changed, 1 insertion(+) diff --git a/README_ZH.md b/README_ZH.md index 6218b87c..f715d3c3 100644 --- a/README_ZH.md +++ b/README_ZH.md @@ -137,6 +137,7 @@ async def main(): 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() From 53030de4315075204f4b1e479115c9c828d56884 Mon Sep 17 00:00:00 2001 From: Zhouwk <57825291+nitwtog@users.noreply.github.com> Date: Fri, 20 Mar 2026 16:30:15 +0800 Subject: [PATCH 40/59] =?UTF-8?q?=E6=9B=B4=E6=96=B0=E5=AE=9E=E9=AA=8C?= =?UTF-8?q?=E7=BB=93=E6=9E=9C=E5=9C=A8README=E4=B8=AD=E4=BD=8D=E7=BD=AE=20?= =?UTF-8?q?(#171)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * Update README_ZH.md * Update README.md * Update README.md * Update README_ZH.md --- README.md | 68 +++++++++++++++++++++++++--------------------------- README_ZH.md | 67 +++++++++++++++++++++++++-------------------------- 2 files changed, 66 insertions(+), 69 deletions(-) diff --git a/README.md b/README.md index 47196a58..ca2538b6 100644 --- a/README.md +++ b/README.md @@ -34,6 +34,8 @@ conversations) and **stateless sessions** (new sessions cannot inherit history a ReMe gives agents **real memory** — old conversations are automatically compacted, important information is persistently stored, and relevant context is automatically recalled in future interactions. +ReMe achieves state-of-the-art results on the LoCoMo and HaluMem benchmarks; see the [Experimental results](#experimental-results). +
What you can do with ReMe @@ -440,40 +442,6 @@ memories: Installation and environment configuration are the same as [ReMeLight](#installation). API keys are configured via environment variables and can be stored in a `.env` file at the project root. -## 🧪 Experiments - -Evaluations are conducted on three benchmarks: **LoCoMo** and **HaluMem**. Experimental settings: - -1. **ReMe backbone**: as specified in each table. -2. **Evaluation protocol**: LLM-as-a-Judge following MemOS — each answer is scored by GPT-4o-mini. - -Baseline results are reproduced from their respective papers under aligned settings where possible. - -### LoCoMo - -| Method | Single Hop | Multi Hop | Temporal | Open Domain | Overall | -|----------|------------|-----------|-----------|-------------|-----------| -| MemoryOS | 62.43 | 56.50 | 37.18 | 40.28 | 54.70 | -| Mem0 | 66.71 | 58.16 | 55.45 | 40.62 | 61.00 | -| MemU | 72.77 | 62.41 | 33.96 | 46.88 | 61.15 | -| MemOS | 81.45 | 69.15 | 72.27 | 60.42 | 75.87 | -| HiMem | 89.22 | 70.92 | 74.77 | 54.86 | 80.71 | -| Zep | 88.11 | 71.99 | 74.45 | 66.67 | 81.06 | -| TiMem | 81.43 | 62.20 | 77.63 | 52.08 | 75.30 | -| TSM | 84.30 | 66.67 | 71.03 | 58.33 | 76.69 | -| MemR3 | 89.44 | 71.39 | 76.22 | 61.11 | 81.55 | -| **ReMe** | **89.89** | **82.98** | **83.80** | **71.88** | **86.23** | - -### HaluMem - -| Method | Memory Integrity | Memory Accuracy | QA Accuracy | -|-------------|------------------|-----------------|-------------| -| MemoBase | 14.55 | 92.24 | 35.53 | -| Supermemory | 41.53 | 90.32 | 54.07 | -| Mem0 | 42.91 | 86.26 | 53.02 | -| ProMem | **73.80** | 89.47 | 62.26 | -| **ReMe** | 67.72 | **94.06** | **88.78** | - ### Python usage ```python @@ -589,7 +557,37 @@ graph LR ### Experimental results -Coming soon... +Evaluations are conducted on two benchmarks: **LoCoMo** and **HaluMem**. Experimental settings: + +1. **ReMe backbone**: as specified in each table. +2. **Evaluation protocol**: LLM-as-a-Judge following MemOS — each answer is scored by GPT-4o-mini. + +Baseline results are reproduced from their respective papers under aligned settings where possible. + +### LoCoMo + +| Method | Single Hop | Multi Hop | Temporal | Open Domain | Overall | +|----------|------------|-----------|-----------|-------------|-----------| +| MemoryOS | 62.43 | 56.50 | 37.18 | 40.28 | 54.70 | +| Mem0 | 66.71 | 58.16 | 55.45 | 40.62 | 61.00 | +| MemU | 72.77 | 62.41 | 33.96 | 46.88 | 61.15 | +| MemOS | 81.45 | 69.15 | 72.27 | 60.42 | 75.87 | +| HiMem | 89.22 | 70.92 | 74.77 | 54.86 | 80.71 | +| Zep | 88.11 | 71.99 | 74.45 | 66.67 | 81.06 | +| TiMem | 81.43 | 62.20 | 77.63 | 52.08 | 75.30 | +| TSM | 84.30 | 66.67 | 71.03 | 58.33 | 76.69 | +| MemR3 | 89.44 | 71.39 | 76.22 | 61.11 | 81.55 | +| **ReMe** | **89.89** | **82.98** | **83.80** | **71.88** | **86.23** | + +### HaluMem + +| Method | Memory Integrity | Memory Accuracy | QA Accuracy | +|-------------|------------------|-----------------|-------------| +| MemoBase | 14.55 | 92.24 | 35.53 | +| Supermemory | 41.53 | 90.32 | 54.07 | +| Mem0 | 42.91 | 86.26 | 53.02 | +| ProMem | **73.80** | 89.47 | 62.26 | +| **ReMe** | 67.72 | **94.06** | **88.78** | --- diff --git a/README_ZH.md b/README_ZH.md index f715d3c3..32a89039 100644 --- a/README_ZH.md +++ b/README_ZH.md @@ -32,6 +32,8 @@ ReMe 让智能体拥有**真正的记忆力**——旧对话自动浓缩，重要信息持久保存，下次对话自动想起来。 +在 LoCoMo 与 HaluMem 基准测试中，ReMe 取得了领先结果，详见[实验效果](#实验效果)。 +
你可以用 ReMe 做什么 @@ -421,39 +423,6 @@ graph LR 安装和环境变量配置与 [ReMeLight 一致](#安装)，通过环境变量设置 API 密钥，可写在项目根目录的 `.env` 文件中。 -## 🧪 实验 - -本实验部分在 LoCoMo、LongMemEval、HaluMem 三个数据集上进行评测，实验设置如下： - -1. **ReMe 使用模型**：如各表 backbone 列所示。 -2. **评估使用模型**：采用 LLM-as-a-Judge 协议（参照 MemOS）——每条回答由 GPT-4o-mini 裁判模型打分。 - -实验设置尽量与各基线论文保持一致，以复用其公开结果。 - -### LoCoMo - -| Method | Single Hop | Multi Hop | Temporal | Open Domain | Overall | -|----------|------------|-----------|-----------|-------------|-----------| -| MemoryOS | 62.43 | 56.50 | 37.18 | 40.28 | 54.70 | -| Mem0 | 66.71 | 58.16 | 55.45 | 40.62 | 61.00 | -| MemU | 72.77 | 62.41 | 33.96 | 46.88 | 61.15 | -| MemOS | 81.45 | 69.15 | 72.27 | 60.42 | 75.87 | -| HiMem | 89.22 | 70.92 | 74.77 | 54.86 | 80.71 | -| Zep | 88.11 | 71.99 | 74.45 | 66.67 | 81.06 | -| TiMem | 81.43 | 62.20 | 77.63 | 52.08 | 75.30 | -| TSM | 84.30 | 66.67 | 71.03 | 58.33 | 76.69 | -| MemR3 | 89.44 | 71.39 | 76.22 | 61.11 | 81.55 | -| **ReMe** | **89.89** | **82.98** | **83.80** | **71.88** | **86.23** | - -### HaluMem - -| Method | Memory Integrity | Memory Accuracy | QA Accuracy | -|-------------|------------------|-----------------|-------------| -| MemoBase | 14.55 | 92.24 | 35.53 | -| Supermemory | 41.53 | 90.32 | 54.07 | -| Mem0 | 42.91 | 86.26 | 53.02 | -| ProMem | **73.80** | 89.47 | 62.26 | -| **ReMe** | 67.72 | **94.06** | **88.78** | ### Python 使用 @@ -570,7 +539,37 @@ graph LR ### 实验效果 -Coming soon... +本实验部分在 LoCoMo和HaluMem 两个数据集上进行评测，实验设置如下： + +1. **ReMe 使用模型**：如各表 backbone 列所示。 +2. **评估使用模型**：采用 LLM-as-a-Judge 协议（参照 MemOS）——每条回答由 GPT-4o-mini 裁判模型打分。 + +实验设置尽量与各基线论文保持一致，以复用其公开结果。 + +### LoCoMo + +| Method | Single Hop | Multi Hop | Temporal | Open Domain | Overall | +|----------|------------|-----------|-----------|-------------|-----------| +| MemoryOS | 62.43 | 56.50 | 37.18 | 40.28 | 54.70 | +| Mem0 | 66.71 | 58.16 | 55.45 | 40.62 | 61.00 | +| MemU | 72.77 | 62.41 | 33.96 | 46.88 | 61.15 | +| MemOS | 81.45 | 69.15 | 72.27 | 60.42 | 75.87 | +| HiMem | 89.22 | 70.92 | 74.77 | 54.86 | 80.71 | +| Zep | 88.11 | 71.99 | 74.45 | 66.67 | 81.06 | +| TiMem | 81.43 | 62.20 | 77.63 | 52.08 | 75.30 | +| TSM | 84.30 | 66.67 | 71.03 | 58.33 | 76.69 | +| MemR3 | 89.44 | 71.39 | 76.22 | 61.11 | 81.55 | +| **ReMe** | **89.89** | **82.98** | **83.80** | **71.88** | **86.23** | + +### HaluMem + +| Method | Memory Integrity | Memory Accuracy | QA Accuracy | +|-------------|------------------|-----------------|-------------| +| MemoBase | 14.55 | 92.24 | 35.53 | +| Supermemory | 41.53 | 90.32 | 54.07 | +| Mem0 | 42.91 | 86.26 | 53.02 | +| ProMem | **73.80** | 89.47 | 62.26 | +| **ReMe** | 67.72 | **94.06** | **88.78** | --- From 0beaa035cbbc886dcb4fdf3bfe6243aea85ee6c1 Mon Sep 17 00:00:00 2001 From: jinliyl <6469360+jinliyl@users.noreply.github.com> Date: Fri, 20 Mar 2026 17:55:18 +0800 Subject: [PATCH 41/59] fix(file-store): handle embedding API errors gracefully with fallback mechanism (#173) --- reme/core/file_store/base_file_store.py | 40 ++++++++++++++++++++++--- 1 file changed, 36 insertions(+), 4 deletions(-) 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): From 7b02c4521826c3f497f95d3077d676cd31b11a74 Mon Sep 17 00:00:00 2001 From: jinliyl <6469360+jinliyl@users.noreply.github.com> Date: Tue, 24 Mar 2026 00:20:15 +0800 Subject: [PATCH 42/59] style(memory): update message formatting and improve logging (#175) * style(memory): update message formatting and improve logging - Change default include_thinking parameter to True in as_msg_handler.py - Replace angle brackets with square brackets for block formatting in as_msg_stat.py - Add newline replacement in text truncation method in as_msg_stat.py - Add loading duration timing to embedding cache loading in base_embedding_model.py - Replace XML-style tags with markdown headers in compactor.py conversation format - Update compactor.yaml prompts to reference markdown-style headers instead of XML tags - Modify summarizer.py to use markdown-style conversation header format * refactor(file-watcher): replace scan_on_start with rebuild_index_on_start parameter - Replace scan_on_start and clear_on_start boolean parameters with single rebuild_index_on_start - Update BaseFileWatcher constructor to use rebuild_index_on_start instead of two separate flags - Modify initialization logic to clear and rescan when rebuild_index_on_start is True - Remove scan_on_start parameter from CLI and light configuration files - Update documentation to remove scan_on_start from quick start guides - Rename all test methods and classes from scan_on_start to rebuild_index_on_start - Add timezone-aware datetime helper method to summarizer component - Format log message with proper line breaks for readability * fix(core): resolve file watcher initialization issue and update version - Fixed file watcher task creation to properly handle rebuild index on start logic - Moved initialization and watch loop into async function to ensure proper execution order - Updated package version from 0.3.1.1 to 0.3.1.2 - Added missing comma in embedding model logging statement * fix(core): reduce max formatter text length limit - Changed _DEFAULT_MAX_FORMATTER_TEXT_LENGTH from 2000 to 1000 - Updated constant value in as_msg_stat.py schema module * fix(file-watcher): change default rebuild index behavior on start - Changed rebuild_index_on_start parameter default from False to True - This ensures index is rebuilt by default when file watcher starts - Maintains consistent state initialization for file watching operations * feat(compactor): add return_dict option and improve summary validation - Add _is_valid_summary function to validate summary content format - Introduce return_dict parameter to return structured results with validation - Update prompt templates with clearer task descriptions and formatting rules - Refactor update_user_message prompts to combine prefix and suffix logic - Return dictionary with user_message, history_compact, and is_valid fields when enabled - Add proper error handling for exception cases in memory compaction - Maintain backward compatibility with string return when return_dict=False * feat(memory): add thinking block configuration option - Add add_thinking_block parameter to compactor component - Pass include_thinking flag to message formatting in compactor - Add add_thinking_block parameter to reme_light compact function - Add add_thinking_block parameter to reme_light summarize function - Add add_thinking_block parameter to summarizer component - Pass include_thinking flag to message formatting in summarizer - Remove previous-summary tags from compressed summary format --- docs/cli/quick_start_en.md | 1 - docs/cli/quick_start_zh.md | 1 - reme/__init__.py | 2 +- reme/config/cli.yaml | 1 - reme/config/light.yaml | 1 - reme/core/embedding/base_embedding_model.py | 5 +- reme/core/file_watcher/base_file_watcher.py | 27 +++++----- reme/core/schema/as_msg_stat.py | 15 +++--- .../memory/file_based/components/compactor.py | 49 +++++++++++++++---- .../file_based/components/compactor.yaml | 49 +++++++++++-------- .../file_based/components/summarizer.py | 22 ++++++--- .../file_based/reme_in_memory_memory.py | 5 +- .../memory/file_based/utils/as_msg_handler.py | 2 +- reme/reme_light.py | 19 +++++-- test/test/test_fs_file_watch_integration.py | 1 - tests/test_base_file_watcher.py | 46 ++++++++--------- 16 files changed, 146 insertions(+), 100 deletions(-) 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/reme/__init__.py b/reme/__init__.py index 97905514..adb4cbe4 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.1.1" +__version__ = "0.3.1.2" __all__ = [ "config", diff --git a/reme/config/cli.yaml b/reme/config/cli.yaml index a6914170..e7cab227 100644 --- a/reme/config/cli.yaml +++ b/reme/config/cli.yaml @@ -47,5 +47,4 @@ file_watchers: watch_paths: [ ".reme", ".reme/memory" ] suffix_filters: [ ".md" ] recursive: false - scan_on_start: true diff --git a/reme/config/light.yaml b/reme/config/light.yaml index e18a1302..242df716 100644 --- a/reme/config/light.yaml +++ b/reme/config/light.yaml @@ -34,4 +34,3 @@ file_watchers: file_store: default suffix_filters: [ ".md" ] recursive: false - scan_on_start: true diff --git a/reme/core/embedding/base_embedding_model.py b/reme/core/embedding/base_embedding_model.py index 28efddce..aeea6869 100644 --- a/reme/core/embedding/base_embedding_model.py +++ b/reme/core/embedding/base_embedding_model.py @@ -159,6 +159,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() @@ -201,7 +202,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: diff --git a/reme/core/file_watcher/base_file_watcher.py b/reme/core/file_watcher/base_file_watcher.py index 61e8a147..c9063997 100644 --- a/reme/core/file_watcher/base_file_watcher.py +++ b/reme/core/file_watcher/base_file_watcher.py @@ -34,8 +34,7 @@ class BaseFileWatcher: 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 = True, - clear_on_start: bool = True, + rebuild_index_on_start: bool = True, **kwargs, ): """ @@ -50,9 +49,8 @@ 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 - clear_on_start: If True, clear all indexed data on start before scanning. - Useful for full rebuild of the index. + rebuild_index_on_start: If True, clear all indexed data on start and rescan existing files. + If False, only monitor new changes without initialization. **kwargs: Additional keyword arguments """ self.watch_paths: list[str] = [watch_paths] if isinstance(watch_paths, str) else watch_paths @@ -63,8 +61,7 @@ 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.clear_on_start: bool = clear_on_start + self.rebuild_index_on_start: bool = rebuild_index_on_start self.kwargs: dict = kwargs self._stop_event = asyncio.Event() @@ -78,16 +75,14 @@ class BaseFileWatcher: self._running = True - # Clear all indexed data if requested - if self.clear_on_start and self.file_store is not None: - await self.file_store.clear_all() - logger.info("Cleared all indexed data on start") + 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() - # Scan existing files if requested - if self.scan_on_start: - await self._scan_existing_files() - - 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): diff --git a/reme/core/schema/as_msg_stat.py b/reme/core/schema/as_msg_stat.py index b4ef02d4..b502dd02 100644 --- a/reme/core/schema/as_msg_stat.py +++ b/reme/core/schema/as_msg_stat.py @@ -3,7 +3,7 @@ from pydantic import BaseModel, Field _DEFAULT_MAX_BLOCK_TEXT_PREVIEW_LENGTH = 100 -_DEFAULT_MAX_FORMATTER_TEXT_LENGTH = 2000 +_DEFAULT_MAX_FORMATTER_TEXT_LENGTH = 1000 class AsBlockStat(BaseModel): @@ -27,7 +27,8 @@ class AsBlockStat(BaseModel): return self.format(_DEFAULT_MAX_BLOCK_TEXT_PREVIEW_LENGTH) def _truncate(self, text: str, max_length: int) -> str: - """Simple truncation with ellipsis.""" + """Truncate text with ellipsis, replacing newlines with spaces.""" + text = text.replace("\n", " ") if len(text) <= max_length: return text return text[:max_length] + "..." @@ -46,22 +47,22 @@ class AsBlockStat(BaseModel): if self.block_type == "text": if not self.text: return "" - return f"{self._truncate(self.text, max_length)}" + 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"{self._truncate(self.text, max_length)}" + 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}" + return f"[{self.block_type}]: {content}" if self.block_type == "tool_use": content = f"{self.tool_name} params={self._truncate(self.tool_input, max_length)}" - return f"{content}" + return f"[tool_use]: {content}" if self.block_type == "tool_result": if not self.tool_output: return "" content = f"{self.tool_name} output={self._truncate(self.tool_output, max_length)}" - return f"{content}" + return f"[tool_result]: {content}" return "" diff --git a/reme/memory/file_based/components/compactor.py b/reme/memory/file_based/components/compactor.py index b6725b0b..7e8b31ad 100644 --- a/reme/memory/file_based/components/compactor.py +++ b/reme/memory/file_based/components/compactor.py @@ -10,6 +10,22 @@ 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.""" @@ -17,17 +33,24 @@ class Compactor(BaseOp): self, memory_compact_threshold: int, console_enabled: bool = False, + return_dict: bool = False, + add_thinking_block: bool = True, **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 + # 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) @@ -35,12 +58,15 @@ class Compactor(BaseOp): 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( @@ -52,18 +78,12 @@ class Compactor(BaseOp): agent.set_console_output_enabled(self.console_enabled) 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}" + 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"\n{history_formatted_str}\n\n\n" + self.get_prompt( - "initial_user_message", - ) + user_message: str = f"# conversation\n{history_formatted_str}\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( @@ -75,5 +95,16 @@ class Compactor(BaseOp): ) 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/components/compactor.yaml b/reme/memory/file_based/components/compactor.yaml index 83c4f952..6fef5383 100644 --- a/reme/memory/file_based/components/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/summarizer.py b/reme/memory/file_based/components/summarizer.py index 855d64c1..9733eaa5 100644 --- a/reme/memory/file_based/components/summarizer.py +++ b/reme/memory/file_based/components/summarizer.py @@ -25,6 +25,7 @@ class Summarizer(BaseOp): toolkit: Toolkit | None = None, console_enabled: bool = False, timezone: str | None = None, + add_thinking_block: bool = True, **kwargs, ): super().__init__(**kwargs) @@ -34,6 +35,16 @@ class Summarizer(BaseOp): 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", []) @@ -46,6 +57,7 @@ class Summarizer(BaseOp): 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}") @@ -63,15 +75,9 @@ class Summarizer(BaseOp): ) agent.set_console_output_enabled(self.console_enabled) - user_message: str = f"\n{history_formatted_str}\n\n" + self.prompt_format( + user_message: str = f"# conversation\n{history_formatted_str}\n\n" + self.prompt_format( "user_message", - date=( - datetime.datetime.now( - zoneinfo.ZoneInfo(self.timezone), - ) - if self.timezone - else datetime.datetime.now() - ).strftime("%Y-%m-%d"), + date=self._get_current_datetime().strftime("%Y-%m-%d"), working_dir=self.working_dir, memory_dir=self.memory_dir, ) diff --git a/reme/memory/file_based/reme_in_memory_memory.py b/reme/memory/file_based/reme_in_memory_memory.py index 5882b900..cd4d2d8b 100644 --- a/reme/memory/file_based/reme_in_memory_memory.py +++ b/reme/memory/file_based/reme_in_memory_memory.py @@ -126,11 +126,8 @@ class ReMeInMemoryMemory(InMemoryMemory): 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. +The above is a summary of previous conversation, use it as context to maintain continuity. """.strip() return [ diff --git a/reme/memory/file_based/utils/as_msg_handler.py b/reme/memory/file_based/utils/as_msg_handler.py index 0692ebce..e175a402 100644 --- a/reme/memory/file_based/utils/as_msg_handler.py +++ b/reme/memory/file_based/utils/as_msg_handler.py @@ -204,7 +204,7 @@ class AsMsgHandler: self, messages: list[Msg], memory_compact_threshold: int, - include_thinking: bool = False, + include_thinking: bool = True, ) -> str: """Format list of messages to a single formatted string. diff --git a/reme/reme_light.py b/reme/reme_light.py index 58008120..d9835f5f 100644 --- a/reme/reme_light.py +++ b/reme/reme_light.py @@ -347,7 +347,9 @@ class ReMeLight(Application): max_input_length: float = 128 * 1024, compact_ratio: float = 0.7, previous_summary: str = "", - ) -> str: + return_dict: bool = False, + add_thinking_block: bool = True, + ) -> str | dict: """ Compact a list of messages into a condensed summary. @@ -371,10 +373,13 @@ class ReMeLight(Application): 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. Returns: - str: The condensed summary of the messages, or an empty string if - an error occurred during compaction. + 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: compactor = Compactor( @@ -383,6 +388,8 @@ class ReMeLight(Application): 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, ) return await compactor.call( @@ -392,8 +399,10 @@ 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( @@ -407,6 +416,7 @@ class ReMeLight(Application): 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. @@ -459,6 +469,7 @@ class ReMeLight(Application): as_token_counter=as_token_counter, language=language if language == "zh" else "", timezone=timezone, + add_thinking_block=add_thinking_block, ) return await summarizer.call(messages=messages, service_context=self.service_context) diff --git a/test/test/test_fs_file_watch_integration.py b/test/test/test_fs_file_watch_integration.py index 9f679d38..a3534337 100644 --- a/test/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/tests/test_base_file_watcher.py b/tests/test_base_file_watcher.py index ca310eb3..234d03c2 100644 --- a/tests/test_base_file_watcher.py +++ b/tests/test_base_file_watcher.py @@ -5,7 +5,7 @@ Async unit tests for BaseFileWatcher covering: - File suffix filtering - Start/stop lifecycle - Callback functionality -- scan_on_start feature +- rebuild_index_on_start feature Usage: pytest tests/test_base_file_watcher.py -v @@ -369,12 +369,12 @@ class TestCallbackFunctionality: # ==================== Test Scan on Start ==================== -class TestScanOnStart: - """Tests for scan_on_start feature.""" +class TestRebuildIndexOnStart: + """Tests for rebuild_index_on_start feature.""" @pytest.mark.asyncio - async def test_scan_on_start_false(self, temp_files, temp_dir: Path): - """Test that scan_on_start=False doesn't scan existing files.""" + 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): @@ -387,7 +387,7 @@ class TestScanOnStart: watcher = BaseFileWatcher( watch_paths=str(temp_dir), - scan_on_start=False, + rebuild_index_on_start=False, callback=callback, file_store=mock_file_store, ) @@ -400,8 +400,8 @@ class TestScanOnStart: assert len(callback_called) == 0 @pytest.mark.asyncio - async def test_scan_on_start_true_with_files(self, temp_files, temp_dir: Path): - """Test that scan_on_start=True scans existing files.""" + 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): @@ -414,7 +414,7 @@ class TestScanOnStart: watcher = BaseFileWatcher( watch_paths=str(temp_dir), - scan_on_start=True, + rebuild_index_on_start=True, callback=callback, file_store=mock_file_store, ) @@ -434,8 +434,8 @@ class TestScanOnStart: assert all(change == Change.added for change, _ in all_changes) @pytest.mark.asyncio - async def test_scan_on_start_with_suffix_filter(self, temp_files, temp_dir: Path): - """Test scan_on_start respects suffix filters.""" + 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): @@ -447,7 +447,7 @@ class TestScanOnStart: watcher = BaseFileWatcher( watch_paths=str(temp_dir), - scan_on_start=True, + rebuild_index_on_start=True, suffix_filters=[".txt"], callback=callback, file_store=mock_file_store, @@ -467,8 +467,8 @@ class TestScanOnStart: assert path.endswith(".txt"), f"Expected .txt file, got {path}" @pytest.mark.asyncio - async def test_scan_on_start_recursive(self, temp_nested_dir: Path): - """Test scan_on_start with recursive=True.""" + 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): @@ -480,7 +480,7 @@ class TestScanOnStart: watcher = BaseFileWatcher( watch_paths=str(temp_nested_dir), - scan_on_start=True, + rebuild_index_on_start=True, recursive=True, suffix_filters=[".txt"], callback=callback, @@ -503,8 +503,8 @@ class TestScanOnStart: assert nested_found, "Should find files in nested directories" @pytest.mark.asyncio - async def test_scan_on_start_non_recursive(self, temp_nested_dir: Path): - """Test scan_on_start with recursive=False.""" + 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): @@ -516,7 +516,7 @@ class TestScanOnStart: watcher = BaseFileWatcher( watch_paths=str(temp_nested_dir), - scan_on_start=True, + rebuild_index_on_start=True, recursive=False, suffix_filters=[".txt"], callback=callback, @@ -538,8 +538,8 @@ class TestScanOnStart: assert not nested_found, "Should not find files in nested directories" @pytest.mark.asyncio - async def test_scan_on_start_nonexistent_path(self): - """Test scan_on_start with non-existent path.""" + 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): @@ -551,7 +551,7 @@ class TestScanOnStart: watcher = BaseFileWatcher( watch_paths="/nonexistent/path", - scan_on_start=True, + rebuild_index_on_start=True, callback=callback, file_store=mock_file_store, ) @@ -709,7 +709,7 @@ class TestEdgeCases: watcher = BaseFileWatcher( watch_paths=str(file_path), - scan_on_start=True, + rebuild_index_on_start=True, callback=callback, file_store=mock_file_store, ) @@ -742,7 +742,7 @@ class TestEdgeCases: watcher = BaseFileWatcher( watch_paths=str(empty_dir), - scan_on_start=True, + rebuild_index_on_start=True, callback=callback, file_store=mock_file_store, ) From cc11b77b27ca5a6d8ccae62d9939b46f7eb6d3ca Mon Sep 17 00:00:00 2001 From: jinliyl <6469360+jinliyl@users.noreply.github.com> Date: Tue, 24 Mar 2026 22:15:02 +0800 Subject: [PATCH 43/59] chore(deps): update version and move litellm to dev dependencies (#176) * chore(deps): update version and move litellm to dev dependencies - Updated package version from 0.3.1.2 to 0.3.1.3 - Removed litellm from main dependencies in pyproject.toml - Added litellm as fixed version dependency in dev group - Maintained litellm requirement while reorganizing dependency structure * chore(deps): move litellm dependency to full extras - Moved litellm==1.80.0 from main dependencies to full extra - Kept litellm as optional dependency for users needing full feature set - Maintains backward compatibility for light installation option * feat(pyproject): add litellm dependency to project configuration - Added litellm==1.80.0 as optional dependency in pyproject.toml - Created new litellm extra group for LiteLLM integration - Updated full dependency group to include the new litellm option --- pyproject.toml | 5 ++++- reme/__init__.py | 2 +- 2 files changed, 5 insertions(+), 2 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 4b91b329..5a43f021 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -44,7 +44,6 @@ 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", @@ -80,6 +79,10 @@ full = [ "reme_ai[dev,ray,light]", ] +litellm = [ + "litellm==1.80.0", +] + light = [ "agentscope==1.0.17", ] diff --git a/reme/__init__.py b/reme/__init__.py index adb4cbe4..28d87f9e 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.1.2" +__version__ = "0.3.1.3" __all__ = [ "config", From 7f6bf11aab11576581d4c8398ebd53a5bd2f2b98 Mon Sep 17 00:00:00 2001 From: "jinli.yl" Date: Tue, 24 Mar 2026 22:18:12 +0800 Subject: [PATCH 44/59] refactor(pyproject.toml): move flowllm dependency from core to litellm extra --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 5a43f021..ececfe6d 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", @@ -81,6 +80,7 @@ full = [ litellm = [ "litellm==1.80.0", + "flowllm[reme]>=0.2.0.10", ] light = [ From 5b801c0d3e1aa9c0a1038b8a4eb613fab6863a7b Mon Sep 17 00:00:00 2001 From: jinliyl <6469360+jinliyl@users.noreply.github.com> Date: Wed, 25 Mar 2026 20:21:37 +0800 Subject: [PATCH 45/59] refactor(file_io): update file I/O operations and truncation logic (#177) * refactor(file_io): update file I/O operations and truncation logic * refactor(memory): update file-based memory compaction logic --- docs/copaw_context_design.md | 122 +++++++++++ reme/config/light.yaml | 2 + reme/core/application.py | 6 +- reme/core/file_watcher/base_file_watcher.py | 9 +- reme/core/op/base_op.py | 2 +- reme/core/schema/service_config.py | 5 +- reme/memory/file_based/components/cli.py | 6 +- .../file_based/components/summarizer.yaml | 39 ++-- .../components/tool_result_compactor.py | 145 ++++++------- .../file_based/reme_in_memory_memory.py | 5 +- reme/memory/file_based/tools/file_io.py | 72 ++++--- reme/memory/file_based/tools/shell.py | 6 - reme/memory/file_based/utils/__init__.py | 7 +- reme/memory/file_based/utils/file_utils.py | 175 ++++++++------- reme/reme_light.py | 41 ++-- tests/light/test_summarizer.py | 6 +- tests/light/test_tool_result_compactor.py | 18 +- tests/light/test_tools.py | 199 ++++++++++++++---- 18 files changed, 579 insertions(+), 286 deletions(-) create mode 100644 docs/copaw_context_design.md diff --git a/docs/copaw_context_design.md b/docs/copaw_context_design.md new file mode 100644 index 00000000..01da3764 --- /dev/null +++ b/docs/copaw_context_design.md @@ -0,0 +1,122 @@ +## Copaw Context Management V2 + +> 注：不涉及长期记忆 + +### 上下文数据结构 + +#### 1. 上下文-内存 + +- **compact_summary**（可选）： + - **历史对话原始数据引导**：存储于 `dialog/YYYY-MM-DD.jsonl`，共 N 行，按时间顺序排列；回顾时建议从后往前读。 + - **历史对话摘要**：包含 `Goal + Constraints + Progress + KeyDecisions + NextSteps`。 +- **messages**：当前对话上下文（完整消息列表）。 + +#### 2. 上下文-缓存到文件系统 + +- **历史对话原始数据**：`dialog/YYYY-MM-DD.jsonl` +- **工具调用结果原始数据**：`tool_result/{uuid}.txt`（保留 N 天） + +```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] --> G[dialog/YYYY-MM-DD.jsonl] + F --> H[tool_result/uuid.txt N-day TTL] +``` + +--- + +### 上下文机制（Pre-Reasoning Hook） + +1. **工具结果 Offload** (`ToolCallResultCompact`) +2. **上下文检查** (`ContextChecker`) +3. **若 Token 超阈值**： + - 保留最近 **X%** 的 Token（保障连贯性） + - 其余历史对话生成摘要 (`Compactor`) +4. **被摘要的上下文 Offload 到文件系统** (`SaveDialog`) + +```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 & Summary old context] + F --> G[SaveDialog: offload to file] + D -->|No| H[Proceed normally] +``` + +--- + +### 工具结果 Offload 机制 + +1. 所有工具调用结果先放入上下文，等待 Pre-Reasoning Hook 处理。 +2. 根据是否属于 **recent_n** 范围，决定截断策略： + - **recent_n 内**：近期内容 → 低截断比例 + - **recent_n 外**：远期内容 → 高截断比例 + +#### 示例：Browser Use 类工具 + +| 阶段 | 行为 | +|----|------------------------------------------------------------------------------| +| 1 | 原始工具调用结果 | +| 2 | 保存原始内容到文件：– 若在 recent_n 内：截断较少– 附注：“FullText saved to xxxx”– 提示：“请从第 N 行开始读” | +| 3 | 若再次引用且超出 recent_n：– 二次截断（更激进）– 仍指向原文件路径 | + +```mermaid +flowchart LR + A[Tool Call Result] --> B{Within recent_n?} + B -->|Yes| C[Low truncation
Save full text to tool_result/uuid.txt
Hint: 'Read from line N'] + B -->|No| D[High truncation
Reference existing file
More aggressive truncation] + C --> E[Context includes snippet + file ref] + D --> E +``` + +--- + +### ReadFile 工具调用结果变化示例 + +| 阶段 | 行为 | +|----|---------------------------------------------| +| 1 | 原始工具调用结果 | +| 2 | 若在 recent_n 内：– 不截断– 不保存文件（因内容已由用户指定） | +| 3 | 若超出 recent_n：– 二次截断（更小）– 保存 FullText 到文件并引用 | + +> 注：ReadFile 本身读取的是外部文件，因此首次调用通常无需重复保存。 + +```mermaid +flowchart LR + A[ReadFile Result] --> B{Within recent_n?} + B -->|Yes| C[No truncation
No file save needed] + B -->|No| D[Apply secondary truncation
Save FullText to tool_result/uuid.txt] + C --> E[Include full content in context] + D --> F[Include snippet + file ref] +``` + +--- + +## Copaw Memory + +### 触发逻辑 + +1. **主 Agent 主动写入**： + - `Memory.md`（长期记忆主干） + - `YYYY-MM-DD.md`（当日日志） +2. **触发阈值时**，由 **Summarizer（React Agent）** 写日志： + - 个性化信息（如偏好、习惯） + - Try-error 信息（失败尝试与修正） +3. **定时任务**（每日 00:00）： + - 汇总最近的 `YYYY-MM-DD.md` 文件 + - 更新 `Memory.md` + +```mermaid +flowchart TD + A[Main Agent] --> B[Write Memory.md] + A --> C[Write YYYY-MM-DD.md] + D[Context Threshold Reached?] -->|Yes| E[Summarizer Agent] + E --> F[Log: Personalization] + E --> G[Log: Try-Error Info] + H[Cron @ 00:00 daily] --> I[Aggregate recent YYYY-MM-DD.md] +I --> J[Update Memory.md] +``` diff --git a/reme/config/light.yaml b/reme/config/light.yaml index 242df716..3500f9a5 100644 --- a/reme/config/light.yaml +++ b/reme/config/light.yaml @@ -3,6 +3,8 @@ as_llms: backend: openai model_name: qwen3.5-plus +thread_pool_max_workers: -1 + as_llm_formatters: default: backend: openai diff --git a/reme/core/application.py b/reme/core/application.py index 6f04e663..e2e5b3c0 100644 --- a/reme/core/application.py +++ b/reme/core/application.py @@ -156,13 +156,15 @@ 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) @@ -518,7 +520,7 @@ class Application: 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/file_watcher/base_file_watcher.py b/reme/core/file_watcher/base_file_watcher.py index c9063997..383d5294 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, rebuild_index_on_start: bool = True, + poll_delay_ms: int = 1000, **kwargs, ): """ @@ -51,6 +52,7 @@ class BaseFileWatcher: callback: Callback function for changes 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 @@ -62,6 +64,7 @@ class BaseFileWatcher: self.file_store: BaseFileStore = file_store self.callback = callback 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() @@ -178,11 +181,15 @@ class BaseFileWatcher: try: logger.info(f"Starting watch on valid paths: {valid_paths}") + # Enable force_polling if poll_delay_ms > default 300ms to reduce CPU usage + force_polling = self.poll_delay_ms > 300 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, + force_polling=force_polling, stop_event=self._stop_event, ): if self._stop_event.is_set(): diff --git a/reme/core/op/base_op.py b/reme/core/op/base_op.py index a84ba18d..b016bb47 100644 --- a/reme/core/op/base_op.py +++ b/reme/core/op/base_op.py @@ -301,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/schema/service_config.py b/reme/core/schema/service_config.py index 758944b2..48b97e0a 100644 --- a/reme/core/schema/service_config.py +++ b/reme/core/schema/service_config.py @@ -116,7 +116,10 @@ class ServiceConfig(BasicConfig): 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) diff --git a/reme/memory/file_based/components/cli.py b/reme/memory/file_based/components/cli.py index 11d45035..138a3aad 100644 --- a/reme/memory/file_based/components/cli.py +++ b/reme/memory/file_based/components/cli.py @@ -113,9 +113,9 @@ class CliAgent(BaseOp): toolkit = Toolkit() file_io = FileIO(working_dir=self.working_dir) - toolkit.register_tool_function(file_io.read) - toolkit.register_tool_function(file_io.write) - toolkit.register_tool_function(file_io.edit) + 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 diff --git a/reme/memory/file_based/components/summarizer.yaml b/reme/memory/file_based/components/summarizer.yaml index 00dbdc36..14599f71 100644 --- a/reme/memory/file_based/components/summarizer.yaml +++ b/reme/memory/file_based/components/summarizer.yaml @@ -1,28 +1,29 @@ user_message: | - Memory Pre-compression Flush Cycle Initiated + 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. First, `read` {memory_dir}/YYYY-MM-DD.md (if the file doesn’t exist, an error message will be returned). - 2. Extract and synthesize content from the current session: + # 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. - 3. Intelligently merge new information with existing content (skip merging if the file doesn’t exist): + 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. - 4. 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. - Reflections MUST focus on forming reusable cognitive frameworks based on user feedback, aiming to improve future task execution. @@ -37,23 +38,23 @@ user_message_zh: | 当前日期：{date} 工作目录：{working_dir} + # 任务 立即存储持久化记忆与反思（使用路径 {memory_dir}/YYYY-MM-DD.md）。 - 工作流程： - 1. 先 `read` {memory_dir}/YYYY-MM-DD.md（如文件不存在，会返回错误提示） - 2. 从当前会话中提取并综合两类内容： + # 工作流程 + 1. 从当前会话中提取并综合两类内容： - 持久化记忆：客观事实、用户信息更新、项目状态及重要事件。 - 经验反思：基于用户反馈形成的可复用思考逻辑、成功的问题解决策略、犯下的错误/应避免的陷阱，以及对未来交互有帮助的行动指南。 - 3. 智能合并新信息与现有内容（若文件不存在则跳过合并）： + 2. `read` {memory_dir}/YYYY-MM-DD.md（如文件不存在，会返回错误提示） + - 若文件不存在，直接使用 `write` 工具写入。 + - 若文件已存在，智能合并新信息与现有内容，尽可能使用 `edit` 更新特定部分，仅在需要大幅重构时使用 `write` 覆盖整个文件。 + + # 原则 + - 智能合并新信息与现有内容： - 将内容进行清晰的分类（例如明确区分“事实记忆”与“反思与逻辑”）。 - 避免重复已记录的信息。 - 在相关时丰富现有条目的新细节。 - 在适用时保持时间顺序。 - 4. 写入更新后的内容： - - 尽可能使用 `edit` 更新特定部分。 - - 如需大幅重构则使用 `write` 覆盖整个文件。 - - 原则： - 始终保留时间戳、日期和时间相关上下文。 - 仅添加真正新的或有丰富价值的信息。 - 反思内容必须侧重于根据用户反馈构建可复用的思维逻辑，以改善未来的任务执行。 diff --git a/reme/memory/file_based/components/tool_result_compactor.py b/reme/memory/file_based/components/tool_result_compactor.py index 002bb8ae..cb53381f 100644 --- a/reme/memory/file_based/components/tool_result_compactor.py +++ b/reme/memory/file_based/components/tool_result_compactor.py @@ -1,33 +1,19 @@ """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 -from ....core.utils import truncate_text_head, TRUNCATION_MARKER_START logger = get_logger() -MAX_LINE_LENGTH = 10000 - - -def _split_long_lines(text: str, max_len: int = MAX_LINE_LENGTH) -> str: - """Split lines that exceed max_len by inserting newlines.""" - lines = text.split("\n") - result = [] - for line in lines: - if len(line) <= max_len: - result.append(line) - else: - # Split line into chunks of max_len - for i in range(0, len(line), max_len): - result.append(line[i : i + max_len]) - return "\n".join(result) - class ToolResultCompactor(BaseOp): """Truncate large tool_result outputs and save full content to files.""" @@ -35,62 +21,50 @@ class ToolResultCompactor(BaseOp): def __init__( self, tool_result_dir: str | Path, - retention_days: int = 7, + retention_days: int = 3, + old_max_bytes: int = 3000, + recent_max_bytes: int = DEFAULT_MAX_BYTES, recent_n: int = 1, - old_threshold: int = 500, - recent_threshold: int = 30000, + 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.old_threshold = old_threshold - self.recent_threshold = recent_threshold - - def _save_and_truncate(self, content: str, tool_name: str, threshold: int) -> str: - """Save full content to file and return truncated version with file reference.""" - if not content: - return content - - # Check if content was previously truncated - if TRUNCATION_MARKER_START in content: - parts = content.split(TRUNCATION_MARKER_START, 1) - if len(parts[0]) <= threshold: - return content - return f"{truncate_text_head(parts[0], threshold)}{parts[1]}" - - # Not truncated before - if len(content) <= threshold: - return content - - # Save full content with long lines split + self.encoding = encoding 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() - processed_content = _split_long_lines(content) - file_path.write_text( - f"# tool_name: {tool_name}\n# created_at: {created_at}\n# ---\n{processed_content}", - encoding="utf-8", - ) - logger.debug("Saved tool result to %s (len=%d)", file_path, len(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.""" - # Return truncated with file reference - return f"{truncate_text_head(content, threshold)}\n\n[Full content saved to: {file_path}]" + def _truncate(content: str) -> str: + if not content: + return content + + if TRUNCATION_NOTICE_MARKER in content: + return truncate_text_output(content, max_bytes=max_bytes) + + if len(content.encode(self.encoding)) <= max_bytes + 100: + return content + + saved_path: str | None = None + try: + fp = self.tool_result_dir / f"{uuid.uuid4().hex}.txt" + fp.write_text(content, encoding=self.encoding) + saved_path = str(fp) + except Exception as e: + logger.warning("Failed to save full tool result to file: %s", e) + + return truncate_text_output(content, 1, content.count("\n") + 1, max_bytes, file_path=saved_path) - def _process_output(self, output: str | list[dict], tool_name: str, threshold: int) -> str | list[dict]: - """Process tool result output, truncating if necessary.""" if isinstance(output, str): - return self._save_and_truncate(output, tool_name, threshold) - + return _truncate(output) if isinstance(output, list): return [ - ( - {**b, "text": self._save_and_truncate(b.get("text", ""), tool_name, threshold)} - if isinstance(b, dict) and b.get("type") == "text" - else b - ) + {**b, "text": _truncate(b.get("text", ""))} if isinstance(b, dict) and b.get("type") == "text" else b for b in output ] return output @@ -101,43 +75,54 @@ class ToolResultCompactor(BaseOp): if not messages: return messages - # Split messages into old and recent parts - split_index = max(0, len(messages) - self.recent_n) + 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)) for idx, msg in enumerate(messages): if not isinstance(msg.content, list): continue - - # Determine threshold based on message position - threshold = self.recent_threshold if idx >= split_index else self.old_threshold - + 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": - output = block.get("output") - if output: - block["output"] = self._process_output(output, block.get("name", "unknown"), threshold) + if isinstance(block, dict) and block.get("type") == "tool_result" and block.get("output"): + block["output"] = self._compact(block["output"], max_bytes) return messages def cleanup_expired_files(self) -> int: - """Clean up files older than retention_days.""" + """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 = 0 + deleted = failed = 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 + 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: - logger.warning("Failed to process %s: %s", fp, e) + failed += 1 + logger.warning("Failed to delete %s: %s", fp, e) - if deleted: - logger.info("Cleaned up %d expired files", deleted) + if deleted or failed: + logger.info("Cleaned up %d expired files (%d failed)", deleted, failed) return deleted diff --git a/reme/memory/file_based/reme_in_memory_memory.py b/reme/memory/file_based/reme_in_memory_memory.py index cd4d2d8b..1ee687fd 100644 --- a/reme/memory/file_based/reme_in_memory_memory.py +++ b/reme/memory/file_based/reme_in_memory_memory.py @@ -126,9 +126,10 @@ class ReMeInMemoryMemory(InMemoryMemory): if prepend_summary and self._compressed_summary: previous_summary = f""" +Raw conversation logs are in dialog/YYYY-MM-DD.jsonl (or nearby date files). +Entries are chronological; read from the end for recent history. {self._compressed_summary} -The above is a summary of previous conversation, use it as context to maintain continuity. - """.strip() +The above is a summary of previous conversation, use it as context to maintain continuity.""".strip() return [ Msg( diff --git a/reme/memory/file_based/tools/file_io.py b/reme/memory/file_based/tools/file_io.py index 2b792475..6f4ea9cf 100644 --- a/reme/memory/file_based/tools/file_io.py +++ b/reme/memory/file_based/tools/file_io.py @@ -7,7 +7,7 @@ from typing import Optional from agentscope.message import TextBlock from agentscope.tool import ToolResponse -from ..utils import DEFAULT_MAX_BYTES, read_file_safe, truncate_output +from ..utils import read_file_safe, truncate_text_output class FileIO: @@ -38,13 +38,13 @@ class FileIO: else: return str(self.working_dir / file_path) - async def read( # pylint: disable=too-many-return-statements + 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. + """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. @@ -57,6 +57,34 @@ class FileIO: 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): @@ -111,29 +139,21 @@ class FileIO: # Extract selected lines selected_content = "\n".join(all_lines[s - 1 : e]) - # Apply smart truncation (keep head for file reading) - truncated, was_truncated, output_lines, reason = truncate_output(selected_content, keep="head") + # Apply smart truncation (consistent with shell output format) + text = truncate_text_output( + selected_content, + start_line=s, + total_lines=total, + file_path=file_path, + ) - # Build response with truncation hints - if was_truncated: - end_display = s + output_lines - 1 - next_line = end_display + 1 - if reason == "lines": - hint = f"\n\n[Lines {s}-{end_display} of {total}. Use start_line={next_line} to continue.]" - else: - hint = ( - f"\n\n[Lines {s}-{end_display} of {total} ({DEFAULT_MAX_BYTES // 1024}KB limit). " - f"Use start_line={next_line} to continue.]" - ) - text = truncated + hint - elif e < total: + # Add continuation hint if partial read without truncation + if text == selected_content and e < total: remaining = total - e text = ( - f"{file_path} (lines {s}-{e} of {total})\n{truncated}\n\n[{remaining} more lines. " - f"Use start_line={e + 1} to continue.]" + f"{file_path} (lines {s}-{e} of {total})\n{text}\n\n" + f"[{remaining} more lines. Use start_line={e + 1} to continue.]" ) - else: - text = truncated return ToolResponse( content=[TextBlock(type="text", text=text)], @@ -149,7 +169,7 @@ class FileIO: ], ) - async def write( + async def write_file( self, file_path: str, content: str, @@ -195,7 +215,7 @@ class FileIO: ], ) - async def edit( + async def edit_file( self, file_path: str, old_text: str, @@ -212,7 +232,7 @@ class FileIO: new_text (`str`): Replacement text. """ - response = await self.read(file_path=file_path) + response = await self.read_file(file_path=file_path) if response.content and len(response.content) > 0: error_text = response.content[0].get("text", "") if error_text.startswith("Error:"): @@ -239,7 +259,7 @@ class FileIO: ) new_content = content.replace(old_text, new_text) - write_response = await self.write(file_path=file_path, content=new_content) + write_response = await self.write_file(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", "") diff --git a/reme/memory/file_based/tools/shell.py b/reme/memory/file_based/tools/shell.py index 2bee1bd6..cce138f9 100644 --- a/reme/memory/file_based/tools/shell.py +++ b/reme/memory/file_based/tools/shell.py @@ -12,8 +12,6 @@ from pathlib import Path from agentscope.message import TextBlock from agentscope.tool import ToolResponse -from ..utils import truncate_shell_output - def _execute_subprocess_sync( cmd: str, @@ -189,10 +187,6 @@ class Shell: stdout_str = "" stderr_str = stderr_suffix - # Apply output truncation - stdout_str = truncate_shell_output(stdout_str) - stderr_str = truncate_shell_output(stderr_str) - # Format the response in a human-friendly way if returncode == 0: # Success case: just show the output diff --git a/reme/memory/file_based/utils/__init__.py b/reme/memory/file_based/utils/__init__.py index 48231232..769a8a66 100644 --- a/reme/memory/file_based/utils/__init__.py +++ b/reme/memory/file_based/utils/__init__.py @@ -1,13 +1,12 @@ """utils""" from .as_msg_handler import AsMsgHandler -from .file_utils import truncate_output, truncate_shell_output, read_file_safe, DEFAULT_MAX_BYTES, DEFAULT_MAX_LINES +from .file_utils import truncate_text_output, read_file_safe, DEFAULT_MAX_BYTES, TRUNCATION_NOTICE_MARKER __all__ = [ "AsMsgHandler", - "truncate_output", - "truncate_shell_output", + "truncate_text_output", "read_file_safe", "DEFAULT_MAX_BYTES", - "DEFAULT_MAX_LINES", + "TRUNCATION_NOTICE_MARKER", ] diff --git a/reme/memory/file_based/utils/file_utils.py b/reme/memory/file_based/utils/file_utils.py index 17f58877..f6324296 100644 --- a/reme/memory/file_based/utils/file_utils.py +++ b/reme/memory/file_based/utils/file_utils.py @@ -1,112 +1,141 @@ +# -*- coding: utf-8 -*- """Shared utilities for file and shell tools.""" -# Default truncation limits -DEFAULT_MAX_LINES = 1000 -DEFAULT_MAX_BYTES = 30 * 1024 # 30KB +import re + +from ....core.utils import get_logger + +logger = get_logger() + +# Default truncation limit +DEFAULT_MAX_BYTES = 100 * 1024 + +# Maximum file size to read into memory (1GB) +MAX_FILE_READ_BYTES = 1024 * 1024 * 1024 + +# Marker prepended to every truncation notice. +# Format: <<>> +# File: file_path +# Content from start_line=X, next N bytes. +# total_lines=Z +# Use start_line=Y to continue. +# Split on this to recover the original (un-truncated) portion: +# original = output.split(TRUNCATION_NOTICE_MARKER)[0] +TRUNCATION_NOTICE_MARKER = "<<>>" -def truncate_output( +# pylint: disable=too-many-return-statements +def truncate_text_output( text: str, - max_lines: int = DEFAULT_MAX_LINES, + start_line: int = 0, + total_lines: int = 0, max_bytes: int = DEFAULT_MAX_BYTES, - keep: str = "head", -) -> tuple[str, bool, int, str]: - """Smart truncation for large content. + file_path: str | None = None, +) -> str: + """Truncate file output by bytes with line integrity. - Args: - text: Text content to truncate. - max_lines: Maximum number of lines. - max_bytes: Maximum size in bytes. - keep: Which part to keep - "head" (first lines) or "tail" (last lines). + 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. - Returns: - (truncated_content, was_truncated, output_line_count, truncate_reason) - """ - if not text: - return text, False, 0, "" - - lines = text.split("\n") - total_lines = len(lines) - - # No truncation needed - if total_lines <= max_lines and len(text.encode("utf-8")) <= max_bytes: - return text, False, total_lines, "" - - # Apply line limit - if total_lines > max_lines: - if keep == "tail": - lines = lines[-max_lines:] - else: - lines = lines[:max_lines] - reason = "lines" - else: - reason = "" - - # Apply byte limit - if len("\n".join(lines).encode("utf-8")) > max_bytes: - if keep == "tail": - while lines and len("\n".join(lines).encode("utf-8")) > max_bytes: - lines.pop(0) - else: - truncated = [] - current_bytes = 0 - for line in lines: - line_bytes = len(line.encode("utf-8")) + 1 - if current_bytes + line_bytes > max_bytes: - break - truncated.append(line) - current_bytes += line_bytes - lines = truncated - reason = "bytes" - - return "\n".join(lines), True, len(lines), reason - - -def truncate_shell_output(text: str) -> str: - """Truncate shell output to last N lines or M bytes, with truncation notice. + If TRUNCATION_NOTICE_MARKER is already in text (previously truncated), + extract the original content, re-truncate it, and update the + truncation notice using regex. 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. Returns: Truncated text with notice if truncated. """ if not text: return text + if max_bytes <= 0: + return text try: - total_lines = len(text.split("\n")) - truncated, was_truncated, output_lines, reason = truncate_output(text, keep="tail") + if TRUNCATION_NOTICE_MARKER in text: + parts = text.split(TRUNCATION_NOTICE_MARKER, 1) + original_content = parts[0] + old_notice = parts[1] - if not was_truncated: - return text + text_bytes = original_content.encode("utf-8") + + # Allow a small slack to avoid re-truncating near-limit content + if len(text_bytes) <= max_bytes + 100: + return text + + # Parse start_line and total_lines from notice; return text unchanged if not found + start_match = re.search(r"start_line=(\d+),", old_notice) + total_match = re.search(r"total_lines=(\d+)", old_notice) + if not start_match or not total_match: + return text + start_line_parsed = int(start_match.group(1)) + total_lines_parsed = int(total_match.group(1)) + + truncated_bytes = text_bytes[:max_bytes] + result = truncated_bytes.decode("utf-8", errors="ignore") + newline_count = result.count("\n") + + next_line = start_line_parsed + max(1, newline_count) + + if not re.search(r"next \d+ bytes", old_notice): + return text + has_continuation = bool(re.search(r"Use start_line=\d+", old_notice)) + new_notice = re.sub(r"next \d+ bytes", f"next {max_bytes} bytes", old_notice) + if has_continuation: + new_notice = re.sub(r"Use start_line=\d+", f"Use start_line={next_line}", new_notice) + elif next_line <= total_lines_parsed: + new_notice = re.sub(r"(total_lines=\d+)", f"\\1\nUse start_line={next_line} to continue.", new_notice) + + return result + TRUNCATION_NOTICE_MARKER + new_notice - start_line = total_lines - output_lines + 1 - if reason == "lines": - notice = f"\n\n[Output truncated: showing lines {start_line}-{total_lines} of {total_lines} total]" else: + text_bytes = text.encode("utf-8") + + if len(text_bytes) <= max_bytes: + return text + + truncated = text_bytes[:max_bytes] + result = truncated.decode("utf-8", errors="ignore") + + newline_count = result.count("\n") + + next_line = start_line + max(1, newline_count) + + continuation = f"\nUse start_line={next_line} to continue." if next_line <= total_lines else "" notice = ( - f"\n\n[Output truncated: showing lines {start_line}-{total_lines} of {total_lines} " - f"({DEFAULT_MAX_BYTES // 1024}KB limit)]" + TRUNCATION_NOTICE_MARKER + + f"\n\nFile: {file_path or ''}\nContent from start_line={start_line}, next {max_bytes} bytes." + f"\ntotal_lines={total_lines}{continuation}" ) - return truncated + notice + return result + notice + except Exception: + logger.warning("truncate_text_output failed, returning original text", exc_info=True) return text -def read_file_safe(file_path: str) -> str: - """Read file with Unicode error handling. +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. + File content as string (up to max_bytes). """ try: with open(file_path, "r", encoding="utf-8") as f: - return f.read() + return f.read(max_bytes) except UnicodeDecodeError: with open(file_path, "r", encoding="utf-8", errors="ignore") as f: - return f.read() + return f.read(max_bytes) diff --git a/reme/reme_light.py b/reme/reme_light.py index d9835f5f..11f9ce56 100644 --- a/reme/reme_light.py +++ b/reme/reme_light.py @@ -168,7 +168,7 @@ class ReMeLight(Application): Returns: Computed compaction threshold as an integer. """ - return int(max_input_length * compact_ratio * 0.9) + return int(max_input_length * compact_ratio * 0.95) def _cleanup_tool_results(self) -> int: """ @@ -231,10 +231,10 @@ class ReMeLight(Application): 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, - old_threshold: int = 500, - recent_threshold: int = 30000, - retention_days: int = 7, ) -> list[Msg]: """ Compact tool results by truncating large outputs and saving full content to files. @@ -247,30 +247,37 @@ class ReMeLight(Application): Args: messages (list[Msg]): List of messages potentially containing tool results that may need compaction. - recent_n (int): Number of recent messages to use recent_threshold for. - Default 1. - old_threshold (int): Character threshold for old messages. Default 500. - recent_threshold (int): Character threshold for recent messages. Default 30000. + 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 7. + 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 list of messages with large tool results compacted. If an error occurs, returns the original unmodified messages. Note: - - Tool results are truncated based on old_threshold/recent_threshold - - Full content of truncated results is saved to tool_result_path - - Expired files are automatically cleaned up during this operation + - 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, retention_days=retention_days, + old_max_bytes=old_max_bytes, + recent_max_bytes=recent_max_bytes, recent_n=recent_n, - old_threshold=old_threshold, - recent_threshold=recent_threshold, ) # Execute compaction and get processed messages @@ -455,9 +462,9 @@ class ReMeLight(Application): if toolkit is None: toolkit = Toolkit() file_io = FileIO(working_dir=str(self.working_path)) - toolkit.register_tool_function(file_io.read) - toolkit.register_tool_function(file_io.write) - toolkit.register_tool_function(file_io.edit) + 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), diff --git a/tests/light/test_summarizer.py b/tests/light/test_summarizer.py index 475d81cd..5f8a67a1 100644 --- a/tests/light/test_summarizer.py +++ b/tests/light/test_summarizer.py @@ -101,9 +101,9 @@ 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) - toolkit.register_tool_function(file_io.write) - toolkit.register_tool_function(file_io.edit) + 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 diff --git a/tests/light/test_tool_result_compactor.py b/tests/light/test_tool_result_compactor.py index 1697911d..ac812c35 100644 --- a/tests/light/test_tool_result_compactor.py +++ b/tests/light/test_tool_result_compactor.py @@ -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,7 +45,7 @@ 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)] @@ -68,7 +68,7 @@ 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) + op = ToolResultCompactor(tool_result_dir=tmpdir, recent_max_bytes=100) truncated_content = "head<<>>(100 chars omitted)<<>>tail" messages = [create_tool_result_msg(truncated_content)] @@ -80,7 +80,7 @@ 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)] @@ -93,7 +93,7 @@ class TestToolResultCompactor: 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"}, @@ -124,7 +124,7 @@ class TestToolResultCompactor: 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"}}, @@ -141,7 +141,7 @@ class TestToolResultCompactor: 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 index e35bde22..779c59a2 100644 --- a/tests/light/test_tools.py +++ b/tests/light/test_tools.py @@ -4,6 +4,7 @@ import asyncio import os +import re import shutil import tempfile @@ -11,7 +12,7 @@ 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_LINES, DEFAULT_MAX_BYTES +from reme.memory.file_based.utils import DEFAULT_MAX_BYTES # ============ Shell Tests ============ @@ -76,23 +77,6 @@ def test_shell_multiline_output(shell_env): assert "line3" in text -def test_shell_truncated_output(shell_env): - """Test output truncation for large output.""" - lines_to_generate = DEFAULT_MAX_LINES + 500 - cmd = f"seq 1 {lines_to_generate}" - result = asyncio.run(shell_env["shell"].execute_shell_command(cmd)) - text = result.content[0].get("text", "") - - # Should contain truncation notice - assert "truncated" in text.lower() - # Should contain the last line (tail is kept) - assert str(lines_to_generate) in text - # Verify first numeric line is > 1 (truncated from head) - numeric_lines = [tl for tl in text.strip().split("\n") if tl.isdigit()] - if numeric_lines: - assert int(numeric_lines[0]) > 1 - - def test_shell_timeout(shell_env): """Test command timeout handling.""" result = asyncio.run( @@ -116,13 +100,17 @@ def fileio_env(): with open(simple_file, "w", encoding="utf-8") as f: f.write("line1\nline2\nline3\nline4\nline5") - # Create large file (exceeds DEFAULT_MAX_LINES) + # 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: - for i in range(1, DEFAULT_MAX_LINES + 500): + # 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" @@ -130,19 +118,31 @@ def fileio_env(): 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(fileio_env["simple_file"])) + 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 @@ -150,14 +150,14 @@ def test_read_file_success(fileio_env): def test_read_file_relative_path(fileio_env): """Test reading file with relative path.""" - result = asyncio.run(fileio_env["file_io"].read("simple.txt")) + 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("nonexistent.txt")) + 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 @@ -166,7 +166,7 @@ def test_read_file_not_exists(fileio_env): def test_read_file_with_line_range(fileio_env): """Test reading specific line range.""" result = asyncio.run( - fileio_env["file_io"].read(fileio_env["simple_file"], start_line=2, end_line=4), + 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 @@ -177,7 +177,7 @@ def test_read_file_with_line_range(fileio_env): def test_read_file_start_line_exceeds(fileio_env): """Test start_line exceeding file length.""" result = asyncio.run( - fileio_env["file_io"].read(fileio_env["simple_file"], start_line=100), + fileio_env["file_io"].read_file(fileio_env["simple_file"], start_line=100), ) text = result.content[0].get("text", "") assert "Error" in text @@ -187,15 +187,15 @@ def test_read_file_start_line_exceeds(fileio_env): def test_read_file_invalid_range(fileio_env): """Test invalid line range (start > end).""" result = asyncio.run( - fileio_env["file_io"].read(fileio_env["simple_file"], start_line=4, end_line=2), + 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_by_lines(fileio_env): - """Test file truncation by line limit.""" - result = asyncio.run(fileio_env["file_io"].read(fileio_env["large_file"])) +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() @@ -203,19 +203,140 @@ def test_read_file_truncated_by_lines(fileio_env): def test_read_file_truncated_by_bytes(fileio_env): """Test file truncation by byte limit.""" - result = asyncio.run(fileio_env["file_io"].read(fileio_env["large_bytes_file"])) + 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() or "KB" in 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(fileio_env["dir"])) + 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 ============ @@ -231,7 +352,7 @@ def write_env(): 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_path, "test content")) + result = asyncio.run(write_env["file_io"].write_file(file_path, "test content")) text = result.content[0].get("text", "") assert "Wrote" in text @@ -245,7 +366,7 @@ def test_write_overwrite_file(write_env): with open(file_path, "w", encoding="utf-8") as f: f.write("old content") - result = asyncio.run(write_env["file_io"].write(file_path, "new content")) + result = asyncio.run(write_env["file_io"].write_file(file_path, "new content")) text = result.content[0].get("text", "") assert "Wrote" in text @@ -255,14 +376,14 @@ def test_write_overwrite_file(write_env): def test_write_empty_path(write_env): """Test writing with empty path.""" - result = asyncio.run(write_env["file_io"].write("", "content")) + 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("relative.txt", "relative content")) + result = asyncio.run(write_env["file_io"].write_file("relative.txt", "relative content")) text = result.content[0].get("text", "") assert "Wrote" in text @@ -290,7 +411,7 @@ def edit_env(): def test_edit_replace_text(edit_env): """Test replacing text in file.""" result = asyncio.run( - edit_env["file_io"].edit(edit_env["edit_file"], "Hello", "Hi"), + edit_env["file_io"].edit_file(edit_env["edit_file"], "Hello", "Hi"), ) text = result.content[0].get("text", "") assert "Successfully" in text @@ -305,7 +426,7 @@ def test_edit_replace_text(edit_env): def test_edit_text_not_found(edit_env): """Test editing when text not found.""" result = asyncio.run( - edit_env["file_io"].edit(edit_env["edit_file"], "NotExists", "Replacement"), + edit_env["file_io"].edit_file(edit_env["edit_file"], "NotExists", "Replacement"), ) text = result.content[0].get("text", "") assert "Error" in text @@ -315,7 +436,7 @@ def test_edit_text_not_found(edit_env): def test_edit_nonexistent_file(edit_env): """Test editing non-existent file.""" result = asyncio.run( - edit_env["file_io"].edit("nonexistent.txt", "old", "new"), + edit_env["file_io"].edit_file("nonexistent.txt", "old", "new"), ) text = result.content[0].get("text", "") assert "Error" in text From f17028e1b22e885e1d95ee7976850073af8ddbee Mon Sep 17 00:00:00 2001 From: "jinli.yl" Date: Wed, 25 Mar 2026 20:22:46 +0800 Subject: [PATCH 46/59] chore(release): bump version to 0.3.1.4 --- reme/__init__.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/reme/__init__.py b/reme/__init__.py index 28d87f9e..85d6b7b7 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.1.3" +__version__ = "0.3.1.4" __all__ = [ "config", From dc8eab56a127708e22752dfe1202421744ffd137 Mon Sep 17 00:00:00 2001 From: jinliyl <6469360+jinliyl@users.noreply.github.com> Date: Thu, 26 Mar 2026 12:10:24 +0800 Subject: [PATCH 47/59] refactor(core): replace text truncation utilities with new marker system (#179) * refactor(core): replace text truncation utilities with new marker system - Remove old truncate_text_utils module and its exports - Replace TRUNCATION_MARKER_START with _TRUNCATION_NOTICE_MARKER constant - Update as_msg_stat.py to split content using new marker format - Modify FileIO tool to use TRUNCATION_NOTICE_MARKER for continuation hints - Change is_truncated function checks to use marker presence detection - Move transformers dependency from main deps to light extra dependencies - Update tool result compactor tests to verify marker instead of is_truncated calls * feat(file_io): enhance file operations with path resolution and append functionality - Add expanduser() to resolve file paths with ~ symbol - Implement proper file existence and type validation in update_file - Add new append_file method to append content to files - Update truncation notice format for better readability - Fix typo in error message from "provide" to "provided" - Update transformers dependency in pyproject.toml - Remove duplicate transformers dependency from light extras * refactor(file_io): disable pylint too-many-return-statements warning * perf(file_watcher): increase default polling delay and optimize watcher configuration - Increased default poll_delay_ms from 1000ms to 2000ms to reduce CPU usage - Removed force_polling parameter as it's no longer needed with updated polling strategy - Simplified async watch configuration by removing conditional force_polling logic - Reduced overall system resource consumption during file watching operations * refactor(memory): update conversation log documentation in memory summary - Changed "Raw conversation logs" to "Earlier conversation logs" for clarity - Added warning note about potentially large dialog file sizes - Improved formatting with additional line break for better readability - Maintained existing compressed summary integration unchanged * feat(memory): add long-term memory support to file-based memory system - Initialize _long_term_memory attribute as empty string - Add memories section to content when long-term memory exists - Consolidate summary and memories into single user message - Format memories with markdown header # Memories - Maintain existing compressed summary functionality - Join multiple content parts with double newlines --- pyproject.toml | 2 +- reme/core/file_watcher/base_file_watcher.py | 5 +- reme/core/schema/as_msg_stat.py | 5 +- reme/core/utils/__init__.py | 5 - reme/core/utils/truncate_text_utils.py | 82 ------------- .../file_based/reme_in_memory_memory.py | 26 ++-- reme/memory/file_based/tools/file_io.py | 115 +++++++++++++++--- reme/memory/file_based/utils/file_utils.py | 24 ++-- tests/light/test_tool_result_compactor.py | 12 +- 9 files changed, 136 insertions(+), 140 deletions(-) delete mode 100644 reme/core/utils/truncate_text_utils.py diff --git a/pyproject.toml b/pyproject.toml index ececfe6d..36b12e98 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -80,11 +80,11 @@ full = [ litellm = [ "litellm==1.80.0", - "flowllm[reme]>=0.2.0.10", ] light = [ "agentscope==1.0.17", + "flowllm[reme]>=0.2.0.10", ] [tool.setuptools.packages.find] diff --git a/reme/core/file_watcher/base_file_watcher.py b/reme/core/file_watcher/base_file_watcher.py index 383d5294..b4644dd7 100644 --- a/reme/core/file_watcher/base_file_watcher.py +++ b/reme/core/file_watcher/base_file_watcher.py @@ -35,7 +35,7 @@ class BaseFileWatcher: file_store: BaseFileStore | None = None, callback: Callable[[set[tuple[Change, str]]], None | Coroutine[Any, Any, None]] | None = None, rebuild_index_on_start: bool = True, - poll_delay_ms: int = 1000, + poll_delay_ms: int = 2000, **kwargs, ): """ @@ -181,15 +181,12 @@ class BaseFileWatcher: try: logger.info(f"Starting watch on valid paths: {valid_paths}") - # Enable force_polling if poll_delay_ms > default 300ms to reduce CPU usage - force_polling = self.poll_delay_ms > 300 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, - force_polling=force_polling, stop_event=self._stop_event, ): if self._stop_event.is_set(): diff --git a/reme/core/schema/as_msg_stat.py b/reme/core/schema/as_msg_stat.py index b502dd02..a8d2daaa 100644 --- a/reme/core/schema/as_msg_stat.py +++ b/reme/core/schema/as_msg_stat.py @@ -2,6 +2,8 @@ from pydantic import BaseModel, Field +_TRUNCATION_NOTICE_MARKER = "<<>>" + _DEFAULT_MAX_BLOCK_TEXT_PREVIEW_LENGTH = 100 _DEFAULT_MAX_FORMATTER_TEXT_LENGTH = 1000 @@ -61,7 +63,8 @@ class AsBlockStat(BaseModel): if self.block_type == "tool_result": if not self.tool_output: return "" - content = f"{self.tool_name} output={self._truncate(self.tool_output, max_length)}" + 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 "" diff --git a/reme/core/utils/__init__.py b/reme/core/utils/__init__.py index b51936c7..26ca2827 100644 --- a/reme/core/utils/__init__.py +++ b/reme/core/utils/__init__.py @@ -19,7 +19,6 @@ 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 -from .truncate_text_utils import truncate_text, truncate_text_head, is_truncated, TRUNCATION_MARKER_START __all__ = [ "convert_dashscope_to_agentscope", @@ -51,8 +50,4 @@ __all__ = [ "timer", "get_now_time", "get_hf_token_counter", - "truncate_text", - "truncate_text_head", - "is_truncated", - "TRUNCATION_MARKER_START", ] diff --git a/reme/core/utils/truncate_text_utils.py b/reme/core/utils/truncate_text_utils.py deleted file mode 100644 index 1f60d19f..00000000 --- a/reme/core/utils/truncate_text_utils.py +++ /dev/null @@ -1,82 +0,0 @@ -"""Utility functions for truncating long text strings.""" - -from .std_logger import get_logger - -logger = get_logger() - -TRUNCATION_MARKER_START = "<<>>" -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 truncate_text_head(text: str, max_length: int) -> str: - """Truncate text from the beginning, keeping only the head portion. - - Args: - text: The text to truncate - max_length: Maximum allowed length - - Returns: - Truncated text with marker indicating truncation at the end - """ - text = str(text) if text else "" - if not text: - return text - - if len(text) <= max_length: - return text - - truncated_chars = len(text) - max_length - logger.debug( - "Text truncated from head: original %d chars, kept %d, removed %d chars from tail.", - len(text), - max_length, - truncated_chars, - ) - return f"{text[:max_length]}{TRUNCATION_MARKER_START}" - - -def is_truncated(text: str) -> bool: - """Check if the text has been truncated (contains truncation marker). - - Args: - text: The text to check - - Returns: - bool: True if text contains truncation marker, False otherwise - """ - if not text: - return False - return TRUNCATION_MARKER_START in text diff --git a/reme/memory/file_based/reme_in_memory_memory.py b/reme/memory/file_based/reme_in_memory_memory.py index 1ee687fd..d8293e0e 100644 --- a/reme/memory/file_based/reme_in_memory_memory.py +++ b/reme/memory/file_based/reme_in_memory_memory.py @@ -34,6 +34,7 @@ class ReMeInMemoryMemory(InMemoryMemory): self._token_counter: HuggingFaceTokenCounter = token_counter 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. @@ -124,21 +125,20 @@ class ReMeInMemoryMemory(InMemoryMemory): """ 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""" -Raw conversation logs are in dialog/YYYY-MM-DD.jsonl (or nearby date files). -Entries are chronological; read from the end for recent history. -{self._compressed_summary} -The above is a summary of 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] diff --git a/reme/memory/file_based/tools/file_io.py b/reme/memory/file_based/tools/file_io.py index 6f4ea9cf..d876407c 100644 --- a/reme/memory/file_based/tools/file_io.py +++ b/reme/memory/file_based/tools/file_io.py @@ -7,7 +7,7 @@ from typing import Optional from agentscope.message import TextBlock from agentscope.tool import ToolResponse -from ..utils import read_file_safe, truncate_text_output +from ..utils import read_file_safe, truncate_text_output, TRUNCATION_NOTICE_MARKER class FileIO: @@ -32,7 +32,7 @@ class FileIO: Returns: The resolved absolute file path as string. """ - path = Path(file_path) + path = Path(file_path).expanduser() if path.is_absolute(): return str(path) else: @@ -147,13 +147,18 @@ class FileIO: file_path=file_path, ) - # Add continuation hint if partial read without truncation + # 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: - remaining = total - e - text = ( - f"{file_path} (lines {s}-{e} of {total})\n{text}\n\n" - f"[{remaining} more lines. Use start_line={e + 1} to continue.]" + content_bytes = len(text.encode("utf-8")) + notice = ( + TRUNCATION_NOTICE_MARKER + + f"\nFile: {file_path}\nStarting at start_line={s}, next {content_bytes} bytes." + f"\nTotal lines: {total}" + f"\nUse start_line={e + 1} to continue." ) + text = text + notice return ToolResponse( content=[TextBlock(type="text", text=text)], @@ -187,7 +192,7 @@ class FileIO: content=[ TextBlock( type="text", - text="Error: No `file_path` provide.", + text="Error: No `file_path` provided.", ), ], ) @@ -215,6 +220,7 @@ class FileIO: ], ) + # pylint: disable=too-many-return-statements async def edit_file( self, file_path: str, @@ -232,22 +238,50 @@ class FileIO: new_text (`str`): Replacement text. """ - response = await self.read_file(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: + if not file_path: return ToolResponse( content=[ TextBlock( type="text", - text=f"Error: Failed to read file {file_path}.", + 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}", ), ], ) - content = response.content[0].get("text", "") if old_text not in content: return ToolResponse( content=[ @@ -259,7 +293,7 @@ class FileIO: ) new_content = content.replace(old_text, new_text) - write_response = await self.write_file(file_path=file_path, content=new_content) + 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", "") @@ -274,3 +308,50 @@ class FileIO: ), ], ) + + 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/file_based/utils/file_utils.py b/reme/memory/file_based/utils/file_utils.py index f6324296..ee556bba 100644 --- a/reme/memory/file_based/utils/file_utils.py +++ b/reme/memory/file_based/utils/file_utils.py @@ -14,12 +14,14 @@ DEFAULT_MAX_BYTES = 100 * 1024 MAX_FILE_READ_BYTES = 1024 * 1024 * 1024 # Marker prepended to every truncation notice. -# Format: <<>> -# File: file_path -# Content from start_line=X, next N bytes. -# total_lines=Z -# Use start_line=Y to continue. -# Split on this to recover the original (un-truncated) portion: +# Format: +# <<>> +# File: +# Starting at start_line=X, next N bytes. +# Total lines: Z +# Use start_line=Y to continue. +# +# Split output on this marker to recover the original (untruncated) portion: # original = output.split(TRUNCATION_NOTICE_MARKER)[0] TRUNCATION_NOTICE_MARKER = "<<>>" @@ -72,8 +74,8 @@ def truncate_text_output( return text # Parse start_line and total_lines from notice; return text unchanged if not found - start_match = re.search(r"start_line=(\d+),", old_notice) - total_match = re.search(r"total_lines=(\d+)", old_notice) + start_match = re.search(r"Starting at start_line=(\d+)", old_notice) + total_match = re.search(r"Total lines: (\d+)", old_notice) if not start_match or not total_match: return text start_line_parsed = int(start_match.group(1)) @@ -92,7 +94,7 @@ def truncate_text_output( if has_continuation: new_notice = re.sub(r"Use start_line=\d+", f"Use start_line={next_line}", new_notice) elif next_line <= total_lines_parsed: - new_notice = re.sub(r"(total_lines=\d+)", f"\\1\nUse start_line={next_line} to continue.", new_notice) + new_notice = re.sub(r"(Total lines: \d+)", f"\\1\nUse start_line={next_line} to continue.", new_notice) return result + TRUNCATION_NOTICE_MARKER + new_notice @@ -112,8 +114,8 @@ def truncate_text_output( continuation = f"\nUse start_line={next_line} to continue." if next_line <= total_lines else "" notice = ( TRUNCATION_NOTICE_MARKER - + f"\n\nFile: {file_path or ''}\nContent from start_line={start_line}, next {max_bytes} bytes." - f"\ntotal_lines={total_lines}{continuation}" + + f"\nFile: {file_path or ''}\nStarting at start_line={start_line}, next {max_bytes} bytes." + f"\nTotal lines: {total_lines}{continuation}" ) return result + notice diff --git a/tests/light/test_tool_result_compactor.py b/tests/light/test_tool_result_compactor.py index ac812c35..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.core.utils import is_truncated 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: @@ -52,7 +52,7 @@ class TestToolResultCompactor: _ = asyncio.run(op.call(messages=messages)) output = messages[0].content[0]["output"] - assert is_truncated(output) + assert TRUNCATION_NOTICE_MARKER in output assert "[Full content saved to:" in output # Verify file was created @@ -87,7 +87,7 @@ class TestToolResultCompactor: asyncio.run(op.call(messages=messages)) text_block = messages[0].content[0]["output"][0] - assert is_truncated(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): @@ -116,9 +116,9 @@ class TestToolResultCompactor: asyncio.run(op.call(messages=messages)) output = messages[0].content[0]["output"] - assert is_truncated(output[0]["text"]) + assert TRUNCATION_NOTICE_MARKER in output[0]["text"] assert output[1]["text"] == "short" # unchanged - assert is_truncated(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): @@ -134,7 +134,7 @@ class TestToolResultCompactor: asyncio.run(op.call(messages=messages)) output = messages[0].content[0]["output"] - assert is_truncated(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 From 03cbc42b25e096341c331c4a61e74b344a213ae9 Mon Sep 17 00:00:00 2001 From: Sen Huang <48879559+ployts@users.noreply.github.com> Date: Thu, 26 Mar 2026 14:57:27 +0800 Subject: [PATCH 48/59] docs(README): add Trendshift repository badge (#180) --- README.md | 4 ++++ README_ZH.md | 4 ++++ 2 files changed, 8 insertions(+) diff --git a/README.md b/README.md index ca2538b6..1d0790c9 100644 --- a/README.md +++ b/README.md @@ -17,6 +17,10 @@ DeepWiki

+

+agentscope-ai%2FReMe | Trendshift +

+

A memory management toolkit for AI agents — Remember Me, Refine Me.

diff --git a/README_ZH.md b/README_ZH.md index 32a89039..ee455318 100644 --- a/README_ZH.md +++ b/README_ZH.md @@ -17,6 +17,10 @@ DeepWiki

+

+agentscope-ai%2FReMe | Trendshift +

+

面向智能体的记忆管理工具包，Remember Me, Refine Me.

From 9cb8dc834eb6e94b9b2dbb1c375014d8e1583242 Mon Sep 17 00:00:00 2001 From: jinliyl <6469360+jinliyl@users.noreply.github.com> Date: Thu, 26 Mar 2026 15:46:42 +0800 Subject: [PATCH 49/59] feat(reme): add configurable file watcher support (#181) - Add default_file_watcher_config parameter to RemeLight constructor - Document file watcher configuration options in docstring - Implement logic to merge custom watch paths with default paths - Set up default watch paths including MEMORY.md, memory.md and memory directory - Replace hardcoded file watcher config with dynamic merged configuration --- reme/reme_light.py | 27 ++++++++++++++++++++------- 1 file changed, 20 insertions(+), 7 deletions(-) diff --git a/reme/reme_light.py b/reme/reme_light.py index 11f9ce56..5cad2f64 100644 --- a/reme/reme_light.py +++ b/reme/reme_light.py @@ -74,6 +74,7 @@ class ReMeLight(Application): 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, enable_load_env: bool = False, @@ -101,6 +102,10 @@ class ReMeLight(Application): 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. @@ -130,6 +135,20 @@ class ReMeLight(Application): self.vector_weight: float = vector_weight self.candidate_multiplier: float = candidate_multiplier + # 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: + _merged_file_watcher_config = { + **(default_file_watcher_config or {}), + "watch_paths": _default_watch_paths, + } + # Initialize the parent Application class with comprehensive configuration super().__init__( llm_api_key=llm_api_key, @@ -145,13 +164,7 @@ class ReMeLight(Application): 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={ - "watch_paths": [ - str(self.working_path / "MEMORY.md"), - str(self.working_path / "memory.md"), - str(self.memory_path), - ], - }, + default_file_watcher_config=_merged_file_watcher_config, ) # Initialize list to track background summarization tasks From bf79986f9cd8e702473b11143fed28950ef6379d Mon Sep 17 00:00:00 2001 From: Xinmin Zeng <135568692+fancyboi999@users.noreply.github.com> Date: Fri, 27 Mar 2026 12:22:29 +0800 Subject: [PATCH 50/59] fix: surface summarize/retrieve failures instead of masking them (#160) * fix(memory): surface summarize and retrieve failures clearly * fix(memory): make raise_exception configurable * fix(tests): resolve flake8 and pylint errors in error handling tests - Remove unnecessary sys.path.insert hack - Add module/class/function docstrings - Initialize call_kwargs in __init__ to fix W0201 - Suppress W0212 with inline pylint disable for _started access - Remove unnecessary lambda wrappers (W0108) --- reme/reme.py | 66 ++++++++++-- tests/test_reme_memory_error_handling.py | 132 +++++++++++++++++++++++ 2 files changed, 188 insertions(+), 10 deletions(-) create mode 100644 tests/test_reme_memory_error_handling.py diff --git a/reme/reme.py b/reme/reme.py index 685cdcdd..5e2aa5b2 100644 --- a/reme/reme.py +++ b/reme/reme.py @@ -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): @@ -192,12 +216,14 @@ 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, ), ] 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/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", + ) From ff49a77f18aba82039bd9fb1c89d73a0895591d2 Mon Sep 17 00:00:00 2001 From: jinliyl <6469360+jinliyl@users.noreply.github.com> Date: Fri, 27 Mar 2026 21:14:44 +0800 Subject: [PATCH 51/59] feat(memory): improve skills tool result truncation (#182) * fix(memory): correct line numbering and improve tool result truncation - Changed default start_line from 0 to 1 in truncate_text_output function - Refactored _truncate method to be a standalone method in ToolResultCompactor - Improved tool result compaction logic to handle text blocks more efficiently - Added detection of skill-related tool calls for special handling - Implemented conditional byte limits based on tool type for better memory management - Updated version number from 0.3.1.4 to 0.3.1.5 * feat(file_utils): add encoding parameter to truncate_text_output function - Added encoding parameter with default value "utf-8" to truncate_text_output function - Updated all encode/decode calls to use the specified encoding parameter - Modified ToolResultCompactor to pass encoding parameter when calling truncate_text_output - Added error handling for skill tool ID detection in message processing loop - Fixed potential AttributeError when accessing raw_input field that might be None * fix(file-store): handle corrupted ChromaDB initialization and improve tool result truncation - Add shutil import for directory removal operations - Extract ChromaDB client creation into separate _create_chroma_client method - Implement retry mechanism with database wipe on ChromaDB initialization failure - Add proper exception handling in tool result compaction to prevent truncation errors - Move file writing logic outside of exception handling scope for better error management - Add warning log when truncation fails and return original content as fallback --- reme/__init__.py | 2 +- reme/core/file_store/chroma_file_store.py | 28 +++++-- .../components/tool_result_compactor.py | 73 ++++++++++++++----- reme/memory/file_based/utils/file_utils.py | 12 +-- 4 files changed, 82 insertions(+), 33 deletions(-) diff --git a/reme/__init__.py b/reme/__init__.py index 85d6b7b7..1b211a70 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.1.4" +__version__ = "0.3.1.5" __all__ = [ "config", diff --git a/reme/core/file_store/chroma_file_store.py b/reme/core/file_store/chroma_file_store.py index d87f5fef..4deb07d0 100644 --- a/reme/core/file_store/chroma_file_store.py +++ b/reme/core/file_store/chroma_file_store.py @@ -2,6 +2,7 @@ import json import random +import shutil import time from pathlib import Path @@ -108,13 +109,9 @@ class ChromaFileStore(BaseFileStore): except Exception as e: logger.error(f"Failed to save file metadata to {self._metadata_file}: {e}") - async def start(self) -> None: - """Initialize ChromaDB client and collection.""" - if self.client is not None: - return - - # Initialize persistent ChromaDB client - self.client = chromadb.PersistentClient( + def _create_chroma_client(self): + """Create a ChromaDB PersistentClient instance.""" + return chromadb.PersistentClient( path=str(self.db_path), settings=Settings( anonymized_telemetry=False, @@ -122,6 +119,23 @@ class ChromaFileStore(BaseFileStore): ), ) + async def start(self) -> None: + """Initialize ChromaDB client and collection.""" + if self.client is not None: + return + + # Initialize persistent ChromaDB client, retry once after wiping db_path on failure + try: + self.client = self._create_chroma_client() + except Exception as e: + logger.warning( + f"ChromaDB failed to initialize at {self.db_path} ({e}). " f"Deleting corrupted database and retrying.", + ) + if self.db_path.exists(): + shutil.rmtree(self.db_path) + logger.info(f"Deleted ChromaDB directory: {self.db_path}") + self.client = self._create_chroma_client() + # Get or create the chunks collection # ChromaDB uses cosine distance by default for similarity self.chunks_collection = self.client.get_or_create_collection( diff --git a/reme/memory/file_based/components/tool_result_compactor.py b/reme/memory/file_based/components/tool_result_compactor.py index cb53381f..d808974c 100644 --- a/reme/memory/file_based/components/tool_result_compactor.py +++ b/reme/memory/file_based/components/tool_result_compactor.py @@ -37,36 +37,43 @@ class ToolResultCompactor(BaseOp): self.encoding = encoding self.tool_result_dir.mkdir(parents=True, exist_ok=True) - 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.""" - - def _truncate(content: str) -> str: - if not content: - return content + 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) + 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 - try: - fp = self.tool_result_dir / f"{uuid.uuid4().hex}.txt" - fp.write_text(content, encoding=self.encoding) - saved_path = str(fp) - except Exception as e: - logger.warning("Failed to save full tool result to file: %s", e) + 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) + 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 _truncate(output) + return self._truncate(output, max_bytes) if isinstance(output, list): - return [ - {**b, "text": _truncate(b.get("text", ""))} if isinstance(b, dict) and b.get("type") == "text" else b - for b in output - ] + 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]: @@ -84,6 +91,27 @@ class ToolResultCompactor(BaseOp): recent_n += 1 split_index = max(0, len(messages) - max(recent_n, self.recent_n)) + skills_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 "skill.md" in (block.get("raw_input") or "").lower() + ): + skills_tool_ids.add(tool_id) + except Exception as e: + logger.warning("Failed to detect skill tool ids: %s", e) + logger.info(f"skills_tool_ids: {skills_tool_ids}") + for idx, msg in enumerate(messages): if not isinstance(msg.content, list): continue @@ -91,7 +119,12 @@ class ToolResultCompactor(BaseOp): 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"): - block["output"] = self._compact(block["output"], max_bytes) + tool_use_id = block.get("id", "") + if tool_use_id in skills_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 diff --git a/reme/memory/file_based/utils/file_utils.py b/reme/memory/file_based/utils/file_utils.py index ee556bba..e7204537 100644 --- a/reme/memory/file_based/utils/file_utils.py +++ b/reme/memory/file_based/utils/file_utils.py @@ -29,10 +29,11 @@ TRUNCATION_NOTICE_MARKER = "<<>>" # pylint: disable=too-many-return-statements def truncate_text_output( text: str, - start_line: int = 0, + 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. @@ -52,6 +53,7 @@ def truncate_text_output( 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. @@ -67,7 +69,7 @@ def truncate_text_output( original_content = parts[0] old_notice = parts[1] - text_bytes = original_content.encode("utf-8") + text_bytes = original_content.encode(encoding) # Allow a small slack to avoid re-truncating near-limit content if len(text_bytes) <= max_bytes + 100: @@ -82,7 +84,7 @@ def truncate_text_output( total_lines_parsed = int(total_match.group(1)) truncated_bytes = text_bytes[:max_bytes] - result = truncated_bytes.decode("utf-8", errors="ignore") + result = truncated_bytes.decode(encoding, errors="ignore") newline_count = result.count("\n") next_line = start_line_parsed + max(1, newline_count) @@ -99,13 +101,13 @@ def truncate_text_output( return result + TRUNCATION_NOTICE_MARKER + new_notice else: - text_bytes = text.encode("utf-8") + text_bytes = text.encode(encoding) if len(text_bytes) <= max_bytes: return text truncated = text_bytes[:max_bytes] - result = truncated.decode("utf-8", errors="ignore") + result = truncated.decode(encoding, errors="ignore") newline_count = result.count("\n") From 37628ba524c9c06a2ad2c21445c8c835a42e4418 Mon Sep 17 00:00:00 2001 From: jinliyl <6469360+jinliyl@users.noreply.github.com> Date: Sat, 28 Mar 2026 18:41:09 +0800 Subject: [PATCH 52/59] refactor(truncation): improve file truncation logic (#184) * refactor(file_store): simplify ChromaDB client initialization and improve file truncation logic - Remove shutil import and _create_chroma_client method from chroma_file_store.py - Directly initialize ChromaDB PersistentClient in start method without retry logic - Reduce DEFAULT_MAX_BYTES from 100KB to 50KB in file_utils.py - Update truncation notice format to provide clearer continuation instructions - Add _truncate_fresh and _retruncate functions for better text truncation handling - Replace inline truncation logic with dedicated function calls in file_utils.py - Rename skills_tool_ids to md_file_tool_ids in tool_result_compactor.py - Update file detection logic to identify any .md files instead of only skill.md - Create comprehensive unit tests for truncation functionality in test_truncate_text_output.py * chore(version): bump version to 0.3.1.6 - Update __version__ from 0.3.1.5 to 0.3.1.6 in __init__.py --- reme/__init__.py | 2 +- reme/core/file_store/chroma_file_store.py | 28 +- .../components/tool_result_compactor.py | 12 +- reme/memory/file_based/tools/file_io.py | 11 +- reme/memory/file_based/utils/file_utils.py | 200 ++++++--- tests/light/test_truncate_text_output.py | 381 ++++++++++++++++++ 6 files changed, 538 insertions(+), 96 deletions(-) create mode 100644 tests/light/test_truncate_text_output.py diff --git a/reme/__init__.py b/reme/__init__.py index 1b211a70..f9ade852 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.1.5" +__version__ = "0.3.1.6" __all__ = [ "config", diff --git a/reme/core/file_store/chroma_file_store.py b/reme/core/file_store/chroma_file_store.py index 4deb07d0..d87f5fef 100644 --- a/reme/core/file_store/chroma_file_store.py +++ b/reme/core/file_store/chroma_file_store.py @@ -2,7 +2,6 @@ import json import random -import shutil import time from pathlib import Path @@ -109,9 +108,13 @@ class ChromaFileStore(BaseFileStore): except Exception as e: logger.error(f"Failed to save file metadata to {self._metadata_file}: {e}") - def _create_chroma_client(self): - """Create a ChromaDB PersistentClient instance.""" - return chromadb.PersistentClient( + async def start(self) -> None: + """Initialize ChromaDB client and collection.""" + if self.client is not None: + return + + # Initialize persistent ChromaDB client + self.client = chromadb.PersistentClient( path=str(self.db_path), settings=Settings( anonymized_telemetry=False, @@ -119,23 +122,6 @@ class ChromaFileStore(BaseFileStore): ), ) - async def start(self) -> None: - """Initialize ChromaDB client and collection.""" - if self.client is not None: - return - - # Initialize persistent ChromaDB client, retry once after wiping db_path on failure - try: - self.client = self._create_chroma_client() - except Exception as e: - logger.warning( - f"ChromaDB failed to initialize at {self.db_path} ({e}). " f"Deleting corrupted database and retrying.", - ) - if self.db_path.exists(): - shutil.rmtree(self.db_path) - logger.info(f"Deleted ChromaDB directory: {self.db_path}") - self.client = self._create_chroma_client() - # Get or create the chunks collection # ChromaDB uses cosine distance by default for similarity self.chunks_collection = self.client.get_or_create_collection( diff --git a/reme/memory/file_based/components/tool_result_compactor.py b/reme/memory/file_based/components/tool_result_compactor.py index d808974c..d80b2f61 100644 --- a/reme/memory/file_based/components/tool_result_compactor.py +++ b/reme/memory/file_based/components/tool_result_compactor.py @@ -91,7 +91,7 @@ class ToolResultCompactor(BaseOp): recent_n += 1 split_index = max(0, len(messages) - max(recent_n, self.recent_n)) - skills_tool_ids = set() + md_file_tool_ids = set() try: for msg in messages: if not isinstance(msg.content, list): @@ -105,12 +105,12 @@ class ToolResultCompactor(BaseOp): if ( block.get("name", "").lower() == "read_file" - and "skill.md" in (block.get("raw_input") or "").lower() + and ".md" in (block.get("raw_input") or "").lower() ): - skills_tool_ids.add(tool_id) + md_file_tool_ids.add(tool_id) except Exception as e: - logger.warning("Failed to detect skill tool ids: %s", e) - logger.info(f"skills_tool_ids: {skills_tool_ids}") + 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): @@ -120,7 +120,7 @@ class ToolResultCompactor(BaseOp): 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 skills_tool_ids: + if tool_use_id in md_file_tool_ids: effective_max_bytes = self.recent_max_bytes else: effective_max_bytes = max_bytes diff --git a/reme/memory/file_based/tools/file_io.py b/reme/memory/file_based/tools/file_io.py index d876407c..3e28addc 100644 --- a/reme/memory/file_based/tools/file_io.py +++ b/reme/memory/file_based/tools/file_io.py @@ -153,10 +153,13 @@ class FileIO: if text == selected_content and e < total: content_bytes = len(text.encode("utf-8")) notice = ( - TRUNCATION_NOTICE_MARKER - + f"\nFile: {file_path}\nStarting at start_line={s}, next {content_bytes} bytes." - f"\nTotal lines: {total}" - f"\nUse start_line={e + 1} to continue." + 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 diff --git a/reme/memory/file_based/utils/file_utils.py b/reme/memory/file_based/utils/file_utils.py index e7204537..50248554 100644 --- a/reme/memory/file_based/utils/file_utils.py +++ b/reme/memory/file_based/utils/file_utils.py @@ -8,7 +8,7 @@ from ....core.utils import get_logger logger = get_logger() # Default truncation limit -DEFAULT_MAX_BYTES = 100 * 1024 +DEFAULT_MAX_BYTES = 50 * 1024 # Maximum file size to read into memory (1GB) MAX_FILE_READ_BYTES = 1024 * 1024 * 1024 @@ -16,17 +16,136 @@ MAX_FILE_READ_BYTES = 1024 * 1024 * 1024 # Marker prepended to every truncation notice. # Format: # <<>> -# File: -# Starting at start_line=X, next N bytes. -# Total lines: Z -# Use start_line=Y to continue. +# 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 = "<<>>" -# pylint: disable=too-many-return-statements +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, @@ -41,9 +160,9 @@ def truncate_text_output( If over limit, truncate at the last complete line that fits, allowing the next read to start from a fresh line. - If TRUNCATION_NOTICE_MARKER is already in text (previously truncated), - extract the original content, re-truncate it, and update the - truncation notice using regex. + 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. @@ -65,63 +184,16 @@ def truncate_text_output( try: if TRUNCATION_NOTICE_MARKER in text: - 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 re-truncating near-limit content - if len(text_bytes) <= max_bytes + 100: - return text - - # Parse start_line and total_lines from notice; return text unchanged if not found - start_match = re.search(r"Starting at start_line=(\d+)", old_notice) - total_match = re.search(r"Total lines: (\d+)", old_notice) - if not start_match or not total_match: - return text - start_line_parsed = int(start_match.group(1)) - total_lines_parsed = int(total_match.group(1)) - - truncated_bytes = text_bytes[:max_bytes] - result = truncated_bytes.decode(encoding, errors="ignore") - newline_count = result.count("\n") - - next_line = start_line_parsed + max(1, newline_count) - - if not re.search(r"next \d+ bytes", old_notice): - return text - has_continuation = bool(re.search(r"Use start_line=\d+", old_notice)) - new_notice = re.sub(r"next \d+ bytes", f"next {max_bytes} bytes", old_notice) - if has_continuation: - new_notice = re.sub(r"Use start_line=\d+", f"Use start_line={next_line}", new_notice) - elif next_line <= total_lines_parsed: - new_notice = re.sub(r"(Total lines: \d+)", f"\\1\nUse start_line={next_line} to continue.", new_notice) - - return result + TRUNCATION_NOTICE_MARKER + new_notice - + return _retruncate(text, max_bytes=max_bytes, encoding=encoding) else: - text_bytes = text.encode(encoding) - - if len(text_bytes) <= max_bytes: - return text - - truncated = text_bytes[:max_bytes] - result = truncated.decode(encoding, errors="ignore") - - newline_count = result.count("\n") - - next_line = start_line + max(1, newline_count) - - continuation = f"\nUse start_line={next_line} to continue." if next_line <= total_lines else "" - notice = ( - TRUNCATION_NOTICE_MARKER - + f"\nFile: {file_path or ''}\nStarting at start_line={start_line}, next {max_bytes} bytes." - f"\nTotal lines: {total_lines}{continuation}" + return _truncate_fresh( + text, + start_line=start_line, + total_lines=total_lines, + max_bytes=max_bytes, + file_path=file_path, + encoding=encoding, ) - - return result + notice - except Exception: logger.warning("truncate_text_output failed, returning original text", exc_info=True) return 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) From d845cff1e31b4026e6c9da216c76ec53a362a2bf Mon Sep 17 00:00:00 2001 From: jinliyl <6469360+jinliyl@users.noreply.github.com> Date: Mon, 30 Mar 2026 16:09:49 +0800 Subject: [PATCH 53/59] docs(memory): update ReMeLight memory system documentation (#186) - Add context data structure diagram showing compact_summary and file system cache - Update ToolResultCompactor section with detailed truncation strategies for recent vs old messages - Add parameter tables for tool result compaction with recent_max_bytes and old_max_bytes settings - Update execution flow steps with detailed descriptions of each memory operation - Add key parameters table including tool_result_compact_keep_n and memory_compact_reserve - Include thinking enhancement feature description for summary generation quality improvement - Update both English and Chinese README documentation consistently --- README.md | 74 ++++++++++++++++++++++++++++++++++++---------------- README_ZH.md | 51 +++++++++++++++++++++++++++--------- 2 files changed, 91 insertions(+), 34 deletions(-) diff --git a/README.md b/README.md index 1d0790c9..0a5edc94 100644 --- a/README.md +++ b/README.md @@ -223,9 +223,22 @@ if __name__ == "__main__": ### Architecture of the file-based ReMeLight memory system -[CoPaw MemoryManager](https://github.com/agentscope-ai/CoPaw/blob/main/src/copaw/agents/memory/memory_manager.py) -inherits -`ReMeLight` and integrates its memory capabilities into the agent reasoning loop: +#### Context data structure + +```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] +``` + +--- + +[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: ```mermaid graph LR @@ -283,16 +296,18 @@ graph LR **Summary structure** (context checkpoints): -| Field | Description | -|-----------------------|------------------------------------------------------------------------| -| `## Goal` | User goals | -| `## Constraints` | Constraints and preferences | -| `## Progress` | Task progress | -| `## Key Decisions` | Key decisions | -| `## Next Steps` | Next step plans | -| `## Critical Context` | Critical data such as file paths, function names, error messages, etc. | +| Field | Description | +|-----------------------|-----------------------------------------------------------------------------------------| +| `## Goal` | User goals | +| `## Constraints` | Constraints and preferences | +| `## Progress` | Task progress | +| `## Key Decisions` | Key decisions | +| `## Next Steps` | Next step plans | +| `## Critical Context` | Critical data such as file paths, function names, error messages, etc. | - **Incremental updates**: when `previous_summary` is provided, new conversations are merged into the existing summary. +- **Thinking enhancement**: with `add_thinking_block=True` (default), a reasoning step is added before generating the + summary to improve quality. --- @@ -325,18 +340,25 @@ graph LR #### 4. `compact_tool_result` — tool result compaction [ToolResultCompactor](reme/memory/file_based/components/tool_result_compactor.py) addresses the problem of long tool -outputs bloating the context. +outputs bloating the context. It applies two different truncation strategies depending on whether a message falls within +the `recent_n` window: ```mermaid graph LR - M[messages] --> L{Iterate tool_result
len > threshold?} - L -->|No| K[Keep as-is] - L -->|Yes| T[truncate_text
Truncate to threshold] - T --> S[Write full content
tool_result/uuid.txt] - S --> R[Append file path reference
to message] - R --> C[cleanup_expired_files
Delete expired files] + M[messages] --> B{Within recent_n?} + B -->|Yes - recent| C[Low truncation recent_max_bytes=100KB
Save full content to tool_result/uuid.txt
Hint: 'Read from line N'] + B -->|No - old| D[High truncation old_max_bytes=3KB
Reference existing file
More aggressive truncation] + C --> E[cleanup_expired_files
Delete expired files] + D --> E ``` +| Parameter | Default | Description | +|--------------------|-----------------------|-------------------------------------------------------------------------------------------------------------------------------| +| `recent_n` | `1` | Minimum number of trailing consecutive tool-result messages treated as "recent" (use low truncation) | +| `recent_max_bytes` | `100 * 1024` (100 KB) | Truncation threshold for recent messages; content beyond this is saved to `tool_result/` with a file path and start-line hint | +| `old_max_bytes` | `3000` (3 KB) | Truncation threshold for older messages; truncation is more aggressive | +| `retention_days` | `3` | Number of days to retain tool result files; expired files are auto-cleaned | + - **Auto cleanup**: expired files (older than `retention_days`) are deleted automatically during `start` / `close` / `compact_tool_result`. @@ -411,10 +433,18 @@ graph LR **Execution flow**: -1. `compact_tool_result` — compact long tool outputs. -2. `check_context` — check whether the context exceeds limits. -3. `compact_memory` — generate compact summary (sync). -4. `summary_memory` — persist memory (async in the background). +1. `compact_tool_result` — compact long tool outputs for all messages except the most recent + `tool_result_compact_keep_n`. +2. `check_context` — check whether the context exceeds limits (remaining space = threshold minus tokens used by system + prompt and compressed summary). +3. `compact_memory` — generate compact summary (sync), appended into `compact_summary`. +4. `summary_memory` — persist memory to `memory/*.md` (async in the background, non-blocking). + +| Key parameter | Default | Description | +|------------------------------|---------|-------------------------------------------------------------------------------------| +| `tool_result_compact_keep_n` | `3` | Skip tool result compaction for the most recent N messages (preserve full content) | +| `memory_compact_reserve` | `10000` | Token count to reserve for recent messages; messages beyond this trigger compaction | +| `compact_ratio` | `0.7` | Compaction threshold ratio: `max_input_length × compact_ratio × 0.95` | --- diff --git a/README_ZH.md b/README_ZH.md index ee455318..306dc66a 100644 --- a/README_ZH.md +++ b/README_ZH.md @@ -215,7 +215,21 @@ if __name__ == "__main__": ### 基于文件的 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 @@ -281,6 +295,7 @@ graph LR | `## Critical Context` | 文件路径、函数名、错误信息等关键数据 | - **增量更新**：传入 `previous_summary` 时，自动将新对话与旧摘要合并 +- **思考增强**：`add_thinking_block=True`（默认）时，在生成摘要前加入思考步骤，提升摘要质量 --- @@ -311,18 +326,24 @@ graph LR #### 4. compact_tool_result — 工具结果压缩 -[ToolResultCompactor](reme/memory/file_based/components/tool_result_compactor.py) 解决工具输出过长导致上下文膨胀的问题。 +[ToolResultCompactor](reme/memory/file_based/components/tool_result_compactor.py) 解决工具输出过长导致上下文膨胀的问题。根据消息是否在 `recent_n` 范围内，采用不同的截断策略： ```mermaid graph LR - M[messages] --> L{遍历 tool_result
len > threshold?} - L -->|否| K[保留原样] - L -->|是| T[truncate_text
截断到 threshold] - T --> S[完整内容写入
tool_result/uuid.txt] - S --> R[消息追加文件路径引用] - R --> C[cleanup_expired_files
清理过期文件] + M[messages] --> B{属于 recent_n 范围?} + B -->|是 近期消息| C[低截断 recent_max_bytes=100KB
完整内容写入 tool_result/uuid.txt
消息追加: 从第N行开始读] + B -->|否 历史消息| D[高截断 old_max_bytes=3KB
引用已有文件路径
更激进截断] + C --> E[cleanup_expired_files
清理过期文件] + D --> E ``` +| 参数 | 默认值 | 说明 | +|--------------------|---------------------|----------------------------------------------| +| `recent_n` | `1` | 末尾连续工具结果消息的最小数量，视为"近期"，使用低截断阈值 | +| `recent_max_bytes` | `100 * 1024`（100KB） | 近期消息的截断阈值；超出部分转存到 `tool_result/` 并附注文件路径和起始行 | +| `old_max_bytes` | `3000`（3KB） | 历史消息的截断阈值，截断更激进 | +| `retention_days` | `3` | 工具结果文件的保留天数，过期自动清理 | + - **自动清理**：过期文件（超过 `retention_days`）在 `start`/`close`/`compact_tool_result` 时自动删除 --- @@ -394,10 +415,16 @@ graph LR **执行流程**： -1. `compact_tool_result` — 压缩超长工具输出 -2. `check_context` — 检查上下文是否超限 -3. `compact_memory` — 生成压缩摘要（同步） -4. `summary_memory` — 持久化记忆（异步后台） +1. `compact_tool_result` — 对除最近 `tool_result_compact_keep_n` 条消息之外的历史消息压缩超长工具输出 +2. `check_context` — 检查上下文是否超限（扣除 system_prompt 和 compressed_summary 的 token 后计算剩余空间） +3. `compact_memory` — 生成压缩摘要（同步），结果追加到 `compact_summary` +4. `summary_memory` — 持久化记忆到 `memory/*.md`（异步后台，不阻塞推理） + +| 关键参数 | 默认值 | 说明 | +|------------------------------|---------|--------------------------------------------------| +| `tool_result_compact_keep_n` | `3` | 最近 N 条消息跳过工具结果压缩（保留完整内容） | +| `memory_compact_reserve` | `10000` | 保留近期消息的 token 数，超出部分触发压缩 | +| `compact_ratio` | `0.7` | 压缩阈值比例：`max_input_length × compact_ratio × 0.95` | --- From 2a999ce4f4b52f7d6b5c90a19f0dc1a60932d23c Mon Sep 17 00:00:00 2001 From: jinliyl <6469360+jinliyl@users.noreply.github.com> Date: Mon, 30 Mar 2026 20:21:30 +0800 Subject: [PATCH 54/59] docs(context): add comprehensive context management design documentation (#187) * docs(context): add comprehensive context management design documentation - Create detailed Chinese documentation for CoPaw context management V2 - Document memory layer and file system cache architecture - Explain Pre-Reasoning Hook workflow with four-step process - Detail two-stage truncation strategy for tool results - Add examples for Browser Use and ReadFile tools - Include Mermaid diagrams for visual flow representation - Update README with link to new context design document - Fix minor formatting issues in existing documentation - Add protection thresholds for Markdown files in truncation - Document long-term memory trigger mechanisms * docs(README): add latest articles section and CoPaw context management design doc - Added "Latest Articles" section to README with table format - Included link to CoPaw Context Management Design document - Created comprehensive documentation for CoPaw context management V2 - Documented in-memory and file system layer architecture - Explained pre-reasoning hook and context compaction process - Detailed two-phase truncation strategy for tool results - Described special handling for readFile tool and markdown files - Added long-term memory trigger logic overview - Included mermaid diagrams for visualizing context flow --- README.md | 8 + README_ZH.md | 17 +- docs/copaw_context_design.md | 302 ++++++++++++++++++++++++-------- docs/copaw_context_design_zh.md | 289 ++++++++++++++++++++++++++++++ 4 files changed, 539 insertions(+), 77 deletions(-) create mode 100644 docs/copaw_context_design_zh.md diff --git a/README.md b/README.md index 0a5edc94..7c93e489 100644 --- a/README.md +++ b/README.md @@ -29,6 +29,14 @@ --- +## 📰 Latest Articles + +| Date | Title | +|------------|-----------------------------------------------------------------| +| 2026-03-30 | [CoPaw Context Management Design](docs/copaw_context_design.md) | + +--- + 🧠 ReMe is a memory management framework designed for **AI agents**, providing both [file-based](#-file-based-memory-system-remelight) and [vector-based](#-vector-based-memory-system) memory systems. diff --git a/README_ZH.md b/README_ZH.md index 306dc66a..5fa2e4d2 100644 --- a/README_ZH.md +++ b/README_ZH.md @@ -29,6 +29,14 @@ --- +## 📰 最新文章 + +| 日期 | 标题 | +|------------|----------------------------------------------------| +| 2026-03-30 | [CoPaw 上下文管理设计解析](docs/copaw_context_design_zh.md) | + +--- + 🧠 ReMe 是一个专为 **AI 智能体** 打造的记忆管理框架，同时提供基于[文件系统](#-基于文件的记忆系统-remelight) 和基于[向量库](#-基于向量库的记忆系统)的记忆系统。 @@ -224,12 +232,13 @@ flowchart TD A --> E[messages: 完整对话历史] A --> F[文件系统缓存] F --> G[dialog/YYYY-MM-DD.jsonl] - F --> H[tool_result/uuid.txt N天TTL] + 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) 继承 +[CoPaw MemoryManager](https://github.com/agentscope-ai/CoPaw/blob/main/src/copaw/agents/memory/reme_light_memory_manager.py) +继承 `ReMeLight`，将记忆能力集成到 Agent 推理流程中： ```mermaid @@ -326,7 +335,8 @@ graph LR #### 4. compact_tool_result — 工具结果压缩 -[ToolResultCompactor](reme/memory/file_based/components/tool_result_compactor.py) 解决工具输出过长导致上下文膨胀的问题。根据消息是否在 `recent_n` 范围内，采用不同的截断策略： +[ToolResultCompactor](reme/memory/file_based/components/tool_result_compactor.py) 解决工具输出过长导致上下文膨胀的问题。根据消息是否在 +`recent_n` 范围内，采用不同的截断策略： ```mermaid graph LR @@ -454,7 +464,6 @@ graph LR 安装和环境变量配置与 [ReMeLight 一致](#安装)，通过环境变量设置 API 密钥，可写在项目根目录的 `.env` 文件中。 - ### Python 使用 ```python diff --git a/docs/copaw_context_design.md b/docs/copaw_context_design.md index 01da3764..13fa5fc6 100644 --- a/docs/copaw_context_design.md +++ b/docs/copaw_context_design.md @@ -1,122 +1,278 @@ -## Copaw Context Management V2 +# CoPaw Context Management Design -> 注：不涉及长期记忆 +> This article focuses on **short-term context management** and does not cover the long-term memory module. -### 上下文数据结构 +--- -#### 1. 上下文-内存 +Sooner or later, every AI Agent hits the same wall: **the context window fills up**. -- **compact_summary**（可选）： - - **历史对话原始数据引导**：存储于 `dialog/YYYY-MM-DD.jsonl`，共 N 行，按时间顺序排列；回顾时建议从后往前读。 - - **历史对话摘要**：包含 `Goal + Constraints + Progress + KeyDecisions + NextSteps`。 -- **messages**：当前对话上下文（完整消息列表）。 +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. -#### 2. 上下文-缓存到文件系统 +[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**. -- **历史对话原始数据**：`dialog/YYYY-MM-DD.jsonl` -- **工具调用结果原始数据**：`tool_result/{uuid}.txt`（保留 N 天） +--- + +## 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] --> G[dialog/YYYY-MM-DD.jsonl] + 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. + --- -### 上下文机制（Pre-Reasoning Hook） +## What happens before reasoning? The Pre-Reasoning Hook -1. **工具结果 Offload** (`ToolCallResultCompact`) -2. **上下文检查** (`ContextChecker`) -3. **若 Token 超阈值**： - - 保留最近 **X%** 的 Token（保障连贯性） - - 其余历史对话生成摘要 (`Compactor`) -4. **被摘要的上下文 Offload 到文件系统** (`SaveDialog`) +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?} + C --> D{Token > threshold?} D -->|Yes| E[Keep recent X% tokens] - E --> F[Compact & Summary old context] - F --> G[SaveDialog: offload to file] - D -->|No| H[Proceed normally] + 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. + --- -### 工具结果 Offload 机制 +## Tool Result Offload: Unified Two-Phase Truncation -1. 所有工具调用结果先放入上下文，等待 Pre-Reasoning Hook 处理。 -2. 根据是否属于 **recent_n** 范围，决定截断策略： - - **recent_n 内**：近期内容 → 低截断比例 - - **recent_n 外**：远期内容 → 高截断比例 +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: -#### 示例：Browser Use 类工具 - -| 阶段 | 行为 | -|----|------------------------------------------------------------------------------| -| 1 | 原始工具调用结果 | -| 2 | 保存原始内容到文件：– 若在 recent_n 内：截断较少– 附注：“FullText saved to xxxx”– 提示：“请从第 N 行开始读” | -| 3 | 若再次引用且超出 recent_n：– 二次截断（更激进）– 仍指向原文件路径 | +- **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 Result] --> B{Within recent_n?} - B -->|Yes| C[Low truncation
Save full text to tool_result/uuid.txt
Hint: 'Read from line N'] - B -->|No| D[High truncation
Reference existing file
More aggressive truncation] - C --> E[Context includes snippet + file ref] + A[Tool call completes] --> T[First truncation
executed immediately on write] + T --> S[Full content written to tool_result/uuid.txt
Message annotated with file path + start line] + S --> B{Pre-Reasoning Hook
Is message within recent_n?} + B -->|Yes| C[No action
Keep first-truncation result] + B -->|No| D[Second truncation
More aggressive compression
file_path unchanged] +``` + +The benefit of this design is: the first truncation ensures that no tool result can blow up the context from the moment it is written; the second truncation automatically "fades out" older messages as the conversation progresses, always leaving enough room for recent content. + +### Browser Use Tools as an Example + +| Phase | Behavior | +|----------------------------|-----------------------------------------------------------------------------------------------------------------------------| +| First truncation | Result is truncated immediately on return; full content written to `tool_result/uuid.txt`; message annotated with "FullText saved to xxxx, please read from line N" | +| Within `recent_n` | Pre-Reasoning Hook makes no additional changes; keeps first-truncation result | +| Outside `recent_n` (second truncation) | Parses the existing message result, applies more aggressive truncation to the original content, updates meta info (e.g. line-number hint); **`file_path` is unchanged**, still pointing to the original file | + +The key insight of second truncation: **the original full content is always saved under the same file path**. No matter how many rounds of truncation have occurred, the Agent can always retrieve the original content via the file reference. Truncation only affects the message fragment and meta info in the context — never the file itself. + +```mermaid +flowchart LR + A[Browser Use Result] -->|First truncation| B[Context: fragment + file_path + start line] + B --> C{Outside recent_n?} + C -->|No| D[Unchanged] + C -->|Yes| E[Parse existing message
Apply second truncation to original content
Update meta info] + E --> F[Context: shorter fragment + same file_path] +``` + +### Code Implementation and Examples of Two-Phase Truncation + +The entry point for truncation logic is `truncate_text_output`, which dispatches to two different functions depending on whether the text already contains the `<<>>` 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
original file path known] + B --> C{Outside recent_n?} + C -->|No| D[No changes needed] + C -->|Yes| E[Second truncation
Update meta info
Same behavior as other tools] +``` + +### Special Protection for Markdown Files + +For Markdown files such as `skill.md` and rule files, CoPaw applies a **higher protection threshold** during truncation. + +Markdown files typically carry structured knowledge or instructions; over-aggressive truncation would break their semantic integrity. Therefore, both the first and second truncation thresholds for Markdown files are set higher than those for regular tool outputs, ensuring the Agent can read as complete a structured content as possible. + +```mermaid +flowchart LR + A[Tool result] --> B{Is it a Markdown file?} + B -->|Yes| C[Higher truncation threshold
Greater protection] + B -->|No| D[Standard truncation threshold] + C --> E[First / second truncation logic] D --> E ``` --- -### ReadFile 工具调用结果变化示例 +## Long-term Memory Trigger Logic -| 阶段 | 行为 | -|----|---------------------------------------------| -| 1 | 原始工具调用结果 | -| 2 | 若在 recent_n 内：– 不截断– 不保存文件（因内容已由用户指定） | -| 3 | 若超出 recent_n：– 二次截断（更小）– 保存 FullText 到文件并引用 | +> This section goes beyond the core scope of context management and briefly introduces CoPaw's long-term memory write mechanism. -> 注：ReadFile 本身读取的是外部文件，因此首次调用通常无需重复保存。 +Long-term memory is driven by three trigger paths: -```mermaid -flowchart LR - A[ReadFile Result] --> B{Within recent_n?} - B -->|Yes| C[No truncation
No file save needed] - B -->|No| D[Apply secondary truncation
Save FullText to tool_result/uuid.txt] - C --> E[Include full content in context] - D --> F[Include snippet + file ref] -``` +1. **Explicitly written by the Main Agent**: + - `Memory.md` (the backbone of long-term memory, recording persistent information such as user preferences) + - `YYYY-MM-DD.md` (daily log) ---- +2. **Triggered by context compaction**, written by the **Summarizer (ReAct Agent)**: + - Personalization information (user preferences, habits, etc.) + - Try-error information (failed attempts and corrective lessons) -## Copaw Memory - -### 触发逻辑 - -1. **主 Agent 主动写入**： - - `Memory.md`（长期记忆主干） - - `YYYY-MM-DD.md`（当日日志） -2. **触发阈值时**，由 **Summarizer（React Agent）** 写日志： - - 个性化信息（如偏好、习惯） - - Try-error 信息（失败尝试与修正） -3. **定时任务**（每日 00:00）： - - 汇总最近的 `YYYY-MM-DD.md` 文件 - - 更新 `Memory.md` +3. **Scheduled task** (daily at 00:00): + - Aggregates recent `YYYY-MM-DD.md` files + - Merges and updates the journal into `Memory.md` ```mermaid flowchart TD A[Main Agent] --> B[Write Memory.md] A --> C[Write YYYY-MM-DD.md] - D[Context Threshold Reached?] -->|Yes| E[Summarizer Agent] - E --> F[Log: Personalization] - E --> G[Log: Try-Error Info] - H[Cron @ 00:00 daily] --> I[Aggregate recent YYYY-MM-DD.md] -I --> J[Update Memory.md] + D[Context reaches compaction threshold?] -->|Yes| E[Summarizer Agent] + E --> F[Record: personalization info] + E --> G[Record: try-error experience] + H[Scheduled task 00:00] --> I[Aggregate recent YYYY-MM-DD.md] + I --> J[Update Memory.md] ``` + +This mechanism ensures that important information from short-term conversations is distilled into long-term memory and not lost when the session ends. + +--- + +## Summary + +The core design philosophy of CoPaw's context management can be summed up in one sentence: + +**Keep only "what is needed now" in memory; let the file system hold "what might be needed later".** + +Through the four-step Pre-Reasoning Hook flow, a unified two-phase truncation strategy, and persistent file system backing, CoPaw maximizes information availability for the Agent within a limited context window — no matter how long the conversation runs, the Agent can always find the context it needs. + +--- + +*The design described in this article is implemented in [CoPaw MemoryManager](https://github.com/agentscope-ai/CoPaw/blob/main/src/copaw/agents/memory/reme_light_memory_manager.py) and [ReMe ReMeLight](https://github.com/agentscope-ai/ReMe).* diff --git a/docs/copaw_context_design_zh.md b/docs/copaw_context_design_zh.md new file mode 100644 index 00000000..8a72f169 --- /dev/null +++ b/docs/copaw_context_design_zh.md @@ -0,0 +1,289 @@ +# CoPaw 上下文管理设计解析 + +> 本文聚焦**短期上下文管理**，不涉及长期记忆模块。 + +--- + +AI Agent 在使用过程中，迟早会遇到一个让人头疼的问题：**上下文窗口被塞满了**。 + +工具调用返回了一大段 HTML、几千行日志、或者完整的文件内容——这些都会急剧消耗宝贵的 Token +配额。随着对话轮次增加，早期的信息要么被截断，要么把整个窗口撑爆，Agent 的表现开始下滑。 + +[CoPaw](https://github.com/agentscope-ai/CoPaw) 在设计上下文管理时，围绕这个问题给出了一套系统性的答案。本文将完整拆解 * +*CoPaw Context Management V2** 的数据结构与运行机制。 + +--- + +## 上下文长什么样？ + +在讨论"如何管理"之前，先看清楚"管理的是什么"。 + +CoPaw 的上下文分为两层：**内存层**与**文件系统层**。 + +### 内存层（In-Memory） + +内存中维护两个核心字段： + +- **`compact_summary`**（可选）：当历史对话被压缩后，这里存放结构化摘要，包含 `Goal`、`Constraints`、`Progress`、`KeyDecisions`、 + `NextSteps` 五个维度——相当于一份精炼的"工作备忘录"。同时还包含一个**历史对话原始数据的路径引导**，告诉 Agent 去哪里找 + `dialog/YYYY-MM-DD.jsonl`，以及建议"从后往前读"。 +- **`messages`**：当前对话的完整消息列表，是 Agent 实际推理时消费的数据。 + +### 文件系统层（File Cache） + +对于体积较大、不适合长期驻留内存的内容，CoPaw 将其 offload 到文件系统： + +- **历史对话原始数据**：`dialog/YYYY-MM-DD.jsonl`，按日期分文件存储 +- **工具调用结果**：`tool_result/{uuid}.txt`，设有 N 天 TTL，过期自动清理 + +```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] +``` + +这个设计让 Agent 既能在内存中快速访问近期对话，又能在需要时按需回溯历史——而不是把所有历史内容硬塞进上下文。 + +--- + +## 推理前做什么？Pre-Reasoning Hook + +每轮推理正式开始前，CoPaw 会执行一个 **Pre-Reasoning Hook**，自动完成上下文的整理工作。整个流程分四步： + +1. **工具结果压缩**（`ToolCallResultCompact`）：先处理工具调用结果，将超长内容截断并 offload 到文件系统 +2. **上下文检查**（`ContextChecker`）：计算当前上下文的 Token 使用量，判断是否超出阈值 +3. **若超出阈值**： + - 保留最近 **X%** 的 Token（保障对话连贯性） + - 对更早的历史对话调用 `Compactor` 生成结构化摘要 +4. **历史对话持久化**（`SaveDialog`）：将被压缩的原始对话保存到文件系统 + +```mermaid +flowchart LR + A[Pre-Reasoning Hook] --> B[ToolCallResultCompact] + B --> C[ContextChecker] + C --> D{Token > 阈值?} + D -->|是| E[保留近期 X% Token] + E --> F[Compact & 生成摘要] + F --> G[SaveDialog: 持久化到文件] + D -->|否| H[正常推理] +``` + +这个流程确保每次推理开始前，上下文都处于一个"干净"的状态。 + +--- + +## 工具结果 Offload：统一的两阶段截断 + +工具调用结果是上下文膨胀的主要来源之一。CoPaw 采用**两阶段截断**策略，将截断时机与截断力度分离： + +- **一次截断**：在工具调用结果**写入上下文时**立即触发，所有工具（包括 `read_file`）统一适用。截断后将完整原始内容保存到 + `tool_result/{uuid}.txt`，并在消息中附注文件路径与起始行提示。 +- **二次截断**：在 Pre-Reasoning Hook 处理时，对**已滑出 `recent_n` 范围**的历史消息触发，截断更为激进，进一步压缩上下文占用。 + +```mermaid +flowchart LR + A[工具调用完成] --> T[一次截断
写入上下文时立即执行] + T --> S[完整内容写入 tool_result/uuid.txt
消息附注文件路径 + 起始行] + S --> B{Pre-Reasoning Hook
该消息在 recent_n 内?} + B -->|是| C[无需处理
保持一次截断结果] + B -->|否| D[二次截断
更激进压缩
file_path 不变] +``` + +这样设计的好处在于：一次截断保证所有工具结果从写入那刻起就不会撑爆上下文；二次截断则随着对话推进自动"淡化" +历史信息，始终为近期内容留出充足空间。 + +### 以 Browser Use 类工具为例 + +| 阶段 | 行为 | +|-------------------|------------------------------------------------------------------------------------| +| 一次截断 | 工具返回结果后立即截断，完整内容写入 `tool_result/uuid.txt`，消息附注 "FullText saved to xxxx，请从第 N 行开始读" | +| 在 recent_n 内 | Pre-Reasoning Hook 不做额外处理，保持一次截断结果 | +| 超出 recent_n（二次截断） | 解析已有的消息结果，对原始内容做更激进的截断，同步更新消息中的 meta 信息（如行号提示）；**`file_path` 不变**，仍指向原文件 | + +二次截断的关键在于：**原始完整内容始终保存在同一个文件路径下**，无论经过多少轮截断，Agent 都能通过文件引用找到原始内容；截断只影响上下文中的消息片段和 +meta 信息，不改变文件。 + +```mermaid +flowchart LR + A[Browser Use Result] -->|一次截断| B[上下文: 片段 + file_path + 起始行] + B --> C{超出 recent_n?} + C -->|否| D[保持不变] + C -->|是| E[解析现有消息
对原始内容二次截断
更新 meta 信息] + E --> F[上下文: 更短片段 + 同一 file_path] +``` + +### 两阶段截断的代码实现与示例 + +截断逻辑的入口是 `truncate_text_output`，它根据文本中是否已包含 `<<>>` 标记来分发到两个不同的函数： + +```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[上下文: 截断内容
原始文件路径已知] + B --> C{超出 recent_n?} + C -->|否| D[无需修改] + C -->|是| E[二次截断
更新 meta 信息
与其他工具行为一致] +``` + +### Markdown 文件的特殊保护 + +对于 `skill.md`、规则文件等 Markdown 文件，CoPaw 在截断时给予**更大的保护阈值**。 + +Markdown 文件通常承载结构化的知识或指令，过度截断会破坏其完整语义。因此，在一次截断和二次截断时，Markdown +文件的截断触发上限均高于普通工具输出，确保 Agent 能读到尽可能完整的结构化内容。 + +```mermaid +flowchart LR + A[工具结果] --> B{是 Markdown 文件?} + B -->|是| C[更高截断阈值
更大保护] + B -->|否| D[标准截断阈值] + C --> E[一次 / 二次截断逻辑] + D --> E +``` + +--- + +## 长期记忆的触发逻辑 + +> 本节超出上下文管理的核心范畴，简要介绍 CoPaw 的长期记忆写入机制。 + +长期记忆由三个触发路径驱动： + +1. **主 Agent 主动写入**： + - `Memory.md`（长期记忆主干，记录用户偏好等持久信息） + - `YYYY-MM-DD.md`（当日日志） + +2. **上下文压缩触发时**，由 **Summarizer（ReAct Agent）** 写入： + - 个性化信息（用户偏好、习惯等） + - Try-error 信息（失败尝试与修正经验） + +3. **定时任务**（每日 00:00）： + - 汇总最近的 `YYYY-MM-DD.md` 文件 + - 将日志整合更新到 `Memory.md` + +```mermaid +flowchart TD + A[Main Agent] --> B[写入 Memory.md] + A --> C[写入 YYYY-MM-DD.md] + D[上下文达到压缩阈值?] -->|是| E[Summarizer Agent] + E --> F[记录: 个性化信息] + E --> G[记录: Try-Error 经验] + H[定时任务 00:00] --> I[汇总最近 YYYY-MM-DD.md] + I --> J[更新 Memory.md] +``` + +这套机制确保了短期对话中的重要信息能够沉淀为长期记忆，不因对话结束而丢失。 + +--- + +## 小结 + +CoPaw 上下文管理的核心设计哲学可以用一句话概括： + +**让内存只放"现在需要的"，让文件系统保管"之后可能需要的"。** + +通过 Pre-Reasoning Hook 的四步流程、统一的两阶段截断策略，以及文件系统的持久化支撑，CoPaw 在有限的上下文窗口内为 Agent +提供了最大程度的信息可用性——无论对话持续多久，Agent 总能找到它需要的上下文。 + +--- + +*本文设计对应实现可参考 [CoPaw MemoryManager](https://github.com/agentscope-ai/CoPaw/blob/main/src/copaw/agents/memory/reme_light_memory_manager.py) +与 [ReMe ReMeLight](https://github.com/agentscope-ai/ReMe)。* From 9ad81209592062320aa49d53b90ae64801b95a00 Mon Sep 17 00:00:00 2001 From: jinliyl <6469360+jinliyl@users.noreply.github.com> Date: Tue, 31 Mar 2026 16:22:39 +0800 Subject: [PATCH 55/59] feat(compactor): add extra instruction support and improve error handing (#190) * feat(compactor): add extra instruction support and improve error handling - Add extra_instruction parameter to compactor for custom guidance during message compaction - Implement try-catch blocks around AS LLM initialization with detailed error logging - Add extra_instruction parameter to ReMe.compact method with comprehensive documentation - Update agentscope dependency from 1.0.17 to 1.0.18 in light installation - Bump version number from 0.3.1.6 to 0.3.1.7 - Pass extra_instruction parameter through compactor instantiation and execution flow * fix(core): add error handling for AS LLM formatters and token counters initialization - Wrapped AS LLM formatters initialization in try-except blocks - Added specific error logging for failed AS LLM formatter initialization - Wrapped AS token counters initialization in try-except blocks - Added specific error logging for failed AS token counter initialization - Applied same error handling pattern to both initial setup and restart operations - Maintained existing warning logs for unsupported backends --- pyproject.toml | 2 +- reme/__init__.py | 2 +- reme/core/application.py | 72 ++++++++++++------- .../memory/file_based/components/compactor.py | 5 ++ reme/reme_light.py | 8 +++ 5 files changed, 60 insertions(+), 29 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 36b12e98..91956ce2 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -83,7 +83,7 @@ litellm = [ ] light = [ - "agentscope==1.0.17", + "agentscope==1.0.18", "flowllm[reme]>=0.2.0.10", ] diff --git a/reme/__init__.py b/reme/__init__.py index f9ade852..3c235eb0 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.1.6" +__version__ = "0.3.1.7" __all__ = [ "config", diff --git a/reme/core/application.py b/reme/core/application.py index e2e5b3c0..8b5b1932 100644 --- a/reme/core/application.py +++ b/reme/core/application.py @@ -173,28 +173,37 @@ class Application: if config.backend not in R.as_llms: logger.warning(f"AS LLM backend {config.backend} is not supported.") else: - 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) + 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: - config_dict = config.model_dump(exclude={"backend"}) - self.service_context.as_llm_formatters[name] = R.as_llm_formatters[config.backend](**config_dict) + 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: - config_dict = config.model_dump(exclude={"backend"}) - self.service_context.as_token_counters[name] = R.as_token_counters[config.backend](**config_dict) + 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: @@ -287,15 +296,18 @@ class Application: logger.warning(f"AS LLM backend {config.get('backend')} is not supported.") continue - 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}") + 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: @@ -308,9 +320,12 @@ class Application: if config.get("backend") not in R.as_llm_formatters: logger.warning(f"AS LLM formatter backend {config.get('backend')} is not supported.") continue - 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}") + 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: @@ -323,9 +338,12 @@ class Application: if config.get("backend") not in R.as_token_counters: logger.warning(f"Token counter backend {config.get('backend')} is not supported.") continue - 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}") + 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: diff --git a/reme/memory/file_based/components/compactor.py b/reme/memory/file_based/components/compactor.py index 7e8b31ad..e8fe9383 100644 --- a/reme/memory/file_based/components/compactor.py +++ b/reme/memory/file_based/components/compactor.py @@ -35,6 +35,7 @@ class Compactor(BaseOp): console_enabled: bool = False, return_dict: bool = False, add_thinking_block: bool = True, + extra_instruction: str = "", **kwargs, ): super().__init__(**kwargs) @@ -42,6 +43,7 @@ class Compactor(BaseOp): 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): @@ -84,6 +86,9 @@ class Compactor(BaseOp): ) 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( diff --git a/reme/reme_light.py b/reme/reme_light.py index 5cad2f64..91e738ed 100644 --- a/reme/reme_light.py +++ b/reme/reme_light.py @@ -369,6 +369,7 @@ class ReMeLight(Application): 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. @@ -395,6 +396,12 @@ class ReMeLight(Application): 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 | dict: The condensed summary string, or a dict containing @@ -410,6 +417,7 @@ class ReMeLight(Application): language=language if language == "zh" else "", return_dict=return_dict, add_thinking_block=add_thinking_block, + extra_instruction=extra_instruction, ) return await compactor.call( From a97635752b41ebbf3e8787b3c659fb130d4a9e4d Mon Sep 17 00:00:00 2001 From: jinliyl <6469360+jinliyl@users.noreply.github.com> Date: Tue, 31 Mar 2026 20:07:00 +0800 Subject: [PATCH 56/59] fix(core): handle chromadb import error gracefully (#192) - Changed CHROMADB_AVAILABLE flag to _CHROMADB_IMPORT_ERROR exception storage - Updated version from 0.3.1.7 to 0.3.1.8 - Modified import error handling to preserve original exception details - Removed hardcoded ImportError message in favor of dynamic exception raising - Added proper logger initialization using get_logger utility --- reme/__init__.py | 2 +- reme/core/file_store/chroma_file_store.py | 17 ++++++++--------- 2 files changed, 9 insertions(+), 10 deletions(-) diff --git a/reme/__init__.py b/reme/__init__.py index 3c235eb0..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.1.7" +__version__ = "0.3.1.8" __all__ = [ "config", diff --git a/reme/core/file_store/chroma_file_store.py b/reme/core/file_store/chroma_file_store.py index d87f5fef..9892f7d1 100644 --- a/reme/core/file_store/chroma_file_store.py +++ b/reme/core/file_store/chroma_file_store.py @@ -5,19 +5,20 @@ 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: ImportError | None = None +except ImportError as e: + _CHROMADB_IMPORT_ERROR = e chromadb = None Settings = None @@ -39,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 From d5c929722bfec6bb4268d5430d62c69792582098 Mon Sep 17 00:00:00 2001 From: jinliyl <6469360+jinliyl@users.noreply.github.com> Date: Tue, 31 Mar 2026 20:59:59 +0800 Subject: [PATCH 57/59] refactor(file_store): move sqlite3 imports inside initialization methods (#193) * fix(core): handle chromadb import error gracefully - Changed CHROMADB_AVAILABLE flag to _CHROMADB_IMPORT_ERROR exception storage - Updated version from 0.3.1.7 to 0.3.1.8 - Modified import error handling to preserve original exception details - Removed hardcoded ImportError message in favor of dynamic exception raising - Added proper logger initialization using get_logger utility * refactor(file_store): move sqlite3 imports inside initialization methods - Moved sqlite3 import from module level to inside init methods - Removed unused import statement at top of file - Maintains same functionality while improving import organization - Prevents potential issues with early sqlite3 dependency loading * refactor(core): update import error handling with broader exception types - Changed ImportError to Exception for ray import error handling - Updated chromadb import error to use Exception instead of ImportError - Modified elasticsearch import error to catch general exceptions - Changed asyncpg import error handling from ImportError to Exception - Updated qdrant import error to use Exception instead of ImportError - Added explicit type hints for all import error variables as Exception | None --- reme/core/file_store/chroma_file_store.py | 4 ++-- reme/core/file_store/sqlite_file_store.py | 4 +++- reme/core/op/base_ray_op.py | 4 ++-- reme/core/vector_store/chroma_vector_store.py | 4 ++-- reme/core/vector_store/es_vector_store.py | 4 ++-- reme/core/vector_store/pgvector_store.py | 4 ++-- reme/core/vector_store/qdrant_vector_store.py | 4 ++-- 7 files changed, 15 insertions(+), 13 deletions(-) diff --git a/reme/core/file_store/chroma_file_store.py b/reme/core/file_store/chroma_file_store.py index 9892f7d1..6dda41b8 100644 --- a/reme/core/file_store/chroma_file_store.py +++ b/reme/core/file_store/chroma_file_store.py @@ -16,8 +16,8 @@ try: import chromadb from chromadb.config import Settings - _CHROMADB_IMPORT_ERROR: ImportError | None = None -except ImportError as e: + _CHROMADB_IMPORT_ERROR: Exception | None = None +except Exception as e: _CHROMADB_IMPORT_ERROR = e chromadb = None Settings = None 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/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/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/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 From 935e886af37ef357a844d07152dd4df26710d0eb Mon Sep 17 00:00:00 2001 From: Zhouwk <57825291+nitwtog@users.noreply.github.com> Date: Tue, 7 Apr 2026 14:20:30 +0800 Subject: [PATCH 58/59] =?UTF-8?q?=E6=9B=B4=E6=96=B0longmemeval=E5=92=8Chal?= =?UTF-8?q?umem=E7=9A=84quick=20start=20(#194)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * feat(reme): 添加配置选项以启用或禁用个人资料功能 - 在 ReMe 初始化方法中添加 enable_profile 参数，默认值为 True - 根据 enable_profile 设置决定是否创建 profile 目录和设置 profile_dir - 在 PersonalSummarizer 中根据 enable_profile 条件性地添加个人资料相关工具 - 在 PersonalRetriever 中根据 enable_profile 条件性地添加 ReadAllProfiles 工具 - 修改 profile_path 属性以在禁用个人资料时返回 None - 修改 get_profile_handler 方法以在禁用个人资料时返回 None - 为 enable_profile 参数添加文档说明其用于云向量存储场景 * refactor(benchmark): 重构LongMemEval基准测试中的ReMe实例管理 - 移除未使用的shutil导入 - 将固定的ReMe实例改为每个问题创建独立实例以实现隔离 - 更新LLM配置名称从qwen3-max-think到qwen-max-t - 修改模型调用逻辑使用正确的model_name参数 - 添加qwen-flash和GPT-4o-mini等新模型配置 - 统一使用"User"作为用户名，通过集合名实现隔离 - 调整并发处理数从4降至1，批处理大小从10增至30 - 每个问题类型采样数从2增至4 - 添加异步上下文管理确保资源正确释放 * reformat 2 files * refactor(benchmark): 重构长记忆评估中的模型配置 - 将原有的 eval_model_name 替换为专门的 retrieve_model_name 用于检索操作 - 添加对 qwen-max 模型配置的支持 - 更新参数解析器以支持新的检索模型参数 - 修改最大并发数默认值从 1 提升到 4 - 调整样本数量默认值从 4 减少到 1 - 统一模型参数命名规范，区分摘要、检索和评估模型 - 优化内存处理器初始化逻辑，支持独立的检索模型配置 * fix(benchmark): 移除数据路径默认值并设为必填参数 - 将LongMemEval评估脚本中的data_path参数改为必需参数 - 将HaluMem评估脚本中的data_path参数改为必需参数 - 删除了硬编码的默认文件路径配置 - 强制用户显式指定数据集文件路径以避免路径错误 * Update __init__.py * Update __init__.py * fix(benchmark): 修复ReMe评估中的模型配置和空值处理问题 - 移除了retrieve_memory调用中不需要的llm_config_name参数 - 修复了长字符串打印的换行格式问题 - 添加了eval_result为空时的初始化处理 - 在accuracy评估中加入了eval_model_name参数传递 * style(benchmark): 格式化模型名称打印输出 - 移除了多行字符串中的换行符和多余空格 - 将模型名称信息合并为单行连续显示 - 保持了原有的打印格式和信息完整性 * docs(readme): 更新文档添加实验结果表格 - 在英文版 README 中添加 🧪 Experiments 章节 - 添加 LoCoMo 和 HaluMem 两个基准测试的结果表格 - 在中文版 README_ZH 中添加 🧪 实验 章节 - 添加 LoCoMo 和 HaluMem 测试集的实验配置说明 - 添加完整的实验数据对比表格和评估协议说明 * docs(readme): 更新文档中的内存系统链接 - 为基于文件的记忆系统添加锚点链接 - 为基于向量库的记忆系统添加锚点链接 - 修复英文文档中的链接格式 - 修复中文文档中的链接格式和空行问题 * docs(readme): update experimental results section in documentation - Remove outdated experimental data placeholder "Coming soon..." - Add complete evaluation results for LoCoMo and HaluMem benchmarks - Include detailed performance metrics tables for all memory methods - Update experimental settings description with ReMe backbone details - Align evaluation protocol information with LLM-as-a-Judge approach - Maintain consistent formatting between English and Chinese documentation * docs(benchmark): add quick start guides for halumem and longmemeval experiments - Created HaluMem experiment quick start guide with ReMe integration setup - Added detailed steps for installing ReMe environment using conda - Included repository cloning instructions for HaluMem benchmark - Provided complete command examples for running HaluMem experiments - Created LongMeMEval quick start guide with data download procedures - Added wget commands for downloading cleaned dataset files - Included evaluation script instructions for computing experiment statistics - Documented parameter configurations for different model types and batch sizes * docs(longmemeval): update quickstart guide documentation - Changed project name from Halumem to Longmemeval in title - Updated description to reference Longmemeval experiments instead of Halumem - Maintained existing ReMe integration instructions unchanged --- benchmark/halumem/quickstart.md | 33 ++++++++++++++++++++ benchmark/longmemeval/quickstart.md | 48 +++++++++++++++++++++++++++++ 2 files changed, 81 insertions(+) create mode 100644 benchmark/halumem/quickstart.md create mode 100644 benchmark/longmemeval/quickstart.md 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/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 From f3d09aaa3899cf48f19644f9f98e0220526056e6 Mon Sep 17 00:00:00 2001 From: Chojan Shang Date: Thu, 9 Apr 2026 09:17:52 +0100 Subject: [PATCH 59/59] feat(vector_store): add OceanBase/seekdb vector store implementation (#201) * feat(vector_store): add OceanBase as a VectorStore * refactor(obvec): make it cleaner * docs: add obvec related info * refactor: minor update * refactor: clean code and pass lint * docs: remove unrelated edit * docs: minor update --- README.md | 2 +- README_ZH.md | 2 +- docs/index.md | 2 +- docs/vector_store_api_guide.md | 59 ++- pyproject.toml | 3 + reme/core/vector_store/__init__.py | 3 + reme/core/vector_store/obvec_vector_store.py | 453 +++++++++++++++++++ tests/test_vector_store.py | 72 ++- 8 files changed, 578 insertions(+), 18 deletions(-) create mode 100644 reme/core/vector_store/obvec_vector_store.py diff --git a/README.md b/README.md index 7c93e489..781f5412 100644 --- a/README.md +++ b/README.md @@ -506,7 +506,7 @@ async def main(): "dimensions": 1024, }, default_vector_store_config={ - "backend": "local", # Supports local/chroma/qdrant/elasticsearch + "backend": "local", # Supports local/chroma/qdrant/elasticsearch/obvec }, ) await reme.start() diff --git a/README_ZH.md b/README_ZH.md index 5fa2e4d2..7e6ccd58 100644 --- a/README_ZH.md +++ b/README_ZH.md @@ -486,7 +486,7 @@ async def main(): "dimensions": 1024, }, default_vector_store_config={ - "backend": "local", # 支持 local/chroma/qdrant/elasticsearch + "backend": "local", # 支持 local/chroma/qdrant/elasticsearch/obvec }, ) await reme.start() diff --git a/docs/index.md b/docs/index.md index 5cbbcb56..d7e2e716 100644 --- a/docs/index.md +++ b/docs/index.md @@ -139,7 +139,7 @@ response = requests.post("http://localhost:8002/retrieve_task_memory", json={ ## 📚 Resources - **[Installation Guide](installation.md)**, **[Quick Start](quick_start.md)**: Get started quickly with practical examples -- **[Vector Storage Setup](vector_store_api_guide.md)**: Configure local/vector databases and usage +- **[Vector Storage Setup](vector_store_api_guide.md)**: Configure local, Elasticsearch, Qdrant, ChromaDB, or ObVec (OceanBase / seekdb via pyobvector) storage and usage - **[MCP Guide](mcp_quick_start.md)**: Create MCP services - **[Personal Memory](personal_memory/personal_memory.md)**, **[Task Memory](task_memory/task_memory.md)** & **[Tool Memory](tool_memory/tool_memory.md)**: Operators used in personal memory, task memory and tool memory. You can modify the config to customize the pipelines. - **[Example Collection](./cookbook/appworld/quickstart.md)**: Real use cases and best practices diff --git a/docs/vector_store_api_guide.md b/docs/vector_store_api_guide.md index df0800ef..b0a06eab 100644 --- a/docs/vector_store_api_guide.md +++ b/docs/vector_store_api_guide.md @@ -33,8 +33,9 @@ FlowLLM provides multiple Vector Store implementations tailored to different use - **QdrantVectorStore** ([source code](https://github.com/flowllm-ai/flowllm/blob/main/flowllm/core/vector_store/qdrant_vector_store.py)): Built on the Qdrant vector database, supporting high-performance vector search. Recommended for large-scale production environments. - **ChromaVectorStore** ([source code](https://github.com/flowllm-ai/flowllm/blob/main/flowllm/core/vector_store/chroma_vector_store.py)): Based on ChromaDB, providing persistent storage and metadata filtering capabilities. - **EsVectorStore** ([source code](https://github.com/flowllm-ai/flowllm/blob/main/flowllm/core/vector_store/es_vector_store.py)): Built on Elasticsearch, enabling powerful combined full-text and vector search functionalities. +- **ObVecVectorStore** ([source code](https://github.com/agentscope-ai/ReMe/blob/main/reme/core/vector_store/obvec_vector_store.py)): Uses [pyobvector](https://pypi.org/project/pyobvector/) against **OceanBase** or **seekdb** (MySQL-compatible wire protocol). Suitable when you already run OceanBase/seekdb or need a SQL-native vector table with HNSW-style ANN search and JSON metadata filters. -All Vector Store implementations inherit from **BaseVectorStore** ([source code](https://github.com/flowllm-ai/flowllm/blob/main/flowllm/core/vector_store/base_vector_store.py)), ensuring a consistent interface specification. +All Vector Store implementations inherit from **BaseVectorStore** ([source code](https://github.com/agentscope-ai/ReMe/blob/main/reme/core/vector_store/base_vector_store.py)) in ReMe, ensuring a consistent interface specification. ## Core Features @@ -108,6 +109,28 @@ The asynchronous interface is particularly useful in the following scenarios: - **hosts**: Elasticsearch host address(es), either a string or a list (default: `http://localhost:9200`). - **basic_auth**: Basic authentication credentials (username and password). +### ObVecVectorStore Configuration + +- **uri**: Server address as `host:port` (default: `127.0.0.1:2881`). +- **user**: MySQL-compatible user. seekdb single-tenant images often use `root`; OceanBase multi-tenant setups typically use `root@` (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/pyproject.toml b/pyproject.toml index 91956ce2..11177ffd 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -49,6 +49,9 @@ dependencies = [ "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", diff --git a/reme/core/vector_store/__init__.py b/reme/core/vector_store/__init__.py index 9d411d03..8429b911 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,5 +22,6 @@ __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/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/tests/test_vector_store.py b/tests/test_vector_store.py index ccdd9c28..7f7264be 100644 --- a/tests/test_vector_store.py +++ b/tests/test_vector_store.py @@ -2,8 +2,8 @@ """Unified test suite for vector store implementations. This module provides comprehensive test coverage for LocalVectorStore, ESVectorStore, -PGVectorStore, QdrantVectorStore, and ChromaVectorStore implementations. Tests can be -run for specific vector stores or all implementations. +PGVectorStore, QdrantVectorStore, ChromaVectorStore, and ObVecVectorStore implementations. +Tests can be run for specific vector stores or all implementations. Usage: python test_vector_store.py --local # Test LocalVectorStore only @@ -11,12 +11,13 @@ 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 --all # Test all vector stores - """ import argparse import asyncio +import os import shutil import tempfile from pathlib import Path @@ -32,12 +33,19 @@ from reme.core.vector_store import ( ChromaVectorStore, LocalVectorStore, ESVectorStore, + ObVecVectorStore, PGVectorStore, QdrantVectorStore, ) 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 ==================== @@ -73,6 +81,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 @@ -182,7 +199,7 @@ def get_store_type(store: BaseVectorStore) -> str: store: Vector store instance Returns: - str: Type identifier ("local", "es", "pgvector", "qdrant", or "chroma") + str: Type identifier ("local", "es", "pgvector", "qdrant", "chroma", or "obvec") """ if isinstance(store, LocalVectorStore): return "local" @@ -194,6 +211,8 @@ def get_store_type(store: BaseVectorStore) -> str: return "pgvector" elif isinstance(store, ChromaVectorStore): return "chroma" + elif isinstance(store, ObVecVectorStore): + return "obvec" else: raise ValueError(f"Unknown vector store type: {type(store)}") @@ -202,7 +221,7 @@ 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", or "chroma") + store_type: Type of vector store ("local", "es", "pgvector", "qdrant", "chroma", or "obvec") collection_name: Name of the collection Returns: @@ -264,6 +283,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, + ) else: raise ValueError(f"Unknown store type: {store_type}") @@ -327,7 +358,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", "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" @@ -581,9 +612,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() # Clean up if exists @@ -1010,7 +1041,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", "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]}...") @@ -1028,7 +1059,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", "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") @@ -1718,7 +1749,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) @@ -1752,6 +1783,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}") + logger.info("✓ Cleanup completed") except Exception as e: logger.error(f"Cleanup error: {e}") @@ -1772,6 +1810,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 --all # Test all vector stores """, ) @@ -1800,6 +1839,11 @@ Examples: action="store_true", help="Test ChromaVectorStore", ) + parser.add_argument( + "--obvec", + action="store_true", + help="Test ObVecVectorStore", + ) parser.add_argument( "--all", action="store_true", @@ -1818,6 +1862,7 @@ Examples: ("pgvector", "PGVectorStore"), ("qdrant", "QdrantVectorStore"), ("chroma", "ChromaVectorStore"), + ("obvec", "ObVecVectorStore"), ] else: # Build list based on individual flags @@ -1831,6 +1876,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 not stores_to_test: # Default to all vector stores if no argument provided @@ -1840,10 +1887,11 @@ Examples: ("pgvector", "PGVectorStore"), ("qdrant", "QdrantVectorStore"), ("chroma", "ChromaVectorStore"), + ("obvec", "ObVecVectorStore"), ] print("No vector store specified, defaulting to test all vector stores") print( - "Use --local/--es/--pgvector/--qdrant/--chroma to test specific ones\n", + "Use --local/--es/--pgvector/--qdrant/--chroma/--obvec to test specific ones\n", ) # Run tests for each vector store