|
|
|
@@ -31,6 +31,65 @@ class TokenBufferMemory: |
|
|
|
self.conversation = conversation |
|
|
|
self.model_instance = model_instance |
|
|
|
|
|
|
|
def _build_prompt_message_with_files( |
|
|
|
self, message_files: list[MessageFile], text_content: str, message: Message, app_record, is_user_message: bool |
|
|
|
) -> PromptMessage: |
|
|
|
""" |
|
|
|
Build prompt message with files. |
|
|
|
:param message_files: list of MessageFile objects |
|
|
|
:param text_content: text content of the message |
|
|
|
:param message: Message object |
|
|
|
:param app_record: app record |
|
|
|
:param is_user_message: whether this is a user message |
|
|
|
:return: PromptMessage |
|
|
|
""" |
|
|
|
if self.conversation.mode in {AppMode.AGENT_CHAT, AppMode.COMPLETION, AppMode.CHAT}: |
|
|
|
file_extra_config = FileUploadConfigManager.convert(self.conversation.model_config) |
|
|
|
elif self.conversation.mode in {AppMode.ADVANCED_CHAT, AppMode.WORKFLOW}: |
|
|
|
workflow_run = db.session.scalar(select(WorkflowRun).where(WorkflowRun.id == message.workflow_run_id)) |
|
|
|
if not workflow_run: |
|
|
|
raise ValueError(f"Workflow run not found: {message.workflow_run_id}") |
|
|
|
workflow = db.session.scalar(select(Workflow).where(Workflow.id == workflow_run.workflow_id)) |
|
|
|
if not workflow: |
|
|
|
raise ValueError(f"Workflow not found: {workflow_run.workflow_id}") |
|
|
|
file_extra_config = FileUploadConfigManager.convert(workflow.features_dict, is_vision=False) |
|
|
|
else: |
|
|
|
raise AssertionError(f"Invalid app mode: {self.conversation.mode}") |
|
|
|
|
|
|
|
detail = ImagePromptMessageContent.DETAIL.HIGH |
|
|
|
if file_extra_config and app_record: |
|
|
|
# Build files directly without filtering by belongs_to |
|
|
|
file_objs = [ |
|
|
|
file_factory.build_from_message_file( |
|
|
|
message_file=message_file, tenant_id=app_record.tenant_id, config=file_extra_config |
|
|
|
) |
|
|
|
for message_file in message_files |
|
|
|
] |
|
|
|
if file_extra_config.image_config and file_extra_config.image_config.detail: |
|
|
|
detail = file_extra_config.image_config.detail |
|
|
|
else: |
|
|
|
file_objs = [] |
|
|
|
|
|
|
|
if not file_objs: |
|
|
|
if is_user_message: |
|
|
|
return UserPromptMessage(content=text_content) |
|
|
|
else: |
|
|
|
return AssistantPromptMessage(content=text_content) |
|
|
|
else: |
|
|
|
prompt_message_contents: list[PromptMessageContentUnionTypes] = [] |
|
|
|
for file in file_objs: |
|
|
|
prompt_message = file_manager.to_prompt_message_content( |
|
|
|
file, |
|
|
|
image_detail_config=detail, |
|
|
|
) |
|
|
|
prompt_message_contents.append(prompt_message) |
|
|
|
prompt_message_contents.append(TextPromptMessageContent(data=text_content)) |
|
|
|
|
|
|
|
if is_user_message: |
|
|
|
return UserPromptMessage(content=prompt_message_contents) |
|
|
|
else: |
|
|
|
return AssistantPromptMessage(content=prompt_message_contents) |
|
|
|
|
|
|
|
def get_history_prompt_messages( |
|
|
|
self, max_token_limit: int = 2000, message_limit: Optional[int] = None |
|
|
|
) -> Sequence[PromptMessage]: |
|
|
|
@@ -67,52 +126,46 @@ class TokenBufferMemory: |
|
|
|
|
|
|
|
prompt_messages: list[PromptMessage] = [] |
|
|
|
for message in messages: |
|
|
|
files = db.session.query(MessageFile).where(MessageFile.message_id == message.id).all() |
|
|
|
if files: |
|
|
|
file_extra_config = None |
|
|
|
if self.conversation.mode in {AppMode.AGENT_CHAT, AppMode.COMPLETION, AppMode.CHAT}: |
|
|
|
file_extra_config = FileUploadConfigManager.convert(self.conversation.model_config) |
|
|
|
elif self.conversation.mode in {AppMode.ADVANCED_CHAT, AppMode.WORKFLOW}: |
|
|
|
workflow_run = db.session.scalar( |
|
|
|
select(WorkflowRun).where(WorkflowRun.id == message.workflow_run_id) |
|
|
|
) |
|
|
|
if not workflow_run: |
|
|
|
raise ValueError(f"Workflow run not found: {message.workflow_run_id}") |
|
|
|
workflow = db.session.scalar(select(Workflow).where(Workflow.id == workflow_run.workflow_id)) |
|
|
|
if not workflow: |
|
|
|
raise ValueError(f"Workflow not found: {workflow_run.workflow_id}") |
|
|
|
file_extra_config = FileUploadConfigManager.convert(workflow.features_dict, is_vision=False) |
|
|
|
else: |
|
|
|
raise AssertionError(f"Invalid app mode: {self.conversation.mode}") |
|
|
|
|
|
|
|
detail = ImagePromptMessageContent.DETAIL.LOW |
|
|
|
if file_extra_config and app_record: |
|
|
|
file_objs = file_factory.build_from_message_files( |
|
|
|
message_files=files, tenant_id=app_record.tenant_id, config=file_extra_config |
|
|
|
) |
|
|
|
if file_extra_config.image_config and file_extra_config.image_config.detail: |
|
|
|
detail = file_extra_config.image_config.detail |
|
|
|
else: |
|
|
|
file_objs = [] |
|
|
|
|
|
|
|
if not file_objs: |
|
|
|
prompt_messages.append(UserPromptMessage(content=message.query)) |
|
|
|
else: |
|
|
|
prompt_message_contents: list[PromptMessageContentUnionTypes] = [] |
|
|
|
for file in file_objs: |
|
|
|
prompt_message = file_manager.to_prompt_message_content( |
|
|
|
file, |
|
|
|
image_detail_config=detail, |
|
|
|
) |
|
|
|
prompt_message_contents.append(prompt_message) |
|
|
|
prompt_message_contents.append(TextPromptMessageContent(data=message.query)) |
|
|
|
|
|
|
|
prompt_messages.append(UserPromptMessage(content=prompt_message_contents)) |
|
|
|
|
|
|
|
# Process user message with files |
|
|
|
user_files = ( |
|
|
|
db.session.query(MessageFile) |
|
|
|
.where( |
|
|
|
MessageFile.message_id == message.id, |
|
|
|
(MessageFile.belongs_to == "user") | (MessageFile.belongs_to.is_(None)), |
|
|
|
) |
|
|
|
.all() |
|
|
|
) |
|
|
|
|
|
|
|
if user_files: |
|
|
|
user_prompt_message = self._build_prompt_message_with_files( |
|
|
|
message_files=user_files, |
|
|
|
text_content=message.query, |
|
|
|
message=message, |
|
|
|
app_record=app_record, |
|
|
|
is_user_message=True, |
|
|
|
) |
|
|
|
prompt_messages.append(user_prompt_message) |
|
|
|
else: |
|
|
|
prompt_messages.append(UserPromptMessage(content=message.query)) |
|
|
|
|
|
|
|
prompt_messages.append(AssistantPromptMessage(content=message.answer)) |
|
|
|
# Process assistant message with files |
|
|
|
assistant_files = ( |
|
|
|
db.session.query(MessageFile) |
|
|
|
.where(MessageFile.message_id == message.id, MessageFile.belongs_to == "assistant") |
|
|
|
.all() |
|
|
|
) |
|
|
|
|
|
|
|
if assistant_files: |
|
|
|
assistant_prompt_message = self._build_prompt_message_with_files( |
|
|
|
message_files=assistant_files, |
|
|
|
text_content=message.answer, |
|
|
|
message=message, |
|
|
|
app_record=app_record, |
|
|
|
is_user_message=False, |
|
|
|
) |
|
|
|
prompt_messages.append(assistant_prompt_message) |
|
|
|
else: |
|
|
|
prompt_messages.append(AssistantPromptMessage(content=message.answer)) |
|
|
|
|
|
|
|
if not prompt_messages: |
|
|
|
return [] |