feat: 多项功能优化与问题修复

- 添加高频轮询端点日志抑制,减少 debug 日志噪音
- Gemini 格式转换支持 responseModalities、thinkingConfig 透传和图片生成输出
- OpenAI 格式转换支持流式图片内容块,区分 URL 引用和 base64 内嵌图片
- Provider 余额缓存认证失败时使用短 TTL,避免前端无限加载中
- Usage 服务支持更新 pending/streaming 记录,处理重复 request_id 冲突
- Usage 超时检测增强,避免错误标记已完成请求为超时
This commit is contained in:
fawney19
2026-01-28 23:07:58 +08:00
parent 91986baa57
commit 6b34a7e3f5
9 changed files with 540 additions and 57 deletions

View File

@@ -45,6 +45,9 @@ class ApiRequestContext:
extra: Dict[str, Any] = field(default_factory=dict)
audit_metadata: Dict[str, Any] = field(default_factory=dict)
# 高频轮询端点日志抑制标志
quiet_logging: bool = False
def ensure_json_body(self) -> Dict[str, Any]:
"""确保请求体已解析为JSON并返回。"""
if self.json_body is not None:

View File

@@ -22,6 +22,14 @@ if TYPE_CHECKING:
from .adapter import ApiAdapter, ApiMode
from .context import ApiRequestContext
# 高频轮询端点,抑制其 debug 日志以减少噪音
QUIET_POLLING_PATHS: set[str] = {
"/api/admin/usage/active",
"/api/admin/usage/records",
"/api/admin/usage/stats",
"/api/admin/usage/aggregation/stats",
"/api/admin/health/status",
}
class ApiRequestPipeline:
@@ -47,9 +55,12 @@ class ApiRequestPipeline:
api_format_hint: Optional[str] = None,
path_params: Optional[dict[str, Any]] = None,
):
logger.debug(f"[Pipeline] START | path={http_request.url.path}")
logger.debug(f"[Pipeline] Running with mode={mode}, adapter={adapter.__class__.__name__}, "
f"adapter.mode={adapter.mode}, path={http_request.url.path}")
# 高频轮询端点抑制 debug 日志
is_quiet = http_request.url.path in QUIET_POLLING_PATHS
if not is_quiet:
logger.debug(f"[Pipeline] START | path={http_request.url.path}")
logger.debug(f"[Pipeline] Running with mode={mode}, adapter={adapter.__class__.__name__}, "
f"adapter.mode={adapter.mode}, path={http_request.url.path}")
if mode == ApiMode.ADMIN:
user, management_token = await self._authenticate_admin(http_request, db)
api_key = None
@@ -64,10 +75,12 @@ class ApiRequestPipeline:
user, management_token = await self._authenticate_management(http_request, db)
api_key = None
else:
logger.debug("[Pipeline] 调用 _authenticate_client")
user, api_key = self._authenticate_client(http_request, db, adapter)
if not is_quiet:
logger.debug("[Pipeline] 调用 _authenticate_client")
user, api_key = self._authenticate_client(http_request, db, adapter, quiet=is_quiet)
management_token = None
logger.debug(f"[Pipeline] 认证完成 | user={user.username if user else None}")
if not is_quiet:
logger.debug(f"[Pipeline] 认证完成 | user={user.username if user else None}")
raw_body = None
if http_request.method in {"POST", "PUT", "PATCH"}:
@@ -78,7 +91,8 @@ class ApiRequestPipeline:
raw_body = await asyncio.wait_for(
http_request.body(), timeout=config.request_body_timeout
)
logger.debug(f"[Pipeline] Raw body读取完成 | size={len(raw_body) if raw_body is not None else 0} bytes")
if not is_quiet:
logger.debug(f"[Pipeline] Raw body读取完成 | size={len(raw_body) if raw_body is not None else 0} bytes")
except asyncio.TimeoutError:
timeout_sec = int(config.request_body_timeout)
logger.error(f"读取请求体超时({timeout_sec}s),可能客户端未发送完整请求体")
@@ -87,7 +101,8 @@ class ApiRequestPipeline:
detail=f"Request timeout: body not received within {timeout_sec} seconds",
)
else:
logger.debug(f"[Pipeline] 非写请求跳过读取Body | method={http_request.method}")
if not is_quiet:
logger.debug(f"[Pipeline] 非写请求跳过读取Body | method={http_request.method}")
context = ApiRequestContext.build(
request=http_request,
@@ -102,14 +117,17 @@ class ApiRequestPipeline:
# 存储 management_token 到 context用于权限检查
if management_token:
context.management_token = management_token
logger.debug(f"[Pipeline] Context构建完成 | adapter={adapter.name} | request_id={context.request_id}")
# 存储 quiet 标志到 context用于审计日志判断
context.quiet_logging = is_quiet
if not is_quiet:
logger.debug(f"[Pipeline] Context构建完成 | adapter={adapter.name} | request_id={context.request_id}")
if mode != ApiMode.ADMIN and user:
context.quota_remaining = self._calculate_quota_remaining(user)
logger.debug(f"[Pipeline] Adapter={adapter.name} | RequestID={context.request_id}")
logger.debug(f"[Pipeline] Calling authorize on {adapter.__class__.__name__}, user={context.user}")
if not is_quiet:
logger.debug(f"[Pipeline] Adapter={adapter.name} | RequestID={context.request_id}")
logger.debug(f"[Pipeline] Calling authorize on {adapter.__class__.__name__}, user={context.user}")
# authorize 可能是异步的,需要检查并 await
authorize_result = adapter.authorize(context)
if hasattr(authorize_result, "__await__"):
@@ -145,18 +163,22 @@ class ApiRequestPipeline:
# --------------------------------------------------------------------- #
def _authenticate_client(
self, request: Request, db: Session, adapter: ApiAdapter
self, request: Request, db: Session, adapter: ApiAdapter, *, quiet: bool = False
) -> Tuple[User, ApiKey]:
logger.debug("[Pipeline._authenticate_client] 开始")
if not quiet:
logger.debug("[Pipeline._authenticate_client] 开始")
# 使用 adapter 的 extract_api_key 方法,支持不同 API 格式的认证头
client_api_key = adapter.extract_api_key(request)
logger.debug(f"[Pipeline._authenticate_client] 提取API密钥完成 | key_prefix={client_api_key[:8] if client_api_key else None}...")
if not quiet:
logger.debug(f"[Pipeline._authenticate_client] 提取API密钥完成 | key_prefix={client_api_key[:8] if client_api_key else None}...")
if not client_api_key:
raise HTTPException(status_code=401, detail="请提供API密钥")
logger.debug("[Pipeline._authenticate_client] 调用 auth_service.authenticate_api_key")
if not quiet:
logger.debug("[Pipeline._authenticate_client] 调用 auth_service.authenticate_api_key")
auth_result = self.auth_service.authenticate_api_key(db, client_api_key)
logger.debug(f"[Pipeline._authenticate_client] 认证结果 | result={bool(auth_result)}")
if not quiet:
logger.debug(f"[Pipeline._authenticate_client] 认证结果 | result={bool(auth_result)}")
if not auth_result:
raise HTTPException(status_code=401, detail="无效的API密钥")
@@ -435,6 +457,8 @@ class ApiRequestPipeline:
"request_content_type": request.headers.get("content-type"),
"quota_remaining": context.quota_remaining,
"success": success,
# 传递 quiet_logging 标志给审计服务,用于抑制高频轮询日志
"quiet_logging": getattr(context, "quiet_logging", False),
}
if status_code is not None:
metadata["status_code"] = status_code

