From 819443813bc2437fc872bad57fb7caf88c3fc2cf Mon Sep 17 00:00:00 2001 From: "jinli.yl" Date: Thu, 16 Apr 2026 11:15:23 +0800 Subject: [PATCH] feat(components): add token counter and file-based utility components - Introduce BaseAsTokenCounter and EstimatedAsTokenCounter for token estimation - Add AsMsgStat and AsBlockStat schema for message statistics tracking - Implement FileIO class with read/write/append/edit operations - Create file utility functions for safe async file reading and truncation - Add MemorySearch component for semantic search in memory files - Register new component types in ComponentEnum and update imports - Add constants for default host, port, and truncation limits - Create BaseService abstract base class for service implementations - Implement BaseStep with component accessors and lifecycle management - Add proper __all__ exports for all new modules and components --- pyproject.toml | 1 + reme_cli/component/__init__.py | 2 + .../component/as_token_counter/__init__.py | 60 +++ .../estimate_token_counter.py | 22 ++ reme_cli/component/base_step.py | 27 +- reme_cli/component/runtime_context.py | 52 +++ reme_cli/component/service/base_service.py | 31 +- reme_cli/constants.py | 18 + reme_cli/enumeration/component_enum.py | 2 + reme_cli/file_based/__init__.py | 17 + reme_cli/file_based/file_io.py | 362 ++++++++++++++++++ reme_cli/file_based/file_utils.py | 202 ++++++++++ reme_cli/file_based/memory_search.py | 52 +++ reme_cli/file_based/summarizer.py | 324 ++++++++++++++++ reme_cli/file_based/summarizer.yaml | 62 +++ reme_cli/jobs/__init__.py | 1 - reme_cli/reme.py | 117 ++++-- reme_cli/schema/__init__.py | 3 + reme_cli/schema/as_msg_stat.py | 97 +++++ reme_cli/vector_based/__init__.py | 0 20 files changed, 1393 insertions(+), 59 deletions(-) create mode 100644 reme_cli/component/as_token_counter/__init__.py create mode 100644 reme_cli/component/as_token_counter/estimate_token_counter.py create mode 100644 reme_cli/file_based/__init__.py create mode 100644 reme_cli/file_based/file_io.py create mode 100644 reme_cli/file_based/file_utils.py create mode 100644 reme_cli/file_based/memory_search.py create mode 100644 reme_cli/file_based/summarizer.py create mode 100644 reme_cli/file_based/summarizer.yaml delete mode 100644 reme_cli/jobs/__init__.py create mode 100644 reme_cli/schema/as_msg_stat.py create mode 100644 reme_cli/vector_based/__init__.py diff --git a/pyproject.toml b/pyproject.toml index 7d1f73a2..bdea30f2 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -33,6 +33,7 @@ classifiers = [ keywords = ["llm", "memory", "experience", "memoryscope", "ai", "mcp", "http", "reme", "personal"] dependencies = [ + "aiofiles>=24.1.0", "sqlite-vec>=0.1.6", "prompt_toolkit>=3.0.52", "rich>=14.2.0", diff --git a/reme_cli/component/__init__.py b/reme_cli/component/__init__.py index dcd07fb6..5998242c 100644 --- a/reme_cli/component/__init__.py +++ b/reme_cli/component/__init__.py @@ -9,6 +9,7 @@ from .runtime_context import RuntimeContext from . import as_llm from . import as_llm_formatter +from . import as_token_counter from . import client from . import embedding from . import file_store @@ -27,6 +28,7 @@ __all__ = [ # base components "as_llm", "as_llm_formatter", + "client", "embedding", "file_store", diff --git a/reme_cli/component/as_token_counter/__init__.py b/reme_cli/component/as_token_counter/__init__.py new file mode 100644 index 00000000..a39b2acf --- /dev/null +++ b/reme_cli/component/as_token_counter/__init__.py @@ -0,0 +1,60 @@ +"""AgentScope TokenCounter wrappers.""" + +from .estimate_token_counter import EstimatedTokenCounter +from ..base_component import BaseComponent +from ..component_registry import R +from ...enumeration import ComponentEnum + + +class BaseAsTokenCounter(BaseComponent): + """Base wrapper for token counters. + + Subclasses should implement _start() to initialize self.token_counter. + """ + + component_type = ComponentEnum.AS_TOKEN_COUNTER + + def __init__(self, **kwargs) -> None: + """Initialize with token counter configuration kwargs.""" + super().__init__(**kwargs) + self.token_counter: EstimatedTokenCounter | None = None + + async def _start(self, app_context=None) -> None: + """Initialize the token counter. Override in subclasses.""" + + async def _close(self) -> None: + """Release token counter resources.""" + self.token_counter = None + + async def count(self, messages: list[dict], **kwargs) -> int: + """Count tokens in messages. + + Args: + messages: List of message dictionaries. + **kwargs: Additional arguments passed to the token counter. + + Returns: + Estimated token count. + + Raises: + RuntimeError: If token counter is not initialized. + """ + if self.token_counter is None: + raise RuntimeError("Token counter not initialized. Call start() first.") + return await self.token_counter.count(messages, **kwargs) + + +@R.register("estimated") +class EstimatedAsTokenCounter(BaseAsTokenCounter): + """Estimated token counter using character-based estimation.""" + + async def _start(self, app_context=None) -> None: + """Initialize the estimated token counter.""" + self.token_counter = EstimatedTokenCounter(**self.kwargs) + + +__all__ = [ + "BaseAsTokenCounter", + "EstimatedAsTokenCounter", + "EstimatedTokenCounter", +] \ No newline at end of file diff --git a/reme_cli/component/as_token_counter/estimate_token_counter.py b/reme_cli/component/as_token_counter/estimate_token_counter.py new file mode 100644 index 00000000..868e4702 --- /dev/null +++ b/reme_cli/component/as_token_counter/estimate_token_counter.py @@ -0,0 +1,22 @@ +from typing import Any + +from agentscope.token import TokenCounterBase + + +class EstimatedTokenCounter(TokenCounterBase): + + def __init__(self, estimate_divisor: float = 4): + if estimate_divisor == 0: + raise ValueError("estimate_divisor cannot be zero") + self.estimate_divisor: float = estimate_divisor + + async def count( + self, + messages: list[dict], + text: str | None = None, + **kwargs: Any, + ) -> int: + if not text: + return 0 + else: + return int(len(text.encode("utf-8")) / self.estimate_divisor + 0.5) diff --git a/reme_cli/component/base_step.py b/reme_cli/component/base_step.py index 2c6a08f8..e78e1e9c 100644 --- a/reme_cli/component/base_step.py +++ b/reme_cli/component/base_step.py @@ -6,6 +6,7 @@ from abc import abstractmethod from .application_context import ApplicationContext from .as_llm import BaseAsLLM from .as_llm_formatter import BaseAsLLMFormatter +from .as_token_counter import BaseAsTokenCounter from .base_component import BaseComponent from .embedding import BaseEmbeddingModel from .file_store import BaseFileStore @@ -29,13 +30,13 @@ class BaseStep(BaseComponent): return instance def __init__( - self, - name: str = "", - language: str = "", - prompt_dict: dict[str, str] | None = None, - input_mapping: dict[str, str] | None = None, - output_mapping: dict[str, str] | None = None, - **kwargs, + self, + name: str = "", + language: str = "", + prompt_dict: dict[str, str] | None = None, + input_mapping: dict[str, str] | None = None, + output_mapping: dict[str, str] | None = None, + **kwargs, ): """Initialize step configurations.""" super().__init__(**kwargs) @@ -106,6 +107,18 @@ class BaseStep(BaseComponent): raise TypeError(f"{name} is not a BaseAsLLMFormatter instance.") return formatter + @property + def as_token_counter(self): + """Get the TokenCounter instance by name.""" + name: str = self.kwargs.get("as_token_counter", "default") + counters = self.application_context.components[ComponentEnum.AS_TOKEN_COUNTER] + if name not in counters: + raise ValueError(f"AsTokenCounter {name} not found.") + counter = counters[name] + if not isinstance(counter, BaseAsTokenCounter): + raise TypeError(f"{name} is not a BaseAsTokenCounter instance.") + return counter + @property def file_store(self) -> BaseFileStore: """Get the FileStore instance by name.""" diff --git a/reme_cli/component/runtime_context.py b/reme_cli/component/runtime_context.py index b56fb232..af455536 100644 --- a/reme_cli/component/runtime_context.py +++ b/reme_cli/component/runtime_context.py @@ -14,6 +14,58 @@ class RuntimeContext: """Initialize the context with all keyword arguments stored in data.""" self.data: dict = kwargs + def get(self, key: str, default=None): + """Get a value from data by key, with optional default.""" + return self.data.get(key, default) + + def set(self, key: str, value) -> "RuntimeContext": + """Set a value in data by key.""" + self.data[key] = value + return self + + def delete(self, key: str) -> "RuntimeContext": + """Delete a key from data.""" + if key in self.data: + del self.data[key] + return self + + def contains(self, key: str) -> bool: + """Check if a key exists in data.""" + return key in self.data + + def update(self, data: dict) -> "RuntimeContext": + """Update data with a dictionary.""" + self.data.update(data) + return self + + def keys(self) -> list[str]: + """Get all keys in data.""" + return list(self.data.keys()) + + def values(self) -> list: + """Get all values in data.""" + return list(self.data.values()) + + def items(self) -> list[tuple]: + """Get all key-value pairs in data.""" + return list(self.data.items()) + + def __getitem__(self, key: str): + """Get a value using bracket syntax.""" + return self.data[key] + + def __setitem__(self, key: str, value): + """Set a value using bracket syntax.""" + self.data[key] = value + + def __delitem__(self, key: str): + """Delete a key using bracket syntax.""" + del self.data[key] + + def __contains__(self, key: str) -> bool: + """Check if a key exists using 'in' operator.""" + return key in self.data + @property def response(self) -> Response: """Get or create the response object.""" diff --git a/reme_cli/component/service/base_service.py b/reme_cli/component/service/base_service.py index 897efd39..d70b4939 100644 --- a/reme_cli/component/service/base_service.py +++ b/reme_cli/component/service/base_service.py @@ -1,11 +1,14 @@ """Abstract base class for service implementations.""" - from abc import abstractmethod +from typing import TYPE_CHECKING from ..base_component import BaseComponent from ..job.base_job import BaseJob from ...enumeration import ComponentEnum +if TYPE_CHECKING: + from ...application import Application + class BaseService(BaseComponent): """Abstract base class for services that expose jobs. @@ -16,14 +19,8 @@ class BaseService(BaseComponent): component_type = ComponentEnum.SERVICE - from ...application import Application - def __init__(self, **kwargs): - """Initialize the service. - - Args: - **kwargs: Additional service-specific configuration. - """ + """Initialize the service.""" super().__init__(**kwargs) self.service = None @@ -34,20 +31,12 @@ class BaseService(BaseComponent): """Default empty implementation for sync services.""" @abstractmethod - def add_job(self, job: BaseJob) -> None: - """Register a job with the service. - - Args: - job: The job to register. - """ + def build_service(self, app: "Application") -> None: + """Build the service.""" @abstractmethod - def build_service(self, app: "Application") -> None: - """Build the service. - - Args: - app: The application instance. - """ + def add_job(self, job: BaseJob) -> None: + """Register a job with the service.""" @abstractmethod def start_service(self, app: "Application") -> None: @@ -58,7 +47,7 @@ class BaseService(BaseComponent): for name, job in app.context.jobs.values(): try: self.add_job(job) - self.logger.info(f"Added job {name}") + self.logger.info(f"Successfully Added job {name}") except Exception as e: self.logger.error(f"Failed to add job {name}: {e}") diff --git a/reme_cli/constants.py b/reme_cli/constants.py index fce66e64..d866b97b 100644 --- a/reme_cli/constants.py +++ b/reme_cli/constants.py @@ -5,3 +5,21 @@ REME_SERVICE_INFO = "REME_SERVICE_INFO" REME_DEFAULT_HOST = "127.0.0.1" REME_DEFAULT_PORT = 2333 + +# Default truncation limit for text output +DEFAULT_MAX_BYTES = 50 * 1024 + +# Maximum file size to read into memory (1GB) +MAX_FILE_READ_BYTES = 1024 * 1024 * 1024 + +# Marker prepended to every truncation notice. +# Format: +# <<>> +# The output above was truncated. +# The full content is saved to the file and contains Z lines in total. +# This excerpt starts at line X and covers the next N bytes. +# If the current content is not enough, call `read_file` with file_path= start_line=Y to read more. +# +# Split output on this marker to recover the original (untruncated) portion: +# original = output.split(TRUNCATION_NOTICE_MARKER)[0] +TRUNCATION_NOTICE_MARKER = "<<>>" diff --git a/reme_cli/enumeration/component_enum.py b/reme_cli/enumeration/component_enum.py index bce70d7b..2bf13579 100644 --- a/reme_cli/enumeration/component_enum.py +++ b/reme_cli/enumeration/component_enum.py @@ -20,6 +20,8 @@ class ComponentEnum(str, Enum): AS_LLM_FORMATTER = "as_llm_formatter" + AS_TOKEN_COUNTER = "as_token_counter" + EMBEDDING_MODEL = "embedding_model" FILE_STORE = "file_store" diff --git a/reme_cli/file_based/__init__.py b/reme_cli/file_based/__init__.py new file mode 100644 index 00000000..81b58f52 --- /dev/null +++ b/reme_cli/file_based/__init__.py @@ -0,0 +1,17 @@ +"""File-based components and utilities.""" + +from .file_io import FileIO +from .file_utils import ( + async_read_file_safe, + truncate_text_output, +) +from .memory_search import MemorySearch +from .summarizer import Summarizer + +__all__ = [ + "FileIO", + "async_read_file_safe", + "truncate_text_output", + "MemorySearch", + "Summarizer", +] diff --git a/reme_cli/file_based/file_io.py b/reme_cli/file_based/file_io.py new file mode 100644 index 00000000..a7459440 --- /dev/null +++ b/reme_cli/file_based/file_io.py @@ -0,0 +1,362 @@ +"""File I/O operations with a configurable working directory.""" + +import os +from pathlib import Path + +import aiofiles +from agentscope.message import TextBlock +from agentscope.tool import ToolResponse + +from .file_utils import async_read_file_safe, truncate_text_output +from ..constants import TRUNCATION_NOTICE_MARKER + + +class FileIO: + """File I/O operations with a configurable working directory.""" + + def __init__(self, working_dir: str | Path): + """Initialize FileIO with a working directory. + + Args: + working_dir (`str`): + The working directory for resolving relative paths. + """ + self.working_dir = Path(working_dir) + + def _resolve_file_path(self, file_path: str) -> str: + """Resolve file path: use absolute path as-is, + resolve relative path from working_dir. + + Args: + file_path: The input file path (absolute or relative). + + Returns: + The resolved absolute file path as string. + """ + path = Path(file_path).expanduser() + if path.is_absolute(): + return str(path) + else: + return str(self.working_dir / file_path) + + async def read_file( # pylint: disable=too-many-return-statements + self, + file_path: str, + start_line: int | None = None, + end_line: int | None = None, + ) -> ToolResponse: + """Read a file. Relative paths resolve from WORKING_DIR. + + Use start_line/end_line to read a specific line range (output includes + line numbers). Omit both to read the full file. + + Args: + file_path (`str`): + Path to the file. + start_line (`int`, optional): + First line to read (1-based, inclusive). + end_line (`int`, optional): + Last line to read (1-based, inclusive). + """ + + # Convert start_line/end_line to int if they are strings + if start_line is not None: + try: + start_line = int(start_line) + except (ValueError, TypeError): + return ToolResponse( + content=[ + TextBlock( + type="text", + text=f"Error: start_line must be an integer, got {start_line!r}.", + ), + ], + ) + + if end_line is not None: + try: + end_line = int(end_line) + except (ValueError, TypeError): + return ToolResponse( + content=[ + TextBlock( + type="text", + text=f"Error: end_line must be an integer, got {end_line!r}.", + ), + ], + ) + + file_path = self._resolve_file_path(file_path) + + if not os.path.exists(file_path): + return ToolResponse( + content=[ + TextBlock( + type="text", + text=f"Error: The file {file_path} does not exist.", + ), + ], + ) + + if not os.path.isfile(file_path): + return ToolResponse( + content=[ + TextBlock( + type="text", + text=f"Error: The path {file_path} is not a file.", + ), + ], + ) + + try: + content = await async_read_file_safe(file_path) + all_lines = content.split("\n") + total = len(all_lines) + + # Determine read range + s = max(1, start_line if start_line is not None else 1) + e = min(total, end_line if end_line is not None else total) + + if s > total: + return ToolResponse( + content=[ + TextBlock( + type="text", + text=f"Error: start_line {s} exceeds file length ({total} lines).", + ), + ], + ) + + if s > e: + return ToolResponse( + content=[ + TextBlock( + type="text", + text=f"Error: start_line ({s}) > end_line ({e}).", + ), + ], + ) + + # Extract selected lines + selected_content = "\n".join(all_lines[s - 1: e]) + + # Apply smart truncation (consistent with shell output format) + text = truncate_text_output( + selected_content, + start_line=s, + total_lines=total, + file_path=file_path, + ) + + # Add continuation hint if partial read without truncation. + # Use TRUNCATION_NOTICE_MARKER format so ToolResultCompactor can + # re-truncate with the correct start_line when compacting old messages. + if text == selected_content and e < total: + content_bytes = len(text.encode("utf-8")) + notice = ( + TRUNCATION_NOTICE_MARKER + + f"\nThe output above was truncated." + f"\nThe full content is saved to the file " + f"and contains {total} lines in total." + f"\nThis excerpt starts at line {s} and " + f"covers the next {content_bytes} bytes." + "\nIf the current content is not enough, " + f"call `read_file` with file_path={file_path} start_line={e + 1} to read more." + ) + text = text + notice + + return ToolResponse( + content=[TextBlock(type="text", text=text)], + ) + + except Exception as e: + return ToolResponse( + content=[ + TextBlock( + type="text", + text=f"Error: Read file failed due to \n{e}", + ), + ], + ) + + async def write_file( + self, + file_path: str, + content: str, + ) -> ToolResponse: + """Create or overwrite a file. Relative paths resolve from working_dir. + + Args: + file_path (`str`): + Path to the file. + content (`str`): + Content to write. + """ + if not file_path: + return ToolResponse( + content=[ + TextBlock( + type="text", + text="Error: No `file_path` provided.", + ), + ], + ) + + file_path = self._resolve_file_path(file_path) + + try: + async with aiofiles.open(file_path, "w", encoding="utf-8") as file: + await file.write(content) + return ToolResponse( + content=[ + TextBlock( + type="text", + text=f"Wrote {len(content)} bytes to {file_path}.", + ), + ], + ) + except Exception as e: + return ToolResponse( + content=[ + TextBlock( + type="text", + text=f"Error: Write file failed due to \n{e}", + ), + ], + ) + + # pylint: disable=too-many-return-statements + async def edit_file( + self, + file_path: str, + old_text: str, + new_text: str, + ) -> ToolResponse: + """Find-and-replace text in a file. All occurrences of old_text are + replaced with new_text. Relative paths resolve from working_dir. + + Args: + file_path (`str`): + Path to the file. + old_text (`str`): + Exact text to find. + new_text (`str`): + Replacement text. + """ + if not file_path: + return ToolResponse( + content=[ + TextBlock( + type="text", + text="Error: No `file_path` provided.", + ), + ], + ) + + resolved_path = self._resolve_file_path(file_path) + + if not os.path.exists(resolved_path): + return ToolResponse( + content=[ + TextBlock( + type="text", + text=f"Error: The file {resolved_path} does not exist.", + ), + ], + ) + + if not os.path.isfile(resolved_path): + return ToolResponse( + content=[ + TextBlock( + type="text", + text=f"Error: The path {resolved_path} is not a file.", + ), + ], + ) + + try: + content = await async_read_file_safe(resolved_path) + except Exception as e: + return ToolResponse( + content=[ + TextBlock( + type="text", + text=f"Error: Read file failed due to \n{e}", + ), + ], + ) + + if old_text not in content: + return ToolResponse( + content=[ + TextBlock( + type="text", + text=f"Error: The text to replace was not found in {file_path}.", + ), + ], + ) + + new_content = content.replace(old_text, new_text) + write_response = await self.write_file(file_path=resolved_path, content=new_content) + + if write_response.content and len(write_response.content) > 0: + write_text = write_response.content[0].get("text", "") + if write_text.startswith("Error:"): + return write_response + + return ToolResponse( + content=[ + TextBlock( + type="text", + text=f"Successfully replaced text in {file_path}.", + ), + ], + ) + + async def append_file( + self, + file_path: str, + content: str, + ) -> ToolResponse: + """Append content to the end of a file. Relative paths resolve from + working_dir. + + Args: + file_path (`str`): + Path to the file. + content (`str`): + Content to append. + """ + if not file_path: + return ToolResponse( + content=[ + TextBlock( + type="text", + text="Error: No `file_path` provided.", + ), + ], + ) + + file_path = self._resolve_file_path(file_path) + + try: + async with aiofiles.open(file_path, "a", encoding="utf-8") as file: + await 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_cli/file_based/file_utils.py b/reme_cli/file_based/file_utils.py new file mode 100644 index 00000000..89981f63 --- /dev/null +++ b/reme_cli/file_based/file_utils.py @@ -0,0 +1,202 @@ +"""Shared utilities for file and shell tools.""" + +import re + +import aiofiles + +from ..constants import ( + DEFAULT_MAX_BYTES, + MAX_FILE_READ_BYTES, + TRUNCATION_NOTICE_MARKER, +) + + +def _truncate_fresh( + text: str, + start_line: int, + total_lines: int, + max_bytes: int, + file_path: str | None, + encoding: str, +) -> str: + """Truncate fresh text (no prior truncation marker) by bytes with line integrity. + + Slices at the byte boundary and appends a truncation notice with a continuation + hint so callers know which line to read next. + + Returns the original text unchanged when it fits within max_bytes, or when the + last line itself exceeds max_bytes (unhandled edge case). + """ + text_bytes = text.encode(encoding) + + # Under the byte limit — return as-is without any modification. + if len(text_bytes) <= max_bytes: + return text + + # Slice at the byte boundary. + # Assuming every single line is shorter than DEFAULT_MAX_BYTES, this cut always + # lands mid-line, guaranteeing at least one complete line before the boundary. + # Lines that exceed DEFAULT_MAX_BYTES are not handled and may be skipped entirely. + truncated = text_bytes[:max_bytes] + # Decode back to str; errors="ignore" drops any split multibyte 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 multibyte character at the cut boundary. + result = truncated_bytes.decode(encoding, errors="ignore") + # Each '\n' in result corresponds to one fully-included line; + # anything after the last '\n' is a partial line that was cut off. + newline_count = result.count("\n") + + # The next read should start at the line immediately after all complete lines. + # max(1, ...) guards against the theoretical zero-newline case + # (impossible when every line is shorter than DEFAULT_MAX_BYTES). + next_line = start_line_parsed + max(1, newline_count) + + if not re.search(r"covers the next \d+ bytes", old_notice): + return text + # _truncate_fresh always includes a continuation hint, so both fields are always present. + new_notice = re.sub(r"covers the next \d+ bytes", f"covers the next {max_bytes} bytes", old_notice) + new_notice = re.sub(r"start_line=\d+ to read more", f"start_line={next_line} to read more", new_notice) + + return result + TRUNCATION_NOTICE_MARKER + new_notice + + +def truncate_text_output( + text: str, + start_line: int = 1, + total_lines: int = 0, + max_bytes: int = DEFAULT_MAX_BYTES, + file_path: str | None = None, + encoding: str = "utf-8", +) -> str: + """Truncate file output by bytes with line integrity. + + If text is under byte limit, return as-is. + If over limit, truncate at the last complete line that fits, + allowing the next read to start from a fresh line. + + Dispatches to :func:`_truncate_fresh` for text seen for the first time, or to + :func:`_retruncate` when the text already contains a TRUNCATION_NOTICE_MARKER + from a previous pass. + + Args: + text: The output text to truncate. + start_line: The starting line number (1-based). Ignored when text already + contains a truncation notice (values are parsed from the notice instead). + total_lines: Total lines in the original file. Ignored when text already + contains a truncation notice (values are parsed from the notice instead). + max_bytes: Maximum size in bytes. + file_path: Optional file path to include in the truncation notice. + encoding: Character encoding used for byte-length calculation and decoding. + + Returns: + Truncated text with notice if truncated. + """ + if not text: + return text + if max_bytes <= 0: + return text + + try: + if TRUNCATION_NOTICE_MARKER in text: + return _retruncate(text, max_bytes=max_bytes, encoding=encoding) + else: + return _truncate_fresh( + text, + start_line=start_line, + total_lines=total_lines, + max_bytes=max_bytes, + file_path=file_path, + encoding=encoding, + ) + except Exception: + return text + + +async def async_read_file_safe(file_path: str, max_bytes: int = MAX_FILE_READ_BYTES) -> str: + """Async version of read_file_safe with Unicode error handling and memory protection. + + Args: + file_path: Path to the file. + max_bytes: Maximum bytes to read into memory (default 1GB). + + Returns: + File content as string (up to max_bytes). + """ + try: + async with aiofiles.open(file_path, "r", encoding="utf-8") as f: + return await f.read(max_bytes) + except UnicodeDecodeError: + async with aiofiles.open(file_path, "r", encoding="utf-8", errors="ignore") as f: + return await f.read(max_bytes) diff --git a/reme_cli/file_based/memory_search.py b/reme_cli/file_based/memory_search.py new file mode 100644 index 00000000..77cab3a6 --- /dev/null +++ b/reme_cli/file_based/memory_search.py @@ -0,0 +1,52 @@ +"""Memory search step for semantic search in memory files.""" + +import json + +from ..component import R +from ..component.base_step import BaseStep +from ..enumeration import ComponentEnum + + +@R.register("memory_search") +class MemorySearch(BaseStep): + """Semantically search MEMORY.md and memory files.""" + + component_type = ComponentEnum.STEP + + def __init__(self, vector_weight: float = 0.7, candidate_multiplier: float = 3.0, **kwargs): + """Initialize memory search step. + + Args: + vector_weight: Weight for vector search vs keyword search. + candidate_multiplier: Multiplier for candidate count before filtering. + **kwargs: Additional arguments passed to BaseStep. + """ + super().__init__(**kwargs) + self.vector_weight = vector_weight + self.candidate_multiplier = candidate_multiplier + + async def execute(self): + """Execute the memory search operation.""" + assert self.context is not None, "Context is not set" + query: str = self.context.get("query", "").strip() + min_score: float = self.context.get("min_score", 0.1) + max_results: int = self.context.get("max_results", 5) + + assert query, "Query cannot be empty" + assert isinstance(min_score, float | int) and 0.0 <= min_score <= 1.0, \ + f"min_score must be between 0 and 1, got {min_score}" + assert isinstance(max_results, int) and max_results > 0, \ + f"max_results must be a positive integer, got {max_results}" + + # Use hybrid_search from file_store + results = await self.file_store.hybrid_search( + query=query, + limit=max_results, + vector_weight=self.vector_weight, + candidate_multiplier=self.candidate_multiplier, + ) + + # Filter by min_score + results = [r for r in results if r.score >= min_score] + + return json.dumps([result.model_dump(exclude_none=True) for result in results], indent=2, ensure_ascii=False) diff --git a/reme_cli/file_based/summarizer.py b/reme_cli/file_based/summarizer.py new file mode 100644 index 00000000..53284e2f --- /dev/null +++ b/reme_cli/file_based/summarizer.py @@ -0,0 +1,324 @@ +"""Summarizer module for memory summarization operations.""" + +import datetime +import json +import zoneinfo + +from agentscope.agent import ReActAgent +from agentscope.message import Msg +from agentscope.token import HuggingFaceTokenCounter +from agentscope.tool import Toolkit + +from ..component import BaseStep +from ..schema import AsMsgStat, AsBlockStat + + +class Summarizer(BaseStep): + """Summarizer step for summarizing memory messages.""" + + def __init__( + self, + working_dir: str, + memory_dir: str, + memory_compact_threshold: int, + toolkit: Toolkit | None = None, + console_enabled: bool = False, + timezone: str | None = None, + add_thinking_block: bool = True, + as_token_counter: HuggingFaceTokenCounter | None = None, + **kwargs, + ): + """Initialize the summarizer step. + + Args: + working_dir: Working directory path. + memory_dir: Memory directory path for storing summaries. + memory_compact_threshold: Token threshold for memory compaction. + toolkit: Optional toolkit for the agent. + console_enabled: Whether to enable console output. + timezone: Optional timezone string for date formatting. + add_thinking_block: Whether to include thinking blocks in output. + as_token_counter: Optional token counter instance. + **kwargs: Additional keyword arguments passed to BaseStep. + """ + super().__init__(**kwargs) + self.working_dir: str = working_dir + self.memory_dir: str = memory_dir + self.memory_compact_threshold: int = memory_compact_threshold + self.toolkit: Toolkit | None = toolkit + self.console_enabled: bool = console_enabled + self.timezone: str | None = timezone + self.add_thinking_block: bool = add_thinking_block + self._as_token_counter: HuggingFaceTokenCounter | None = as_token_counter + + + 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: + self.logger.error(f"Invalid timezone: {self.timezone}, falling back to local time error={e}") + return datetime.datetime.now() + + async def _count_str_token(self, text: str) -> int: + """Count tokens in a string.""" + return await self.as_token_counter.count(messages=[], text=text) + + async def _format_tool_result_output(self, output: str | list[dict]) -> tuple[str, int]: + """Convert tool result output to string.""" + if isinstance(output, str): + return output, await self._count_str_token(output) + + textual_parts = [] + total_token_count = 0 + for block in output: + try: + if not isinstance(block, dict) or "type" not in block: + logger.warning( + f"Invalid block: {block}, expected a dict with 'type' key, skipped.", + ) + continue + + block_type = block["type"] + + if block_type == "text": + textual_parts.append(block.get("text", "")) + total_token_count += await self._count_str_token(textual_parts[-1]) + + elif block_type in ["image", "audio", "video"]: + source = block.get("source", {}) + if source.get("type") == "base64": + data = source.get("data", "") + total_token_count += len(data) // 4 if data else 10 + else: + url = source.get("url", "") + total_token_count += await self._count_str_token(url) if url else 10 + textual_parts.append(f"[{block_type}] {url}") + + elif block_type == "file": + file_path = block.get("path", "") or block.get("url", "") + file_name = block.get("name", file_path) + textual_parts.append(f"[file] {file_name}: {file_path}") + total_token_count += await self._count_str_token(file_path) + + else: + logger.warning( + f"Unsupported block type '{block_type}' in tool result, skipped.", + ) + + except Exception as e: + logger.warning( + f"Failed to process block {block}: {e}, skipped.", + ) + + return "\n".join(textual_parts), total_token_count + + async def _stat_message(self, message: Msg) -> AsMsgStat: + """Analyze a message and generate block statistics.""" + blocks = [] + if isinstance(message.content, str): + blocks.append( + AsBlockStat( + block_type="text", + text=message.content, + token_count=await self._count_str_token(message.content), + ), + ) + return AsMsgStat( + name=message.name or message.role, + role=message.role, + content=blocks, + timestamp=message.timestamp or "", + metadata=message.metadata or {}, + ) + + for block in message.content: + block_type = block.get("type", "unknown") + + if block_type == "text": + text = block.get("text", "") + token_count = await self._count_str_token(text) + blocks.append( + AsBlockStat( + block_type=block_type, + text=text, + token_count=token_count, + ), + ) + + elif block_type == "thinking": + thinking = block.get("thinking", "") + token_count = await self._count_str_token(thinking) + blocks.append( + AsBlockStat( + block_type=block_type, + text=thinking, + token_count=token_count, + ), + ) + + elif block_type in ("image", "audio", "video"): + source = block.get("source", {}) + url = source.get("url", "") + if source.get("type") == "base64": + data = source.get("data", "") + token_count = len(data) // 4 if data else 10 + else: + token_count = await self._count_str_token(url) if url else 10 + blocks.append( + AsBlockStat( + block_type=block_type, + text="", + token_count=token_count, + media_url=url, + ), + ) + + elif block_type == "tool_use": + tool_name = block.get("name", "") + tool_input = block.get("input", "") + try: + input_str = json.dumps(tool_input, ensure_ascii=False) + except (TypeError, ValueError): + input_str = str(tool_input) + token_count = await self._count_str_token(tool_name + input_str) + blocks.append( + AsBlockStat( + block_type=block_type, + text="", + token_count=token_count, + tool_name=tool_name, + tool_input=input_str, + ), + ) + + elif block_type == "tool_result": + tool_name = block.get("name", "") + output = block.get("output", "") + formatted_output, token_count = await self._format_tool_result_output(output) + blocks.append( + AsBlockStat( + block_type=block_type, + text="", + token_count=token_count, + tool_name=tool_name, + tool_output=formatted_output, + ), + ) + + else: + logger.warning(f"Unsupported block type {block_type}, skipped.") + + return AsMsgStat( + name=message.name or message.role, + role=message.role, + content=blocks, + timestamp=message.timestamp or "", + metadata=message.metadata or {}, + ) + + async def _count_msgs_token(self, messages: list[Msg]) -> int: + """Count total token count of a list of messages.""" + total = 0 + for msg in messages: + stat = await self._stat_message(msg) + total += stat.total_tokens + return total + + async def _format_msgs_to_str( + self, + messages: list[Msg], + memory_compact_threshold: int, + include_thinking: bool = True, + ) -> str: + """Format list of messages to a single formatted string. + + Messages are processed in reverse order (newest first) and older + messages are skipped when token count exceeds memory_compact_threshold. + + Args: + messages: List of Msg objects to format. + memory_compact_threshold: Maximum token count before skipping older messages. + include_thinking: Whether to include thinking blocks in output. + """ + if not messages: + return "" + + formatted_parts: list[str] = [] + total_token_count = 0 + + for i in range(len(messages) - 1, -1, -1): + stat = await self._stat_message(messages[i]) + formatted_content = stat.format(include_thinking=include_thinking) + content_token_count = await self._count_str_token(formatted_content) + + is_latest = i == len(messages) - 1 + if not is_latest and total_token_count + content_token_count > memory_compact_threshold: + logger.info( + f"Skipping older messages: adding {content_token_count} tokens would exceed threshold " + f"{memory_compact_threshold} (current: {total_token_count})", + ) + break + + if is_latest and content_token_count > memory_compact_threshold: + logger.warning( + f"Latest message alone ({content_token_count} tokens) exceeds threshold " + f"{memory_compact_threshold}, including it anyway.", + ) + + formatted_parts.append(formatted_content) + total_token_count += content_token_count + + formatted_parts.reverse() + return "\n\n".join(formatted_parts) + + async def execute(self): + """Execute the summarization step.""" + messages: list[Msg] = self.context.data.get("messages", []) + + if not messages: + return "" + + before_token_count = await self._count_msgs_token(messages) + history_formatted_str: str = await self._format_msgs_to_str( + messages=messages, + memory_compact_threshold=self.memory_compact_threshold, + include_thinking=self.add_thinking_block, + ) + after_token_count = await self._count_str_token(history_formatted_str) + logger.info(f"Summarizer before_token_count={before_token_count} after_token_count={after_token_count}") + + if not history_formatted_str: + logger.warning(f"No history to summarize. messages={messages}") + return "" + + agent = ReActAgent( + name="reme_summarizer", + model=self.as_llm.model, + sys_prompt="You are a helpful assistant.", + formatter=self.as_llm_formatter.formatter, + toolkit=self.toolkit, + ) + agent.set_console_output_enabled(self.console_enabled) + + user_message: str = f"# conversation\n{history_formatted_str}\n\n" + self.prompt_format( + "user_message", + date=self._get_current_datetime().strftime("%Y-%m-%d"), + working_dir=self.working_dir, + memory_dir=self.memory_dir, + ) + + summary_msg: Msg = await agent.reply( + Msg( + name="reme", + role="user", + content=user_message, + ), + ) + for i, (msg, _) in enumerate(agent.memory.content): + logger.info(f"Summarizer memory[{i}]: {msg.content}") + + history_summary: str = summary_msg.get_text_content() + logger.info(f"Summarizer Result:\n{history_summary}") + return history_summary diff --git a/reme_cli/file_based/summarizer.yaml b/reme_cli/file_based/summarizer.yaml new file mode 100644 index 00000000..14599f71 --- /dev/null +++ b/reme_cli/file_based/summarizer.yaml @@ -0,0 +1,62 @@ +user_message: | + Memory Pre-compression Flush Cycle. + + The current session is about to enter the automatic compression phase. Please capture persistent memory AND session reflections, then write them to disk. + + Current date: {date} + Working directory: {working_dir} + + # Task + Immediately store persistent memory and reflections to: {memory_dir}/YYYY-MM-DD.md + + # Workflow + 1. Extract and synthesize content from the current session: + - Persistent Memory: Facts, user profile updates, project states, and important events. + - Experience Reflection: Reusable thinking logic derived from user feedback, successful problem-solving strategies, mistakes made/pitfalls to avoid, and actionable insights for future interactions. + 2. `read` {memory_dir}/YYYY-MM-DD.md (if the file doesn’t exist, an error message will be returned) + - If the file doesn’t exist, use `write` tool directly. + - If the file exists, intelligently merge new information with existing content, prefer using `edit` to update specific sections. + - Use `write` to overwrite the entire file only if substantial restructuring is required. + + # Principles + - Intelligently merge new information with existing content: + - Categorize clearly (e.g., separate "Factual Memory" from "Reflections & Logic"). + - Avoid duplicating already recorded information. + - Enrich existing entries with new details where relevant. + - Maintain chronological order wherever applicable. + - Always preserve timestamps and any date/time-related context. + - Add only genuinely new or meaningfully enriching information. + - Reflections MUST focus on forming reusable cognitive frameworks based on user feedback, aiming to improve future task execution. + - Keep entries concise yet complete. + - If there’s nothing to store or reflect on, respond with [SILENT]. + +user_message_zh: | + 预压缩内存刷新轮次。 + + 当前会话即将进入自动压缩阶段；请将持久化记忆与经验反思捕获并写入磁盘。 + + 当前日期：{date} + 工作目录：{working_dir} + + # 任务 + 立即存储持久化记忆与反思（使用路径 {memory_dir}/YYYY-MM-DD.md）。 + + # 工作流程 + 1. 从当前会话中提取并综合两类内容： + - 持久化记忆：客观事实、用户信息更新、项目状态及重要事件。 + - 经验反思：基于用户反馈形成的可复用思考逻辑、成功的问题解决策略、犯下的错误/应避免的陷阱，以及对未来交互有帮助的行动指南。 + 2. `read` {memory_dir}/YYYY-MM-DD.md（如文件不存在，会返回错误提示） + - 若文件不存在，直接使用 `write` 工具写入。 + - 若文件已存在，智能合并新信息与现有内容，尽可能使用 `edit` 更新特定部分，仅在需要大幅重构时使用 `write` 覆盖整个文件。 + + # 原则 + - 智能合并新信息与现有内容： + - 将内容进行清晰的分类（例如明确区分“事实记忆”与“反思与逻辑”）。 + - 避免重复已记录的信息。 + - 在相关时丰富现有条目的新细节。 + - 在适用时保持时间顺序。 + - 始终保留时间戳、日期和时间相关上下文。 + - 仅添加真正新的或有丰富价值的信息。 + - 反思内容必须侧重于根据用户反馈构建可复用的思维逻辑，以改善未来的任务执行。 + - 保持条目简洁但完整。 + - 若无任何新内容可存储或反思，请回复 [SILENT]。 \ No newline at end of file diff --git a/reme_cli/jobs/__init__.py b/reme_cli/jobs/__init__.py deleted file mode 100644 index c4d57960..00000000 --- a/reme_cli/jobs/__init__.py +++ /dev/null @@ -1 +0,0 @@ -"""Jobs""" diff --git a/reme_cli/reme.py b/reme_cli/reme.py index 03fd34fd..9a8b1235 100644 --- a/reme_cli/reme.py +++ b/reme_cli/reme.py @@ -2,59 +2,116 @@ import asyncio import sys +from pathlib import Path from agentscope.formatter import FormatterBase from agentscope.message import Msg from agentscope.model import ChatModelBase -from agentscope.token import TokenCounterBase -from agentscope.tool import Toolkit, ToolResponse +from agentscope.token import HuggingFaceTokenCounter, TokenCounterBase +from agentscope.tool import Toolkit from .application import Application from .component import R +from .component.runtime_context import RuntimeContext from .config import parse_args from .enumeration import ComponentEnum +from .file_based.summarizer import Summarizer class ReMe(Application): """ReMe memory management application.""" async def summary_memory( - self, - messages: list[Msg], - as_llm: str | ChatModelBase = "default", - as_llm_formatter: str | FormatterBase = "default", - as_token_counter: str | TokenCounterBase = "default", - toolkit: Toolkit | None = None, - language: str = "zh", - max_input_length: float = 128 * 1024, - compact_ratio: float = 0.7, - timezone: str | None = None, - add_thinking_block: bool = True, + self, + messages: list[Msg], + as_llm: str | ChatModelBase = "default", + as_llm_formatter: str | FormatterBase = "default", + as_token_counter: str | TokenCounterBase | HuggingFaceTokenCounter = "default", + toolkit: Toolkit | None = None, + language: str = "zh", + max_input_length: float = 128 * 1024, + compact_ratio: float = 0.7, + timezone: str | None = None, + add_thinking_block: bool = True, ) -> str: - """Summarize and compact memory messages.""" + """Summarize and compact memory messages. - async def memory_search(self, query: str, max_results: int = 5, min_score: float = 0.1) -> ToolResponse: + Args: + messages: List of AgentScope messages to summarize. + as_llm: LLM model name or instance. + as_llm_formatter: Formatter name or instance. + as_token_counter: Token counter name or instance. + toolkit: Optional toolkit for the summarizer agent. + language: Language for prompts (zh or en). + max_input_length: Maximum input token length. + compact_ratio: Ratio of max_input_length to use as compact threshold. + timezone: Optional timezone for date formatting. + add_thinking_block: Whether to include thinking blocks. + + Returns: + Summarized memory string. + """ + working_dir = Path(self.config.working_dir).absolute() + memory_dir = working_dir / "memory" + memory_compact_threshold = int(max_input_length * compact_ratio) + + # Resolve token counter - use provided instance or create default + token_counter_instance = None + if isinstance(as_token_counter, HuggingFaceTokenCounter): + token_counter_instance = as_token_counter + else: + token_counter_instance = HuggingFaceTokenCounter() + + summarizer = Summarizer( + working_dir=str(working_dir), + memory_dir=str(memory_dir), + memory_compact_threshold=memory_compact_threshold, + toolkit=toolkit, + timezone=timezone, + add_thinking_block=add_thinking_block, + as_token_counter=token_counter_instance, + language=language, + as_llm=as_llm if isinstance(as_llm, str) else "default", + as_llm_formatter=as_llm_formatter if isinstance(as_llm_formatter, str) else "default", + ) + + context = RuntimeContext( + messages=messages, + application_context=self.context, + ) + + result = await summarizer(context=context) + return result or "" + + async def memory_search(self, query: str, max_results: int = 5, min_score: float = 0.1) -> str: """Search memory for relevant entries.""" + from .file_based.memory_search import MemorySearch + try: + search_step = MemorySearch() + self.logger.info(f"Running memory search with {query} {max_results} {min_score}") + return await search_step(query=query, max_results=max_results, min_score=min_score) + except Exception as e: + return str(e) async def dream( - self, - as_llm: str | ChatModelBase = "default", - as_llm_formatter: str | FormatterBase = "default", - as_token_counter: str | TokenCounterBase = "default", - toolkit: Toolkit | None = None, - language: str = "zh", - timezone: str | None = None, + self, + as_llm: str | ChatModelBase = "default", + as_llm_formatter: str | FormatterBase = "default", + as_token_counter: str | TokenCounterBase = "default", + toolkit: Toolkit | None = None, + language: str = "zh", + timezone: str | None = None, ) -> str: """Process and consolidate memories in background.""" async def proactive( - self, - as_llm: str | ChatModelBase = "default", - as_llm_formatter: str | FormatterBase = "default", - as_token_counter: str | TokenCounterBase = "default", - toolkit: Toolkit | None = None, - language: str = "zh", - timezone: str | None = None, + self, + as_llm: str | ChatModelBase = "default", + as_llm_formatter: str | FormatterBase = "default", + as_token_counter: str | TokenCounterBase = "default", + toolkit: Toolkit | None = None, + language: str = "zh", + timezone: str | None = None, ) -> str: """Generate proactive memory insights.""" @@ -62,7 +119,7 @@ class ReMe(Application): def main(): """Entry point for ReMe CLI.""" action, config = parse_args(sys.argv[1:]) - if action == "app": + if action == "start": reme = ReMe(**config) reme.run_app() diff --git a/reme_cli/schema/__init__.py b/reme_cli/schema/__init__.py index 89b7fe3d..0e6f16ef 100644 --- a/reme_cli/schema/__init__.py +++ b/reme_cli/schema/__init__.py @@ -1,6 +1,7 @@ """Schema""" from .application_config import ApplicationConfig, ComponentConfig, JobConfig +from .as_msg_stat import AsBlockStat, AsMsgStat from .base_node import BaseNode from .file_chunk import FileChunk from .file_metadata import FileMetadata @@ -12,6 +13,8 @@ __all__ = [ "ApplicationConfig", "ComponentConfig", "JobConfig", + "AsBlockStat", + "AsMsgStat", "BaseNode", "FileChunk", "FileMetadata", diff --git a/reme_cli/schema/as_msg_stat.py b/reme_cli/schema/as_msg_stat.py new file mode 100644 index 00000000..fa02496d --- /dev/null +++ b/reme_cli/schema/as_msg_stat.py @@ -0,0 +1,97 @@ +"""Schema definitions for AgentScope message statistics.""" + +from pydantic import BaseModel, Field + +_TRUNCATION_NOTICE_MARKER = "<<>>" + +_DEFAULT_MAX_BLOCK_TEXT_PREVIEW_LENGTH = 100 +_DEFAULT_MAX_FORMATTER_TEXT_LENGTH = 1000 + + +class AsBlockStat(BaseModel): + """Statistics and metadata for a single content block in an AgentScope message.""" + + block_type: str = Field(default=...) + text: str = Field(default="", description="Text content of the block") + token_count: int = Field(default=0, description="Token count of the block, including base64 data") + + # For tool_use and tool_result blocks + tool_name: str = Field(default="", description="Tool name for tool_use/tool_result blocks") + tool_input: str = Field(default="", description="Tool input arguments for tool_use blocks") + tool_output: str = Field(default="", description="Tool output for tool_result blocks") + + # For media blocks + media_url: str = Field(default="", description="URL for image/audio/video blocks") + + @property + def preview(self) -> str: + """Return a short preview of the block content.""" + return self.format(_DEFAULT_MAX_BLOCK_TEXT_PREVIEW_LENGTH) + + @staticmethod + def _truncate(text: str, max_length: int) -> str: + """Truncate text with ellipsis, replacing newlines with spaces.""" + text = text.replace("\n", " ") + if len(text) <= max_length: + return text + return text[:max_length] + "..." + + # pylint: disable=too-many-return-statements + def format(self, max_length: int = _DEFAULT_MAX_FORMATTER_TEXT_LENGTH, include_thinking: bool = True) -> str: + """Format block content to string representation. + + Args: + max_length: Maximum length of text content in the output. + include_thinking: Whether to include thinking block content. + + Returns: + Formatted string representation of the block. + """ + if self.block_type == "text": + if not self.text: + return "" + return f"[text]: {self._truncate(self.text, max_length)}" + if self.block_type == "thinking": + if not include_thinking or not self.text: + return "" + return f"[think]: {self._truncate(self.text, max_length)}" + if self.block_type in ("image", "audio", "video"): + content = self.media_url if self.media_url else "" + return f"[{self.block_type}]: {content}" + if self.block_type == "tool_use": + content = f"{self.tool_name} params={self._truncate(self.tool_input, max_length)}" + return f"[tool_use]: {content}" + if self.block_type == "tool_result": + if not self.tool_output: + return "" + display_output = self.tool_output.split(_TRUNCATION_NOTICE_MARKER)[0] + content = f"{self.tool_name} output={self._truncate(display_output, max_length)}" + return f"[tool_result]: {content}" + return "" + + +class AsMsgStat(BaseModel): + """Statistics and metadata for a complete AgentScope message.""" + + name: str = Field(default=...) + role: str = Field(default="") + content: list[AsBlockStat] = Field(default_factory=list) + timestamp: str = Field(default="") + metadata: dict = Field(default_factory=dict) + + @property + def total_tokens(self) -> int: + """Return the total token count across all content blocks.""" + return sum(block.token_count for block in self.content) + + @property + def preview(self) -> str: + """Return a short preview of the message content.""" + return self.format(_DEFAULT_MAX_BLOCK_TEXT_PREVIEW_LENGTH) + + def format(self, max_length: int = _DEFAULT_MAX_FORMATTER_TEXT_LENGTH, include_thinking: bool = True) -> str: + """Format message to string representation.""" + time_str = f"[{self.timestamp}] " if self.timestamp else "" + header = f"{time_str}{self.name or self.role}:" + blocks = [block.format(max_length, include_thinking) for block in self.content] + return "\n".join([header] + [b for b in blocks if b]) diff --git a/reme_cli/vector_based/__init__.py b/reme_cli/vector_based/__init__.py new file mode 100644 index 00000000..e69de29b