From d522308a29c31b5f89ca88bbebe8a763724e28cb Mon Sep 17 00:00:00 2001 From: takatost Date: Sun, 7 Jul 2024 08:54:24 +0800 Subject: [PATCH] chore: optimize memory fetch performance (#6039) --- api/core/memory/token_buffer_memory.py | 32 ++++++++++++++++++-------- 1 file changed, 22 insertions(+), 10 deletions(-) diff --git a/api/core/memory/token_buffer_memory.py b/api/core/memory/token_buffer_memory.py index f38175020cca70..21f1965e93adac 100644 --- a/api/core/memory/token_buffer_memory.py +++ b/api/core/memory/token_buffer_memory.py @@ -12,7 +12,8 @@ UserPromptMessage, ) from extensions.ext_database import db -from models.model import AppMode, Conversation, Message +from models.model import AppMode, Conversation, Message, MessageFile +from models.workflow import WorkflowRun class TokenBufferMemory: @@ -30,7 +31,13 @@ def get_history_prompt_messages(self, max_token_limit: int = 2000, app_record = self.conversation.app # fetch limited messages, and return reversed - query = db.session.query(Message).filter( + query = db.session.query( + Message.id, + Message.query, + Message.answer, + Message.created_at, + Message.workflow_run_id + ).filter( Message.conversation_id == self.conversation.id, Message.answer != '' ).order_by(Message.created_at.desc()) @@ -47,18 +54,23 @@ def get_history_prompt_messages(self, max_token_limit: int = 2000, tenant_id=app_record.tenant_id, app_id=app_record.id ) - prompt_messages = [] for message in messages: - files = message.message_files + files = db.session.query(MessageFile).filter(MessageFile.message_id == message.id).all() if files: + file_extra_config = None if self.conversation.mode not in [AppMode.ADVANCED_CHAT.value, AppMode.WORKFLOW.value]: - file_extra_config = FileUploadConfigManager.convert(message.app_model_config.to_dict()) + file_extra_config = FileUploadConfigManager.convert(self.conversation.model_config) else: - file_extra_config = FileUploadConfigManager.convert( - message.workflow_run.workflow.features_dict, - is_vision=False - ) + if message.workflow_run_id: + workflow_run = (db.session.query(WorkflowRun) + .filter(WorkflowRun.id == message.workflow_run_id).first()) + + if workflow_run: + file_extra_config = FileUploadConfigManager.convert( + workflow_run.workflow.features_dict, + is_vision=False + ) if file_extra_config: file_objs = message_file_parser.transform_message_files( @@ -138,4 +150,4 @@ def get_history_prompt_text(self, human_prefix: str = "Human", message = f"{role}: {m.content}" string_messages.append(message) - return "\n".join(string_messages) \ No newline at end of file + return "\n".join(string_messages)