mirror of
https://github.com/fawney19/Aether.git
synced 2026-01-03 08:12:26 +08:00
refactor: 重构流式处理模块，提取 StreamContext/Processor/Telemetry
- 将 chat_handler_base.py 中的流式处理逻辑拆分为三个独立模块： - StreamContext: 类型安全的流式上下文数据类，替代原有的 ctx dict - StreamProcessor: SSE 解析、预读、嵌套错误检测 - StreamTelemetryRecorder: 统计记录（Usage/Audit/Candidate） - 将硬编码配置外置到 settings.py，支持环境变量覆盖： - HTTP 超时配置（connect/write/pool） - 流式处理配置（预读行数、统计延迟） - 并发控制配置（槽位 TTL、缓存预留比例）
This commit is contained in:
fawney19
293
src/api/handlers/base/stream_telemetry.py
Normal file
293
src/api/handlers/base/stream_telemetry.py
Normal file
@@ -0,0 +1,293 @@
|
||||
"""
|
||||
流式遥测记录器 - 从 ChatHandlerBase 提取的统计记录逻辑
|
||||
|
||||
职责：
|
||||
1. 记录流式请求的成功/失败统计
|
||||
2. 更新 Usage 状态
|
||||
3. 更新候选记录状态
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
from typing import Any, Dict, Optional
|
||||
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
from src.api.handlers.base.base_handler import MessageTelemetry
|
||||
from src.api.handlers.base.stream_context import StreamContext
|
||||
from src.config.settings import config
|
||||
from src.core.logger import logger
|
||||
from src.database import get_db
|
||||
from src.models.database import ApiKey, User
|
||||
|
||||
|
||||
class StreamTelemetryRecorder:
|
||||
"""
|
||||
流式遥测记录器
|
||||
|
||||
负责在流式请求完成后记录统计信息。
|
||||
从 ChatHandlerBase 中提取的 _record_stream_stats 逻辑。
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
request_id: str,
|
||||
user_id: str,
|
||||
api_key_id: str,
|
||||
client_ip: str,
|
||||
format_id: str,
|
||||
):
|
||||
"""
|
||||
初始化遥测记录器
|
||||
|
||||
Args:
|
||||
request_id: 请求 ID
|
||||
user_id: 用户 ID
|
||||
api_key_id: API Key ID
|
||||
client_ip: 客户端 IP
|
||||
format_id: API 格式标识
|
||||
"""
|
||||
self.request_id = request_id
|
||||
self.user_id = user_id
|
||||
self.api_key_id = api_key_id
|
||||
self.client_ip = client_ip
|
||||
self.format_id = format_id
|
||||
|
||||
async def record_stream_stats(
|
||||
self,
|
||||
ctx: StreamContext,
|
||||
original_headers: Dict[str, str],
|
||||
original_request_body: Dict[str, Any],
|
||||
response_time_ms: int,
|
||||
) -> None:
|
||||
"""
|
||||
记录流式统计信息
|
||||
|
||||
Args:
|
||||
ctx: 流式上下文
|
||||
original_headers: 原始请求头
|
||||
original_request_body: 原始请求体
|
||||
response_time_ms: 响应时间（毫秒）
|
||||
"""
|
||||
bg_db = None
|
||||
|
||||
try:
|
||||
await asyncio.sleep(config.stream_stats_delay) # 等待流完全关闭
|
||||
|
||||
if not ctx.provider_name:
|
||||
await self._update_usage_status_on_error(
|
||||
response_time_ms=response_time_ms,
|
||||
error_message="Provider name not available",
|
||||
)
|
||||
return
|
||||
|
||||
db_gen = get_db()
|
||||
bg_db = next(db_gen)
|
||||
|
||||
try:
|
||||
user = bg_db.query(User).filter(User.id == self.user_id).first()
|
||||
api_key_obj = bg_db.query(ApiKey).filter(ApiKey.id == self.api_key_id).first()
|
||||
|
||||
if not user or not api_key_obj:
|
||||
logger.warning(
|
||||
f"[{self.request_id}] User or ApiKey not found, updating status directly"
|
||||
)
|
||||
await self._update_usage_status_directly(
|
||||
bg_db,
|
||||
status="completed" if ctx.is_success() else "failed",
|
||||
response_time_ms=response_time_ms,
|
||||
status_code=ctx.status_code,
|
||||
)
|
||||
return
|
||||
|
||||
bg_telemetry = MessageTelemetry(
|
||||
bg_db, user, api_key_obj, self.request_id, self.client_ip
|
||||
)
|
||||
|
||||
actual_request_body = ctx.provider_request_body or original_request_body
|
||||
response_body = ctx.build_response_body(response_time_ms)
|
||||
|
||||
if ctx.is_success():
|
||||
await self._record_success(
|
||||
bg_telemetry,
|
||||
ctx,
|
||||
original_headers,
|
||||
actual_request_body,
|
||||
response_body,
|
||||
response_time_ms,
|
||||
)
|
||||
else:
|
||||
await self._record_failure(
|
||||
bg_telemetry,
|
||||
ctx,
|
||||
original_headers,
|
||||
actual_request_body,
|
||||
response_body,
|
||||
response_time_ms,
|
||||
)
|
||||
|
||||
# 更新候选记录状态
|
||||
await self._update_candidate_status(bg_db, ctx, response_time_ms)
|
||||
|
||||
finally:
|
||||
if bg_db:
|
||||
bg_db.close()
|
||||
|
||||
except Exception as e:
|
||||
logger.exception("记录流式统计信息时出错")
|
||||
await self._update_usage_status_on_error(
|
||||
response_time_ms=response_time_ms,
|
||||
error_message=f"记录统计信息失败: {str(e)[:200]}",
|
||||
)
|
||||
|
||||
async def _record_success(
|
||||
self,
|
||||
telemetry: MessageTelemetry,
|
||||
ctx: StreamContext,
|
||||
original_headers: Dict[str, str],
|
||||
actual_request_body: Dict[str, Any],
|
||||
response_body: Dict[str, Any],
|
||||
response_time_ms: int,
|
||||
) -> None:
|
||||
"""记录成功的请求"""
|
||||
await telemetry.record_success(
|
||||
provider=ctx.provider_name or "unknown",
|
||||
model=ctx.model,
|
||||
input_tokens=ctx.input_tokens,
|
||||
output_tokens=ctx.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,
|
||||
response_body=response_body,
|
||||
cache_creation_tokens=ctx.cache_creation_tokens,
|
||||
cache_read_tokens=ctx.cached_tokens,
|
||||
is_stream=True,
|
||||
provider_request_headers=ctx.provider_request_headers,
|
||||
api_format=ctx.api_format,
|
||||
provider_id=ctx.provider_id,
|
||||
provider_endpoint_id=ctx.endpoint_id,
|
||||
provider_api_key_id=ctx.key_id,
|
||||
target_model=ctx.mapped_model,
|
||||
)
|
||||
|
||||
logger.debug(f"{self.format_id} 流式响应完成")
|
||||
logger.info(ctx.get_log_summary(self.request_id, response_time_ms))
|
||||
|
||||
async def _record_failure(
|
||||
self,
|
||||
telemetry: MessageTelemetry,
|
||||
ctx: StreamContext,
|
||||
original_headers: Dict[str, str],
|
||||
actual_request_body: Dict[str, Any],
|
||||
response_body: Dict[str, Any],
|
||||
response_time_ms: int,
|
||||
) -> None:
|
||||
"""记录失败的请求"""
|
||||
await telemetry.record_failure(
|
||||
provider=ctx.provider_name or "unknown",
|
||||
model=ctx.model,
|
||||
response_time_ms=response_time_ms,
|
||||
status_code=ctx.status_code,
|
||||
error_message=ctx.error_message or f"HTTP {ctx.status_code}",
|
||||
request_headers=original_headers,
|
||||
request_body=actual_request_body,
|
||||
is_stream=True,
|
||||
api_format=ctx.api_format,
|
||||
provider_request_headers=ctx.provider_request_headers,
|
||||
input_tokens=ctx.input_tokens,
|
||||
output_tokens=ctx.output_tokens,
|
||||
cache_creation_tokens=ctx.cache_creation_tokens,
|
||||
cache_read_tokens=ctx.cached_tokens,
|
||||
response_body=response_body,
|
||||
target_model=ctx.mapped_model,
|
||||
)
|
||||
|
||||
logger.debug(f"{self.format_id} 流式响应中断")
|
||||
log_summary = ctx.get_log_summary(self.request_id, response_time_ms)
|
||||
# 对于失败日志，添加缓存信息
|
||||
logger.info(f"{log_summary} cache:{ctx.cached_tokens}")
|
||||
|
||||
async def _update_candidate_status(
|
||||
self,
|
||||
db: Session,
|
||||
ctx: StreamContext,
|
||||
response_time_ms: int,
|
||||
) -> None:
|
||||
"""更新候选记录状态"""
|
||||
if not ctx.attempt_id:
|
||||
return
|
||||
|
||||
from src.services.request.candidate import RequestCandidateService
|
||||
|
||||
if ctx.is_success():
|
||||
RequestCandidateService.mark_candidate_success(
|
||||
db=db,
|
||||
candidate_id=ctx.attempt_id,
|
||||
status_code=ctx.status_code,
|
||||
latency_ms=response_time_ms,
|
||||
extra_data={
|
||||
"stream_completed": True,
|
||||
"data_count": ctx.data_count,
|
||||
},
|
||||
)
|
||||
else:
|
||||
error_type = "client_disconnected" if ctx.status_code == 499 else "stream_error"
|
||||
RequestCandidateService.mark_candidate_failed(
|
||||
db=db,
|
||||
candidate_id=ctx.attempt_id,
|
||||
error_type=error_type,
|
||||
error_message=ctx.error_message or f"HTTP {ctx.status_code}",
|
||||
status_code=ctx.status_code,
|
||||
latency_ms=response_time_ms,
|
||||
extra_data={
|
||||
"stream_completed": False,
|
||||
"data_count": ctx.data_count,
|
||||
},
|
||||
)
|
||||
|
||||
async def _update_usage_status_on_error(
|
||||
self,
|
||||
response_time_ms: int,
|
||||
error_message: str,
|
||||
) -> None:
|
||||
"""在记录失败时更新 Usage 状态"""
|
||||
try:
|
||||
db_gen = get_db()
|
||||
error_db = next(db_gen)
|
||||
try:
|
||||
await self._update_usage_status_directly(
|
||||
error_db,
|
||||
status="failed",
|
||||
response_time_ms=response_time_ms,
|
||||
status_code=500,
|
||||
error_message=error_message,
|
||||
)
|
||||
finally:
|
||||
error_db.close()
|
||||
except Exception as inner_e:
|
||||
logger.error(f"[{self.request_id}] 更新 Usage 状态失败: {inner_e}")
|
||||
|
||||
async def _update_usage_status_directly(
|
||||
self,
|
||||
db: Session,
|
||||
status: str,
|
||||
response_time_ms: int,
|
||||
status_code: int = 200,
|
||||
error_message: Optional[str] = None,
|
||||
) -> None:
|
||||
"""直接更新 Usage 表的状态字段"""
|
||||
try:
|
||||
from src.models.database import Usage
|
||||
|
||||
usage = db.query(Usage).filter(Usage.request_id == self.request_id).first()
|
||||
if usage:
|
||||
setattr(usage, "status", status)
|
||||
setattr(usage, "status_code", status_code)
|
||||
setattr(usage, "response_time_ms", response_time_ms)
|
||||
if error_message:
|
||||
setattr(usage, "error_message", error_message)
|
||||
db.commit()
|
||||
logger.debug(f"[{self.request_id}] Usage 状态已更新: {status}")
|
||||
except Exception as e:
|
||||
logger.error(f"[{self.request_id}] 直接更新 Usage 状态失败: {e}")
|
||||
Reference in New Issue
Block a user