mirror of
https://github.com/agentscope-ai/ReMe.git
synced 2026-09-20 00:12:50 +00:00
351 lines
13 KiB
Python
351 lines
13 KiB
Python
import json
|
|
|
|
from agentscope.message import Msg
|
|
from agentscope.token import HuggingFaceTokenCounter
|
|
|
|
from ...core.schema import AsMsgStat, AsBlockStat
|
|
from ...core.utils import get_std_logger
|
|
|
|
logger = get_std_logger()
|
|
|
|
|
|
class AsMsgHandler:
|
|
|
|
def __init__(self, token_counter: HuggingFaceTokenCounter):
|
|
self._token_counter = token_counter
|
|
|
|
def count_str_token(self, text: str) -> int:
|
|
"""Count tokens in a string.
|
|
|
|
Args:
|
|
text: The text to count tokens for.
|
|
|
|
Returns:
|
|
The number of tokens in the text.
|
|
"""
|
|
if not text:
|
|
return 0
|
|
|
|
try:
|
|
token_ids = self._token_counter.tokenizer.encode(text)
|
|
token_count = len(token_ids)
|
|
return token_count
|
|
|
|
except Exception as e:
|
|
estimated_tokens = len(text.encode("utf-8")) // 4
|
|
logger.warning(f"Failed to count string tokens: {text}, using estimated_tokens={estimated_tokens}")
|
|
return estimated_tokens
|
|
|
|
@staticmethod
|
|
def _format_tool_result_output(output: str | list[dict]) -> str:
|
|
"""Convert tool result output to string.
|
|
|
|
Args:
|
|
output: Tool result output, either string or list of content blocks.
|
|
|
|
Returns:
|
|
Formatted string representation of the tool result.
|
|
"""
|
|
if isinstance(output, str):
|
|
return output
|
|
|
|
textual_parts = []
|
|
|
|
for block in output:
|
|
try:
|
|
if not isinstance(block, dict) or "type" not in block:
|
|
logger.warning(
|
|
"Invalid block: %s, expected a dict with 'type' key, skipped.",
|
|
block,
|
|
)
|
|
continue
|
|
|
|
block_type = block["type"]
|
|
|
|
if block_type == "text":
|
|
textual_parts.append(block.get("text", ""))
|
|
|
|
elif block_type in ["image", "audio", "video"]:
|
|
source = block.get("source", {})
|
|
url = source.get("url", "")
|
|
if url:
|
|
textual_parts.append(f"[{block_type}] {url}")
|
|
else:
|
|
textual_parts.append(f"[{block_type}]")
|
|
|
|
elif block_type == "file":
|
|
file_path = block.get("path", "") or block.get("url", "")
|
|
file_name = block.get("name", file_path)
|
|
textual_parts.append(f"[file] {file_name}: {file_path}")
|
|
|
|
else:
|
|
logger.warning(
|
|
"Unsupported block type '%s' in tool result, skipped.",
|
|
block_type,
|
|
)
|
|
|
|
except Exception as e:
|
|
logger.warning(
|
|
"Failed to process block %s: %s, skipped.",
|
|
block,
|
|
e,
|
|
)
|
|
|
|
if not textual_parts:
|
|
return ""
|
|
if len(textual_parts) == 1:
|
|
return textual_parts[0]
|
|
return "\n".join(f"- {part}" for part in textual_parts)
|
|
|
|
def stat_message(self, message: Msg) -> AsMsgStat:
|
|
"""Analyze a message and generate block statistics."""
|
|
blocks = []
|
|
|
|
for block in message.get_content_blocks():
|
|
block_type = block.get("type", "unknown")
|
|
|
|
if block_type == "text":
|
|
text = block.get("text", "")
|
|
token_count = self.count_str_token(text)
|
|
blocks.append(AsBlockStat(
|
|
block_type=block_type,
|
|
text=text,
|
|
token_count=token_count,
|
|
))
|
|
|
|
elif block_type == "thinking":
|
|
thinking = block.get("thinking", "")
|
|
token_count = self.count_str_token(thinking)
|
|
blocks.append(AsBlockStat(
|
|
block_type=block_type,
|
|
text=thinking,
|
|
token_count=token_count,
|
|
))
|
|
|
|
elif block_type in ("image", "audio", "video"):
|
|
source = block.get("source", {})
|
|
url = source.get("url", "")
|
|
# For media, estimate fixed token cost or count URL
|
|
if source.get("type") == "base64":
|
|
data = source.get("data", "")
|
|
token_count = len(data) // 4 if data else 10
|
|
else:
|
|
token_count = self.count_str_token(url) if url else 10
|
|
blocks.append(AsBlockStat(
|
|
block_type=block_type,
|
|
text="",
|
|
token_count=token_count,
|
|
media_url=url,
|
|
))
|
|
|
|
elif block_type == "tool_use":
|
|
tool_name = block.get("name", "")
|
|
tool_input = block.get("input", {})
|
|
try:
|
|
input_str = json.dumps(tool_input, ensure_ascii=False)
|
|
except (TypeError, ValueError):
|
|
input_str = str(tool_input)
|
|
token_count = self.count_str_token(tool_name + input_str)
|
|
blocks.append(AsBlockStat(
|
|
block_type=block_type,
|
|
text="",
|
|
token_count=token_count,
|
|
tool_name=tool_name,
|
|
tool_input=input_str,
|
|
))
|
|
|
|
elif block_type == "tool_result":
|
|
tool_name = block.get("name", "")
|
|
output = block.get("output", "")
|
|
formatted_output = self._format_tool_result_output(output)
|
|
token_count = self.count_str_token(formatted_output)
|
|
blocks.append(AsBlockStat(
|
|
block_type=block_type,
|
|
text="",
|
|
token_count=token_count,
|
|
tool_name=tool_name,
|
|
tool_output=formatted_output,
|
|
))
|
|
|
|
else:
|
|
logger.warning("Unsupported block type %s, skipped.", block_type)
|
|
|
|
return AsMsgStat(
|
|
name=message.name or message.role,
|
|
role=message.role,
|
|
content=blocks,
|
|
timestamp=message.timestamp or "",
|
|
metadata=message.metadata or {},
|
|
)
|
|
|
|
def format_msgs_to_str(
|
|
self,
|
|
messages: list[Msg],
|
|
memory_compact_threshold: int,
|
|
include_thinking: bool = False,
|
|
) -> str:
|
|
"""Format list of messages to a single formatted string.
|
|
|
|
Messages are processed in reverse order (newest first) and older
|
|
messages are skipped when token count exceeds memory_compact_threshold.
|
|
|
|
Args:
|
|
messages: List of Msg objects to format.
|
|
memory_compact_threshold: Maximum token count before skipping older messages.
|
|
include_thinking: Whether to include thinking blocks in output.
|
|
"""
|
|
if not messages:
|
|
return ""
|
|
|
|
formatted_parts: list[str] = []
|
|
total_token_count = 0
|
|
|
|
for i in range(len(messages) - 1, -1, -1):
|
|
stat = self.stat_message(messages[i])
|
|
|
|
if total_token_count + stat.total_tokens > memory_compact_threshold:
|
|
logger.info(
|
|
"Skipping older messages: adding %d tokens would exceed threshold %d (current: %d)",
|
|
stat.total_tokens,
|
|
memory_compact_threshold,
|
|
total_token_count,
|
|
)
|
|
break
|
|
|
|
formatted_parts.append(stat.format(include_thinking=include_thinking))
|
|
total_token_count += stat.total_tokens
|
|
|
|
formatted_parts.reverse()
|
|
return "\n\n".join(formatted_parts)
|
|
|
|
def context_check(
|
|
self,
|
|
messages: list[Msg],
|
|
memory_compact_threshold: int,
|
|
memory_compact_reserve: int,
|
|
) -> tuple[list[Msg], list[Msg]]:
|
|
"""Check if context exceeds threshold and split messages accordingly.
|
|
|
|
This method checks if the total token count of messages exceeds the
|
|
memory_compact_threshold. If not, returns empty list and original messages.
|
|
If exceeded, uses memory_compact_reserve as the limit to keep messages
|
|
from the end, ensuring tool_use and tool_result blocks are properly paired.
|
|
|
|
Args:
|
|
messages: List of Msg objects to check.
|
|
memory_compact_threshold: Maximum token count threshold to trigger compaction.
|
|
memory_compact_reserve: Token limit for messages to keep after compaction.
|
|
|
|
Returns:
|
|
A tuple of (messages_to_compact, messages_to_keep):
|
|
- messages_to_compact: Older messages that need to be compacted
|
|
- messages_to_keep: Recent messages within the reserve limit
|
|
"""
|
|
if not messages:
|
|
return [], []
|
|
|
|
# Calculate total tokens and stats for all messages
|
|
msg_stats: list[tuple[Msg, AsMsgStat]] = []
|
|
total_tokens = 0
|
|
for msg in messages:
|
|
stat = self.stat_message(msg)
|
|
msg_stats.append((msg, stat))
|
|
total_tokens += stat.total_tokens
|
|
|
|
# If total tokens don't exceed threshold, no compaction needed
|
|
if total_tokens <= memory_compact_threshold:
|
|
return [], messages
|
|
|
|
# Collect all tool_use ids and their message indices
|
|
# tool_use_id -> message index
|
|
tool_use_locations: dict[str, int] = {}
|
|
# tool_result_id -> message index
|
|
tool_result_locations: dict[str, int] = {}
|
|
|
|
for idx, (msg, _) in enumerate(msg_stats):
|
|
for block in msg.get_content_blocks("tool_use"):
|
|
tool_id = block.get("id", "")
|
|
if tool_id:
|
|
tool_use_locations[tool_id] = idx
|
|
|
|
for block in msg.get_content_blocks("tool_result"):
|
|
tool_id = block.get("id", "")
|
|
if tool_id:
|
|
tool_result_locations[tool_id] = idx
|
|
|
|
# Iterate from the end, accumulating messages to keep within reserve limit
|
|
keep_indices: set[int] = set()
|
|
accumulated_tokens = 0
|
|
|
|
for i in range(len(msg_stats) - 1, -1, -1):
|
|
msg, stat = msg_stats[i]
|
|
|
|
# Check if adding this message would exceed reserve limit
|
|
if accumulated_tokens + stat.total_tokens > memory_compact_reserve:
|
|
logger.info(
|
|
"Context check: adding message %d with %d tokens would exceed reserve %d (current: %d)",
|
|
i,
|
|
stat.total_tokens,
|
|
memory_compact_reserve,
|
|
accumulated_tokens,
|
|
)
|
|
break
|
|
|
|
# Check tool_result dependencies - if this message has tool_result,
|
|
# we need to ensure the corresponding tool_use is also included
|
|
tool_result_ids = [
|
|
block.get("id", "")
|
|
for block in msg.get_content_blocks("tool_result")
|
|
if block.get("id", "")
|
|
]
|
|
|
|
# Calculate extra tokens needed for dependent tool_use messages
|
|
extra_tokens = 0
|
|
dependent_indices: set[int] = set()
|
|
|
|
for tool_id in tool_result_ids:
|
|
if tool_id in tool_use_locations:
|
|
tool_use_idx = tool_use_locations[tool_id]
|
|
if tool_use_idx not in keep_indices and tool_use_idx != i:
|
|
dependent_indices.add(tool_use_idx)
|
|
_, dep_stat = msg_stats[tool_use_idx]
|
|
extra_tokens += dep_stat.total_tokens
|
|
|
|
# Check if we can fit this message plus its dependencies within reserve
|
|
if accumulated_tokens + stat.total_tokens + extra_tokens > memory_compact_reserve:
|
|
logger.info(
|
|
"Context check: message %d requires %d extra tokens for tool_use dependencies, "
|
|
"total would exceed reserve %d",
|
|
i,
|
|
extra_tokens,
|
|
memory_compact_reserve,
|
|
)
|
|
break
|
|
|
|
# Add this message and its dependencies
|
|
keep_indices.add(i)
|
|
keep_indices.update(dependent_indices)
|
|
accumulated_tokens += stat.total_tokens + extra_tokens
|
|
|
|
# Build final lists based on keep_indices (preserve original order)
|
|
messages_to_compact = []
|
|
messages_to_keep = []
|
|
|
|
for idx, (msg, _) in enumerate(msg_stats):
|
|
if idx in keep_indices:
|
|
messages_to_keep.append(msg)
|
|
else:
|
|
messages_to_compact.append(msg)
|
|
|
|
logger.info(
|
|
"Context check result: %d messages to compact, %d messages to keep, "
|
|
"total tokens: %d, threshold: %d, reserve: %d, kept tokens: %d",
|
|
len(messages_to_compact),
|
|
len(messages_to_keep),
|
|
total_tokens,
|
|
memory_compact_threshold,
|
|
memory_compact_reserve,
|
|
accumulated_tokens,
|
|
)
|
|
|
|
return messages_to_compact, messages_to_keep
|