From dba67cd87a8f40d0656850cfe201777bd1cb5311 Mon Sep 17 00:00:00 2001 From: -LAN- Date: Tue, 15 Oct 2024 12:01:00 +0800 Subject: [PATCH] fix(memory): filter non-image file types in prompt message content - Skip non-image files when converting file objects to prompt message content. - Ensures only image files are processed, improving the accuracy and relevance of prompt messages. --- api/core/memory/token_buffer_memory.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/api/core/memory/token_buffer_memory.py b/api/core/memory/token_buffer_memory.py index 98704a9f43..d0b5f67ae5 100644 --- a/api/core/memory/token_buffer_memory.py +++ b/api/core/memory/token_buffer_memory.py @@ -1,7 +1,7 @@ from typing import Optional from core.app.app_config.features.file_upload.manager import FileUploadConfigManager -from core.file import file_manager +from core.file import FileType, file_manager from core.model_manager import ModelInstance from core.model_runtime.entities import ( AssistantPromptMessage, @@ -98,7 +98,10 @@ class TokenBufferMemory: prompt_message_contents: list[PromptMessageContent] = [] prompt_message_contents.append(TextPromptMessageContent(data=message.query)) for file_obj in file_objs: - prompt_message_contents.append(file_manager.to_prompt_message_content(file_obj)) + if file_obj.type != FileType.IMAGE: + continue + prompt_message = file_manager.to_prompt_message_content(file_obj) + prompt_message_contents.append(prompt_message) prompt_messages.append(UserPromptMessage(content=prompt_message_contents)) else: