mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-02 01:10:23 +08:00
1d9c77522a
- 删除全部 Python 源码 (src/) 及 Alembic 迁移脚本，归档至 _deprecated_py_src/ - 重构 Rust gateway ai_pipeline: 拆分 planner/finalize 模块，新增 contracts/adaptation 层 - 重组 handlers 模块为 admin/public/proxy/internal/shared 子模块结构 - 新增 executor 模块，引入 Rust 原生数据库迁移 (aether-data/migrations) - 简化 CI/Docker 构建流程，移除 base image 二级构建，统一为单一 app image - 移除 Python 相关基础设施文件 (entrypoint.sh, gunicorn_conf.py, Dockerfile.base)
474 lines
17 KiB
Python
474 lines
17 KiB
Python
"""
|
||
流式处理上下文 - 类型安全的数据类替代 dict
|
||
|
||
提供流式请求处理过程中的状态跟踪，包括：
|
||
- Provider/Endpoint/Key 信息
|
||
- Token 统计
|
||
- 响应状态
|
||
- 请求/响应数据
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import json
|
||
import time
|
||
from contextlib import contextmanager
|
||
from dataclasses import dataclass, field
|
||
from typing import TYPE_CHECKING, Any, Iterator
|
||
|
||
if TYPE_CHECKING:
|
||
from src.core.api_format.conversion.stream_state import StreamState
|
||
|
||
|
||
def extract_proxy_timing(proxy_info: dict[str, Any] | None, headers: dict[str, str]) -> None:
|
||
"""从响应头中提取代理分阶段耗时（X-Proxy-Timing）写入 proxy_info"""
|
||
if proxy_info is None:
|
||
return
|
||
timing_raw = headers.get("x-proxy-timing")
|
||
if not timing_raw:
|
||
return
|
||
try:
|
||
timing = json.loads(timing_raw)
|
||
if isinstance(timing, dict):
|
||
proxy_info["timing"] = timing
|
||
except (json.JSONDecodeError, TypeError):
|
||
pass
|
||
|
||
|
||
def is_format_converted(
|
||
provider_api_format: str | None,
|
||
client_api_format: str | None,
|
||
) -> bool:
|
||
"""client 与 provider 的 api_format 是否真正不同（用于 usage 展示层）"""
|
||
return bool(
|
||
provider_api_format
|
||
and client_api_format
|
||
and provider_api_format.strip().lower() != client_api_format.strip().lower()
|
||
)
|
||
|
||
|
||
_MAX_COLLECTED_TEXT_CHARS = 16 * 1024
|
||
|
||
|
||
@dataclass
|
||
class RecordedStreamBodies:
|
||
"""统一封装 telemetry/usage 使用的流式响应体引用。"""
|
||
|
||
response_body: dict[str, Any] | None
|
||
client_response_body: dict[str, Any] | None
|
||
|
||
def ensure_populated(self, ctx: StreamContext, response_time_ms: int) -> None:
|
||
"""在 fallback 到需要 body 的路径时按需补建响应体。"""
|
||
if self.response_body is None:
|
||
self.response_body = ctx.build_response_body(response_time_ms)
|
||
if self.client_response_body is None:
|
||
self.client_response_body = ctx.build_client_response_body(response_time_ms)
|
||
|
||
|
||
@dataclass
|
||
class StreamContext:
|
||
"""
|
||
流式处理上下文
|
||
|
||
用于在流式请求处理过程中跟踪状态，替代原有的 ctx dict。
|
||
所有字段都有类型注解，提供更好的 IDE 支持和运行时类型安全。
|
||
"""
|
||
|
||
# 请求基本信息
|
||
model: str
|
||
api_format: str
|
||
api_family: str | None = None # 协议族（从 Adapter 层透传）
|
||
endpoint_kind: str | None = None # 端点类型（从 Adapter 层透传）
|
||
|
||
# 请求标识信息（CLI handler 需要）
|
||
request_id: str = ""
|
||
user_id: int = 0
|
||
api_key_id: int = 0
|
||
|
||
# Provider 信息（在请求执行时填充）
|
||
provider_name: str | None = None
|
||
provider_id: str | None = None
|
||
provider_type: str | None = None # Provider 类型（如 codex），用于元数据采集
|
||
# Transport 层选中的 base_url（用于 URL 可用性更新/故障转移等场景）
|
||
selected_base_url: str | None = None
|
||
endpoint_id: str | None = None
|
||
key_id: str | None = None
|
||
attempt_id: str | None = None
|
||
attempt_synced: bool = False
|
||
provider_api_format: str | None = None # Provider 的响应格式
|
||
|
||
# 模型映射
|
||
mapped_model: str | None = None
|
||
|
||
# Token 统计
|
||
input_tokens: int = 0
|
||
output_tokens: int = 0
|
||
cached_tokens: int = 0
|
||
cache_creation_tokens: int = 0
|
||
cache_creation_tokens_5m: int = 0 # 5min TTL 缓存创建
|
||
cache_creation_tokens_1h: int = 0 # 1h TTL 缓存创建
|
||
|
||
# 响应内容
|
||
_collected_text_parts: list[str] = field(default_factory=list, repr=False)
|
||
_collected_text_chars: int = field(default=0, repr=False)
|
||
_stored_collected_text_chars: int = field(default=0, repr=False)
|
||
response_id: str | None = None
|
||
final_usage: dict[str, Any] | None = None
|
||
final_response: dict[str, Any] | None = None
|
||
|
||
# 时间指标
|
||
first_byte_time_ms: int | None = None # 首字时间 (TTFB - Time To First Byte)
|
||
start_time: float = field(default_factory=time.time)
|
||
|
||
# 响应状态
|
||
status_code: int = 200
|
||
error_message: str | None = None # 客户端友好的错误消息
|
||
upstream_response: str | None = None # 原始 Provider 响应（用于请求链路追踪）
|
||
has_completion: bool = False
|
||
|
||
# 请求/响应数据
|
||
response_headers: dict[str, str] = field(default_factory=dict) # 提供商响应头
|
||
client_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
|
||
|
||
# 格式转换信息（CLI handler 需要）
|
||
client_api_format: str = ""
|
||
needs_conversion: bool = False # 是否需要跨格式转换（由 handler 层设置）
|
||
|
||
# Provider 响应元数据（CLI handler 需要）
|
||
response_metadata: dict[str, Any] = field(default_factory=dict)
|
||
|
||
# 整流标记（Thinking Rectifier）
|
||
rectified: bool = False # 请求是否经过整流（移除 thinking 块后重试）
|
||
|
||
# 流式处理统计
|
||
data_count: int = 0
|
||
chunk_count: int = 0
|
||
parsed_chunks: list[dict[str, Any]] = field(default_factory=list)
|
||
# 格式转换时保留提供商原始 chunks（转换前的数据）
|
||
provider_parsed_chunks: list[dict[str, Any]] = field(default_factory=list)
|
||
# 是否记录 parsed_chunks（可用于降低高并发/长流式响应的内存占用）
|
||
record_parsed_chunks: bool = True
|
||
|
||
# 性能采集（可选）
|
||
perf_sampled: bool = False
|
||
perf_metrics: dict[str, Any] = field(default_factory=dict)
|
||
|
||
# 代理信息（用于 usage 记录和日志，含 ttfb_ms）
|
||
proxy_info: dict[str, Any] | None = None
|
||
|
||
# 号池调度摘要（来自 ExecutionResult.pool_summary）
|
||
pool_summary: dict[str, Any] | None = None
|
||
# 候选轨迹（来自 ExecutionResult.candidate_keys，写入 usage metadata）
|
||
candidate_keys: list[dict[str, Any]] = field(default_factory=list)
|
||
# 内部调度审计摘要（重试/故障转移/账号使用轨迹）
|
||
scheduling_audit: dict[str, Any] | None = None
|
||
|
||
# 流式格式转换状态（跨 chunk 追踪）
|
||
stream_conversion_state: StreamState | None = None
|
||
stream_conversion_event_count: int = 0 # 流式转换成功的 event 计数
|
||
|
||
def reset_for_retry(self) -> None:
|
||
"""
|
||
重试时重置状态
|
||
|
||
在故障转移重试时调用，清除之前的数据避免累积。
|
||
保留 model 和 api_format，重置其他所有状态。
|
||
"""
|
||
self.release_recorded_chunks()
|
||
self.chunk_count = 0
|
||
self.data_count = 0
|
||
self.has_completion = False
|
||
self._collected_text_parts = []
|
||
self._collected_text_chars = 0
|
||
self._stored_collected_text_chars = 0
|
||
self.input_tokens = 0
|
||
self.output_tokens = 0
|
||
self.cached_tokens = 0
|
||
self.cache_creation_tokens = 0
|
||
self.cache_creation_tokens_5m = 0
|
||
self.cache_creation_tokens_1h = 0
|
||
self.error_message = None
|
||
self.upstream_response = None
|
||
self.status_code = 200
|
||
self.first_byte_time_ms = None
|
||
self.response_headers = {}
|
||
self.client_response_headers = {}
|
||
self.provider_request_headers = {}
|
||
self.provider_request_body = None
|
||
self.response_id = None
|
||
self.final_usage = None
|
||
self.final_response = None
|
||
self.proxy_info = None
|
||
self.pool_summary = None
|
||
self.candidate_keys = []
|
||
self.scheduling_audit = None
|
||
self.stream_conversion_state = None
|
||
self.stream_conversion_event_count = 0
|
||
self.needs_conversion = False
|
||
self.selected_base_url = None
|
||
|
||
def release_recorded_chunks(self) -> None:
|
||
"""释放 telemetry/usage 已消费完的 chunk 列表，避免后台任务继续持有大对象。"""
|
||
self.parsed_chunks = []
|
||
self.provider_parsed_chunks = []
|
||
|
||
@contextmanager
|
||
def managed_recorded_bodies(
|
||
self,
|
||
response_time_ms: int,
|
||
*,
|
||
include_bodies: bool = True,
|
||
) -> Iterator[RecordedStreamBodies]:
|
||
"""统一管理响应体构建与 chunk 释放，避免 telemetry 路径重复写 finally。"""
|
||
recorded_bodies = RecordedStreamBodies(
|
||
response_body=self.build_response_body(response_time_ms) if include_bodies else None,
|
||
client_response_body=(
|
||
self.build_client_response_body(response_time_ms) if include_bodies else None
|
||
),
|
||
)
|
||
try:
|
||
yield recorded_bodies
|
||
finally:
|
||
self.release_recorded_chunks()
|
||
recorded_bodies.response_body = None
|
||
recorded_bodies.client_response_body = None
|
||
|
||
@property
|
||
def collected_text(self) -> str:
|
||
"""已收集的文本内容（按需拼接，避免在流式过程中频繁做字符串拷贝）"""
|
||
return "".join(self._collected_text_parts)
|
||
|
||
@property
|
||
def collected_text_length(self) -> int:
|
||
"""已收集文本的总字符数（包含未保留到内存的截断部分）"""
|
||
return self._collected_text_chars
|
||
|
||
def append_text(self, text: str) -> None:
|
||
"""追加文本内容（仅保留有限前缀，避免长流导致内存增长）"""
|
||
if not text:
|
||
return
|
||
|
||
text_len = len(text)
|
||
self._collected_text_chars += text_len
|
||
|
||
remaining = _MAX_COLLECTED_TEXT_CHARS - self._stored_collected_text_chars
|
||
if remaining <= 0:
|
||
return
|
||
|
||
if text_len <= remaining:
|
||
self._collected_text_parts.append(text)
|
||
self._stored_collected_text_chars += text_len
|
||
return
|
||
|
||
self._collected_text_parts.append(text[:remaining])
|
||
self._stored_collected_text_chars += remaining
|
||
|
||
def update_provider_info(
|
||
self,
|
||
provider_name: str,
|
||
provider_id: str,
|
||
endpoint_id: str,
|
||
key_id: str,
|
||
provider_api_format: str | None = None,
|
||
) -> None:
|
||
"""更新 Provider 信息"""
|
||
self.provider_name = provider_name
|
||
self.provider_id = provider_id
|
||
self.endpoint_id = endpoint_id
|
||
self.key_id = key_id
|
||
self.provider_api_format = provider_api_format
|
||
|
||
def update_usage(
|
||
self,
|
||
input_tokens: int | None = None,
|
||
output_tokens: int | None = None,
|
||
cached_tokens: int | None = None,
|
||
cache_creation_tokens: int | None = None,
|
||
) -> None:
|
||
"""
|
||
更新 Token 使用统计
|
||
|
||
采用防御性更新策略：只有当新值 > 0 或当前值为 0 时才更新，避免用 0 覆盖已有的正确值。
|
||
|
||
设计原理：
|
||
- 在流式响应中，某些事件可能不包含完整的 usage 信息（字段为 0 或不存在）
|
||
- 后续事件可能会提供完整的统计数据
|
||
- 通过这种策略，确保一旦获得非零值就保留它，不会被后续的 0 值覆盖
|
||
|
||
示例场景：
|
||
- message_start 事件：input_tokens=100, output_tokens=0
|
||
- message_delta 事件：input_tokens=0, output_tokens=50
|
||
- 最终结果：input_tokens=100, output_tokens=50
|
||
|
||
注意事项：
|
||
- 此策略假设初始值为 0 是正确的默认状态
|
||
- 如果需要将已有值重置为 0，请直接修改实例属性（不使用此方法）
|
||
|
||
Args:
|
||
input_tokens: 输入 tokens 数量
|
||
output_tokens: 输出 tokens 数量
|
||
cached_tokens: 缓存命中 tokens 数量
|
||
cache_creation_tokens: 缓存创建 tokens 数量
|
||
"""
|
||
if input_tokens is not None and (input_tokens > 0 or self.input_tokens == 0):
|
||
self.input_tokens = input_tokens
|
||
if output_tokens is not None and (output_tokens > 0 or self.output_tokens == 0):
|
||
self.output_tokens = output_tokens
|
||
if cached_tokens is not None and (cached_tokens > 0 or self.cached_tokens == 0):
|
||
self.cached_tokens = cached_tokens
|
||
if cache_creation_tokens is not None and (
|
||
cache_creation_tokens > 0 or self.cache_creation_tokens == 0
|
||
):
|
||
self.cache_creation_tokens = cache_creation_tokens
|
||
|
||
def mark_failed(
|
||
self,
|
||
status_code: int,
|
||
error_message: str,
|
||
upstream_response: str | None = None,
|
||
) -> None:
|
||
"""
|
||
标记请求失败
|
||
|
||
Args:
|
||
status_code: HTTP 状态码
|
||
error_message: 客户端友好的错误消息
|
||
upstream_response: 原始 Provider 响应（用于请求链路追踪）
|
||
"""
|
||
self.status_code = status_code
|
||
self.error_message = error_message
|
||
if upstream_response:
|
||
self.upstream_response = upstream_response
|
||
|
||
def record_first_byte_time(self, start_time: float) -> None:
|
||
"""
|
||
记录首字时间 (TTFB - Time To First Byte)
|
||
|
||
应在第一次向客户端发送数据时调用。
|
||
如果已记录过,则不会覆盖(避免重试时重复记录)。
|
||
|
||
Args:
|
||
start_time: 请求开始时间 (time.time())
|
||
"""
|
||
if self.first_byte_time_ms is None:
|
||
self.first_byte_time_ms = int((time.time() - start_time) * 1000)
|
||
|
||
@property
|
||
def has_format_conversion(self) -> bool:
|
||
"""是否发生了真正的格式转换（client 和 provider 的 api_format 不同）
|
||
|
||
区别于 needs_conversion：后者包含 envelope rewrite（如 Antigravity v1internal），
|
||
不代表客户端与上游的数据格式真正不同。此属性用于 usage 展示层。
|
||
"""
|
||
return is_format_converted(self.provider_api_format, self.client_api_format)
|
||
|
||
def is_success(self) -> bool:
|
||
"""检查请求是否成功"""
|
||
return self.status_code < 400
|
||
|
||
def is_client_disconnected(self) -> bool:
|
||
"""检查是否因客户端断开连接而结束"""
|
||
return self.status_code == 499
|
||
|
||
def has_partial_response(self) -> bool:
|
||
"""是否已收到部分流式响应数据。"""
|
||
return self.data_count > 0 or self.chunk_count > 0 or self.collected_text_length > 0
|
||
|
||
def ensure_estimated_output_tokens(self) -> bool:
|
||
"""在缺少 usage 时，基于已收集文本补充输出 tokens。"""
|
||
if self.output_tokens > 0 or self.collected_text_length <= 0:
|
||
return False
|
||
self.output_tokens = max(1, self.collected_text_length // 4)
|
||
return True
|
||
|
||
def should_estimate_incomplete_tokens(self) -> bool:
|
||
"""流异常结束且尚无 usage 时，是否应做兜底 token 估算。
|
||
|
||
使用 or 而非 and：CancelledError 路径中 ensure_estimated_output_tokens
|
||
可能已补了 output_tokens，但 input_tokens 仍为 0，此时仍需估算。
|
||
"""
|
||
return (
|
||
not self.has_completion
|
||
and (self.input_tokens == 0 or self.output_tokens == 0)
|
||
and self.has_partial_response()
|
||
)
|
||
|
||
def set_ttfb_ms(self, ms: int) -> None:
|
||
"""将首字节响应耗时（TTFB）注入到 proxy_info 中"""
|
||
if self.proxy_info is not None:
|
||
self.proxy_info["ttfb_ms"] = ms
|
||
|
||
def set_proxy_timing(self, headers: dict[str, str]) -> None:
|
||
"""从代理响应头中提取分阶段耗时信息（X-Proxy-Timing）"""
|
||
extract_proxy_timing(self.proxy_info, headers)
|
||
|
||
def build_response_body(self, response_time_ms: int) -> dict[str, Any]:
|
||
"""
|
||
构建响应体元数据
|
||
|
||
用于记录到 Usage 表的 response_body 字段。
|
||
当有格式转换时，返回提供商原始 chunks；否则返回 parsed_chunks。
|
||
"""
|
||
chunks = self.provider_parsed_chunks if self.provider_parsed_chunks else self.parsed_chunks
|
||
return {
|
||
"chunks": chunks,
|
||
"metadata": {
|
||
"stream": True,
|
||
"total_chunks": len(chunks),
|
||
"data_count": self.data_count,
|
||
"has_completion": self.has_completion,
|
||
"response_time_ms": response_time_ms,
|
||
},
|
||
}
|
||
|
||
def build_client_response_body(self, response_time_ms: int) -> dict[str, Any] | None:
|
||
"""
|
||
构建客户端侧响应体元数据
|
||
|
||
仅当有格式转换时返回（parsed_chunks 是转换后的客户端格式）；
|
||
无格式转换时返回 None（此时 parsed_chunks 已在 response_body 中）。
|
||
"""
|
||
if not self.provider_parsed_chunks:
|
||
return None
|
||
return {
|
||
"chunks": self.parsed_chunks,
|
||
"metadata": {
|
||
"stream": True,
|
||
"total_chunks": len(self.parsed_chunks),
|
||
"data_count": self.data_count,
|
||
"has_completion": self.has_completion,
|
||
"response_time_ms": response_time_ms,
|
||
},
|
||
}
|
||
|
||
def get_log_summary(self, request_id: str, response_time_ms: int) -> str:
|
||
"""
|
||
获取日志摘要
|
||
|
||
用于请求完成/失败时的日志输出。
|
||
包含首字时间 (TTFB) 和总响应时间,分两行显示。
|
||
"""
|
||
if self.is_success():
|
||
status = "OK"
|
||
elif self.is_client_disconnected():
|
||
status = "CANCEL"
|
||
else:
|
||
status = "FAIL"
|
||
|
||
# 第一行:基本信息 + 首字时间
|
||
line1 = (
|
||
f"[{status}] {request_id[:8]} | {self.model} | " f"{self.provider_name or 'unknown'}"
|
||
)
|
||
if self.first_byte_time_ms is not None:
|
||
line1 += f" | TTFB: {self.first_byte_time_ms}ms"
|
||
|
||
# 第二行:总响应时间 + tokens
|
||
line2 = (
|
||
f" Total: {response_time_ms}ms | "
|
||
f"in:{self.input_tokens} out:{self.output_tokens}"
|
||
)
|
||
|
||
return f"{line1}\n{line2}"
|