diff --git a/backend/src/ea_chatbot/history/manager.py b/backend/src/ea_chatbot/history/manager.py index a11b25d..3adb1d3 100644 --- a/backend/src/ea_chatbot/history/manager.py +++ b/backend/src/ea_chatbot/history/manager.py @@ -185,4 +185,28 @@ class HistoryManager: # Pre-load plots for each message for m in messages: _ = m.plots - return messages \ No newline at end of file + 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 diff --git a/backend/tests/test_history_window.py b/backend/tests/test_history_window.py new file mode 100644 index 0000000..d356e95 --- /dev/null +++ b/backend/tests/test_history_window.py @@ -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