chore: optimize memory fetch performance (#6039)

langgenius · Jul 7, 2024 · d522308 · d522308
1 parent 85744b7
commit d522308
Showing 1 changed file with 22 additions and 10 deletions.
diff --git 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)
+        return "\n".join(string_messages)