youngkingdom/dify
1
0
Fork 0
mirror of https://github.com/langgenius/dify.git synced 2026-05-03 08:58:09 +08:00
Code Issues Packages Projects Releases Wiki Activity

refactor: tool message transformer

Browse Source
This commit is contained in:
Yeuoly
2024-08-29 13:42:31 +08:00
parent 4b4741f7ed
commit c28998a6f0
5 changed files with 52 additions and 31 deletions
Download Patch File Download Diff File Expand all files Collapse all files

44
api/core/tools/utils/message_transformer.py
View File

@ -1,8 +1,9 @@
import logging
from collections.abc import Generator
from mimetypes import guess_extension
from typing import Optional
from core.file.file_obj import FileTransferMethod, FileType
from core.file.file_obj import FileTransferMethod, FileType, FileVar
from core.tools.entities.tool_entities import ToolInvokeMessage
from core.tools.tool_file_manager import ToolFileManager
@ -13,7 +14,7 @@ class ToolFileMessageTransformer:
def transform_tool_invoke_messages(cls, messages: Generator[ToolInvokeMessage, None, None],
user_id: str,
tenant_id: str,
conversation_id: str) -> Generator[ToolInvokeMessage, None, None]:
conversation_id: Optional[str] = None) -> Generator[ToolInvokeMessage, None, None]:
"""
Transform tool message and handle file download
"""
@ -25,18 +26,23 @@ class ToolFileMessageTransformer:
elif message.type == ToolInvokeMessage.MessageType.IMAGE:
# try to download image
try:
if not conversation_id:
raise
assert isinstance(message.message, ToolInvokeMessage.TextMessage)
file = ToolFileManager.create_file_by_url(
user_id=user_id,
tenant_id=tenant_id,
file_url=message.message.text,
conversation_id=conversation_id,
file_url=message.message
)
url = f'/files/tools/{file.id}{guess_extension(file.mimetype) or ".png"}'
yield ToolInvokeMessage(
type=ToolInvokeMessage.MessageType.IMAGE_LINK,
message=url,
message=ToolInvokeMessage.TextMessage(text=url),
save_as=message.save_as,
meta=message.meta.copy() if message.meta is not None else {},
)
@ -44,57 +50,67 @@ class ToolFileMessageTransformer:
logger.exception(e)
yield ToolInvokeMessage(
type=ToolInvokeMessage.MessageType.TEXT,
message=f"Failed to download image: {message.message}, you can try to download it yourself.",
message=ToolInvokeMessage.TextMessage(
text=f"Failed to download image: {message.message}, you can try to download it yourself."
),
meta=message.meta.copy() if message.meta is not None else {},
save_as=message.save_as,
)
elif message.type == ToolInvokeMessage.MessageType.BLOB:
# get mime type and save blob to storage
assert message.meta
mimetype = message.meta.get('mime_type', 'octet/stream')
# if message is str, encode it to bytes
if isinstance(message.message, str):
message.message = message.message.encode('utf-8')
if not isinstance(message.message, ToolInvokeMessage.BlobMessage):
raise ValueError("unexpected message type")
file = ToolFileManager.create_file_by_raw(
user_id=user_id, tenant_id=tenant_id,
conversation_id=conversation_id,
file_binary=message.message,
file_binary=message.message.blob,
mimetype=mimetype
)
url = cls.get_tool_file_url(file.id, guess_extension(file.mimetype))
extension = guess_extension(file.mimetype) or ".bin"
url = cls.get_tool_file_url(file.id, extension)
# check if file is image
if 'image' in mimetype:
yield ToolInvokeMessage(
type=ToolInvokeMessage.MessageType.IMAGE_LINK,
message=url,
message=ToolInvokeMessage.TextMessage(text=url),
save_as=message.save_as,
meta=message.meta.copy() if message.meta is not None else {},
)
else:
yield ToolInvokeMessage(
type=ToolInvokeMessage.MessageType.LINK,
message=url,
message=ToolInvokeMessage.TextMessage(text=url),
save_as=message.save_as,
meta=message.meta.copy() if message.meta is not None else {},
)
elif message.type == ToolInvokeMessage.MessageType.FILE_VAR:
file_var = message.meta.get('file_var')
assert message.meta
file_var: FileVar | None = message.meta.get('file_var')
if file_var:
if file_var.transfer_method == FileTransferMethod.TOOL_FILE:
assert file_var.related_id and file_var.extension
url = cls.get_tool_file_url(file_var.related_id, file_var.extension)
if file_var.type == FileType.IMAGE:
yield ToolInvokeMessage(
type=ToolInvokeMessage.MessageType.IMAGE_LINK,
message=url,
message=ToolInvokeMessage.TextMessage(text=url),
save_as=message.save_as,
meta=message.meta.copy() if message.meta is not None else {},
)
else:
yield ToolInvokeMessage(
type=ToolInvokeMessage.MessageType.LINK,
message=url,
message=ToolInvokeMessage.TextMessage(text=url),
save_as=message.save_as,
meta=message.meta.copy() if message.meta is not None else {},
)