2023-06-11 00:59:16 +03:00
import os
2023-06-28 20:39:27 +03:00
from uuid import UUID
2023-06-11 00:59:16 +03:00
2023-06-14 22:21:13 +03:00
from auth . auth_bearer import AuthBearer , get_current_user
2023-06-28 20:39:27 +03:00
from fastapi import APIRouter , Depends , Query , Request , UploadFile
from models . brains import Brain
from models . files import File
from models . settings import common_dependencies
2023-06-11 00:59:16 +03:00
from models . users import User
from utils . file import convert_bytes , get_file_size
from utils . processors import filter_file
upload_router = APIRouter ( )
2023-06-15 15:43:40 +03:00
@upload_router.post ( " /upload " , dependencies = [ Depends ( AuthBearer ( ) ) ] , tags = [ " Upload " ] )
2023-06-28 20:39:27 +03:00
async def upload_file ( request : Request , uploadFile : UploadFile , brain_id : UUID = Query ( . . . , description = " The ID of the brain " ) , enable_summarization : bool = False , current_user : User = Depends ( get_current_user ) ) :
2023-06-15 15:43:40 +03:00
"""
Upload a file to the user ' s storage.
- ` file ` : The file to be uploaded .
- ` enable_summarization ` : Flag to enable summarization of the file ' s content.
- ` current_user ` : The current authenticated user .
- Returns the response message indicating the success or failure of the upload .
This endpoint allows users to upload files to their storage ( brain ) . It checks the remaining free space in the user ' s storage (brain)
and ensures that the file size does not exceed the maximum capacity . If the file is within the allowed size limit ,
it can optionally apply summarization to the file ' s content. The response message will indicate the status of the upload.
"""
2023-06-28 20:39:27 +03:00
print ( brain_id , " brain_id " )
# [TODO] check if the user is the owner/editor of the brain
brain = Brain ( id = brain_id )
print ( " brain " , brain )
2023-06-19 23:54:01 +03:00
commons = common_dependencies ( )
2023-06-11 00:59:16 +03:00
2023-06-28 20:39:27 +03:00
if request . headers . get ( ' Openai-Api-Key ' ) :
brain . max_brain_size = os . getenv ( " MAX_BRAIN_SIZE_WITH_KEY " , 209715200 )
remaining_free_space = brain . remaining_brain_size
2023-06-12 18:58:05 +03:00
2023-06-28 20:39:27 +03:00
file_size = get_file_size ( uploadFile )
file = File ( file = uploadFile )
2023-06-11 00:59:16 +03:00
if remaining_free_space - file_size < 0 :
message = { " message " : f " ❌ User ' s brain will exceed maximum capacity with this upload. Maximum file allowed is : { convert_bytes ( remaining_free_space ) } " , " type " : " error " }
else :
2023-06-28 20:39:27 +03:00
message = await filter_file ( commons , file , enable_summarization , brain_id = brain_id , openai_api_key = request . headers . get ( ' Openai-Api-Key ' , None ) )
2023-06-11 00:59:16 +03:00
return message