File size: 5,722 Bytes
a8b3f00
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
from collections.abc import Mapping, Sequence
from typing import Optional

from pydantic import BaseModel, Field, model_validator

from core.model_runtime.entities.message_entities import ImagePromptMessageContent

from . import helpers
from .constants import FILE_MODEL_IDENTITY
from .enums import FileTransferMethod, FileType
from .tool_file_parser import ToolFileParser


class ImageConfig(BaseModel):
    """
    NOTE: This part of validation is deprecated, but still used in app features "Image Upload".
    """

    number_limits: int = 0
    transfer_methods: Sequence[FileTransferMethod] = Field(default_factory=list)
    detail: ImagePromptMessageContent.DETAIL | None = None


class FileExtraConfig(BaseModel):
    """
    File Upload Entity.
    """

    image_config: Optional[ImageConfig] = None
    allowed_file_types: Sequence[FileType] = Field(default_factory=list)
    allowed_extensions: Sequence[str] = Field(default_factory=list)
    allowed_upload_methods: Sequence[FileTransferMethod] = Field(default_factory=list)
    number_limits: int = 0


class File(BaseModel):
    dify_model_identity: str = FILE_MODEL_IDENTITY

    id: Optional[str] = None  # message file id
    tenant_id: str
    type: FileType
    transfer_method: FileTransferMethod
    remote_url: Optional[str] = None  # remote url
    related_id: Optional[str] = None
    filename: Optional[str] = None
    extension: Optional[str] = Field(default=None, description="File extension, should contains dot")
    mime_type: Optional[str] = None
    size: int = -1
    _extra_config: FileExtraConfig | None = None

    def to_dict(self) -> Mapping[str, str | int | None]:
        data = self.model_dump(mode="json")
        return {
            **data,
            "url": self.generate_url(),
        }

    @property
    def markdown(self) -> str:
        url = self.generate_url()
        if self.type == FileType.IMAGE:
            text = f'![{self.filename or ""}]({url})'
        else:
            text = f"[{self.filename or url}]({url})"

        return text

    def generate_url(self) -> Optional[str]:
        if self.type == FileType.IMAGE:
            if self.transfer_method == FileTransferMethod.REMOTE_URL:
                return self.remote_url
            elif self.transfer_method == FileTransferMethod.LOCAL_FILE:
                if self.related_id is None:
                    raise ValueError("Missing file related_id")
                return helpers.get_signed_file_url(upload_file_id=self.related_id)
            elif self.transfer_method == FileTransferMethod.TOOL_FILE:
                assert self.related_id is not None
                assert self.extension is not None
                return ToolFileParser.get_tool_file_manager().sign_file(
                    tool_file_id=self.related_id, extension=self.extension
                )
        else:
            if self.transfer_method == FileTransferMethod.REMOTE_URL:
                return self.remote_url
            elif self.transfer_method == FileTransferMethod.LOCAL_FILE:
                if self.related_id is None:
                    raise ValueError("Missing file related_id")
                return helpers.get_signed_file_url(upload_file_id=self.related_id)
            elif self.transfer_method == FileTransferMethod.TOOL_FILE:
                assert self.related_id is not None
                assert self.extension is not None
                return ToolFileParser.get_tool_file_manager().sign_file(
                    tool_file_id=self.related_id, extension=self.extension
                )

    @model_validator(mode="after")
    def validate_after(self):
        match self.transfer_method:
            case FileTransferMethod.REMOTE_URL:
                if not self.remote_url:
                    raise ValueError("Missing file url")
                if not isinstance(self.remote_url, str) or not self.remote_url.startswith("http"):
                    raise ValueError("Invalid file url")
            case FileTransferMethod.LOCAL_FILE:
                if not self.related_id:
                    raise ValueError("Missing file related_id")
            case FileTransferMethod.TOOL_FILE:
                if not self.related_id:
                    raise ValueError("Missing file related_id")

        # Validate the extra config.
        if not self._extra_config:
            return self

        if self._extra_config.allowed_file_types:
            if self.type not in self._extra_config.allowed_file_types and self.type != FileType.CUSTOM:
                raise ValueError(f"Invalid file type: {self.type}")

        if self._extra_config.allowed_extensions and self.extension not in self._extra_config.allowed_extensions:
            raise ValueError(f"Invalid file extension: {self.extension}")

        if (
            self._extra_config.allowed_upload_methods
            and self.transfer_method not in self._extra_config.allowed_upload_methods
        ):
            raise ValueError(f"Invalid transfer method: {self.transfer_method}")

        match self.type:
            case FileType.IMAGE:
                # NOTE: This part of validation is deprecated, but still used in app features "Image Upload".
                if not self._extra_config.image_config:
                    return self
                # TODO: skip check if transfer_methods is empty, because many test cases are not setting this field
                if (
                    self._extra_config.image_config.transfer_methods
                    and self.transfer_method not in self._extra_config.image_config.transfer_methods
                ):
                    raise ValueError(f"Invalid transfer method: {self.transfer_method}")

        return self