mirror of
https://github.com/langgenius/dify.git
synced 2024-11-16 11:42:29 +08:00
80 lines
2.4 KiB
Python
80 lines
2.4 KiB
Python
import enum
|
|
from typing import Optional
|
|
|
|
from pydantic import BaseModel
|
|
|
|
from core.file.upload_file_parser import UploadFileParser
|
|
from core.model_providers.models.entity.message import PromptMessageFile, ImagePromptMessageFile
|
|
from extensions.ext_database import db
|
|
from models.model import UploadFile
|
|
|
|
|
|
class FileType(enum.Enum):
|
|
IMAGE = 'image'
|
|
|
|
@staticmethod
|
|
def value_of(value):
|
|
for member in FileType:
|
|
if member.value == value:
|
|
return member
|
|
raise ValueError(f"No matching enum found for value '{value}'")
|
|
|
|
|
|
class FileTransferMethod(enum.Enum):
|
|
REMOTE_URL = 'remote_url'
|
|
LOCAL_FILE = 'local_file'
|
|
|
|
@staticmethod
|
|
def value_of(value):
|
|
for member in FileTransferMethod:
|
|
if member.value == value:
|
|
return member
|
|
raise ValueError(f"No matching enum found for value '{value}'")
|
|
|
|
|
|
class FileObj(BaseModel):
|
|
id: Optional[str]
|
|
tenant_id: str
|
|
type: FileType
|
|
transfer_method: FileTransferMethod
|
|
url: Optional[str]
|
|
upload_file_id: Optional[str]
|
|
file_config: dict
|
|
|
|
@property
|
|
def data(self) -> Optional[str]:
|
|
return self._get_data()
|
|
|
|
@property
|
|
def preview_url(self) -> Optional[str]:
|
|
return self._get_data(force_url=True)
|
|
|
|
@property
|
|
def prompt_message_file(self) -> PromptMessageFile:
|
|
if self.type == FileType.IMAGE:
|
|
image_config = self.file_config.get('image')
|
|
|
|
return ImagePromptMessageFile(
|
|
data=self.data,
|
|
detail=ImagePromptMessageFile.DETAIL.HIGH
|
|
if image_config.get("detail") == "high" else ImagePromptMessageFile.DETAIL.LOW
|
|
)
|
|
|
|
def _get_data(self, force_url: bool = False) -> Optional[str]:
|
|
if self.type == FileType.IMAGE:
|
|
if self.transfer_method == FileTransferMethod.REMOTE_URL:
|
|
return self.url
|
|
elif self.transfer_method == FileTransferMethod.LOCAL_FILE:
|
|
upload_file = (db.session.query(UploadFile)
|
|
.filter(
|
|
UploadFile.id == self.upload_file_id,
|
|
UploadFile.tenant_id == self.tenant_id
|
|
).first())
|
|
|
|
return UploadFileParser.get_image_data(
|
|
upload_file=upload_file,
|
|
force_url=force_url
|
|
)
|
|
|
|
return None
|