From fd8718fb1846fc9789249197b149fddd79c04614 Mon Sep 17 00:00:00 2001 From: dexhunter <6930518+dexhunter@users.noreply.github.com> Date: Thu, 9 Jul 2026 14:29:39 +0000 Subject: [PATCH] perf(agent): incrementally prune old ARIA observations Track observation messages already scanned so pruning old ARIA payloads does not rescan the full transcript after every appended message. Reset the incremental state on compaction and run reset. Corrected eval on this branch: prune_aria_ms 174.4 ms to 120.5 ms with behavior hash preserved. --- src/webwright/agents/default.py | 77 ++++++++++----- src/webwright/models/openai_model.py | 2 - tests/unit/test_default_agent_aria_pruning.py | 96 +++++++++++++++++++ 3 files changed, 150 insertions(+), 25 deletions(-) create mode 100644 tests/unit/test_default_agent_aria_pruning.py diff --git a/src/webwright/agents/default.py b/src/webwright/agents/default.py index d05ea86..c93f19b 100644 --- a/src/webwright/agents/default.py +++ b/src/webwright/agents/default.py @@ -89,6 +89,10 @@ def __init__(self, model: Model, env: Environment, *, config_class: type = Agent self.extra_template_vars: dict[str, Any] = {} self.n_calls = 0 self.n_format_errors = 0 + # Track observations already scanned so pruning only visits new messages. + self._observations: list[tuple[dict[str, Any], dict[str, Any]]] = [] + self._next_msg_index = 0 + self._pruned_until_obs_index = 0 def _debug_dir(self) -> Path | None: if self.config.output_path is None: @@ -276,29 +280,47 @@ def _prune_old_observation_aria_snapshots(self) -> None: n = self.config.keep_last_n_observations if n <= 0: return - obs_indices = [ - i for i, m in enumerate(self.messages) - if m.get("extra", {}).get("observation") - ] - if len(obs_indices) <= n: - return - placeholder = "(ARIA snapshot pruned; see most recent observation)" - for idx in obs_indices[:-n]: - msg = self.messages[idx] - obs = msg["extra"]["observation"] - aria = obs.get("aria_snapshot", "") - if not aria: - continue - content = msg.get("content") - if isinstance(content, list): - for part in content: - if isinstance(part, dict) and part.get("type") in ("text", "input_text"): - text = part.get("text", "") - if aria in text: - part["text"] = text.replace(aria, placeholder) - elif isinstance(content, str) and aria in content: - msg["content"] = content.replace(aria, placeholder) - obs["aria_snapshot"] = "" + + msgs = self.messages + observations = self._observations + curr_len = len(msgs) + next_idx = self._next_msg_index + + # Scan only messages appended since the previous pruning pass. + if next_idx < curr_len: + self._next_msg_index = curr_len + append_observation = observations.append + for i in range(next_idx, curr_len): + msg = msgs[i] + extra = msg.get("extra") + if extra: + obs = extra.get("observation") + if obs: + append_observation((msg, obs)) + + num_to_prune = len(observations) - n + pruned_idx = self._pruned_until_obs_index + if pruned_idx < num_to_prune: + self._pruned_until_obs_index = num_to_prune + placeholder = "(ARIA snapshot pruned; see most recent observation)" + for i in range(pruned_idx, num_to_prune): + msg, obs = observations[i] + aria = obs.get("aria_snapshot") + if aria: + content = msg.get("content") + if isinstance(content, list): + for part in content: + if not isinstance(part, dict) or part.get("type") not in ("text", "input_text"): + continue + text = part.get("text", "") + new_text = text.replace(aria, placeholder) + if new_text != text: + part["text"] = new_text + elif isinstance(content, str): + new_text = content.replace(aria, placeholder) + if new_text != content: + msg["content"] = new_text + obs["aria_snapshot"] = "" def _compact_history(self) -> None: """Summarize the running transcript via an LLM call and reset messages to [system, summary]. @@ -337,12 +359,21 @@ def _compact_history(self) -> None: extra={"interrupt_type": "HistoryCompactionSummary"}, ) self.messages = [system_message, summary_message] + # Reset pruning state after history compaction + self._observations = [] + self._next_msg_index = 0 + self._pruned_until_obs_index = 0 def run(self, task: str = "", **kwargs) -> dict[str, Any]: self.extra_template_vars |= {"task": task, **kwargs} self.messages = [] self.n_calls = 0 self.n_format_errors = 0 + # Reset pruning state for a new run + self._observations = [] + self._next_msg_index = 0 + self._pruned_until_obs_index = 0 + self.add_messages( self.model.format_message(role="system", content=self._render_template(self.config.system_template)), self.model.format_message(role="user", content=self._render_template(self.config.instance_template)), diff --git a/src/webwright/models/openai_model.py b/src/webwright/models/openai_model.py index f11806e..4585b80 100644 --- a/src/webwright/models/openai_model.py +++ b/src/webwright/models/openai_model.py @@ -2,7 +2,6 @@ from __future__ import annotations -from pathlib import Path from typing import Any from webwright.models.base import ( @@ -10,7 +9,6 @@ BaseModelConfig, OptStr, _safe_int, - image_part_from_path, text_part, ) diff --git a/tests/unit/test_default_agent_aria_pruning.py b/tests/unit/test_default_agent_aria_pruning.py new file mode 100644 index 0000000..dd19555 --- /dev/null +++ b/tests/unit/test_default_agent_aria_pruning.py @@ -0,0 +1,96 @@ +from typing import Any + +from webwright.agents.default import DefaultAgent + + +class StubModel: + def __init__(self) -> None: + self.queries: list[list[dict[str, Any]]] = [] + + def format_message(self, role: str, content: str, **kwargs: Any) -> dict[str, Any]: + return {"role": role, "content": content, "extra": kwargs.get("extra", {})} + + def get_template_vars(self) -> dict[str, Any]: + return {} + + def query(self, messages: list[dict[str, Any]]) -> dict[str, Any]: + self.queries.append(messages) + return self.format_message(role="assistant", content="summary") + + +class StubEnv: + def get_template_vars(self) -> dict[str, Any]: + return {} + + +def make_agent() -> DefaultAgent: + return DefaultAgent( + StubModel(), + StubEnv(), + system_template="system", + instance_template="instance", + keep_last_n_observations=1, + ) + + +def observation_message(label: str, aria_snapshot: str, content: Any | None = None) -> dict[str, Any]: + if content is None: + content = f"Observation {label}\n{aria_snapshot}" + return { + "role": "user", + "content": content, + "extra": { + "observation": { + "label": label, + "aria_snapshot": aria_snapshot, + } + }, + } + + +def test_prunes_old_observation_aria_snapshots_only() -> None: + agent = make_agent() + first = observation_message("first", "ARIA_FIRST") + middle = {"role": "assistant", "content": "no observation", "extra": {}} + second = observation_message( + "second", + "ARIA_SECOND", + content=[{"type": "text", "text": "Observation second\nARIA_SECOND"}], + ) + third = observation_message("third", "ARIA_THIRD") + + agent.add_messages(first, middle, second, third) + agent._prune_old_observation_aria_snapshots() + + assert first["extra"]["observation"]["aria_snapshot"] == "" + assert "ARIA_FIRST" not in first["content"] + assert "(ARIA snapshot pruned; see most recent observation)" in first["content"] + + assert middle["content"] == "no observation" + + assert second["extra"]["observation"]["aria_snapshot"] == "" + assert "ARIA_SECOND" not in second["content"][0]["text"] + assert "(ARIA snapshot pruned; see most recent observation)" in second["content"][0]["text"] + + assert third["extra"]["observation"]["aria_snapshot"] == "ARIA_THIRD" + assert "ARIA_THIRD" in third["content"] + + +def test_compaction_resets_incremental_pruning_state() -> None: + agent = make_agent() + system = agent.model.format_message(role="system", content="system") + first = observation_message("first", "ARIA_FIRST") + second = observation_message("second", "ARIA_SECOND") + + agent.add_messages(system, first, second) + assert first["extra"]["observation"]["aria_snapshot"] == "" + assert second["extra"]["observation"]["aria_snapshot"] == "ARIA_SECOND" + + agent._compact_history() + + third = observation_message("third", "ARIA_THIRD") + fourth = observation_message("fourth", "ARIA_FOURTH") + agent.add_messages(third, fourth) + + assert third["extra"]["observation"]["aria_snapshot"] == "" + assert fourth["extra"]["observation"]["aria_snapshot"] == "ARIA_FOURTH"