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

perf: 优化流式响应内存占用和正则预编译

Browse Source 
- 添加 record_parsed_chunks 标志，仅在 FULL 日志级别时累积 parsed_chunks
- 预编译 CJK 和敏感信息正则表达式，避免重复编译
- 优化 header 脱敏循环，预计算 lowercase 集合
- 移除 consumer_streams.py 中未使用的 message_fields 变量
- 将 SystemConfigService import 移至文件顶部
This commit is contained in:
fawney19
2026-02-03 21:32:53 +08:00
parent d6597121f5
commit 6f363a7703
9 changed files with 40 additions and 23 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -73,6 +73,7 @@ from src.services.provider.transport import (
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:
@@ -541,6 +542,8 @@ class ChatHandlerBase(BaseMessageHandler, ABC):
ctx.client_api_format = (
api_format.value if hasattr(api_format, "value") else str(api_format)
)
# 仅在 FULL 级别才需要保留 parsed_chunks，避免长流式响应导致的内存占用
ctx.record_parsed_chunks = SystemConfigService.should_log_body(self.db)
# 创建更新状态的回调闭包（可以访问 ctx）
def update_streaming_status() -> None:

11
src/api/handlers/base/cli_handler_base.py
View File

@@ -73,6 +73,7 @@ from src.models.database import (
)
from src.services.cache.aware_scheduler import ProviderCandidate
from src.services.provider.transport import build_provider_url
from src.services.system.config import SystemConfigService
from src.utils.sse_parser import SSEEventParser
from src.utils.timeout import read_first_chunk_with_ttfb_timeout
@@ -544,6 +545,8 @@ class CliMessageHandlerBase(BaseMessageHandler):
user_id=self.user.id,
api_key_id=self.api_key.id,
)
# 仅在 FULL 级别才需要保留 parsed_chunks，避免长流式响应导致的内存占用
ctx.record_parsed_chunks = SystemConfigService.should_log_body(self.db)
# 定义请求函数
async def stream_request_func(
@@ -1530,11 +1533,12 @@ class CliMessageHandlerBase(BaseMessageHandler):
except json.JSONDecodeError:
return
# 当不需要格式转换时，记录原始数据到 parsed_chunks 并更新 data_count
# 当不需要格式转换时，更新 data_count；需要记录时再写入 parsed_chunks。
# 当需要格式转换时（record_chunk=False），data_count 由 _record_converted_chunks 更新
if record_chunk and isinstance(data, dict):
ctx.parsed_chunks.append(data)
ctx.data_count += 1
if ctx.record_parsed_chunks:
ctx.parsed_chunks.append(data)
if not isinstance(data, dict):
return
@@ -1641,8 +1645,9 @@ class CliMessageHandlerBase(BaseMessageHandler):
"""
for evt in converted_events:
if isinstance(evt, dict):
ctx.parsed_chunks.append(evt)
ctx.data_count += 1
if ctx.record_parsed_chunks:
ctx.parsed_chunks.append(evt)
# 检测完成事件（根据客户端格式判断）
# OpenAI 格式: choices[].finish_reason

2
src/api/handlers/base/stream_context.py
View File

@@ -90,6 +90,8 @@ class StreamContext:
data_count: int = 0
chunk_count: int = 0
parsed_chunks: list[dict[str, Any]] = field(default_factory=list)
# 是否记录 parsed_chunks（可用于降低高并发/长流式响应的内存占用）
record_parsed_chunks: bool = True
# 流式格式转换状态（跨 chunk 追踪）
stream_conversion_state: StreamState | None = None

8
src/api/handlers/base/stream_processor.py
View File

@@ -139,10 +139,11 @@ class StreamProcessor:
if not isinstance(data, dict):
return
# 收集原始 chunk 数据（当需要格式转换时跳过，由 _emit_converted_line 记录转换后的数据）
# 统计数据事件数量（当需要格式转换时跳过，由 _emit_converted_line 统计/记录转换后的数据）
if not skip_record:
ctx.parsed_chunks.append(data)
ctx.data_count += 1
if ctx.record_parsed_chunks:
ctx.parsed_chunks.append(data)
# 根据 Provider 格式选择解析器
parser = self.get_parser_for_provider(ctx)
@@ -565,8 +566,9 @@ class StreamProcessor:
for evt in converted_events:
# 记录转换后的数据到 parsed_chunks（这是客户端实际收到的格式）
if isinstance(evt, dict):
ctx.parsed_chunks.append(evt)
ctx.data_count += 1
if ctx.record_parsed_chunks:
ctx.parsed_chunks.append(evt)
# 统一使用 SSE 格式输出（Gemini streamGenerateContent 也使用 SSE）
# 参考: https://ai.google.dev/api/generate-content

17
src/api/handlers/base/stream_telemetry.py
View File

@@ -20,6 +20,7 @@ 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
from src.services.system.config import SystemConfigService
from src.services.usage.telemetry_writer import (
DbTelemetryWriter,
QueueTelemetryWriter,
@@ -99,9 +100,15 @@ class StreamTelemetryRecorder:
if writer is None:
return
actual_request_body = ctx.provider_request_body or original_request_body
response_body = None
if not isinstance(writer, QueueTelemetryWriter) or writer.include_bodies:
response_body = ctx.build_response_body(response_time_ms)
should_log_body = SystemConfigService.should_log_body(bg_db)
include_bodies = (
writer.include_bodies
if isinstance(writer, QueueTelemetryWriter)
else should_log_body
)
response_body = (
ctx.build_response_body(response_time_ms) if include_bodies else None
)
try:
await self._dispatch_record(
@@ -127,7 +134,7 @@ class StreamTelemetryRecorder:
status_code=ctx.status_code,
)
return
if response_body is None:
if response_body is None and should_log_body:
response_body = ctx.build_response_body(response_time_ms)
await self._dispatch_record(
db_writer,
@@ -400,8 +407,6 @@ class StreamTelemetryRecorder:
self, bg_db: Session, ctx: StreamContext, response_time_ms: int
) -> TelemetryWriter | None:
if config.usage_queue_enabled and self.user_id and self.api_key_id:
from src.services.system.config import SystemConfigService
# Queue payload detail follows system config request_record_level.
log_level = SystemConfigService.get_request_record_level(bg_db).value
sensitive_headers = SystemConfigService.get_sensitive_headers(bg_db) or []

5
src/plugins/token/claude_counter.py
View File

@@ -11,6 +11,8 @@ from typing import Any
from .base import TokenCounterPlugin
_CJK_PATTERN = re.compile(r"[\u4e00-\u9fff\u3040-\u30ff\uac00-\ud7af]")
class ClaudeTokenCounterPlugin(TokenCounterPlugin):
"""
@@ -138,8 +140,7 @@ class ClaudeTokenCounterPlugin(TokenCounterPlugin):
# 考虑不同语言的特点
# 检测是否包含中文/日文/韩文
cjk_pattern = re.compile(r"[\u4e00-\u9fff\u3040-\u30ff\uac00-\ud7af]")
cjk_count = len(cjk_pattern.findall(text))
cjk_count = len(_CJK_PATTERN.findall(text))
if cjk_count > len(text) * 0.3: # 超过30%是CJK字符
# CJK字符通常每个字符1-2个token

3
src/services/system/config.py
View File

@@ -500,10 +500,11 @@ class SystemConfigService:
return headers
sensitive_headers = cls.get_sensitive_headers(db)
sensitive_lower = {h.lower() for h in sensitive_headers if isinstance(h, str) and h}
masked_headers = {}
for key, value in headers.items():
if key.lower() in [h.lower() for h in sensitive_headers]:
if key.lower() in sensitive_lower:
# 保留前后各4个字符，中间用星号替换
if len(str(value)) > 8:
masked_value = str(value)[:4] + "****" + str(value)[-4:]

12
src/services/task/service.py
View File

@@ -1,5 +1,6 @@
from __future__ import annotations
import re
from collections.abc import Callable
from typing import Any
@@ -18,6 +19,11 @@ from src.services.task.exceptions import TaskNotFoundError
from src.services.task.protocol import AttemptKind, AttemptResult
from src.services.task.schema import ExecutionResult, TaskStatusResult
_SENSITIVE_PATTERN = re.compile(
r"(api[_-]?key|token|bearer|authorization)[=:\s]+\S+",
re.IGNORECASE,
)
class TaskService:
"""
@@ -883,7 +889,6 @@ class TaskService:
- stop on "client error" (raise UpstreamClientRequestError)
- if all failed, raise AllCandidatesFailedError
"""
import re
from datetime import datetime, timezone
import httpx
@@ -901,11 +906,6 @@ class TaskService:
from src.services.orchestration.error_classifier import ErrorClassifier
from src.services.system.config import SystemConfigService
_SENSITIVE_PATTERN = re.compile(
r"(api[_-]?key|token|bearer|authorization)[=:\s]+\S+",
re.IGNORECASE,
)
def _sanitize(message: str, max_length: int = 200) -> str:
if not message:
return "request_failed"

2
src/services/usage/consumer_streams.py
View File

@@ -283,12 +283,10 @@ class UsageQueueConsumer:
# 准备批量记录数据
records: list[dict[str, Any]] = []
message_ids: list[str] = []
message_fields: list[dict[str, Any]] = []
for message_id, fields, event in messages:
records.append(_event_to_record(event))
message_ids.append(message_id)
message_fields.append(fields)
# 批量写入
await UsageService.record_usage_batch(db, records)