refactor: 重构异步任务系统和计费服务架构

- 重构任务系统:新增 lifecycle (TaskStatus/BillingStatus)、context、application 模块
- 将 video tasks 泛化为 async tasks,支持更通用的异步任务管理
- 新增 Gemini Files 管理模块和管理界面
- 重构 billing 服务:拆分 schema.py 和 service.py
- 新增 candidate 服务模块用于请求候选管理
- 数据库迁移:添加 billing_status、request_id、gemini_file_mappings 表和索引
- 移除废弃的 video_telemetry、task orchestrator 等模块
This commit is contained in:
fawney19
2026-02-02 03:16:52 +08:00
parent feb7484fda
commit 9e31efe26c
75 changed files with 7511 additions and 2068 deletions

View File

@@ -14,13 +14,14 @@ from .system_catalog import router as system_catalog_router
from .videos import router as videos_router
router = APIRouter()
# Models API 需要在最前面注册,避免被其他路由的 path 参数捕获
# Video API 路由需要在 Models API 之前注册,因为 Models API 有 /v1beta/models/{path} 通配符路由
# 会错误匹配 /v1beta/models/{model}/operations/{id}/content 等视频路由
router.include_router(videos_router, tags=["Video Generation"])
router.include_router(models_router)
router.include_router(claude_router, tags=["Claude API"])
router.include_router(openai_router)
router.include_router(gemini_router, tags=["Gemini API"])
router.include_router(gemini_files_router, tags=["Gemini Files API"])
router.include_router(videos_router, tags=["Video Generation"])
router.include_router(system_catalog_router, tags=["System Catalog"])
router.include_router(catalog_router)
router.include_router(capabilities_router)

View File

