添加Google file upload api

* 添加Google file upload api * WIP : 增加Gemini 上传接口 --------- Merge request URL: https://g-ldyi2063.coding.net/p/dev/d/modalDeploy/git/merge/4741 Co-authored-by: shuohigh@gmail.com
2025-05-23 14:28:05 +08:00
parent 225a003d45
commit dc7b140c2d
4 changed files with 105 additions and 2 deletions
--- a/src/BowongModalFunctions/router/google.py
+++ 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)