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"