api/core/memory/token_buffer_memory.py
Rama main · réplica de lectura del repositorio original.
from collections import defaultdict
from collections.abc import Sequence
from dataclasses import dataclass
from sqlalchemy import select
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
from core.model_manager import ModelInstance
from core.prompt.utils.extract_thread_messages import extract_thread_messages
from extensions.ext_database import db
from factories import file_factory
from graphon.file import FileTransferMethod, FileType, FileUploadConfig, file_manager
from graphon.model_runtime.entities import (
AssistantPromptMessage,
ImagePromptMessageContent,
PromptMessage,
PromptMessageRole,
TextPromptMessageContent,
UserPromptMessage,
)
from graphon.model_runtime.entities.message_entities import PromptMessageContentUnionTypes
from models.agent_config_entities import AgentSoulConfig
from models.enums import ConversationFromSource
from models.model import App, AppMode, Conversation, Message, MessageFile
from models.workflow import Workflow
from repositories.api_workflow_run_repository import APIWorkflowRunRepository
from repositories.factory import DifyAPIRepositoryFactory
from services.agent.runtime_config_service import AgentRuntimeConfigService
_file_access_controller = DatabaseFileAccessController()
@dataclass(frozen=True)
class HistoryFile:
"""Persisted attachment reference; metadata and content are resolved after loading history."""
id: str
type: FileType
transfer_method: FileTransferMethod
url: str | None
upload_file_id: str | None
@classmethod
def from_record(cls, record: MessageFile) -> "HistoryFile":
return cls(record.id, record.type, record.transfer_method, record.url, record.upload_file_id)
@dataclass(frozen=True)
class HistoryPrompt:
text: str
is_user_message: bool
files: tuple[HistoryFile, ...]
tenant_id: str | None
image_detail: ImagePromptMessageContent.DETAIL
def to_prompt_message(self) -> PromptMessage:
contents: list[PromptMessageContentUnionTypes] = []
if self.files and self.tenant_id is not None:
for reference in self.files:
file = file_factory.build_from_mapping(
mapping={
"id": reference.id,
"type": reference.type,
"transfer_method": reference.transfer_method,
"url": reference.url,
"tool_file_id"
if reference.transfer_method == FileTransferMethod.TOOL_FILE
else "upload_file_id": reference.upload_file_id,
},
tenant_id=self.tenant_id,
access_controller=_file_access_controller,
)
contents.append(file_manager.to_prompt_message_content(file, image_detail_config=self.image_detail))
if contents:
contents.append(TextPromptMessageContent(data=self.text))
content = contents or self.text
if self.is_user_message:
return UserPromptMessage(content=content)
return AssistantPromptMessage(content=content)
@dataclass(frozen=True)
class PreparedHistory:
"""Detached history that can restore attachments and count tokens without its read session.
File factories retain ownership of their short metadata lookups. Rendering may
perform storage/network I/O, so callers should first close the history session.
"""
prompts: tuple[HistoryPrompt, ...]
def get_prompt_messages(self, *, model_instance: ModelInstance, max_token_limit: int) -> Sequence[PromptMessage]:
prompt_messages = [prompt.to_prompt_message() for prompt in self.prompts]
if not prompt_messages:
return []
curr_message_tokens = model_instance.get_llm_num_tokens(prompt_messages)
while curr_message_tokens > max_token_limit and len(prompt_messages) > 1:
prompt_messages.pop(0)
curr_message_tokens = model_instance.get_llm_num_tokens(prompt_messages)
return prompt_messages
def get_prompt_text(
self,
*,
model_instance: ModelInstance,
max_token_limit: int,
human_prefix: str = "Human",
ai_prefix: str = "Assistant",
) -> str:
prompt_messages = self.get_prompt_messages(model_instance=model_instance, max_token_limit=max_token_limit)
return _prompt_messages_to_text(prompt_messages, human_prefix=human_prefix, ai_prefix=ai_prefix)
def _prompt_messages_to_text(prompt_messages: Sequence[PromptMessage], *, human_prefix: str, ai_prefix: str) -> str:
string_messages = []
for message in prompt_messages:
if message.role == PromptMessageRole.USER:
role = human_prefix
elif message.role == PromptMessageRole.ASSISTANT:
role = ai_prefix
else:
continue
if isinstance(message.content, list):
inner_msg = ""
for content in message.content:
match content:
case TextPromptMessageContent():
inner_msg += f"{content.data}\n"
case ImagePromptMessageContent():
inner_msg += "[image]\n"
string_messages.append(f"{role}: {inner_msg.strip()}")
else:
string_messages.append(f"{role}: {message.content}")
return "\n".join(string_messages)
class TokenBufferMemory:
def __init__(
self,
conversation: Conversation,
model_instance: ModelInstance,
) -> None:
self.conversation = conversation
self.model_instance = model_instance
@staticmethod
def _file_config(
*,
conversation: Conversation,
app_record: App | None,
message: Message,
session: Session,
workflow_run_repo: APIWorkflowRunRepository | None,
) -> FileUploadConfig | None:
"""Read attachment settings using the caller's session without file I/O.
Agent and workflow modes require the owning app. Agent settings come
from the conversation's Soul, never the legacy model configuration.
Disabled uploads return None; missing required records raise ValueError
and unsupported modes raise AssertionError. No records are changed.
"""
match conversation.mode:
case AppMode.AGENT:
if not app_record:
raise ValueError("App not found for conversation")
soul = AgentRuntimeConfigService(session).resolve_conversation_soul(
app_model=app_record,
conversation=conversation,
account_id=conversation.from_account_id,
use_debug_draft=conversation.from_source == ConversationFromSource.CONSOLE,
)
features = (soul or AgentSoulConfig()).app_features.model_dump(mode="json", exclude_none=True)
return FileUploadConfigManager.convert(features)
case AppMode.AGENT_CHAT | AppMode.COMPLETION | AppMode.CHAT:
return FileUploadConfigManager.convert(conversation.model_config_with_session(session=session))
case AppMode.ADVANCED_CHAT | AppMode.WORKFLOW:
if not app_record:
raise ValueError("App not found for conversation")
if not message.workflow_run_id:
raise ValueError("Workflow run ID not found")
assert workflow_run_repo is not None
workflow_run = workflow_run_repo.get_workflow_run_by_id(
tenant_id=app_record.tenant_id, app_id=app_record.id, run_id=message.workflow_run_id
)
if not workflow_run:
raise ValueError(f"Workflow run not found: {message.workflow_run_id}")
workflow = session.scalar(select(Workflow).where(Workflow.id == workflow_run.workflow_id))
if not workflow:
raise ValueError(f"Workflow not found: {workflow_run.workflow_id}")
return FileUploadConfigManager.convert(workflow.normalized_features_dict, is_vision=False)
case _:
raise AssertionError(f"Invalid app mode: {conversation.mode}")
@staticmethod
def _prepare_prompt(
*,
message_files: Sequence[MessageFile],
text_content: str,
app_record: App | None,
is_user_message: bool,
file_config: FileUploadConfig | None,
) -> HistoryPrompt:
detail = ImagePromptMessageContent.DETAIL.HIGH
files: tuple[HistoryFile, ...] = ()
if file_config and app_record:
files = tuple(HistoryFile.from_record(message_file) for message_file in message_files)
if file_config.image_config and file_config.image_config.detail:
detail = file_config.image_config.detail
return HistoryPrompt(
text=text_content,
is_user_message=is_user_message,
files=files,
tenant_id=app_record.tenant_id if app_record else None,
image_detail=detail,
)
@classmethod
def load_history(
cls,
*,
conversation: Conversation,
app_record: App | None,
session: Session,
message_limit: int | None,
workflow_run_repo: APIWorkflowRunRepository | None = None,
) -> PreparedHistory:
"""Read ordered history and attachment configuration without provider or file I/O.
``None`` for message_limit uses the existing 500-message ceiling. A missing
app keeps text-only legacy history available; Agent and workflow attachments still
require the app. An omitted workflow repository uses the configured factory.
"""
# fetch limited messages, and return reversed
# 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
msg_limit_stmt = stmt.limit(message_limit)
messages = session.scalars(msg_limit_stmt).all()
# instead of all messages from the conversation, we only need to extract messages
# that belong to the thread of last message
thread_messages = extract_thread_messages(messages)
# for newly created message, its answer is temporarily empty, we don't need to add it to memory
if thread_messages and not thread_messages[0].answer and thread_messages[0].answer_tokens == 0:
thread_messages.pop(0)
messages = list(reversed(thread_messages))
# Batch-load message files for the whole thread to avoid an N+1 query pattern.
# Previously each message issued two MessageFile queries (user + assistant),
# i.e. 2N+1 round-trips for N messages. We now use two batched queries keyed by
# message_id, preserving the exact filter semantics (user files include rows
# whose belongs_to is NULL).
message_ids = [message.id for message in messages]
user_files_by_message: dict[str, list[MessageFile]] = defaultdict(list)
assistant_files_by_message: dict[str, list[MessageFile]] = defaultdict(list)
if message_ids:
for message_file in session.scalars(
select(MessageFile).where(
MessageFile.message_id.in_(message_ids),
(MessageFile.belongs_to == "user") | (MessageFile.belongs_to.is_(None)),
)
).all():
user_files_by_message[message_file.message_id].append(message_file)
for message_file in session.scalars(
select(MessageFile).where(
MessageFile.message_id.in_(message_ids),
MessageFile.belongs_to == "assistant",
)
).all():
assistant_files_by_message[message_file.message_id].append(message_file)
prompts: list[HistoryPrompt] = []
for message in messages:
user_files = user_files_by_message.get(message.id, [])
assistant_files = assistant_files_by_message.get(message.id, [])
if (
(user_files or assistant_files)
and conversation.mode in {AppMode.ADVANCED_CHAT, AppMode.WORKFLOW}
and app_record is not None
and message.workflow_run_id
and workflow_run_repo is None
):
session_maker = sessionmaker(bind=session.get_bind(), expire_on_commit=False)
workflow_run_repo = DifyAPIRepositoryFactory.create_api_workflow_run_repository(session_maker)
file_config = (
cls._file_config(
conversation=conversation,
app_record=app_record,
message=message,
session=session,
workflow_run_repo=workflow_run_repo,
)
if user_files or assistant_files
else None
)
for files, text_content, is_user_message in (
(user_files, message.query, True),
(assistant_files, message.answer, False),
):
prompts.append(
cls._prepare_prompt(
message_files=files,
text_content=text_content,
app_record=app_record,
is_user_message=is_user_message,
file_config=file_config,
)
)
return PreparedHistory(prompts=tuple(prompts))
def get_history_prompt_messages(
self, max_token_limit: int = 2000, message_limit: int | None = None
) -> Sequence[PromptMessage]:
session = db.session()
history = self.load_history(
conversation=self.conversation,
app_record=self.conversation.app(session=session),
session=session,
message_limit=message_limit,
)
return history.get_prompt_messages(model_instance=self.model_instance, max_token_limit=max_token_limit)
def get_history_prompt_text(
self,
human_prefix: str = "Human",
ai_prefix: str = "Assistant",
max_token_limit: int = 2000,
message_limit: int | None = None,
) -> str:
"""
Get history prompt text.
:param human_prefix: human prefix
:param ai_prefix: ai prefix
:param max_token_limit: max token limit
:param message_limit: message limit
:return:
"""
prompt_messages = self.get_history_prompt_messages(max_token_limit=max_token_limit, message_limit=message_limit)
return _prompt_messages_to_text(prompt_messages, human_prefix=human_prefix, ai_prefix=ai_prefix)