View File

@@ -1851,6 +1851,9 @@ class CliMessageHandlerBase(BaseMessageHandler):
api_key = bg_db.query(ApiKeyModel).filter(ApiKeyModel.id == ctx.api_key_id).first()
if not user or not api_key:
logger.warning(
f"[{ctx.request_id}] 无法记录统计: user={user is not None}, api_key={api_key is not None}"
)
return
bg_telemetry = MessageTelemetry(
@@ -1925,6 +1928,11 @@ class CliMessageHandlerBase(BaseMessageHandler):
}
)
logger.debug(
f"[{ctx.request_id}] 开始记录 Usage: "
f"provider={ctx.provider_name}, model={ctx.model}, "
f"in={actual_input_tokens}, out={ctx.output_tokens}"
)
total_cost = await bg_telemetry.record_success(
provider=ctx.provider_name,
model=ctx.model,
@@ -1955,7 +1963,9 @@ class CliMessageHandlerBase(BaseMessageHandler):
# Provider 响应元数据(如 Gemini 的 modelVersion
response_metadata=ctx.response_metadata if ctx.response_metadata else None,
)
logger.debug(f"{self.FORMAT_ID} 流式响应完成")
logger.debug(
f"[{ctx.request_id}] Usage 记录完成: cost=${total_cost:.6f}"
)
# 简洁的请求完成摘要(两行格式)
line1 = f"[OK] {self.request_id[:8]} | {ctx.model} | {ctx.provider_name}"
if ctx.first_byte_time_ms:

View File

