Underground-Digital/Workflow-Engine

1from collections.abc import Mapping, Sequence2from typing import Optional3 4from pydantic import BaseModel, Field, model_validator5 6from core.model_runtime.entities.message_entities import ImagePromptMessageContent7 8from . import helpers9from .constants import FILE_MODEL_IDENTITY10from .enums import FileTransferMethod, FileType11from .tool_file_parser import ToolFileParser12 13 14class ImageConfig(BaseModel):15    """16    NOTE: This part of validation is deprecated, but still used in app features "Image Upload".17    """18 19    number_limits: int = 020    transfer_methods: Sequence[FileTransferMethod] = Field(default_factory=list)21    detail: ImagePromptMessageContent.DETAIL | None = None22 23 24class FileExtraConfig(BaseModel):25    """26    File Upload Entity.27    """28 29    image_config: Optional[ImageConfig] = None30    allowed_file_types: Sequence[FileType] = Field(default_factory=list)31    allowed_extensions: Sequence[str] = Field(default_factory=list)32    allowed_upload_methods: Sequence[FileTransferMethod] = Field(default_factory=list)33    number_limits: int = 034 35 36class File(BaseModel):37    dify_model_identity: str = FILE_MODEL_IDENTITY38 39    id: Optional[str] = None  # message file id40    tenant_id: str41    type: FileType42    transfer_method: FileTransferMethod43    remote_url: Optional[str] = None  # remote url44    related_id: Optional[str] = None45    filename: Optional[str] = None46    extension: Optional[str] = Field(default=None, description="File extension, should contains dot")47    mime_type: Optional[str] = None48    size: int = -149    _extra_config: FileExtraConfig | None = None50 51    def to_dict(self) -> Mapping[str, str | int | None]:52        data = self.model_dump(mode="json")53        return {54            **data,55            "url": self.generate_url(),56        }57 58    @property59    def markdown(self) -> str:60        url = self.generate_url()61        if self.type == FileType.IMAGE:62            text = f'![{self.filename or ""}]({url})'63        else:64            text = f"[{self.filename or url}]({url})"65 66        return text67 68    def generate_url(self) -> Optional[str]:69        if self.type == FileType.IMAGE:70            if self.transfer_method == FileTransferMethod.REMOTE_URL:71                return self.remote_url72            elif self.transfer_method == FileTransferMethod.LOCAL_FILE:73                if self.related_id is None:74                    raise ValueError("Missing file related_id")75                return helpers.get_signed_file_url(upload_file_id=self.related_id)76            elif self.transfer_method == FileTransferMethod.TOOL_FILE:77                assert self.related_id is not None78                assert self.extension is not None79                return ToolFileParser.get_tool_file_manager().sign_file(80                    tool_file_id=self.related_id, extension=self.extension81                )82        else:83            if self.transfer_method == FileTransferMethod.REMOTE_URL:84                return self.remote_url85            elif self.transfer_method == FileTransferMethod.LOCAL_FILE:86                if self.related_id is None:87                    raise ValueError("Missing file related_id")88                return helpers.get_signed_file_url(upload_file_id=self.related_id)89            elif self.transfer_method == FileTransferMethod.TOOL_FILE:90                assert self.related_id is not None91                assert self.extension is not None92                return ToolFileParser.get_tool_file_manager().sign_file(93                    tool_file_id=self.related_id, extension=self.extension94                )95 96    @model_validator(mode="after")97    def validate_after(self):98        match self.transfer_method:99            case FileTransferMethod.REMOTE_URL:100                if not self.remote_url:101                    raise ValueError("Missing file url")102                if not isinstance(self.remote_url, str) or not self.remote_url.startswith("http"):103                    raise ValueError("Invalid file url")104            case FileTransferMethod.LOCAL_FILE:105                if not self.related_id:106                    raise ValueError("Missing file related_id")107            case FileTransferMethod.TOOL_FILE:108                if not self.related_id:109                    raise ValueError("Missing file related_id")110 111        # Validate the extra config.112        if not self._extra_config:113            return self114 115        if self._extra_config.allowed_file_types:116            if self.type not in self._extra_config.allowed_file_types and self.type != FileType.CUSTOM:117                raise ValueError(f"Invalid file type: {self.type}")118 119        if self._extra_config.allowed_extensions and self.extension not in self._extra_config.allowed_extensions:120            raise ValueError(f"Invalid file extension: {self.extension}")121 122        if (123            self._extra_config.allowed_upload_methods124            and self.transfer_method not in self._extra_config.allowed_upload_methods125        ):126            raise ValueError(f"Invalid transfer method: {self.transfer_method}")127 128        match self.type:129            case FileType.IMAGE:130                # NOTE: This part of validation is deprecated, but still used in app features "Image Upload".131                if not self._extra_config.image_config:132                    return self133                # TODO: skip check if transfer_methods is empty, because many test cases are not setting this field134                if (135                    self._extra_config.image_config.transfer_methods136                    and self.transfer_method not in self._extra_config.image_config.transfer_methods137                ):138                    raise ValueError(f"Invalid transfer method: {self.transfer_method}")139 140        return self141