mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-02 17:30:23 +08:00
438f16094f
- 新增 has_partial_response / ensure_estimated_output_tokens / should_estimate_incomplete_tokens 方法 - CancelledError 路径在归因前即补充 output_tokens，确保日志包含估算值 - CLI Handler 和 Chat Handler 的兜底估算统一使用 should_estimate_incomplete_tokens - 移除 cli_monitor_mixin 和 stream_telemetry 中重复的条件判断 - 新增对应单元测试
67 lines
2.0 KiB
Python
67 lines
2.0 KiB
Python
from __future__ import annotations
|
|
|
|
import time
|
|
from types import SimpleNamespace
|
|
from unittest.mock import AsyncMock
|
|
|
|
import pytest
|
|
|
|
from src.api.handlers.base import stream_telemetry as stream_telemetry_module
|
|
from src.api.handlers.base.stream_context import StreamContext
|
|
from src.api.handlers.base.stream_telemetry import StreamTelemetryRecorder
|
|
|
|
|
|
class _DummyDb:
|
|
def close(self) -> None:
|
|
pass
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_record_stream_stats_estimates_tokens_for_failed_partial_stream(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
recorder = StreamTelemetryRecorder(
|
|
request_id="req-telemetry",
|
|
user_id="1",
|
|
api_key_id="2",
|
|
client_ip="127.0.0.1",
|
|
format_id="openai:chat",
|
|
)
|
|
recorder._get_telemetry_writer = AsyncMock( # type: ignore[method-assign]
|
|
return_value=SimpleNamespace(include_bodies=False)
|
|
)
|
|
recorder._dispatch_record = AsyncMock() # type: ignore[method-assign]
|
|
recorder._update_candidate_status = AsyncMock() # type: ignore[method-assign]
|
|
|
|
ctx = StreamContext(
|
|
model="test-model",
|
|
api_format="openai:chat",
|
|
request_id="req-telemetry",
|
|
user_id=1,
|
|
api_key_id=2,
|
|
)
|
|
ctx.provider_name = "test-provider"
|
|
ctx.status_code = 503
|
|
ctx.data_count = 2
|
|
ctx.chunk_count = 4
|
|
ctx.append_text("partial output")
|
|
|
|
monkeypatch.setattr(stream_telemetry_module, "get_db", lambda: iter([_DummyDb()]))
|
|
monkeypatch.setattr(
|
|
stream_telemetry_module.SystemConfigService,
|
|
"should_log_body",
|
|
lambda _db: False,
|
|
)
|
|
monkeypatch.setattr(stream_telemetry_module.config, "stream_stats_delay", 0)
|
|
|
|
await recorder.record_stream_stats(
|
|
ctx,
|
|
original_headers={},
|
|
original_request_body={"input": [{"content": "hello world"}]},
|
|
start_time=time.time(),
|
|
)
|
|
|
|
assert ctx.input_tokens > 0
|
|
assert ctx.output_tokens == max(1, len("partial output") // 4)
|
|
recorder._dispatch_record.assert_awaited_once() # type: ignore[attr-defined]
|