From 929d9e0b3f87aabf742d5ba5a282397b3d76cc87 Mon Sep 17 00:00:00 2001 From: QIN2DIM <62018067+QIN2DIM@users.noreply.github.com> Date: Fri, 29 Aug 2025 15:19:55 +0800 Subject: [PATCH] feat(api): maintain assistant content parts and file handling in advanced chat (#24663) Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com> --- .../advanced_chat/generate_task_pipeline.py | 12 +- api/core/memory/token_buffer_memory.py | 137 ++++++++++++------ api/factories/file_factory.py | 13 +- 3 files changed, 117 insertions(+), 45 deletions(-) diff --git a/api/core/app/apps/advanced_chat/generate_task_pipeline.py b/api/core/app/apps/advanced_chat/generate_task_pipeline.py index 347fed4a1..a61bba512 100644 --- a/api/core/app/apps/advanced_chat/generate_task_pipeline.py +++ b/api/core/app/apps/advanced_chat/generate_task_pipeline.py @@ -1,4 +1,5 @@ import logging +import re import time from collections.abc import Callable, Generator, Mapping from contextlib import contextmanager @@ -373,7 +374,7 @@ class AdvancedChatAppGenerateTaskPipeline: ) -> Generator[StreamResponse, None, None]: """Handle node succeeded events.""" # Record files if it's an answer node or end node - if event.node_type in [NodeType.ANSWER, NodeType.END]: + if event.node_type in [NodeType.ANSWER, NodeType.END, NodeType.LLM]: self._recorded_files.extend( self._workflow_response_converter.fetch_files_from_node_outputs(event.outputs or {}) ) @@ -896,7 +897,14 @@ class AdvancedChatAppGenerateTaskPipeline: def _save_message(self, *, session: Session, graph_runtime_state: Optional[GraphRuntimeState] = None) -> None: message = self._get_message(session=session) - message.answer = self._task_state.answer + + # If there are assistant files, remove markdown image links from answer + answer_text = self._task_state.answer + if self._recorded_files: + # Remove markdown image links since we're storing files separately + answer_text = re.sub(r"!\[.*?\]\(.*?\)", "", answer_text).strip() + + message.answer = answer_text message.updated_at = naive_utc_now() message.provider_response_latency = time.perf_counter() - self._base_task_pipeline._start_at message.message_metadata = self._task_state.metadata.model_dump_json() diff --git a/api/core/memory/token_buffer_memory.py b/api/core/memory/token_buffer_memory.py index 2a76b1f41..36f8c606b 100644 --- a/api/core/memory/token_buffer_memory.py +++ b/api/core/memory/token_buffer_memory.py @@ -31,6 +31,65 @@ class TokenBufferMemory: self.conversation = conversation self.model_instance = model_instance + def _build_prompt_message_with_files( + self, message_files: list[MessageFile], text_content: str, message: Message, app_record, is_user_message: bool + ) -> PromptMessage: + """ + Build prompt message with files. + :param message_files: list of MessageFile objects + :param text_content: text content of the message + :param message: Message object + :param app_record: app record + :param is_user_message: whether this is a user message + :return: PromptMessage + """ + if self.conversation.mode in {AppMode.AGENT_CHAT, AppMode.COMPLETION, AppMode.CHAT}: + file_extra_config = FileUploadConfigManager.convert(self.conversation.model_config) + elif self.conversation.mode in {AppMode.ADVANCED_CHAT, AppMode.WORKFLOW}: + workflow_run = db.session.scalar(select(WorkflowRun).where(WorkflowRun.id == message.workflow_run_id)) + if not workflow_run: + raise ValueError(f"Workflow run not found: {message.workflow_run_id}") + workflow = db.session.scalar(select(Workflow).where(Workflow.id == workflow_run.workflow_id)) + if not workflow: + raise ValueError(f"Workflow not found: {workflow_run.workflow_id}") + file_extra_config = FileUploadConfigManager.convert(workflow.features_dict, is_vision=False) + else: + raise AssertionError(f"Invalid app mode: {self.conversation.mode}") + + detail = ImagePromptMessageContent.DETAIL.HIGH + if file_extra_config and app_record: + # Build files directly without filtering by belongs_to + file_objs = [ + file_factory.build_from_message_file( + message_file=message_file, tenant_id=app_record.tenant_id, config=file_extra_config + ) + for message_file in message_files + ] + if file_extra_config.image_config and file_extra_config.image_config.detail: + detail = file_extra_config.image_config.detail + else: + file_objs = [] + + if not file_objs: + if is_user_message: + return UserPromptMessage(content=text_content) + else: + return AssistantPromptMessage(content=text_content) + else: + prompt_message_contents: list[PromptMessageContentUnionTypes] = [] + for file in file_objs: + prompt_message = file_manager.to_prompt_message_content( + file, + image_detail_config=detail, + ) + prompt_message_contents.append(prompt_message) + prompt_message_contents.append(TextPromptMessageContent(data=text_content)) + + if is_user_message: + return UserPromptMessage(content=prompt_message_contents) + else: + return AssistantPromptMessage(content=prompt_message_contents) + def get_history_prompt_messages( self, max_token_limit: int = 2000, message_limit: Optional[int] = None ) -> Sequence[PromptMessage]: @@ -67,52 +126,46 @@ class TokenBufferMemory: prompt_messages: list[PromptMessage] = [] for message in messages: - files = db.session.query(MessageFile).where(MessageFile.message_id == message.id).all() - if files: - file_extra_config = None - if self.conversation.mode in {AppMode.AGENT_CHAT, AppMode.COMPLETION, AppMode.CHAT}: - file_extra_config = FileUploadConfigManager.convert(self.conversation.model_config) - elif self.conversation.mode in {AppMode.ADVANCED_CHAT, AppMode.WORKFLOW}: - workflow_run = db.session.scalar( - select(WorkflowRun).where(WorkflowRun.id == message.workflow_run_id) - ) - if not workflow_run: - raise ValueError(f"Workflow run not found: {message.workflow_run_id}") - workflow = db.session.scalar(select(Workflow).where(Workflow.id == workflow_run.workflow_id)) - if not workflow: - raise ValueError(f"Workflow not found: {workflow_run.workflow_id}") - file_extra_config = FileUploadConfigManager.convert(workflow.features_dict, is_vision=False) - else: - raise AssertionError(f"Invalid app mode: {self.conversation.mode}") - - detail = ImagePromptMessageContent.DETAIL.LOW - if file_extra_config and app_record: - file_objs = file_factory.build_from_message_files( - message_files=files, tenant_id=app_record.tenant_id, config=file_extra_config - ) - if file_extra_config.image_config and file_extra_config.image_config.detail: - detail = file_extra_config.image_config.detail - else: - file_objs = [] - - if not file_objs: - prompt_messages.append(UserPromptMessage(content=message.query)) - else: - prompt_message_contents: list[PromptMessageContentUnionTypes] = [] - for file in file_objs: - prompt_message = file_manager.to_prompt_message_content( - file, - image_detail_config=detail, - ) - prompt_message_contents.append(prompt_message) - prompt_message_contents.append(TextPromptMessageContent(data=message.query)) - - prompt_messages.append(UserPromptMessage(content=prompt_message_contents)) + # Process user message with files + user_files = ( + db.session.query(MessageFile) + .where( + MessageFile.message_id == message.id, + (MessageFile.belongs_to == "user") | (MessageFile.belongs_to.is_(None)), + ) + .all() + ) + if user_files: + user_prompt_message = self._build_prompt_message_with_files( + message_files=user_files, + text_content=message.query, + message=message, + app_record=app_record, + is_user_message=True, + ) + prompt_messages.append(user_prompt_message) else: prompt_messages.append(UserPromptMessage(content=message.query)) - prompt_messages.append(AssistantPromptMessage(content=message.answer)) + # Process assistant message with files + assistant_files = ( + db.session.query(MessageFile) + .where(MessageFile.message_id == message.id, MessageFile.belongs_to == "assistant") + .all() + ) + + if assistant_files: + assistant_prompt_message = self._build_prompt_message_with_files( + message_files=assistant_files, + text_content=message.answer, + message=message, + app_record=app_record, + is_user_message=False, + ) + prompt_messages.append(assistant_prompt_message) + else: + prompt_messages.append(AssistantPromptMessage(content=message.answer)) if not prompt_messages: return [] diff --git a/api/factories/file_factory.py b/api/factories/file_factory.py index a0ff33ab6..0ea7d3ae1 100644 --- a/api/factories/file_factory.py +++ b/api/factories/file_factory.py @@ -41,8 +41,14 @@ def build_from_message_file( "url": message_file.url, "id": message_file.id, "type": message_file.type, - "upload_file_id": message_file.upload_file_id, } + + # Set the correct ID field based on transfer method + if message_file.transfer_method == FileTransferMethod.TOOL_FILE.value: + mapping["tool_file_id"] = message_file.upload_file_id + else: + mapping["upload_file_id"] = message_file.upload_file_id + return build_from_mapping( mapping=mapping, tenant_id=tenant_id, @@ -318,6 +324,11 @@ def _is_file_valid_with_config( file_transfer_method: FileTransferMethod, config: FileUploadConfig, ) -> bool: + # FIXME(QIN2DIM): Always allow tool files (files generated by the assistant/model) + # These are internally generated and should bypass user upload restrictions + if file_transfer_method == FileTransferMethod.TOOL_FILE: + return True + if ( config.allowed_file_types and input_file_type not in config.allowed_file_types