feat(memory): add incremental transcript compaction
This commit is contained in:
somew1nd committed
1 parent
cbe8f4f0c6
commit
41ae386a3a
65 files changed
+9971
-815
No files matched your search
@@ -0,0 +1,364 @@
|
||||
# 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-level message transcript ledger for the Pro operator (M2).
|
||||
|
||||
Implements the four-region discipline of the history-module redesign §3.2:
|
||||
|
||||
- **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)**: in M2 only the cold-start restored-history block; L2
|
||||
chunks arrive in M3.
|
||||
- **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`,
|
||||
generalized in M2) strips old UI lists / 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.
|
||||
"""
|
||||
|
||||
import json
|
||||
import time
|
||||
from typing import Any, Callable
|
||||
|
||||
from langchain_core.messages import BaseMessage, HumanMessage
|
||||
|
||||
from artemis.agents.flash.context_compressor import ScrubEdgeCompressor
|
||||
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}"
|
||||
|
||||
|
||||
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,
|
||||
):
|
||||
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 (M3). ``_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] = {}
|
||||
|
||||
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 (M2: cold-start restored history only)
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
@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,
|
||||
) -> 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.
|
||||
"""
|
||||
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)
|
||||
|
||||
self._turns.append(
|
||||
{
|
||||
"step_key": str(step_key) if step_key is not None else None,
|
||||
"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 (M3): chunk compression primitives
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
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 text characters of the given turns' messages (size trigger)."""
|
||||
total = 0
|
||||
for turn in turns:
|
||||
for msg in self._active[turn["start"] : turn["end"]]:
|
||||
content = getattr(msg, "content", None)
|
||||
if isinstance(content, str):
|
||||
total += len(content)
|
||||
elif isinstance(content, list):
|
||||
for block in content:
|
||||
if isinstance(block, dict) and block.get("type") == "text":
|
||||
total += len(block.get("text", ""))
|
||||
return total
|
||||
|
||||
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,
|
||||
]
|
||||
Reference in new issue
Block a user