338 lines
14 KiB
Python
338 lines
14 KiB
Python
# Copyright 2026 Google LLC
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
|
|
"""Universal Context Compressor for Artemis Flash profile.
|
|
|
|
:class:`ScrubEdgeCompressor` implements the scrub-edge (擦洗沿) discipline
|
|
from the history-module redesign §3.2: each message is mutated a bounded
|
|
number of times near the tail (XML stripped at depth 1, screenshot resolved
|
|
at depth K), then frozen forever. A late summary never backfills a frozen
|
|
message. FlashRunner uses this implementation; the Pro transcript ledger
|
|
reuses it over its active region (see the class docstring).
|
|
|
|
The legacy per-turn full-rescan ``compress_flash_messages`` was removed in M5
|
|
(2026-09-01) after serving as the M1/M2 semantic reference — its surviving
|
|
semantics are asserted directly against ``ScrubEdgeCompressor`` in
|
|
``tests/unit/agents/test_flash_scrub_edge.py``.
|
|
"""
|
|
|
|
from typing import Any, Callable
|
|
|
|
from langchain_core.messages import BaseMessage
|
|
|
|
from artemis.memory.step_memory import StepMemoryService
|
|
from artemis.utils.logger import get_logger
|
|
|
|
logger = get_logger(__name__)
|
|
|
|
_HISTORY_SUMMARY_PREFIX = "--- Historical Visual Transition ---\n"
|
|
_UI_LIST_MARKER = "--- UI Element List ---"
|
|
|
|
|
|
def _without_marked_suffix(text: str, markers: tuple[str, ...] = (_UI_LIST_MARKER,)) -> str:
|
|
"""Remove a marked heavy block without discarding text before its marker."""
|
|
marker_index = -1
|
|
for marker in markers:
|
|
idx = text.find(marker)
|
|
if idx >= 0 and (marker_index < 0 or idx < marker_index):
|
|
marker_index = idx
|
|
if marker_index < 0:
|
|
return text
|
|
return text[:marker_index].rstrip()
|
|
|
|
|
|
class ScrubEdgeCompressor:
|
|
"""Scrub-edge context compressor (history-module redesign §3.2, M1).
|
|
|
|
Mutation discipline — every message is touched at most twice, always near
|
|
the tail, and is then frozen forever:
|
|
|
|
1. Depth 1 (XML): once a newer observation exists, historical UI Element
|
|
lists are stripped (same semantics as the legacy ``prune_history_xml``).
|
|
2. Depth K (image, ``image_scrub_depth``): the K-th most recent
|
|
image-bearing message has its screenshot resolved — replaced by the
|
|
ready visual summary; if the summary is still pending the image is
|
|
retained for up to ``pending_grace_steps`` further image-depths, after
|
|
which it is replaced by a pending placeholder. A failed summary is
|
|
replaced by an unavailable placeholder immediately. Both placeholders
|
|
carry the DataEngine step reference.
|
|
3. Frozen: after resolution the message is never read or written again —
|
|
a late summary never backfills a frozen message.
|
|
|
|
Contract: the message list is append-only (FlashRunner never removes or
|
|
reorders entries), so bookkeeping is index-based and each pass only
|
|
touches the few messages near the scrub edge instead of rescanning the
|
|
whole history.
|
|
|
|
Profile generalization (M2): the Pro transcript ledger reuses this exact
|
|
discipline over its active region. Its observation messages carry no
|
|
``tool_call_id``, so ``summary_key_getter`` supplies the summary-job key
|
|
(the DataEngine step id) per message; ``strip_markers`` extends the
|
|
depth-1 strip to the Pro tail blocks (UI element list + plan recitation);
|
|
and ``tail_offset=1`` accounts for the live observation living outside
|
|
the tracked list (the ledger renders it as a separate tail), so depth
|
|
arithmetic and the keep-window stay aligned with the Flash semantics
|
|
where the live observation is *inside* the list. Defaults leave the
|
|
Flash behavior byte-identical.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
summarizer: StepMemoryService | None = None,
|
|
*,
|
|
prune_history_xml: bool = True,
|
|
image_scrub_depth: int = 3,
|
|
pending_grace_steps: int = 3,
|
|
xml_scrub_depth: int = 1,
|
|
summary_key_getter: Callable[[BaseMessage], str | None] | None = None,
|
|
strip_markers: tuple[str, ...] = (_UI_LIST_MARKER,),
|
|
tail_offset: int = 0,
|
|
):
|
|
self._summarizer = summarizer
|
|
self._prune_history_xml = prune_history_xml
|
|
self._image_scrub_depth = max(1, image_scrub_depth)
|
|
self._pending_grace_steps = max(0, pending_grace_steps)
|
|
self._xml_scrub_depth = max(1, xml_scrub_depth)
|
|
self._summary_key_getter = summary_key_getter
|
|
self._strip_markers = tuple(strip_markers) or (_UI_LIST_MARKER,)
|
|
self._tail_offset = max(0, tail_offset)
|
|
|
|
self._scanned_until = 0
|
|
self._tool_msg_count = 0
|
|
# Append-only records of image-bearing messages, in message order:
|
|
# {"idx", "key" (str tool_call_id | None), "legacy" (ordinal | None), "is_tool"}
|
|
self._tracked: list[dict[str, Any]] = []
|
|
self._frozen: set[int] = set()
|
|
self._xml_candidates: list[int] = []
|
|
|
|
def compress(self, messages: list[BaseMessage]) -> None:
|
|
"""Advance the scrub edge over newly appended messages in-place."""
|
|
self._discover(messages)
|
|
if not self._tracked:
|
|
return
|
|
self._scrub_xml(messages)
|
|
self._scrub_images(messages)
|
|
|
|
# ------------------------------------------------------------------
|
|
# Discovery (incremental; frozen prefix is never rescanned)
|
|
# ------------------------------------------------------------------
|
|
|
|
def _discover(self, messages: list[BaseMessage]) -> None:
|
|
for idx in range(self._scanned_until, len(messages)):
|
|
msg = messages[idx]
|
|
tool_call_id = getattr(msg, "tool_call_id", None)
|
|
if tool_call_id is not None:
|
|
self._tool_msg_count += 1
|
|
content = getattr(msg, "content", None)
|
|
if not isinstance(content, list):
|
|
continue
|
|
|
|
has_image = False
|
|
has_xml = False
|
|
for block in content:
|
|
if not isinstance(block, dict):
|
|
continue
|
|
if block.get("type") in ("image_url", "image"):
|
|
has_image = True
|
|
elif block.get("type") == "text" and self._has_strip_marker(
|
|
block.get("text", "")
|
|
):
|
|
has_xml = True
|
|
|
|
if has_xml:
|
|
self._xml_candidates.append(idx)
|
|
if has_image:
|
|
key = str(tool_call_id) if tool_call_id is not None else None
|
|
if key is None and self._summary_key_getter is not None:
|
|
try:
|
|
getter_key = self._summary_key_getter(msg)
|
|
except Exception:
|
|
getter_key = None
|
|
if getter_key is not None:
|
|
key = str(getter_key)
|
|
self._tracked.append(
|
|
{
|
|
"idx": idx,
|
|
"key": key,
|
|
"legacy": self._tool_msg_count if tool_call_id is not None else None,
|
|
"is_tool": tool_call_id is not None,
|
|
}
|
|
)
|
|
self._scanned_until = len(messages)
|
|
|
|
def _has_strip_marker(self, text: str) -> bool:
|
|
return any(marker in text for marker in self._strip_markers)
|
|
|
|
# ------------------------------------------------------------------
|
|
# Depth-1 XML scrub
|
|
# ------------------------------------------------------------------
|
|
|
|
def _scrub_xml(self, messages: list[BaseMessage]) -> None:
|
|
if not self._prune_history_xml:
|
|
return
|
|
# Keep the UI list only in the newest xml_scrub_depth observation
|
|
# messages; everything older (or non-image XML carriers) is stripped.
|
|
# With a live tail outside the tracked list (tail_offset > 0) the tail
|
|
# itself occupies the newest slots of the keep window.
|
|
keep_count = max(0, self._xml_scrub_depth - self._tail_offset)
|
|
keep_indices = (
|
|
{rec["idx"] for rec in self._tracked[-keep_count:]} if keep_count else set()
|
|
)
|
|
still_pending: list[int] = []
|
|
for idx in self._xml_candidates:
|
|
if idx in keep_indices:
|
|
still_pending.append(idx)
|
|
continue
|
|
if idx in self._frozen:
|
|
continue
|
|
self._strip_xml_blocks(messages[idx])
|
|
self._xml_candidates = still_pending
|
|
|
|
def _strip_xml_blocks(self, msg: BaseMessage) -> None:
|
|
content = getattr(msg, "content", None)
|
|
if not isinstance(content, list):
|
|
return
|
|
new_blocks: list[Any] = []
|
|
for block in content:
|
|
if (
|
|
isinstance(block, dict)
|
|
and block.get("type") == "text"
|
|
and self._has_strip_marker(block.get("text", ""))
|
|
):
|
|
retained_text = _without_marked_suffix(
|
|
block.get("text", ""), self._strip_markers
|
|
)
|
|
if retained_text:
|
|
retained_block = dict(block)
|
|
retained_block["text"] = retained_text
|
|
new_blocks.append(retained_block)
|
|
continue
|
|
new_blocks.append(block)
|
|
msg.content = new_blocks
|
|
|
|
# ------------------------------------------------------------------
|
|
# Depth-K image scrub
|
|
# ------------------------------------------------------------------
|
|
|
|
def _scrub_images(self, messages: list[BaseMessage]) -> None:
|
|
total = len(self._tracked)
|
|
scrub_depth = self._image_scrub_depth
|
|
grace_limit = scrub_depth + self._pending_grace_steps
|
|
|
|
for pos, rec in enumerate(self._tracked):
|
|
if rec["idx"] in self._frozen:
|
|
continue
|
|
# 1 == most recent image-bearing message; a live tail outside the
|
|
# tracked list (tail_offset > 0) counts as the newest depths.
|
|
depth = total - pos + self._tail_offset
|
|
if depth < scrub_depth:
|
|
continue # active window
|
|
if self._tail_offset == 0 and pos == total - 1:
|
|
continue # the live observation is never scrubbed
|
|
|
|
summary, pending, failed, step_no = self._lookup(rec)
|
|
if summary is not None:
|
|
replacement = {
|
|
"type": "text",
|
|
"text": f"{_HISTORY_SUMMARY_PREFIX}{summary}",
|
|
}
|
|
elif failed:
|
|
replacement = {
|
|
"type": "text",
|
|
"text": f"[visual summary unavailable; evidence at DataEngine step {step_no}]",
|
|
}
|
|
elif pending and depth <= grace_limit:
|
|
continue # grace period: retain the image, revisit next pass
|
|
elif pending:
|
|
replacement = {
|
|
"type": "text",
|
|
"text": f"[visual summary pending; evidence at DataEngine step {step_no}]",
|
|
}
|
|
else:
|
|
replacement = None # no summary job exists; drop the image silently
|
|
|
|
self._resolve(messages, rec, replacement)
|
|
|
|
def _lookup(self, rec: dict[str, Any]) -> tuple[str | None, bool, bool, int | None]:
|
|
"""Mirror the legacy keying: tool_call_id first, ordinal fallback."""
|
|
summarizer = self._summarizer
|
|
if summarizer is None:
|
|
return None, False, False, None
|
|
|
|
effective_key: Any = rec["key"]
|
|
if effective_key is None:
|
|
return None, False, False, None
|
|
|
|
summary = summarizer.get_summary(effective_key)
|
|
if (
|
|
summary is None
|
|
and not summarizer.has_job(effective_key)
|
|
and rec["legacy"] is not None
|
|
):
|
|
effective_key = rec["legacy"]
|
|
summary = summarizer.get_summary(effective_key)
|
|
|
|
failed = summarizer.has_failed(effective_key)
|
|
pending = summarizer.is_pending(effective_key) and not failed
|
|
step_no = summarizer.get_step_number(effective_key)
|
|
if step_no is None:
|
|
step_no = rec["legacy"]
|
|
return summary, pending, failed, step_no
|
|
|
|
def _resolve(
|
|
self,
|
|
messages: list[BaseMessage],
|
|
rec: dict[str, Any],
|
|
replacement: dict[str, Any] | None,
|
|
) -> None:
|
|
"""Replace the message's image blocks and freeze it permanently."""
|
|
msg = messages[rec["idx"]]
|
|
content = getattr(msg, "content", None)
|
|
new_blocks: list[Any] = []
|
|
if isinstance(content, list):
|
|
for block in content:
|
|
if isinstance(block, dict) and block.get("type") in ("image_url", "image"):
|
|
continue
|
|
if (
|
|
self._prune_history_xml
|
|
and isinstance(block, dict)
|
|
and block.get("type") == "text"
|
|
and self._has_strip_marker(block.get("text", ""))
|
|
):
|
|
# Freezing implies fully scrubbed: no stale UI list survives
|
|
# even when xml_scrub_depth exceeds the image depth.
|
|
retained_text = _without_marked_suffix(
|
|
block.get("text", ""), self._strip_markers
|
|
)
|
|
if retained_text:
|
|
retained_block = dict(block)
|
|
retained_block["text"] = retained_text
|
|
new_blocks.append(retained_block)
|
|
continue
|
|
new_blocks.append(block)
|
|
|
|
if replacement is not None:
|
|
new_blocks.append(replacement)
|
|
|
|
# Defensive fallback: never leave ToolMessage or HumanMessage content empty
|
|
if not new_blocks:
|
|
new_blocks = [
|
|
{"type": "text", "text": "Action completed." if rec["is_tool"] else ""}
|
|
]
|
|
|
|
msg.content = new_blocks
|
|
self._frozen.add(rec["idx"])
|
|
if rec["idx"] in self._xml_candidates:
|
|
self._xml_candidates.remove(rec["idx"])
|