feat: add Vertex AI authentication support for provider API keys

- Add auth_type field to ProviderAPIKey model (api_key or vertex_ai)
- Implement Vertex AI OAuth token generation with service account
- Update transport layer to handle Vertex AI authentication
- Add Vertex AI endpoint URL generation in request builder
- Update frontend KeyFormDialog to support auth_type selection
- Add migration for auth_type column in provider_api_keys table
This commit is contained in:
fawney19
2026-01-30 02:43:50 +08:00
parent 3e75bc8964
commit 32b293ef3e
16 changed files with 956 additions and 65 deletions

View File

@@ -2,6 +2,7 @@
Provider API Keys 管理
"""
import json
import uuid
from dataclasses import dataclass
from datetime import datetime, timezone
@@ -229,8 +230,35 @@ class AdminUpdateEndpointKeyAdapter(AdminApiAdapter):
exclude_patterns_before = key.model_exclude_patterns
update_data = self.key_data.model_dump(exclude_unset=True)
if "api_key" in update_data:
# 验证 auth_type 切换
current_auth_type = getattr(key, "auth_type", "api_key") or "api_key"
target_auth_type = update_data.get("auth_type", current_auth_type) or current_auth_type
# auth_type 切换校验 + 字段归一化
if "auth_type" in update_data:
if target_auth_type == "api_key":
if current_auth_type == "vertex_ai" and not update_data.get("api_key"):
raise InvalidRequestException(
"从 Vertex AI 切换到 API Key 认证模式时,必须提供新的 API Key"
)
# 切换回 API Key清理 Service Account 配置
update_data["auth_config"] = None
elif target_auth_type == "vertex_ai":
if current_auth_type != "vertex_ai" and not update_data.get("auth_config"):
raise InvalidRequestException(
"从 API Key 切换到 Vertex AI 认证模式时,必须提供 Service Account JSON"
)
# Vertex AI 不使用 api_key写入占位符若未提供 api_key
if "api_key" not in update_data:
update_data["api_key"] = "__placeholder__"
# 加密 api_key非 None 时)
if "api_key" in update_data and update_data["api_key"] is not None:
update_data["api_key"] = crypto_service.encrypt(update_data["api_key"])
# 加密 auth_config包含敏感的 Service Account 凭证)
if "auth_config" in update_data and update_data["auth_config"]:
update_data["auth_config"] = crypto_service.encrypt(json.dumps(update_data["auth_config"]))
# 特殊处理 rpm_limit需要区分"未提供"和"显式设置为 null"
if "rpm_limit" in self.key_data.model_fields_set:
@@ -347,7 +375,7 @@ class AdminUpdateEndpointKeyAdapter(AdminApiAdapter):
@dataclass
class AdminRevealEndpointKeyAdapter(AdminApiAdapter):
"""获取完整的 API Key用于查看和复制"""
"""获取完整的 API Key 或 Auth Config(用于查看和复制)"""
key_id: str
@@ -357,6 +385,42 @@ class AdminRevealEndpointKeyAdapter(AdminApiAdapter):
if not key:
raise NotFoundException(f"Key {self.key_id} 不存在")
auth_type = getattr(key, "auth_type", "api_key") or "api_key"
# Vertex AI 类型返回 auth_config需要解密
if auth_type == "vertex_ai":
encrypted_auth_config = getattr(key, "auth_config", None)
if encrypted_auth_config:
try:
decrypted_config = crypto_service.decrypt(encrypted_auth_config)
auth_config = json.loads(decrypted_config)
logger.info(f"[REVEAL] 查看 Auth Config: ID={self.key_id}, Name={key.name}")
return {"auth_type": "vertex_ai", "auth_config": auth_config}
except Exception as e:
logger.error(f"解密 Auth Config 失败: ID={self.key_id}, Error={e}")
raise InvalidRequestException(
"无法解密认证配置,可能是加密密钥已更改。请重新添加该密钥。"
)
# 兼容auth_config 为空时尝试从 api_key 解密(仅对迁移前的旧数据有效)
try:
decrypted_key = crypto_service.decrypt(key.api_key)
# 检查是否是新格式的占位符(表示 auth_config 丢失)
if decrypted_key == "__placeholder__":
logger.error(f"Vertex AI Key 缺少 auth_config: ID={self.key_id}")
raise InvalidRequestException(
"认证配置丢失,请重新添加该密钥。"
)
logger.info(f"[REVEAL] 查看完整 Key (legacy vertex_ai): ID={self.key_id}, Name={key.name}")
return {"auth_type": "vertex_ai", "auth_config": decrypted_key}
except InvalidRequestException:
raise
except Exception as e:
logger.error(f"解密 Key 失败: ID={self.key_id}, Error={e}")
raise InvalidRequestException(
"无法解密认证配置,可能是加密密钥已更改。请重新添加该密钥。"
)
# API Key 类型返回 api_key
try:
decrypted_key = crypto_service.decrypt(key.api_key)
except Exception as e:
@@ -366,7 +430,7 @@ class AdminRevealEndpointKeyAdapter(AdminApiAdapter):
)
logger.info(f"[REVEAL] 查看完整 Key: ID={self.key_id}, Name={key.name}")
return {"api_key": decrypted_key}
return {"auth_type": "api_key", "api_key": decrypted_key}
@dataclass
@@ -451,12 +515,16 @@ class AdminGetKeysGroupedByFormatAdapter(AdminApiAdapter):
if not api_formats:
continue # 跳过没有 API 格式的 Key
try:
decrypted_key = crypto_service.decrypt(key.api_key)
masked_key = f"{decrypted_key[:8]}***{decrypted_key[-4:]}"
except Exception as e:
logger.error(f"解密 Key 失败: key_id={key.id}, error={e}")
masked_key = "***ERROR***"
auth_type = getattr(key, "auth_type", "api_key") or "api_key"
if auth_type == "vertex_ai":
masked_key = "[Service Account]"
else:
try:
decrypted_key = crypto_service.decrypt(key.api_key)
masked_key = f"{decrypted_key[:8]}***{decrypted_key[-4:]}"
except Exception as e:
logger.error(f"解密 Key 失败: key_id={key.id}, error={e}")
masked_key = "***ERROR***"
# 计算健康度指标
success_rate = key.success_count / key.request_count if key.request_count > 0 else None
@@ -478,6 +546,7 @@ class AdminGetKeysGroupedByFormatAdapter(AdminApiAdapter):
key_info = {
"id": key.id,
"name": key.name,
"auth_type": auth_type,
"api_key_masked": masked_key,
"internal_priority": key.internal_priority,
"global_priority_by_format": key.global_priority_by_format,
@@ -525,11 +594,17 @@ def _build_key_response(
key: ProviderAPIKey, api_key_plain: str | None = None
) -> EndpointAPIKeyResponse:
"""构建 Key 响应对象的辅助函数"""
try:
decrypted_key = crypto_service.decrypt(key.api_key)
masked_key = f"{decrypted_key[:8]}***{decrypted_key[-4:]}"
except Exception:
masked_key = "***ERROR***"
auth_type = getattr(key, "auth_type", "api_key") or "api_key"
if auth_type == "vertex_ai":
# Vertex AI 使用 Service Account不显示占位符
masked_key = "[Service Account]"
else:
try:
decrypted_key = crypto_service.decrypt(key.api_key)
masked_key = f"{decrypted_key[:8]}***{decrypted_key[-4:]}"
except Exception:
masked_key = "***ERROR***"
success_rate = key.success_count / key.request_count if key.request_count > 0 else 0.0
avg_response_time_ms = (
@@ -539,6 +614,8 @@ def _build_key_response(
is_adaptive = key.rpm_limit is None
key_dict = key.__dict__.copy()
key_dict.pop("_sa_instance_state", None)
key_dict.pop("api_key", None) # 移除敏感字段,避免泄露
key_dict.pop("auth_config", None) # 移除敏感字段,避免泄露
# 从 health_by_format 计算汇总字段(便于列表展示)
health_by_format = key.health_by_format or {}
@@ -636,17 +713,38 @@ class AdminCreateProviderKeyAdapter(AdminApiAdapter):
if not self.key_data.api_formats:
raise InvalidRequestException("api_formats 为必填字段")
# 验证认证配置
auth_type = self.key_data.auth_type or "api_key"
if auth_type == "api_key":
if not self.key_data.api_key:
raise InvalidRequestException("API Key 认证模式下 api_key 为必填字段")
elif auth_type == "vertex_ai":
if not self.key_data.auth_config:
raise InvalidRequestException("Service Account 认证模式下 auth_config 为必填字段")
# 允许同一个 API Key 在同一 Provider 下添加多次
# 用户可以为不同的 API 格式创建独立的配置记录,便于分开管理
encrypted_key = crypto_service.encrypt(self.key_data.api_key)
# 加密 API Key如果有
encrypted_key = (
crypto_service.encrypt(self.key_data.api_key)
if self.key_data.api_key
else crypto_service.encrypt("__placeholder__") # 占位符,保持 NOT NULL 约束
)
now = datetime.now(timezone.utc)
# 加密 auth_config包含敏感的 Service Account 凭证)
encrypted_auth_config = None
if self.key_data.auth_config:
encrypted_auth_config = crypto_service.encrypt(json.dumps(self.key_data.auth_config))
new_key = ProviderAPIKey(
id=str(uuid.uuid4()),
provider_id=self.provider_id,
api_formats=self.key_data.api_formats,
auth_type=auth_type,
api_key=encrypted_key,
auth_config=encrypted_auth_config,
name=self.key_data.name,
note=self.key_data.note,
rate_multipliers=self.key_data.rate_multipliers,

View File

@@ -35,7 +35,7 @@ from src.api.handlers.base.base_handler import (
wait_for_with_disconnect_detection,
)
from src.api.handlers.base.parsers import get_parser_for_format
from src.api.handlers.base.request_builder import PassthroughRequestBuilder
from src.api.handlers.base.request_builder import PassthroughRequestBuilder, get_provider_auth
from src.api.handlers.base.response_parser import ResponseParser
from src.api.handlers.base.stream_context import StreamContext
from src.api.handlers.base.stream_processor import StreamProcessor
@@ -681,6 +681,9 @@ class ChatHandlerBase(BaseMessageHandler, ABC):
# 同格式:按原逻辑做轻量清理(子类可覆盖以移除不需要的字段)
request_body = self.prepare_provider_request_body(request_body)
# 获取认证信息(处理 Service Account 等异步认证场景)
auth_info = await get_provider_auth(endpoint, key)
# 构建请求(上游始终使用 header 认证,不跟随客户端的 query 方式)
provider_payload, provider_headers = self._request_builder.build(
request_body,
@@ -688,6 +691,7 @@ class ChatHandlerBase(BaseMessageHandler, ABC):
endpoint,
key,
is_stream=True,
pre_computed_auth=auth_info.as_tuple() if auth_info else None,
)
ctx.provider_request_headers = provider_headers
@@ -701,6 +705,8 @@ class ChatHandlerBase(BaseMessageHandler, ABC):
query_params=query_params,
path_params={"model": url_model},
is_stream=True,
key=key,
decrypted_auth_config=auth_info.decrypted_auth_config if auth_info else None,
)
logger.debug(
@@ -986,6 +992,9 @@ class ChatHandlerBase(BaseMessageHandler, ABC):
# 同格式:按原逻辑做轻量清理(子类可覆盖以移除不需要的字段)
request_body = self.prepare_provider_request_body(request_body)
# 获取认证信息(处理 Service Account 等异步认证场景)
auth_info = await get_provider_auth(endpoint, key)
# 构建请求(上游始终使用 header 认证,不跟随客户端的 query 方式)
provider_payload, provider_hdrs = self._request_builder.build(
request_body,
@@ -993,6 +1002,7 @@ class ChatHandlerBase(BaseMessageHandler, ABC):
endpoint,
key,
is_stream=False,
pre_computed_auth=auth_info.as_tuple() if auth_info else None,
)
provider_request_headers = provider_hdrs
@@ -1006,6 +1016,8 @@ class ChatHandlerBase(BaseMessageHandler, ABC):
query_params=query_params,
path_params={"model": url_model},
is_stream=False,
key=key,
decrypted_auth_config=auth_info.decrypted_auth_config if auth_info else None,
)
logger.info(

View File

@@ -40,7 +40,7 @@ from src.api.handlers.base.base_handler import (
wait_for_with_disconnect_detection,
)
from src.api.handlers.base.parsers import get_parser_for_format
from src.api.handlers.base.request_builder import PassthroughRequestBuilder
from src.api.handlers.base.request_builder import PassthroughRequestBuilder, get_provider_auth
# 直接从具体模块导入,避免循环依赖
from src.api.handlers.base.response_parser import (
@@ -718,6 +718,9 @@ class CliMessageHandlerBase(BaseMessageHandler):
request_body = self.prepare_provider_request_body(request_body)
url_model = self.get_model_for_url(request_body, mapped_model) or mapped_model or ctx.model
# 获取认证信息(处理 Service Account 等异步认证场景)
auth_info = await get_provider_auth(endpoint, key)
# 使用 RequestBuilder 构建请求体和请求头
# 注意mapped_model 已经应用到 request_body这里不再传递
# 上游始终使用 header 认证,不跟随客户端的 query 方式
@@ -727,6 +730,7 @@ class CliMessageHandlerBase(BaseMessageHandler):
endpoint,
key,
is_stream=True,
pre_computed_auth=auth_info.as_tuple() if auth_info else None,
)
# 保存发送给 Provider 的请求信息(用于调试和统计)
@@ -738,6 +742,8 @@ class CliMessageHandlerBase(BaseMessageHandler):
query_params=query_params,
path_params={"model": url_model},
is_stream=True, # CLI handler 处理流式请求
key=key,
decrypted_auth_config=auth_info.decrypted_auth_config if auth_info else None,
)
# 配置 HTTP 超时
@@ -2182,6 +2188,9 @@ class CliMessageHandlerBase(BaseMessageHandler):
request_body = self.prepare_provider_request_body(request_body)
url_model = self.get_model_for_url(request_body, mapped_model) or mapped_model or model
# 获取认证信息(处理 Service Account 等异步认证场景)
auth_info = await get_provider_auth(endpoint, key)
# 使用 RequestBuilder 构建请求体和请求头
# 注意mapped_model 已经应用到 request_body这里不再传递
# 上游始终使用 header 认证,不跟随客户端的 query 方式
@@ -2191,6 +2200,7 @@ class CliMessageHandlerBase(BaseMessageHandler):
endpoint,
key,
is_stream=False,
pre_computed_auth=auth_info.as_tuple() if auth_info else None,
)
# 保存发送给 Provider 的请求信息(用于调试和统计)
@@ -2202,6 +2212,8 @@ class CliMessageHandlerBase(BaseMessageHandler):
query_params=query_params,
path_params={"model": url_model},
is_stream=False, # 非流式请求
key=key,
decrypted_auth_config=auth_info.decrypted_auth_config if auth_info else None,
)
logger.info(

View File

@@ -13,12 +13,36 @@
from __future__ import annotations
import json
from abc import ABC, abstractmethod
from typing import Any, Dict, FrozenSet, Optional, Tuple
from dataclasses import dataclass
from typing import TYPE_CHECKING, Any, Dict, FrozenSet, Optional, Tuple
from src.core.crypto import crypto_service
from src.core.api_format import HeaderBuilder, UPSTREAM_DROP_HEADERS
if TYPE_CHECKING:
from src.models.database import ProviderAPIKey, ProviderEndpoint
# ==============================================================================
# Service Account 认证结果类型
# ==============================================================================
@dataclass
class ProviderAuthInfo:
"""Provider 认证信息(用于 Service Account 等异步认证场景)"""
auth_header: str
auth_value: str
# 解密后的认证配置(用于 URL 构建等场景,避免重复解密)
decrypted_auth_config: Optional[Dict[str, Any]] = None
def as_tuple(self) -> Tuple[str, str]:
"""返回 (auth_header, auth_value) 元组"""
return (self.auth_header, self.auth_value)
# ==============================================================================
# 统一的头部配置常量
# ==============================================================================
@@ -119,6 +143,7 @@ class RequestBuilder(ABC):
key: Any,
*,
extra_headers: Optional[Dict[str, str]] = None,
pre_computed_auth: Optional[Tuple[str, str]] = None,
) -> Dict[str, str]:
"""构建请求头"""
pass
@@ -133,6 +158,7 @@ class RequestBuilder(ABC):
mapped_model: Optional[str] = None,
is_stream: bool = False,
extra_headers: Optional[Dict[str, str]] = None,
pre_computed_auth: Optional[Tuple[str, str]] = None,
) -> Tuple[Dict[str, Any], Dict[str, str]]:
"""
构建完整的请求(请求体 + 请求头)
@@ -145,6 +171,7 @@ class RequestBuilder(ABC):
mapped_model: 映射后的模型名
is_stream: 是否为流式请求
extra_headers: 额外请求头
pre_computed_auth: 预先计算的认证信息 (auth_header, auth_value)
Returns:
Tuple[payload, headers]
@@ -159,6 +186,7 @@ class RequestBuilder(ABC):
endpoint,
key,
extra_headers=extra_headers,
pre_computed_auth=pre_computed_auth,
)
return payload, headers
@@ -195,6 +223,7 @@ class PassthroughRequestBuilder(RequestBuilder):
key: Any,
*,
extra_headers: Optional[Dict[str, str]] = None,
pre_computed_auth: Optional[Tuple[str, str]] = None,
) -> Dict[str, str]:
"""
透传请求头 - 清理敏感头部(黑名单),透传其他所有头部
@@ -204,18 +233,24 @@ class PassthroughRequestBuilder(RequestBuilder):
endpoint: 端点配置
key: Provider API Key
extra_headers: 额外请求头
pre_computed_auth: 预先计算的认证信息 (auth_header, auth_value)
用于 Service Account 等异步获取 token 的场景
"""
from src.core.api_format import get_auth_config, resolve_api_format
# 1. 根据 API 格式自动设置认证头
decrypted_key = crypto_service.decrypt(key.api_key)
api_format = getattr(endpoint, "api_format", None)
resolved_format = resolve_api_format(api_format)
auth_header, auth_type = (
get_auth_config(resolved_format) if resolved_format else ("Authorization", "bearer")
)
auth_value = f"Bearer {decrypted_key}" if auth_type == "bearer" else decrypted_key
if pre_computed_auth:
# 使用预先计算的认证信息Service Account 等场景)
auth_header, auth_value = pre_computed_auth
else:
# 标准 API Key 认证
decrypted_key = crypto_service.decrypt(key.api_key)
api_format = getattr(endpoint, "api_format", None)
resolved_format = resolve_api_format(api_format)
auth_header, auth_type = (
get_auth_config(resolved_format) if resolved_format else ("Authorization", "bearer")
)
auth_value = f"Bearer {decrypted_key}" if auth_type == "bearer" else decrypted_key
# 认证头始终受保护,防止 header_rules 覆盖
protected_keys = {auth_header.lower(), "content-type"}
@@ -272,3 +307,81 @@ def build_passthrough_request(
endpoint,
key,
)
# ==============================================================================
# Service Account 认证支持
# ==============================================================================
async def get_provider_auth(
endpoint: "ProviderEndpoint",
key: "ProviderAPIKey",
) -> Optional[ProviderAuthInfo]:
"""
获取 Provider 的认证信息
对于标准 API Key返回 None由 build_headers 自动处理)。
对于 Service Account异步获取 Access Token 并返回认证信息。
Args:
endpoint: 端点配置
key: Provider API Key
Returns:
Service Account 场景: ProviderAuthInfo 对象(包含认证信息和解密后的配置)
API Key 场景: None由 build_headers 处理)
Raises:
InvalidRequestException: 认证配置无效或认证失败
"""
from src.core.exceptions import InvalidRequestException
auth_type = getattr(key, "auth_type", "api_key")
if auth_type == "vertex_ai":
from src.core.vertex_auth import VertexAuthError, VertexAuthService
try:
# 优先从 auth_config 读取,兼容从 api_key 读取(过渡期)
encrypted_auth_config = getattr(key, "auth_config", None)
if encrypted_auth_config:
# auth_config 是加密存储的,需要解密
decrypted_config = crypto_service.decrypt(encrypted_auth_config)
sa_json = json.loads(decrypted_config)
else:
# 兼容旧数据:从 api_key 读取
decrypted_key = crypto_service.decrypt(key.api_key)
# 检查是否是占位符(表示 auth_config 丢失)
if decrypted_key == "__placeholder__":
raise InvalidRequestException("认证配置丢失,请重新添加该密钥。")
sa_json = json.loads(decrypted_key)
if not isinstance(sa_json, dict):
raise InvalidRequestException("Service Account JSON 无效,请重新添加该密钥。")
# 获取 Access Token
service = VertexAuthService(sa_json)
access_token = await service.get_access_token()
# Vertex AI 使用 Bearer token
return ProviderAuthInfo(
auth_header="Authorization",
auth_value=f"Bearer {access_token}",
decrypted_auth_config=sa_json,
)
except InvalidRequestException:
raise
except VertexAuthError as e:
raise InvalidRequestException(f"Vertex AI 认证失败:{e}")
except json.JSONDecodeError:
raise InvalidRequestException("Service Account JSON 格式无效,请重新添加该密钥。")
except Exception:
raise InvalidRequestException("Vertex AI 认证失败,请检查 Key 的 auth_config")
# 其他认证类型可在此扩展
# elif auth_type == "oauth2":
# ...
# 标准 API Key返回 None由 build_headers 处理
return None