@@ -15,14 +15,17 @@ Gemini Files API 代理端点
参考文档:
https://ai.google.dev/api/files
优化HTTP 代理请求期间不持有数据库连接,避免阻塞其他请求。
"""
from __future__ import annotations
from dataclasses import dataclass
from typing import Any, Dict, Optional, Tuple
from urllib.parse import urlencode
from fastapi import APIRouter, Depends, HTTPException, Request, Response
from fastapi import APIRouter, HTTPException, Request, Response
from fastapi.responses import JSONResponse
from sqlalchemy.orm import Session
@@ -30,20 +33,30 @@ from src.clients.http_client import HTTPClientPool
from src.core.api_format import get_auth_handler, get_default_auth_method_for_endpoint
from src.core.crypto import crypto_service
from src.core.logger import logger
from src.database import get_db
from src.database import create_session
from src.models.database import ApiKey, GlobalModel, Model, Provider, ProviderEndpoint, User
from src.services.auth.service import AuthService
from src.services.cache.aware_scheduler import CacheAwareScheduler, ProviderCandidate
from src.services.gemini_files_mapping import delete_file_key_mapping, store_file_key_mapping
from src.services.provider.transport import redact_url_for_log
@dataclass
class UpstreamContext:
"""上游请求上下文(不依赖数据库会话)"""
upstream_key: str
base_url: str
file_key_id: str
user_id: str
router = APIRouter(tags=["Gemini Files API"])
# Gemini Files API 基础 URL
GEMINI_FILES_BASE_URL = "https://generativelanguage.googleapis.com"
# Gemini Files API 能力标签
REQUIRED_CAPABILITIES = {"gemini_files_api": True}
# Gemini Files API 能力限制(任何 Gemini key 都可用)
# 需要从客户端请求中移除的头部(这些会由代理重新设置或不应转发)
HEADERS_TO_REMOVE = frozenset(
@@ -184,9 +197,25 @@ async def _select_provider_candidate(
db: Session,
user_api_key: ApiKey,
model_name: str,
require_files_capability: bool = True,
) -> ProviderCandidate | None:
"""选择支持 Files API 的 Provider/Endpoint/Key 组合"""
"""
选择可用的 Provider/Endpoint/Key 组合
Args:
db: 数据库会话
user_api_key: 用户 API Key
model_name: 模型名称
require_files_capability: 是否要求 gemini_files 能力(默认 True
Returns:
匹配的候选,如果没有则返回 None
"""
scheduler = CacheAwareScheduler()
# 要求 gemini_files 能力:只有 Google 官方 API 才支持 Files API
capability_requirements = {"gemini_files": True} if require_files_capability else None
candidates, _global_model_id = await scheduler.list_all_candidates(
db=db,
api_format="gemini:chat",
@@ -194,7 +223,7 @@ async def _select_provider_candidate(
affinity_key=str(user_api_key.id),
user_api_key=user_api_key,
max_candidates=10,
capability_requirements=REQUIRED_CAPABILITIES,
capability_requirements=capability_requirements,
)
for candidate in candidates:
auth_type = getattr(candidate.key, "auth_type", "api_key") or "api_key"
@@ -206,12 +235,20 @@ async def _select_provider_candidate(
async def _resolve_upstream_context(
request: Request,
db: Session,
) -> tuple[str, str, str]:
) -> tuple[str, str, str, str]:
"""
解析上游 Key 与 Base URL
解析上游 Key 与 Base URL(需要外部提供 db session
仅允许系统 API Key通过能力标签选择支持 Files API 的 Provider Key。
仅允许系统 API Key选择可用的 Gemini Provider Key(无能力限制)
Args:
request: HTTP 请求
db: 数据库会话
Returns:
(upstream_key, base_url, key_id, user_id)
"""
client_key = _extract_gemini_api_key(request)
if not client_key:
raise HTTPException(
@@ -252,14 +289,19 @@ async def _resolve_upstream_context(
},
)
candidate = await _select_provider_candidate(db, user_api_key, model_name)
# 选择可用的 provider candidate(要求 gemini_files 能力)
candidate = await _select_provider_candidate(
db, user_api_key, model_name, require_files_capability=True
)
if not candidate:
raise HTTPException(
status_code=503,
detail={
"error": {
"code": 503,
"message": "No available key with gemini_files_api capability",
"message": "No available Gemini key with 'gemini_files' capability. "
"Please ensure at least one Provider Key has the 'gemini_files' capability enabled.",
"status": "UNAVAILABLE",
}
},
@@ -268,7 +310,7 @@ async def _resolve_upstream_context(
try:
upstream_key = crypto_service.decrypt(candidate.key.api_key)
except Exception as exc:
logger.error(f"Failed to decrypt provider key for Gemini Files API: {exc}")
logger.error("Failed to decrypt provider key for Gemini Files API: %s", exc)
raise HTTPException(
status_code=500,
detail={
@@ -281,7 +323,29 @@ async def _resolve_upstream_context(
)
base_url = candidate.endpoint.base_url or GEMINI_FILES_BASE_URL
return upstream_key, base_url, str(candidate.key.id)
return upstream_key, base_url, str(candidate.key.id), str(user.id)
async def _resolve_upstream_context_standalone(request: Request) -> UpstreamContext:
"""
解析上游上下文(自管理数据库连接,适用于 HTTP 代理场景)
优化在返回上下文后立即释放数据库连接HTTP 请求期间不持有连接。
Args:
request: HTTP 请求
Returns:
UpstreamContext: 包含所有必要信息的上下文对象
"""
with create_session() as db:
upstream_key, base_url, file_key_id, user_id = await _resolve_upstream_context(request, db)
return UpstreamContext(
upstream_key=upstream_key,
base_url=base_url,
file_key_id=file_key_id,
user_id=user_id,
)
async def _proxy_request(
@@ -291,6 +355,7 @@ async def _proxy_request(
content: bytes | None = None,
json_body: dict[str, Any] | None = None,
file_key_id: str | None = None,
user_id: str | None = None,
) -> Response:
"""
代理请求到上游 Gemini API
@@ -302,6 +367,7 @@ async def _proxy_request(
content: 原始请求体(二进制)
json_body: JSON 请求体
file_key_id: 上游 Provider Key ID用于成功响应时存储 file→key 映射
user_id: 用户 ID用于文件映射的权限验证
Returns:
FastAPI Response 对象
@@ -338,12 +404,28 @@ async def _proxy_request(
try:
payload = response.json()
file_name = None
file_obj = None
if isinstance(payload, dict):
# 单文件上传响应
file_name = payload.get("name")
file_obj = payload
# 嵌套格式:{"file": {...}}
if not file_name and isinstance(payload.get("file"), dict):
file_name = payload["file"].get("name")
if file_name:
await store_file_key_mapping(file_name, file_key_id)
file_obj = payload["file"]
if file_name and file_obj:
display_name = file_obj.get("displayName") or file_obj.get("display_name")
mime_type = file_obj.get("mimeType") or file_obj.get("mime_type")
await store_file_key_mapping(
file_name,
file_key_id,
user_id=user_id,
display_name=display_name,
mime_type=mime_type,
)
logger.debug(
f"Gemini file→key 映射已存储: {file_name} → key_id={file_key_id}"
)
@@ -355,14 +437,26 @@ async def _proxy_request(
mapped_count = 0
for item in files_list:
if isinstance(item, dict) and item.get("name"):
await store_file_key_mapping(item["name"], file_key_id)
item_display_name = item.get("displayName") or item.get(
"display_name"
)
item_mime_type = item.get("mimeType") or item.get("mime_type")
await store_file_key_mapping(
item["name"],
file_key_id,
user_id=user_id,
display_name=item_display_name,
mime_type=item_mime_type,
)
mapped_count += 1
if mapped_count > 0:
logger.debug(
f"Gemini list_files 批量映射已存储: {mapped_count} 个文件 → key_id={file_key_id}"
"Gemini list_files 批量映射已存储: %d 个文件 → key_id=%s",
mapped_count,
file_key_id,
)
except (ValueError, KeyError) as e:
logger.debug(f"Failed to store Gemini file mapping: {e}")
logger.debug("Failed to store Gemini file mapping: %s", e)
return Response(
content=response.content,
@@ -373,7 +467,7 @@ async def _proxy_request(
except Exception as e:
sanitized_error = redact_url_for_log(str(e))
logger.error(f"Gemini Files API proxy error: {sanitized_error}")
logger.error("Gemini Files API proxy error: %s", sanitized_error)
return JSONResponse(
status_code=502,
content={
@@ -394,7 +488,6 @@ async def _proxy_request(
@router.post("/upload/v1beta/files")
async def upload_file(
request: Request,
db: Session = Depends(get_db),
) -> Any:
"""
上传文件到 Gemini Files API
@@ -421,26 +514,34 @@ async def upload_file(
}
}
```
"""
upstream_key, base_url, file_key_id = await _resolve_upstream_context(request, db)
# 读取请求体
优化HTTP 代理期间不持有数据库连接
"""
# 阶段 1解析上下文短暂持有数据库连接
ctx = await _resolve_upstream_context_standalone(request)
# 阶段 2读取请求体
body = await request.body()
# 构建上游请求
# 阶段 3代理请求不持有数据库连接
upstream_url = _build_upstream_url(
base_url,
ctx.base_url,
"/v1beta/files",
dict(request.query_params),
is_upload=True,
)
headers = _build_upstream_headers(dict(request.headers), upstream_key)
headers = _build_upstream_headers(dict(request.headers), ctx.upstream_key)
logger.debug(f"Gemini Files upload proxy: POST {redact_url_for_log(upstream_url)}")
logger.debug("Gemini Files upload proxy: POST %s", redact_url_for_log(upstream_url))
return await _proxy_request(
"POST", upstream_url, headers, content=body, file_key_id=file_key_id
"POST",
upstream_url,
headers,
content=body,
file_key_id=ctx.file_key_id,
user_id=ctx.user_id,
)
@@ -452,13 +553,14 @@ async def upload_file(
@router.get("/v1beta/files")
async def list_files(
request: Request,
db: Session = Depends(get_db),
pageSize: int | None = None,
pageToken: str | None = None,
) -> Any:
"""
列出已上传的文件
优化HTTP 代理期间不持有数据库连接
**认证方式**:
- `x-goog-api-key` 请求头,或
- `?key=` URL 参数
@@ -488,22 +590,222 @@ async def list_files(
}
```
"""
upstream_key, base_url, file_key_id = await _resolve_upstream_context(request, db)
# 阶段 1解析上下文短暂持有数据库连接
ctx = await _resolve_upstream_context_standalone(request)
# 构建查询参数
# 阶段 2代理请求不持有数据库连接
query_params = dict(request.query_params)
if pageSize is not None:
query_params["pageSize"] = pageSize
if pageToken is not None:
query_params["pageToken"] = pageToken
# 构建上游请求
upstream_url = _build_upstream_url(base_url, "/v1beta/files", query_params)
upstream_url = _build_upstream_url(ctx.base_url, "/v1beta/files", query_params)
headers = _build_upstream_headers(dict(request.headers), ctx.upstream_key)
logger.debug("Gemini Files list proxy: GET %s", redact_url_for_log(upstream_url))
return await _proxy_request(
"GET", upstream_url, headers, file_key_id=ctx.file_key_id, user_id=ctx.user_id
)
# ==============================================================================
# 下载文件内容端点(用于视频等媒体文件)
# 注意:必须在 /v1beta/files/{file_name:path} 之前注册,否则会被通配符路由捕获
# ==============================================================================
async def _find_video_task_by_id(
db: Session, short_id: str, user_id: str
) -> tuple[str | None, str | None]:
"""
通过短 ID 查找视频任务,返回其 provider key 和 video_url
Args:
db: 数据库会话
short_id: 视频任务的短 IDVideoTask.short_idGemini 风格)
user_id: 用户 ID用于权限验证
Returns:
(upstream_key, video_url) - 如果找到任务返回 key 和 url否则返回 (None, None)
"""
from src.models.database import ProviderAPIKey, VideoTask
logger.debug(
"[Files Download] Searching video task: short_id=%s, user_id=%s", short_id, user_id
)
# 通过 short_id 查找,同时验证用户权限
task = (
db.query(VideoTask)
.filter(VideoTask.short_id == short_id, VideoTask.user_id == user_id)
.first()
)
if not task:
logger.debug("[Files Download] No video task found: short_id=%s", short_id)
return None, None
if not task.video_url:
logger.debug("[Files Download] Task found but no video_url: short_id=%s", short_id)
return None, None
if not task.key_id:
logger.debug("[Files Download] Task found but no key_id: short_id=%s", short_id)
return None, task.video_url
# 获取 provider key
provider_key = db.query(ProviderAPIKey).filter(ProviderAPIKey.id == task.key_id).first()
if not provider_key or not provider_key.api_key:
logger.debug("[Files Download] Provider key not found: key_id=%s", task.key_id)
return None, task.video_url
try:
upstream_key = crypto_service.decrypt(provider_key.api_key)
logger.debug("[Files Download] Found key for task: short_id=%s", short_id)
return upstream_key, task.video_url
except Exception as e:
logger.error("[Files Download] Failed to decrypt key: %s", e)
return None, task.video_url
@router.get("/v1beta/files/{file_id}:download")
async def download_file(
file_id: str,
request: Request,
) -> Any:
"""
下载文件(官方 Gemini API 格式)
**认证方式**:
- `x-goog-api-key` 请求头,或
- `?key=` URL 参数
**路径参数**:
- `file_id`: 文件 ID
- 以 `aev_` 开头:视频任务下载(如 `aev_sknuzqlo8sds`Gemini 风格短 ID
- 其他:普通 Gemini 文件下载(透传到上游)
**查询参数**:
- `alt=media`: 可选,保持与官方 API 兼容
**示例**:
```
GET /v1beta/files/aev_{short_id}:download?alt=media # 视频任务
GET /v1beta/files/{gemini_file_id}:download?alt=media # 普通文件
```
优化HTTP 下载期间不持有数据库连接
"""
import httpx
from fastapi import HTTPException
from fastapi.responses import JSONResponse, Response
# ========== 阶段 1数据库操作短暂持有连接==========
client_key = _extract_gemini_api_key(request)
if not client_key:
raise HTTPException(
status_code=401,
detail={
"error": {"code": 401, "message": "API key required", "status": "UNAUTHENTICATED"}
},
)
# 在数据库会话内完成所有查询
with create_session() as db:
auth_result = AuthService.authenticate_api_key(db, client_key)
if not auth_result:
raise HTTPException(
status_code=401,
detail={
"error": {
"code": 401,
"message": "API key not valid",
"status": "UNAUTHENTICATED",
}
},
)
user, _user_api_key = auth_result
# 根据前缀判断处理方式
if file_id.startswith("aev_"):
# 视频任务下载:使用短 ID 查找
short_id = file_id[4:] # 去掉 "aev_" 前缀
logger.debug("[Files Download] Video task: short_id=%s, user_id=%s", short_id, user.id)
upstream_key, video_url = await _find_video_task_by_id(db, short_id, user.id)
if not upstream_key or not video_url:
raise HTTPException(
status_code=404,
detail={
"error": {
"code": 404,
"message": f"Video not found or not ready: {file_id}",
"status": "NOT_FOUND",
}
},
)
upstream_url = video_url
else:
# 普通文件下载:透传到 Gemini
try:
upstream_key, base_url, _file_key_id, _user_id = await _resolve_upstream_context(
request, db
)
except HTTPException:
raise HTTPException(
status_code=404,
detail={
"error": {
"code": 404,
"message": f"File not found: {file_id}",
"status": "NOT_FOUND",
}
},
)
file_name = f"files/{file_id}" if not file_id.startswith("files/") else file_id
upstream_url = _build_upstream_url(
base_url,
f"/v1beta/{file_name}:download",
dict(request.query_params),
)
# ========== 阶段 2HTTP 下载(不持有数据库连接)==========
headers = _build_upstream_headers(dict(request.headers), upstream_key)
logger.debug(f"Gemini Files list proxy: GET {redact_url_for_log(upstream_url)}")
logger.debug("Gemini Files download proxy: GET %s", redact_url_for_log(upstream_url))
return await _proxy_request("GET", upstream_url, headers, file_key_id=file_key_id)
# 使用 follow_redirects=True 跟随重定向Gemini 文件下载会重定向)
try:
async with httpx.AsyncClient(follow_redirects=True, timeout=httpx.Timeout(300.0)) as client:
response = await client.get(upstream_url, headers=headers)
except Exception as exc:
logger.error("Gemini Files download failed: %s", exc)
raise HTTPException(status_code=502, detail="Failed to download file")
if response.status_code >= 400:
content: dict[str, Any]
if response.headers.get("content-type", "").startswith("application/json"):
try:
content = response.json()
except Exception:
content = {"error": response.text}
else:
content = {"error": response.text}
return JSONResponse(content=content, status_code=response.status_code)
# 返回文件内容
return Response(
content=response.content,
status_code=response.status_code,
headers={
k: v
for k, v in response.headers.items()
if k.lower() not in {"transfer-encoding", "connection", "keep-alive"}
},
media_type=response.headers.get("content-type", "application/octet-stream"),
)
# ==============================================================================
@@ -515,7 +817,6 @@ async def list_files(
async def get_file(
file_name: str,
request: Request,
db: Session = Depends(get_db),
) -> Any:
"""
获取指定文件的元数据
@@ -542,24 +843,29 @@ async def get_file(
"state": "ACTIVE"
}
```
"""
upstream_key, base_url, file_key_id = await _resolve_upstream_context(request, db)
优化HTTP 代理期间不持有数据库连接
"""
# 阶段 1解析上下文短暂持有数据库连接
ctx = await _resolve_upstream_context_standalone(request)
# 阶段 2代理请求不持有数据库连接
# 规范化文件名(确保以 files/ 开头)
if not file_name.startswith("files/"):
file_name = f"files/{file_name}"
# 构建上游请求
upstream_url = _build_upstream_url(
base_url,
ctx.base_url,
f"/v1beta/{file_name}",
dict(request.query_params),
)
headers = _build_upstream_headers(dict(request.headers), upstream_key)
headers = _build_upstream_headers(dict(request.headers), ctx.upstream_key)
logger.debug(f"Gemini Files get proxy: GET {redact_url_for_log(upstream_url)}")
logger.debug("Gemini Files get proxy: GET %s", redact_url_for_log(upstream_url))
return await _proxy_request("GET", upstream_url, headers, file_key_id=file_key_id)
return await _proxy_request(
"GET", upstream_url, headers, file_key_id=ctx.file_key_id, user_id=ctx.user_id
)
# ==============================================================================
@@ -571,7 +877,6 @@ async def get_file(
async def delete_file(
file_name: str,
request: Request,
db: Session = Depends(get_db),
) -> Any:
"""
删除指定文件
@@ -585,29 +890,31 @@ async def delete_file(
**响应格式**:
成功时返回空 JSON 对象:`{}`
"""
upstream_key, base_url, _file_key_id = await _resolve_upstream_context(request, db)
优化HTTP 代理期间不持有数据库连接
"""
# 阶段 1解析上下文短暂持有数据库连接
ctx = await _resolve_upstream_context_standalone(request)
# 阶段 2代理请求不持有数据库连接
# 规范化文件名(确保以 files/ 开头)
if not file_name.startswith("files/"):
file_name = f"files/{file_name}"
# 构建上游请求
upstream_url = _build_upstream_url(
base_url,
ctx.base_url,
f"/v1beta/{file_name}",
dict(request.query_params),
)
headers = _build_upstream_headers(dict(request.headers), upstream_key)
headers = _build_upstream_headers(dict(request.headers), ctx.upstream_key)
logger.debug(f"Gemini Files delete proxy: DELETE {redact_url_for_log(upstream_url)}")
logger.debug("Gemini Files delete proxy: DELETE %s", redact_url_for_log(upstream_url))
del _file_key_id # 显式标记delete 端点不需要存储映射
response = await _proxy_request("DELETE", upstream_url, headers)
if response.status_code < 300:
await delete_file_key_mapping(file_name)
else:
logger.debug(
f"Gemini Files delete failed, skip mapping cleanup: status={response.status_code}"
"Gemini Files delete failed, skip mapping cleanup: status=%s", response.status_code
)
return response

View File

@@ -61,9 +61,10 @@ async def list_video_tasks_sora(http_request: Request, db: Session = Depends(get
@router.delete("/v1/videos/{task_id}")
async def cancel_video_task_sora(
async def delete_video_task_sora(
task_id: str, http_request: Request, db: Session = Depends(get_db)
) -> Any:
"""删除已完成或失败的视频及其存储资源"""
adapter = OpenAIVideoAdapter()
return await pipeline.run(
adapter=adapter,
@@ -71,7 +72,7 @@ async def cancel_video_task_sora(
db=db,
mode=adapter.mode,
api_format_hint=adapter.allowed_api_formats[0],
path_params={"task_id": task_id, "action": "cancel"},
path_params={"task_id": task_id},
)
@@ -121,6 +122,47 @@ async def create_video_veo(model: str, http_request: Request, db: Session = Depe
)
# Gemini Veo operation routes - support both formats:
# 1. models/{model}/operations/{id} (official Gemini Veo format)
# 2. operations/{...} (legacy format for compatibility)
@router.get("/v1beta/models/{model}/operations/{operation_id}")
async def get_video_veo_by_model(
model: str, operation_id: str, http_request: Request, db: Session = Depends(get_db)
) -> Any:
"""Get video task status (Gemini Veo format: models/{model}/operations/{id})"""
adapter = GeminiVeoAdapter()
# Reconstruct full operation name
full_operation_name = f"models/{model}/operations/{operation_id}"
return await pipeline.run(
adapter=adapter,
http_request=http_request,
db=db,
mode=adapter.mode,
api_format_hint=adapter.allowed_api_formats[0],
path_params={"task_id": full_operation_name},
)
@router.post("/v1beta/models/{model}/operations/{operation_id}:cancel")
async def cancel_video_veo_by_model(
model: str, operation_id: str, http_request: Request, db: Session = Depends(get_db)
) -> Any:
"""Cancel video task (Gemini Veo format: models/{model}/operations/{id}:cancel)"""
adapter = GeminiVeoAdapter()
full_operation_name = f"models/{model}/operations/{operation_id}"
return await pipeline.run(
adapter=adapter,
http_request=http_request,
db=db,
mode=adapter.mode,
api_format_hint=adapter.allowed_api_formats[0],
path_params={"task_id": full_operation_name, "action": "cancel"},
)
# Legacy routes for backward compatibility
@router.get("/v1beta/operations/{operation_id:path}")
async def get_video_veo(
operation_id: str, http_request: Request, db: Session = Depends(get_db)
@@ -163,16 +205,4 @@ async def cancel_video_veo(
)
@router.get("/v1beta/operations/{operation_id:path}/content")
async def download_video_content_veo(
operation_id: str, http_request: Request, db: Session = Depends(get_db)
) -> Any:
adapter = GeminiVeoAdapter()
return await pipeline.run(
adapter=adapter,
http_request=http_request,
db=db,
mode=adapter.mode,
api_format_hint=adapter.allowed_api_formats[0],
path_params={"task_id": operation_id},
)
# Video download is now handled by /v1beta/files/{task_id}:download in gemini_files.py