feat(history): Implement get_messages_by_window in HistoryManager

This commit is contained in:
Yunxiao Xu
2026-02-15 03:20:12 -08:00
parent 68c0985482
commit be687d0fa5
2 changed files with 69 additions and 1 deletions

View File

@@ -185,4 +185,28 @@ class HistoryManager:
# Pre-load plots for each message # Pre-load plots for each message
for m in messages: for m in messages:
_ = m.plots _ = m.plots
return messages return messages
def get_messages_by_window(self, conversation_id: str, window_size: int = 10) -> List[Message]:
"""
Get the last N messages for a conversation, ordered by creation time (ascending).
"""
with self.get_session() as session:
# 1. Get the last N messages in descending order
stmt = (
select(Message)
.where(Message.conversation_id == conversation_id)
.order_by(Message.created_at.desc())
.limit(window_size)
)
result = session.execute(stmt)
messages = list(result.scalars().all())
# 2. Reverse to get ascending chronological order
messages.reverse()
# Pre-load plots if needed (though graph usually only needs text)
for m in messages:
_ = m.plots
return messages

View File

@@ -0,0 +1,44 @@
import pytest
from ea_chatbot.history.manager import HistoryManager
from ea_chatbot.history.models import Base
from sqlalchemy import create_engine
@pytest.fixture
def history_manager():
# Use in-memory SQLite for testing
# Note: HistoryManager creates its own engine if we pass a URL,
# but we need to ensure the schema is created on the underlying DB.
hm = HistoryManager("sqlite:///:memory:")
# Get the engine from the manager's sessionmaker
engine = hm.engine
Base.metadata.create_all(engine)
return hm
def test_get_messages_by_window(history_manager):
"""
Test that get_messages_by_window returns the last N messages
in chronological order.
"""
# 1. Setup a user and conversation
user = history_manager.create_user("test@test.com", "password", "Test User")
conv = history_manager.create_conversation(user.id, "nj", "Test Conv")
# 2. Add 15 messages
for i in range(15):
history_manager.add_message(conv.id, "user" if i % 2 == 0 else "assistant", f"Message {i}")
# 3. Fetch last 10 messages
messages = history_manager.get_messages_by_window(conv.id, window_size=10)
# 4. Assertions
assert len(messages) == 10
# The first message in the list should be "Message 5" (15 - 10 = 5)
assert messages[0].content == "Message 5"
# The last message should be "Message 14"
assert messages[-1].content == "Message 14"
# Ensure they are in ascending order of creation (or at least order of content here)
for i in range(9):
# Extract number from "Message X"
current_val = int(messages[i].content.split(" ")[1])
next_val = int(messages[i+1].content.split(" ")[1])
assert current_val < next_val