diff --git a/.runtime.env b/.runtime.env index 6595d0c..dbc5e5b 100644 --- a/.runtime.env +++ b/.runtime.env @@ -1,4 +1,4 @@ -MODAL_ENVIRONMENT=test +MODAL_ENVIRONMENT=dev modal_app_name=bowong-ai-video S3_mount_dir=/mntS3 S3_bucket_name=modal-media-cache diff --git a/src/BowongModalFunctions/api.py b/src/BowongModalFunctions/api.py index d2b796c..e2617ab 100644 --- a/src/BowongModalFunctions/api.py +++ b/src/BowongModalFunctions/api.py @@ -5,7 +5,7 @@ from sentry_sdk.integrations.loguru import LoguruIntegration, LoggingLevels from sentry_sdk.integrations.fastapi import FastApiIntegration from fastapi.middleware.cors import CORSMiddleware from .utils.KVCache import KVCache -from .router import ffmpeg, cache, comfyui +from .router import ffmpeg, cache, comfyui, google from .config import WorkerConfig config = WorkerConfig() @@ -82,3 +82,4 @@ async def scalar(): web_app.include_router(ffmpeg.router) web_app.include_router(cache.router) web_app.include_router(comfyui.router) +web_app.include_router(google.router) diff --git a/src/BowongModalFunctions/router/google.py b/src/BowongModalFunctions/router/google.py new file mode 100644 index 0000000..7715b3d --- /dev/null +++ b/src/BowongModalFunctions/router/google.py @@ -0,0 +1,98 @@ +import os +from typing import Annotated, Optional + +from loguru import logger +import httpx +from fastapi import APIRouter, UploadFile, Header, HTTPException +from pydantic import BaseModel, Field +from starlette import status +from starlette.responses import JSONResponse + +from BowongModalFunctions.config import WorkerConfig + +config = WorkerConfig + +router = APIRouter(prefix="/google", tags=["google"]) + + +class GoogleAPIKeyHeaders(BaseModel): + x_google_api_key: Optional[str] = Field(description="Google API Key", default=None) + + +@router.post("/upload", + summary="上传文件到Google File", + description="上传文件到Google File, 换取Google File URI, 不同Google API Key之间的URI不互通, 最多可为每个项目存储 20 GB 的文件，每个文件的大小上限为 2 GB。文件会存储 48 小时") +async def upload_file_multipart(file: UploadFile, + headers: Annotated[GoogleAPIKeyHeaders, Header()]): + google_api_key = headers.x_google_api_key or os.environ.get("GOOGLE_API_KEY") + if not google_api_key: + raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail="Missing Google API Key") + content_length = file.size + content_type = file.content_type + if content_type not in ['video/mp4', 'video/mpeg', 'video/mov', 'video/avi', 'video/x-flv', 'video/mpg', + 'video/webm', 'video/wmv', 'video/3gpp']: + raise HTTPException(status_code=status.HTTP_415_UNSUPPORTED_MEDIA_TYPE) + logger.info(f"Uploading name = {file.filename}, size = {content_length}, type = {content_type} to google file") + with httpx.Client() as client: + pre_upload_response = client.post( + url=f"https://generativelanguage.googleapis.com/upload/v1beta/files?key={google_api_key}", + headers={ + "X-Goog-Upload-Protocol": "resumable", + "X-Goog-Upload-Command": "start", + "X-Goog-Upload-Header-Content-Length": str(content_length), + "X-Goog-Upload-Header-Content-Type": content_type + }, + json={ + "file": { + "display_name": file.filename.split(".")[0], + } + }) + pre_upload_response.raise_for_status() + + upload_url = pre_upload_response.headers.get("X-Goog-Upload-Url") + + upload_response = client.post(url=upload_url, content=file.file.read(), headers={ + "X-Goog-Upload-Offset": "0", + "X-Goog-Upload-Command": "upload, finalize", + "Content-Type": content_type + }) + upload_response.raise_for_status() + + return JSONResponse(content=upload_response.json(), status_code=upload_response.status_code) + + +@router.get("/status", summary="获取已上传文件的处理状态") +async def uploaded_file_status(filename: str, + headers: Annotated[GoogleAPIKeyHeaders, Header()]): + google_api_key = headers.x_google_api_key or os.environ.get("GOOGLE_API_KEY") + if not google_api_key: + raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail="Missing Google API Key") + with httpx.Client() as client: + response = client.get( + url=f"https://generativelanguage.googleapis.com/v1beta/files/{filename}?key={google_api_key}") + response.raise_for_status() + return JSONResponse(content=response.json(), status_code=response.status_code) + + +@router.delete('/delete', summary="删除已上传的文件") +async def delete_file(filename: str, headers: Annotated[GoogleAPIKeyHeaders, Header()]): + google_api_key = headers.x_google_api_key or os.environ.get("GOOGLE_API_KEY") + if not google_api_key: + raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail="Missing Google API Key") + with httpx.Client() as client: + response = client.delete( + url=f"https://generativelanguage.googleapis.com/v1beta/files/{filename}?key={google_api_key}") + response.raise_for_status() + return JSONResponse(content=response.json(), status_code=response.status_code) + + +@router.get('/list', summary="列出已上传的文件") +async def list_files(headers: Annotated[GoogleAPIKeyHeaders, Header()]): + google_api_key = headers.x_google_api_key or os.environ.get("GOOGLE_API_KEY") + if not google_api_key: + raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail="Missing Google API Key") + with httpx.Client() as client: + response = client.get( + url=f"https://generativelanguage.googleapis.com/v1beta/files?key={google_api_key}") + response.raise_for_status() + return JSONResponse(content=response.json(), status_code=response.status_code) diff --git a/src/cluster/web.py b/src/cluster/web.py index 75d7842..37f2a54 100644 --- a/src/cluster/web.py +++ b/src/cluster/web.py @@ -13,6 +13,9 @@ fastapi_image = ( app = modal.App( name="web_app", image=fastapi_image, + secrets=[ + modal.Secret.from_name('google-secret') + ], include_source=False) with fastapi_image.imports(): @@ -21,6 +24,7 @@ with fastapi_image.imports(): config = WorkerConfig() + @app.function(scaledown_window=60, secrets=[ modal.Secret.from_name("cf-kv-secret", environment_name='dev'),