204 lines
11 KiB
Python
204 lines
11 KiB
Python
import json
|
|
import os
|
|
from typing import Annotated, Optional, cast, List
|
|
|
|
import modal
|
|
import sentry_sdk
|
|
from fastapi.responses import Response
|
|
from loguru import logger
|
|
import httpx
|
|
from fastapi import APIRouter, UploadFile, Header, HTTPException, Depends
|
|
from pydantic import BaseModel, Field
|
|
from starlette import status
|
|
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
|
|
from BowongModalFunctions.utils.ModalUtils import ModalUtils
|
|
from BowongModalFunctions.utils.HTTPUtils import GoogleAuthUtils
|
|
|
|
config = WorkerConfig()
|
|
|
|
router = APIRouter(prefix="/google", tags=["Google"])
|
|
|
|
|
|
class GoogleAPIKeyHeaders(BaseModel):
|
|
x_google_api_key: Optional[str] = Field(description="Google API Key")
|
|
|
|
|
|
class BundleHeaders(BaseModel):
|
|
x_google_api_key: Optional[str] = Field(description="Google API Key")
|
|
x_trace_id: str = Field(description="Sentry Transaction ID", default=None)
|
|
x_baggage: str = Field(description="Sentry Transaction baggage", 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
|
|
# 允许上传的视频
|
|
google_content_type_allow_list = ['video/mp4', 'video/mpeg', 'video/mov', 'video/avi', 'video/x-flv', 'video/mpg',
|
|
'video/webm', 'video/wmv', 'video/3gpp', ]
|
|
# 允许上传的图片
|
|
google_content_type_allow_list.extend(["image/png", "image/jpeg", "image/webp", "image/heic", "image/heif"])
|
|
if content_type not in google_content_type_allow_list:
|
|
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(timeout=1800) 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, 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.delete('/delete_all', summary="删除所有已上传的文件/第一页")
|
|
async def delete_all(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(timeout=60) as client:
|
|
resp = client.get(f"https://generativelanguage.googleapis.com/v1beta/files?key={google_api_key}")
|
|
if resp.status_code != 200:
|
|
return JSONResponse(content={"msg": "列出文件失败"}, status_code=resp.status_code)
|
|
filename_list = [i["name"].split("/")[1] for i in resp.json()["files"]]
|
|
for filename in filename_list:
|
|
logger.info(f"删除 {filename}...")
|
|
response = client.delete(
|
|
url=f"https://generativelanguage.googleapis.com/v1beta/files/{filename}?key={google_api_key}")
|
|
response.raise_for_status()
|
|
if resp.status_code != 200:
|
|
return JSONResponse(content={"msg": f"删除文件{filename}失败"}, status_code=resp.status_code)
|
|
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(timeout=60) as client:
|
|
response = client.get(
|
|
url=f"https://generativelanguage.googleapis.com/v1beta/files?key={google_api_key}")
|
|
if response.status_code != 200:
|
|
return JSONResponse(content={}, status_code=response.status_code)
|
|
return JSONResponse(content=response.json() if len(response.text) > 0 else "", status_code=response.status_code)
|
|
|
|
|
|
@router.get('/access-token', summary="获取一个新的Google Access Token", dependencies=[Depends(verify_token)])
|
|
async def get_access_token() -> GoogleAuthUtils.GoogleAuthResponse:
|
|
service_account_json = os.environ.get("GOOGLE_AUTH_JSON", None)
|
|
if not service_account_json:
|
|
raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail="Missing Google IAM Service Account JSON")
|
|
service_account_info = json.loads(service_account_json)
|
|
access_token = await GoogleAuthUtils.get_google_auth_jwt(service_account_info=service_account_info,
|
|
scopes=['https://www.googleapis.com/auth/cloud-platform'])
|
|
return access_token
|
|
|
|
|
|
@router.post("/make_grid_gemini", summary="将输入图拼为网格上传到Gemini网盘")
|
|
async def make_grid_gemini_upload(data: MakeGridGeminiRequest, headers: Annotated[BundleHeaders, 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")
|
|
fn = modal.Function.from_name(config.modal_app_name, "make_image_grid_upload",
|
|
environment_name=config.modal_environment)
|
|
image_grid_uri = await fn.remote.aio(data.pic_info_list, data.image_size, data.text_height, data.font_size,
|
|
data.padding, data.separator, google_api_key,
|
|
SentryTransactionInfo(x_trace_id=sentry_sdk.get_traceparent(),
|
|
x_baggage=sentry_sdk.get_baggage())
|
|
if headers.x_trace_id is None else SentryTransactionInfo(
|
|
x_trace_id=headers.x_trace_id, x_baggage=headers.x_baggage))
|
|
return JSONResponse(content={"uri": image_grid_uri}, status_code=200)
|
|
|
|
|
|
@router.post('/inference_gemini', summary="使用Gemini推理hls视频流指定时间段的打点情况")
|
|
async def inference_gemini(
|
|
data: GeminiRequest,
|
|
headers: Annotated[BundleHeaders, Header()],
|
|
) -> ModalTaskResponse:
|
|
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")
|
|
fn = modal.Function.from_name(config.modal_app_name, "video_hls_slice_inference",
|
|
environment_name=config.modal_environment)
|
|
fn_call = fn.spawn(data.media_hls_url, google_api_key, data.product_cover_grid_uri_list, data.product_list,
|
|
data.start_time, data.end_time, data.options,
|
|
SentryTransactionInfo(x_trace_id=sentry_sdk.get_traceparent(),
|
|
x_baggage=sentry_sdk.get_baggage())
|
|
if headers.x_trace_id is None else SentryTransactionInfo(x_trace_id=headers.x_trace_id,
|
|
x_baggage=headers.x_baggage),
|
|
data.webhook, 2, data.scale)
|
|
return ModalTaskResponse(success=True, taskId=fn_call.object_id)
|
|
|
|
|
|
@router.get("/inference_gemini/{task_id}", summary="查询Gemini推理hls视频流指定时间段的打点任务")
|
|
async def gemini_status(task_id: str, response: Response) -> GeminiResultResponse:
|
|
task_info = await ModalUtils.get_modal_task_status(task_id)
|
|
if task_info.transaction:
|
|
response.headers["x-trace-id"] = task_info.transaction.x_trace_id
|
|
response.headers["x-baggage"] = task_info.transaction.x_baggage
|
|
try:
|
|
return GeminiResultResponse(taskId=task_id, status=task_info.status, code=cast(int, task_info.error_code.value),
|
|
error=task_info.error_reason,
|
|
result=json.dumps(task_info.results, ensure_ascii=False)
|
|
if task_info.results is not None else "")
|
|
except Exception as e:
|
|
logger.exception(f"获取Gemini状态发生错误 {e}")
|
|
return GeminiResultResponse(taskId=task_id, status=task_info.status, code=cast(int, task_info.error_code.value),
|
|
error=task_info.error_reason + f"获取Gemini状态发生错误 {e}"
|
|
if task_info.error_reason else f"获取Gemini状态发生错误 {e}", result="")
|