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,
)
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)
return "\n".join(string_messages)