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

refactor: 继续拆分大型模块并增强模块注册健壮性

Browse Source 
后端:
- chat_handler_base 错误处理函数提取到 chat_error_utils 子模块
- CLI mixin 引入 CliHandlerProtocol 协议类改善类型标注
- aware_scheduler 拆分为 _candidate_builder 和 _candidate_sorter 子模块
- usage recording 拆分为 _billing_integration 和 _recording_helpers 子模块
- ModuleRegistry 添加循环依赖检测，将写操作从查询方法分离到 reconcile_module_state
- 修正 plugin manager 入度注释

前端:
- 路由守卫逻辑拆分为独立 guards 模块
- ProviderManagement 拆分为 TableHeader/TableRow/BalanceCell/MobileCard 子组件
- SystemSettings 拆分为多个 Section 子组件和 composables
- 提取 useEndpointStatus/useProviderBalance/useProviderFilters composables

测试适配重构后的子模块结构
This commit is contained in:
fawney19
2026-02-14 20:06:38 +08:00
parent 676e918edc
commit 8a670f5524
51 changed files with 7095 additions and 5359 deletions
Download Patch File Download Diff File Expand all files Collapse all files

150
src/api/handlers/base/chat_error_utils.py Normal file
View File

@@ -0,0 +1,150 @@
"""
Chat Error Utils - Chat Handler 错误处理工具函数
从 chat_handler_base.py 提取的模块级工具函数，用于错误响应的构建和转换。
"""
from __future__ import annotations
import json
from typing import Any
from src.api.handlers.base.utils import get_format_converter_registry
from src.core.exceptions import ThinkingSignatureException, UpstreamClientException
from src.core.logger import logger
from src.models.database import ProviderAPIKey
from src.services.cache.aware_scheduler import ProviderCandidate
from src.services.provider.transport import get_vertex_ai_effective_format
def _get_error_status_code(e: Exception, default: int = 400) -> int:
"""从异常中提取 HTTP 状态码"""
code = getattr(e, "status_code", None)
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,
target_format: str,
) -> dict[str, Any]:
"""
将上游错误响应 best-effort 转换为客户端格式。
说明：错误转换走 Canonical registry。转换失败时构造安全的通用错误响应，
避免泄露上游原始错误详情。
"""
try:
registry = get_format_converter_registry()
return registry.convert_error_response(error_response, source_format, target_format)
except Exception as e:
logger.debug(f"错误响应转换失败 ({source_format} -> {target_format}): {e}")
# 转换失败时构造安全的通用错误，避免泄露上游详情
return _build_client_error_response_best_effort("upstream error", target_format)
def _build_client_error_response_best_effort(
message: str,
target_format: str,
) -> dict[str, Any]:
"""
当无法解析上游错误 body 时，构造一个目标格式的错误响应（best-effort）。
"""
try:
from src.core.api_format.conversion.internal import ErrorType, InternalError
registry = get_format_converter_registry()
normalizer = registry.get_normalizer(target_format)
if normalizer and normalizer.capabilities.supports_error_conversion:
return normalizer.error_from_internal(
InternalError(type=ErrorType.INVALID_REQUEST, message=message, retryable=False)
)
except Exception as e:
logger.debug(f"构建客户端错误响应失败 (target={target_format}): {e}")
return {"error": {"type": "upstream_client_error", "message": message}}
def _build_error_json_payload(
e: ThinkingSignatureException | UpstreamClientException,
client_format: str,
provider_format: str,
needs_conversion: bool = True,
) -> dict[str, Any]:
"""
构建错误 JSON 响应 payload（公共逻辑）。
从异常中提取上游错误信息，尝试转换为客户端格式。
Args:
e: ThinkingSignatureException 或 UpstreamClientException
client_format: 客户端 API 格式
provider_format: Provider API 格式
needs_conversion: 是否需要格式转换
Returns:
格式化的错误响应字典
"""
raw = getattr(e, "upstream_error", None)
message = getattr(e, "message", str(e))
if isinstance(raw, str) and raw:
try:
parsed = json.loads(raw)
except json.JSONDecodeError:
parsed = None
if isinstance(parsed, dict):
if needs_conversion:
return _convert_error_response_best_effort(parsed, provider_format, client_format)
return parsed
return _build_client_error_response_best_effort(message, client_format)

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