View File

@@ -270,9 +270,15 @@ async def test_connection(
# 定义请求函数
async def test_request_func(_prov, endpoint, key, _candidate):
from src.api.handlers.base.request_builder import get_provider_auth
# 获取认证信息(处理 Service Account 等异步认证场景)
auth_info = await get_provider_auth(endpoint, key)
request_builder = PassthroughRequestBuilder()
provider_payload, provider_headers = request_builder.build(
payload, {}, endpoint, key, is_stream=False
payload, {}, endpoint, key, is_stream=False,
pre_computed_auth=auth_info.as_tuple() if auth_info else None,
)
url = build_provider_url(
@@ -280,6 +286,8 @@ async def test_connection(
query_params=dict(request.query_params),
path_params={"model": model},
is_stream=False,
key=key,
decrypted_auth_config=auth_info.decrypted_auth_config if auth_info else None,
)
async with httpx.AsyncClient(timeout=30.0, verify=get_ssl_context()) as client:

194
src/core/vertex_auth.py Normal file
View File

@@ -0,0 +1,194 @@
"""
Vertex AI Service Account 认证服务
用于处理 Google Service Account 凭证的 JWT 签名和 Access Token 获取。
Access Token 会被缓存,直到过期前 60 秒才刷新。
"""
import json
import time
from collections import OrderedDict
from typing import Optional, Tuple
import httpx
import jwt
from src.core.logger import logger
class VertexAuthError(Exception):
"""Vertex AI 认证错误"""
pass
def _mask_email(email: str) -> str:
"""脱敏邮箱地址,如 foo@bar.iam.gserviceaccount.com -> foo@***.com"""
if "@" not in email:
return email[:8] + "***" if len(email) > 8 else "***"
local, domain = email.rsplit("@", 1)
# 保留 local 部分前几个字符和域名后缀
masked_local = local[:6] + "***" if len(local) > 6 else local
parts = domain.rsplit(".", 1)
suffix = f".{parts[-1]}" if len(parts) > 1 else ""
return f"{masked_local}@***{suffix}"
class VertexAuthService:
"""
Vertex AI Service Account 认证服务
用于将 Service Account JSON 凭证转换为 Access Token。
使用方式:
service = VertexAuthService(service_account_json)
token = await service.get_access_token()
project_id = service.project_id
# 使用 token 和 project_id 构建请求
"""
# Token 缓存:使用 OrderedDict 实现 LRU
# key = client_email, value = (token, expires_at)
_token_cache: OrderedDict[str, Tuple[str, float]] = OrderedDict()
_cache_max_size: int = 100 # 最多缓存 100 个 Service Account 的 Token
# Token 请求端点
TOKEN_URL = "https://oauth2.googleapis.com/token"
# OAuth2 scope
SCOPE = "https://www.googleapis.com/auth/cloud-platform"
def __init__(self, service_account_json: str):
"""
初始化认证服务
Args:
service_account_json: Service Account JSON 字符串或已解析的字典
"""
if isinstance(service_account_json, str):
try:
self.sa_info = json.loads(service_account_json)
except json.JSONDecodeError as e:
raise VertexAuthError(f"Invalid Service Account JSON: {e}")
else:
self.sa_info = service_account_json
# 验证必需字段
required_fields = ["client_email", "private_key", "project_id"]
missing = [f for f in required_fields if f not in self.sa_info]
if missing:
raise VertexAuthError(f"Service Account JSON missing required fields: {missing}")
self.client_email = self.sa_info["client_email"]
self.private_key = self.sa_info["private_key"]
self.project_id = self.sa_info["project_id"]
def _create_jwt(self) -> str:
"""
创建签名的 JWT
Returns:
签名的 JWT 字符串
"""
now = int(time.time())
payload = {
"iss": self.client_email,
"sub": self.client_email,
"aud": self.TOKEN_URL,
"iat": now,
"exp": now + 3600, # 1 小时有效期
"scope": self.SCOPE,
}
return jwt.encode(payload, self.private_key, algorithm="RS256")
async def get_access_token(self) -> str:
"""
获取 Access Token带 LRU 缓存)
如果缓存中有有效的 Token距离过期超过 60 秒),直接返回。
否则重新获取 Token。缓存采用 LRU 策略,超过 100 个条目时淘汰最旧的。
Returns:
Access Token 字符串
Raises:
VertexAuthError: 获取 Token 失败
"""
# 检查缓存
cache_key = self.client_email
if cache_key in self._token_cache:
token, expires_at = self._token_cache[cache_key]
# 距离过期还有超过 60 秒,使用缓存
if time.time() < expires_at - 60:
# LRU: 移动到末尾(最近使用)
self._token_cache.move_to_end(cache_key)
return token
# 获取新 Token
try:
signed_jwt = self._create_jwt()
async with httpx.AsyncClient(timeout=30) as client:
resp = await client.post(
self.TOKEN_URL,
data={
"grant_type": "urn:ietf:params:oauth:grant-type:jwt-bearer",
"assertion": signed_jwt,
},
)
resp.raise_for_status()
data = resp.json()
access_token = data["access_token"]
expires_in = data.get("expires_in", 3600)
expires_at = time.time() + expires_in
# 缓存 TokenLRU新条目放在末尾
self._token_cache[cache_key] = (access_token, expires_at)
self._token_cache.move_to_end(cache_key)
# LRU 淘汰:超过最大缓存数时移除最旧的条目
while len(self._token_cache) > self._cache_max_size:
oldest_key = next(iter(self._token_cache))
del self._token_cache[oldest_key]
logger.debug(f"[VertexAuth] Evicted oldest cache entry: {_mask_email(oldest_key)}")
logger.debug(
f"[VertexAuth] Obtained access token for {_mask_email(self.client_email)}, "
f"expires in {expires_in}s (cache size: {len(self._token_cache)})"
)
return access_token
except httpx.HTTPStatusError as e:
error_body = e.response.text[:500] if e.response.text else "(empty)"
raise VertexAuthError(f"Failed to get access token: HTTP {e.response.status_code}: {error_body}")
except Exception as e:
raise VertexAuthError(f"Failed to get access token: {e}")
@classmethod
def clear_cache(cls, client_email: Optional[str] = None) -> None:
"""
清除 Token 缓存
Args:
client_email: 指定要清除的账号None 表示清除全部
"""
if client_email:
cls._token_cache.pop(client_email, None)
else:
cls._token_cache.clear()
async def get_vertex_access_token(service_account_json: str) -> Tuple[str, str]:
"""
便捷函数:获取 Vertex AI Access Token 和 Project ID
Args:
service_account_json: Service Account JSON 字符串
Returns:
(access_token, project_id) 元组
"""
service = VertexAuthService(service_account_json)
token = await service.get_access_token()
return token, service.project_id

View File

@@ -1111,8 +1111,22 @@ class ProviderAPIKey(Base):
# None 表示支持所有格式(兼容历史数据),空列表 [] 表示不支持任何格式
api_formats = Column(JSON, nullable=True, default=list) # ["CLAUDE", "CLAUDE_CLI"]
# API密钥信息
api_key = Column(String(500), nullable=False) # API密钥加密存储
# 认证类型
# - "api_key": 标准 API Key 认证(默认
# - "vertex_ai": Google Vertex AI 认证Service Account JSON
# - 未来可扩展oauth2, azure_ad, aws_iam 等
auth_type = Column(String(20), default="api_key", nullable=False)
# API密钥加密存储
# - auth_type="api_key" 时:存储 API Key 字符串
# - auth_type="vertex_ai" 等:可为空,敏感凭证存在 auth_config 中
api_key = Column(String(500), nullable=False) # 保持 NOT NULL 兼容历史数据
# 认证配置(加密存储)
# - auth_type="api_key" 时:可为空
# - auth_type="vertex_ai" 时:存储加密后的 Service Account JSON
# - auth_type="oauth2" 时:存储加密后的 {client_id, client_secret, token_url, scope}
auth_config = Column(Text, nullable=True)
name = Column(String(100), nullable=False) # 密钥名称(必填,用于识别)
note = Column(String(500), nullable=True) # 备注说明(可选)

View File

@@ -4,7 +4,7 @@ ProviderEndpoint 相关的 API 模型定义
import re
from datetime import datetime
from typing import Any, Dict, List, Optional
from typing import Any, Dict, List, Literal, Optional
from pydantic import BaseModel, ConfigDict, Field, field_validator
@@ -166,7 +166,15 @@ class EndpointAPIKeyCreate(BaseModel):
default=None, min_length=1, description="支持的 API 格式列表(必填,路由层校验)"
)
api_key: str = Field(..., min_length=3, max_length=500, description="API Key将自动加密")
api_key: str = Field(default="", max_length=500, description="API Key标准认证时必填,将自动加密)")
auth_type: Literal["api_key", "vertex_ai"] = Field(
default="api_key",
description="认证类型api_key标准 API Key或 vertex_aiVertex AI Service Account"
)
auth_config: Optional[Dict[str, Any]] = Field(
default=None,
description="认证配置JSONvertex_ai 时存储完整 Service Account JSON"
)
name: str = Field(..., min_length=1, max_length=100, description="密钥名称(必填,用于识别)")
# 成本计算
@@ -313,7 +321,15 @@ class EndpointAPIKeyUpdate(BaseModel):
)
api_key: Optional[str] = Field(
default=None, min_length=3, max_length=500, description="API Key将自动加密"
default=None, min_length=3, max_length=500, description="API Key标准认证时使用,将自动加密)"
)
auth_type: Optional[Literal["api_key", "vertex_ai"]] = Field(
default=None,
description="认证类型api_key标准 API Key或 vertex_aiVertex AI Service Account"
)
auth_config: Optional[Dict[str, Any]] = Field(
default=None,
description="认证配置JSONvertex_ai 时存储完整 Service Account JSON"
)
name: Optional[str] = Field(default=None, min_length=1, max_length=100, description="密钥名称")
rate_multipliers: Optional[Dict[str, float]] = Field(
@@ -445,6 +461,8 @@ class EndpointAPIKeyResponse(BaseModel):
# Key 信息(脱敏)
api_key_masked: str = Field(..., description="脱敏后的 Key")
api_key_plain: Optional[str] = Field(default=None, description="完整的 Key")
auth_type: str = Field(default="api_key", description="认证类型api_key 或 vertex_ai")
# auth_config 不在响应中返回(包含敏感信息),前端通过 auth_type 判断类型
name: str = Field(..., description="密钥名称")
# 成本计算

View File

@@ -74,6 +74,87 @@ class GeminiUsageMetadata(BaseModelWithExtras):
total_token_count: int = Field(default=0, alias="totalTokenCount")
# ---------------------------------------------------------------------------
# 文件 API 模型
# ---------------------------------------------------------------------------
class GeminiFileMetadata(BaseModelWithExtras):
"""
Gemini 文件元数据
用于上传文件时指定的元数据信息
"""
display_name: Optional[str] = Field(default=None, alias="displayName")
class GeminiFileUploadRequest(BaseModelWithExtras):
"""
Gemini 文件上传请求
用于 media.upload API 的请求体
"""
file: Optional[GeminiFileMetadata] = None
class GeminiFile(BaseModelWithExtras):
"""
Gemini 文件资源
表示已上传到 Gemini API 的文件
"""
name: Optional[str] = None # 文件名格式files/xxx
display_name: Optional[str] = Field(default=None, alias="displayName")
mime_type: Optional[str] = Field(default=None, alias="mimeType")
size_bytes: Optional[str] = Field(default=None, alias="sizeBytes")
create_time: Optional[str] = Field(default=None, alias="createTime")
update_time: Optional[str] = Field(default=None, alias="updateTime")
expiration_time: Optional[str] = Field(default=None, alias="expirationTime")
sha256_hash: Optional[str] = Field(default=None, alias="sha256Hash")
uri: Optional[str] = None # 文件 URI用于在请求中引用
download_uri: Optional[str] = Field(default=None, alias="downloadUri")
state: Optional[str] = None # PROCESSING, ACTIVE, FAILED
error: Optional[Dict[str, Any]] = None
# 视频文件元数据
video_metadata: Optional[Dict[str, Any]] = Field(default=None, alias="videoMetadata")
class GeminiFileListResponse(BaseModelWithExtras):
"""
Gemini 文件列表响应
用于 files.list API 的响应体
"""
files: Optional[List["GeminiFile"]] = None
next_page_token: Optional[str] = Field(default=None, alias="nextPageToken")
class GeminiFileUploadResponse(BaseModelWithExtras):
"""
Gemini 文件上传响应
用于 media.upload API 的响应体
"""
file: Optional[GeminiFile] = None
class GeminiFilePart(BaseModelWithExtras):
"""
Gemini 文件引用部分
用于在请求内容中引用已上传的文件
使用 file_data 字段引用文件 URI
"""
file_data: Optional[Dict[str, Any]] = Field(default=None, alias="fileData")
# fileData 格式:{"mimeType": "...", "fileUri": "..."}
# ---------------------------------------------------------------------------
# Thought Signature 常量
# ---------------------------------------------------------------------------

View File

@@ -311,6 +311,14 @@ class ModelFetchScheduler:
key.last_models_fetch_at = now
return "error"
# Vertex AI 类型不支持自动获取模型(需要使用 Service Account 认证)
auth_type = getattr(key, "auth_type", "api_key") or "api_key"
if auth_type == "vertex_ai":
key.last_models_fetch_error = "auto_fetch_models 暂不支持 Vertex AI 类型的 Key"
key.last_models_fetch_at = now
logger.info(f"Key {key.id} 为 Vertex AI 类型,跳过自动获取模型")
return "skip"
# 解密 API Key
if not key.api_key:
logger.warning(f"Key {key.id} 没有 API Key跳过")

View File

@@ -4,6 +4,7 @@
负责:
- 根据 API 格式或端点配置生成请求 URL
- URL 脱敏(用于日志记录)
- Vertex AI URL 自动构建
"""
import re
@@ -14,7 +15,7 @@ from src.core.api_format import APIFormat, get_default_path, resolve_api_format
from src.core.logger import logger
if TYPE_CHECKING:
from src.models.database import ProviderEndpoint
from src.models.database import ProviderAPIKey, ProviderEndpoint
# URL 中需要脱敏的查询参数(正则模式)
@@ -69,20 +70,36 @@ def build_provider_url(
query_params: Optional[Dict[str, Any]] = None,
path_params: Optional[Dict[str, Any]] = None,
is_stream: bool = False,
key: Optional["ProviderAPIKey"] = None,
decrypted_auth_config: Optional[Dict[str, Any]] = None,
) -> str:
"""
根据 endpoint 配置生成请求 URL
优先级:
1. endpoint.custom_path - 自定义路径(支持模板变量如 {model}
2. API 格式默认路径 - 根据 api_format 自动选择
1. Vertex AI 自动构建 - 当 key.auth_type == "vertex_ai"
2. endpoint.custom_path - 自定义路径(支持模板变量如 {model}
3. API 格式默认路径 - 根据 api_format 自动选择
Args:
endpoint: 端点配置
query_params: 查询参数
path_params: 路径模板参数 (如 {model})
is_stream: 是否为流式请求,用于 Gemini API 选择正确的操作方法
key: Provider API Key用于 Vertex AI 等需要从密钥配置读取信息的场景)
decrypted_auth_config: 已解密的认证配置(避免重复解密,由 get_provider_auth 提供)
"""
# 检查是否为 Vertex AI 认证类型
auth_type = getattr(key, "auth_type", "api_key") if key else "api_key"
if auth_type == "vertex_ai":
return _build_vertex_ai_url(
key=key,
path_params=path_params,
query_params=query_params,
is_stream=is_stream,
decrypted_auth_config=decrypted_auth_config,
)
# 准备路径参数,添加 Gemini API 所需的 action 参数
effective_path_params = dict(path_params) if path_params else {}
@@ -152,3 +169,131 @@ def _resolve_default_path(api_format: Optional[str]) -> str:
logger.warning(f"Unknown api_format '{api_format}' for endpoint, fallback to '/'")
return "/"
# Vertex AI 模型默认 region 映射
# 用户可以通过 auth_config.model_regions 覆盖
VERTEX_AI_DEFAULT_MODEL_REGIONS: Dict[str, str] = {
# Gemini 3 系列(使用 global
"gemini-3-pro-image-preview": "global",
# Gemini 2.0 系列
"gemini-2.0-flash": "us-central1",
"gemini-2.0-flash-exp": "us-central1",
"gemini-2.0-flash-001": "us-central1",
"gemini-2.0-pro-exp": "us-central1",
"gemini-2.0-flash-exp-image-generation": "us-central1",
# Gemini 1.5 系列
"gemini-1.5-pro": "us-central1",
"gemini-1.5-pro-001": "us-central1",
"gemini-1.5-pro-002": "us-central1",
"gemini-1.5-flash": "us-central1",
"gemini-1.5-flash-001": "us-central1",
"gemini-1.5-flash-002": "us-central1",
# Imagen 系列
"imagen-3.0-generate-001": "us-central1",
"imagen-3.0-fast-generate-001": "us-central1",
}
def _build_vertex_ai_url(
key: "ProviderAPIKey",
*,
path_params: Optional[Dict[str, Any]] = None,
query_params: Optional[Dict[str, Any]] = None,
is_stream: bool = False,
decrypted_auth_config: Optional[Dict[str, Any]] = None,
) -> str:
"""
构建 Vertex AI URL
Vertex AI URL 格式:
https://{region}-aiplatform.googleapis.com/v1/projects/{project_id}/locations/{region}/publishers/google/models/{model}:{action}
从 auth_config 中读取:
- project_id: GCP 项目 ID必需
- region: 默认 GCP 区域(覆盖内置默认值)
- model_regions: 模型到区域的映射(可选),覆盖内置和默认配置
Region 优先级:
1. auth_config.model_regions[model] - 用户为该模型指定的区域
2. VERTEX_AI_DEFAULT_MODEL_REGIONS[model] - 内置的模型默认区域
3. auth_config.region - 用户配置的默认区域
4. global - 最终兜底
Args:
key: Provider API Key包含 auth_config
path_params: 路径参数(需要 model
query_params: 查询参数
is_stream: 是否为流式请求
decrypted_auth_config: 已解密的认证配置(由 get_provider_auth 提供,避免重复解密)
Returns:
完整的 Vertex AI URL
"""
import json
from src.core.crypto import crypto_service
# 优先使用传入的已解密配置,避免重复解密
auth_config: Dict[str, Any] = {}
if decrypted_auth_config:
auth_config = decrypted_auth_config
else:
# 兜底:从 key.auth_config 解密(理论上不应走到这里)
encrypted_auth_config = getattr(key, "auth_config", None)
if encrypted_auth_config:
try:
decrypted_config = crypto_service.decrypt(encrypted_auth_config)
auth_config = json.loads(decrypted_config)
except Exception as e:
logger.error(f"解密 Vertex AI auth_config 失败: {e}")
auth_config = {}
from src.core.exceptions import InvalidRequestException
# 获取必需的配置
project_id = auth_config.get("project_id")
if not project_id:
raise InvalidRequestException("Vertex AI 配置缺少 project_id请在 Key 的 auth_config 中提供)")
# 获取模型名
model = (path_params or {}).get("model", "")
if not model:
raise InvalidRequestException("Vertex AI 请求缺少 model 参数")
# 确定 region优先级用户配置 > 内置默认 > 用户默认 > 兜底)
user_model_regions = auth_config.get("model_regions", {})
user_default_region = auth_config.get("region")
if model in user_model_regions:
region = user_model_regions[model]
elif model in VERTEX_AI_DEFAULT_MODEL_REGIONS:
region = VERTEX_AI_DEFAULT_MODEL_REGIONS[model]
elif user_default_region:
region = user_default_region
else:
region = "global"
# 确定 action
action = "streamGenerateContent" if is_stream else "generateContent"
# 构建 URLglobal region 使用不同的 URL 格式)
if region == "global":
base_url = "https://aiplatform.googleapis.com"
else:
base_url = f"https://{region}-aiplatform.googleapis.com"
path = f"/v1/projects/{project_id}/locations/{region}/publishers/google/models/{model}:{action}"
url = f"{base_url}{path}"
# 添加查询参数
effective_query_params = dict(query_params) if query_params else {}
# Vertex AI 流式请求使用 SSE 格式
if is_stream:
effective_query_params.setdefault("alt", "sse")
if effective_query_params:
query_string = urlencode(effective_query_params, doseq=True)
if query_string:
url = f"{url}?{query_string}"
logger.debug(f"Vertex AI URL: {redact_url_for_log(url)} (region={region})")
return url