471 lines
19 KiB
Python
471 lines
19 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.
|
|
|
|
"""Session message ledger shared by the Pro operator and Flash runner.
|
|
|
|
Messages are organized into four regions:
|
|
|
|
- **S (stable prefix)**: the static system message — byte-identical for the
|
|
whole session (the plan+history section is moved out by the template split).
|
|
- **F (frozen history)**: restored history and compressed chunks.
|
|
- **A (active window)**: raw per-turn messages, append-only. Each committed
|
|
turn contributes its tail observation HumanMessage, the operator AIMessages
|
|
(tool_calls and native thinking preserved by reference), the in-turn
|
|
ToolMessages, and the turn's validator result message. The scrub edge
|
|
(:class:`~artemis.agents.flash.context_compressor.ScrubEdgeCompressor`) strips
|
|
old UI lists and plan recitations at depth 1 and resolves screenshots to
|
|
visual summaries at depth K — messages are never
|
|
removed or reordered, so tool-call/response pairs are never split.
|
|
- **T (current tail)**: built fresh every turn by the operator; passed to
|
|
:meth:`render` and only enters A when the turn is committed.
|
|
|
|
All timestamps inside the transcript use the session-start offset ``T+mm:ss``
|
|
(byte-stable once frozen); "ago" wording is reserved for the auxiliary agents'
|
|
per-call compiled views. When the caller supplies the DataEngine
|
|
``session_start`` epoch, the ledger clock is anchored to it so the tail
|
|
offsets, the chunk ledger lines and the video analyzer's action timeline all
|
|
share one origin.
|
|
|
|
Profiles: the Pro operator (one committed turn per graph step) and the Flash
|
|
runner (one committed turn per reactive turn; a multi-action turn registers
|
|
every recorded step id via ``extra_step_keys`` so the chunk ledger lists each
|
|
step) build their prompts from the same ledger.
|
|
"""
|
|
|
|
import json
|
|
import time
|
|
from typing import Any, Callable
|
|
|
|
from langchain_core.messages import BaseMessage, HumanMessage
|
|
|
|
from artemis.memory.step_memory import StepMemoryService
|
|
from artemis.utils.logger import get_logger
|
|
|
|
logger = get_logger(__name__)
|
|
|
|
#: Marker for the per-turn task-plan recitation block in the observation tail.
|
|
PLAN_RECITATION_MARKER = "--- Task Plan (recited) ---"
|
|
|
|
#: Marker of the Pro observation's UI element list block (see
|
|
#: ``ObservationPromptComponent``).
|
|
PRO_UI_LIST_MARKER = "--- Visible UI Elements ---"
|
|
|
|
#: Marker prefix of a committed turn's validator result message.
|
|
EXECUTION_RESULT_MARKER = "--- Action Execution Result"
|
|
|
|
#: Header prefix of the cold-start restored-history block.
|
|
RESTORED_HISTORY_HEADER = "[Restored history]"
|
|
|
|
|
|
def format_session_offset(seconds: float) -> str:
|
|
"""Render a session-start offset as ``T+mm:ss`` (minutes never wrap)."""
|
|
total = max(0, int(seconds))
|
|
return f"T+{total // 60:02d}:{total % 60:02d}"
|
|
|
|
|
|
#: Placeholder for an image block that is still inside the scrub edge's
|
|
#: pending-grace window when a turn transcript is rendered.
|
|
SCREENSHOT_PLACEHOLDER = "[screenshot]"
|
|
|
|
_ROLE_LABELS = {
|
|
"human": "observation",
|
|
"ai": "operator",
|
|
"tool": "tool result",
|
|
"system": "system",
|
|
}
|
|
|
|
|
|
def _render_content_text(content: Any) -> str:
|
|
"""Flatten message content to plain text.
|
|
|
|
Text blocks pass through verbatim, native thinking blocks are labelled,
|
|
image blocks become :data:`SCREENSHOT_PLACEHOLDER`.
|
|
"""
|
|
if content is None:
|
|
return ""
|
|
if isinstance(content, str):
|
|
return content
|
|
if not isinstance(content, list):
|
|
return str(content)
|
|
parts: list[str] = []
|
|
for block in content:
|
|
if isinstance(block, str):
|
|
parts.append(block)
|
|
elif not isinstance(block, dict):
|
|
parts.append(str(block))
|
|
elif block.get("type") == "text":
|
|
parts.append(str(block.get("text", "")))
|
|
elif block.get("type") == "thinking":
|
|
parts.append(f"(thinking) {block.get('thinking', '')}")
|
|
elif block.get("type") in ("image_url", "image"):
|
|
parts.append(SCREENSHOT_PLACEHOLDER)
|
|
else:
|
|
parts.append(json.dumps(block, ensure_ascii=False, default=str))
|
|
return "\n".join(part for part in parts if part)
|
|
|
|
|
|
def render_turn_transcript(messages: list[BaseMessage]) -> str:
|
|
"""Mechanical plain-text rendering of one committed turn's messages.
|
|
|
|
Lossless with respect to what the operator model was shown: every text
|
|
block, native thinking block, tool call (name and arguments) and tool
|
|
result is carried verbatim, in message order. Image blocks — present only
|
|
while a screenshot is still inside the scrub edge's pending-grace window —
|
|
become a placeholder; a resolved visual-transition summary is already
|
|
plain text. The segment capsule uses this transcript as its source when
|
|
compressing the turn.
|
|
|
|
A tool result is labelled with its tool's name, resolved through the
|
|
turn's own preceding tool calls when the message carries no name (the Pro
|
|
operator's ToolMessages don't); the opaque call id is shown only when no
|
|
name can be resolved.
|
|
"""
|
|
lines: list[str] = []
|
|
names_by_id: dict[str, str] = {}
|
|
for msg in messages:
|
|
kind = str(getattr(msg, "type", "") or "message")
|
|
role = _ROLE_LABELS.get(kind, kind)
|
|
if kind == "tool":
|
|
call_id = getattr(msg, "tool_call_id", None)
|
|
name = getattr(msg, "name", None) or names_by_id.get(str(call_id)) or call_id
|
|
if name:
|
|
role = f"{role} {name}"
|
|
status = getattr(msg, "status", None)
|
|
if status and status != "success":
|
|
role = f"{role} ({status})"
|
|
body = _render_content_text(getattr(msg, "content", None))
|
|
if body:
|
|
lines.append(f"[{role}]\n{body}")
|
|
for call in getattr(msg, "tool_calls", None) or []:
|
|
if not isinstance(call, dict):
|
|
continue
|
|
if call.get("id") and call.get("name"):
|
|
names_by_id[str(call["id"])] = str(call["name"])
|
|
args = json.dumps(call.get("args") or {}, ensure_ascii=False, default=str)
|
|
lines.append(f"[tool call] {call.get('name')}({args})")
|
|
return "\n".join(lines)
|
|
|
|
|
|
class TranscriptLedger:
|
|
"""Append-only four-region message ledger for one Pro session.
|
|
|
|
The ledger owns no LLM calls: summary lookups go through the shared
|
|
:class:`StepMemoryService` (``ctx.step_memory``), whose visual-transition
|
|
lens is fed by the Pro SummarizerNode each step.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
*,
|
|
step_memory: StepMemoryService | None = None,
|
|
prune_history_xml: bool = True,
|
|
image_scrub_depth: int = 3,
|
|
pending_grace_steps: int = 3,
|
|
xml_scrub_depth: int = 1,
|
|
clock: Callable[[], float] | None = None,
|
|
session_start: float | None = None,
|
|
):
|
|
# ``session_start`` is an epoch-seconds origin (the DataEngine's
|
|
# ``session_start_time``); with it the ledger reads the wall clock so
|
|
# every ``T+mm:ss`` label is relative to the recorded session start.
|
|
# Without it the ledger keeps its own monotonic origin.
|
|
if session_start is not None:
|
|
self._clock = clock or time.time
|
|
self._session_start = float(session_start)
|
|
else:
|
|
self._clock = clock or time.monotonic
|
|
self._session_start = self._clock()
|
|
|
|
self._static: list[BaseMessage] = []
|
|
self._restored: list[BaseMessage] = []
|
|
self._active: list[BaseMessage] = []
|
|
self._staged: list[BaseMessage] | None = None
|
|
self._turn_count = 0
|
|
|
|
# L2/L3 chunk compression. ``_turns`` records each committed
|
|
# turn's message span in ``_active``; a compression event advances
|
|
# ``_active_start`` past whole turns (never splitting tool-call pairs)
|
|
# and replaces the frozen chunk blocks wholesale. The underlying
|
|
# ``_active`` list stays append-only so scrub-edge indices stay valid.
|
|
self._turns: list[dict] = []
|
|
self._chunked_turn_count = 0
|
|
self._active_start = 0
|
|
self._frozen_blocks: list[BaseMessage] = []
|
|
self._chunker: Any | None = None
|
|
|
|
# id(message) -> summary-job key (DataEngine step id). A side map keeps
|
|
# the key out of the serialized message payload entirely.
|
|
self._step_keys: dict[int, str] = {}
|
|
|
|
# Imported here, not at module level: context_compressor imports
|
|
# artemis.memory.step_memory, so a top-level import forms a cycle
|
|
# whose failure depends on which side is imported first.
|
|
from artemis.agents.flash.context_compressor import ScrubEdgeCompressor
|
|
|
|
self._compressor = ScrubEdgeCompressor(
|
|
summarizer=step_memory,
|
|
prune_history_xml=prune_history_xml,
|
|
image_scrub_depth=image_scrub_depth,
|
|
pending_grace_steps=pending_grace_steps,
|
|
xml_scrub_depth=xml_scrub_depth,
|
|
summary_key_getter=lambda msg: self._step_keys.get(id(msg)),
|
|
strip_markers=(PRO_UI_LIST_MARKER, PLAN_RECITATION_MARKER),
|
|
tail_offset=1,
|
|
)
|
|
|
|
# ------------------------------------------------------------------
|
|
# Session clock
|
|
# ------------------------------------------------------------------
|
|
|
|
def elapsed_seconds(self) -> float:
|
|
return max(0.0, self._clock() - self._session_start)
|
|
|
|
def elapsed_label(self) -> str:
|
|
"""The current session offset as ``T+mm:ss``."""
|
|
return format_session_offset(self.elapsed_seconds())
|
|
|
|
# ------------------------------------------------------------------
|
|
# S region
|
|
# ------------------------------------------------------------------
|
|
|
|
@property
|
|
def has_static_prefix(self) -> bool:
|
|
return bool(self._static)
|
|
|
|
def set_static_prefix(self, messages: list[BaseMessage]) -> None:
|
|
"""Install the byte-stable system prefix; only settable once."""
|
|
if self._static:
|
|
raise RuntimeError("TranscriptLedger static prefix is already set.")
|
|
self._static = list(messages)
|
|
|
|
# ------------------------------------------------------------------
|
|
# F region: restored history
|
|
# ------------------------------------------------------------------
|
|
|
|
@property
|
|
def has_restored_history(self) -> bool:
|
|
return bool(self._restored)
|
|
|
|
def set_restored_history(self, text: str) -> None:
|
|
"""Install the cold-start frozen history block (empty ledger only)."""
|
|
if self._restored:
|
|
raise RuntimeError("TranscriptLedger restored history is already set.")
|
|
if self._active or self._staged or self._turn_count:
|
|
raise RuntimeError("Restored history can only seed an empty ledger (cold start).")
|
|
self._restored = [HumanMessage(content=[{"type": "text", "text": text}])]
|
|
|
|
# ------------------------------------------------------------------
|
|
# A region (append-only turn commits)
|
|
# ------------------------------------------------------------------
|
|
|
|
@property
|
|
def turn_count(self) -> int:
|
|
return self._turn_count
|
|
|
|
@property
|
|
def has_staged_turn(self) -> bool:
|
|
return self._staged is not None
|
|
|
|
@property
|
|
def active_messages(self) -> tuple[BaseMessage, ...]:
|
|
return tuple(self._active)
|
|
|
|
def stage_turn(self, messages: list[BaseMessage]) -> None:
|
|
"""Hold a finished turn's messages until the next build commits them.
|
|
|
|
Committing is deferred because the turn's DataEngine step id and its
|
|
validator result only exist after the operator returns.
|
|
"""
|
|
if self._staged is not None:
|
|
logger.warning(
|
|
"TranscriptLedger: previous staged turn was never committed;"
|
|
" committing it without step metadata."
|
|
)
|
|
self.commit_staged()
|
|
self._staged = [m for m in messages if isinstance(m, BaseMessage)]
|
|
|
|
def commit_staged(
|
|
self,
|
|
*,
|
|
step_key: str | None = None,
|
|
validator_result: Any | None = None,
|
|
extra_step_keys: tuple[str, ...] | list[str] = (),
|
|
) -> None:
|
|
"""Move the staged turn into the active region.
|
|
|
|
Args:
|
|
step_key: The DataEngine step id recorded for this turn; keys the
|
|
turn's observation screenshot to its visual-transition summary
|
|
job for the depth-K scrub.
|
|
validator_result: The turn's validator report (``None`` when the
|
|
turn executed no terminal action); appended as a frozen result
|
|
message carrying the ``T+mm:ss`` session offset.
|
|
extra_step_keys: Further DataEngine step ids recorded during the
|
|
same turn (a Flash turn that executed several actions records
|
|
one step per action). They join ``step_key`` in the turn's
|
|
``step_keys`` so chunk compression lists every step.
|
|
"""
|
|
if self._staged is None:
|
|
return
|
|
staged = self._staged
|
|
self._staged = None
|
|
|
|
if step_key is not None:
|
|
observation = self._first_image_message(staged)
|
|
if observation is not None:
|
|
self._step_keys[id(observation)] = str(step_key)
|
|
|
|
span_start = len(self._active)
|
|
self._active.extend(staged)
|
|
|
|
if validator_result is not None:
|
|
result_message = self._build_result_message(validator_result)
|
|
if result_message is not None:
|
|
self._active.append(result_message)
|
|
|
|
step_keys: list[str] = []
|
|
for key in (step_key, *extra_step_keys):
|
|
if key is not None and str(key) not in step_keys:
|
|
step_keys.append(str(key))
|
|
self._turns.append(
|
|
{
|
|
"step_key": str(step_key) if step_key is not None else None,
|
|
"step_keys": step_keys,
|
|
"start": span_start,
|
|
"end": len(self._active),
|
|
}
|
|
)
|
|
self._turn_count += 1
|
|
|
|
@staticmethod
|
|
def _first_image_message(messages: list[BaseMessage]) -> BaseMessage | None:
|
|
for msg in messages:
|
|
content = getattr(msg, "content", None)
|
|
if not isinstance(content, list):
|
|
continue
|
|
for block in content:
|
|
if isinstance(block, dict) and block.get("type") in ("image_url", "image"):
|
|
return msg
|
|
return None
|
|
|
|
def _build_result_message(self, result: Any) -> HumanMessage | None:
|
|
rendered = None
|
|
if isinstance(result, dict):
|
|
status = result.get("status") or "unknown"
|
|
detail = None
|
|
try:
|
|
from artemis.utils.task_tree import format_result_clean
|
|
|
|
# format_result_clean reports only errors/repairs; success
|
|
# renders as the bare status line.
|
|
detail = format_result_clean(result)
|
|
except Exception:
|
|
detail = None
|
|
rendered = f"Status: {status}" + (f"\n{detail}" if detail else "")
|
|
else:
|
|
try:
|
|
rendered = json.dumps(result, ensure_ascii=False, default=str)
|
|
except Exception:
|
|
rendered = str(result)
|
|
if not rendered:
|
|
return None
|
|
return HumanMessage(
|
|
content=[
|
|
{
|
|
"type": "text",
|
|
"text": (f"{EXECUTION_RESULT_MARKER} ({self.elapsed_label()}) ---\n{rendered}"),
|
|
}
|
|
]
|
|
)
|
|
|
|
# ------------------------------------------------------------------
|
|
# F region: chunk compression
|
|
# ------------------------------------------------------------------
|
|
|
|
def attach_chunker(self, chunker: Any) -> None:
|
|
"""Install the L2/L3 :class:`HistoryChunkManager` consulted at render."""
|
|
self._chunker = chunker
|
|
|
|
@property
|
|
def chunker(self) -> Any | None:
|
|
return self._chunker
|
|
|
|
@property
|
|
def frozen_blocks(self) -> tuple[BaseMessage, ...]:
|
|
return tuple(self._frozen_blocks)
|
|
|
|
def unchunked_turns(self) -> list[dict]:
|
|
"""Committed turns not yet consumed by a compression event (copies)."""
|
|
return [dict(t) for t in self._turns[self._chunked_turn_count :]]
|
|
|
|
def turn_text_chars(self, turns: list[dict]) -> int:
|
|
"""Total characters of the given turns' rendered transcripts (size
|
|
trigger). Measured on the very text the chunk capsule receives, so
|
|
the trigger sizes exactly what gets compressed."""
|
|
return sum(len(self.turn_transcript(turn)) for turn in turns)
|
|
|
|
def turn_transcript(self, turn: dict) -> str:
|
|
"""Plain-text rendering of one committed turn exactly as it currently
|
|
stands in the active region (i.e. after the scrub edge), see
|
|
:func:`render_turn_transcript`. Chunk compression feeds this to the
|
|
segment capsule so the capsule digests precisely what it replaces."""
|
|
return render_turn_transcript(self._active[turn["start"] : turn["end"]])
|
|
|
|
def freeze_turns(self, turn_count: int, frozen_blocks: list[BaseMessage]) -> None:
|
|
"""Compression event: consume the oldest ``turn_count`` unchunked turns
|
|
and replace the frozen chunk blocks wholesale.
|
|
|
|
The boundary always advances to a committed turn's end, so a tool-call
|
|
message and its responses are never split across the F/A boundary.
|
|
"""
|
|
if turn_count <= 0:
|
|
self._frozen_blocks = list(frozen_blocks)
|
|
return
|
|
available = self._turns[self._chunked_turn_count :]
|
|
if turn_count > len(available):
|
|
raise ValueError(
|
|
f"freeze_turns({turn_count}) exceeds the {len(available)}"
|
|
" unchunked committed turns."
|
|
)
|
|
self._chunked_turn_count += turn_count
|
|
self._active_start = self._turns[self._chunked_turn_count - 1]["end"]
|
|
self._frozen_blocks = list(frozen_blocks)
|
|
|
|
# ------------------------------------------------------------------
|
|
# Rendering
|
|
# ------------------------------------------------------------------
|
|
|
|
def render(self, tail: list[BaseMessage]) -> list[BaseMessage]:
|
|
"""Advance compression and return ``S + F + A + tail``.
|
|
|
|
Order per turn: the chunker runs first (a triggered compression event
|
|
deep-mutates the frozen region and advances the F/A boundary), then
|
|
the scrub edge advances over the active window. The returned list is a
|
|
fresh container: appending to it (the operator's in-turn tool loop)
|
|
never mutates the ledger regions.
|
|
"""
|
|
if self._chunker is not None:
|
|
try:
|
|
self._chunker.on_render(self)
|
|
except Exception as e:
|
|
logger.error(f"History chunker render hook failed: {e}")
|
|
self._compressor.compress(self._active)
|
|
return [
|
|
*self._static,
|
|
*self._restored,
|
|
*self._frozen_blocks,
|
|
*self._active[self._active_start :],
|
|
*tail,
|
|
]
|