diff --git a/enterprise/integrations/utils.py b/enterprise/integrations/utils.py index 6ead2e642d..5cdf38666d 100644 --- a/enterprise/integrations/utils.py +++ b/enterprise/integrations/utils.py @@ -20,6 +20,7 @@ from openhands.events.action import ( AgentFinishAction, MessageAction, ) +from openhands.events.event_filter import EventFilter from openhands.events.event_store_abc import EventStoreABC from openhands.events.observation.agent import AgentStateChangedObservation from openhands.integrations.service_types import Repository @@ -203,18 +204,35 @@ def get_summary_for_agent_state( def get_final_agent_observation( event_store: EventStoreABC, ) -> list[AgentStateChangedObservation]: - return event_store.get_matching_events( - source=EventSource.ENVIRONMENT, - event_types=(AgentStateChangedObservation,), - limit=1, - reverse=True, + events = list( + event_store.search_events( + filter=EventFilter( + source=EventSource.ENVIRONMENT, + include_types=(AgentStateChangedObservation,), + ), + limit=1, + reverse=True, + ) ) + result = [e for e in events if isinstance(e, AgentStateChangedObservation)] + assert len(result) == len(events) + return result def get_last_user_msg(event_store: EventStoreABC) -> list[MessageAction]: - return event_store.get_matching_events( - source=EventSource.USER, event_types=(MessageAction,), limit=1, reverse='true' + events = list( + event_store.search_events( + filter=EventFilter( + source=EventSource.USER, + include_types=(MessageAction,), + ), + limit=1, + reverse=True, + ) ) + result = [e for e in events if isinstance(e, MessageAction)] + assert len(result) == len(events) + return result def extract_summary_from_event_store( @@ -226,18 +244,22 @@ def extract_summary_from_event_store( conversation_link = CONVERSATION_URL.format(conversation_id) summary_instruction = get_summary_instruction() - instruction_event: list[MessageAction] = event_store.get_matching_events( - query=json.dumps(summary_instruction), - source=EventSource.USER, - event_types=(MessageAction,), - limit=1, - reverse=True, + instruction_events = list( + event_store.search_events( + filter=EventFilter( + query=json.dumps(summary_instruction), + source=EventSource.USER, + include_types=(MessageAction,), + ), + limit=1, + reverse=True, + ) ) final_agent_observation = get_final_agent_observation(event_store) # Find summary instruction event ID - if len(instruction_event) == 0: + if not instruction_events: logger.warning( 'no_instruction_event_found', extra={'conversation_id': conversation_id} ) @@ -245,19 +267,19 @@ def extract_summary_from_event_store( final_agent_observation, conversation_link ) # Agent did not receive summary instruction - event_id: int = instruction_event[0].id - - agent_messages: list[MessageAction | AgentFinishAction] = ( - event_store.get_matching_events( - start_id=event_id, - source=EventSource.AGENT, - event_types=(MessageAction, AgentFinishAction), - reverse=True, + summary_events = list( + event_store.search_events( + filter=EventFilter( + source=EventSource.AGENT, + include_types=(MessageAction, AgentFinishAction), + ), limit=1, + reverse=True, + start_id=instruction_events[0].id, ) ) - if len(agent_messages) == 0: + if not summary_events: logger.warning( 'no_agent_messages_found', extra={'conversation_id': conversation_id} ) @@ -265,10 +287,11 @@ def extract_summary_from_event_store( final_agent_observation, conversation_link ) # Agent failed to generate summary - summary_event: MessageAction | AgentFinishAction = agent_messages[0] + summary_event = summary_events[0] if isinstance(summary_event, MessageAction): return summary_event.content + assert isinstance(summary_event, AgentFinishAction) return summary_event.final_thought