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..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.6b3"
+__version__ = "0.3.0.8"
__all__ = [
"config",
diff --git a/reme/config/light.yaml b/reme/config/light.yaml
index bc85c10d..e18a1302 100644
--- a/reme/config/light.yaml
+++ b/reme/config/light.yaml
@@ -7,6 +7,11 @@ as_llm_formatters:
default:
backend: openai
+as_token_counters:
+ default:
+ backend: rule
+ token_count_estimate_divisor: 3.75
+
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..95910b8e 100644
--- a/reme/core/as_token_counter/__init__.py
+++ b/reme/core/as_token_counter/__init__.py
@@ -1,9 +1,8 @@
"""Module for registering AgentScope token counters."""
-from agentscope.token import OpenAITokenCounter
-from agentscope.token import HuggingFaceTokenCounter
-
+from .reme_token_counter import ReMeTokenCounter
+from .rule_token_counter import RuleTokenCounter
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)
+R.as_token_counters.register("rule")(RuleTokenCounter)
diff --git a/reme/core/as_token_counter/reme_token_counter.py b/reme/core/as_token_counter/reme_token_counter.py
new file mode 100644
index 00000000..39ad8f6d
--- /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.error(f"Failed to initialize tokenizer {e}")
+ self._tokenizer_available = False
+
+ async def count(
+ self,
+ messages: list[dict],
+ tools: list[dict] | None = None,
+ text: str | None = None,
+ **kwargs: Any,
+ ) -> int:
+ """Count tokens in messages or text.
+
+ If text is provided, counts tokens directly in the text string.
+ Otherwise, counts tokens in the messages using the parent class method.
+
+ Args:
+ messages: List of message dictionaries in chat format.
+ tools: Optional list of tool definitions for token counting.
+ text: Optional text string to count tokens directly.
+ **kwargs: Additional keyword arguments passed to parent count method.
+
+ Returns:
+ The number of tokens, guaranteed to be at least the estimated minimum.
+ """
+ if text:
+ if self._tokenizer_available:
+ try:
+ token_ids = self.tokenizer.encode(text)
+ return max(len(token_ids), self.estimate_tokens(text))
+ except Exception as e:
+ logger.exception("Failed to encode text with tokenizer: %s", e)
+ return self.estimate_tokens(text)
+ else:
+ return self.estimate_tokens(text)
+ else:
+ return await super().count(messages, tools, **kwargs)
+
+ def estimate_tokens(self, text: str) -> int:
+ """Estimate the number of tokens in a text string.
+
+ Provides a fast character-based estimation as a fallback or lower bound.
+ Uses the configured divisor from instance settings.
+
+ Args:
+ text: The text string to estimate tokens for.
+
+ Returns:
+ The estimated number of tokens in the text string.
+ """
+ return int(len(text.encode("utf-8")) / self.token_count_estimate_divisor + 0.5)
diff --git a/reme/core/as_token_counter/rule_token_counter.py b/reme/core/as_token_counter/rule_token_counter.py
new file mode 100644
index 00000000..1201e185
--- /dev/null
+++ b/reme/core/as_token_counter/rule_token_counter.py
@@ -0,0 +1,78 @@
+"""Rule-based token counter for fast estimation without loading tokenizer."""
+
+from typing import Any
+
+from agentscope.token import HuggingFaceTokenCounter
+
+
+class RuleTokenCounter(HuggingFaceTokenCounter):
+ """Lightweight token counter using rule-based estimation only.
+
+ This class provides fast token estimation without loading any tokenizer,
+ useful when exact token counts are not critical or for quick approximations.
+
+ Attributes:
+ token_count_estimate_divisor: Divisor for token estimation.
+ """
+
+ def __init__(
+ self,
+ token_count_estimate_divisor: float = 3.75,
+ **_kwargs,
+ ):
+ """Initialize the rule-based token counter.
+
+ Args:
+ token_count_estimate_divisor: Divisor for estimating tokens.
+ Defaults to 3.75 (approximately 4 characters per token).
+ **kwargs: Additional keyword arguments (ignored).
+ """
+ self.token_count_estimate_divisor = token_count_estimate_divisor
+ # Skip tokenizer initialization from parent
+ self._tokenizer_available = False
+
+ async def count(
+ self,
+ messages: list[dict],
+ _tools: list[dict] | None = None,
+ text: str | None = None,
+ **_kwargs: Any,
+ ) -> int:
+ """Count tokens using rule-based estimation.
+
+ Args:
+ messages: List of message dictionaries in chat format.
+ _tools: Optional list of tool definitions (ignored).
+ text: Optional text string to count tokens directly.
+ **_kwargs: Additional keyword arguments (ignored).
+
+ Returns:
+ The estimated number of tokens.
+ """
+ if text:
+ return self.estimate_tokens(text)
+
+ # Estimate from messages
+ total_text = ""
+ for msg in messages:
+ content = msg.get("content", "")
+ if isinstance(content, str):
+ total_text += content
+ elif isinstance(content, list):
+ for part in content:
+ if isinstance(part, dict) and "text" in part:
+ total_text += part["text"]
+ return self.estimate_tokens(total_text)
+
+ def estimate_tokens(self, text: str) -> int:
+ """Estimate the number of tokens in a text string.
+
+ Uses character-based estimation with the configured divisor.
+
+ Args:
+ text: The text string to estimate tokens for.
+
+ Returns:
+ The estimated number of tokens in the text string.
+ """
+ return int(len(text.encode("utf-8")) / self.token_count_estimate_divisor + 0.5)
diff --git a/reme/core/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]:
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..4fcec627 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,
@@ -56,7 +145,8 @@ class ReMeInMemoryMemory(InMemoryMemory):
{self._compressed_summary}
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 [
@@ -105,19 +195,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 +264,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..0692ebce 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
@@ -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
@@ -59,7 +58,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,31 +67,28 @@ 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(
- "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
- 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 +96,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 +107,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 +123,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 +139,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 +156,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 +170,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,
@@ -199,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,
@@ -209,11 +192,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,25 +223,22 @@ 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:
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)
@@ -286,7 +270,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 +299,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
@@ -354,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
@@ -383,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
@@ -410,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
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 "