mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-01 17:00:21 +08:00
refactor(usage): 统一 cache token 提取逻辑,新增请求缓存指纹记录
- 新增 extract_cache_read_tokens() 兼容 OpenAI/Claude/Gemini 多种字段命名 - parsers/stream_processor/cli_event_mixin 统一使用提取函数替换内联逻辑 - 同时兼容 prompt_tokens/completion_tokens (OpenAI) 和 input_tokens/output_tokens (Claude) - 新增 cache_fingerprint 模块,在 telemetry 记录时自动计算并附带请求缓存指纹 - 新增对应单元测试
This commit is contained in:
113
tests/api/handlers/base/test_openai_usage_parsing.py
Normal file
113
tests/api/handlers/base/test_openai_usage_parsing.py
Normal file
@@ -0,0 +1,113 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
from src.api.handlers.base.parsers import OpenAICliResponseParser, OpenAIResponseParser
|
||||
from src.api.handlers.base.response_parser import (
|
||||
ParsedChunk,
|
||||
ParsedResponse,
|
||||
ResponseParser,
|
||||
StreamStats,
|
||||
)
|
||||
from src.api.handlers.base.stream_context import StreamContext
|
||||
from src.api.handlers.base.stream_processor import StreamProcessor
|
||||
|
||||
|
||||
class _DummyParser(ResponseParser):
|
||||
def parse_sse_line(self, line: str, stats: StreamStats) -> ParsedChunk | None:
|
||||
return None
|
||||
|
||||
def parse_response(self, response: dict[str, Any], status_code: int) -> ParsedResponse:
|
||||
return ParsedResponse(raw_response=response, status_code=status_code)
|
||||
|
||||
def extract_usage_from_response(self, response: dict[str, Any]) -> dict[str, int]:
|
||||
return {}
|
||||
|
||||
def extract_text_content(self, response: dict[str, Any]) -> str:
|
||||
return ""
|
||||
|
||||
|
||||
def test_openai_response_parser_extracts_cached_tokens_from_prompt_tokens_details() -> None:
|
||||
parser = OpenAIResponseParser()
|
||||
|
||||
usage = parser.extract_usage_from_response(
|
||||
{
|
||||
"usage": {
|
||||
"prompt_tokens": 120,
|
||||
"completion_tokens": 18,
|
||||
"prompt_tokens_details": {"cached_tokens": 96},
|
||||
}
|
||||
}
|
||||
)
|
||||
|
||||
assert usage["input_tokens"] == 120
|
||||
assert usage["output_tokens"] == 18
|
||||
assert usage["cache_read_tokens"] == 96
|
||||
|
||||
|
||||
def test_openai_cli_response_parser_extracts_cached_tokens_from_input_tokens_details() -> None:
|
||||
parser = OpenAICliResponseParser()
|
||||
|
||||
usage = parser.extract_usage_from_response(
|
||||
{
|
||||
"type": "response.completed",
|
||||
"response": {
|
||||
"usage": {
|
||||
"input_tokens": 2048,
|
||||
"output_tokens": 128,
|
||||
"input_tokens_details": {"cached_tokens": 1792},
|
||||
}
|
||||
},
|
||||
}
|
||||
)
|
||||
|
||||
assert usage["input_tokens"] == 2048
|
||||
assert usage["output_tokens"] == 128
|
||||
assert usage["cache_read_tokens"] == 1792
|
||||
|
||||
|
||||
def test_stream_processor_extracts_cached_tokens_from_openai_cli_converted_event() -> None:
|
||||
processor = StreamProcessor(request_id="req_test", default_parser=_DummyParser())
|
||||
ctx = StreamContext(model="gpt-5", api_format="openai:chat")
|
||||
|
||||
processor._extract_usage_from_converted_event(
|
||||
ctx,
|
||||
{
|
||||
"type": "response.completed",
|
||||
"response": {
|
||||
"usage": {
|
||||
"input_tokens": 4096,
|
||||
"output_tokens": 64,
|
||||
"input_tokens_details": {"cached_tokens": 3584},
|
||||
}
|
||||
},
|
||||
},
|
||||
"response.completed",
|
||||
)
|
||||
|
||||
assert ctx.input_tokens == 4096
|
||||
assert ctx.output_tokens == 64
|
||||
assert ctx.cached_tokens == 3584
|
||||
|
||||
|
||||
def test_stream_processor_extracts_cached_tokens_from_openai_chat_converted_event() -> None:
|
||||
processor = StreamProcessor(request_id="req_test", default_parser=_DummyParser())
|
||||
ctx = StreamContext(model="gpt-5", api_format="openai:chat")
|
||||
|
||||
processor._extract_usage_from_converted_event(
|
||||
ctx,
|
||||
{
|
||||
"object": "chat.completion.chunk",
|
||||
"choices": [],
|
||||
"usage": {
|
||||
"prompt_tokens": 512,
|
||||
"completion_tokens": 21,
|
||||
"prompt_tokens_details": {"cached_tokens": 480},
|
||||
},
|
||||
},
|
||||
"chat.completion.chunk",
|
||||
)
|
||||
|
||||
assert ctx.input_tokens == 512
|
||||
assert ctx.output_tokens == 21
|
||||
assert ctx.cached_tokens == 480
|
||||
207
tests/services/test_request_cache_fingerprint.py
Normal file
207
tests/services/test_request_cache_fingerprint.py
Normal file
@@ -0,0 +1,207 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from types import SimpleNamespace
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
|
||||
from src.config.settings import config
|
||||
from src.services.provider.cache_fingerprint import build_request_cache_fingerprint
|
||||
from src.services.usage._recording_helpers import sanitize_request_metadata
|
||||
from src.services.usage.service import UsageService
|
||||
from src.services.usage.telemetry import MessageTelemetry
|
||||
|
||||
|
||||
def _build_openai_cli_body() -> dict[str, Any]:
|
||||
return {
|
||||
"model": "gpt-5.4",
|
||||
"instructions": "You are precise.",
|
||||
"input": [{"role": "user", "content": [{"type": "input_text", "text": "hello"}]}],
|
||||
"tools": [
|
||||
{
|
||||
"type": "function",
|
||||
"name": "lookup_weather",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"required": ["city", "country"],
|
||||
"properties": {
|
||||
"country": {"type": "string"},
|
||||
"city": {"type": "string"},
|
||||
},
|
||||
},
|
||||
}
|
||||
],
|
||||
"temperature": 0.2,
|
||||
"prompt_cache_key": "pcache-123",
|
||||
}
|
||||
|
||||
|
||||
def test_build_request_cache_fingerprint_is_stable_for_dict_key_reordering() -> None:
|
||||
body_a = _build_openai_cli_body()
|
||||
body_b = {
|
||||
"prompt_cache_key": "pcache-123",
|
||||
"temperature": 0.2,
|
||||
"tools": [
|
||||
{
|
||||
"parameters": {
|
||||
"properties": {
|
||||
"city": {"type": "string"},
|
||||
"country": {"type": "string"},
|
||||
},
|
||||
"required": ["city", "country"],
|
||||
"type": "object",
|
||||
},
|
||||
"name": "lookup_weather",
|
||||
"type": "function",
|
||||
}
|
||||
],
|
||||
"input": [{"content": [{"text": "hello", "type": "input_text"}], "role": "user"}],
|
||||
"instructions": "You are precise.",
|
||||
"model": "gpt-5.4",
|
||||
}
|
||||
|
||||
fingerprint_a = build_request_cache_fingerprint(body_a, provider_api_format="openai:cli")
|
||||
fingerprint_b = build_request_cache_fingerprint(body_b, provider_api_format="openai:cli")
|
||||
|
||||
assert fingerprint_a is not None
|
||||
assert fingerprint_b is not None
|
||||
assert fingerprint_a["payload_sha256"] == fingerprint_b["payload_sha256"]
|
||||
assert fingerprint_a["cache_relevant_sha256"] == fingerprint_b["cache_relevant_sha256"]
|
||||
assert fingerprint_a["prompt_cache_key"] == "pcache-123"
|
||||
assert fingerprint_a["cache_relevant_keys"] == [
|
||||
"input",
|
||||
"instructions",
|
||||
"model",
|
||||
"prompt_cache_key",
|
||||
"tools",
|
||||
]
|
||||
|
||||
|
||||
def test_build_request_cache_fingerprint_ignores_non_prompt_fields_in_cache_hash() -> None:
|
||||
body_a = _build_openai_cli_body()
|
||||
body_b = _build_openai_cli_body()
|
||||
body_b["temperature"] = 0.9
|
||||
|
||||
fingerprint_a = build_request_cache_fingerprint(body_a, provider_api_format="openai:cli")
|
||||
fingerprint_b = build_request_cache_fingerprint(body_b, provider_api_format="openai:cli")
|
||||
|
||||
assert fingerprint_a is not None
|
||||
assert fingerprint_b is not None
|
||||
assert fingerprint_a["payload_sha256"] != fingerprint_b["payload_sha256"]
|
||||
assert fingerprint_a["cache_relevant_sha256"] == fingerprint_b["cache_relevant_sha256"]
|
||||
|
||||
|
||||
def test_build_request_cache_fingerprint_tracks_prompt_changes() -> None:
|
||||
body_a = _build_openai_cli_body()
|
||||
body_b = _build_openai_cli_body()
|
||||
body_b["instructions"] = "You are terse."
|
||||
|
||||
fingerprint_a = build_request_cache_fingerprint(body_a, provider_api_format="openai:cli")
|
||||
fingerprint_b = build_request_cache_fingerprint(body_b, provider_api_format="openai:cli")
|
||||
|
||||
assert fingerprint_a is not None
|
||||
assert fingerprint_b is not None
|
||||
assert fingerprint_a["cache_relevant_sha256"] != fingerprint_b["cache_relevant_sha256"]
|
||||
|
||||
|
||||
def test_sanitize_request_metadata_preserves_cache_fingerprint(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
monkeypatch.setattr(config, "usage_metadata_max_bytes", 120, raising=False)
|
||||
|
||||
metadata = {
|
||||
"trace": {"payload": "x" * 400},
|
||||
"debug": {"payload": "y" * 400},
|
||||
"cache_fingerprint": {
|
||||
"payload_sha256": "a" * 64,
|
||||
"cache_relevant_sha256": "b" * 64,
|
||||
},
|
||||
}
|
||||
|
||||
sanitized = sanitize_request_metadata(metadata)
|
||||
|
||||
assert sanitized["_metadata_truncated"] is True
|
||||
assert sanitized["cache_fingerprint"]["payload_sha256"] == "a" * 64
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_message_telemetry_record_success_keeps_response_shape_and_adds_fingerprint(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
captured: dict[str, Any] = {}
|
||||
|
||||
async def _fake_record_usage(**kwargs: Any) -> Any:
|
||||
captured.update(kwargs)
|
||||
return SimpleNamespace(total_cost_usd=0.0, input_tokens=1, output_tokens=2)
|
||||
|
||||
monkeypatch.setattr(UsageService, "record_usage", _fake_record_usage)
|
||||
|
||||
telemetry = MessageTelemetry(
|
||||
db=SimpleNamespace(), # type: ignore[arg-type]
|
||||
user=None,
|
||||
api_key=None,
|
||||
request_id="req-cache-fingerprint",
|
||||
client_ip="127.0.0.1",
|
||||
)
|
||||
|
||||
await telemetry.record_success(
|
||||
provider="openai",
|
||||
model="gpt-5.4",
|
||||
input_tokens=1,
|
||||
output_tokens=2,
|
||||
response_time_ms=10,
|
||||
status_code=200,
|
||||
request_body={"messages": [{"role": "user", "content": "hello"}]},
|
||||
request_headers={"user-agent": "codex desktop"},
|
||||
response_body={"id": "resp-1"},
|
||||
response_headers={"x-test": "1"},
|
||||
provider_request_body=_build_openai_cli_body(),
|
||||
response_metadata={"model_version": "gpt-5.4-2026-03-01"},
|
||||
endpoint_api_format="openai:cli",
|
||||
)
|
||||
|
||||
metadata = captured["metadata"]
|
||||
assert metadata["model_version"] == "gpt-5.4-2026-03-01"
|
||||
assert "response" not in metadata
|
||||
assert metadata["cache_fingerprint"]["provider_api_format"] == "openai:cli"
|
||||
assert metadata["cache_fingerprint"]["prompt_cache_key"] == "pcache-123"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_message_telemetry_record_failure_keeps_request_metadata_and_adds_fingerprint(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
captured: dict[str, Any] = {}
|
||||
|
||||
async def _fake_record_usage(**kwargs: Any) -> Any:
|
||||
captured.update(kwargs)
|
||||
return SimpleNamespace()
|
||||
|
||||
monkeypatch.setattr(UsageService, "record_usage", _fake_record_usage)
|
||||
|
||||
telemetry = MessageTelemetry(
|
||||
db=SimpleNamespace(), # type: ignore[arg-type]
|
||||
user=None,
|
||||
api_key=None,
|
||||
request_id="req-cache-fingerprint-fail",
|
||||
client_ip="127.0.0.1",
|
||||
)
|
||||
|
||||
await telemetry.record_failure(
|
||||
provider="openai",
|
||||
model="gpt-5.4",
|
||||
response_time_ms=10,
|
||||
status_code=502,
|
||||
error_message="upstream failed",
|
||||
request_body={"messages": [{"role": "user", "content": "hello"}]},
|
||||
request_headers={"user-agent": "codex desktop"},
|
||||
is_stream=False,
|
||||
provider_request_body=_build_openai_cli_body(),
|
||||
request_metadata={"perf": {"ttfb_ms": 12}},
|
||||
endpoint_api_format="openai:cli",
|
||||
)
|
||||
|
||||
metadata = captured["metadata"]
|
||||
assert metadata["perf"]["ttfb_ms"] == 12
|
||||
assert metadata["cache_fingerprint"]["provider_api_format"] == "openai:cli"
|
||||
assert metadata["cache_fingerprint"]["prompt_cache_key"] == "pcache-123"
|
||||
Reference in New Issue
Block a user