Skip to content

Commit

Permalink
pass lint and truncation tests
Browse files Browse the repository at this point in the history
  • Loading branch information
ryx2 committed Nov 16, 2024
1 parent 81ddb82 commit ba9bc74
Show file tree
Hide file tree
Showing 2 changed files with 17 additions and 5 deletions.
21 changes: 16 additions & 5 deletions openhands/controller/agent_controller.py
Original file line number Diff line number Diff line change
Expand Up @@ -228,11 +228,11 @@ async def on_event(self, event: Event):
# Create temporary history with new event
temp_history = self.state.history + [event]
token_count = self.agent.llm.get_token_count(temp_history)

if token_count > self.context_window:
# Truncate existing history before adding new event
self.state.history = self._apply_conversation_window(self.state.history)

# Now add the new event
self.state.history.append(event)

Expand Down Expand Up @@ -839,6 +839,10 @@ def _apply_conversation_window(self, events: list[Event]) -> list[Event]:
None,
)

# Always set start_id to first user message id if found, regardless of truncation
if first_user_msg:
self.state.start_id = first_user_msg.id

# cut in half
mid_point = max(1, len(events) // 2)
kept_events = events[mid_point:]
Expand Down Expand Up @@ -889,9 +893,16 @@ def _apply_conversation_window(self, events: list[Event]) -> list[Event]:
if first_user_msg and first_user_msg not in kept_events:
kept_events = [first_user_msg] + kept_events

# start_id points to first user message
if first_user_msg:
self.state.start_id = first_user_msg.id
# Verify truncated history fits in context window
while (
self.context_window is not None
and self.agent.llm.get_token_count(kept_events) > self.context_window
):
# Need to truncate more - remove oldest non-first-message event
if len(kept_events) > 2: # Keep at least first message and one more event
kept_events.pop(1)
else:
break

return kept_events

Expand Down
1 change: 1 addition & 0 deletions tests/unit/test_truncation.py
Original file line number Diff line number Diff line change
Expand Up @@ -220,6 +220,7 @@ def test_context_window_parameter_truncation(self, mock_event_stream, mock_agent

# Set initial history
controller.state.history = events[:3] # Start with a few events
controller.state.start_id = first_msg._id # Explicitly set start_id
initial_history_len = len(controller.state.history)

# Add a new event that should trigger truncation due to token count
Expand Down

0 comments on commit ba9bc74

Please sign in to comment.