@@ -151,6 +151,22 @@ class GeminiNormalizer(FormatNormalizer):
else request.get("toolConfig")
)
# 构建 extra保留原始 gemini 字段
extra: Dict[str, Any] = {"gemini": self._extract_extra(request, {"contents"})}
# 保留 generationConfig 中的特殊字段responseModalities, thinkingConfig 等)
# 这些字段在 _get_generation_config 中已提取,需要单独存储以便转换时使用
if isinstance(generation_config, dict):
response_modalities = generation_config.get("response_modalities")
thinking_config = generation_config.get("thinking_config")
if response_modalities or thinking_config:
google_extra: Dict[str, Any] = {}
if response_modalities:
google_extra["response_modalities"] = response_modalities
if thinking_config:
google_extra["thinking_config"] = thinking_config
extra["google"] = google_extra
internal = InternalRequest(
model=model,
messages=messages,
@@ -164,7 +180,7 @@ class GeminiNormalizer(FormatNormalizer):
stream=bool(request.get("stream") or False),
tools=tools,
tool_choice=tool_choice,
extra={"gemini": self._extract_extra(request, {"contents"})},
extra=extra,
)
if dropped:
@@ -208,6 +224,47 @@ class GeminiNormalizer(FormatNormalizer):
if internal.stop_sequences:
generation_config["stop_sequences"] = list(internal.stop_sequences)
# 从 internal.extra["google"] 读取 OpenAI extra_body.google 透传的配置
google_extra = internal.extra.get("google", {})
if isinstance(google_extra, dict):
# 处理 thinking_config -> thinkingConfig
thinking_config = google_extra.get("thinking_config")
if isinstance(thinking_config, dict):
# snake_case -> camelCase 转换
gemini_thinking: Dict[str, Any] = {}
if "thinking_budget" in thinking_config:
gemini_thinking["thinkingBudget"] = thinking_config["thinking_budget"]
if "include_thoughts" in thinking_config:
gemini_thinking["includeThoughts"] = thinking_config["include_thoughts"]
# 保留其他可能的字段
for k, v in thinking_config.items():
if k not in ("thinking_budget", "include_thoughts"):
gemini_thinking[k] = v
if gemini_thinking:
generation_config["thinkingConfig"] = gemini_thinking
# 处理 response_modalities -> responseModalities
response_modalities = google_extra.get("response_modalities")
if response_modalities:
generation_config["responseModalities"] = response_modalities
# 从 internal.extra["gemini"] 读取原生 Gemini 配置Gemini -> Gemini 场景)
gemini_extra = internal.extra.get("gemini", {})
if isinstance(gemini_extra, dict):
# 保留原生 Gemini generationConfig 中的额外字段
orig_gc = gemini_extra.get("generation_config") or gemini_extra.get("generationConfig")
if isinstance(orig_gc, dict):
# responseModalities
if "responseModalities" in orig_gc and "responseModalities" not in generation_config:
generation_config["responseModalities"] = orig_gc["responseModalities"]
if "response_modalities" in orig_gc and "responseModalities" not in generation_config:
generation_config["responseModalities"] = orig_gc["response_modalities"]
# thinkingConfig
if "thinkingConfig" in orig_gc and "thinkingConfig" not in generation_config:
generation_config["thinkingConfig"] = orig_gc["thinkingConfig"]
if "thinking_config" in orig_gc and "thinkingConfig" not in generation_config:
generation_config["thinkingConfig"] = orig_gc["thinking_config"]
contents: List[Dict[str, Any]] = []
for msg in internal.messages:
contents.append(self._internal_message_to_content(msg))
@@ -431,6 +488,35 @@ class GeminiNormalizer(FormatNormalizer):
events.append(ContentBlockStopEvent(block_index=block_index))
continue
# inlineData图片生成等多模态输出
inline_data = part.get("inlineData")
if inline_data is None:
inline_data = part.get("inline_data")
if isinstance(inline_data, dict):
mime_type = str(inline_data.get("mimeType") or inline_data.get("mime_type") or "").strip()
data = str(inline_data.get("data") or "").strip()
# 确保 mime_type 和 data 都非空
if mime_type and data and len(data) > 10: # base64 图片数据至少几十个字符
block_index = int(ss.get("next_block_index") or 1)
ss["next_block_index"] = block_index + 1
# 使用 ContentBlockStartEvent 传递图片数据
# 图片数据存储在 extra 中,供 target normalizer 处理
events.append(
ContentBlockStartEvent(
block_index=block_index,
block_type=ContentType.IMAGE,
extra={
"image_data": data,
"image_media_type": mime_type,
},
)
)
events.append(ContentBlockStopEvent(block_index=block_index))
continue
finish_reason = candidate0.get("finishReason")
if finish_reason is not None:
stop_reason = self._FINISH_REASON_TO_STOP.get(str(finish_reason), StopReason.UNKNOWN)
@@ -488,6 +574,25 @@ class GeminiNormalizer(FormatNormalizer):
}
return out
# 图片内容块(来自其他格式的图像生成输出)
if isinstance(event, ContentBlockStartEvent) and event.block_type == ContentType.IMAGE:
image_data = event.extra.get("image_data")
image_media_type = event.extra.get("image_media_type")
if image_data and image_media_type:
out.append(
base_chunk(
[
{
"inlineData": {
"mimeType": image_media_type,
"data": image_data,
}
}
]
)
)
return out
if isinstance(event, ToolCallDeltaEvent):
tool_blocks = ss.get("tool_blocks")
if isinstance(tool_blocks, dict):
@@ -784,6 +889,17 @@ class GeminiNormalizer(FormatNormalizer):
normalized["top_p"] = pick("top_p", "topP")
normalized["top_k"] = pick("top_k", "topK")
normalized["stop_sequences"] = pick("stop_sequences", "stopSequences")
# 保留 responseModalities图像生成等多模态输出必需
response_modalities = pick("response_modalities", "responseModalities")
if response_modalities:
normalized["response_modalities"] = response_modalities
# 保留 thinkingConfig思考模式配置
thinking_config = pick("thinking_config", "thinkingConfig")
if thinking_config:
normalized["thinking_config"] = thinking_config
return {k: v for k, v in normalized.items() if v is not None}
def _gemini_tools_to_internal(self, tools: Any) -> Optional[List[ToolDefinition]]:

