i/Aether
1
0
Fork 0
mirror of https://github.com/fawney19/Aether.git synced 2026-09-03 01:40:21 +08:00
Code Issues Packages Projects Releases Wiki Activity

feat: add Vertex AI authentication support for provider API keys

Browse Source 
- 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
Download Patch File Download Diff File Expand all files Collapse all files

14
src/api/handlers/base/chat_handler_base.py
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(

14
src/api/handlers/base/cli_handler_base.py
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(

131
src/api/handlers/base/request_builder.py
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