Underground-Digital/Workflow-Engine

1import logging2from mimetypes import guess_extension3from typing import Optional4 5from core.file import File, FileTransferMethod, FileType6from core.tools.entities.tool_entities import ToolInvokeMessage7from core.tools.tool_file_manager import ToolFileManager8 9logger = logging.getLogger(__name__)10 11 12class ToolFileMessageTransformer:13    @classmethod14    def transform_tool_invoke_messages(15        cls, messages: list[ToolInvokeMessage], user_id: str, tenant_id: str, conversation_id: str | None16    ) -> list[ToolInvokeMessage]:17        """18        Transform tool message and handle file download19        """20        result = []21 22        for message in messages:23            if message.type in {ToolInvokeMessage.MessageType.TEXT, ToolInvokeMessage.MessageType.LINK}:24                result.append(message)25            elif message.type == ToolInvokeMessage.MessageType.IMAGE and isinstance(message.message, str):26                # try to download image27                try:28                    file = ToolFileManager.create_file_by_url(29                        user_id=user_id, tenant_id=tenant_id, conversation_id=conversation_id, file_url=message.message30                    )31 32                    url = f'/files/tools/{file.id}{guess_extension(file.mimetype) or ".png"}'33 34                    result.append(35                        ToolInvokeMessage(36                            type=ToolInvokeMessage.MessageType.IMAGE_LINK,37                            message=url,38                            save_as=message.save_as,39                            meta=message.meta.copy() if message.meta is not None else {},40                        )41                    )42                except Exception as e:43                    logger.exception(e)44                    result.append(45                        ToolInvokeMessage(46                            type=ToolInvokeMessage.MessageType.TEXT,47                            message=f"Failed to download image: {message.message}, please try to download it manually.",48                            meta=message.meta.copy() if message.meta is not None else {},49                            save_as=message.save_as,50                        )51                    )52            elif message.type == ToolInvokeMessage.MessageType.BLOB:53                # get mime type and save blob to storage54                assert message.meta is not None55                mimetype = message.meta.get("mime_type", "octet/stream")56                # if message is str, encode it to bytes57                if isinstance(message.message, str):58                    message.message = message.message.encode("utf-8")59 60                # FIXME: should do a type check here.61                assert isinstance(message.message, bytes)62                file = ToolFileManager.create_file_by_raw(63                    user_id=user_id,64                    tenant_id=tenant_id,65                    conversation_id=conversation_id,66                    file_binary=message.message,67                    mimetype=mimetype,68                )69 70                url = cls.get_tool_file_url(tool_file_id=file.id, extension=guess_extension(file.mimetype))71 72                # check if file is image73                if "image" in mimetype:74                    result.append(75                        ToolInvokeMessage(76                            type=ToolInvokeMessage.MessageType.IMAGE_LINK,77                            message=url,78                            save_as=message.save_as,79                            meta=message.meta.copy() if message.meta is not None else {},80                        )81                    )82                else:83                    result.append(84                        ToolInvokeMessage(85                            type=ToolInvokeMessage.MessageType.LINK,86                            message=url,87                            save_as=message.save_as,88                            meta=message.meta.copy() if message.meta is not None else {},89                        )90                    )91            elif message.type == ToolInvokeMessage.MessageType.FILE:92                assert message.meta is not None93                file = message.meta.get("file")94                if isinstance(file, File):95                    if file.transfer_method == FileTransferMethod.TOOL_FILE:96                        assert file.related_id is not None97                        url = cls.get_tool_file_url(tool_file_id=file.related_id, extension=file.extension)98                        if file.type == FileType.IMAGE:99                            result.append(100                                ToolInvokeMessage(101                                    type=ToolInvokeMessage.MessageType.IMAGE_LINK,102                                    message=url,103                                    save_as=message.save_as,104                                    meta=message.meta.copy() if message.meta is not None else {},105                                )106                            )107                        else:108                            result.append(109                                ToolInvokeMessage(110                                    type=ToolInvokeMessage.MessageType.LINK,111                                    message=url,112                                    save_as=message.save_as,113                                    meta=message.meta.copy() if message.meta is not None else {},114                                )115                            )116                    else:117                        result.append(message)118            else:119                result.append(message)120 121        return result122 123    @classmethod124    def get_tool_file_url(cls, tool_file_id: str, extension: Optional[str]) -> str:125        return f'/files/tools/{tool_file_id}{extension or ".bin"}'126