Fix get_last_k_turns ordering issues and add robust event sorting

[DennisTraub]

Root Cause

The original implementation had a fundamental misunderstanding of the boto3 list_events API behavior. The code mistakenly assumed that list_events returns events in chronological order, when it actually returns them in reverse chronological order (newest first). This incorrect assumption caused cascading issues throughout MemoryClient.list_events() and MemoryClient.get_last_k_turns(). The problems were never caught because all tests mocked the boto3 list_events behavior based on mistaken assumptions.

Problems Caused by Root Cause

1. Incorrect event ordering: Code assumed chronological order but received reverse chronological
2. Mixed message ordering: Messages within single events are in chronological order, but events themselves were in reverse chronological order, creating inconsistent ordering logic
3. Pagination issues: When using pagination, the reverse chronological order from boto3 meant that older events (which should be processed first) were returned later, causing incorrect turn grouping
4. Test coverage gap: All tests mocked list_events, so the real API behavior was never validated

Solution

1. Fixed MemoryClient's list_events Method

• Added robust timestamp-based sorting in the list_events method to ensure events are always returned in the correct chronological order regardless of what boto3 returns
• Made ordering independent of boto3 behavior by sorting events by eventTimestamp after retrieval
• Fixed pagination handling to work correctly with the new sorting approach

2. Refactored get_last_k_turns Method

• Changed event ordering from REVERSE_CHRONOLOGICAL to CHRONOLOGICAL (now that we guarantee correct order)
• Implemented flatten-then-group approach:
  • Flatten all messages from events into a single list
  • Reverse the message list to process from newest to oldest
  • Group consecutive USER-ASSISTANT message pairs into turns
• Added efficiency improvements by stopping once k turns are collected
• Handles mixed scenarios correctly: Works whether turns are within single events or span multiple events

Key Changes

MemoryClient.list_events:

• Added timestamp-based sorting to ensure chronological order regardless of boto3 behavior
• Added optional parameter order with an Enum EventOrdering.CHRONOLOGICAL|REVERSE_CHRONOLOGICAL, with CHRONOLOGICAL as default
• Made the method robust against irregularities or future changes in boto3 API ordering

MemoryClient.get_last_k_turns:

• Replaced incremental turn building with flatten-then-group approach to maintain ordering inside and across events
• Added proper message ordering logic for all scenarios
• Maintained efficiency through early termination when k turns are collected

Test Coverage

• Added comprehensive test case covering mixed scenarios (turns within single events and across multiple events)
• Verified correct chronological ordering of messages within turns
• Fixed test cases for event pagination
• Note: Tests still mock boto3, but now the implementation is robust against real API behavior

Impact

This fix ensures that list_events and get_last_k_turns works correctly with the actual boto3 API behavior, handles pagination properly, and provides consistent results regardless of how events are distributed across API responses. The robust sorting approach also makes the code more resilient to potential future changes in the boto3 API.

[bergjaak]

the ordering comes from the server, so we can't specify this from the client

[DennisTraub]

Right, but can change the order once you've received the events from the server. That's what I added to the function as part of the overall bug fix.

[bergjaak]

I think that would be confusing, though. That we say "you are paginating from the oldest event first" and then the 2nd page has a list of events that are older than the first page