feat: Combining condensers (#7867)

Co-authored-by: Calvin Smith <calvin@all-hands.dev>

openhands/core/config/condenser_config.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from typing import Literal, cast
 
 from pydantic import BaseModel, Field, ValidationError
@@ -153,6 +155,21 @@ class StructuredSummaryCondenserConfig(BaseModel):
     model_config = {'extra': 'forbid'}
 
 
+class CondenserPipelineConfig(BaseModel):
+    """Configuration for the CondenserPipeline.
+
+    Not currently supported by the TOML or ENV_VAR configuration strategies.
+    """
+
+    type: Literal['pipeline'] = Field('pipeline')
+    condensers: list[CondenserConfig] = Field(
+        default_factory=list,
+        description='List of condenser configurations to be used in the pipeline.',
+    )
+
+    model_config = {'extra': 'forbid'}
+
+
 # Type alias for convenience
 CondenserConfig = (
     NoOpCondenserConfig
@@ -163,6 +180,7 @@ CondenserConfig = (
     | AmortizedForgettingCondenserConfig
     | LLMAttentionCondenserConfig
     | StructuredSummaryCondenserConfig
+    | CondenserPipelineConfig
 )
 
 

openhands/memory/condenser/impl/__init__.py

@@ -15,6 +15,7 @@ from openhands.memory.condenser.impl.no_op_condenser import NoOpCondenser
 from openhands.memory.condenser.impl.observation_masking_condenser import (
     ObservationMaskingCondenser,
 )
+from openhands.memory.condenser.impl.pipeline import CondenserPipeline
 from openhands.memory.condenser.impl.recent_events_condenser import (
     RecentEventsCondenser,
 )
@@ -32,4 +33,5 @@ __all__ = [
     'BrowserOutputCondenser',
     'RecentEventsCondenser',
     'StructuredSummaryCondenser',
+    'CondenserPipeline',
 ]

openhands/memory/condenser/impl/pipeline.py

@@ -0,0 +1,47 @@
+from __future__ import annotations
+
+from contextlib import contextmanager
+
+from openhands.controller.state.state import State
+from openhands.core.config.condenser_config import CondenserPipelineConfig
+from openhands.memory.condenser.condenser import Condensation, Condenser
+from openhands.memory.view import View
+
+
+class CondenserPipeline(Condenser):
+    """Combines multiple condensers into a single condenser.
+
+    This is useful for creating a pipeline of condensers that can be chained together to achieve very specific condensation aims. Each condenser is run in sequence, passing the output view of one to the next, until we reach the end or a `CondensationAction` is returned instead.
+    """
+
+    def __init__(self, *condenser: Condenser) -> None:
+        self.condensers = list(condenser)
+        super().__init__()
+
+    @contextmanager
+    def metadata_batch(self, state: State):
+        try:
+            yield
+        finally:
+            # The parent class assumes the metadata is stored in the "calling
+            # condenser" -- since we're not threading a State through to each
+            # step in the pipeline, we need to walk back through the pipeline
+            # and manually collect the relevant metadata.
+            for condenser in self.condensers:
+                condenser.write_metadata(state)
+
+    def condense(self, view: View) -> View | Condensation:
+        result: View | Condensation = view
+        for condenser in self.condensers:
+            result = condenser.condense(result)
+            if isinstance(result, Condensation):
+                break
+        return result
+
+    @classmethod
+    def from_config(cls, config: CondenserPipelineConfig) -> CondenserPipeline:
+        condensers = [Condenser.from_config(c) for c in config.condensers]
+        return CondenserPipeline(*condensers)
+
+
+CondenserPipeline.register_config(CondenserPipelineConfig)

openhands/server/session/session.py

@@ -7,7 +7,11 @@ import socketio
 
 from openhands.controller.agent import Agent
 from openhands.core.config import AppConfig
-from openhands.core.config.condenser_config import LLMSummarizingCondenserConfig
+from openhands.core.config.condenser_config import (
+    BrowserOutputCondenserConfig,
+    CondenserPipelineConfig,
+    LLMSummarizingCondenserConfig,
+)
 from openhands.core.logger import OpenHandsLoggerAdapter
 from openhands.core.schema import AgentState
 from openhands.events.action import MessageAction, NullAction
@@ -126,8 +130,18 @@ class Session:
         agent_config = self.config.get_agent_config(agent_cls)
 
         if settings.enable_default_condenser:
-            default_condenser_config = LLMSummarizingCondenserConfig(
-                llm_config=llm.config, keep_first=3, max_size=80
+            # Default condenser chains a condenser that limits browser the total
+            # size of browser observations with a condenser that limits the size
+            # of the view given to the LLM. The order matters: with the browser
+            # output first, the summarizer will only see the most recent browser
+            # output, which should keep the summarization cost down.
+            default_condenser_config = CondenserPipelineConfig(
+                condensers=[
+                    BrowserOutputCondenserConfig(),
+                    LLMSummarizingCondenserConfig(
+                        llm_config=llm.config, keep_first=3, max_size=80
+                    ),
+                ]
             )
 
             self.logger.info(f'Enabling default condenser: {default_condenser_config}')

tests/unit/test_condenser.py

@@ -8,6 +8,7 @@ from openhands.controller.state.state import State
 from openhands.core.config.condenser_config import (
     AmortizedForgettingCondenserConfig,
     BrowserOutputCondenserConfig,
+    CondenserPipelineConfig,
     LLMAttentionCondenserConfig,
     LLMSummarizingCondenserConfig,
     NoOpCondenserConfig,
@@ -17,6 +18,7 @@ from openhands.core.config.condenser_config import (
 )
 from openhands.core.config.llm_config import LLMConfig
 from openhands.core.message import Message, TextContent
+from openhands.core.schema.action import ActionType
 from openhands.events.event import Event, EventSource
 from openhands.events.observation import BrowserOutputObservation
 from openhands.events.observation.agent import AgentCondensationObservation
@@ -35,6 +37,7 @@ from openhands.memory.condenser.impl import (
     RecentEventsCondenser,
     StructuredSummaryCondenser,
 )
+from openhands.memory.condenser.impl.pipeline import CondenserPipeline
 
 
 def create_test_event(
@@ -117,6 +120,11 @@ class RollingCondenserTestHarness:
         state = State()
 
         for event in events:
+            # Set the event's ID -- this is normally done by the event stream,
+            # but this harness is intended to act as a testing stand-in.
+            if not hasattr(event, '_id'):
+                event._id = len(state.history)
+
             state.history.append(event)
             for callback in self.callbacks:
                 callback(state.history)
@@ -705,3 +713,66 @@ def test_structured_summary_condenser_keeps_first_and_summary_events(mock_llm):
         # If we've condensed, ensure that the summary event is present
         if i > max_size:
             assert isinstance(view[keep_first], AgentCondensationObservation)
+
+
+def test_condenser_pipeline_from_config():
+    """Test that CondenserPipeline condensers can be created from configuration objects."""
+    config = CondenserPipelineConfig(
+        condensers=[
+            NoOpCondenserConfig(),
+            BrowserOutputCondenserConfig(attention_window=2),
+            LLMSummarizingCondenserConfig(
+                max_size=50,
+                keep_first=10,
+                llm_config=LLMConfig(model='gpt-4o', api_key='test_key'),
+            ),
+        ]
+    )
+    condenser = Condenser.from_config(config)
+
+    assert isinstance(condenser, CondenserPipeline)
+    assert len(condenser.condensers) == 3
+    assert isinstance(condenser.condensers[0], NoOpCondenser)
+    assert isinstance(condenser.condensers[1], BrowserOutputCondenser)
+    assert isinstance(condenser.condensers[2], LLMSummarizingCondenser)
+
+
+def test_condenser_pipeline_chains_sub_condensers():
+    """Test that the CondenserPipeline chains sub-condensers and combines their behavior."""
+    MAX_SIZE = 10
+    ATTENTION_WINDOW = 2
+    NUMBER_OF_CONDENSATIONS = 3
+
+    condenser = CondenserPipeline(
+        AmortizedForgettingCondenser(max_size=MAX_SIZE),
+        BrowserOutputCondenser(attention_window=ATTENTION_WINDOW),
+    )
+
+    harness = RollingCondenserTestHarness(condenser)
+    events = [
+        BrowserOutputObservation(
+            f'Observation {i}', url='', trigger_by_action=ActionType.BROWSE
+        )
+        if i % 3 == 0
+        else create_test_event(f'Event {i}')
+        for i in range(0, MAX_SIZE * NUMBER_OF_CONDENSATIONS)
+    ]
+
+    for index, view in enumerate(harness.views(events)):
+        # The amortized forgetting condenser is responsible for keeping the size
+        # bounded despite the large number of events.
+        assert len(view) == harness.expected_size(index, MAX_SIZE)
+
+        # The browser output condenser should mask out the content of all the
+        # browser observations outside the attention window (which is relative
+        # to the number of browser outputs in the view, not the whole view or
+        # the event stream).
+        browser_outputs = [
+            event for event in view if isinstance(event, BrowserOutputObservation)
+        ]
+
+        for event in browser_outputs[:-ATTENTION_WINDOW]:
+            assert 'Content Omitted' in str(event)
+
+        for event in browser_outputs[-ATTENTION_WINDOW:]:
+            assert 'Content Omitted' not in str(event)