mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-02 01:10:23 +08:00
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:
@@ -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,
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
194
src/core/vertex_auth.py
Normal 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
|
||||
|
||||
# 缓存 Token(LRU:新条目放在末尾)
|
||||
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
|
||||
@@ -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) # 备注说明(可选)
|
||||
|
||||
|
||||
@@ -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_ai(Vertex AI Service Account)"
|
||||
)
|
||||
auth_config: Optional[Dict[str, Any]] = Field(
|
||||
default=None,
|
||||
description="认证配置(JSON):vertex_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_ai(Vertex AI Service Account)"
|
||||
)
|
||||
auth_config: Optional[Dict[str, Any]] = Field(
|
||||
default=None,
|
||||
description="认证配置(JSON):vertex_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="密钥名称")
|
||||
|
||||
# 成本计算
|
||||
|
||||
@@ -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 常量
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
@@ -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,跳过")
|
||||
|
||||
@@ -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"
|
||||
|
||||
# 构建 URL(global 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
|
||||
|
||||
Reference in New Issue
Block a user