@@ -38,19 +38,19 @@ from src.api.handlers.base.base_handler import (
ClientDisconnectedException,
wait_for_with_disconnect_detection,
)
from src.api.handlers.base.chat_error_utils import (
_build_error_json_payload,
_get_error_status_code,
_resolve_vertex_ai_format,
)
from src.api.handlers.base.parsers import get_parser_for_format
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,
extract_proxy_timing,
is_format_converted,
)
from src.api.handlers.base.stream_processor import StreamProcessor
from src.api.handlers.base.stream_telemetry import StreamTelemetryRecorder
from src.api.handlers.base.upstream_stream_bridge import (
aggregate_upstream_stream_to_internal_response,
)
from src.api.handlers.base.utils import (
build_sse_headers,
filter_proxy_response_headers,
@@ -60,12 +60,9 @@ from src.config.settings import config
from src.core.api_format.conversion.stream_bridge import (
iter_internal_response_as_stream_events,
)
from src.core.error_utils import extract_client_error_message
from src.core.exceptions import (
EmbeddedErrorException,
ProviderAuthException,
ProviderNotAvailableException,
ProviderRateLimitException,
ProviderTimeoutException,
ThinkingSignatureException,
UpstreamClientException,
@@ -87,145 +84,10 @@ from src.services.provider.stream_policy import (
)
from src.services.provider.transport import (
build_provider_url,
get_vertex_ai_effective_format,
redact_url_for_log,
)
from src.services.system.config import SystemConfigService
def _get_error_status_code(e: Exception, default: int = 400) -> int:
"""从异常中提取 HTTP 状态码"""
code = getattr(e, "status_code", None)
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,
target_format: str,
) -> dict[str, Any]:
"""
将上游错误响应 best-effort 转换为客户端格式。
说明：错误转换走 Canonical registry。转换失败时构造安全的通用错误响应，
避免泄露上游原始错误详情。
"""
try:
registry = get_format_converter_registry()
return registry.convert_error_response(error_response, source_format, target_format)
except Exception as e:
logger.debug(f"错误响应转换失败 ({source_format} -> {target_format}): {e}")
# 转换失败时构造安全的通用错误，避免泄露上游详情
return _build_client_error_response_best_effort("upstream error", target_format)
def _build_client_error_response_best_effort(
message: str,
target_format: str,
) -> dict[str, Any]:
"""
当无法解析上游错误 body 时，构造一个目标格式的错误响应（best-effort）。
"""
try:
from src.core.api_format.conversion.internal import ErrorType, InternalError
registry = get_format_converter_registry()
normalizer = registry.get_normalizer(target_format)
if normalizer and normalizer.capabilities.supports_error_conversion:
return normalizer.error_from_internal(
InternalError(type=ErrorType.INVALID_REQUEST, message=message, retryable=False)
)
except Exception as e:
logger.debug(f"构建客户端错误响应失败 (target={target_format}): {e}")
return {"error": {"type": "upstream_client_error", "message": message}}
def _build_error_json_payload(
e: ThinkingSignatureException | UpstreamClientException,
client_format: str,
provider_format: str,
needs_conversion: bool = True,
) -> dict[str, Any]:
"""
构建错误 JSON 响应 payload（公共逻辑）。
从异常中提取上游错误信息，尝试转换为客户端格式。
Args:
e: ThinkingSignatureException 或 UpstreamClientException
client_format: 客户端 API 格式
provider_format: Provider API 格式
needs_conversion: 是否需要格式转换
Returns:
格式化的错误响应字典
"""
raw = getattr(e, "upstream_error", None)
message = getattr(e, "message", str(e))
if isinstance(raw, str) and raw:
try:
parsed = json.loads(raw)
except json.JSONDecodeError:
parsed = None
if isinstance(parsed, dict):
if needs_conversion:
return _convert_error_response_best_effort(parsed, provider_format, client_format)
return parsed
return _build_client_error_response_best_effort(message, client_format)
@dataclass
class ProviderRequestResult:
"""_prepare_provider_request() 的返回结果，封装请求构建阶段的所有产出。"""
@@ -756,7 +618,11 @@ class ChatHandlerBase(BaseMessageHandler, ABC):
"签名错误" 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)
from src.api.handlers.base.chat_sync_executor import ChatSyncExecutor
await ChatSyncExecutor(self)._record_stream_failure(
ctx, e, original_headers, original_request_body
)
client_format = (ctx.client_api_format or "").upper()
provider_format = (ctx.provider_api_format or client_format).upper()
payload = _build_error_json_payload(
@@ -769,7 +635,11 @@ class ChatHandlerBase(BaseMessageHandler, ABC):
except Exception as e:
self._log_request_error("流式请求失败", e)
await self._record_stream_failure(ctx, e, original_headers, original_request_body)
from src.api.handlers.base.chat_sync_executor import ChatSyncExecutor
await ChatSyncExecutor(self)._record_stream_failure(
ctx, e, original_headers, original_request_body
)
raise
async def _prepare_provider_request(
@@ -1351,7 +1221,9 @@ class ChatHandlerBase(BaseMessageHandler, ABC):
response_ctx = None
continue
error_text = await self._extract_error_text(e)
from src.api.handlers.base.chat_sync_executor import ChatSyncExecutor
error_text = await ChatSyncExecutor(self)._extract_error_text(e)
logger.error(
f"Provider 返回错误: {e.response.status_code}\n Response: {error_text}"
)
@@ -1384,57 +1256,6 @@ class ChatHandlerBase(BaseMessageHandler, ABC):
start_time=self.start_time,
)
async def _record_stream_failure(
self,
ctx: StreamContext,
error: Exception,
original_headers: dict[str, str],
original_request_body: dict[str, Any],
) -> None:
"""记录流式请求失败"""
response_time_ms = self.elapsed_ms()
status_code = 503
if isinstance(error, ThinkingSignatureException):
status_code = 400
elif isinstance(error, UpstreamClientException):
status_code = _get_error_status_code(error)
elif isinstance(error, ProviderAuthException):
status_code = 503
elif isinstance(error, ProviderRateLimitException):
status_code = 429
elif isinstance(error, ProviderTimeoutException):
status_code = 504
actual_request_body = ctx.provider_request_body or original_request_body
# 失败时返回给客户端的是 JSON 错误响应
client_response_headers = {"content-type": "application/json"}
stream_fail_metadata: dict[str, Any] | None = None
if ctx.proxy_info:
stream_fail_metadata = {"proxy": ctx.proxy_info}
await self.telemetry.record_failure(
provider=ctx.provider_name or "unknown",
model=ctx.model,
response_time_ms=response_time_ms,
status_code=status_code,
error_message=extract_client_error_message(error),
request_headers=original_headers,
request_body=actual_request_body,
is_stream=True,
api_format=ctx.api_format,
provider_request_headers=ctx.provider_request_headers,
response_headers=ctx.response_headers,
client_response_headers=client_response_headers,
# 格式转换追踪
endpoint_api_format=ctx.provider_api_format or None,
has_format_conversion=ctx.has_format_conversion,
target_model=ctx.mapped_model,
request_metadata=stream_fail_metadata,
)
# ==================== 非流式处理 ====================
async def process_sync(
@@ -1446,561 +1267,9 @@ class ChatHandlerBase(BaseMessageHandler, ABC):
query_params: dict[str, str] | None = None,
) -> JSONResponse:
"""处理非流式响应"""
logger.debug(f"开始非流式响应处理 ({self.FORMAT_ID})")
from src.api.handlers.base.chat_sync_executor import ChatSyncExecutor
# 转换请求格式
converted_request = await self._convert_request(request)
model = getattr(converted_request, "model", original_request_body.get("model", "unknown"))
api_format = self.allowed_api_formats[0]
# 提前创建 pending 记录，让前端可以立即看到"处理中"
self._create_pending_usage(
model=model,
is_stream=False,
request_type="chat",
api_format=self.FORMAT_ID,
request_headers=original_headers,
request_body=original_request_body,
executor = ChatSyncExecutor(self)
return await executor.execute(
request, http_request, original_headers, original_request_body, query_params
)
# 可变请求体容器：允许 TaskService 在遇到 Thinking 签名错误时整流请求体后重试
# 结构: {"body": 实际请求体, "_rectified": 是否已整流, "_rectified_this_turn": 本轮是否整流}
request_body_ref: dict[str, Any] = {"body": original_request_body}
# 用于跟踪的变量
provider_name: str | None = None
response_json: dict[str, Any] | None = None
status_code = 200
response_headers: dict[str, str] = {}
provider_request_headers: dict[str, str] = {}
provider_request_body: dict[str, Any] | None = None
provider_api_format_for_error: str | None = None
client_api_format_for_error: str | None = None
needs_conversion_for_error: bool = False # 用于构建错误 payload（含 envelope rewrite）
provider_id: str | None = None # Provider ID（用于失败记录）
endpoint_id: str | None = None # Endpoint ID（用于失败记录）
key_id: str | None = None # Key ID（用于失败记录）
mapped_model_result: str | None = None # 映射后的目标模型名（用于 Usage 记录）
sync_proxy_info: dict[str, Any] | None = None # 代理信息（用于 Usage 记录）
async def sync_request_func(
provider: Provider,
endpoint: ProviderEndpoint,
key: ProviderAPIKey,
candidate: ProviderCandidate,
) -> dict[str, Any]:
nonlocal provider_name, response_json, status_code, response_headers
nonlocal provider_request_headers, provider_request_body, mapped_model_result
nonlocal provider_api_format_for_error, client_api_format_for_error, needs_conversion_for_error
nonlocal sync_proxy_info
provider_name = str(provider.name)
provider_api_format = str(endpoint.api_format or api_format)
client_api_format = (
api_format.value if hasattr(api_format, "value") else str(api_format)
)
# 构建 Provider 请求（模型映射、格式转换、envelope 包装）
prep = await self._prepare_provider_request(
model=model,
provider=provider,
endpoint=endpoint,
key=key,
original_request_body=request_body_ref["body"],
client_api_format=client_api_format,
provider_api_format=provider_api_format,
candidate=candidate,
client_is_stream=False,
)
provider_api_format = prep.provider_api_format
needs_conversion = prep.needs_conversion
provider_api_format_for_error = provider_api_format
client_api_format_for_error = client_api_format
needs_conversion_for_error = needs_conversion
mapped_model = prep.mapped_model
if mapped_model:
mapped_model_result = mapped_model
request_body = prep.request_body
url_model = prep.url_model
envelope = prep.envelope
upstream_is_stream = prep.upstream_is_stream
auth_info = prep.auth_info
# 构建请求（上游始终使用 header 认证，不跟随客户端的 query 方式）
provider_payload, provider_hdrs = self._request_builder.build(
request_body,
original_headers,
endpoint,
key,
is_stream=upstream_is_stream,
extra_headers=prep.extra_headers if prep.extra_headers else None,
pre_computed_auth=auth_info.as_tuple() if auth_info else None,
)
if upstream_is_stream:
from src.core.api_format.headers import set_accept_if_absent
set_accept_if_absent(provider_hdrs)
provider_request_headers = provider_hdrs
provider_request_body = provider_payload
url = build_provider_url(
endpoint,
query_params=query_params,
path_params={"model": url_model},
is_stream=upstream_is_stream, # sync handler may still force upstream streaming
key=key,
decrypted_auth_config=auth_info.decrypted_auth_config if auth_info else None,
)
# 非流式：必须在 build_provider_url 调用后立即缓存（避免 contextvar 被后续调用覆盖）
selected_base_url_cached = envelope.capture_selected_base_url() if envelope else None
# 解析有效代理（Key 级别优先于 Provider 级别）
from src.services.proxy_node.resolver import (
get_proxy_label,
resolve_effective_proxy,
resolve_proxy_info,
)
_effective_proxy = resolve_effective_proxy(provider.proxy, getattr(key, "proxy", None))
sync_proxy_info = resolve_proxy_info(_effective_proxy)
_proxy_label = get_proxy_label(sync_proxy_info)
provider_type = str(getattr(provider, "provider_type", "") or "").lower()
logger.info(
f" [{self.request_id}] 发送{'上游流式(聚合)' if upstream_is_stream else '非流式'}请求: "
f"Provider={provider.name}, 模型={model} -> {mapped_model or '无映射'}, "
f"代理={_proxy_label}"
)
logger.debug(f" [{self.request_id}] 请求URL: {redact_url_for_log(url)}")
# 获取复用的 HTTP 客户端（支持代理配置，Key 级别优先于 Provider 级别）
# 注意：使用 get_proxy_client 复用连接池，不再每次创建新客户端
from src.clients.http_client import HTTPClientPool
from src.services.proxy_node.resolver import (
build_post_kwargs,
build_stream_kwargs,
resolve_delegate_config,
)
# 非流式请求使用 http_request_timeout 作为整体超时
# 优先使用 Provider 配置，否则使用全局配置
request_timeout = provider.request_timeout or config.http_request_timeout
delegate_cfg = resolve_delegate_config(_effective_proxy)
http_client = await HTTPClientPool.get_upstream_client(
delegate_cfg, proxy_config=_effective_proxy
)
# 注意：不使用 async with，因为复用的客户端不应该被关闭
# 超时通过 timeout 参数控制
resp: httpx.Response | None = None
if not upstream_is_stream:
try:
_pkw = build_post_kwargs(
delegate_cfg,
url=url,
headers=provider_hdrs,
payload=provider_payload,
timeout=request_timeout,
)
resp = await http_client.post(**_pkw)
except (httpx.ConnectError, httpx.ConnectTimeout, httpx.TimeoutException) as e:
if envelope:
envelope.on_connection_error(base_url=selected_base_url_cached, exc=e)
if selected_base_url_cached:
logger.warning(
f"[{envelope.name}] Connection error: {selected_base_url_cached} ({e})"
)
raise
else:
# Forced upstream streaming: aggregate SSE to a sync JSON response.
provider_parser = (
get_parser_for_format(provider_api_format) if provider_api_format else None
)
try:
_stream_args = build_stream_kwargs(
delegate_cfg,
url=url,
headers=provider_hdrs,
payload=provider_payload,
timeout=request_timeout,
)
async with http_client.stream(**_stream_args) as stream_resp:
resp = stream_resp
status_code = stream_resp.status_code
response_headers = dict(stream_resp.headers)
extract_proxy_timing(sync_proxy_info, response_headers)
if envelope:
envelope.on_http_status(
base_url=selected_base_url_cached,
status_code=status_code,
)
stream_resp.raise_for_status()
byte_iter = stream_resp.aiter_bytes()
if provider_type == "kiro" and envelope and envelope.force_stream_rewrite():
from src.services.provider.adapters.kiro.eventstream_rewriter import (
apply_kiro_stream_rewrite,
)
byte_iter = apply_kiro_stream_rewrite(byte_iter, model=str(model or ""))
internal_resp = await aggregate_upstream_stream_to_internal_response(
byte_iter,
provider_api_format=provider_api_format,
provider_name=str(provider.name),
model=str(model or ""),
request_id=str(self.request_id or ""),
envelope=envelope,
provider_parser=provider_parser,
)
registry = get_format_converter_registry()
tgt_norm = (
registry.get_normalizer(client_api_format)
if client_api_format
else None
)
if tgt_norm is None:
raise RuntimeError(f"未注册 Normalizer: {client_api_format}")
response_json = tgt_norm.response_from_internal(
internal_resp,
requested_model=model,
)
response_json = response_json if isinstance(response_json, dict) else {}
except (httpx.ConnectError, httpx.ConnectTimeout, httpx.TimeoutException) as e:
if envelope:
envelope.on_connection_error(base_url=selected_base_url_cached, exc=e)
if selected_base_url_cached:
logger.warning(
f"[{envelope.name}] Connection error: {selected_base_url_cached} ({e})"
)
raise
status_code = resp.status_code
response_headers = dict(resp.headers)
extract_proxy_timing(sync_proxy_info, response_headers)
if envelope:
envelope.on_http_status(base_url=selected_base_url_cached, status_code=status_code)
# Forced upstream streaming already built response_json via aggregator.
if upstream_is_stream:
return response_json if isinstance(response_json, dict) else {}
# 统一使用 HTTPStatusError，让 TaskService/error_classifier 负责分类（客户端错误/兼容性错误/限流等）
try:
resp.raise_for_status()
except httpx.HTTPStatusError as e:
error_body = ""
try:
error_body = resp.text[:4000] if resp.text else ""
except Exception:
error_body = ""
# 供 ErrorClassifier 优先读取
e.upstream_response = error_body # type: ignore[attr-defined]
raise
# 安全解析 JSON 响应，处理可能的编码错误
try:
response_json = resp.json()
except (UnicodeDecodeError, json.JSONDecodeError) as e:
# 获取原始响应内容用于调试（存入 upstream_response）
raw_content = ""
try:
raw_content = resp.text[:500] if resp.text else "(empty)"
except Exception:
try:
raw_content = repr(resp.content[:500]) if resp.content else "(empty)"
except Exception:
raw_content = "(unable to read)"
logger.error(f"[{self.request_id}] 无法解析响应 JSON: {e}, 原始内容: {raw_content}")
# 判断错误类型，生成友好的客户端错误消息（不暴露提供商信息）
if raw_content == "(empty)" or not raw_content.strip():
client_message = "上游服务返回了空响应"
elif raw_content.strip().startswith(("<", "<!doctype", "<!DOCTYPE")):
client_message = "上游服务返回了非预期的响应格式"
else:
client_message = "上游服务返回了无效的响应"
raise ProviderNotAvailableException(
client_message,
provider_name=str(provider.name),
upstream_status=resp.status_code,
upstream_response=raw_content,
)
if envelope:
response_json = envelope.unwrap_response(response_json)
envelope.postprocess_unwrapped_response(model=model, data=response_json)
# 检查响应体中的嵌套错误（HTTP 200 但响应体包含错误）
if isinstance(response_json, dict):
parser = get_parser_for_format(provider_api_format)
if parser.is_error_response(response_json):
parsed = parser.parse_response(response_json, 200)
logger.warning(
f" [{self.request_id}] 非流式检测到嵌套错误: "
f"Provider={provider.name}, "
f"error_type={parsed.error_type}, "
f"embedded_status={parsed.embedded_status_code}, "
f"message={parsed.error_message}"
)
raise EmbeddedErrorException(
provider_name=str(provider.name),
error_code=parsed.embedded_status_code,
error_message=parsed.error_message,
error_status=parsed.error_type,
)
# 跨格式：响应转换回 client_format（失败触发 failover）
if needs_conversion and isinstance(response_json, dict):
registry = get_format_converter_registry()
response_json = registry.convert_response(
response_json,
provider_api_format,
client_api_format,
requested_model=model, # 使用用户请求的原始模型名
)
return response_json if isinstance(response_json, dict) else {}
try:
# 解析能力需求
capability_requirements = self._resolve_capability_requirements(
model_name=model,
request_headers=original_headers,
request_body=original_request_body,
)
preferred_key_ids = await self._resolve_preferred_key_ids(
model_name=model,
request_body=original_request_body,
)
# 统一入口：总是通过 TaskService
from src.services.task import TaskService
from src.services.task.context import TaskMode
exec_result = await TaskService(self.db, self.redis).execute(
task_type="chat",
task_mode=TaskMode.SYNC,
api_format=api_format,
model_name=model,
user_api_key=self.api_key,
request_func=sync_request_func,
request_id=self.request_id,
is_stream=False,
capability_requirements=capability_requirements or None,
preferred_key_ids=preferred_key_ids or None,
request_body_ref=request_body_ref,
)
result = exec_result.response
actual_provider_name = exec_result.provider_name or "unknown"
attempt_id = exec_result.request_candidate_id
provider_id = exec_result.provider_id
endpoint_id = exec_result.endpoint_id
key_id = exec_result.key_id
provider_name = actual_provider_name
response_time_ms = self.elapsed_ms()
# 确保 response_json 不为 None
if response_json is None:
response_json = {}
# 规范化响应
response_json = self._normalize_response(response_json)
# 提取 usage
usage_info = self._extract_usage(response_json)
input_tokens = usage_info.get("input_tokens", 0)
output_tokens = usage_info.get("output_tokens", 0)
cache_creation_tokens = usage_info.get("cache_creation_input_tokens", 0)
cached_tokens = usage_info.get("cache_read_input_tokens", 0)
actual_request_body = provider_request_body or original_request_body
# 非流式成功时，返回给客户端的是提供商响应头（透传）
# JSONResponse 会自动设置 content-type，但我们记录实际返回的完整头
client_response_headers = filter_proxy_response_headers(response_headers)
client_response_headers["content-type"] = "application/json"
request_metadata = self._build_request_metadata() or {}
if sync_proxy_info:
request_metadata["proxy"] = sync_proxy_info
total_cost = await self.telemetry.record_success(
provider=provider_name,
model=model,
input_tokens=input_tokens,
output_tokens=output_tokens,
response_time_ms=response_time_ms,
status_code=status_code,
request_headers=original_headers,
request_body=actual_request_body,
response_headers=response_headers,
client_response_headers=client_response_headers,
response_body=response_json,
cache_creation_tokens=cache_creation_tokens,
cache_read_tokens=cached_tokens,
is_stream=False,
provider_request_headers=provider_request_headers,
api_format=api_format,
# 格式转换追踪
endpoint_api_format=provider_api_format_for_error or None,
has_format_conversion=is_format_converted(
provider_api_format_for_error, client_api_format_for_error
),
provider_id=provider_id,
provider_endpoint_id=endpoint_id,
provider_api_key_id=key_id,
# 模型映射信息
target_model=mapped_model_result,
request_metadata=request_metadata or None,
)
logger.debug(f"{self.FORMAT_ID} 非流式响应完成")
# 简洁的请求完成摘要
logger.info(
f"[OK] {self.request_id[:8]} | {model} | {provider_name or 'unknown'} | {response_time_ms}ms | "
f"in:{input_tokens or 0} out:{output_tokens or 0}"
)
# 透传提供商的响应头
return JSONResponse(
status_code=status_code,
content=response_json,
headers=client_response_headers,
)
except ThinkingSignatureException as e:
# Thinking 签名错误：TaskService 层已处理整流重试但仍失败
# 记录实际发送给 Provider 的请求体，便于排查问题根因
response_time_ms = self.elapsed_ms()
actual_request_body = provider_request_body or original_request_body
request_metadata = self._build_request_metadata() or {}
if sync_proxy_info:
request_metadata["proxy"] = sync_proxy_info
await self.telemetry.record_failure(
provider=provider_name or "unknown",
model=model,
response_time_ms=response_time_ms,
status_code=e.status_code or 400,
request_headers=original_headers,
request_body=actual_request_body,
error_message=str(e),
is_stream=False,
request_metadata=request_metadata or None,
)
client_format = (client_api_format_for_error or "").upper()
provider_format = (provider_api_format_for_error or client_format).upper()
payload = _build_error_json_payload(
e, client_format, provider_format, needs_conversion=needs_conversion_for_error
)
return JSONResponse(
status_code=_get_error_status_code(e),
content=payload,
)
except UpstreamClientException as e:
response_time_ms = self.elapsed_ms()
actual_request_body = provider_request_body or original_request_body
request_metadata = self._build_request_metadata() or {}
if sync_proxy_info:
request_metadata["proxy"] = sync_proxy_info
await self.telemetry.record_failure(
provider=provider_name or "unknown",
model=model,
response_time_ms=response_time_ms,
status_code=_get_error_status_code(e),
request_headers=original_headers,
request_body=actual_request_body,
error_message=str(e),
is_stream=False,
api_format=api_format,
provider_request_headers=provider_request_headers,
response_headers=response_headers,
client_response_headers={"content-type": "application/json"},
# 格式转换追踪
endpoint_api_format=provider_api_format_for_error or None,
has_format_conversion=is_format_converted(
provider_api_format_for_error, client_api_format_for_error
),
target_model=mapped_model_result,
request_metadata=request_metadata,
)
client_format = (client_api_format_for_error or "").upper()
provider_format = (provider_api_format_for_error or client_format).upper()
payload = _build_error_json_payload(
e, client_format, provider_format, needs_conversion=needs_conversion_for_error
)
return JSONResponse(
status_code=_get_error_status_code(e),
content=payload,
)
except Exception as e:
response_time_ms = self.elapsed_ms()
status_code = 503
if isinstance(e, ProviderAuthException):
status_code = 503
elif isinstance(e, ProviderRateLimitException):
status_code = 429
elif isinstance(e, ProviderTimeoutException):
status_code = 504
actual_request_body = provider_request_body or original_request_body
# 尝试从异常中提取响应头
error_response_headers: dict[str, str] = {}
if isinstance(e, ProviderRateLimitException) and e.response_headers:
error_response_headers = e.response_headers
elif isinstance(e, httpx.HTTPStatusError) and hasattr(e, "response"):
error_response_headers = dict(e.response.headers)
request_metadata = self._build_request_metadata() or {}
if sync_proxy_info:
request_metadata["proxy"] = sync_proxy_info
await self.telemetry.record_failure(
provider=provider_name or "unknown",
model=model,
response_time_ms=response_time_ms,
status_code=status_code,
error_message=extract_client_error_message(e),
request_headers=original_headers,
request_body=actual_request_body,
is_stream=False,
api_format=api_format,
provider_request_headers=provider_request_headers,
response_headers=error_response_headers,
# 非流式失败返回给客户端的是 JSON 错误响应
client_response_headers={"content-type": "application/json"},
# 格式转换追踪
endpoint_api_format=provider_api_format_for_error or None,
has_format_conversion=is_format_converted(
provider_api_format_for_error, client_api_format_for_error
),
# 模型映射信息
target_model=mapped_model_result,
request_metadata=request_metadata,
)
raise
async def _extract_error_text(self, e: httpx.HTTPStatusError) -> str:
"""从 HTTP 错误中提取错误文本"""
try:
if hasattr(e.response, "is_stream_consumed") and not e.response.is_stream_consumed:
error_bytes = await e.response.aread()
return error_bytes.decode("utf-8", errors="replace")
else:
return e.response.text if hasattr(e.response, "_content") else "Unable to read"
except Exception as decode_error:
return f"Unable to read error: {decode_error}"

729
src/api/handlers/base/chat_sync_executor.py Normal file
View File

@@ -0,0 +1,729 @@
"""
ChatSyncExecutor - 非流式请求执行器
从 ChatHandlerBase.process_sync() 提取的独立类，负责：
- 非流式请求的完整执行流程（请求构建、发送、响应解析）
- 通过 SyncRequestContext 管理可变状态（替代原来的 nonlocal 变量）
- 异常处理与 telemetry 记录
- 流式失败记录（_record_stream_failure）
- HTTP 错误文本提取（_extract_error_text）
"""
from __future__ import annotations
import json
from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Any
import httpx
from fastapi.responses import JSONResponse
from src.api.handlers.base.chat_error_utils import (
_build_error_json_payload,
_get_error_status_code,
)
from src.api.handlers.base.parsers import get_parser_for_format
from src.api.handlers.base.stream_context import (
StreamContext,
extract_proxy_timing,
is_format_converted,
)
from src.api.handlers.base.utils import (
filter_proxy_response_headers,
get_format_converter_registry,
)
from src.core.error_utils import extract_client_error_message
from src.core.exceptions import (
EmbeddedErrorException,
ProviderAuthException,
ProviderNotAvailableException,
ProviderRateLimitException,
ProviderTimeoutException,
ThinkingSignatureException,
UpstreamClientException,
)
from src.core.logger import logger
if TYPE_CHECKING:
from fastapi import Request
from src.api.handlers.base.chat_handler_base import ChatHandlerBase
from src.models.database import Provider, ProviderAPIKey, ProviderEndpoint
from src.services.cache.aware_scheduler import ProviderCandidate
@dataclass
class SyncRequestContext:
"""同步请求的可变状态容器，替代原来的 nonlocal 变量"""
provider_name: str | None = None
response_json: dict[str, Any] | None = None
status_code: int = 200
response_headers: dict[str, str] = field(default_factory=dict)
provider_request_headers: dict[str, str] = field(default_factory=dict)
provider_request_body: dict[str, Any] | None = None
provider_api_format_for_error: str | None = None
client_api_format_for_error: str | None = None
needs_conversion_for_error: bool = False
provider_id: str | None = None
endpoint_id: str | None = None
key_id: str | None = None
mapped_model_result: str | None = None
sync_proxy_info: dict[str, Any] | None = None
class ChatSyncExecutor:
"""非流式请求执行器，从 ChatHandlerBase 提取"""
def __init__(self, handler: ChatHandlerBase) -> None:
self._handler = handler
self._ctx = SyncRequestContext()
async def execute(
self,
request: Any,
http_request: Request,
original_headers: dict[str, Any],
original_request_body: dict[str, Any],
query_params: dict[str, str] | None = None,
) -> JSONResponse:
"""处理非流式响应（原 process_sync 的完整逻辑）"""
handler = self._handler
logger.debug(f"开始非流式响应处理 ({handler.FORMAT_ID})")
# 转换请求格式
converted_request = await handler._convert_request(request)
model = getattr(converted_request, "model", original_request_body.get("model", "unknown"))
api_format = handler.allowed_api_formats[0]
# 提前创建 pending 记录，让前端可以立即看到"处理中"
handler._create_pending_usage(
model=model,
is_stream=False,
request_type="chat",
api_format=handler.FORMAT_ID,
request_headers=original_headers,
request_body=original_request_body,
)
# 可变请求体容器：允许 TaskService 在遇到 Thinking 签名错误时整流请求体后重试
# 结构: {"body": 实际请求体, "_rectified": 是否已整流, "_rectified_this_turn": 本轮是否整流}
request_body_ref: dict[str, Any] = {"body": original_request_body}
# 捕获的上下文变量
ctx = self._ctx
async def sync_request_func(
provider: Provider,
endpoint: ProviderEndpoint,
key: ProviderAPIKey,
candidate: ProviderCandidate,
) -> dict[str, Any]:
return await self._sync_request_func(
provider,
endpoint,
key,
candidate,
model=model,
api_format=api_format,
original_headers=original_headers,
request_body_ref=request_body_ref,
query_params=query_params,
)
try:
# 解析能力需求
capability_requirements = handler._resolve_capability_requirements(
model_name=model,
request_headers=original_headers,
request_body=original_request_body,
)
preferred_key_ids = await handler._resolve_preferred_key_ids(
model_name=model,
request_body=original_request_body,
)
# 统一入口：总是通过 TaskService
from src.services.task import TaskService
from src.services.task.context import TaskMode
exec_result = await TaskService(handler.db, handler.redis).execute(
task_type="chat",
task_mode=TaskMode.SYNC,
api_format=api_format,
model_name=model,
user_api_key=handler.api_key,
request_func=sync_request_func,
request_id=handler.request_id,
is_stream=False,
capability_requirements=capability_requirements or None,
preferred_key_ids=preferred_key_ids or None,
request_body_ref=request_body_ref,
)
actual_provider_name = exec_result.provider_name or "unknown"
ctx.provider_id = exec_result.provider_id
ctx.endpoint_id = exec_result.endpoint_id
ctx.key_id = exec_result.key_id
ctx.provider_name = actual_provider_name
response_time_ms = handler.elapsed_ms()
# 确保 response_json 不为 None
if ctx.response_json is None:
ctx.response_json = {}
# 规范化响应
ctx.response_json = handler._normalize_response(ctx.response_json)
# 提取 usage
usage_info = handler._extract_usage(ctx.response_json)
input_tokens = usage_info.get("input_tokens", 0)
output_tokens = usage_info.get("output_tokens", 0)
cache_creation_tokens = usage_info.get("cache_creation_input_tokens", 0)
cached_tokens = usage_info.get("cache_read_input_tokens", 0)
actual_request_body = ctx.provider_request_body or original_request_body
# 非流式成功时，返回给客户端的是提供商响应头（透传）
# JSONResponse 会自动设置 content-type，但我们记录实际返回的完整头
client_response_headers = filter_proxy_response_headers(ctx.response_headers)
client_response_headers["content-type"] = "application/json"
request_metadata = handler._build_request_metadata() or {}
if ctx.sync_proxy_info:
request_metadata["proxy"] = ctx.sync_proxy_info
total_cost = await handler.telemetry.record_success( # noqa: F841
provider=ctx.provider_name,
model=model,
input_tokens=input_tokens,
output_tokens=output_tokens,
response_time_ms=response_time_ms,
status_code=ctx.status_code,
request_headers=original_headers,
request_body=actual_request_body,
response_headers=ctx.response_headers,
client_response_headers=client_response_headers,
response_body=ctx.response_json,
cache_creation_tokens=cache_creation_tokens,
cache_read_tokens=cached_tokens,
is_stream=False,
provider_request_headers=ctx.provider_request_headers,
api_format=api_format,
# 格式转换追踪
endpoint_api_format=ctx.provider_api_format_for_error or None,
has_format_conversion=is_format_converted(
ctx.provider_api_format_for_error, ctx.client_api_format_for_error
),
provider_id=ctx.provider_id,
provider_endpoint_id=ctx.endpoint_id,
provider_api_key_id=ctx.key_id,
# 模型映射信息
target_model=ctx.mapped_model_result,
request_metadata=request_metadata or None,
)
logger.debug(f"{handler.FORMAT_ID} 非流式响应完成")
# 简洁的请求完成摘要
logger.info(
f"[OK] {handler.request_id[:8]} | {model} | "
f"{ctx.provider_name or 'unknown'} | {response_time_ms}ms | "
f"in:{input_tokens or 0} out:{output_tokens or 0}"
)
# 透传提供商的响应头
return JSONResponse(
status_code=ctx.status_code,
content=ctx.response_json,
headers=client_response_headers,
)
except ThinkingSignatureException as e:
# Thinking 签名错误：TaskService 层已处理整流重试但仍失败
# 记录实际发送给 Provider 的请求体，便于排查问题根因
response_time_ms = handler.elapsed_ms()
actual_request_body = ctx.provider_request_body or original_request_body
request_metadata = handler._build_request_metadata() or {}
if ctx.sync_proxy_info:
request_metadata["proxy"] = ctx.sync_proxy_info
await handler.telemetry.record_failure(
provider=ctx.provider_name or "unknown",
model=model,
response_time_ms=response_time_ms,
status_code=e.status_code or 400,
request_headers=original_headers,
request_body=actual_request_body,
error_message=str(e),
is_stream=False,
request_metadata=request_metadata or None,
)
client_format = (ctx.client_api_format_for_error or "").upper()
provider_format = (ctx.provider_api_format_for_error or client_format).upper()
payload = _build_error_json_payload(
e,
client_format,
provider_format,
needs_conversion=ctx.needs_conversion_for_error,
)
return JSONResponse(
status_code=_get_error_status_code(e),
content=payload,
)
except UpstreamClientException as e:
response_time_ms = handler.elapsed_ms()
actual_request_body = ctx.provider_request_body or original_request_body
request_metadata = handler._build_request_metadata() or {}
if ctx.sync_proxy_info:
request_metadata["proxy"] = ctx.sync_proxy_info
await handler.telemetry.record_failure(
provider=ctx.provider_name or "unknown",
model=model,
response_time_ms=response_time_ms,
status_code=_get_error_status_code(e),
request_headers=original_headers,
request_body=actual_request_body,
error_message=str(e),
is_stream=False,
api_format=api_format,
provider_request_headers=ctx.provider_request_headers,
response_headers=ctx.response_headers,
client_response_headers={"content-type": "application/json"},
# 格式转换追踪
endpoint_api_format=ctx.provider_api_format_for_error or None,
has_format_conversion=is_format_converted(
ctx.provider_api_format_for_error, ctx.client_api_format_for_error
),
target_model=ctx.mapped_model_result,
request_metadata=request_metadata,
)
client_format = (ctx.client_api_format_for_error or "").upper()
provider_format = (ctx.provider_api_format_for_error or client_format).upper()
payload = _build_error_json_payload(
e,
client_format,
provider_format,
needs_conversion=ctx.needs_conversion_for_error,
)
return JSONResponse(
status_code=_get_error_status_code(e),
content=payload,
)
except Exception as e:
response_time_ms = handler.elapsed_ms()
status_code = 503
if isinstance(e, ProviderAuthException):
status_code = 503
elif isinstance(e, ProviderRateLimitException):
status_code = 429
elif isinstance(e, ProviderTimeoutException):
status_code = 504
actual_request_body = ctx.provider_request_body or original_request_body
# 尝试从异常中提取响应头
error_response_headers: dict[str, str] = {}
if isinstance(e, ProviderRateLimitException) and e.response_headers:
error_response_headers = e.response_headers
elif isinstance(e, httpx.HTTPStatusError) and hasattr(e, "response"):
error_response_headers = dict(e.response.headers)
request_metadata = handler._build_request_metadata() or {}
if ctx.sync_proxy_info:
request_metadata["proxy"] = ctx.sync_proxy_info
await handler.telemetry.record_failure(
provider=ctx.provider_name or "unknown",
model=model,
response_time_ms=response_time_ms,
status_code=status_code,
error_message=extract_client_error_message(e),
request_headers=original_headers,
request_body=actual_request_body,
is_stream=False,
api_format=api_format,
provider_request_headers=ctx.provider_request_headers,
response_headers=error_response_headers,
# 非流式失败返回给客户端的是 JSON 错误响应
client_response_headers={"content-type": "application/json"},
# 格式转换追踪
endpoint_api_format=ctx.provider_api_format_for_error or None,
has_format_conversion=is_format_converted(
ctx.provider_api_format_for_error, ctx.client_api_format_for_error
),
# 模型映射信息
target_model=ctx.mapped_model_result,
request_metadata=request_metadata,
)
raise
async def _sync_request_func(
self,
provider: Provider,
endpoint: ProviderEndpoint,
key: ProviderAPIKey,
candidate: ProviderCandidate,
*,
model: str,
api_format: Any,
original_headers: dict[str, Any],
request_body_ref: dict[str, Any],
query_params: dict[str, str] | None = None,
) -> dict[str, Any]:
"""单次同步请求（原 sync_request_func 内嵌函数）"""
handler = self._handler
ctx = self._ctx
ctx.provider_name = str(provider.name)
provider_api_format = str(endpoint.api_format or api_format)
client_api_format = api_format.value if hasattr(api_format, "value") else str(api_format)
# 构建 Provider 请求（模型映射、格式转换、envelope 包装）
prep = await handler._prepare_provider_request(
model=model,
provider=provider,
endpoint=endpoint,
key=key,
original_request_body=request_body_ref["body"],
client_api_format=client_api_format,
provider_api_format=provider_api_format,
candidate=candidate,
client_is_stream=False,
)
provider_api_format = prep.provider_api_format
needs_conversion = prep.needs_conversion
ctx.provider_api_format_for_error = provider_api_format
ctx.client_api_format_for_error = client_api_format
ctx.needs_conversion_for_error = needs_conversion
mapped_model = prep.mapped_model
if mapped_model:
ctx.mapped_model_result = mapped_model
request_body = prep.request_body
url_model = prep.url_model
envelope = prep.envelope
upstream_is_stream = prep.upstream_is_stream
auth_info = prep.auth_info
# 构建请求（上游始终使用 header 认证，不跟随客户端的 query 方式）
provider_payload, provider_hdrs = handler._request_builder.build(
request_body,
original_headers,
endpoint,
key,
is_stream=upstream_is_stream,
extra_headers=prep.extra_headers if prep.extra_headers else None,
pre_computed_auth=auth_info.as_tuple() if auth_info else None,
)
if upstream_is_stream:
from src.core.api_format.headers import set_accept_if_absent
set_accept_if_absent(provider_hdrs)
ctx.provider_request_headers = provider_hdrs
ctx.provider_request_body = provider_payload
from src.services.provider.transport import (
build_provider_url,
redact_url_for_log,
)
url = build_provider_url(
endpoint,
query_params=query_params,
path_params={"model": url_model},
is_stream=upstream_is_stream, # sync handler may still force upstream streaming
key=key,
decrypted_auth_config=auth_info.decrypted_auth_config if auth_info else None,
)
# 非流式：必须在 build_provider_url 调用后立即缓存（避免 contextvar 被后续调用覆盖）
selected_base_url_cached = envelope.capture_selected_base_url() if envelope else None
# 解析有效代理（Key 级别优先于 Provider 级别）
from src.services.proxy_node.resolver import (
get_proxy_label,
resolve_effective_proxy,
resolve_proxy_info,
)
_effective_proxy = resolve_effective_proxy(provider.proxy, getattr(key, "proxy", None))
ctx.sync_proxy_info = resolve_proxy_info(_effective_proxy)
_proxy_label = get_proxy_label(ctx.sync_proxy_info)
provider_type = str(getattr(provider, "provider_type", "") or "").lower()
logger.info(
f" [{handler.request_id}] "
f"发送{'上游流式(聚合)' if upstream_is_stream else '非流式'}请求: "
f"Provider={provider.name}, 模型={model} -> {mapped_model or '无映射'}, "
f"代理={_proxy_label}"
)
logger.debug(f" [{handler.request_id}] 请求URL: {redact_url_for_log(url)}")
# 获取复用的 HTTP 客户端（支持代理配置，Key 级别优先于 Provider 级别）
# 注意：使用 get_proxy_client 复用连接池，不再每次创建新客户端
from src.clients.http_client import HTTPClientPool
from src.config.settings import config
from src.services.proxy_node.resolver import (
build_post_kwargs,
build_stream_kwargs,
resolve_delegate_config,
)
# 非流式请求使用 http_request_timeout 作为整体超时
# 优先使用 Provider 配置，否则使用全局配置
request_timeout = provider.request_timeout or config.http_request_timeout
delegate_cfg = resolve_delegate_config(_effective_proxy)
http_client = await HTTPClientPool.get_upstream_client(
delegate_cfg, proxy_config=_effective_proxy
)
# 注意：不使用 async with，因为复用的客户端不应该被关闭
# 超时通过 timeout 参数控制
resp: httpx.Response | None = None
if not upstream_is_stream:
try:
_pkw = build_post_kwargs(
delegate_cfg,
url=url,
headers=provider_hdrs,
payload=provider_payload,
timeout=request_timeout,
)
resp = await http_client.post(**_pkw)
except (httpx.ConnectError, httpx.ConnectTimeout, httpx.TimeoutException) as e:
if envelope:
envelope.on_connection_error(base_url=selected_base_url_cached, exc=e)
if selected_base_url_cached:
logger.warning(
f"[{envelope.name}] Connection error: "
f"{selected_base_url_cached} ({e})"
)
raise
else:
# Forced upstream streaming: aggregate SSE to a sync JSON response.
provider_parser = (
get_parser_for_format(provider_api_format) if provider_api_format else None
)
try:
_stream_args = build_stream_kwargs(
delegate_cfg,
url=url,
headers=provider_hdrs,
payload=provider_payload,
timeout=request_timeout,
)
async with http_client.stream(**_stream_args) as stream_resp:
resp = stream_resp
ctx.status_code = stream_resp.status_code
ctx.response_headers = dict(stream_resp.headers)
extract_proxy_timing(ctx.sync_proxy_info, ctx.response_headers)
if envelope:
envelope.on_http_status(
base_url=selected_base_url_cached,
status_code=ctx.status_code,
)
stream_resp.raise_for_status()
byte_iter = stream_resp.aiter_bytes()
if provider_type == "kiro" and envelope and envelope.force_stream_rewrite():
from src.services.provider.adapters.kiro.eventstream_rewriter import (
apply_kiro_stream_rewrite,
)
byte_iter = apply_kiro_stream_rewrite(byte_iter, model=str(model or ""))
from src.api.handlers.base.upstream_stream_bridge import (
aggregate_upstream_stream_to_internal_response,
)
internal_resp = await aggregate_upstream_stream_to_internal_response(
byte_iter,
provider_api_format=provider_api_format,
provider_name=str(provider.name),
model=str(model or ""),
request_id=str(handler.request_id or ""),
envelope=envelope,
provider_parser=provider_parser,
)
registry = get_format_converter_registry()
tgt_norm = (
registry.get_normalizer(client_api_format) if client_api_format else None
)
if tgt_norm is None:
raise RuntimeError(f"未注册 Normalizer: {client_api_format}")
ctx.response_json = tgt_norm.response_from_internal(
internal_resp,
requested_model=model,
)
ctx.response_json = (
ctx.response_json if isinstance(ctx.response_json, dict) else {}
)
except (httpx.ConnectError, httpx.ConnectTimeout, httpx.TimeoutException) as e:
if envelope:
envelope.on_connection_error(base_url=selected_base_url_cached, exc=e)
if selected_base_url_cached:
logger.warning(
f"[{envelope.name}] Connection error: "
f"{selected_base_url_cached} ({e})"
)
raise
ctx.status_code = resp.status_code
ctx.response_headers = dict(resp.headers)
extract_proxy_timing(ctx.sync_proxy_info, ctx.response_headers)
if envelope:
envelope.on_http_status(base_url=selected_base_url_cached, status_code=ctx.status_code)
# Forced upstream streaming already built response_json via aggregator.
if upstream_is_stream:
return ctx.response_json if isinstance(ctx.response_json, dict) else {}
# 统一使用 HTTPStatusError，让 TaskService/error_classifier 负责分类
# （客户端错误/兼容性错误/限流等）
try:
resp.raise_for_status()
except httpx.HTTPStatusError as e:
error_body = ""
try:
error_body = resp.text[:4000] if resp.text else ""
except Exception:
error_body = ""
# 供 ErrorClassifier 优先读取
e.upstream_response = error_body # type: ignore[attr-defined]
raise
# 安全解析 JSON 响应，处理可能的编码错误
try:
ctx.response_json = resp.json()
except (UnicodeDecodeError, json.JSONDecodeError) as e:
# 获取原始响应内容用于调试（存入 upstream_response）
raw_content = ""
try:
raw_content = resp.text[:500] if resp.text else "(empty)"
except Exception:
try:
raw_content = repr(resp.content[:500]) if resp.content else "(empty)"
except Exception:
raw_content = "(unable to read)"
logger.error(f"[{handler.request_id}] 无法解析响应 JSON: {e}, 原始内容: {raw_content}")
# 判断错误类型，生成友好的客户端错误消息（不暴露提供商信息）
if raw_content == "(empty)" or not raw_content.strip():
client_message = "上游服务返回了空响应"
elif raw_content.strip().startswith(("<", "<!doctype", "<!DOCTYPE")):
client_message = "上游服务返回了非预期的响应格式"
else:
client_message = "上游服务返回了无效的响应"
raise ProviderNotAvailableException(
client_message,
provider_name=str(provider.name),
upstream_status=resp.status_code,
upstream_response=raw_content,
)
if envelope:
ctx.response_json = envelope.unwrap_response(ctx.response_json)
envelope.postprocess_unwrapped_response(model=model, data=ctx.response_json)
# 检查响应体中的嵌套错误（HTTP 200 但响应体包含错误）
if isinstance(ctx.response_json, dict):
parser = get_parser_for_format(provider_api_format)
if parser.is_error_response(ctx.response_json):
parsed = parser.parse_response(ctx.response_json, 200)
logger.warning(
f" [{handler.request_id}] 非流式检测到嵌套错误: "
f"Provider={provider.name}, "
f"error_type={parsed.error_type}, "
f"embedded_status={parsed.embedded_status_code}, "
f"message={parsed.error_message}"
)
raise EmbeddedErrorException(
provider_name=str(provider.name),
error_code=parsed.embedded_status_code,
error_message=parsed.error_message,
error_status=parsed.error_type,
)
# 跨格式：响应转换回 client_format（失败触发 failover）
if needs_conversion and isinstance(ctx.response_json, dict):
registry = get_format_converter_registry()
ctx.response_json = registry.convert_response(
ctx.response_json,
provider_api_format,
client_api_format,
requested_model=model, # 使用用户请求的原始模型名
)
return ctx.response_json if isinstance(ctx.response_json, dict) else {}
async def _record_stream_failure(
self,
ctx: StreamContext,
error: Exception,
original_headers: dict[str, str],
original_request_body: dict[str, Any],
) -> None:
"""记录流式请求失败"""
handler = self._handler
response_time_ms = handler.elapsed_ms()
status_code = 503
if isinstance(error, ThinkingSignatureException):
status_code = 400
elif isinstance(error, UpstreamClientException):
status_code = _get_error_status_code(error)
elif isinstance(error, ProviderAuthException):
status_code = 503
elif isinstance(error, ProviderRateLimitException):
status_code = 429
elif isinstance(error, ProviderTimeoutException):
status_code = 504
actual_request_body = ctx.provider_request_body or original_request_body
# 失败时返回给客户端的是 JSON 错误响应
client_response_headers = {"content-type": "application/json"}
stream_fail_metadata: dict[str, Any] | None = None
if ctx.proxy_info:
stream_fail_metadata = {"proxy": ctx.proxy_info}
await handler.telemetry.record_failure(
provider=ctx.provider_name or "unknown",
model=ctx.model,
response_time_ms=response_time_ms,
status_code=status_code,
error_message=extract_client_error_message(error),
request_headers=original_headers,
request_body=actual_request_body,
is_stream=True,
api_format=ctx.api_format,
provider_request_headers=ctx.provider_request_headers,
response_headers=ctx.response_headers,
client_response_headers=client_response_headers,
# 格式转换追踪
endpoint_api_format=ctx.provider_api_format or None,
has_format_conversion=ctx.has_format_conversion,
target_model=ctx.mapped_model,
request_metadata=stream_fail_metadata,
)
async def _extract_error_text(self, e: httpx.HTTPStatusError) -> str:
"""从 HTTP 错误中提取错误文本"""
try:
if hasattr(e.response, "is_stream_consumed") and not e.response.is_stream_consumed:
error_bytes = await e.response.aread()
return error_bytes.decode("utf-8", errors="replace")
else:
return e.response.text if hasattr(e.response, "_content") else "Unable to read"
except Exception as decode_error:
return f"Unable to read error: {decode_error}"

7
src/api/handlers/base/cli_event_mixin.py
View File

@@ -3,7 +3,7 @@
from __future__ import annotations
import json
from typing import Any
from typing import TYPE_CHECKING, Any
from src.api.handlers.base.parsers import get_parser_for_format
from src.api.handlers.base.stream_context import StreamContext
@@ -18,12 +18,15 @@ from .cli_sse_helpers import (
_parse_sse_event_data_line,
)
if TYPE_CHECKING:
from src.api.handlers.base.cli_protocol import CliHandlerProtocol
class CliEventMixin:
"""SSE 事件处理和格式转换相关方法的 Mixin"""
def _handle_sse_event(
self,
self: CliHandlerProtocol,
ctx: StreamContext,
event_name: str | None,
data_str: str,

7
src/api/handlers/base/cli_monitor_mixin.py
View File

@@ -5,7 +5,7 @@ from __future__ import annotations
import asyncio
import time
from collections.abc import AsyncGenerator
from typing import Any
from typing import TYPE_CHECKING, Any
import httpx
from fastapi import Request
@@ -25,6 +25,9 @@ from src.database import get_db
from src.models.database import User
from src.services.provider.behavior import get_provider_behavior
if TYPE_CHECKING:
from src.api.handlers.base.cli_protocol import CliHandlerProtocol
class CliMonitorMixin:
"""监控和统计相关方法的 Mixin"""
@@ -157,7 +160,7 @@ class CliMonitorMixin:
raise
async def _record_stream_stats(
self,
self: CliHandlerProtocol,
ctx: StreamContext,
original_headers: dict[str, str],
original_request_body: dict[str, Any],

3
src/api/handlers/base/cli_prefetch_mixin.py
View File

@@ -29,6 +29,7 @@ from src.utils.sse_parser import SSEEventParser
from src.utils.timeout import read_first_chunk_with_ttfb_timeout
if TYPE_CHECKING:
from src.api.handlers.base.cli_protocol import CliHandlerProtocol
from src.models.database import Provider, ProviderEndpoint
@@ -136,7 +137,7 @@ class CliPrefetchMixin:
)
async def _prefetch_and_check_embedded_error(
self,
self: CliHandlerProtocol,
byte_iterator: Any,
provider: "Provider",
endpoint: "ProviderEndpoint",

252
src/api/handlers/base/cli_protocol.py Normal file
View File

@@ -0,0 +1,252 @@
"""
CLI Handler Mixin Protocol -- Mixin 隐式依赖的编译时契约
各 Mixin (CliStreamMixin, CliSyncMixin, CliRequestMixin, CliMonitorMixin,
CliPrefetchMixin, CliEventMixin) 通过 duck typing 访问宿主类的属性和方法。
本模块将这些隐式依赖显式声明为 Protocol，使 mypy/pyright 能在编辑期捕获
缺失属性或类型不匹配的错误。
渐进式采用：仅在各 Mixin 的公开方法签名中标注 `self: CliHandlerProtocol`，
不修改方法体或私有 helper。
"""
from __future__ import annotations
from typing import (
TYPE_CHECKING,
Any,
Protocol,
runtime_checkable,
)
if TYPE_CHECKING:
from redis import Redis
from sqlalchemy.orm import Session
from src.api.handlers.base.base_handler import MessageTelemetry
from src.api.handlers.base.request_builder import RequestBuilder
from src.api.handlers.base.response_parser import ResponseParser
from src.api.handlers.base.stream_context import StreamContext
from src.models.database import ApiKey, User
@runtime_checkable
class CliHandlerProtocol(Protocol):
"""CLI Handler Mixin 宿主需要满足的属性/方法契约。
声明范围仅覆盖 Mixin 实际引用的 self.xxx，不要求宿主实现全部
BaseMessageHandler 接口。
"""
# ------------------------------------------------------------------
# 实例属性 -- 来自 BaseMessageHandler.__init__
# ------------------------------------------------------------------
db: Session
user: User
api_key: ApiKey
request_id: str
client_ip: str
user_agent: str
start_time: float
allowed_api_formats: list[str]
redis: Redis # type: ignore[type-arg]
telemetry: MessageTelemetry
perf_metrics: dict[str, Any] | None
# ------------------------------------------------------------------
# 类属性 -- 来自 CliMessageHandlerBase
# ------------------------------------------------------------------
FORMAT_ID: str
DATA_TIMEOUT: int
EMPTY_CHUNK_THRESHOLD: int
# ------------------------------------------------------------------
# 属性/方法 -- 来自 CliMessageHandlerBase / BaseMessageHandler
# ------------------------------------------------------------------
@property
def parser(self) -> ResponseParser: ...
_request_builder: RequestBuilder
# ------------------------------------------------------------------
# 方法 -- 来自 BaseMessageHandler (被多个 Mixin 引用)
# ------------------------------------------------------------------
def _create_pending_usage(
self,
model: str,
is_stream: bool,
request_type: str = ...,
api_format: str | None = ...,
request_headers: dict[str, Any] | None = ...,
request_body: dict[str, Any] | None = ...,
) -> None: ...
def _build_request_metadata(
self,
http_request: Any | None = ...,
) -> dict[str, Any] | None: ...
def _resolve_capability_requirements(
self,
model_name: str,
request_headers: dict[str, str] | None = ...,
request_body: dict[str, Any] | None = ...,
) -> dict[str, bool]: ...
async def _resolve_preferred_key_ids(
self,
model_name: str,
request_body: dict[str, Any] | None = ...,
) -> list[str] | None: ...
def _update_usage_to_streaming(
self,
request_id: str | None = ...,
) -> None: ...
def _update_usage_to_streaming_with_ctx(
self,
ctx: StreamContext,
) -> None: ...
def _log_request_error(
self,
message: str,
error: Exception,
) -> None: ...
# ------------------------------------------------------------------
# 方法 -- 来自 CliRequestMixin (被 CliStreamMixin / CliSyncMixin 引用)
# ------------------------------------------------------------------
def extract_model_from_request(
self,
request_body: dict[str, Any],
path_params: dict[str, Any] | None = ...,
) -> str: ...
async def _get_mapped_model(
self,
source_model: str,
provider_id: str,
) -> str | None: ...
def apply_mapped_model(
self,
request_body: dict[str, Any],
mapped_model: str,
) -> dict[str, Any]: ...
def prepare_provider_request_body(
self,
request_body: dict[str, Any],
) -> dict[str, Any]: ...
def finalize_provider_request(
self,
request_body: dict[str, Any],
*,
mapped_model: str | None,
provider_api_format: str | None,
) -> dict[str, Any]: ...
def get_model_for_url(
self,
request_body: dict[str, Any],
mapped_model: str | None,
) -> str | None: ...
def _convert_request_for_cross_format(
self,
request_body: dict[str, Any],
client_api_format: str,
provider_api_format: str,
mapped_model: str | None,
fallback_model: str,
is_stream: bool,
*,
target_variant: str | None = ...,
) -> tuple[dict[str, Any], str]: ...
def _extract_response_metadata(
self,
response: dict[str, Any],
) -> dict[str, Any]: ...
# ------------------------------------------------------------------
# 方法 -- 来自 CliEventMixin (被 CliStreamMixin / CliPrefetchMixin 引用)
# ------------------------------------------------------------------
def _handle_sse_event(
self,
ctx: StreamContext,
event_name: str | None,
data_str: str,
record_chunk: bool = ...,
) -> None: ...
def _mark_first_output(
self,
ctx: StreamContext,
state: dict[str, bool],
) -> None: ...
def _convert_sse_line(
self,
ctx: StreamContext,
line: str,
events: list[Any],
) -> tuple[list[str], list[dict[str, Any]]]: ...
def _record_converted_chunks(
self,
ctx: StreamContext,
converted_events: list[dict[str, Any]],
) -> None: ...
def _finalize_stream_metadata(
self,
ctx: StreamContext,
) -> None: ...
# ------------------------------------------------------------------
# 方法 -- 来自 CliPrefetchMixin (被 CliStreamMixin 引用)
# ------------------------------------------------------------------
def _flush_remaining_sse_data(
self,
ctx: StreamContext,
buffer: bytes,
decoder: Any,
sse_parser: Any,
*,
record_chunk: bool = ...,
) -> None: ...
def _estimate_tokens_for_incomplete_stream(
self,
ctx: StreamContext,
request_body: dict[str, Any],
) -> None: ...
# ------------------------------------------------------------------
# 方法 -- 来自 CliMonitorMixin (被 CliStreamMixin 引用)
# ------------------------------------------------------------------
async def _create_monitored_stream(
self,
ctx: StreamContext,
stream_generator: Any,
http_request: Any | None = ...,
) -> Any: ...
async def _record_stream_stats(
self,
ctx: StreamContext,
original_headers: dict[str, str],
original_request_body: dict[str, Any],
) -> None: ...
async def _record_stream_failure(
self,
ctx: StreamContext,
error: Exception,
original_headers: dict[str, str],
original_request_body: dict[str, Any],
) -> None: ...

3
src/api/handlers/base/cli_request_mixin.py
View File

@@ -11,6 +11,7 @@ from src.api.handlers.base.utils import get_format_converter_registry
from src.core.logger import logger
if TYPE_CHECKING:
from src.api.handlers.base.cli_protocol import CliHandlerProtocol
from src.core.api_format import EndpointDefinition
@@ -63,7 +64,7 @@ class CliRequestMixin:
return None
def extract_model_from_request(
self,
self: CliHandlerProtocol,
request_body: dict[str, Any],
path_params: dict[str, Any] | None = None, # noqa: ARG002 - 子类使用
) -> str:

3
src/api/handlers/base/cli_stream_mixin.py
View File

@@ -53,6 +53,7 @@ from src.utils.timeout import read_first_chunk_with_ttfb_timeout
from .cli_sse_helpers import _format_converted_events_to_sse
if TYPE_CHECKING:
from src.api.handlers.base.cli_protocol import CliHandlerProtocol
from src.models.database import Provider, ProviderAPIKey, ProviderEndpoint
@@ -60,7 +61,7 @@ class CliStreamMixin:
"""流式处理核心方法的 Mixin"""
async def process_stream(
self,
self: CliHandlerProtocol,
original_request_body: dict[str, Any],
original_headers: dict[str, str],
query_params: dict[str, str] | None = None,

3
src/api/handlers/base/cli_sync_mixin.py
View File

@@ -39,6 +39,7 @@ from src.services.provider.stream_policy import (
from src.services.provider.transport import build_provider_url
if TYPE_CHECKING:
from src.api.handlers.base.cli_protocol import CliHandlerProtocol
from src.models.database import Provider, ProviderAPIKey, ProviderEndpoint
@@ -46,7 +47,7 @@ class CliSyncMixin:
"""同步处理相关方法的 Mixin"""
async def process_sync(
self,
self: CliHandlerProtocol,
original_request_body: dict[str, Any],
original_headers: dict[str, str],
query_params: dict[str, str] | None = None,

50
src/core/modules/registry.py
View File

@@ -158,12 +158,26 @@ class ModuleRegistry:
# ========== 激活状态检查 ==========
def is_active(self, name: str, db: Session) -> bool:
def is_active(self, name: str, db: Session, _visited: set[str] | None = None) -> bool:
"""
检查模块是否最终激活
激活条件：available && enabled && 依赖模块都激活
Args:
name: 模块名称
db: 数据库会话
_visited: 内部递归防御，防止循环依赖导致无限递归
"""
if _visited is None:
_visited = set()
if name in _visited:
logger.warning(f"Circular dependency detected in module activation chain: {name}")
return False
_visited.add(name)
if not self.is_available(name):
return False
if not self.is_enabled(name, db):
@@ -172,7 +186,7 @@ class ModuleRegistry:
# 检查依赖模块
module = self._modules[name]
for dep in module.metadata.dependencies:
if not self.is_active(dep, db):
if not self.is_active(dep, db, _visited):
return False
return True
@@ -205,6 +219,28 @@ class ModuleRegistry:
logger.warning(f"Module [{name}] config validation error: {e}")
return False, f"配置验证出错: {str(e)}"
def reconcile_module_state(self, name: str, db: Session) -> None:
"""
修复模块启用状态与配置的一致性
如果模块已启用但配置验证失败（例如依赖的 Provider Key 被删除），
则自动禁用该模块以保证状态一致性。
此方法是显式的写操作，应在需要状态修复的场景中调用，
而非在纯查询方法中隐式执行。
"""
if name not in self._modules:
return
if not self.is_available(name):
return
config_validated, config_error = self.validate_config(name, db)
if self.is_enabled(name, db) and not config_validated:
self.set_enabled(name, False, db)
logger.info(
f"Module [{name}] auto-disabled: config validation failed" f" ({config_error})"
)
# ========== 状态查询 ==========
def get_module_status(
@@ -236,14 +272,6 @@ class ModuleRegistry:
# 获取启用状态
enabled = self.is_enabled(name, db) if available else False
# 配置验证失败时自动禁用模块
# 注意：此处故意在 get_status() 中写入，以确保模块状态与配置同步
# 场景：用户删除了模块所依赖的 Provider Key 后，模块应自动关闭
# 权衡：查询方法中的写操作副作用 vs 状态一致性保证
if enabled and not config_validated:
self.set_enabled(name, False, db)
enabled = False
# 计算激活状态：available && enabled && config_validated && 依赖模块都激活
is_active = self.is_active(name, db) if available else False
active = is_active and config_validated
@@ -293,6 +321,7 @@ class ModuleRegistry:
if name not in self._modules:
return None
self.reconcile_module_state(name, db)
health = await self.check_health(name) if self.is_available(name) else ModuleHealth.UNKNOWN
return self.get_module_status(name, db, health=health)
@@ -309,6 +338,7 @@ class ModuleRegistry:
"""获取所有模块状态（同步版本，不含健康检查）"""
result = {}
for name in self._modules:
self.reconcile_module_state(name, db)
status = self.get_module_status(name, db)
if status:
result[name] = status

2
src/plugins/manager.py
View File

@@ -434,7 +434,7 @@ class PluginManager:
# 创建插件名称到插件对象的映射
plugin_map = {plugin.name: plugin for plugin in plugins}
# 计算每个插件的入度（被依赖的次数）
# 计算每个插件的入度（未满足的依赖数量）
in_degree = {plugin.name: 0 for plugin in plugins}
# 构建依赖图

591
src/services/cache/_candidate_builder.py vendored Normal file
View File

@@ -0,0 +1,591 @@
"""
候选构建器 (CandidateBuilder)
从 CacheAwareScheduler 拆分出的候选构建逻辑，负责:
- 查询活跃 Provider
- 检查模型支持
- 检查 Key 可用性
- 构建候选列表
"""
from __future__ import annotations
import re
from collections.abc import Sequence
from typing import TYPE_CHECKING
from sqlalchemy.orm import Session, selectinload
from src.core.api_format.conversion.compatibility import is_format_compatible
from src.core.api_format.enums import EndpointKind
from src.core.api_format.signature import make_signature_key, parse_signature_key
from src.core.key_capabilities import check_capability_match
from src.core.logger import logger
from src.core.model_permissions import check_model_allowed_with_mappings
from src.models.database import (
Model,
Provider,
ProviderAPIKey,
ProviderEndpoint,
)
from src.services.cache.quota_skipper import is_key_quota_exhausted
from src.services.health.monitor import health_monitor
from src.services.provider.format import normalize_endpoint_signature
if TYPE_CHECKING:
from src.models.database import GlobalModel
from src.services.cache.aware_scheduler import CacheAwareScheduler, ProviderCandidate
from src.services.cache.model_cache import ModelCacheService
def _sort_endpoints_by_family_priority(
eps: Sequence[ProviderEndpoint],
) -> list[ProviderEndpoint]:
"""按 ApiFamily 优先级对端点排序（同分组内使用）。"""
from src.core.api_format.enums import ApiFamily
def sort_key(ep: ProviderEndpoint) -> int:
family_str = str(getattr(ep, "api_family", "") or "").strip().lower()
try:
return ApiFamily(family_str).priority
except ValueError:
return 99
return sorted(eps, key=sort_key)
class CandidateBuilder:
"""候选构建器，负责查询 Provider、检查模型支持和 Key 可用性、构建候选列表。"""
def __init__(self, scheduler: CacheAwareScheduler) -> None:
self._scheduler = scheduler
def _query_providers(
self,
db: Session,
provider_offset: int = 0,
provider_limit: int | None = None,
) -> list[Provider]:
"""
查询活跃的 Providers（带预加载）
Args:
db: 数据库会话
provider_offset: 分页偏移
provider_limit: 分页限制
Returns:
Provider 列表
"""
provider_query = (
db.query(Provider)
.options(
# 预加载 Provider 级别的 api_keys
selectinload(Provider.api_keys),
# 预加载 endpoints（用于按 api_format 选择请求配置）
selectinload(Provider.endpoints),
# 同时加载 models 和 global_model 关系
selectinload(Provider.models).selectinload(Model.global_model),
)
.filter(Provider.is_active.is_(True))
.order_by(Provider.provider_priority.asc())
)
if provider_offset:
provider_query = provider_query.offset(provider_offset)
if provider_limit:
provider_query = provider_query.limit(provider_limit)
return provider_query.all()
async def _check_model_support(
self,
db: Session,
provider: Provider,
model_name: str,
api_format: str | None = None,
is_stream: bool = False,
capability_requirements: dict[str, bool] | None = None,
) -> tuple[bool, str | None, list[str] | None, set[str] | None]:
"""
检查 Provider 是否支持指定模型（可选检查流式支持和能力需求）
模型能力检查在这里进行（而不是在 Key 级别），因为：
- 模型支持的能力是全局的，与具体的 Key 无关
- 如果模型不支持某能力，整个 Provider 的所有 Key 都应该被跳过
仅支持直接匹配 GlobalModel.name（外部请求不接受映射名）
Args:
db: 数据库会话
provider: Provider 对象
model_name: 模型名称（必须是 GlobalModel.name）
is_stream: 是否是流式请求，如果为 True 则同时检查流式支持
capability_requirements: 能力需求（可选），用于检查模型是否支持所需能力
Returns:
(is_supported, skip_reason, supported_capabilities, provider_model_names)
- is_supported: 是否支持
- skip_reason: 跳过原因
- supported_capabilities: 模型支持的能力列表
- provider_model_names: Provider 侧可用的模型名称集合（主名称 + 映射名称，按 api_format 过滤）
"""
# Avoid holding a DB connection while awaiting cache/Redis inside ModelCacheService.
self._scheduler._release_db_connection_before_await(db)
# 仅接受 GlobalModel.name（不允许映射名）
normalized_name = model_name.strip() if isinstance(model_name, str) else ""
if not normalized_name:
return False, "模型不存在或名称无效", None, None
global_model = await ModelCacheService.get_global_model_by_name(db, normalized_name)
if not global_model or not global_model.is_active:
return False, "模型不存在或已停用", None, None
# 找到 GlobalModel 后，检查当前 Provider 是否支持
is_supported, skip_reason, caps, provider_model_names = (
await self._check_model_support_for_global_model(
db,
provider,
global_model,
model_name,
api_format,
is_stream,
capability_requirements,
)
)
return is_supported, skip_reason, caps, provider_model_names
async def _check_model_support_for_global_model(
self,
db: Session,
provider: Provider,
global_model: GlobalModel,
model_name: str,
api_format: str | None = None,
is_stream: bool = False,
capability_requirements: dict[str, bool] | None = None,
) -> tuple[bool, str | None, list[str] | None, set[str] | None]:
"""
检查 Provider 是否支持指定的 GlobalModel
Args:
db: 数据库会话
provider: Provider 对象
global_model: GlobalModel 对象
model_name: 用户请求的模型名称（用于错误消息）
is_stream: 是否是流式请求
capability_requirements: 能力需求
Returns:
(is_supported, skip_reason, supported_capabilities, provider_model_names)
"""
# 确保 global_model 附加到当前 Session
# 注意：从缓存重建的对象是 transient 状态，不能使用 load=False
# 使用 load=True（默认）允许 SQLAlchemy 正确处理 transient 对象
from sqlalchemy import inspect
insp = inspect(global_model)
if insp.transient or insp.detached:
# transient/detached 对象：使用默认 merge（会查询 DB 检查是否存在）
global_model = db.merge(global_model)
else:
# persistent 对象：已经附加到 session，无需 merge
pass
# 获取模型支持的能力列表
model_supported_capabilities: list[str] = list(global_model.supported_capabilities or [])
# 查询该 Provider 是否有实现这个 GlobalModel
for model in provider.models:
if model.global_model_id == global_model.id and model.is_active:
# 检查流式支持
if is_stream:
supports_streaming = model.get_effective_supports_streaming()
if not supports_streaming:
return False, f"模型 {model_name} 在此 Provider 不支持流式", None, None
# 检查模型是否支持所需的能力（在 Provider 级别检查，而不是 Key 级别）
# 只有当 model_supported_capabilities 非空时才进行检查
# 空列表意味着模型没有配置能力限制，默认支持所有能力
if capability_requirements and model_supported_capabilities:
for cap_name, is_required in capability_requirements.items():
if is_required and cap_name not in model_supported_capabilities:
return (
False,
f"模型 {model_name} 不支持能力: {cap_name}",
list(model_supported_capabilities),
None,
)
provider_model_names: set[str] = {model.provider_model_name}
raw_mappings = model.provider_model_mappings
if isinstance(raw_mappings, list):
for raw in raw_mappings:
if not isinstance(raw, dict):
continue
name = raw.get("name")
if not isinstance(name, str) or not name.strip():
continue
mapping_api_formats = raw.get("api_formats")
if api_format and mapping_api_formats:
# 新模式：endpoint signature（family:kind），按小写 canonical 比较
if isinstance(mapping_api_formats, list):
target = str(api_format).strip().lower()
allowed = {
str(fmt).strip().lower() for fmt in mapping_api_formats if fmt
}
if target not in allowed:
continue
provider_model_names.add(name.strip())
return True, None, list(model_supported_capabilities), provider_model_names
return False, "Provider 未实现此模型", None, None
def _check_key_availability(
self,
key: ProviderAPIKey,
api_format: str | None,
model_name: str,
capability_requirements: dict[str, bool] | None = None,
model_mappings: list[str] | None = None,
candidate_models: set[str] | None = None,
*,
provider_type: str | None = None,
) -> tuple[bool, str | None, str | None]:
"""
检查 API Key 的可用性
注意：模型能力检查已移到 _check_model_support 中进行（Provider 级别），
这里只检查 Key 级别的能力匹配。
Args:
key: API Key 对象
model_name: 模型名称（GlobalModel.name）
capability_requirements: 能力需求（可选）
model_mappings: GlobalModel 的映射列表（用于通配符匹配）
candidate_models: Provider 侧可用的模型名称集合（用于限制映射匹配范围）
Returns:
(is_available, skip_reason, mapping_matched_model)
- is_available: Key 是否可用
- skip_reason: 不可用时的原因
- mapping_matched_model: 通过映射匹配到的模型名（用于实际请求）
"""
# 检查熔断器状态（使用详细状态方法获取更丰富的跳过原因，按 API 格式）
is_available, circuit_reason = health_monitor.get_circuit_breaker_status(
key, api_format=api_format
)
if not is_available:
return False, circuit_reason or "熔断器已打开", None
# 模型权限检查：使用 allowed_models 白名单
# None = 允许所有模型，[] = 拒绝所有模型，["a","b"] = 只允许指定模型
# 支持通配符映射匹配（通过 model_mappings）
try:
is_allowed, mapping_matched_model = check_model_allowed_with_mappings(
model_name=model_name,
allowed_models=key.allowed_models,
model_mappings=model_mappings,
candidate_models=candidate_models,
)
if mapping_matched_model:
logger.debug(
"[Scheduler] Key {}... 模型名匹配: model={} -> {}, allowed_models={}",
key.id[:8],
model_name,
mapping_matched_model,
key.allowed_models,
)
except TimeoutError:
# 正则匹配超时（可能是 ReDoS 攻击或复杂模式）
logger.warning("映射匹配超时: key_id={}, model={}", key.id, model_name)
return False, "映射匹配超时，请简化配置", None
except re.error as e:
# 正则语法错误（配置问题）
logger.warning("映射规则无效: key_id={}, model={}, error={}", key.id, model_name, e)
return False, f"映射规则无效: {str(e)}", None
except Exception as e:
# 其他未知异常
logger.error(
"映射匹配异常: key_id={}, model={}, error={}", key.id, model_name, e, exc_info=True
)
# 异常时保守处理：不允许使用该 Key
return False, "映射匹配失败", None
if not is_allowed:
return (
False,
f"Key 不支持 {model_name}",
None,
)
# Key 级别的能力匹配检查
# 注意：模型级别的能力检查已在 _check_model_support 中完成
# 始终执行检查，即使 capability_requirements 为空
# 因为 check_capability_match 会检查 Key 的 EXCLUSIVE 能力是否被浪费
key_caps: dict[str, bool] = dict(key.capabilities or {})
is_match, skip_reason = check_capability_match(key_caps, capability_requirements)
if not is_match:
return False, skip_reason, None
effective_model_name = mapping_matched_model or model_name
quota_exhausted, quota_reason = is_key_quota_exhausted(
provider_type,
key,
model_name=effective_model_name,
)
if quota_exhausted:
return False, quota_reason, mapping_matched_model
return True, None, mapping_matched_model
async def _build_candidates(
self,
db: Session,
providers: list[Provider],
client_format: str,
model_name: str,
affinity_key: str | None,
model_mappings: list[str] | None = None,
max_candidates: int | None = None,
is_stream: bool = False,
capability_requirements: dict[str, bool] | None = None,
global_conversion_enabled: bool = True,
) -> "list[ProviderCandidate]":
"""
构建候选列表
Key 直属 Provider，通过 api_formats 筛选符合端点格式的 Key。
Args:
db: 数据库会话
providers: Provider 列表
client_format: 客户端请求的 API 格式
model_name: 模型名称（GlobalModel.name）
affinity_key: 亲和性标识符（通常为API Key ID）
model_mappings: GlobalModel 的映射列表（用于 Key.allowed_models 通配符匹配）
max_candidates: 最大候选数
is_stream: 是否是流式请求，如果为 True 则过滤不支持流式的 Provider
capability_requirements: 能力需求（可选）
global_conversion_enabled: 格式转换总开关（数据库配置），关闭时禁止任何跨格式转换
Returns:
候选列表
"""
from src.services.cache.aware_scheduler import ProviderCandidate
candidates: list[ProviderCandidate] = []
client_format_str = normalize_endpoint_signature(client_format)
client_sig = parse_signature_key(client_format_str)
client_family, client_kind = client_sig.api_family, client_sig.endpoint_kind
# chat/cli 互相可回退（用于同协议族下的端点变体），video/image 等不跨类回退
if client_kind in {EndpointKind.CHAT, EndpointKind.CLI}:
allowed_kinds = {EndpointKind.CHAT, EndpointKind.CLI}
else:
allowed_kinds = {client_kind}
for provider in providers:
logger.debug(
"[Scheduler] Checking provider: {}, endpoints={}",
provider.name,
len(provider.endpoints) if provider.endpoints else 0,
)
# 按端点格式分别判断兼容性与模型/Key 可用性：
# - 同格式端点优先（needs_conversion=False）
# - 跨格式端点次之（needs_conversion=True）
model_support_cache: dict[
str, tuple[bool, str | None, list[str] | None, set[str] | None]
] = {}
exact_candidates: list[ProviderCandidate] = []
convertible_candidates: list[ProviderCandidate] = []
# 使用新架构字段 (api_family, endpoint_kind) 进行预过滤与排序：
# - family/kind 匹配的 endpoint 排在前面（但不做硬过滤，避免破坏格式转换路径）
# - chat/cli 请求允许互相回退（优先同 kind）
# - video 等请求只允许同 kind
endpoints = list(provider.endpoints or [])
allowed_kind_values = {k.value for k in allowed_kinds}
preferred: list[ProviderEndpoint] = []
preferred_other_family: list[ProviderEndpoint] = []
fallback: list[ProviderEndpoint] = []
fallback_other_family: list[ProviderEndpoint] = []
for ep in endpoints:
if not getattr(ep, "is_active", False):
continue
raw_family = getattr(ep, "api_family", None)
raw_kind = getattr(ep, "endpoint_kind", None)
if not isinstance(raw_family, str) or not raw_family.strip():
continue
if not isinstance(raw_kind, str) or not raw_kind.strip():
continue
ep_family = raw_family.strip().lower()
ep_kind = raw_kind.strip().lower()
if allowed_kind_values and ep_kind not in allowed_kind_values:
continue
same_family = ep_family == client_family.value
same_kind = ep_kind == client_kind.value
if same_kind and same_family:
preferred.append(ep)
elif same_kind:
preferred_other_family.append(ep)
elif same_family:
fallback.append(ep)
else:
fallback_other_family.append(ep)
endpoints = (
_sort_endpoints_by_family_priority(preferred)
+ _sort_endpoints_by_family_priority(preferred_other_family)
+ _sort_endpoints_by_family_priority(fallback)
+ _sort_endpoints_by_family_priority(fallback_other_family)
)
for endpoint in endpoints:
logger.debug(
"[Scheduler] Checking endpoint: family={}, kind={}, is_active={}, base_url={}",
getattr(endpoint, "api_family", None),
getattr(endpoint, "endpoint_kind", None),
getattr(endpoint, "is_active", None),
(endpoint.base_url[:50] if endpoint.base_url else "N/A"),
)
if not endpoint.is_active:
logger.debug("[Scheduler] Endpoint skipped: not active")
continue
endpoint_format_str = make_signature_key(
str(getattr(endpoint, "api_family", "")).strip().lower(),
str(getattr(endpoint, "endpoint_kind", "")).strip().lower(),
)
# 计算格式转换开关状态（三层优先级）
#
# 1) 全局开关（数据库配置）关闭 -> 禁止任何跨格式转换
# 2) 全局开关开启 -> 允许跨格式转换
# 3) 提供商覆盖（Provider.enable_format_conversion）开启 -> 强制允许（跳过端点检查）
# 4) 否则 -> 由端点配置 format_acceptance_config 决定是否允许
provider_allows_conversion = getattr(provider, "enable_format_conversion", True)
skip_endpoint_check = global_conversion_enabled or provider_allows_conversion
is_compatible, needs_conversion, _compat_reason = is_format_compatible(
client_format_str,
endpoint_format_str,
getattr(endpoint, "format_acceptance_config", None),
is_stream,
global_conversion_enabled,
skip_endpoint_check=skip_endpoint_check,
)
logger.debug(
"[Scheduler] Format compatibility: client={}, endpoint={}, compatible={}, "
"global={}, provider={}, skip_endpoint={}, reason={}",
client_format_str,
endpoint_format_str,
is_compatible,
global_conversion_enabled,
provider_allows_conversion,
skip_endpoint_check,
_compat_reason,
)
if not is_compatible:
continue
# 检查模型支持（按端点格式过滤 provider_model_mappings）
if endpoint_format_str not in model_support_cache:
model_support_cache[endpoint_format_str] = await self._check_model_support(
db,
provider,
model_name,
api_format=endpoint_format_str,
is_stream=is_stream,
capability_requirements=capability_requirements,
)
supports_model, skip_reason, _model_caps, provider_model_names = (
model_support_cache[endpoint_format_str]
)
logger.debug(
"[Scheduler] Model support: provider={}, model={}, supports={}, reason={}",
provider.name,
model_name,
supports_model,
skip_reason,
)
if not supports_model:
logger.debug(
f"Provider {provider.name} 端点 {endpoint_format_str} "
f"不支持模型 {model_name}: {skip_reason}"
)
continue
# Key 直属 Provider，通过 api_formats 按端点格式筛选
# api_formats=None 视为"全支持"（兼容历史数据）
active_keys = [
key
for key in provider.api_keys
if key.is_active
and (key.api_formats is None or endpoint_format_str in key.api_formats)
]
if not active_keys:
continue
# 检查是否所有 Key 都是 TTL=0（轮换模式）
use_random = all((key.cache_ttl_minutes or 0) == 0 for key in active_keys)
if use_random and len(active_keys) > 1:
logger.debug(
f" Provider {provider.name} 启用 Key 轮换模式 "
f"(endpoint_format={endpoint_format_str}, {len(active_keys)} keys)"
)
keys = self._scheduler._candidate_sorter._shuffle_keys_by_internal_priority(
active_keys, affinity_key, use_random
)
for key in keys:
# Key 级别检查（健康度/熔断按 provider_format bucket）
# 传入 provider_model_names 作为 candidate_models，
# 用于检查 Key 的 allowed_models 是否支持 Provider 定义的模型名称
is_available, key_skip_reason, mapping_matched_model = (
self._check_key_availability(
key,
endpoint_format_str,
model_name,
capability_requirements,
model_mappings=model_mappings,
candidate_models=provider_model_names,
provider_type=getattr(provider, "provider_type", None),
)
)
candidate = ProviderCandidate(
provider=provider,
endpoint=endpoint,
key=key,
is_skipped=not is_available,
skip_reason=key_skip_reason,
mapping_matched_model=mapping_matched_model,
needs_conversion=needs_conversion,
provider_api_format=str(endpoint_format_str or ""),
)
if needs_conversion:
convertible_candidates.append(candidate)
else:
exact_candidates.append(candidate)
candidates.extend(exact_candidates)
candidates.extend(convertible_candidates)
# max_candidates 截断应在所有候选收集完成后统一处理，确保优先级排序正确
if max_candidates and len(candidates) > max_candidates:
candidates = candidates[:max_candidates]
return candidates

268
src/services/cache/_candidate_sorter.py vendored Normal file
View File

@@ -0,0 +1,268 @@
"""
候选排序器 (CandidateSorter)
从 CacheAwareScheduler 拆分出的候选排序逻辑，负责:
- 优先级模式排序（provider / global_key）
- 负载均衡模式排序
- Key 内部按优先级分组打乱
"""
from __future__ import annotations
import random
from collections import defaultdict
from typing import TYPE_CHECKING
from src.core.logger import logger
from src.services.system.config import SystemConfigService
if TYPE_CHECKING:
from sqlalchemy.orm import Session
from src.models.database import ProviderAPIKey
from src.services.cache.aware_scheduler import CacheAwareScheduler, ProviderCandidate
class CandidateSorter:
"""候选排序器，负责优先级模式排序、负载均衡排序和 Key 内部打乱。"""
def __init__(self, scheduler: CacheAwareScheduler) -> None:
self._scheduler = scheduler
def _apply_priority_mode_sort(
self,
candidates: list[ProviderCandidate],
db: Session,
affinity_key: str | None = None,
api_format: str | None = None,
) -> list[ProviderCandidate]:
"""
根据优先级模式对候选列表排序（数字越小越优先）
排序规则（受 keep_priority_on_conversion 配置影响）：
1. 如果全局配置 keep_priority_on_conversion=True，所有候选保持原优先级
2. 否则，按 needs_conversion 和 provider.keep_priority_on_conversion 分组：
- 保持优先级的候选（exact 或 provider.keep_priority_on_conversion=True）按原优先级排序
- 需要降级的候选（convertible 且 provider.keep_priority_on_conversion=False）整体排在后面
3. 在同一组内，按优先级模式排序：
- provider: 按 Provider.provider_priority -> Key.internal_priority 排序
- global_key: 按 Key.global_priority_by_format 排序
"""
if not candidates:
return candidates
s = self._scheduler
# 全局配置：如果开启，所有候选保持原优先级
global_keep_priority = SystemConfigService.is_keep_priority_on_conversion(db)
if global_keep_priority:
# 全局开启：不分组，直接按优先级模式排序
if s.priority_mode == s.PRIORITY_MODE_GLOBAL_KEY:
return self._sort_by_global_priority_with_hash(candidates, affinity_key, api_format)
# 提供商优先模式：保持构建时的顺序（已按 provider_priority 排序）
return candidates
# 全局未开启：按是否需要降级分组
# - 不需要降级：exact 候选 或 provider.keep_priority_on_conversion=True 的 convertible 候选
# - 需要降级：convertible 且 provider.keep_priority_on_conversion=False
keep_priority_candidates: list[ProviderCandidate] = []
demote_candidates: list[ProviderCandidate] = []
for c in candidates:
if not c.needs_conversion:
# exact 候选：不需要降级
keep_priority_candidates.append(c)
elif getattr(c.provider, "keep_priority_on_conversion", False):
# convertible 但提供商配置了保持优先级
keep_priority_candidates.append(c)
else:
# convertible 且未配置保持优先级：降级
demote_candidates.append(c)
if s.priority_mode == s.PRIORITY_MODE_GLOBAL_KEY:
# 全局 Key 优先模式：分别对两组排序后合并
sorted_keep = self._sort_by_global_priority_with_hash(
keep_priority_candidates, affinity_key, api_format
)
sorted_demote = self._sort_by_global_priority_with_hash(
demote_candidates, affinity_key, api_format
)
return sorted_keep + sorted_demote
# 提供商优先模式：保持优先级的在前，降级的在后（各组内部顺序已由构建时保证）
return keep_priority_candidates + demote_candidates
def _sort_by_global_priority_with_hash(
self,
candidates: list[ProviderCandidate],
affinity_key: str | None = None,
api_format: str | None = None,
) -> list[ProviderCandidate]:
"""
按 global_priority_by_format 分组排序，同优先级内通过哈希分散实现负载均衡
排序逻辑：
1. 按 global_priority_by_format[api_format] 分组（数字小的优先，NULL 排后面）
2. 同优先级组内，使用 affinity_key 哈希分散
3. 确保同一用户请求稳定选择同一个 Key（缓存亲和性）
"""
def get_priority(candidate: ProviderCandidate) -> int:
"""获取候选的优先级"""
if not candidate.key:
return 999999
priority_by_format = candidate.key.global_priority_by_format or {}
if api_format and api_format in priority_by_format:
return priority_by_format[api_format]
return 999999 # NULL 排在后面
# 按优先级分组
priority_groups: dict[int, list[ProviderCandidate]] = defaultdict(list)
for candidate in candidates:
priority = get_priority(candidate)
priority_groups[priority].append(candidate)
result = []
for priority in sorted(priority_groups.keys()): # 数字小的优先级高
group = priority_groups[priority]
if len(group) > 1 and affinity_key:
# 同优先级内哈希分散负载均衡
scored_candidates = []
for candidate in group:
key_id = candidate.key.id if candidate.key else ""
hash_value = self._scheduler._affinity_hash(affinity_key, key_id)
scored_candidates.append((hash_value, candidate))
# 按哈希值排序
sorted_group = [c for _, c in sorted(scored_candidates, key=lambda x: x[0])]
result.extend(sorted_group)
else:
# 单个候选或没有 affinity_key，按次要排序条件排序
def secondary_sort(c: ProviderCandidate) -> tuple[int, int, str]:
pp = c.provider.provider_priority
ip = c.key.internal_priority if c.key else None
return (
pp if pp is not None else 999999,
ip if ip is not None else 999999,
c.key.id if c.key else "",
)
result.extend(sorted(group, key=secondary_sort))
return result
def _apply_load_balance(
self, candidates: list[ProviderCandidate], api_format: str | None = None
) -> list[ProviderCandidate]:
"""
负载均衡模式：同优先级内随机轮换
排序逻辑：
1. 按优先级分组（provider_priority, internal_priority 或 global_priority_by_format）
2. 同优先级组内随机打乱
3. 不考虑缓存亲和性
"""
if not candidates:
return candidates
s = self._scheduler
priority_groups: dict[tuple, list[ProviderCandidate]] = defaultdict(list)
# 根据优先级模式选择分组方式
if s.priority_mode == s.PRIORITY_MODE_GLOBAL_KEY:
# 全局 Key 优先模式：按格式特定优先级分组
for candidate in candidates:
priority = 999999
if candidate.key:
priority_by_format = candidate.key.global_priority_by_format or {}
if api_format and api_format in priority_by_format:
priority = priority_by_format[api_format]
priority_groups[(priority,)].append(candidate)
else:
# 提供商优先模式：按 (provider_priority, internal_priority) 分组
for candidate in candidates:
pp = candidate.provider.provider_priority
ip = candidate.key.internal_priority if candidate.key else None
key = (
pp if pp is not None else 999999,
ip if ip is not None else 999999,
)
priority_groups[key].append(candidate)
result: list[ProviderCandidate] = []
for priority in sorted(priority_groups.keys()):
group = priority_groups[priority]
if len(group) > 1:
# 同优先级内随机打乱
shuffled = list(group)
random.shuffle(shuffled)
result.extend(shuffled)
else:
result.extend(group)
return result
def _shuffle_keys_by_internal_priority(
self,
keys: list[ProviderAPIKey],
affinity_key: str | None = None,
use_random: bool = False,
) -> list[ProviderAPIKey]:
"""
对 API Key 按 internal_priority 分组，同优先级内部基于 affinity_key 进行确定性打乱
目的：
- 数字越小越优先使用
- 同优先级 Key 之间实现负载均衡
- 使用 affinity_key 哈希确保同一请求 Key 的请求稳定（避免破坏缓存亲和性）
- 当 use_random=True 时，使用随机排序实现轮换（用于 TTL=0 的场景）
Args:
keys: API Key 列表
affinity_key: 亲和性标识符（通常为 API Key ID，用于确定性打乱）
use_random: 是否使用随机排序（TTL=0 时为 True）
Returns:
排序后的 Key 列表
"""
if not keys:
return []
# 按 internal_priority 分组
priority_groups: dict[int, list[ProviderAPIKey]] = defaultdict(list)
for key in keys:
priority = key.internal_priority if key.internal_priority is not None else 999999
priority_groups[priority].append(key)
# 对每个优先级组内的 Key 进行打乱
result = []
for priority in sorted(priority_groups.keys()): # 数字小的优先级高，排前面
group_keys = priority_groups[priority]
if len(group_keys) > 1:
if use_random:
# TTL=0 模式：使用随机排序实现 Key 轮换
shuffled = list(group_keys)
random.shuffle(shuffled)
result.extend(shuffled)
elif affinity_key:
# 正常模式：使用哈希确定性打乱（保持缓存亲和性）
key_scores = []
for key in group_keys:
hash_value = self._scheduler._affinity_hash(affinity_key, key.id)
key_scores.append((hash_value, key))
# 按哈希值排序
sorted_group = [key for _, key in sorted(key_scores, key=lambda x: x[0])]
result.extend(sorted_group)
else:
# 没有 affinity_key 时按 ID 排序保持稳定性
result.extend(sorted(group_keys, key=lambda k: k.id))
else:
# 单个 Key 直接添加
result.extend(group_keys)
return result

829
src/services/cache/aware_scheduler.py vendored
View File

@@ -32,46 +32,37 @@ from __future__ import annotations
import hashlib
import math
import random
import re
import time
from collections import defaultdict
from collections.abc import Sequence
from dataclasses import dataclass
from typing import TYPE_CHECKING, Any
from typing import Any
from sqlalchemy.orm import Session, selectinload
from sqlalchemy.orm import Session
from src.core.api_format.conversion.compatibility import is_format_compatible
from src.core.api_format.enums import ApiFamily, EndpointKind
from src.core.api_format.signature import make_signature_key, parse_signature_key
from src.core.exceptions import ModelNotSupportedException, ProviderNotAvailableException
from src.core.key_capabilities import check_capability_match
from src.core.logger import logger
from src.core.model_permissions import (
check_model_allowed,
check_model_allowed_with_mappings,
get_allowed_models_preview,
merge_allowed_models,
)
from src.models.database import (
ApiKey,
Model,
Provider,
ProviderAPIKey,
ProviderEndpoint,
)
from src.services.cache.quota_skipper import is_key_quota_exhausted
if TYPE_CHECKING:
from src.models.database import GlobalModel
from src.services.cache._candidate_builder import (
CandidateBuilder,
)
from src.services.cache._candidate_builder import (
_sort_endpoints_by_family_priority as _sort_endpoints_by_family_priority,
)
from src.services.cache._candidate_sorter import CandidateSorter
from src.services.cache.affinity_manager import (
CacheAffinityManager,
get_affinity_manager,
)
from src.services.cache.model_cache import ModelCacheService
from src.services.health.monitor import health_monitor
from src.services.provider.format import normalize_endpoint_signature
from src.services.rate_limit.adaptive_reservation import (
AdaptiveReservationManager,
@@ -154,21 +145,6 @@ class ConcurrencySnapshot:
)
def _sort_endpoints_by_family_priority(
eps: Sequence[ProviderEndpoint],
) -> list[ProviderEndpoint]:
"""按 ApiFamily 优先级对端点排序（同分组内使用）。"""
def sort_key(ep: ProviderEndpoint) -> int:
family_str = str(getattr(ep, "api_family", "") or "").strip().lower()
try:
return ApiFamily(family_str).priority
except ValueError:
return 99
return sorted(eps, key=sort_key)
class CacheAwareScheduler:
"""
缓存感知调度器
@@ -248,6 +224,10 @@ class CacheAwareScheduler:
"last_reservation_result": None,
}
# 初始化拆分出的子模块
self._candidate_builder = CandidateBuilder(self)
self._candidate_sorter = CandidateSorter(self)
@staticmethod
def _release_db_connection_before_await(db: Session) -> None:
"""
@@ -436,7 +416,7 @@ class CacheAwareScheduler:
- 总槽位: 有效 RPM 限制（固定值或学习到的值）
- 预留比例: 由 AdaptiveReservationManager 根据置信度和负载动态计算
- 缓存用户可用: 全部槽位
- 新用户可用: 总槽位 × (1 - 动态预留比例)
- 新用户可用: 总槽位 x (1 - 动态预留比例)
Args:
key: ProviderAPIKey对象
@@ -625,8 +605,8 @@ class CacheAwareScheduler:
预先获取所有可用的 Provider/Endpoint/Key 组合
重构后的方法将逻辑拆分为：
1. _query_providers: 数据库查询逻辑
2. _build_candidates: 候选构建逻辑
1. _query_providers: 数据库查询逻辑（委托给 CandidateBuilder）
2. _build_candidates: 候选构建逻辑（委托给 CandidateBuilder）
3. _apply_cache_affinity: 缓存亲和性处理
Args:
@@ -714,8 +694,8 @@ class CacheAwareScheduler:
)
return [], global_model_id, queried_provider_count
# 1. 查询 Providers
providers = self._query_providers(
# 1. 查询 Providers（委托给 CandidateBuilder）
providers = self._candidate_builder._query_providers(
db=db,
provider_offset=provider_offset,
provider_limit=provider_limit,
@@ -755,12 +735,12 @@ class CacheAwareScheduler:
if not providers:
return [], global_model_id, queried_provider_count
# 2. 构建候选列表（传入 is_stream 和 capability_requirements 用于过滤）
# 2. 构建候选列表（委托给 CandidateBuilder）
# 格式转换总开关（数据库配置）：关闭时禁止任何跨格式候选进入队列
global_conversion_enabled = SystemConfigService.is_format_conversion_enabled(db)
candidates = await self._build_candidates(
candidates = await self._candidate_builder._build_candidates(
db=db,
providers=providers,
client_format=target_format,
@@ -818,8 +798,10 @@ class CacheAwareScheduler:
if not candidates:
return candidates
# 1. 优先级模式排序
candidates = self._apply_priority_mode_sort(candidates, db, affinity_key, api_format)
# 1. 优先级模式排序（委托给 CandidateSorter）
candidates = self._candidate_sorter._apply_priority_mode_sort(
candidates, db, affinity_key, api_format
)
# 2. 调度模式排序
if self.scheduling_mode == self.SCHEDULING_MODE_CACHE_AFFINITY:
@@ -832,7 +814,7 @@ class CacheAwareScheduler:
global_model_id=global_model_id,
)
elif self.scheduling_mode == self.SCHEDULING_MODE_LOAD_BALANCE:
candidates = self._apply_load_balance(candidates, api_format)
candidates = self._candidate_sorter._apply_load_balance(candidates, api_format)
for candidate in candidates:
candidate.is_cached = False
else:
@@ -841,532 +823,6 @@ class CacheAwareScheduler:
return candidates
def _query_providers(
self,
db: Session,
provider_offset: int = 0,
provider_limit: int | None = None,
) -> list[Provider]:
"""
查询活跃的 Providers（带预加载）
Args:
db: 数据库会话
provider_offset: 分页偏移
provider_limit: 分页限制
Returns:
Provider 列表
"""
provider_query = (
db.query(Provider)
.options(
# 预加载 Provider 级别的 api_keys
selectinload(Provider.api_keys),
# 预加载 endpoints（用于按 api_format 选择请求配置）
selectinload(Provider.endpoints),
# 同时加载 models 和 global_model 关系
selectinload(Provider.models).selectinload(Model.global_model),
)
.filter(Provider.is_active == True)
.order_by(Provider.provider_priority.asc())
)
if provider_offset:
provider_query = provider_query.offset(provider_offset)
if provider_limit:
provider_query = provider_query.limit(provider_limit)
return provider_query.all()
async def _check_model_support(
self,
db: Session,
provider: Provider,
model_name: str,
api_format: str | None = None,
is_stream: bool = False,
capability_requirements: dict[str, bool] | None = None,
) -> tuple[bool, str | None, list[str] | None, set[str] | None]:
"""
检查 Provider 是否支持指定模型（可选检查流式支持和能力需求）
模型能力检查在这里进行（而不是在 Key 级别），因为：
- 模型支持的能力是全局的，与具体的 Key 无关
- 如果模型不支持某能力，整个 Provider 的所有 Key 都应该被跳过
仅支持直接匹配 GlobalModel.name（外部请求不接受映射名）
Args:
db: 数据库会话
provider: Provider 对象
model_name: 模型名称（必须是 GlobalModel.name）
is_stream: 是否是流式请求，如果为 True 则同时检查流式支持
capability_requirements: 能力需求（可选），用于检查模型是否支持所需能力
Returns:
(is_supported, skip_reason, supported_capabilities, provider_model_names)
- is_supported: 是否支持
- skip_reason: 跳过原因
- supported_capabilities: 模型支持的能力列表
- provider_model_names: Provider 侧可用的模型名称集合（主名称 + 映射名称，按 api_format 过滤）
"""
# Avoid holding a DB connection while awaiting cache/Redis inside ModelCacheService.
self._release_db_connection_before_await(db)
# 仅接受 GlobalModel.name（不允许映射名）
normalized_name = model_name.strip() if isinstance(model_name, str) else ""
if not normalized_name:
return False, "模型不存在或名称无效", None, None
global_model = await ModelCacheService.get_global_model_by_name(db, normalized_name)
if not global_model or not global_model.is_active:
return False, "模型不存在或已停用", None, None
# 找到 GlobalModel 后，检查当前 Provider 是否支持
is_supported, skip_reason, caps, provider_model_names = (
await self._check_model_support_for_global_model(
db,
provider,
global_model,
model_name,
api_format,
is_stream,
capability_requirements,
)
)
return is_supported, skip_reason, caps, provider_model_names
async def _check_model_support_for_global_model(
self,
db: Session,
provider: Provider,
global_model: GlobalModel,
model_name: str,
api_format: str | None = None,
is_stream: bool = False,
capability_requirements: dict[str, bool] | None = None,
) -> tuple[bool, str | None, list[str] | None, set[str] | None]:
"""
检查 Provider 是否支持指定的 GlobalModel
Args:
db: 数据库会话
provider: Provider 对象
global_model: GlobalModel 对象
model_name: 用户请求的模型名称（用于错误消息）
is_stream: 是否是流式请求
capability_requirements: 能力需求
Returns:
(is_supported, skip_reason, supported_capabilities, provider_model_names)
"""
# 确保 global_model 附加到当前 Session
# 注意：从缓存重建的对象是 transient 状态，不能使用 load=False
# 使用 load=True（默认）允许 SQLAlchemy 正确处理 transient 对象
from sqlalchemy import inspect
insp = inspect(global_model)
if insp.transient or insp.detached:
# transient/detached 对象：使用默认 merge（会查询 DB 检查是否存在）
global_model = db.merge(global_model)
else:
# persistent 对象：已经附加到 session，无需 merge
pass
# 获取模型支持的能力列表
model_supported_capabilities: list[str] = list(global_model.supported_capabilities or [])
# 查询该 Provider 是否有实现这个 GlobalModel
for model in provider.models:
if model.global_model_id == global_model.id and model.is_active:
# 检查流式支持
if is_stream:
supports_streaming = model.get_effective_supports_streaming()
if not supports_streaming:
return False, f"模型 {model_name} 在此 Provider 不支持流式", None, None
# 检查模型是否支持所需的能力（在 Provider 级别检查，而不是 Key 级别）
# 只有当 model_supported_capabilities 非空时才进行检查
# 空列表意味着模型没有配置能力限制，默认支持所有能力
if capability_requirements and model_supported_capabilities:
for cap_name, is_required in capability_requirements.items():
if is_required and cap_name not in model_supported_capabilities:
return (
False,
f"模型 {model_name} 不支持能力: {cap_name}",
list(model_supported_capabilities),
None,
)
provider_model_names: set[str] = {model.provider_model_name}
raw_mappings = model.provider_model_mappings
if isinstance(raw_mappings, list):
for raw in raw_mappings:
if not isinstance(raw, dict):
continue
name = raw.get("name")
if not isinstance(name, str) or not name.strip():
continue
mapping_api_formats = raw.get("api_formats")
if api_format and mapping_api_formats:
# 新模式：endpoint signature（family:kind），按小写 canonical 比较
if isinstance(mapping_api_formats, list):
target = str(api_format).strip().lower()
allowed = {
str(fmt).strip().lower() for fmt in mapping_api_formats if fmt
}
if target not in allowed:
continue
provider_model_names.add(name.strip())
return True, None, list(model_supported_capabilities), provider_model_names
return False, "Provider 未实现此模型", None, None
def _check_key_availability(
self,
key: ProviderAPIKey,
api_format: str | None,
model_name: str,
capability_requirements: dict[str, bool] | None = None,
model_mappings: list[str] | None = None,
candidate_models: set[str] | None = None,
*,
provider_type: str | None = None,
) -> tuple[bool, str | None, str | None]:
"""
检查 API Key 的可用性
注意：模型能力检查已移到 _check_model_support 中进行（Provider 级别），
这里只检查 Key 级别的能力匹配。
Args:
key: API Key 对象
model_name: 模型名称（GlobalModel.name）
capability_requirements: 能力需求（可选）
model_mappings: GlobalModel 的映射列表（用于通配符匹配）
candidate_models: Provider 侧可用的模型名称集合（用于限制映射匹配范围）
Returns:
(is_available, skip_reason, mapping_matched_model)
- is_available: Key 是否可用
- skip_reason: 不可用时的原因
- mapping_matched_model: 通过映射匹配到的模型名（用于实际请求）
"""
# 检查熔断器状态（使用详细状态方法获取更丰富的跳过原因，按 API 格式）
is_available, circuit_reason = health_monitor.get_circuit_breaker_status(
key, api_format=api_format
)
if not is_available:
return False, circuit_reason or "熔断器已打开", None
# 模型权限检查：使用 allowed_models 白名单
# None = 允许所有模型，[] = 拒绝所有模型，["a","b"] = 只允许指定模型
# 支持通配符映射匹配（通过 model_mappings）
try:
is_allowed, mapping_matched_model = check_model_allowed_with_mappings(
model_name=model_name,
allowed_models=key.allowed_models,
model_mappings=model_mappings,
candidate_models=candidate_models,
)
if mapping_matched_model:
logger.debug(
"[Scheduler] Key {}... 模型名匹配: model={} -> {}, allowed_models={}",
key.id[:8],
model_name,
mapping_matched_model,
key.allowed_models,
)
except TimeoutError:
# 正则匹配超时（可能是 ReDoS 攻击或复杂模式）
logger.warning("映射匹配超时: key_id={}, model={}", key.id, model_name)
return False, "映射匹配超时，请简化配置", None
except re.error as e:
# 正则语法错误（配置问题）
logger.warning("映射规则无效: key_id={}, model={}, error={}", key.id, model_name, e)
return False, f"映射规则无效: {str(e)}", None
except Exception as e:
# 其他未知异常
logger.error(
"映射匹配异常: key_id={}, model={}, error={}", key.id, model_name, e, exc_info=True
)
# 异常时保守处理：不允许使用该 Key
return False, "映射匹配失败", None
if not is_allowed:
return (
False,
f"Key 不支持 {model_name}",
None,
)
# Key 级别的能力匹配检查
# 注意：模型级别的能力检查已在 _check_model_support 中完成
# 始终执行检查，即使 capability_requirements 为空
# 因为 check_capability_match 会检查 Key 的 EXCLUSIVE 能力是否被浪费
key_caps: dict[str, bool] = dict(key.capabilities or {})
is_match, skip_reason = check_capability_match(key_caps, capability_requirements)
if not is_match:
return False, skip_reason, None
effective_model_name = mapping_matched_model or model_name
quota_exhausted, quota_reason = is_key_quota_exhausted(
provider_type,
key,
model_name=effective_model_name,
)
if quota_exhausted:
return False, quota_reason, mapping_matched_model
return True, None, mapping_matched_model
async def _build_candidates(
self,
db: Session,
providers: list[Provider],
client_format: str,
model_name: str,
affinity_key: str | None,
model_mappings: list[str] | None = None,
max_candidates: int | None = None,
is_stream: bool = False,
capability_requirements: dict[str, bool] | None = None,
global_conversion_enabled: bool = True,
) -> list[ProviderCandidate]:
"""
构建候选列表
Key 直属 Provider，通过 api_formats 筛选符合端点格式的 Key。
Args:
db: 数据库会话
providers: Provider 列表
client_format: 客户端请求的 API 格式
model_name: 模型名称（GlobalModel.name）
affinity_key: 亲和性标识符（通常为API Key ID）
model_mappings: GlobalModel 的映射列表（用于 Key.allowed_models 通配符匹配）
max_candidates: 最大候选数
is_stream: 是否是流式请求，如果为 True 则过滤不支持流式的 Provider
capability_requirements: 能力需求（可选）
global_conversion_enabled: 格式转换总开关（数据库配置），关闭时禁止任何跨格式转换
Returns:
候选列表
"""
candidates: list[ProviderCandidate] = []
client_format_str = normalize_endpoint_signature(client_format)
client_sig = parse_signature_key(client_format_str)
client_family, client_kind = client_sig.api_family, client_sig.endpoint_kind
# chat/cli 互相可回退（用于同协议族下的端点变体），video/image 等不跨类回退
if client_kind in {EndpointKind.CHAT, EndpointKind.CLI}:
allowed_kinds = {EndpointKind.CHAT, EndpointKind.CLI}
else:
allowed_kinds = {client_kind}
for provider in providers:
logger.debug(
"[Scheduler] Checking provider: {}, endpoints={}",
provider.name,
len(provider.endpoints) if provider.endpoints else 0,
)
# 按端点格式分别判断兼容性与模型/Key 可用性：
# - 同格式端点优先（needs_conversion=False）
# - 跨格式端点次之（needs_conversion=True）
model_support_cache: dict[
str, tuple[bool, str | None, list[str] | None, set[str] | None]
] = {}
exact_candidates: list[ProviderCandidate] = []
convertible_candidates: list[ProviderCandidate] = []
# 使用新架构字段 (api_family, endpoint_kind) 进行预过滤与排序：
# - family/kind 匹配的 endpoint 排在前面（但不做硬过滤，避免破坏格式转换路径）
# - chat/cli 请求允许互相回退（优先同 kind）
# - video 等请求只允许同 kind
endpoints = list(provider.endpoints or [])
allowed_kind_values = {k.value for k in allowed_kinds}
preferred: list[ProviderEndpoint] = []
preferred_other_family: list[ProviderEndpoint] = []
fallback: list[ProviderEndpoint] = []
fallback_other_family: list[ProviderEndpoint] = []
for ep in endpoints:
if not getattr(ep, "is_active", False):
continue
raw_family = getattr(ep, "api_family", None)
raw_kind = getattr(ep, "endpoint_kind", None)
if not isinstance(raw_family, str) or not raw_family.strip():
continue
if not isinstance(raw_kind, str) or not raw_kind.strip():
continue
ep_family = raw_family.strip().lower()
ep_kind = raw_kind.strip().lower()
if allowed_kind_values and ep_kind not in allowed_kind_values:
continue
same_family = ep_family == client_family.value
same_kind = ep_kind == client_kind.value
if same_kind and same_family:
preferred.append(ep)
elif same_kind:
preferred_other_family.append(ep)
elif same_family:
fallback.append(ep)
else:
fallback_other_family.append(ep)
endpoints = (
_sort_endpoints_by_family_priority(preferred)
+ _sort_endpoints_by_family_priority(preferred_other_family)
+ _sort_endpoints_by_family_priority(fallback)
+ _sort_endpoints_by_family_priority(fallback_other_family)
)
for endpoint in endpoints:
logger.debug(
"[Scheduler] Checking endpoint: family={}, kind={}, is_active={}, base_url={}",
getattr(endpoint, "api_family", None),
getattr(endpoint, "endpoint_kind", None),
getattr(endpoint, "is_active", None),
(endpoint.base_url[:50] if endpoint.base_url else "N/A"),
)
if not endpoint.is_active:
logger.debug("[Scheduler] Endpoint skipped: not active")
continue
endpoint_format_str = make_signature_key(
str(getattr(endpoint, "api_family", "")).strip().lower(),
str(getattr(endpoint, "endpoint_kind", "")).strip().lower(),
)
# 计算格式转换开关状态（三层优先级）
#
# 1) 全局开关（数据库配置）关闭 -> 禁止任何跨格式转换
# 2) 全局开关开启 -> 允许跨格式转换
# 3) 提供商覆盖（Provider.enable_format_conversion）开启 -> 强制允许（跳过端点检查）
# 4) 否则 -> 由端点配置 format_acceptance_config 决定是否允许
provider_allows_conversion = getattr(provider, "enable_format_conversion", True)
skip_endpoint_check = global_conversion_enabled or provider_allows_conversion
is_compatible, needs_conversion, _compat_reason = is_format_compatible(
client_format_str,
endpoint_format_str,
getattr(endpoint, "format_acceptance_config", None),
is_stream,
global_conversion_enabled,
skip_endpoint_check=skip_endpoint_check,
)
logger.debug(
"[Scheduler] Format compatibility: client={}, endpoint={}, compatible={}, "
"global={}, provider={}, skip_endpoint={}, reason={}",
client_format_str,
endpoint_format_str,
is_compatible,
global_conversion_enabled,
provider_allows_conversion,
skip_endpoint_check,
_compat_reason,
)
if not is_compatible:
continue
# 检查模型支持（按端点格式过滤 provider_model_mappings）
if endpoint_format_str not in model_support_cache:
model_support_cache[endpoint_format_str] = await self._check_model_support(
db,
provider,
model_name,
api_format=endpoint_format_str,
is_stream=is_stream,
capability_requirements=capability_requirements,
)
supports_model, skip_reason, _model_caps, provider_model_names = (
model_support_cache[endpoint_format_str]
)
logger.debug(
"[Scheduler] Model support: provider={}, model={}, supports={}, reason={}",
provider.name,
model_name,
supports_model,
skip_reason,
)
if not supports_model:
logger.debug(
f"Provider {provider.name} 端点 {endpoint_format_str} 不支持模型 {model_name}: {skip_reason}"
)
continue
# Key 直属 Provider，通过 api_formats 按端点格式筛选
# api_formats=None 视为"全支持"（兼容历史数据）
active_keys = [
key
for key in provider.api_keys
if key.is_active
and (key.api_formats is None or endpoint_format_str in key.api_formats)
]
if not active_keys:
continue
# 检查是否所有 Key 都是 TTL=0（轮换模式）
use_random = all((key.cache_ttl_minutes or 0) == 0 for key in active_keys)
if use_random and len(active_keys) > 1:
logger.debug(
f" Provider {provider.name} 启用 Key 轮换模式 "
f"(endpoint_format={endpoint_format_str}, {len(active_keys)} keys)"
)
keys = self._shuffle_keys_by_internal_priority(
active_keys, affinity_key, use_random
)
for key in keys:
# Key 级别检查（健康度/熔断按 provider_format bucket）
# 传入 provider_model_names 作为 candidate_models，
# 用于检查 Key 的 allowed_models 是否支持 Provider 定义的模型名称
is_available, key_skip_reason, mapping_matched_model = (
self._check_key_availability(
key,
endpoint_format_str,
model_name,
capability_requirements,
model_mappings=model_mappings,
candidate_models=provider_model_names,
provider_type=getattr(provider, "provider_type", None),
)
)
candidate = ProviderCandidate(
provider=provider,
endpoint=endpoint,
key=key,
is_skipped=not is_available,
skip_reason=key_skip_reason,
mapping_matched_model=mapping_matched_model,
needs_conversion=needs_conversion,
provider_api_format=str(endpoint_format_str or ""),
)
if needs_conversion:
convertible_candidates.append(candidate)
else:
exact_candidates.append(candidate)
candidates.extend(exact_candidates)
candidates.extend(convertible_candidates)
# max_candidates 截断应在所有候选收集完成后统一处理，确保优先级排序正确
if max_candidates and len(candidates) > max_candidates:
candidates = candidates[:max_candidates]
return candidates
async def _apply_cache_affinity(
self,
candidates: list[ProviderCandidate],
@@ -1545,241 +1001,6 @@ class CacheAwareScheduler:
self.scheduling_mode = normalized
logger.debug(f"[CacheAwareScheduler] 切换调度模式为: {self.scheduling_mode}")
def _apply_priority_mode_sort(
self,
candidates: list[ProviderCandidate],
db: Session,
affinity_key: str | None = None,
api_format: str | None = None,
) -> list[ProviderCandidate]:
"""
根据优先级模式对候选列表排序（数字越小越优先）
排序规则（受 keep_priority_on_conversion 配置影响）：
1. 如果全局配置 keep_priority_on_conversion=True，所有候选保持原优先级
2. 否则，按 needs_conversion 和 provider.keep_priority_on_conversion 分组：
- 保持优先级的候选（exact 或 provider.keep_priority_on_conversion=True）按原优先级排序
- 需要降级的候选（convertible 且 provider.keep_priority_on_conversion=False）整体排在后面
3. 在同一组内，按优先级模式排序：
- provider: 按 Provider.provider_priority -> Key.internal_priority 排序
- global_key: 按 Key.global_priority_by_format 排序
"""
if not candidates:
return candidates
# 全局配置：如果开启，所有候选保持原优先级
global_keep_priority = SystemConfigService.is_keep_priority_on_conversion(db)
if global_keep_priority:
# 全局开启：不分组，直接按优先级模式排序
if self.priority_mode == self.PRIORITY_MODE_GLOBAL_KEY:
return self._sort_by_global_priority_with_hash(candidates, affinity_key, api_format)
# 提供商优先模式：保持构建时的顺序（已按 provider_priority 排序）
return candidates
# 全局未开启：按是否需要降级分组
# - 不需要降级：exact 候选 或 provider.keep_priority_on_conversion=True 的 convertible 候选
# - 需要降级：convertible 且 provider.keep_priority_on_conversion=False
keep_priority_candidates: list[ProviderCandidate] = []
demote_candidates: list[ProviderCandidate] = []
for c in candidates:
if not c.needs_conversion:
# exact 候选：不需要降级
keep_priority_candidates.append(c)
elif getattr(c.provider, "keep_priority_on_conversion", False):
# convertible 但提供商配置了保持优先级
keep_priority_candidates.append(c)
else:
# convertible 且未配置保持优先级：降级
demote_candidates.append(c)
if self.priority_mode == self.PRIORITY_MODE_GLOBAL_KEY:
# 全局 Key 优先模式：分别对两组排序后合并
sorted_keep = self._sort_by_global_priority_with_hash(
keep_priority_candidates, affinity_key, api_format
)
sorted_demote = self._sort_by_global_priority_with_hash(
demote_candidates, affinity_key, api_format
)
return sorted_keep + sorted_demote
# 提供商优先模式：保持优先级的在前，降级的在后（各组内部顺序已由构建时保证）
return keep_priority_candidates + demote_candidates
def _sort_by_global_priority_with_hash(
self,
candidates: list[ProviderCandidate],
affinity_key: str | None = None,
api_format: str | None = None,
) -> list[ProviderCandidate]:
"""
按 global_priority_by_format 分组排序，同优先级内通过哈希分散实现负载均衡
排序逻辑：
1. 按 global_priority_by_format[api_format] 分组（数字小的优先，NULL 排后面）
2. 同优先级组内，使用 affinity_key 哈希分散
3. 确保同一用户请求稳定选择同一个 Key（缓存亲和性）
"""
def get_priority(candidate: ProviderCandidate) -> int:
"""获取候选的优先级"""
if not candidate.key:
return 999999
priority_by_format = candidate.key.global_priority_by_format or {}
if api_format and api_format in priority_by_format:
return priority_by_format[api_format]
return 999999 # NULL 排在后面
# 按优先级分组
priority_groups: dict[int, list[ProviderCandidate]] = defaultdict(list)
for candidate in candidates:
priority = get_priority(candidate)
priority_groups[priority].append(candidate)
result = []
for priority in sorted(priority_groups.keys()): # 数字小的优先级高
group = priority_groups[priority]
if len(group) > 1 and affinity_key:
# 同优先级内哈希分散负载均衡
scored_candidates = []
for candidate in group:
key_id = candidate.key.id if candidate.key else ""
hash_value = self._affinity_hash(affinity_key, key_id)
scored_candidates.append((hash_value, candidate))
# 按哈希值排序
sorted_group = [c for _, c in sorted(scored_candidates, key=lambda x: x[0])]
result.extend(sorted_group)
else:
# 单个候选或没有 affinity_key，按次要排序条件排序
def secondary_sort(c: ProviderCandidate) -> tuple[int, int, str]:
pp = c.provider.provider_priority
ip = c.key.internal_priority if c.key else None
return (
pp if pp is not None else 999999,
ip if ip is not None else 999999,
c.key.id if c.key else "",
)
result.extend(sorted(group, key=secondary_sort))
return result
def _apply_load_balance(
self, candidates: list[ProviderCandidate], api_format: str | None = None
) -> list[ProviderCandidate]:
"""
负载均衡模式：同优先级内随机轮换
排序逻辑：
1. 按优先级分组（provider_priority, internal_priority 或 global_priority_by_format）
2. 同优先级组内随机打乱
3. 不考虑缓存亲和性
"""
if not candidates:
return candidates
priority_groups: dict[tuple, list[ProviderCandidate]] = defaultdict(list)
# 根据优先级模式选择分组方式
if self.priority_mode == self.PRIORITY_MODE_GLOBAL_KEY:
# 全局 Key 优先模式：按格式特定优先级分组
for candidate in candidates:
priority = 999999
if candidate.key:
priority_by_format = candidate.key.global_priority_by_format or {}
if api_format and api_format in priority_by_format:
priority = priority_by_format[api_format]
priority_groups[(priority,)].append(candidate)
else:
# 提供商优先模式：按 (provider_priority, internal_priority) 分组
for candidate in candidates:
pp = candidate.provider.provider_priority
ip = candidate.key.internal_priority if candidate.key else None
key = (
pp if pp is not None else 999999,
ip if ip is not None else 999999,
)
priority_groups[key].append(candidate)
result: list[ProviderCandidate] = []
for priority in sorted(priority_groups.keys()):
group = priority_groups[priority]
if len(group) > 1:
# 同优先级内随机打乱
shuffled = list(group)
random.shuffle(shuffled)
result.extend(shuffled)
else:
result.extend(group)
return result
def _shuffle_keys_by_internal_priority(
self,
keys: list[ProviderAPIKey],
affinity_key: str | None = None,
use_random: bool = False,
) -> list[ProviderAPIKey]:
"""
对 API Key 按 internal_priority 分组，同优先级内部基于 affinity_key 进行确定性打乱
目的：
- 数字越小越优先使用
- 同优先级 Key 之间实现负载均衡
- 使用 affinity_key 哈希确保同一请求 Key 的请求稳定（避免破坏缓存亲和性）
- 当 use_random=True 时，使用随机排序实现轮换（用于 TTL=0 的场景）
Args:
keys: API Key 列表
affinity_key: 亲和性标识符（通常为 API Key ID，用于确定性打乱）
use_random: 是否使用随机排序（TTL=0 时为 True）
Returns:
排序后的 Key 列表
"""
if not keys:
return []
# 按 internal_priority 分组
priority_groups: dict[int, list[ProviderAPIKey]] = defaultdict(list)
for key in keys:
priority = key.internal_priority if key.internal_priority is not None else 999999
priority_groups[priority].append(key)
# 对每个优先级组内的 Key 进行打乱
result = []
for priority in sorted(priority_groups.keys()): # 数字小的优先级高，排前面
group_keys = priority_groups[priority]
if len(group_keys) > 1:
if use_random:
# TTL=0 模式：使用随机排序实现 Key 轮换
shuffled = list(group_keys)
random.shuffle(shuffled)
result.extend(shuffled)
elif affinity_key:
# 正常模式：使用哈希确定性打乱（保持缓存亲和性）
key_scores = []
for key in group_keys:
hash_value = self._affinity_hash(affinity_key, key.id)
key_scores.append((hash_value, key))
# 按哈希值排序
sorted_group = [key for _, key in sorted(key_scores, key=lambda x: x[0])]
result.extend(sorted_group)
else:
# 没有 affinity_key 时按 ID 排序保持稳定性
result.extend(sorted(group_keys, key=lambda k: k.id))
else:
# 单个 Key 直接添加
result.extend(group_keys)
return result
async def invalidate_cache(
self,
affinity_key: str,

216
src/services/usage/_billing_integration.py Normal file
View File

@@ -0,0 +1,216 @@
from __future__ import annotations
from typing import Any
from src.core.api_format.signature import normalize_signature_key
from src.services.billing.token_normalization import normalize_input_tokens_for_billing
from src.services.usage._recording_helpers import (
build_usage_params,
sanitize_request_metadata,
)
from src.services.usage._types import UsageCostInfo, UsageRecordParams
class UsageBillingIntegrationMixin:
"""计费集成方法 -- 准备用量记录的共享逻辑"""
@classmethod
async def _prepare_usage_record(
cls,
params: UsageRecordParams,
) -> tuple[dict[str, Any], float]:
"""准备用量记录的共享逻辑
此方法提取了 record_usage 和 record_usage_async 的公共处理逻辑：
- 获取费率倍数
- 计算成本
- 构建 Usage 参数
Args:
params: 用量记录参数数据类
Returns:
(usage_params 字典, total_cost 总成本)
"""
# 计费口径以 Provider 为准（优先 endpoint_api_format）
billing_api_format: str | None = None
if params.endpoint_api_format:
try:
billing_api_format = normalize_signature_key(str(params.endpoint_api_format))
except Exception:
billing_api_format = None
if billing_api_format is None and params.api_format:
try:
billing_api_format = normalize_signature_key(str(params.api_format))
except Exception:
billing_api_format = None
input_tokens_for_billing = normalize_input_tokens_for_billing(
billing_api_format,
params.input_tokens,
params.cache_read_input_tokens,
)
# 获取费率倍数和是否免费套餐（传递 api_format 支持按格式配置的倍率）
actual_rate_multiplier, is_free_tier = await cls._get_rate_multiplier_and_free_tier(
params.db, params.provider_api_key_id, params.provider_id, billing_api_format
)
metadata = dict(params.metadata or {})
is_failed_request = params.status_code >= 400 or params.error_message is not None
# Helper: compute billing task_type (billing domain)
billing_task_type = (params.request_type or "").lower()
if billing_task_type not in {"chat", "cli", "video", "image", "audio"}:
billing_task_type = "chat"
# 使用新计费系统计算费用
from src.services.billing.service import BillingService
request_count = 0 if is_failed_request else 1
dims: dict[str, Any] = {
"input_tokens": input_tokens_for_billing,
"output_tokens": params.output_tokens,
"cache_creation_input_tokens": params.cache_creation_input_tokens,
"cache_read_input_tokens": params.cache_read_input_tokens,
"request_count": request_count,
}
if params.cache_ttl_minutes is not None:
dims["cache_ttl_minutes"] = params.cache_ttl_minutes
# If tiered pricing is disabled, force first tier by using tier-key=0.
if not params.use_tiered_pricing:
dims["total_input_context"] = 0
billing = BillingService(params.db)
result = billing.calculate(
task_type=billing_task_type,
model=params.model,
provider_id=params.provider_id or "",
dimensions=dims,
strict_mode=None,
)
snap = result.snapshot
breakdown = snap.cost_breakdown or {}
input_cost = float(breakdown.get("input_cost", 0.0))
output_cost = float(breakdown.get("output_cost", 0.0))
cache_creation_cost = float(breakdown.get("cache_creation_cost", 0.0))
cache_read_cost = float(breakdown.get("cache_read_cost", 0.0))
request_cost = float(breakdown.get("request_cost", 0.0))
cache_cost = cache_creation_cost + cache_read_cost
total_cost = float(snap.total_cost or 0.0)
rv = snap.resolved_variables or {}
def _as_float(v: Any, d: float | None) -> float | None:
try:
if v is None:
return d
return float(v)
except Exception:
return d
input_price = _as_float(rv.get("input_price_per_1m"), 0.0) or 0.0
output_price = _as_float(rv.get("output_price_per_1m"), 0.0) or 0.0
cache_creation_price = _as_float(rv.get("cache_creation_price_per_1m"), None)
cache_read_price = _as_float(rv.get("cache_read_price_per_1m"), None)
request_price = _as_float(rv.get("price_per_request"), None)
# Audit snapshot (pruned later by sanitize_request_metadata)
metadata["billing_snapshot"] = snap.to_dict()
# Best-effort prune metadata to reduce DB/memory pressure.
metadata = sanitize_request_metadata(metadata)
# 构建 Usage 参数
usage_params = build_usage_params(
db=params.db,
user=params.user,
api_key=params.api_key,
provider=params.provider,
model=params.model,
input_tokens=input_tokens_for_billing,
output_tokens=params.output_tokens,
cache_creation_input_tokens=params.cache_creation_input_tokens,
cache_read_input_tokens=params.cache_read_input_tokens,
request_type=params.request_type,
api_format=params.api_format,
endpoint_api_format=params.endpoint_api_format,
has_format_conversion=params.has_format_conversion,
is_stream=params.is_stream,
response_time_ms=params.response_time_ms,
first_byte_time_ms=params.first_byte_time_ms,
status_code=params.status_code,
error_message=params.error_message,
metadata=metadata,
request_headers=params.request_headers,
request_body=params.request_body,
provider_request_headers=params.provider_request_headers,
response_headers=params.response_headers,
client_response_headers=params.client_response_headers,
response_body=params.response_body,
request_id=params.request_id,
provider_id=params.provider_id,
provider_endpoint_id=params.provider_endpoint_id,
provider_api_key_id=params.provider_api_key_id,
status=params.status,
target_model=params.target_model,
cost=UsageCostInfo(
input_cost=input_cost,
output_cost=output_cost,
cache_creation_cost=cache_creation_cost,
cache_read_cost=cache_read_cost,
cache_cost=cache_cost,
request_cost=request_cost,
total_cost=total_cost,
input_price=input_price,
output_price=output_price,
cache_creation_price=cache_creation_price,
cache_read_price=cache_read_price,
request_price=request_price,
actual_rate_multiplier=actual_rate_multiplier,
is_free_tier=is_free_tier,
),
)
return usage_params, total_cost
@classmethod
async def _prepare_usage_records_batch(
cls,
params_list: list[UsageRecordParams],
) -> list[tuple[dict[str, Any], float, Exception | None]]:
"""批量并行准备用量记录（性能优化）
并行调用 _prepare_usage_record，提高批量处理效率。
Args:
params_list: 用量记录参数列表
Returns:
列表，每项为 (usage_params, total_cost, exception)
如果处理成功，exception 为 None
"""
import asyncio
async def prepare_single(
params: UsageRecordParams,
) -> tuple[dict[str, Any], float, Exception | None]:
try:
usage_params, total_cost = await cls._prepare_usage_record(params)
return (usage_params, total_cost, None)
except Exception as e:
return ({}, 0.0, e)
if not params_list:
return []
# 避免一次性创建过多 task（并且 _prepare_usage_record 内部也可能包含并行调用）
# 这里采用分批 gather 来限制并发量。
chunk_size = 50
results: list[tuple[dict[str, Any], float, Exception | None]] = []
for i in range(0, len(params_list), chunk_size):
chunk = params_list[i : i + chunk_size]
chunk_results = await asyncio.gather(*(prepare_single(p) for p in chunk))
results.extend(chunk_results)
return results

310
src/services/usage/_recording_helpers.py Normal file
View File

@@ -0,0 +1,310 @@
from __future__ import annotations
import json
from typing import Any
from sqlalchemy.orm import Session
from src.models.database import ApiKey, Usage, User
from src.services.system.config import SystemConfigService
from src.services.usage._types import UsageCostInfo
from src.services.usage.error_classifier import classify_error
# Metadata pruning configuration (ordered by priority - drop first to last)
METADATA_PRUNE_KEYS: tuple[str, ...] = (
"raw_response_ref",
"poll_raw_response",
"trace",
"debug",
"dimensions",
"provider_response_headers",
"client_response_headers",
)
# Keys to preserve even under aggressive pruning
METADATA_KEEP_KEYS: frozenset[str] = frozenset(
{
"billing_snapshot",
"billing_updated_at",
"perf",
"_metadata_truncated",
}
)
def build_usage_params(
*,
db: Session,
user: User | None,
api_key: ApiKey | None,
provider: str,
model: str,
input_tokens: int,
output_tokens: int,
cache_creation_input_tokens: int,
cache_read_input_tokens: int,
request_type: str,
api_format: str | None,
endpoint_api_format: str | None,
has_format_conversion: bool,
is_stream: bool,
response_time_ms: int | None,
first_byte_time_ms: int | None,
status_code: int,
error_message: str | None,
metadata: dict[str, Any] | None,
request_headers: dict[str, Any] | None,
request_body: Any | None,
provider_request_headers: dict[str, Any] | None,
response_headers: dict[str, Any] | None,
client_response_headers: dict[str, Any] | None,
response_body: Any | None,
request_id: str,
provider_id: str | None,
provider_endpoint_id: str | None,
provider_api_key_id: str | None,
status: str,
target_model: str | None,
cost: UsageCostInfo,
) -> dict[str, Any]:
"""构建 Usage 记录的参数字典（内部方法，避免代码重复）"""
# 展开成本信息
input_cost = cost.input_cost
output_cost = cost.output_cost
cache_creation_cost = cost.cache_creation_cost
cache_read_cost = cost.cache_read_cost
cache_cost = cost.cache_cost
request_cost = cost.request_cost
total_cost = cost.total_cost
input_price = cost.input_price
output_price = cost.output_price
cache_creation_price = cost.cache_creation_price
cache_read_price = cost.cache_read_price
request_price = cost.request_price
actual_rate_multiplier = cost.actual_rate_multiplier
is_free_tier = cost.is_free_tier
# 根据配置决定是否记录请求详情
should_log_headers = SystemConfigService.should_log_headers(db)
should_log_body = SystemConfigService.should_log_body(db)
# 处理请求头（可能需要脱敏）
processed_request_headers = None
if should_log_headers and request_headers:
processed_request_headers = SystemConfigService.mask_sensitive_headers(db, request_headers)
# 处理提供商请求头（可能需要脱敏）
processed_provider_request_headers = None
if should_log_headers and provider_request_headers:
processed_provider_request_headers = SystemConfigService.mask_sensitive_headers(
db, provider_request_headers
)
# 处理请求体和响应体（可能需要截断）
processed_request_body = None
processed_response_body = None
if should_log_body:
if request_body:
processed_request_body = SystemConfigService.truncate_body(
db, request_body, is_request=True
)
if response_body:
processed_response_body = SystemConfigService.truncate_body(
db, response_body, is_request=False
)
# 处理响应头
processed_response_headers = None
if should_log_headers and response_headers:
processed_response_headers = SystemConfigService.mask_sensitive_headers(
db, response_headers
)
# 处理返回给客户端的响应头
processed_client_response_headers = None
if should_log_headers and client_response_headers:
processed_client_response_headers = SystemConfigService.mask_sensitive_headers(
db, client_response_headers
)
# 计算真实成本（表面成本 * 倍率），免费套餐实际费用为 0
if is_free_tier:
actual_input_cost = 0.0
actual_output_cost = 0.0
actual_cache_creation_cost = 0.0
actual_cache_read_cost = 0.0
actual_request_cost = 0.0
actual_total_cost = 0.0
else:
actual_input_cost = input_cost * actual_rate_multiplier
actual_output_cost = output_cost * actual_rate_multiplier
actual_cache_creation_cost = cache_creation_cost * actual_rate_multiplier
actual_cache_read_cost = cache_read_cost * actual_rate_multiplier
actual_request_cost = request_cost * actual_rate_multiplier
actual_total_cost = total_cost * actual_rate_multiplier
error_category = None
if status_code >= 400 or error_message or status in {"failed", "cancelled"}:
error_category = classify_error(status_code, error_message, status).value
return {
"user_id": user.id if user else None,
"api_key_id": api_key.id if api_key else None,
"request_id": request_id,
"provider_name": provider,
"model": model,
"target_model": target_model,
"provider_id": provider_id,
"provider_endpoint_id": provider_endpoint_id,
"provider_api_key_id": provider_api_key_id,
"input_tokens": input_tokens,
"output_tokens": output_tokens,
"total_tokens": input_tokens + output_tokens,
"cache_creation_input_tokens": cache_creation_input_tokens,
"cache_read_input_tokens": cache_read_input_tokens,
"input_cost_usd": input_cost,
"output_cost_usd": output_cost,
"cache_cost_usd": cache_cost,
"cache_creation_cost_usd": cache_creation_cost,
"cache_read_cost_usd": cache_read_cost,
"request_cost_usd": request_cost,
"total_cost_usd": total_cost,
"actual_input_cost_usd": actual_input_cost,
"actual_output_cost_usd": actual_output_cost,
"actual_cache_creation_cost_usd": actual_cache_creation_cost,
"actual_cache_read_cost_usd": actual_cache_read_cost,
"actual_request_cost_usd": actual_request_cost,
"actual_total_cost_usd": actual_total_cost,
"rate_multiplier": actual_rate_multiplier,
"input_price_per_1m": input_price,
"output_price_per_1m": output_price,
"cache_creation_price_per_1m": cache_creation_price,
"cache_read_price_per_1m": cache_read_price,
"price_per_request": request_price,
"request_type": request_type,
"api_format": api_format,
"endpoint_api_format": endpoint_api_format,
"has_format_conversion": has_format_conversion,
"is_stream": is_stream,
"status_code": status_code,
"error_message": error_message,
"error_category": error_category,
"response_time_ms": response_time_ms,
"first_byte_time_ms": first_byte_time_ms,
"status": status,
"request_metadata": metadata,
"request_headers": processed_request_headers,
"request_body": processed_request_body,
"provider_request_headers": processed_provider_request_headers,
"response_headers": processed_response_headers,
"client_response_headers": processed_client_response_headers,
"response_body": processed_response_body,
}
def update_existing_usage(
existing_usage: Usage,
usage_params: dict[str, Any],
target_model: str | None,
) -> None:
"""更新已存在的 Usage 记录（内部方法）"""
# 更新关键字段
existing_usage.provider_name = usage_params["provider_name"]
existing_usage.model = usage_params["model"]
existing_usage.request_type = usage_params["request_type"]
existing_usage.api_format = usage_params["api_format"]
existing_usage.endpoint_api_format = usage_params["endpoint_api_format"]
existing_usage.has_format_conversion = usage_params["has_format_conversion"]
existing_usage.is_stream = usage_params["is_stream"]
existing_usage.status = usage_params["status"]
existing_usage.status_code = usage_params["status_code"]
existing_usage.error_message = usage_params["error_message"]
existing_usage.error_category = usage_params.get("error_category")
existing_usage.response_time_ms = usage_params["response_time_ms"]
existing_usage.first_byte_time_ms = usage_params["first_byte_time_ms"]
# 更新请求头和请求体（如果有新值）
if usage_params["request_headers"] is not None:
existing_usage.request_headers = usage_params["request_headers"]
if usage_params["request_body"] is not None:
existing_usage.request_body = usage_params["request_body"]
if usage_params["provider_request_headers"] is not None:
existing_usage.provider_request_headers = usage_params["provider_request_headers"]
existing_usage.response_body = usage_params["response_body"]
existing_usage.response_headers = usage_params["response_headers"]
existing_usage.client_response_headers = usage_params["client_response_headers"]
# 更新 token 和费用信息
existing_usage.input_tokens = usage_params["input_tokens"]
existing_usage.output_tokens = usage_params["output_tokens"]
existing_usage.total_tokens = usage_params["total_tokens"]
existing_usage.cache_creation_input_tokens = usage_params["cache_creation_input_tokens"]
existing_usage.cache_read_input_tokens = usage_params["cache_read_input_tokens"]
existing_usage.input_cost_usd = usage_params["input_cost_usd"]
existing_usage.output_cost_usd = usage_params["output_cost_usd"]
existing_usage.cache_cost_usd = usage_params["cache_cost_usd"]
existing_usage.cache_creation_cost_usd = usage_params["cache_creation_cost_usd"]
existing_usage.cache_read_cost_usd = usage_params["cache_read_cost_usd"]
existing_usage.request_cost_usd = usage_params["request_cost_usd"]
existing_usage.total_cost_usd = usage_params["total_cost_usd"]
existing_usage.actual_input_cost_usd = usage_params["actual_input_cost_usd"]
existing_usage.actual_output_cost_usd = usage_params["actual_output_cost_usd"]
existing_usage.actual_cache_creation_cost_usd = usage_params["actual_cache_creation_cost_usd"]
existing_usage.actual_cache_read_cost_usd = usage_params["actual_cache_read_cost_usd"]
existing_usage.actual_request_cost_usd = usage_params["actual_request_cost_usd"]
existing_usage.actual_total_cost_usd = usage_params["actual_total_cost_usd"]
existing_usage.rate_multiplier = usage_params["rate_multiplier"]
# 更新 Provider 侧追踪信息
existing_usage.provider_id = usage_params["provider_id"]
existing_usage.provider_endpoint_id = usage_params["provider_endpoint_id"]
existing_usage.provider_api_key_id = usage_params["provider_api_key_id"]
# 更新元数据（如 billing_snapshot/dimensions 等）
if usage_params.get("request_metadata") is not None:
existing_usage.request_metadata = usage_params["request_metadata"]
# 更新模型映射信息
if target_model is not None:
existing_usage.target_model = target_model
def sanitize_request_metadata(metadata: dict[str, Any]) -> dict[str, Any]:
"""
Best-effort metadata pruning to reduce DB/CPU/memory pressure.
This is called right before persisting Usage rows (or updating request_metadata).
Pruning order is defined by `METADATA_PRUNE_KEYS` (first key is dropped first).
"""
if not isinstance(metadata, dict) or not metadata:
return {}
from src.config.settings import config
# Enforce global metadata size limit (best-effort)
max_bytes = int(getattr(config, "usage_metadata_max_bytes", 0) or 0)
if max_bytes <= 0:
return metadata
def _size(d: dict[str, Any]) -> int:
try:
return len(json.dumps(d, ensure_ascii=False, default=str))
except Exception:
return len(str(d))
if _size(metadata) <= max_bytes:
return metadata
# Progressive pruning (configurable order)
metadata["_metadata_truncated"] = True
for k in METADATA_PRUNE_KEYS:
if k in metadata:
metadata.pop(k, None)
if _size(metadata) <= max_bytes:
return metadata
# Fallback: keep only billing-related metadata
reduced = {k: metadata.get(k) for k in METADATA_KEEP_KEYS if k in metadata}
return reduced

522
src/services/usage/recording.py
View File

@@ -1,217 +1,39 @@
from __future__ import annotations
import json
import uuid
from datetime import datetime, timezone
from typing import Any
from sqlalchemy.orm import Session
from src.core.api_format.signature import normalize_signature_key
from src.core.logger import logger
from src.models.database import ApiKey, Provider, Usage, User
from src.services.billing.token_normalization import normalize_input_tokens_for_billing
from src.services.system.config import SystemConfigService
from src.services.usage._billing_integration import UsageBillingIntegrationMixin
from src.services.usage._recording_helpers import (
METADATA_KEEP_KEYS,
METADATA_PRUNE_KEYS,
build_usage_params,
sanitize_request_metadata,
update_existing_usage,
)
from src.services.usage._types import UsageCostInfo, UsageRecordParams
from src.services.usage.error_classifier import classify_error
class UsageRecordingMixin:
class UsageRecordingMixin(UsageBillingIntegrationMixin):
"""记录用量相关方法"""
# Metadata pruning configuration (ordered by priority - drop first to last)
_METADATA_PRUNE_KEYS: tuple[str, ...] = (
"raw_response_ref",
"poll_raw_response",
"trace",
"debug",
"dimensions",
"provider_response_headers",
"client_response_headers",
)
# Metadata pruning configuration -- re-export from helpers for backward compatibility
_METADATA_PRUNE_KEYS: tuple[str, ...] = METADATA_PRUNE_KEYS
_METADATA_KEEP_KEYS: frozenset[str] = METADATA_KEEP_KEYS
# Keys to preserve even under aggressive pruning
_METADATA_KEEP_KEYS: frozenset[str] = frozenset(
{
"billing_snapshot",
"billing_updated_at",
"perf",
"_metadata_truncated",
}
)
# ------------------------------------------------------------------
# Backward-compatible thin wrappers
# ------------------------------------------------------------------
@staticmethod
def _build_usage_params(
*,
db: Session,
user: User | None,
api_key: ApiKey | None,
provider: str,
model: str,
input_tokens: int,
output_tokens: int,
cache_creation_input_tokens: int,
cache_read_input_tokens: int,
request_type: str,
api_format: str | None,
endpoint_api_format: str | None,
has_format_conversion: bool,
is_stream: bool,
response_time_ms: int | None,
first_byte_time_ms: int | None,
status_code: int,
error_message: str | None,
metadata: dict[str, Any] | None,
request_headers: dict[str, Any] | None,
request_body: Any | None,
provider_request_headers: dict[str, Any] | None,
response_headers: dict[str, Any] | None,
client_response_headers: dict[str, Any] | None,
response_body: Any | None,
request_id: str,
provider_id: str | None,
provider_endpoint_id: str | None,
provider_api_key_id: str | None,
status: str,
target_model: str | None,
cost: UsageCostInfo,
) -> dict[str, Any]:
"""构建 Usage 记录的参数字典（内部方法，避免代码重复）"""
# 展开成本信息
input_cost = cost.input_cost
output_cost = cost.output_cost
cache_creation_cost = cost.cache_creation_cost
cache_read_cost = cost.cache_read_cost
cache_cost = cost.cache_cost
request_cost = cost.request_cost
total_cost = cost.total_cost
input_price = cost.input_price
output_price = cost.output_price
cache_creation_price = cost.cache_creation_price
cache_read_price = cost.cache_read_price
request_price = cost.request_price
actual_rate_multiplier = cost.actual_rate_multiplier
is_free_tier = cost.is_free_tier
# 根据配置决定是否记录请求详情
should_log_headers = SystemConfigService.should_log_headers(db)
should_log_body = SystemConfigService.should_log_body(db)
# 处理请求头（可能需要脱敏）
processed_request_headers = None
if should_log_headers and request_headers:
processed_request_headers = SystemConfigService.mask_sensitive_headers(
db, request_headers
)
# 处理提供商请求头（可能需要脱敏）
processed_provider_request_headers = None
if should_log_headers and provider_request_headers:
processed_provider_request_headers = SystemConfigService.mask_sensitive_headers(
db, provider_request_headers
)
# 处理请求体和响应体（可能需要截断）
processed_request_body = None
processed_response_body = None
if should_log_body:
if request_body:
processed_request_body = SystemConfigService.truncate_body(
db, request_body, is_request=True
)
if response_body:
processed_response_body = SystemConfigService.truncate_body(
db, response_body, is_request=False
)
# 处理响应头
processed_response_headers = None
if should_log_headers and response_headers:
processed_response_headers = SystemConfigService.mask_sensitive_headers(
db, response_headers
)
# 处理返回给客户端的响应头
processed_client_response_headers = None
if should_log_headers and client_response_headers:
processed_client_response_headers = SystemConfigService.mask_sensitive_headers(
db, client_response_headers
)
# 计算真实成本（表面成本 * 倍率），免费套餐实际费用为 0
if is_free_tier:
actual_input_cost = 0.0
actual_output_cost = 0.0
actual_cache_creation_cost = 0.0
actual_cache_read_cost = 0.0
actual_request_cost = 0.0
actual_total_cost = 0.0
else:
actual_input_cost = input_cost * actual_rate_multiplier
actual_output_cost = output_cost * actual_rate_multiplier
actual_cache_creation_cost = cache_creation_cost * actual_rate_multiplier
actual_cache_read_cost = cache_read_cost * actual_rate_multiplier
actual_request_cost = request_cost * actual_rate_multiplier
actual_total_cost = total_cost * actual_rate_multiplier
error_category = None
if status_code >= 400 or error_message or status in {"failed", "cancelled"}:
error_category = classify_error(status_code, error_message, status).value
return {
"user_id": user.id if user else None,
"api_key_id": api_key.id if api_key else None,
"request_id": request_id,
"provider_name": provider,
"model": model,
"target_model": target_model,
"provider_id": provider_id,
"provider_endpoint_id": provider_endpoint_id,
"provider_api_key_id": provider_api_key_id,
"input_tokens": input_tokens,
"output_tokens": output_tokens,
"total_tokens": input_tokens + output_tokens,
"cache_creation_input_tokens": cache_creation_input_tokens,
"cache_read_input_tokens": cache_read_input_tokens,
"input_cost_usd": input_cost,
"output_cost_usd": output_cost,
"cache_cost_usd": cache_cost,
"cache_creation_cost_usd": cache_creation_cost,
"cache_read_cost_usd": cache_read_cost,
"request_cost_usd": request_cost,
"total_cost_usd": total_cost,
"actual_input_cost_usd": actual_input_cost,
"actual_output_cost_usd": actual_output_cost,
"actual_cache_creation_cost_usd": actual_cache_creation_cost,
"actual_cache_read_cost_usd": actual_cache_read_cost,
"actual_request_cost_usd": actual_request_cost,
"actual_total_cost_usd": actual_total_cost,
"rate_multiplier": actual_rate_multiplier,
"input_price_per_1m": input_price,
"output_price_per_1m": output_price,
"cache_creation_price_per_1m": cache_creation_price,
"cache_read_price_per_1m": cache_read_price,
"price_per_request": request_price,
"request_type": request_type,
"api_format": api_format,
"endpoint_api_format": endpoint_api_format,
"has_format_conversion": has_format_conversion,
"is_stream": is_stream,
"status_code": status_code,
"error_message": error_message,
"error_category": error_category,
"response_time_ms": response_time_ms,
"first_byte_time_ms": first_byte_time_ms,
"status": status,
"request_metadata": metadata,
"request_headers": processed_request_headers,
"request_body": processed_request_body,
"provider_request_headers": processed_provider_request_headers,
"response_headers": processed_response_headers,
"client_response_headers": processed_client_response_headers,
"response_body": processed_response_body,
}
def _build_usage_params(**kwargs: Any) -> dict[str, Any]:
"""构建 Usage 记录的参数字典（委托到模块级函数）"""
return build_usage_params(**kwargs)
@staticmethod
def _update_existing_usage(
@@ -219,309 +41,17 @@ class UsageRecordingMixin:
usage_params: dict[str, Any],
target_model: str | None,
) -> None:
"""更新已存在的 Usage 记录（内部方法）"""
# 更新关键字段
existing_usage.provider_name = usage_params["provider_name"]
existing_usage.model = usage_params["model"]
existing_usage.request_type = usage_params["request_type"]
existing_usage.api_format = usage_params["api_format"]
existing_usage.endpoint_api_format = usage_params["endpoint_api_format"]
existing_usage.has_format_conversion = usage_params["has_format_conversion"]
existing_usage.is_stream = usage_params["is_stream"]
existing_usage.status = usage_params["status"]
existing_usage.status_code = usage_params["status_code"]
existing_usage.error_message = usage_params["error_message"]
existing_usage.error_category = usage_params.get("error_category")
existing_usage.response_time_ms = usage_params["response_time_ms"]
existing_usage.first_byte_time_ms = usage_params["first_byte_time_ms"]
# 更新请求头和请求体（如果有新值）
if usage_params["request_headers"] is not None:
existing_usage.request_headers = usage_params["request_headers"]
if usage_params["request_body"] is not None:
existing_usage.request_body = usage_params["request_body"]
if usage_params["provider_request_headers"] is not None:
existing_usage.provider_request_headers = usage_params["provider_request_headers"]
existing_usage.response_body = usage_params["response_body"]
existing_usage.response_headers = usage_params["response_headers"]
existing_usage.client_response_headers = usage_params["client_response_headers"]
# 更新 token 和费用信息
existing_usage.input_tokens = usage_params["input_tokens"]
existing_usage.output_tokens = usage_params["output_tokens"]
existing_usage.total_tokens = usage_params["total_tokens"]
existing_usage.cache_creation_input_tokens = usage_params["cache_creation_input_tokens"]
existing_usage.cache_read_input_tokens = usage_params["cache_read_input_tokens"]
existing_usage.input_cost_usd = usage_params["input_cost_usd"]
existing_usage.output_cost_usd = usage_params["output_cost_usd"]
existing_usage.cache_cost_usd = usage_params["cache_cost_usd"]
existing_usage.cache_creation_cost_usd = usage_params["cache_creation_cost_usd"]
existing_usage.cache_read_cost_usd = usage_params["cache_read_cost_usd"]
existing_usage.request_cost_usd = usage_params["request_cost_usd"]
existing_usage.total_cost_usd = usage_params["total_cost_usd"]
existing_usage.actual_input_cost_usd = usage_params["actual_input_cost_usd"]
existing_usage.actual_output_cost_usd = usage_params["actual_output_cost_usd"]
existing_usage.actual_cache_creation_cost_usd = usage_params[
"actual_cache_creation_cost_usd"
]
existing_usage.actual_cache_read_cost_usd = usage_params["actual_cache_read_cost_usd"]
existing_usage.actual_request_cost_usd = usage_params["actual_request_cost_usd"]
existing_usage.actual_total_cost_usd = usage_params["actual_total_cost_usd"]
existing_usage.rate_multiplier = usage_params["rate_multiplier"]
# 更新 Provider 侧追踪信息
existing_usage.provider_id = usage_params["provider_id"]
existing_usage.provider_endpoint_id = usage_params["provider_endpoint_id"]
existing_usage.provider_api_key_id = usage_params["provider_api_key_id"]
# 更新元数据（如 billing_snapshot/dimensions 等）
if usage_params.get("request_metadata") is not None:
existing_usage.request_metadata = usage_params["request_metadata"]
# 更新模型映射信息
if target_model is not None:
existing_usage.target_model = target_model
"""更新已存在的 Usage 记录（委托到模块级函数）"""
update_existing_usage(existing_usage, usage_params, target_model)
@classmethod
def _sanitize_request_metadata(cls, metadata: dict[str, Any]) -> dict[str, Any]:
"""
Best-effort metadata pruning to reduce DB/CPU/memory pressure.
"""元数据清理（委托到模块级函数）"""
return sanitize_request_metadata(metadata)
This is called right before persisting Usage rows (or updating request_metadata).
Pruning order is defined by `_METADATA_PRUNE_KEYS` (first key is dropped first).
"""
if not isinstance(metadata, dict) or not metadata:
return {}
from src.config.settings import config
# Enforce global metadata size limit (best-effort)
max_bytes = int(getattr(config, "usage_metadata_max_bytes", 0) or 0)
if max_bytes <= 0:
return metadata
def _size(d: dict[str, Any]) -> int:
try:
return len(json.dumps(d, ensure_ascii=False, default=str))
except Exception:
return len(str(d))
if _size(metadata) <= max_bytes:
return metadata
# Progressive pruning (configurable order)
metadata["_metadata_truncated"] = True
for k in cls._METADATA_PRUNE_KEYS:
if k in metadata:
metadata.pop(k, None)
if _size(metadata) <= max_bytes:
return metadata
# Fallback: keep only billing-related metadata
reduced = {k: metadata.get(k) for k in cls._METADATA_KEEP_KEYS if k in metadata}
return reduced
@classmethod
async def _prepare_usage_record(
cls,
params: UsageRecordParams,
) -> tuple[dict[str, Any], float]:
"""准备用量记录的共享逻辑
此方法提取了 record_usage 和 record_usage_async 的公共处理逻辑：
- 获取费率倍数
- 计算成本
- 构建 Usage 参数
Args:
params: 用量记录参数数据类
Returns:
(usage_params 字典, total_cost 总成本)
"""
# 计费口径以 Provider 为准（优先 endpoint_api_format）
billing_api_format: str | None = None
if params.endpoint_api_format:
try:
billing_api_format = normalize_signature_key(str(params.endpoint_api_format))
except Exception:
billing_api_format = None
if billing_api_format is None and params.api_format:
try:
billing_api_format = normalize_signature_key(str(params.api_format))
except Exception:
billing_api_format = None
input_tokens_for_billing = normalize_input_tokens_for_billing(
billing_api_format,
params.input_tokens,
params.cache_read_input_tokens,
)
# 获取费率倍数和是否免费套餐（传递 api_format 支持按格式配置的倍率）
actual_rate_multiplier, is_free_tier = await cls._get_rate_multiplier_and_free_tier(
params.db, params.provider_api_key_id, params.provider_id, billing_api_format
)
metadata = dict(params.metadata or {})
is_failed_request = params.status_code >= 400 or params.error_message is not None
# Helper: compute billing task_type (billing domain)
billing_task_type = (params.request_type or "").lower()
if billing_task_type not in {"chat", "cli", "video", "image", "audio"}:
billing_task_type = "chat"
# 使用新计费系统计算费用
from src.services.billing.service import BillingService
request_count = 0 if is_failed_request else 1
dims: dict[str, Any] = {
"input_tokens": input_tokens_for_billing,
"output_tokens": params.output_tokens,
"cache_creation_input_tokens": params.cache_creation_input_tokens,
"cache_read_input_tokens": params.cache_read_input_tokens,
"request_count": request_count,
}
if params.cache_ttl_minutes is not None:
dims["cache_ttl_minutes"] = params.cache_ttl_minutes
# If tiered pricing is disabled, force first tier by using tier-key=0.
if not params.use_tiered_pricing:
dims["total_input_context"] = 0
billing = BillingService(params.db)
result = billing.calculate(
task_type=billing_task_type,
model=params.model,
provider_id=params.provider_id or "",
dimensions=dims,
strict_mode=None,
)
snap = result.snapshot
breakdown = snap.cost_breakdown or {}
input_cost = float(breakdown.get("input_cost", 0.0))
output_cost = float(breakdown.get("output_cost", 0.0))
cache_creation_cost = float(breakdown.get("cache_creation_cost", 0.0))
cache_read_cost = float(breakdown.get("cache_read_cost", 0.0))
request_cost = float(breakdown.get("request_cost", 0.0))
cache_cost = cache_creation_cost + cache_read_cost
total_cost = float(snap.total_cost or 0.0)
rv = snap.resolved_variables or {}
def _as_float(v: Any, d: float | None) -> float | None:
try:
if v is None:
return d
return float(v)
except Exception:
return d
input_price = _as_float(rv.get("input_price_per_1m"), 0.0) or 0.0
output_price = _as_float(rv.get("output_price_per_1m"), 0.0) or 0.0
cache_creation_price = _as_float(rv.get("cache_creation_price_per_1m"), None)
cache_read_price = _as_float(rv.get("cache_read_price_per_1m"), None)
request_price = _as_float(rv.get("price_per_request"), None)
# Audit snapshot (pruned later by _sanitize_request_metadata)
metadata["billing_snapshot"] = snap.to_dict()
# Best-effort prune metadata to reduce DB/memory pressure.
metadata = cls._sanitize_request_metadata(metadata)
# 构建 Usage 参数
usage_params = cls._build_usage_params(
db=params.db,
user=params.user,
api_key=params.api_key,
provider=params.provider,
model=params.model,
input_tokens=input_tokens_for_billing,
output_tokens=params.output_tokens,
cache_creation_input_tokens=params.cache_creation_input_tokens,
cache_read_input_tokens=params.cache_read_input_tokens,
request_type=params.request_type,
api_format=params.api_format,
endpoint_api_format=params.endpoint_api_format,
has_format_conversion=params.has_format_conversion,
is_stream=params.is_stream,
response_time_ms=params.response_time_ms,
first_byte_time_ms=params.first_byte_time_ms,
status_code=params.status_code,
error_message=params.error_message,
metadata=metadata,
request_headers=params.request_headers,
request_body=params.request_body,
provider_request_headers=params.provider_request_headers,
response_headers=params.response_headers,
client_response_headers=params.client_response_headers,
response_body=params.response_body,
request_id=params.request_id,
provider_id=params.provider_id,
provider_endpoint_id=params.provider_endpoint_id,
provider_api_key_id=params.provider_api_key_id,
status=params.status,
target_model=params.target_model,
cost=UsageCostInfo(
input_cost=input_cost,
output_cost=output_cost,
cache_creation_cost=cache_creation_cost,
cache_read_cost=cache_read_cost,
cache_cost=cache_cost,
request_cost=request_cost,
total_cost=total_cost,
input_price=input_price,
output_price=output_price,
cache_creation_price=cache_creation_price,
cache_read_price=cache_read_price,
request_price=request_price,
actual_rate_multiplier=actual_rate_multiplier,
is_free_tier=is_free_tier,
),
)
return usage_params, total_cost
@classmethod
async def _prepare_usage_records_batch(
cls,
params_list: list[UsageRecordParams],
) -> list[tuple[dict[str, Any], float, Exception | None]]:
"""批量并行准备用量记录（性能优化）
并行调用 _prepare_usage_record，提高批量处理效率。
Args:
params_list: 用量记录参数列表
Returns:
列表，每项为 (usage_params, total_cost, exception)
如果处理成功，exception 为 None
"""
import asyncio
async def prepare_single(
params: UsageRecordParams,
) -> tuple[dict[str, Any], float, Exception | None]:
try:
usage_params, total_cost = await cls._prepare_usage_record(params)
return (usage_params, total_cost, None)
except Exception as e:
return ({}, 0.0, e)
if not params_list:
return []
# 避免一次性创建过多 task（并且 _prepare_usage_record 内部也可能包含并行调用）
# 这里采用分批 gather 来限制并发量。
chunk_size = 50
results: list[tuple[dict[str, Any], float, Exception | None]] = []
for i in range(0, len(params_list), chunk_size):
chunk = params_list[i : i + chunk_size]
chunk_results = await asyncio.gather(*(prepare_single(p) for p in chunk))
results.extend(chunk_results)
return results
# ------------------------------------------------------------------
# Recording methods
# ------------------------------------------------------------------
@classmethod
async def record_usage_async(
@@ -891,7 +421,7 @@ class UsageRecordingMixin:
)
total_cost = float(total_cost_usd)
usage_params = cls._build_usage_params(
usage_params = build_usage_params(
db=db,
user=user,
api_key=api_key,