feat(history): Implement get_messages_by_window in HistoryManager
This commit is contained in:
@@ -186,3 +186,27 @@ class HistoryManager:
|
||||
for m in messages:
|
||||
_ = m.plots
|
||||
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
|
||||
|
||||
44
backend/tests/test_history_window.py
Normal file
44
backend/tests/test_history_window.py
Normal 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
|
||||
Reference in New Issue
Block a user