Austin/dify
1
0
Fork 0
mirror of https://github.com/langgenius/dify.git synced 2024-11-16 11:42:29 +08:00
Code Issues Actions 4 Packages Projects Releases Wiki Activity

chore: optimize memory fetch performance (#6039)

Browse Source
This commit is contained in:
takatost 2024-07-07 08:54:24 +08:00 committed by GitHub
parent 85744b72e5
commit d522308a29
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 22 additions and 10 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

32
api/core/memory/token_buffer_memory.py
View File

@ -12,7 +12,8 @@ from core.model_runtime.entities.message_entities import (
UserPromptMessage, UserPromptMessage,
) )
from extensions.ext_database import db 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: class TokenBufferMemory:
@ -30,7 +31,13 @@ class TokenBufferMemory:
app_record = self.conversation.app app_record = self.conversation.app
# fetch limited messages, and return reversed # 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.conversation_id == self.conversation.id,
Message.answer != '' Message.answer != ''
).order_by(Message.created_at.desc()) ).order_by(Message.created_at.desc())
@ -47,18 +54,23 @@ class TokenBufferMemory:
tenant_id=app_record.tenant_id, tenant_id=app_record.tenant_id,
app_id=app_record.id app_id=app_record.id
) )
prompt_messages = [] prompt_messages = []
for message in messages: for message in messages:
files = message.message_files files = db.session.query(MessageFile).filter(MessageFile.message_id == message.id).all()
if files: if files:
file_extra_config = None
if self.conversation.mode not in [AppMode.ADVANCED_CHAT.value, AppMode.WORKFLOW.value]: 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: else:
file_extra_config = FileUploadConfigManager.convert( if message.workflow_run_id:
message.workflow_run.workflow.features_dict, workflow_run = (db.session.query(WorkflowRun)
is_vision=False .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: if file_extra_config:
file_objs = message_file_parser.transform_message_files( file_objs = message_file_parser.transform_message_files(
@ -138,4 +150,4 @@ class TokenBufferMemory:
message = f"{role}: {m.content}" message = f"{role}: {m.content}"
string_messages.append(message) string_messages.append(message)
return "\n".join(string_messages) return "\n".join(string_messages)