diff --git a/src/BowongModalFunctions/models/media_model.py b/src/BowongModalFunctions/models/media_model.py index 9224c78..cc5660f 100644 --- a/src/BowongModalFunctions/models/media_model.py +++ b/src/BowongModalFunctions/models/media_model.py @@ -25,6 +25,7 @@ class MediaProtocol(str, Enum): vod = "vod" cos = "cos" hls = "hls" + gs = "gs" class MediaCacheStatus(str, Enum): diff --git a/src/BowongModalFunctions/models/web_model.py b/src/BowongModalFunctions/models/web_model.py index 7e947a9..f6aa0b7 100644 --- a/src/BowongModalFunctions/models/web_model.py +++ b/src/BowongModalFunctions/models/web_model.py @@ -1,4 +1,5 @@ import uuid +from datetime import datetime from enum import Enum from typing import List, Union, Optional, Dict, Any @@ -535,3 +536,31 @@ class MakeGridGeminiRequest(BaseFFMPEGTaskRequest): font_size: int = Field(default=18, description="文本尺寸/像素") padding: int = Field(default=5, description="文本距离文本框边缘距离/像素") separator: int = Field(default=12, description="分割线宽度/像素") + + +class GoogleUploadResponse(BaseModel): + kind: str + id: str + self_link: HttpUrl = Field(alias='selfLink') + media_link: HttpUrl = Field(alias='mediaLink') + name: str + bucket: str + generation: str + metageneration: str + content_type: str = Field(alias='contentType') + storage_class: str = Field(alias='storageClass') + size: int + md5_hash: str = Field(alias='md5Hash') + crc32c: str + etag: str + time_created: datetime = Field(alias='timeCreated') + updated: datetime + time_storage_class_updated: datetime = Field(alias='timeStorageClassUpdated') + time_finalized: datetime = Field(alias='timeFinalized') + + @computed_field(description="适用于Google内部服务的URN") + @property + def urn(self) -> str: + return self.id.replace(self.bucket, "gs:/") + + model_config = ConfigDict(populate_by_name=True) diff --git a/src/BowongModalFunctions/router/google.py b/src/BowongModalFunctions/router/google.py index a6f68f3..4611d46 100644 --- a/src/BowongModalFunctions/router/google.py +++ b/src/BowongModalFunctions/router/google.py @@ -16,7 +16,7 @@ from starlette.responses import JSONResponse from BowongModalFunctions.config import WorkerConfig from BowongModalFunctions.middleware.authorization import verify_token from BowongModalFunctions.models.web_model import SentryTransactionInfo, GeminiResultResponse, GeminiRequest, \ - ModalTaskResponse, MakeGridGeminiRequest + ModalTaskResponse, MakeGridGeminiRequest, GoogleUploadResponse from BowongModalFunctions.utils.ModalUtils import ModalUtils from BowongModalFunctions.utils.HTTPUtils import GoogleAuthUtils @@ -92,8 +92,15 @@ async def upload_file_multipart(file: UploadFile, @router.post("/vertex-ai/upload", summary="上传文件到Vertex AI可访问的Google cloud storage", ) -async def upload_file(file: UploadFile, headers: Annotated[GoogleAuthHeaders, Header()]): - raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED) +async def upload_file(bucket: str, file: UploadFile, + headers: Annotated[GoogleAuthHeaders, Header()], + prefix: Optional[str] = None) -> GoogleUploadResponse: + key = f"{prefix}/{file.filename}" if prefix else file.filename + return await GoogleAuthUtils.google_upload_file(file_stream=file.file, + content_type=file.content_type, + google_api_key=headers.auth_token, + bucket_name=bucket, filename=key) + @router.get("/status", summary="获取已上传文件的处理状态") async def uploaded_file_status(filename: str, diff --git a/src/BowongModalFunctions/utils/HTTPUtils.py b/src/BowongModalFunctions/utils/HTTPUtils.py index efa41ba..63fc949 100644 --- a/src/BowongModalFunctions/utils/HTTPUtils.py +++ b/src/BowongModalFunctions/utils/HTTPUtils.py @@ -1,4 +1,5 @@ from typing import Union, Dict, Any, BinaryIO +from urllib.parse import urlencode, quote import backoff import httpx @@ -9,6 +10,8 @@ from pathlib import Path from pydantic import BaseModel +from BowongModalFunctions.models.web_model import GoogleUploadResponse + class HTTPDownloadUtils: # 创建一个类级别的 Semaphore 来控制并发 @@ -127,6 +130,19 @@ class GoogleAuthUtils: response.raise_for_status() return GoogleAuthUtils.GoogleAuthResponse.model_validate_json(response.text) - # @staticmethod - # async def google_upload_file(file_stream: BinaryIO, google_api_key: str, bucket_name: str): - # + @staticmethod + async def google_upload_file(file_stream: BinaryIO, content_type: str, google_api_key: str, bucket_name: str, + filename: str) -> GoogleUploadResponse: + safe_filename = quote(filename) + + upload_response = httpx.post( + url=f"https://storage.googleapis.com/upload/storage/v1/b/{bucket_name}/o?uploadType=media&name={safe_filename}", + content=file_stream, + headers={ + "Authorization": f"Bearer {google_api_key}", + "Content-Type": content_type + }) + upload_response.raise_for_status() + + # upload_url = f"gs://dy-media-storage/video/{filename}" + return GoogleUploadResponse.model_validate_json(upload_response.text)