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

feat: 添加 Vertex AI Claude 模型支持和动态格式识别

Browse Source 
- 新增 get_vertex_ai_effective_format 函数，根据模型名自动判断 API 格式
- 支持通过 auth_config.model_format_mapping 自定义模型格式映射
- 修改 Vertex AI URL 构建逻辑，Claude 模型使用 anthropic publisher 和 rawPredict
- 修复 auth_config 可能是未加密 dict 的兼容性问题
This commit is contained in:
fawney19
2026-01-31 00:30:07 +08:00
parent 772cb90f64
commit dc4bb25cc2
3 changed files with 194 additions and 29 deletions
Download Patch File Download Diff File Expand all files Collapse all files

88
src/api/handlers/base/chat_handler_base.py
View File

@@ -24,11 +24,9 @@ from __future__ import annotations
import asyncio
import json
from abc import ABC, abstractmethod
from collections.abc import AsyncGenerator, Awaitable, Callable
from typing import Any
from collections.abc import Callable
from collections.abc import AsyncGenerator, Awaitable
import httpx
from fastapi import BackgroundTasks, Request
from fastapi.responses import JSONResponse, StreamingResponse
@@ -70,7 +68,11 @@ from src.models.database import (
User,
)
from src.services.cache.aware_scheduler import ProviderCandidate
from src.services.provider.transport import build_provider_url, redact_url_for_log
from src.services.provider.transport import (
build_provider_url,
get_vertex_ai_effective_format,
redact_url_for_log,
)
def _get_error_status_code(e: Exception, default: int = 400) -> int:
@@ -79,6 +81,54 @@ def _get_error_status_code(e: Exception, default: int = 400) -> int:
return code if isinstance(code, int) and code > 0 else default
def _resolve_vertex_ai_format(
key: ProviderAPIKey,
auth_info: Any,
model: str,
provider_api_format: str,
client_api_format: str,
candidate: ProviderCandidate | None,
) -> tuple[str, bool]:
"""
解析 Vertex AI 动态格式并计算 needs_conversion
当 auth_type=vertex_ai 时，同一个 GCP 项目可以访问 Gemini 和 Claude，
但它们的请求/响应格式不同，需要根据模型名动态选择。
用户可通过 auth_config.model_format_mapping 配置自定义映射。
Args:
key: Provider API Key
auth_info: 认证信息（包含 decrypted_auth_config）
model: 模型名
provider_api_format: 当前 provider API 格式
client_api_format: 客户端 API 格式
candidate: Provider 候选（用于获取原始 needs_conversion）
Returns:
(effective_provider_format, needs_conversion) 元组
"""
key_auth_type = getattr(key, "auth_type", "api_key")
if key_auth_type == "vertex_ai":
vertex_auth_config = auth_info.decrypted_auth_config if auth_info else None
effective_format = get_vertex_ai_effective_format(model, vertex_auth_config)
if effective_format.upper() != provider_api_format.upper():
logger.debug(
f"Vertex AI 动态格式切换: {provider_api_format} -> {effective_format} "
f"(model={model})"
)
provider_api_format = effective_format
# Vertex AI 模式下，根据动态格式与客户端格式比较确定是否需要转换
needs_conversion = provider_api_format.upper() != client_api_format.upper()
else:
# 非 Vertex AI：使用 candidate 的 needs_conversion
needs_conversion = (
bool(getattr(candidate, "needs_conversion", False)) if candidate else False
)
return provider_api_format, needs_conversion
def _convert_error_response_best_effort(
error_response: dict[str, Any],
source_format: str,
@@ -595,7 +645,9 @@ class ChatHandlerBase(BaseMessageHandler, ABC):
except (ThinkingSignatureException, UpstreamClientException) as e:
# ThinkingSignatureException: orchestrator 层已处理整流重试但仍失败
# UpstreamClientException: 上游客户端错误（HTTP 4xx），不重试，直接返回给客户端
error_type = "签名错误" if isinstance(e, ThinkingSignatureException) else "上游客户端错误"
error_type = (
"签名错误" if isinstance(e, ThinkingSignatureException) else "上游客户端错误"
)
self._log_request_error(f"流式请求失败（{error_type}）", e)
await self._record_stream_failure(ctx, e, original_headers, original_request_body)
client_format = (ctx.client_api_format or "").upper()
@@ -645,9 +697,15 @@ class ChatHandlerBase(BaseMessageHandler, ABC):
)
provider_api_format = ctx.provider_api_format or _api_format_str
client_api_format = ctx.client_api_format or _api_format_str
needs_conversion = (
bool(getattr(candidate, "needs_conversion", False)) if candidate else False
# 提前获取认证信息（Vertex AI 格式判断需要使用 auth_config）
auth_info = await get_provider_auth(endpoint, key)
# 解析 Vertex AI 动态格式并计算 needs_conversion
provider_api_format, needs_conversion = _resolve_vertex_ai_format(
key, auth_info, ctx.model, provider_api_format, client_api_format, candidate
)
ctx.provider_api_format = provider_api_format
ctx.needs_conversion = needs_conversion
# 获取模型映射（优先使用映射匹配到的模型，其次是 Provider 级别的映射）
@@ -692,9 +750,6 @@ 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,
@@ -956,7 +1011,15 @@ class ChatHandlerBase(BaseMessageHandler, ABC):
client_api_format = (
api_format.value if hasattr(api_format, "value") else str(api_format)
)
needs_conversion = bool(getattr(candidate, "needs_conversion", False))
# 提前获取认证信息（Vertex AI 格式判断需要使用 auth_config）
auth_info = await get_provider_auth(endpoint, key)
# 解析 Vertex AI 动态格式并计算 needs_conversion
provider_api_format, needs_conversion = _resolve_vertex_ai_format(
key, auth_info, model, provider_api_format, client_api_format, candidate
)
provider_api_format_for_error = provider_api_format
client_api_format_for_error = client_api_format
needs_conversion_for_error = needs_conversion
@@ -1003,9 +1066,6 @@ 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,

14
src/api/handlers/base/request_builder.py
View File

@@ -18,8 +18,8 @@ from abc import ABC, abstractmethod
from dataclasses import dataclass
from typing import TYPE_CHECKING, Any
from src.core.api_format import UPSTREAM_DROP_HEADERS, HeaderBuilder
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
@@ -43,6 +43,7 @@ class ProviderAuthInfo:
"""返回 (auth_header, auth_value) 元组"""
return (self.auth_header, self.auth_value)
# ==============================================================================
# 统一的头部配置常量
# ==============================================================================
@@ -346,9 +347,14 @@ async def get_provider_auth(
# 优先从 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)
# auth_config 可能是加密字符串或未加密的 dict
if isinstance(encrypted_auth_config, dict):
# 已经是 dict，直接使用（兼容未加密存储的情况）
sa_json = encrypted_auth_config
else:
# 是加密字符串，需要解密
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)

121
src/services/provider/transport.py
View File

@@ -173,6 +173,85 @@ def _resolve_default_path(api_format: str | None) -> str:
return "/"
# ==============================================================================
# Vertex AI 配置
# ==============================================================================
# Vertex AI 模型前缀到 API 格式的映射
# 用于 auth_type=vertex_ai 时，根据模型名动态确定实际的请求/响应格式
# 格式：前缀 -> APIFormat 值
VERTEX_AI_MODEL_FORMAT_MAPPING: dict[str, str] = {
"claude-": "CLAUDE", # Anthropic Claude 模型
"gemini-": "GEMINI", # Google Gemini 模型
"imagen-": "GEMINI", # Google Imagen 模型（使用 Gemini 格式）
}
# Vertex AI 默认 API 格式（当模型前缀不匹配时）
VERTEX_AI_DEFAULT_FORMAT: str = "GEMINI"
def get_vertex_ai_effective_format(
model: str,
auth_config: dict[str, Any] | None = None,
) -> str:
"""
获取 Vertex AI 模式下模型的实际 API 格式
优先级：
1. auth_config.model_format_mapping 中的精确匹配
2. auth_config.model_format_mapping 中的前缀匹配
3. 内置 VERTEX_AI_MODEL_FORMAT_MAPPING 前缀匹配
4. auth_config.default_format
5. 内置 VERTEX_AI_DEFAULT_FORMAT
auth_config 配置示例::
{
"project_id": "your-gcp-project-id",
"model_format_mapping": {
"claude-": "CLAUDE", # 前缀匹配
"my-custom-model": "OPENAI" # 精确匹配
},
"default_format": "GEMINI"
}
Args:
model: 模型名称
auth_config: 解密后的认证配置（可选），可包含 model_format_mapping 和 default_format
Returns:
实际应使用的 API 格式（如 "CLAUDE", "GEMINI"）
"""
# 用户配置的模型-格式映射
user_format_mapping: dict[str, str] = {}
user_default_format: str | None = None
if auth_config:
user_format_mapping = auth_config.get("model_format_mapping", {})
user_default_format = auth_config.get("default_format")
# 1. 用户配置：精确匹配
if model in user_format_mapping:
return user_format_mapping[model].upper()
# 2. 用户配置：前缀匹配
for prefix, api_format in user_format_mapping.items():
if prefix.endswith("-") and model.startswith(prefix):
return api_format.upper()
# 3. 内置配置：前缀匹配
for prefix, api_format in VERTEX_AI_MODEL_FORMAT_MAPPING.items():
if model.startswith(prefix):
return api_format
# 4. 用户默认格式
if user_default_format:
return user_default_format.upper()
# 5. 内置默认格式
return VERTEX_AI_DEFAULT_FORMAT
# Vertex AI 模型默认 region 映射
# 用户可以通过 auth_config.model_regions 覆盖
VERTEX_AI_DEFAULT_MODEL_REGIONS: dict[str, str] = {
@@ -209,7 +288,8 @@ def _build_vertex_ai_url(
构建 Vertex AI URL
Vertex AI URL 格式:
https://{region}-aiplatform.googleapis.com/v1/projects/{project_id}/locations/{region}/publishers/google/models/{model}:{action}
- Gemini: https://{region}-aiplatform.googleapis.com/v1/projects/{project_id}/locations/{region}/publishers/google/models/{model}:{action}
- Claude: https://{region}-aiplatform.googleapis.com/v1/projects/{project_id}/locations/{region}/publishers/anthropic/models/{model}:{action}
从 auth_config 中读取:
- project_id: GCP 项目 ID（必需）
@@ -233,6 +313,7 @@ def _build_vertex_ai_url(
完整的 Vertex AI URL
"""
import json
from src.core.crypto import crypto_service
# 优先使用传入的已解密配置，避免重复解密
@@ -241,11 +322,15 @@ def _build_vertex_ai_url(
auth_config = decrypted_auth_config
else:
# 兜底：从 key.auth_config 解密（理论上不应走到这里）
encrypted_auth_config = getattr(key, "auth_config", None)
if encrypted_auth_config:
raw_auth_config = getattr(key, "auth_config", None)
if raw_auth_config:
try:
decrypted_config = crypto_service.decrypt(encrypted_auth_config)
auth_config = json.loads(decrypted_config)
# auth_config 可能是加密字符串或未加密的 dict
if isinstance(raw_auth_config, dict):
auth_config = raw_auth_config
else:
decrypted_config = crypto_service.decrypt(raw_auth_config)
auth_config = json.loads(decrypted_config)
except Exception as e:
logger.error(f"解密 Vertex AI auth_config 失败: {e}")
auth_config = {}
@@ -255,7 +340,9 @@ def _build_vertex_ai_url(
# 获取必需的配置
project_id = auth_config.get("project_id")
if not project_id:
raise InvalidRequestException("Vertex AI 配置缺少 project_id（请在 Key 的 auth_config 中提供）")
raise InvalidRequestException(
"Vertex AI 配置缺少 project_id（请在 Key 的 auth_config 中提供）"
)
# 获取模型名
model = (path_params or {}).get("model", "")
@@ -275,22 +362,34 @@ def _build_vertex_ai_url(
else:
region = "global"
# 确定 action
action = "streamGenerateContent" if is_stream else "generateContent"
# 判断是 Claude 还是 Gemini 模型
is_claude_model = model.startswith("claude-")
# 根据模型类型确定 publisher 和 action
if is_claude_model:
# Claude 模型使用 Anthropic publisher
publisher = "anthropic"
action = "streamRawPredict" if is_stream else "rawPredict"
else:
# Gemini 模型使用 Google publisher
publisher = "google"
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}"
path = f"/v1/projects/{project_id}/locations/{region}/publishers/{publisher}/models/{model}:{action}"
url = f"{base_url}{path}"
# 添加查询参数
effective_query_params = dict(query_params) if query_params else {}
# Vertex AI 流式请求使用 SSE 格式
if is_stream:
# Gemini 流式请求使用 SSE 格式，Claude 不需要
if is_stream and not is_claude_model:
effective_query_params.setdefault("alt", "sse")
# 移除不适用于 Vertex AI 的参数
effective_query_params.pop("beta", None)
if effective_query_params:
query_string = urlencode(effective_query_params, doseq=True)