fix: Context window truncation using CondensationAction (All-Hands-AI#7578)

Co-authored-by: Calvin Smith <[email protected]>
Co-authored-by: Graham Neubig <[email protected]>

Author: 3 people

openhands/agenthub/browsing_agent/browsing_agent.py

@@ -150,13 +150,13 @@ def step(self, state: State) -> Action:
         last_obs = None
         last_action = None
 
-        if EVAL_MODE and len(state.history) == 1:
+        if EVAL_MODE and len(state.view) == 1:
             # for webarena and miniwob++ eval, we need to retrieve the initial observation already in browser env
             # initialize and retrieve the first observation by issuing an noop OP
             # For non-benchmark browsing, the browser env starts with a blank page, and the agent is expected to first navigate to desired websites
             return BrowseInteractiveAction(browser_actions='noop()')
 
-        for event in state.history:
+        for event in state.view:
             if isinstance(event, BrowseInteractiveAction):
                 prev_actions.append(event.browser_actions)
                 last_action = event

openhands/agenthub/dummy_agent/agent.py

@@ -130,7 +130,7 @@ def step(self, state: State) -> Action:
 
             if 'observations' in prev_step and prev_step['observations']:
                 expected_observations = prev_step['observations']
-                hist_events = state.history[-len(expected_observations) :]
+                hist_events = state.view[-len(expected_observations) :]
 
                 if len(hist_events) < len(expected_observations):
                     print(

openhands/agenthub/visualbrowsing_agent/visualbrowsing_agent.py

@@ -204,13 +204,13 @@ def step(self, state: State) -> Action:
         last_action = None
         set_of_marks = None  # Initialize set_of_marks to None
 
-        if len(state.history) == 1:
+        if len(state.view) == 1:
             # for visualwebarena, webarena and miniwob++ eval, we need to retrieve the initial observation already in browser env
             # initialize and retrieve the first observation by issuing an noop OP
             # For non-benchmark browsing, the browser env starts with a blank page, and the agent is expected to first navigate to desired websites
             return BrowseInteractiveAction(browser_actions='noop(1000)')
 
-        for event in state.history:
+        for event in state.view:
             if isinstance(event, BrowseInteractiveAction):
                 prev_actions.append(event)
                 last_action = event

openhands/controller/agent_controller.py

@@ -57,7 +57,6 @@
 from openhands.events.action.agent import CondensationAction, RecallAction
 from openhands.events.event import Event
 from openhands.events.observation import (
-    AgentCondensationObservation,
     AgentDelegateObservation,
     AgentStateChangedObservation,
     ErrorObservation,
@@ -928,12 +927,6 @@ def _init_history(self) -> None:
         - For delegate events (between AgentDelegateAction and AgentDelegateObservation):
             - Excludes all events between the action and observation
             - Includes the delegate action and observation themselves
-
-        The history is loaded in two parts if truncation_id is set:
-        1. First user message from start_id onwards
-        2. Rest of history from truncation_id to the end
-
-        Otherwise loads normally from start_id.
         """
         # define range of events to fetch
         # delegates start with a start_id and initially won't find any events
@@ -956,29 +949,6 @@ def _init_history(self) -> None:
 
         events: list[Event] = []
 
-        # If we have a truncation point, get first user message and then rest of history
-        if hasattr(self.state, 'truncation_id') and self.state.truncation_id > 0:
-            # Find first user message from stream
-            first_user_msg = next(
-                (
-                    e
-                    for e in self.event_stream.get_events(
-                        start_id=start_id,
-                        end_id=end_id,
-                        reverse=False,
-                        filter_out_type=self.filter_out,
-                        filter_hidden=True,
-                    )
-                    if isinstance(e, MessageAction) and e.source == EventSource.USER
-                ),
-                None,
-            )
-            if first_user_msg:
-                events.append(first_user_msg)
-
-            # the rest of the events are from the truncation point
-            start_id = self.state.truncation_id
-
         # Get rest of history
         events_to_add = list(
             self.event_stream.get_events(
@@ -1046,16 +1016,20 @@ def _init_history(self) -> None:
 
     def _handle_long_context_error(self) -> None:
         # When context window is exceeded, keep roughly half of agent interactions
-        self.state.history = self._apply_conversation_window(self.state.history)
+        kept_event_ids = {
+            e.id for e in self._apply_conversation_window(self.state.history)
+        }
+        forgotten_event_ids = {e.id for e in self.state.history} - kept_event_ids
 
         # Save the ID of the first event in our truncated history for future reloading
         if self.state.history:
             self.state.start_id = self.state.history[0].id
 
         # Add an error event to trigger another step by the agent
         self.event_stream.add_event(
-            AgentCondensationObservation(
-                content='Trimming prompt to meet context window limitations'
+            CondensationAction(
+                forgotten_events_start_id=min(forgotten_event_ids),
+                forgotten_events_end_id=max(forgotten_event_ids),
             ),
             EventSource.AGENT,
         )
@@ -1133,10 +1107,6 @@ def _apply_conversation_window(self, events: list[Event]) -> list[Event]:
                     # if it's an action with source == EventSource.AGENT, we're good
                     break
 
-        # Save where to continue from in next reload
-        if kept_events:
-            self.state.truncation_id = kept_events[0].id
-
         # Ensure first user message is included
         if first_user_msg and first_user_msg not in kept_events:
             kept_events = [first_user_msg] + kept_events

openhands/controller/state/state.py

@@ -15,6 +15,7 @@
 from openhands.events.action.agent import AgentFinishAction
 from openhands.events.event import Event, EventSource
 from openhands.llm.metrics import Metrics
+from openhands.memory.view import View
 from openhands.storage.files import FileStore
 from openhands.storage.locations import get_conversation_agent_state_filename
 
@@ -96,8 +97,6 @@ class State:
     # start_id and end_id track the range of events in history
     start_id: int = -1
     end_id: int = -1
-    # truncation_id tracks where to load history after context window truncation
-    truncation_id: int = -1
 
     delegates: dict[tuple[int, int], tuple[str, str]] = field(default_factory=dict)
     # NOTE: This will never be used by the controller, but it can be used by different
@@ -170,6 +169,12 @@ def __getstate__(self):
         # don't pickle history, it will be restored from the event stream
         state = self.__dict__.copy()
         state['history'] = []
+
+        # Remove any view caching attributes. They'll be rebuilt frmo the
+        # history after that gets reloaded.
+        state.pop('_history_checksum', None)
+        state.pop('_view', None)
+
         return state
 
     def __setstate__(self, state):
@@ -183,7 +188,7 @@ def get_current_user_intent(self) -> tuple[str | None, list[str] | None]:
         """Returns the latest user message and image(if provided) that appears after a FinishAction, or the first (the task) if nothing was finished yet."""
         last_user_message = None
         last_user_message_image_urls: list[str] | None = []
-        for event in reversed(self.history):
+        for event in reversed(self.view):
             if isinstance(event, MessageAction) and event.source == 'user':
                 last_user_message = event.content
                 last_user_message_image_urls = event.image_urls
@@ -194,13 +199,13 @@ def get_current_user_intent(self) -> tuple[str | None, list[str] | None]:
         return last_user_message, last_user_message_image_urls
 
     def get_last_agent_message(self) -> MessageAction | None:
-        for event in reversed(self.history):
+        for event in reversed(self.view):
             if isinstance(event, MessageAction) and event.source == EventSource.AGENT:
                 return event
         return None
 
     def get_last_user_message(self) -> MessageAction | None:
-        for event in reversed(self.history):
+        for event in reversed(self.view):
             if isinstance(event, MessageAction) and event.source == EventSource.USER:
                 return event
         return None
@@ -211,7 +216,22 @@ def to_llm_metadata(self, agent_name: str) -> dict:
             'trace_version': openhands.__version__,
             'tags': [
                 f'agent:{agent_name}',
-                f'web_host:{os.environ.get("WEB_HOST", "unspecified")}',
+                f"web_host:{os.environ.get('WEB_HOST', 'unspecified')}",
                 f'openhands_version:{openhands.__version__}',
             ],
         }
+
+    @property
+    def view(self) -> View:
+        # Compute a simple checksum from the history to see if we can re-use any
+        # cached view.
+        history_checksum = len(self.history)
+        old_history_checksum = getattr(self, '_history_checksum', -1)
+
+        # If the history has changed, we need to re-create the view and update
+        # the caching.
+        if history_checksum != old_history_checksum:
+            self._history_checksum = history_checksum
+            self._view = View.from_events(self.history)
+
+        return self._view

openhands/memory/__init__.py

@@ -1,3 +0,0 @@
-from openhands.memory.condenser import Condenser
-
-__all__ = ['Condenser']

openhands/memory/condenser/condenser.py

@@ -2,15 +2,14 @@
 
 from abc import ABC, abstractmethod
 from contextlib import contextmanager
-from typing import Any, overload
+from typing import Any
 
 from pydantic import BaseModel
 
 from openhands.controller.state.state import State
 from openhands.core.config.condenser_config import CondenserConfig
 from openhands.events.action.agent import CondensationAction
-from openhands.events.event import Event
-from openhands.events.observation.agent import AgentCondensationObservation
+from openhands.memory.view import View
 
 CONDENSER_METADATA_KEY = 'condenser_meta'
 """Key identifying where metadata is stored in a `State` object's `extra_data` field."""
@@ -34,69 +33,6 @@ def get_condensation_metadata(state: State) -> list[dict[str, Any]]:
 """Registry of condenser configurations to their corresponding condenser classes."""
 
 
-class View(BaseModel):
-    """Linearly ordered view of events.
-
-    Produced by a condenser to indicate the included events are ready to process as LLM input.
-    """
-
-    events: list[Event]
-
-    def __len__(self) -> int:
-        return len(self.events)
-
-    def __iter__(self):
-        return iter(self.events)
-
-    # To preserve list-like indexing, we ideally support slicing and position-based indexing.
-    # The only challenge with that is switching the return type based on the input type -- we
-    # can mark the different signatures for MyPy with `@overload` decorators.
-
-    @overload
-    def __getitem__(self, key: slice) -> list[Event]: ...
-
-    @overload
-    def __getitem__(self, key: int) -> Event: ...
-
-    def __getitem__(self, key: int | slice) -> Event | list[Event]:
-        if isinstance(key, slice):
-            start, stop, step = key.indices(len(self))
-            return [self[i] for i in range(start, stop, step)]
-        elif isinstance(key, int):
-            return self.events[key]
-        else:
-            raise ValueError(f'Invalid key type: {type(key)}')
-
-    @staticmethod
-    def from_events(events: list[Event]) -> View:
-        """Create a view from a list of events, respecting the semantics of any condensation events."""
-        forgotten_event_ids: set[int] = set()
-        for event in events:
-            if isinstance(event, CondensationAction):
-                forgotten_event_ids.update(event.forgotten)
-
-        kept_events = [event for event in events if event.id not in forgotten_event_ids]
-
-        # If we have a summary, insert it at the specified offset.
-        summary: str | None = None
-        summary_offset: int | None = None
-
-        # The relevant summary is always in the last condensation event (i.e., the most recent one).
-        for event in reversed(events):
-            if isinstance(event, CondensationAction):
-                if event.summary is not None and event.summary_offset is not None:
-                    summary = event.summary
-                    summary_offset = event.summary_offset
-                    break
-
-        if summary is not None and summary_offset is not None:
-            kept_events.insert(
-                summary_offset, AgentCondensationObservation(content=summary)
-            )
-
-        return View(events=kept_events)
-
-
 class Condensation(BaseModel):
     """Produced by a condenser to indicate the history has been condensed."""
 
@@ -150,13 +86,13 @@ def metadata_batch(self, state: State):
             self.write_metadata(state)
 
     @abstractmethod
-    def condense(self, events: list[Event]) -> View | Condensation:
+    def condense(self, View) -> View | Condensation:
         """Condense a sequence of events into a potentially smaller list.
 
         New condenser strategies should override this method to implement their own condensation logic. Call `self.add_metadata` in the implementation to record any relevant per-condensation diagnostic information.
 
         Args:
-            events: A list of events representing the entire history of the agent.
+            View: A view of the history containing all events that should be condensed.
 
         Returns:
             View | Condensation: A condensed view of the events or an event indicating the history has been condensed.
@@ -165,7 +101,7 @@ def condense(self, events: list[Event]) -> View | Condensation:
     def condensed_history(self, state: State) -> View | Condensation:
         """Condense the state's history."""
         with self.metadata_batch(state):
-            return self.condense(state.history)
+            return self.condense(state.view)
 
     @classmethod
     def register_config(cls, configuration_type: type[CondenserConfig]) -> None:
@@ -221,10 +157,7 @@ def should_condense(self, view: View) -> bool:
     def get_condensation(self, view: View) -> Condensation:
         """Get the condensation from a view."""
 
-    def condense(self, events: list[Event]) -> View | Condensation:
-        # Convert the state to a view. This might require some condenser-specific logic.
-        view = View.from_events(events)
-
+    def condense(self, view: View) -> View | Condensation:
         # If we trigger the condenser-specific condensation threshold, compute and return
         # the condensation.
         if self.should_condense(view):

openhands/memory/condenser/impl/browser_output_condenser.py

@@ -17,11 +17,11 @@ def __init__(self, attention_window: int = 1):
         self.attention_window = attention_window
         super().__init__()
 
-    def condense(self, events: list[Event]) -> View | Condensation:
+    def condense(self, view: View) -> View | Condensation:
         """Replace the content of browser observations outside of the attention window with a placeholder."""
         results: list[Event] = []
         cnt: int = 0
-        for event in reversed(events):
+        for event in reversed(view):
             if (
                 isinstance(event, BrowserOutputObservation)
                 and cnt >= self.attention_window

openhands/memory/condenser/impl/no_op_condenser.py

@@ -1,16 +1,15 @@
 from __future__ import annotations
 
 from openhands.core.config.condenser_config import NoOpCondenserConfig
-from openhands.events.event import Event
 from openhands.memory.condenser.condenser import Condensation, Condenser, View
 
 
 class NoOpCondenser(Condenser):
     """A condenser that does nothing to the event sequence."""
 
-    def condense(self, events: list[Event]) -> View | Condensation:
+    def condense(self, view: View) -> View | Condensation:
         """Returns the list of events unchanged."""
-        return View(events=events)
+        return view
 
     @classmethod
     def from_config(cls, config: NoOpCondenserConfig) -> NoOpCondenser:

openhands/memory/condenser/impl/observation_masking_condenser.py

@@ -15,14 +15,11 @@ def __init__(self, attention_window: int = 5):
 
         super().__init__()
 
-    def condense(self, events: list[Event]) -> View | Condensation:
+    def condense(self, view: View) -> View | Condensation:
         """Replace the content of observations outside of the attention window with a placeholder."""
         results: list[Event] = []
-        for i, event in enumerate(events):
-            if (
-                isinstance(event, Observation)
-                and i < len(events) - self.attention_window
-            ):
+        for i, event in enumerate(view):
+            if isinstance(event, Observation) and i < len(view) - self.attention_window:
                 results.append(AgentCondensationObservation('<MASKED>'))
             else:
                 results.append(event)