diff --git a/api/core/memory/token_buffer_memory.py b/api/core/memory/token_buffer_memory.py index bb0d50357ff..ef1a9416803 100644 --- a/api/core/memory/token_buffer_memory.py +++ b/api/core/memory/token_buffer_memory.py @@ -3,7 +3,7 @@ from collections.abc import Sequence from dataclasses import dataclass from sqlalchemy import select -from sqlalchemy.orm import Session, sessionmaker +from sqlalchemy.orm import Session, load_only, sessionmaker from core.app.app_config.features.file_upload.manager import FileUploadConfigManager from core.app.file_access import DatabaseFileAccessController @@ -215,7 +215,24 @@ class TokenBufferMemory: require the app. An omitted workflow repository uses the configured factory. """ # fetch limited messages, and return reversed - stmt = select(Message).where(Message.conversation_id == conversation.id).order_by(Message.created_at.desc()) + # Load only the columns used to rebuild history. Message.message stores the full prompt of each turn, + # earlier history included, so loading it makes memory grow quadratically with conversation length. + stmt = ( + select(Message) + .options( + load_only( + Message.id, + Message.parent_message_id, + Message.query, + Message.answer, + Message.answer_tokens, + Message.workflow_run_id, + Message.created_at, + ) + ) + .where(Message.conversation_id == conversation.id) + .order_by(Message.created_at.desc()) + ) message_limit = min(message_limit, 500) if message_limit and message_limit > 0 else 500 diff --git a/api/core/prompt/utils/get_thread_messages_length.py b/api/core/prompt/utils/get_thread_messages_length.py index 6e4b3b1e550..6052c841342 100644 --- a/api/core/prompt/utils/get_thread_messages_length.py +++ b/api/core/prompt/utils/get_thread_messages_length.py @@ -1,5 +1,5 @@ from sqlalchemy import select -from sqlalchemy.orm import Session +from sqlalchemy.orm import Session, load_only from core.prompt.utils.extract_thread_messages import extract_thread_messages from models.model import Message @@ -9,8 +9,14 @@ def get_thread_messages_length(conversation_id: str, *, session: Session) -> int """ Get the number of thread messages based on the parent message id. """ - # Fetch all messages related to the conversation - stmt = select(Message).where(Message.conversation_id == conversation_id).order_by(Message.created_at.desc()) + # Fetch all messages related to the conversation, loading only the columns needed to walk the thread. + # Message.message stores the full prompt of each turn, so loading it grows quadratically with the conversation. + stmt = ( + select(Message) + .options(load_only(Message.id, Message.parent_message_id, Message.answer)) + .where(Message.conversation_id == conversation_id) + .order_by(Message.created_at.desc()) + ) messages = session.scalars(stmt).all() diff --git a/api/tests/unit_tests/core/memory/test_token_buffer_memory.py b/api/tests/unit_tests/core/memory/test_token_buffer_memory.py index 39c4ac7f14b..7a11d716b5f 100644 --- a/api/tests/unit_tests/core/memory/test_token_buffer_memory.py +++ b/api/tests/unit_tests/core/memory/test_token_buffer_memory.py @@ -1,6 +1,7 @@ """Comprehensive SQLite-backed tests for token-buffer memory.""" import json +import re from collections.abc import Iterator, Mapping, Sequence from dataclasses import dataclass from datetime import UTC, datetime, timedelta @@ -146,6 +147,7 @@ def _persist_message( answer_tokens: int = 5, created_at: datetime | None = None, workflow_run_id: str | None = None, + parent_message_id: str | None = None, ) -> Message: message = Message( id=str(uuid4()), @@ -154,6 +156,7 @@ def _persist_message( _inputs={}, query=query, message={}, + parent_message_id=parent_message_id, message_unit_price=Decimal(0), answer=answer, answer_tokens=answer_tokens, @@ -375,6 +378,38 @@ class TestGetHistoryPromptMessages: assert sum("FROM apps" in sql for sql in selects) == 1 assert len(result) == 10 + def test_history_query_skips_saved_prompt_column(self, database: Database) -> None: + """Message.message stores each turn's full prompt, so history must neither select nor lazy-load it.""" + mem = self._make_memory(database) + base_time = datetime.now(UTC).replace(tzinfo=None) + parent_message_id: str | None = None + for index in range(3): + parent_message_id = _persist_message( + database, + mem.conversation.id, + query=f"query-{index}", + answer=f"answer-{index}", + created_at=base_time + timedelta(seconds=index), + parent_message_id=parent_message_id, + ).id + # Reload rows from the database instead of reusing fully loaded instances from the identity map. + database.session.expunge_all() + before = len(database.statements) + + result = mem.get_history_prompt_messages() + + message_selects = [sql for sql, _ in database.statements[before:] if "FROM messages" in sql] + assert len(message_selects) == 1 + assert re.search(r"\bmessages\.message\b", message_selects[0]) is None + assert [prompt.content for prompt in result] == [ + "query-0", + "answer-0", + "query-1", + "answer-1", + "query-2", + "answer-2", + ] + @pytest.mark.parametrize( ("token_values", "max_token_limit", "expected_length"), [ diff --git a/api/tests/unit_tests/core/prompt/test_extract_thread_messages.py b/api/tests/unit_tests/core/prompt/test_extract_thread_messages.py index a9438e6aa18..da6c7c64981 100644 --- a/api/tests/unit_tests/core/prompt/test_extract_thread_messages.py +++ b/api/tests/unit_tests/core/prompt/test_extract_thread_messages.py @@ -1,8 +1,10 @@ +import re from datetime import datetime, timedelta from decimal import Decimal from uuid import uuid4 import pytest +from sqlalchemy import Engine, event from sqlalchemy.orm import Session from constants import UUID_NIL @@ -198,3 +200,45 @@ def test_get_thread_messages_length_keeps_non_empty_latest_answer(sqlite_session length = get_thread_messages_length("conversation-2", session=sqlite_session) assert length == 2 + + +def test_get_thread_messages_length_skips_saved_prompt_column(sqlite_engine: Engine, sqlite_session: Session): + """Message.message stores each turn's full prompt, so counting the thread must neither select nor lazy-load it.""" + id1, id2 = str(uuid4()), str(uuid4()) + now = datetime.now() + sqlite_session.add_all( + [ + _persisted_message( + message_id=id2, + conversation_id="conversation-3", + parent_message_id=id1, + answer="latest-answer", + created_at=now, + ), + _persisted_message( + message_id=id1, + conversation_id="conversation-3", + parent_message_id=UUID_NIL, + answer="older-answer", + created_at=now - timedelta(seconds=1), + ), + ] + ) + sqlite_session.commit() + # Reload rows from the database instead of reusing fully loaded instances from the identity map. + sqlite_session.expunge_all() + statements: list[str] = [] + + def record_statement(_connection, _cursor, statement, _parameters, _context, _executemany) -> None: + statements.append(statement) + + event.listen(sqlite_engine, "before_cursor_execute", record_statement) + try: + length = get_thread_messages_length("conversation-3", session=sqlite_session) + finally: + event.remove(sqlite_engine, "before_cursor_execute", record_statement) + + assert length == 2 + message_selects = [sql for sql in statements if "FROM messages" in sql] + assert len(message_selects) == 1 + assert re.search(r"\bmessages\.message\b", message_selects[0]) is None