View File

@@ -13,6 +13,7 @@ import json
import time
from typing import Any, Dict, List, Optional, Tuple, Union
from src.core.logger import logger
from src.core.api_format.conversion.field_mappings import (
ERROR_TYPE_MAPPINGS,
RETRYABLE_ERROR_TYPES,
@@ -144,6 +145,16 @@ class OpenAINormalizer(FormatNormalizer):
mct if mct is not None else request.get("max_tokens")
)
# 构建 extra保留未识别字段
extra: Dict[str, Any] = {"openai": self._extract_extra(request, {"messages"})}
# 处理 extra_body.google (用于 Gemini 特定功能透传,如 thinkingConfig, responseModalities)
extra_body = request.get("extra_body")
if isinstance(extra_body, dict):
google_extra = extra_body.get("google")
if isinstance(google_extra, dict) and google_extra:
extra["google"] = google_extra
internal = InternalRequest(
model=model,
messages=messages,
@@ -156,7 +167,7 @@ class OpenAINormalizer(FormatNormalizer):
stream=bool(request.get("stream") or False),
tools=tools,
tool_choice=tool_choice,
extra={"openai": self._extract_extra(request, {"messages"})},
extra=extra,
)
if dropped:
@@ -470,6 +481,50 @@ class OpenAINormalizer(FormatNormalizer):
)
return out
# 图片内容块(来自 Gemini 图像生成等多模态输出)
if isinstance(event, ContentBlockStartEvent) and event.block_type == ContentType.IMAGE:
image_data = event.extra.get("image_data")
image_media_type = event.extra.get("image_media_type")
# 确保图片数据有效base64 数据至少有一定长度)
if image_data and image_media_type and isinstance(image_data, str) and len(image_data) > 10:
# 构造 data URL 格式的图片
data_url = f"data:{image_media_type};base64,{image_data}"
# 存储图片数据,在 ContentBlockStopEvent 时输出
image_blocks = ss.get("image_blocks")
if not isinstance(image_blocks, dict):
image_blocks = {}
ss["image_blocks"] = image_blocks
image_blocks[int(event.block_index)] = {
"url": data_url,
"media_type": image_media_type,
}
return out
# 图片内容块结束时输出
if isinstance(event, ContentBlockStopEvent):
image_blocks = ss.get("image_blocks")
if isinstance(image_blocks, dict):
entry = image_blocks.get(int(event.block_index))
if isinstance(entry, dict):
url = entry.get("url")
# 确保 URL 有效data URL 至少包含 "data:" 前缀 + 一些数据)
if url and isinstance(url, str) and len(url) > 20:
# OpenAI 流式响应中 delta.content 必须是字符串
# 使用 markdown 图片格式,兼容各种客户端渲染
# 注意base64 data URL 可能很长几百KB~几MB客户端需支持长内容
# 典型图片大小100KB 原图 ≈ 130KB base641MB 原图 ≈ 1.3MB base64
if len(url) > 1_000_000: # > 1MB
logger.warning(
f"Large image in stream response: {len(url)} bytes, "
"client may have rendering issues"
)
out.append(base_chunk({"content": f"![image]({url})"}))
# 清理已处理的图片
del image_blocks[int(event.block_index)]
return out
# 其他 ContentBlockStopEvent 不处理
return out
if isinstance(event, ToolCallDeltaEvent):
tool_index = self._ensure_tool_call_index(ss, event.tool_id)
out.append(
@@ -854,20 +909,30 @@ class OpenAINormalizer(FormatNormalizer):
text_parts.append(b.text)
continue
if isinstance(b, ImageBlock):
url = b.url
if not url and b.data and b.media_type:
url = f"data:{b.media_type};base64,{b.data}"
if url:
parts.append({"type": "image_url", "image_url": {"url": url}})
# 区分两种图片来源:
# 1. URL 引用OpenAI 原生格式)-> multipart content
# 2. base64 内嵌数据(格式转换来的)-> markdown 格式
if b.url and not b.data:
# OpenAI 原生格式URL 引用的图片,使用 multipart content
parts.append({"type": "image_url", "image_url": {"url": b.url}})
elif b.data and b.media_type:
# 格式转换来的图片base64 内嵌),使用 markdown 格式
data_url = f"data:{b.media_type};base64,{b.data}"
text_parts.append(f"![image]({data_url})")
elif b.url:
# 有 URL 也有 data优先使用 URL
parts.append({"type": "image_url", "image_url": {"url": b.url}})
continue
# Unknown / Tool blocks 不进入 OpenAI content
# 如果有 OpenAI 原生格式的图片URL 引用),使用 multipart content
if parts:
if text_parts:
parts = [{"type": "text", "text": "\n".join(text_parts)}] + parts
return parts
# 纯文本或包含 markdown 图片
if text_parts:
return "\n".join(text_parts)

View File

@@ -33,6 +33,8 @@ from src.services.provider_ops.types import (
# 余额缓存 TTL24 小时)
BALANCE_CACHE_TTL = 86400
# 认证失败缓存 TTL60 秒,避免频繁重试但允许用户修正后快速重试)
AUTH_FAILED_CACHE_TTL = 60
def _get_batch_balance_concurrency() -> int:
@@ -388,9 +390,9 @@ class ProviderOpsService:
# 成功或 auth_expired 时缓存auth_expired 带有 cookie_expired 信息供前端显示警告)
if result.status in (ActionStatus.SUCCESS, ActionStatus.AUTH_EXPIRED) and result.data:
await self._cache_balance(provider_id, result)
# auth_failed 时清除缓存(配置错误,用户修正后应立即重试)
# auth_failed 时缓存(使用较短 TTL避免前端无限显示"加载中..."
elif result.status == ActionStatus.AUTH_FAILED:
await self._clear_balance_cache(provider_id)
await self._cache_auth_failed(provider_id, result)
return result
@@ -460,11 +462,29 @@ class ProviderOpsService:
pass
async def _clear_balance_cache(self, provider_id: str) -> None:
"""清除余额缓存(认证失败时调用)"""
"""清除余额缓存"""
cache_key = f"provider_ops:balance:{provider_id}"
await CacheService.delete(cache_key)
logger.info(f"余额缓存已清除: provider_id={provider_id}")
async def _cache_auth_failed(self, provider_id: str, result: ActionResult) -> None:
"""
缓存认证失败结果(使用较短 TTL
这样前端可以立即显示错误信息,而不是无限显示"加载中..."
用户修正配置后,等待 60 秒或手动刷新即可重试。
"""
cache_key = f"provider_ops:balance:{provider_id}"
cache_data = {
"status": result.status.value,
"data": None,
"message": result.message,
"executed_at": result.executed_at.isoformat() if result.executed_at else None,
"response_time_ms": result.response_time_ms,
}
await CacheService.set(cache_key, cache_data, AUTH_FAILED_CACHE_TTL)
logger.info(f"余额缓存已写入(认证失败): provider_id={provider_id}, message={result.message}")
async def _cache_balance(self, provider_id: str, result: ActionResult) -> None:
"""缓存余额结果"""
cache_key = f"provider_ops:balance:{provider_id}"
@@ -546,13 +566,17 @@ class ProviderOpsService:
else datetime.now(timezone.utc)
)
status = ActionStatus(cached.get("status", "success"))
# 认证失败使用较短的缓存 TTL
ttl = AUTH_FAILED_CACHE_TTL if status == ActionStatus.AUTH_FAILED else BALANCE_CACHE_TTL
return ActionResult(
status=ActionStatus(cached.get("status", "success")),
status=status,
action_type=ProviderActionType.QUERY_BALANCE,
data=data,
message=cached.get("message"),
executed_at=executed_at,
response_time_ms=cached.get("response_time_ms"),
cache_ttl_seconds=BALANCE_CACHE_TTL,
cache_ttl_seconds=ttl,
)
except Exception as e:
logger.warning(f"解析缓存余额失败: provider_id={provider_id}, error={e}")

View File

@@ -78,20 +78,24 @@ class AuditService:
db.flush()
# 同时记录到系统日志
log_message = (
f"AUDIT [{event_type.value}] - {description} | "
f"user_id={user_id}, ip={ip_address}"
)
# 检查 metadata 中是否有 quiet_logging 标志(由高频轮询端点设置)
quiet_logging = metadata.get("quiet_logging", False) if metadata else False
if event_type in [
AuditEventType.UNAUTHORIZED_ACCESS,
AuditEventType.SUSPICIOUS_ACTIVITY,
]:
logger.warning(log_message)
elif event_type in [AuditEventType.LOGIN_FAILED, AuditEventType.REQUEST_FAILED]:
logger.info(log_message)
else:
logger.debug(log_message)
if not quiet_logging:
log_message = (
f"AUDIT [{event_type.value}] - {description} | "
f"user_id={user_id}, ip={ip_address}"
)
if event_type in [
AuditEventType.UNAUTHORIZED_ACCESS,
AuditEventType.SUSPICIOUS_ACTIVITY,
]:
logger.warning(log_message)
elif event_type in [AuditEventType.LOGIN_FAILED, AuditEventType.REQUEST_FAILED]:
logger.info(log_message)
else:
logger.debug(log_message)
return audit_log

View File

@@ -14,6 +14,7 @@ import time
from typing import Any, Dict, List, Optional, Tuple
from redis.exceptions import ResponseError
from sqlalchemy.exc import IntegrityError
from sqlalchemy.orm import Session
from src.clients.redis_client import get_redis_client
@@ -136,6 +137,12 @@ class UsageQueueConsumer:
self._dlq_maxlen = config.usage_queue_dlq_maxlen
self._metrics_interval = config.usage_queue_metrics_interval_seconds
@staticmethod
def _is_duplicate_key_error(exc: IntegrityError) -> bool:
"""判断是否为重复键错误(唯一约束冲突)"""
err_str = str(exc).lower()
return "unique" in err_str or "duplicate" in err_str
async def start(self) -> None:
if self._running:
return
@@ -303,6 +310,19 @@ class UsageQueueConsumer:
try:
await self._apply_record_event(event, db=db)
success_ids.append(message_id)
except IntegrityError as ie:
# 重复 request_id 导致的唯一约束冲突,视为成功(记录已存在)
try:
db.rollback()
except Exception:
pass
if self._is_duplicate_key_error(ie):
logger.debug(
f"[usage-queue] Duplicate request_id, skipping: {event.request_id}"
)
success_ids.append(message_id)
else:
await self._handle_processing_error(redis_client, message_id, fields, ie)
except Exception as individual_exc:
await self._handle_processing_error(redis_client, message_id, fields, individual_exc)
# 批量 ACK 成功处理的消息

View File

@@ -12,7 +12,15 @@ from sqlalchemy.orm import Session
from src.core.api_format.metadata import can_passthrough
from src.core.logger import logger
from src.models.database import ApiKey, Provider, ProviderAPIKey, Usage, User, UserRole
from src.models.database import (
ApiKey,
Provider,
ProviderAPIKey,
RequestCandidate,
Usage,
User,
UserRole,
)
from src.services.model.cost import ModelCostService
from src.services.system.config import SystemConfigService
@@ -1025,6 +1033,7 @@ class UsageService:
- 批量插入 Usage 记录,减少 commit 次数
- 聚合更新用户/API Key 统计(按 user_id/api_key_id 分组)
- 聚合更新 GlobalModel 和 Provider 统计
- 支持更新已存在的 pending/streaming 状态记录
Args:
db: 数据库会话
@@ -1040,6 +1049,43 @@ class UsageService:
from sqlalchemy import update
from src.models.database import ApiKey as ApiKeyModel, User as UserModel, GlobalModel
# 分离需要更新和需要新建的记录
request_ids = [r.get("request_id") for r in records if r.get("request_id")]
existing_usages: Dict[str, Usage] = {}
records_to_update: List[Dict[str, Any]] = []
records_to_insert: List[Dict[str, Any]] = []
if request_ids:
# 查询已存在的 Usage 记录(包括 pending/streaming 状态)
existing_records = (
db.query(Usage)
.filter(Usage.request_id.in_(request_ids))
.all()
)
existing_usages = {u.request_id: u for u in existing_records}
for record in records:
req_id = record.get("request_id")
if req_id and req_id in existing_usages:
existing_usage = existing_usages[req_id]
# 只更新 pending/streaming 状态的记录
# 已经是 completed/failed/cancelled 的记录跳过
if existing_usage.status in ("pending", "streaming"):
records_to_update.append(record)
else:
logger.debug(
f"批量记录预过滤: 跳过已完成的 request_id={req_id} (status={existing_usage.status})"
)
else:
records_to_insert.append(record)
else:
records_to_insert = list(records)
if records_to_update:
logger.debug(
f"批量记录: 需要更新 {len(records_to_update)} 条已存在的 pending/streaming 记录"
)
usages: List[Usage] = []
user_costs: Dict[str, float] = defaultdict(float) # user_id -> total_cost
apikey_stats: Dict[str, Dict[str, Any]] = defaultdict(
@@ -1048,9 +1094,12 @@ class UsageService:
model_counts: Dict[str, int] = defaultdict(int) # model -> count
provider_costs: Dict[str, float] = defaultdict(float) # provider_id -> cost
# 合并所有需要处理的记录(用于预取 user/api_key
all_records = records_to_insert + records_to_update
# 批量预取 User 和 ApiKey避免 N+1 查询
user_ids = {r.get("user_id") for r in records if r.get("user_id")}
api_key_ids = {r.get("api_key_id") for r in records if r.get("api_key_id")}
user_ids = {r.get("user_id") for r in all_records if r.get("user_id")}
api_key_ids = {r.get("api_key_id") for r in all_records if r.get("api_key_id")}
users_map: Dict[str, User] = {}
if user_ids:
@@ -1063,9 +1112,95 @@ class UsageService:
api_keys_map = {str(k.id): k for k in api_keys}
skipped_count = 0
total_count = len(records)
updated_count = 0
total_count = len(all_records)
for record in records:
# 1. 处理需要更新的记录pending/streaming -> completed/failed/cancelled
for record in records_to_update:
try:
request_id = record.get("request_id")
existing_usage = existing_usages.get(request_id)
if not existing_usage:
skipped_count += 1
continue
# 从预取的 map 中获取 user 和 api_key 对象
user_id = record.get("user_id")
api_key_id = record.get("api_key_id")
user = users_map.get(str(user_id)) if user_id else None
api_key = api_keys_map.get(str(api_key_id)) if api_key_id else None
# 准备记录参数
params = UsageRecordParams(
db=db,
user=user,
api_key=api_key,
provider=record.get("provider") or "unknown",
model=record.get("model") or "unknown",
input_tokens=int(record.get("input_tokens") or 0),
output_tokens=int(record.get("output_tokens") or 0),
cache_creation_input_tokens=int(record.get("cache_creation_input_tokens") or 0),
cache_read_input_tokens=int(record.get("cache_read_input_tokens") or 0),
request_type=record.get("request_type") or "chat",
api_format=record.get("api_format"),
endpoint_api_format=record.get("endpoint_api_format"),
has_format_conversion=bool(record.get("has_format_conversion")),
is_stream=bool(record.get("is_stream", True)),
response_time_ms=record.get("response_time_ms"),
first_byte_time_ms=record.get("first_byte_time_ms"),
status_code=int(record.get("status_code") or 200),
error_message=record.get("error_message"),
metadata=record.get("metadata"),
request_headers=record.get("request_headers"),
request_body=record.get("request_body"),
provider_request_headers=record.get("provider_request_headers"),
response_headers=record.get("response_headers"),
client_response_headers=record.get("client_response_headers"),
response_body=record.get("response_body"),
request_id=request_id,
provider_id=record.get("provider_id"),
provider_endpoint_id=record.get("provider_endpoint_id"),
provider_api_key_id=record.get("provider_api_key_id"),
status=record.get("status") or "completed",
cache_ttl_minutes=record.get("cache_ttl_minutes"),
use_tiered_pricing=record.get("use_tiered_pricing", True),
target_model=record.get("target_model"),
)
usage_params, total_cost = await cls._prepare_usage_record(params)
# 更新已存在的 Usage 记录
cls._update_existing_usage(existing_usage, usage_params, record.get("target_model"))
usages.append(existing_usage)
updated_count += 1
# 聚合统计(更新记录也需要更新统计,因为 pending 状态没有计费)
model_name = record.get("model") or "unknown"
model_counts[model_name] += 1
provider_id = record.get("provider_id")
if provider_id:
actual_cost = usage_params.get("actual_total_cost_usd", 0)
provider_costs[provider_id] += actual_cost
# 用户统计(独立 Key 不计入创建者)
if user and not (api_key and api_key.is_standalone):
user_costs[str(user.id)] += total_cost
# API Key 统计
if api_key:
key_id = str(api_key.id)
apikey_stats[key_id]["requests"] += 1
apikey_stats[key_id]["cost"] += total_cost
apikey_stats[key_id]["is_standalone"] = api_key.is_standalone
except Exception as e:
skipped_count += 1
logger.warning(f"批量记录中更新失败: {e}, request_id={record.get('request_id')}")
continue
# 2. 处理需要新建的记录
for record in records_to_insert:
try:
# 从预取的 map 中获取 user 和 api_key 对象
user_id = record.get("user_id")
@@ -1215,7 +1350,13 @@ class UsageService:
# 单次提交所有更改
try:
db.commit()
logger.debug(f"批量记录 {len(usages)} 条使用记录成功")
inserted_count = len(usages) - updated_count
if updated_count > 0:
logger.debug(
f"批量记录成功: 更新 {updated_count} 条, 新建 {inserted_count}"
)
else:
logger.debug(f"批量记录 {len(usages)} 条使用记录成功")
except Exception as e:
logger.error(f"批量提交使用记录时出错: {e}")
db.rollback()
@@ -1989,7 +2130,8 @@ class UsageService:
records = query.all()
# 检查超时的 pending/streaming 请求
timeout_ids = []
# 收集可能超时的 usage_id 列表
timeout_candidates: List[str] = []
for r in records:
if r.status in ("pending", "streaming") and r.created_at:
# 使用全局配置的超时时间
@@ -2001,15 +2143,90 @@ class UsageService:
created_at = created_at.replace(tzinfo=timezone.utc)
elapsed = (now - created_at).total_seconds()
if elapsed > timeout_seconds:
timeout_ids.append(r.id)
# 需要获取 request_id 以便检查 RequestCandidate 表
# r.id 是 usage_id需要查询 request_id
timeout_candidates.append(r.id)
# 批量更新超时的请求
if timeout_ids:
db.query(Usage).filter(Usage.id.in_(timeout_ids)).update(
{"status": "failed", "error_message": "请求超时(服务器可能已重启)"},
synchronize_session=False,
# 批量更新超时的请求(排除已有成功完成记录的请求)
timeout_ids = []
if timeout_candidates:
# 检查 RequestCandidate 表是否有成功完成的记录
# 如果流已经成功完成stream_completed: true不应该标记为超时
# 先获取这些 Usage 的 request_id
usage_request_ids = (
db.query(Usage.id, Usage.request_id)
.filter(Usage.id.in_(timeout_candidates))
.all()
)
db.commit()
usage_id_to_request_id = {u.id: u.request_id for u in usage_request_ids}
request_id_to_usage_id = {u.request_id: u.id for u in usage_request_ids}
request_ids = list(request_id_to_usage_id.keys())
# 查询这些请求中已有成功完成记录的 request_id
# 包括两种情况:
# 1. status='success' 且 stream_completed=True正常完成
# 2. status='streaming' 且 status_code=200流传输中但 Provider 已返回 200可能是服务重启导致回调丢失
completed_usage_ids = set()
if request_ids:
from sqlalchemy import or_
candidates = (
db.query(
RequestCandidate.request_id,
RequestCandidate.status,
RequestCandidate.status_code,
RequestCandidate.extra_data,
)
.filter(
RequestCandidate.request_id.in_(request_ids),
or_(
RequestCandidate.status == "success",
# streaming 状态且 status_code=200说明 Provider 响应成功
# 但流传输可能因服务重启而中断
(RequestCandidate.status == "streaming")
& (RequestCandidate.status_code == 200),
),
)
.all()
)
for candidate in candidates:
extra_data = candidate.extra_data or {}
# 情况1status='success' 且 stream_completed=True
if candidate.status == "success" and extra_data.get(
"stream_completed", False
):
usage_id = request_id_to_usage_id.get(candidate.request_id)
if usage_id:
completed_usage_ids.add(usage_id)
# 情况2status='streaming' 且 status_code=200
# 这表示 Provider 返回了 200但流传输可能因服务重启而未正常结束
# 此时应该恢复为 completed 而不是标记为 failed
elif candidate.status == "streaming" and candidate.status_code == 200:
usage_id = request_id_to_usage_id.get(candidate.request_id)
if usage_id:
completed_usage_ids.add(usage_id)
# 只对没有成功完成记录的请求标记超时
timeout_ids = [uid for uid in timeout_candidates if uid not in completed_usage_ids]
if timeout_ids:
db.query(Usage).filter(Usage.id.in_(timeout_ids)).update(
{"status": "failed", "error_message": "请求超时(服务器可能已重启)"},
synchronize_session=False,
)
db.commit()
# 对于已完成但状态未更新的请求,主动恢复状态为 completed
# 这处理了遥测回调丢失的情况(例如服务重启、后台任务未执行等)
if completed_usage_ids:
db.query(Usage).filter(Usage.id.in_(list(completed_usage_ids))).update(
{"status": "completed"},
synchronize_session=False,
)
db.commit()
logger.info(
f"[Usage] 恢复 {len(completed_usage_ids)} 个已完成请求的状态(遥测回调丢失)"
)
result: List[Dict[str, Any]] = []
for r in records: