From d522308a29c31b5f89ca88bbebe8a763724e28cb Mon Sep 17 00:00:00 2001
From: takatost <takatost@users.noreply.github.com>
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 f38175020..21f1965e9 100644
--- a/api/core/memory/token_buffer_memory.py
+++ b/api/core/memory/token_buffer_memory.py
@@ -12,7 +12,8 @@ from core.model_runtime.entities.message_entities import (
     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 @@ class TokenBufferMemory:
         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 @@ class TokenBufferMemory:
             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 @@ class TokenBufferMemory:
                 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)