mirror of https://github.com/langgenius/dify.git
refactor: port api/fields/file_fields.py (#30638)
This commit is contained in:
parent
55de731f9c
commit
0294555893
|
|
@ -1,7 +1,7 @@
|
|||
from typing import Literal
|
||||
|
||||
from flask import request
|
||||
from flask_restx import Resource, marshal_with
|
||||
from flask_restx import Resource
|
||||
from werkzeug.exceptions import Forbidden
|
||||
|
||||
import services
|
||||
|
|
@ -15,18 +15,21 @@ from controllers.common.errors import (
|
|||
TooManyFilesError,
|
||||
UnsupportedFileTypeError,
|
||||
)
|
||||
from controllers.common.schema import register_schema_models
|
||||
from controllers.console.wraps import (
|
||||
account_initialization_required,
|
||||
cloud_edition_billing_resource_check,
|
||||
setup_required,
|
||||
)
|
||||
from extensions.ext_database import db
|
||||
from fields.file_fields import file_fields, upload_config_fields
|
||||
from fields.file_fields import FileResponse, UploadConfig
|
||||
from libs.login import current_account_with_tenant, login_required
|
||||
from services.file_service import FileService
|
||||
|
||||
from . import console_ns
|
||||
|
||||
register_schema_models(console_ns, UploadConfig, FileResponse)
|
||||
|
||||
PREVIEW_WORDS_LIMIT = 3000
|
||||
|
||||
|
||||
|
|
@ -35,26 +38,27 @@ class FileApi(Resource):
|
|||
@setup_required
|
||||
@login_required
|
||||
@account_initialization_required
|
||||
@marshal_with(upload_config_fields)
|
||||
@console_ns.response(200, "Success", console_ns.models[UploadConfig.__name__])
|
||||
def get(self):
|
||||
return {
|
||||
"file_size_limit": dify_config.UPLOAD_FILE_SIZE_LIMIT,
|
||||
"batch_count_limit": dify_config.UPLOAD_FILE_BATCH_LIMIT,
|
||||
"file_upload_limit": dify_config.BATCH_UPLOAD_LIMIT,
|
||||
"image_file_size_limit": dify_config.UPLOAD_IMAGE_FILE_SIZE_LIMIT,
|
||||
"video_file_size_limit": dify_config.UPLOAD_VIDEO_FILE_SIZE_LIMIT,
|
||||
"audio_file_size_limit": dify_config.UPLOAD_AUDIO_FILE_SIZE_LIMIT,
|
||||
"workflow_file_upload_limit": dify_config.WORKFLOW_FILE_UPLOAD_LIMIT,
|
||||
"image_file_batch_limit": dify_config.IMAGE_FILE_BATCH_LIMIT,
|
||||
"single_chunk_attachment_limit": dify_config.SINGLE_CHUNK_ATTACHMENT_LIMIT,
|
||||
"attachment_image_file_size_limit": dify_config.ATTACHMENT_IMAGE_FILE_SIZE_LIMIT,
|
||||
}, 200
|
||||
config = UploadConfig(
|
||||
file_size_limit=dify_config.UPLOAD_FILE_SIZE_LIMIT,
|
||||
batch_count_limit=dify_config.UPLOAD_FILE_BATCH_LIMIT,
|
||||
file_upload_limit=dify_config.BATCH_UPLOAD_LIMIT,
|
||||
image_file_size_limit=dify_config.UPLOAD_IMAGE_FILE_SIZE_LIMIT,
|
||||
video_file_size_limit=dify_config.UPLOAD_VIDEO_FILE_SIZE_LIMIT,
|
||||
audio_file_size_limit=dify_config.UPLOAD_AUDIO_FILE_SIZE_LIMIT,
|
||||
workflow_file_upload_limit=dify_config.WORKFLOW_FILE_UPLOAD_LIMIT,
|
||||
image_file_batch_limit=dify_config.IMAGE_FILE_BATCH_LIMIT,
|
||||
single_chunk_attachment_limit=dify_config.SINGLE_CHUNK_ATTACHMENT_LIMIT,
|
||||
attachment_image_file_size_limit=dify_config.ATTACHMENT_IMAGE_FILE_SIZE_LIMIT,
|
||||
)
|
||||
return config.model_dump(mode="json"), 200
|
||||
|
||||
@setup_required
|
||||
@login_required
|
||||
@account_initialization_required
|
||||
@marshal_with(file_fields)
|
||||
@cloud_edition_billing_resource_check("documents")
|
||||
@console_ns.response(201, "File uploaded successfully", console_ns.models[FileResponse.__name__])
|
||||
def post(self):
|
||||
current_user, _ = current_account_with_tenant()
|
||||
source_str = request.form.get("source")
|
||||
|
|
@ -90,7 +94,8 @@ class FileApi(Resource):
|
|||
except services.errors.file.BlockedFileExtensionError as blocked_extension_error:
|
||||
raise BlockedFileExtensionError(blocked_extension_error.description)
|
||||
|
||||
return upload_file, 201
|
||||
response = FileResponse.model_validate(upload_file, from_attributes=True)
|
||||
return response.model_dump(mode="json"), 201
|
||||
|
||||
|
||||
@console_ns.route("/files/<uuid:file_id>/preview")
|
||||
|
|
|
|||
|
|
@ -1,7 +1,7 @@
|
|||
import urllib.parse
|
||||
|
||||
import httpx
|
||||
from flask_restx import Resource, marshal_with
|
||||
from flask_restx import Resource
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
import services
|
||||
|
|
@ -11,19 +11,22 @@ from controllers.common.errors import (
|
|||
RemoteFileUploadError,
|
||||
UnsupportedFileTypeError,
|
||||
)
|
||||
from controllers.common.schema import register_schema_models
|
||||
from core.file import helpers as file_helpers
|
||||
from core.helper import ssrf_proxy
|
||||
from extensions.ext_database import db
|
||||
from fields.file_fields import file_fields_with_signed_url, remote_file_info_fields
|
||||
from fields.file_fields import FileWithSignedUrl, RemoteFileInfo
|
||||
from libs.login import current_account_with_tenant
|
||||
from services.file_service import FileService
|
||||
|
||||
from . import console_ns
|
||||
|
||||
register_schema_models(console_ns, RemoteFileInfo, FileWithSignedUrl)
|
||||
|
||||
|
||||
@console_ns.route("/remote-files/<path:url>")
|
||||
class RemoteFileInfoApi(Resource):
|
||||
@marshal_with(remote_file_info_fields)
|
||||
@console_ns.response(200, "Remote file info", console_ns.models[RemoteFileInfo.__name__])
|
||||
def get(self, url):
|
||||
decoded_url = urllib.parse.unquote(url)
|
||||
resp = ssrf_proxy.head(decoded_url)
|
||||
|
|
@ -31,10 +34,11 @@ class RemoteFileInfoApi(Resource):
|
|||
# failed back to get method
|
||||
resp = ssrf_proxy.get(decoded_url, timeout=3)
|
||||
resp.raise_for_status()
|
||||
return {
|
||||
"file_type": resp.headers.get("Content-Type", "application/octet-stream"),
|
||||
"file_length": int(resp.headers.get("Content-Length", 0)),
|
||||
}
|
||||
info = RemoteFileInfo(
|
||||
file_type=resp.headers.get("Content-Type", "application/octet-stream"),
|
||||
file_length=int(resp.headers.get("Content-Length", 0)),
|
||||
)
|
||||
return info.model_dump(mode="json")
|
||||
|
||||
|
||||
class RemoteFileUploadPayload(BaseModel):
|
||||
|
|
@ -50,7 +54,7 @@ console_ns.schema_model(
|
|||
@console_ns.route("/remote-files/upload")
|
||||
class RemoteFileUploadApi(Resource):
|
||||
@console_ns.expect(console_ns.models[RemoteFileUploadPayload.__name__])
|
||||
@marshal_with(file_fields_with_signed_url)
|
||||
@console_ns.response(201, "Remote file uploaded", console_ns.models[FileWithSignedUrl.__name__])
|
||||
def post(self):
|
||||
args = RemoteFileUploadPayload.model_validate(console_ns.payload)
|
||||
url = args.url
|
||||
|
|
@ -85,13 +89,14 @@ class RemoteFileUploadApi(Resource):
|
|||
except services.errors.file.UnsupportedFileTypeError:
|
||||
raise UnsupportedFileTypeError()
|
||||
|
||||
return {
|
||||
"id": upload_file.id,
|
||||
"name": upload_file.name,
|
||||
"size": upload_file.size,
|
||||
"extension": upload_file.extension,
|
||||
"url": file_helpers.get_signed_file_url(upload_file_id=upload_file.id),
|
||||
"mime_type": upload_file.mime_type,
|
||||
"created_by": upload_file.created_by,
|
||||
"created_at": upload_file.created_at,
|
||||
}, 201
|
||||
payload = FileWithSignedUrl(
|
||||
id=upload_file.id,
|
||||
name=upload_file.name,
|
||||
size=upload_file.size,
|
||||
extension=upload_file.extension,
|
||||
url=file_helpers.get_signed_file_url(upload_file_id=upload_file.id),
|
||||
mime_type=upload_file.mime_type,
|
||||
created_by=upload_file.created_by,
|
||||
created_at=int(upload_file.created_at.timestamp()),
|
||||
)
|
||||
return payload.model_dump(mode="json"), 201
|
||||
|
|
|
|||
|
|
@ -4,18 +4,18 @@ from flask import request
|
|||
from flask_restx import Resource
|
||||
from flask_restx.api import HTTPStatus
|
||||
from pydantic import BaseModel, Field
|
||||
from werkzeug.datastructures import FileStorage
|
||||
from werkzeug.exceptions import Forbidden
|
||||
|
||||
import services
|
||||
from core.file.helpers import verify_plugin_file_signature
|
||||
from core.tools.tool_file_manager import ToolFileManager
|
||||
from fields.file_fields import build_file_model
|
||||
from fields.file_fields import FileResponse
|
||||
|
||||
from ..common.errors import (
|
||||
FileTooLargeError,
|
||||
UnsupportedFileTypeError,
|
||||
)
|
||||
from ..common.schema import register_schema_models
|
||||
from ..console.wraps import setup_required
|
||||
from ..files import files_ns
|
||||
from ..inner_api.plugin.wraps import get_user
|
||||
|
|
@ -35,6 +35,8 @@ files_ns.schema_model(
|
|||
PluginUploadQuery.__name__, PluginUploadQuery.model_json_schema(ref_template=DEFAULT_REF_TEMPLATE_SWAGGER_2_0)
|
||||
)
|
||||
|
||||
register_schema_models(files_ns, FileResponse)
|
||||
|
||||
|
||||
@files_ns.route("/upload/for-plugin")
|
||||
class PluginUploadFileApi(Resource):
|
||||
|
|
@ -51,7 +53,7 @@ class PluginUploadFileApi(Resource):
|
|||
415: "Unsupported file type",
|
||||
}
|
||||
)
|
||||
@files_ns.marshal_with(build_file_model(files_ns), code=HTTPStatus.CREATED)
|
||||
@files_ns.response(HTTPStatus.CREATED, "File uploaded", files_ns.models[FileResponse.__name__])
|
||||
def post(self):
|
||||
"""Upload a file for plugin usage.
|
||||
|
||||
|
|
@ -69,7 +71,7 @@ class PluginUploadFileApi(Resource):
|
|||
"""
|
||||
args = PluginUploadQuery.model_validate(request.args.to_dict(flat=True)) # type: ignore
|
||||
|
||||
file: FileStorage | None = request.files.get("file")
|
||||
file = request.files.get("file")
|
||||
if file is None:
|
||||
raise Forbidden("File is required.")
|
||||
|
||||
|
|
@ -80,8 +82,8 @@ class PluginUploadFileApi(Resource):
|
|||
user_id = args.user_id
|
||||
user = get_user(tenant_id, user_id)
|
||||
|
||||
filename: str | None = file.filename
|
||||
mimetype: str | None = file.mimetype
|
||||
filename = file.filename
|
||||
mimetype = file.mimetype
|
||||
|
||||
if not filename or not mimetype:
|
||||
raise Forbidden("Invalid request.")
|
||||
|
|
@ -111,22 +113,22 @@ class PluginUploadFileApi(Resource):
|
|||
preview_url = ToolFileManager.sign_file(tool_file_id=tool_file.id, extension=extension)
|
||||
|
||||
# Create a dictionary with all the necessary attributes
|
||||
result = {
|
||||
"id": tool_file.id,
|
||||
"user_id": tool_file.user_id,
|
||||
"tenant_id": tool_file.tenant_id,
|
||||
"conversation_id": tool_file.conversation_id,
|
||||
"file_key": tool_file.file_key,
|
||||
"mimetype": tool_file.mimetype,
|
||||
"original_url": tool_file.original_url,
|
||||
"name": tool_file.name,
|
||||
"size": tool_file.size,
|
||||
"mime_type": mimetype,
|
||||
"extension": extension,
|
||||
"preview_url": preview_url,
|
||||
}
|
||||
result = FileResponse(
|
||||
id=tool_file.id,
|
||||
name=tool_file.name,
|
||||
size=tool_file.size,
|
||||
extension=extension,
|
||||
mime_type=mimetype,
|
||||
preview_url=preview_url,
|
||||
source_url=tool_file.original_url,
|
||||
original_url=tool_file.original_url,
|
||||
user_id=tool_file.user_id,
|
||||
tenant_id=tool_file.tenant_id,
|
||||
conversation_id=tool_file.conversation_id,
|
||||
file_key=tool_file.file_key,
|
||||
)
|
||||
|
||||
return result, 201
|
||||
return result.model_dump(mode="json"), 201
|
||||
except services.errors.file.FileTooLargeError as file_too_large_error:
|
||||
raise FileTooLargeError(file_too_large_error.description)
|
||||
except services.errors.file.UnsupportedFileTypeError:
|
||||
|
|
|
|||
|
|
@ -10,13 +10,16 @@ from controllers.common.errors import (
|
|||
TooManyFilesError,
|
||||
UnsupportedFileTypeError,
|
||||
)
|
||||
from controllers.common.schema import register_schema_models
|
||||
from controllers.service_api import service_api_ns
|
||||
from controllers.service_api.wraps import FetchUserArg, WhereisUserArg, validate_app_token
|
||||
from extensions.ext_database import db
|
||||
from fields.file_fields import build_file_model
|
||||
from fields.file_fields import FileResponse
|
||||
from models import App, EndUser
|
||||
from services.file_service import FileService
|
||||
|
||||
register_schema_models(service_api_ns, FileResponse)
|
||||
|
||||
|
||||
@service_api_ns.route("/files/upload")
|
||||
class FileApi(Resource):
|
||||
|
|
@ -31,8 +34,8 @@ class FileApi(Resource):
|
|||
415: "Unsupported file type",
|
||||
}
|
||||
)
|
||||
@validate_app_token(fetch_user_arg=FetchUserArg(fetch_from=WhereisUserArg.FORM))
|
||||
@service_api_ns.marshal_with(build_file_model(service_api_ns), code=HTTPStatus.CREATED)
|
||||
@validate_app_token(fetch_user_arg=FetchUserArg(fetch_from=WhereisUserArg.FORM)) # type: ignore
|
||||
@service_api_ns.response(HTTPStatus.CREATED, "File uploaded", service_api_ns.models[FileResponse.__name__])
|
||||
def post(self, app_model: App, end_user: EndUser):
|
||||
"""Upload a file for use in conversations.
|
||||
|
||||
|
|
@ -64,4 +67,5 @@ class FileApi(Resource):
|
|||
except services.errors.file.UnsupportedFileTypeError:
|
||||
raise UnsupportedFileTypeError()
|
||||
|
||||
return upload_file, 201
|
||||
response = FileResponse.model_validate(upload_file, from_attributes=True)
|
||||
return response.model_dump(mode="json"), 201
|
||||
|
|
|
|||
|
|
@ -1,5 +1,4 @@
|
|||
from flask import request
|
||||
from flask_restx import marshal_with
|
||||
|
||||
import services
|
||||
from controllers.common.errors import (
|
||||
|
|
@ -9,12 +8,15 @@ from controllers.common.errors import (
|
|||
TooManyFilesError,
|
||||
UnsupportedFileTypeError,
|
||||
)
|
||||
from controllers.common.schema import register_schema_models
|
||||
from controllers.web import web_ns
|
||||
from controllers.web.wraps import WebApiResource
|
||||
from extensions.ext_database import db
|
||||
from fields.file_fields import build_file_model
|
||||
from fields.file_fields import FileResponse
|
||||
from services.file_service import FileService
|
||||
|
||||
register_schema_models(web_ns, FileResponse)
|
||||
|
||||
|
||||
@web_ns.route("/files/upload")
|
||||
class FileApi(WebApiResource):
|
||||
|
|
@ -28,7 +30,7 @@ class FileApi(WebApiResource):
|
|||
415: "Unsupported file type",
|
||||
}
|
||||
)
|
||||
@marshal_with(build_file_model(web_ns))
|
||||
@web_ns.response(201, "File uploaded successfully", web_ns.models[FileResponse.__name__])
|
||||
def post(self, app_model, end_user):
|
||||
"""Upload a file for use in web applications.
|
||||
|
||||
|
|
@ -81,4 +83,5 @@ class FileApi(WebApiResource):
|
|||
except services.errors.file.UnsupportedFileTypeError:
|
||||
raise UnsupportedFileTypeError()
|
||||
|
||||
return upload_file, 201
|
||||
response = FileResponse.model_validate(upload_file, from_attributes=True)
|
||||
return response.model_dump(mode="json"), 201
|
||||
|
|
|
|||
|
|
@ -1,7 +1,6 @@
|
|||
import urllib.parse
|
||||
|
||||
import httpx
|
||||
from flask_restx import marshal_with
|
||||
from pydantic import BaseModel, Field, HttpUrl
|
||||
|
||||
import services
|
||||
|
|
@ -14,7 +13,7 @@ from controllers.common.errors import (
|
|||
from core.file import helpers as file_helpers
|
||||
from core.helper import ssrf_proxy
|
||||
from extensions.ext_database import db
|
||||
from fields.file_fields import build_file_with_signed_url_model, build_remote_file_info_model
|
||||
from fields.file_fields import FileWithSignedUrl, RemoteFileInfo
|
||||
from services.file_service import FileService
|
||||
|
||||
from ..common.schema import register_schema_models
|
||||
|
|
@ -26,7 +25,7 @@ class RemoteFileUploadPayload(BaseModel):
|
|||
url: HttpUrl = Field(description="Remote file URL")
|
||||
|
||||
|
||||
register_schema_models(web_ns, RemoteFileUploadPayload)
|
||||
register_schema_models(web_ns, RemoteFileUploadPayload, RemoteFileInfo, FileWithSignedUrl)
|
||||
|
||||
|
||||
@web_ns.route("/remote-files/<path:url>")
|
||||
|
|
@ -41,7 +40,7 @@ class RemoteFileInfoApi(WebApiResource):
|
|||
500: "Failed to fetch remote file",
|
||||
}
|
||||
)
|
||||
@marshal_with(build_remote_file_info_model(web_ns))
|
||||
@web_ns.response(200, "Remote file info", web_ns.models[RemoteFileInfo.__name__])
|
||||
def get(self, app_model, end_user, url):
|
||||
"""Get information about a remote file.
|
||||
|
||||
|
|
@ -65,10 +64,11 @@ class RemoteFileInfoApi(WebApiResource):
|
|||
# failed back to get method
|
||||
resp = ssrf_proxy.get(decoded_url, timeout=3)
|
||||
resp.raise_for_status()
|
||||
return {
|
||||
"file_type": resp.headers.get("Content-Type", "application/octet-stream"),
|
||||
"file_length": int(resp.headers.get("Content-Length", -1)),
|
||||
}
|
||||
info = RemoteFileInfo(
|
||||
file_type=resp.headers.get("Content-Type", "application/octet-stream"),
|
||||
file_length=int(resp.headers.get("Content-Length", -1)),
|
||||
)
|
||||
return info.model_dump(mode="json")
|
||||
|
||||
|
||||
@web_ns.route("/remote-files/upload")
|
||||
|
|
@ -84,7 +84,7 @@ class RemoteFileUploadApi(WebApiResource):
|
|||
500: "Failed to fetch remote file",
|
||||
}
|
||||
)
|
||||
@marshal_with(build_file_with_signed_url_model(web_ns))
|
||||
@web_ns.response(201, "Remote file uploaded", web_ns.models[FileWithSignedUrl.__name__])
|
||||
def post(self, app_model, end_user):
|
||||
"""Upload a file from a remote URL.
|
||||
|
||||
|
|
@ -139,13 +139,14 @@ class RemoteFileUploadApi(WebApiResource):
|
|||
except services.errors.file.UnsupportedFileTypeError:
|
||||
raise UnsupportedFileTypeError
|
||||
|
||||
return {
|
||||
"id": upload_file.id,
|
||||
"name": upload_file.name,
|
||||
"size": upload_file.size,
|
||||
"extension": upload_file.extension,
|
||||
"url": file_helpers.get_signed_file_url(upload_file_id=upload_file.id),
|
||||
"mime_type": upload_file.mime_type,
|
||||
"created_by": upload_file.created_by,
|
||||
"created_at": upload_file.created_at,
|
||||
}, 201
|
||||
payload1 = FileWithSignedUrl(
|
||||
id=upload_file.id,
|
||||
name=upload_file.name,
|
||||
size=upload_file.size,
|
||||
extension=upload_file.extension,
|
||||
url=file_helpers.get_signed_file_url(upload_file_id=upload_file.id),
|
||||
mime_type=upload_file.mime_type,
|
||||
created_by=upload_file.created_by,
|
||||
created_at=int(upload_file.created_at.timestamp()),
|
||||
)
|
||||
return payload1.model_dump(mode="json"), 201
|
||||
|
|
|
|||
|
|
@ -1,93 +1,85 @@
|
|||
from flask_restx import Namespace, fields
|
||||
from __future__ import annotations
|
||||
|
||||
from libs.helper import TimestampField
|
||||
from datetime import datetime
|
||||
|
||||
upload_config_fields = {
|
||||
"file_size_limit": fields.Integer,
|
||||
"batch_count_limit": fields.Integer,
|
||||
"image_file_size_limit": fields.Integer,
|
||||
"video_file_size_limit": fields.Integer,
|
||||
"audio_file_size_limit": fields.Integer,
|
||||
"workflow_file_upload_limit": fields.Integer,
|
||||
"image_file_batch_limit": fields.Integer,
|
||||
"single_chunk_attachment_limit": fields.Integer,
|
||||
}
|
||||
from pydantic import BaseModel, ConfigDict, field_validator
|
||||
|
||||
|
||||
def build_upload_config_model(api_or_ns: Namespace):
|
||||
"""Build the upload config model for the API or Namespace.
|
||||
|
||||
Args:
|
||||
api_or_ns: Flask-RestX Api or Namespace instance
|
||||
|
||||
Returns:
|
||||
The registered model
|
||||
"""
|
||||
return api_or_ns.model("UploadConfig", upload_config_fields)
|
||||
class ResponseModel(BaseModel):
|
||||
model_config = ConfigDict(
|
||||
from_attributes=True,
|
||||
extra="ignore",
|
||||
populate_by_name=True,
|
||||
serialize_by_alias=True,
|
||||
protected_namespaces=(),
|
||||
)
|
||||
|
||||
|
||||
file_fields = {
|
||||
"id": fields.String,
|
||||
"name": fields.String,
|
||||
"size": fields.Integer,
|
||||
"extension": fields.String,
|
||||
"mime_type": fields.String,
|
||||
"created_by": fields.String,
|
||||
"created_at": TimestampField,
|
||||
"preview_url": fields.String,
|
||||
"source_url": fields.String,
|
||||
}
|
||||
def _to_timestamp(value: datetime | int | None) -> int | None:
|
||||
if isinstance(value, datetime):
|
||||
return int(value.timestamp())
|
||||
return value
|
||||
|
||||
|
||||
def build_file_model(api_or_ns: Namespace):
|
||||
"""Build the file model for the API or Namespace.
|
||||
|
||||
Args:
|
||||
api_or_ns: Flask-RestX Api or Namespace instance
|
||||
|
||||
Returns:
|
||||
The registered model
|
||||
"""
|
||||
return api_or_ns.model("File", file_fields)
|
||||
class UploadConfig(ResponseModel):
|
||||
file_size_limit: int
|
||||
batch_count_limit: int
|
||||
file_upload_limit: int | None = None
|
||||
image_file_size_limit: int
|
||||
video_file_size_limit: int
|
||||
audio_file_size_limit: int
|
||||
workflow_file_upload_limit: int
|
||||
image_file_batch_limit: int
|
||||
single_chunk_attachment_limit: int
|
||||
attachment_image_file_size_limit: int | None = None
|
||||
|
||||
|
||||
remote_file_info_fields = {
|
||||
"file_type": fields.String(attribute="file_type"),
|
||||
"file_length": fields.Integer(attribute="file_length"),
|
||||
}
|
||||
class FileResponse(ResponseModel):
|
||||
id: str
|
||||
name: str
|
||||
size: int
|
||||
extension: str | None = None
|
||||
mime_type: str | None = None
|
||||
created_by: str | None = None
|
||||
created_at: int | None = None
|
||||
preview_url: str | None = None
|
||||
source_url: str | None = None
|
||||
original_url: str | None = None
|
||||
user_id: str | None = None
|
||||
tenant_id: str | None = None
|
||||
conversation_id: str | None = None
|
||||
file_key: str | None = None
|
||||
|
||||
@field_validator("created_at", mode="before")
|
||||
@classmethod
|
||||
def _normalize_created_at(cls, value: datetime | int | None) -> int | None:
|
||||
return _to_timestamp(value)
|
||||
|
||||
|
||||
def build_remote_file_info_model(api_or_ns: Namespace):
|
||||
"""Build the remote file info model for the API or Namespace.
|
||||
|
||||
Args:
|
||||
api_or_ns: Flask-RestX Api or Namespace instance
|
||||
|
||||
Returns:
|
||||
The registered model
|
||||
"""
|
||||
return api_or_ns.model("RemoteFileInfo", remote_file_info_fields)
|
||||
class RemoteFileInfo(ResponseModel):
|
||||
file_type: str
|
||||
file_length: int
|
||||
|
||||
|
||||
file_fields_with_signed_url = {
|
||||
"id": fields.String,
|
||||
"name": fields.String,
|
||||
"size": fields.Integer,
|
||||
"extension": fields.String,
|
||||
"url": fields.String,
|
||||
"mime_type": fields.String,
|
||||
"created_by": fields.String,
|
||||
"created_at": TimestampField,
|
||||
}
|
||||
class FileWithSignedUrl(ResponseModel):
|
||||
id: str
|
||||
name: str
|
||||
size: int
|
||||
extension: str | None = None
|
||||
url: str | None = None
|
||||
mime_type: str | None = None
|
||||
created_by: str | None = None
|
||||
created_at: int | None = None
|
||||
|
||||
@field_validator("created_at", mode="before")
|
||||
@classmethod
|
||||
def _normalize_created_at(cls, value: datetime | int | None) -> int | None:
|
||||
return _to_timestamp(value)
|
||||
|
||||
|
||||
def build_file_with_signed_url_model(api_or_ns: Namespace):
|
||||
"""Build the file with signed URL model for the API or Namespace.
|
||||
|
||||
Args:
|
||||
api_or_ns: Flask-RestX Api or Namespace instance
|
||||
|
||||
Returns:
|
||||
The registered model
|
||||
"""
|
||||
return api_or_ns.model("FileWithSignedUrl", file_fields_with_signed_url)
|
||||
__all__ = [
|
||||
"FileResponse",
|
||||
"FileWithSignedUrl",
|
||||
"RemoteFileInfo",
|
||||
"UploadConfig",
|
||||
]
|
||||
|
|
|
|||
|
|
@ -1,7 +1,9 @@
|
|||
import builtins
|
||||
import io
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
from flask.views import MethodView
|
||||
from werkzeug.exceptions import Forbidden
|
||||
|
||||
from controllers.common.errors import (
|
||||
|
|
@ -14,6 +16,9 @@ from controllers.common.errors import (
|
|||
from services.errors.file import FileTooLargeError as ServiceFileTooLargeError
|
||||
from services.errors.file import UnsupportedFileTypeError as ServiceUnsupportedFileTypeError
|
||||
|
||||
if not hasattr(builtins, "MethodView"):
|
||||
builtins.MethodView = MethodView # type: ignore[attr-defined]
|
||||
|
||||
|
||||
class TestFileUploadSecurity:
|
||||
"""Test file upload security logic without complex framework setup"""
|
||||
|
|
@ -128,7 +133,7 @@ class TestFileUploadSecurity:
|
|||
# Test passes if no exception is raised
|
||||
|
||||
# Test 4: Service error handling
|
||||
@patch("services.file_service.FileService.upload_file")
|
||||
@patch("controllers.console.files.FileService.upload_file")
|
||||
def test_should_handle_file_too_large_error(self, mock_upload):
|
||||
"""Test that service FileTooLargeError is properly converted"""
|
||||
mock_upload.side_effect = ServiceFileTooLargeError("File too large")
|
||||
|
|
@ -140,7 +145,7 @@ class TestFileUploadSecurity:
|
|||
with pytest.raises(FileTooLargeError):
|
||||
raise FileTooLargeError(e.description)
|
||||
|
||||
@patch("services.file_service.FileService.upload_file")
|
||||
@patch("controllers.console.files.FileService.upload_file")
|
||||
def test_should_handle_unsupported_file_type_error(self, mock_upload):
|
||||
"""Test that service UnsupportedFileTypeError is properly converted"""
|
||||
mock_upload.side_effect = ServiceUnsupportedFileTypeError()
|
||||
|
|
|
|||
|
|
@ -0,0 +1,78 @@
|
|||
from __future__ import annotations
|
||||
|
||||
from datetime import datetime
|
||||
from types import SimpleNamespace
|
||||
|
||||
from fields.file_fields import FileResponse, FileWithSignedUrl, RemoteFileInfo, UploadConfig
|
||||
|
||||
|
||||
def test_file_response_serializes_datetime() -> None:
|
||||
created_at = datetime(2024, 1, 1, 12, 0, 0)
|
||||
file_obj = SimpleNamespace(
|
||||
id="file-1",
|
||||
name="example.txt",
|
||||
size=1024,
|
||||
extension="txt",
|
||||
mime_type="text/plain",
|
||||
created_by="user-1",
|
||||
created_at=created_at,
|
||||
preview_url="https://preview",
|
||||
source_url="https://source",
|
||||
original_url="https://origin",
|
||||
user_id="user-1",
|
||||
tenant_id="tenant-1",
|
||||
conversation_id="conv-1",
|
||||
file_key="key-1",
|
||||
)
|
||||
|
||||
serialized = FileResponse.model_validate(file_obj, from_attributes=True).model_dump(mode="json")
|
||||
|
||||
assert serialized["id"] == "file-1"
|
||||
assert serialized["created_at"] == int(created_at.timestamp())
|
||||
assert serialized["preview_url"] == "https://preview"
|
||||
assert serialized["source_url"] == "https://source"
|
||||
assert serialized["original_url"] == "https://origin"
|
||||
assert serialized["user_id"] == "user-1"
|
||||
assert serialized["tenant_id"] == "tenant-1"
|
||||
assert serialized["conversation_id"] == "conv-1"
|
||||
assert serialized["file_key"] == "key-1"
|
||||
|
||||
|
||||
def test_file_with_signed_url_builds_payload() -> None:
|
||||
payload = FileWithSignedUrl(
|
||||
id="file-2",
|
||||
name="remote.pdf",
|
||||
size=2048,
|
||||
extension="pdf",
|
||||
url="https://signed",
|
||||
mime_type="application/pdf",
|
||||
created_by="user-2",
|
||||
created_at=datetime(2024, 1, 2, 0, 0, 0),
|
||||
)
|
||||
|
||||
dumped = payload.model_dump(mode="json")
|
||||
|
||||
assert dumped["url"] == "https://signed"
|
||||
assert dumped["created_at"] == int(datetime(2024, 1, 2, 0, 0, 0).timestamp())
|
||||
|
||||
|
||||
def test_remote_file_info_and_upload_config() -> None:
|
||||
info = RemoteFileInfo(file_type="text/plain", file_length=123)
|
||||
assert info.model_dump(mode="json") == {"file_type": "text/plain", "file_length": 123}
|
||||
|
||||
config = UploadConfig(
|
||||
file_size_limit=1,
|
||||
batch_count_limit=2,
|
||||
file_upload_limit=3,
|
||||
image_file_size_limit=4,
|
||||
video_file_size_limit=5,
|
||||
audio_file_size_limit=6,
|
||||
workflow_file_upload_limit=7,
|
||||
image_file_batch_limit=8,
|
||||
single_chunk_attachment_limit=9,
|
||||
attachment_image_file_size_limit=10,
|
||||
)
|
||||
|
||||
dumped = config.model_dump(mode="json")
|
||||
assert dumped["file_upload_limit"] == 3
|
||||
assert dumped["attachment_image_file_size_limit"] == 10
|
||||
Loading…
Reference in New Issue