(fix): Condensation events to reconstruct contexts added to event stream (#7353)

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

openhands/memory/condenser/impl/llm_attention_condenser.py

@@ -4,9 +4,13 @@ from litellm import supports_response_schema
 from pydantic import BaseModel
 
 from openhands.core.config.condenser_config import LLMAttentionCondenserConfig
-from openhands.events.event import Event
+from openhands.events.action.agent import CondensationAction
 from openhands.llm.llm import LLM
-from openhands.memory.condenser.condenser import RollingCondenser
+from openhands.memory.condenser.condenser import (
+    Condensation,
+    RollingCondenser,
+    View,
+)
 
 
 class ImportantEventSelection(BaseModel):
@@ -43,15 +47,11 @@ class LLMAttentionCondenser(RollingCondenser):
 
         super().__init__()
 
-    def condense(self, events: list[Event]) -> list[Event]:
-        """If the history is too long, use an LLM to select the most important events."""
-        if len(events) <= self.max_size:
-            return events
-
+    def get_condensation(self, view: View) -> Condensation:
         target_size = self.max_size // 2
-        head = events[: self.keep_first]
+        head_event_ids = [event.id for event in view.events[: self.keep_first]]
 
-        events_from_tail = target_size - len(head)
+        events_from_tail = target_size - len(head_event_ids)
 
         message: str = """You will be given a list of actions, observations, and thoughts from a coding agent.
         Each item in the list has an identifier. Please sort the identifiers in order of how important the
@@ -66,7 +66,7 @@ class LLMAttentionCondenser(RollingCondenser):
                         'content': f'<ID>{e.id}</ID>\n<CONTENT>{e.message}</CONTENT>',
                         'role': 'user',
                     }
-                    for e in events
+                    for e in view
                 ],
             ],
             response_format={
@@ -82,27 +82,35 @@ class LLMAttentionCondenser(RollingCondenser):
             response.choices[0].message.content
         ).ids
 
-        self.add_metadata('all_event_ids', [event.id for event in events])
-        self.add_metadata('response_ids', response_ids)
         self.add_metadata('metrics', self.llm.metrics.get())
 
         # Filter out any IDs from the head and trim the results down
-        head_ids = [event.id for event in head]
         response_ids = [
-            response_id for response_id in response_ids if response_id not in head_ids
+            response_id
+            for response_id in response_ids
+            if response_id not in head_event_ids
         ][:events_from_tail]
 
         # If the response IDs aren't _long_ enough, iterate backwards through the events and add any unfound IDs to the list.
-        for event in reversed(events):
+        for event in reversed(view):
             if len(response_ids) >= events_from_tail:
                 break
             if event.id not in response_ids:
                 response_ids.append(event.id)
 
-        # Grab the events associated with the response IDs
-        tail = [event for event in events if event.id in response_ids]
+        # Now that we've found the right number of events to keep, convert this into a list of events to forget.
+        event = CondensationAction(
+            forgotten_event_ids=[
+                event.id
+                for event in view
+                if event.id not in response_ids and event.id not in head_event_ids
+            ],
+        )
 
-        return head + tail
+        return Condensation(action=event)
+
+    def should_condense(self, view: View) -> bool:
+        return len(view) > self.max_size
 
     @classmethod
     def from_config(cls, config: LLMAttentionCondenserConfig) -> LLMAttentionCondenser: