somew1nd
a45daf23f8
Report device actions as dispatched and preserve full failure details. Route Flash through McpActionExecutor and use ToolFailure across helper tools and agent loops instead of inferring failure from response text. Require model-provided descriptions for coordinate targets. Keep observed element metadata separate from those descriptions in safety checks, visual summaries, history, and frontend labels. Add Gemini Interactions API video analysis with configurable processing mode, chunk size, and timeouts. Centralize Google provider capabilities and token usage normalization, and update google-genai to 2.22.0. Preserve system instructions on follow-up interactions, replan pending chunks when falling back to static processing, and reject unanswered native video and audio segments instead of caching placeholder results. Keep helper-tool failure status intact through history tools and committee calls. Validation: 1,762 Python unit tests passed (3 skipped), 40 frontend formatter tests passed, and the production UI build completed. Ruff lint, formatting, the quality ratchet, and git diff checks passed.
322 lines
14 KiB
Python
322 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.
|
|
|
|
"""Context compression shared by Flash and Pro.
|
|
|
|
:class:`ScrubEdgeCompressor` strips XML at depth 1 and resolves screenshots
|
|
at depth K, then freezes the message. 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).
|
|
"""
|
|
|
|
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__)
|
|
|
|
#: Header written above a resolved visual summary (public: the chunk capsule
|
|
#: lens checks it to avoid repeating a summary already present verbatim).
|
|
HISTORY_SUMMARY_PREFIX = "--- Historical Visual Transition ---\n"
|
|
_HISTORY_SUMMARY_PREFIX = HISTORY_SUMMARY_PREFIX
|
|
_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:
|
|
"""Replace older observation details with summaries near the message tail.
|
|
|
|
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.
|
|
|
|
The Pro transcript ledger uses this 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 dispatched." 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"])
|