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

fix: 统一缓存token计费口径，避免OpenAI错误计费 (#144)

Browse Source 
* fix: 统一缓存token计费口径，避免OpenAI错误计费

* fix: 添加 Gemini 缓存 token 归一化支持，优化代码结构

- Gemini 的 promptTokenCount 包含 cachedContentTokenCount，需要扣除
- 合并重复的 helper 函数为 _get_api_family()
- 将 import 移到文件顶部
- 补充 Gemini 测试用例

---------

Co-authored-by: fawney19 <elky0401@gmail.com>
This commit is contained in:
Yorha
2026-02-04 16:04:09 +08:00
committed by GitHub
parent f64631a3a3
commit 26a0f99f8f
5 changed files with 130 additions and 32 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -126,16 +126,7 @@ class MessageTelemetry:
# Provider 响应元数据（如 Gemini 的 modelVersion）
response_metadata: dict[str, Any] | None = None,
) -> float:
total_cost = await self.calculate_cost(
provider,
model,
input_tokens=input_tokens,
output_tokens=output_tokens,
cache_creation_tokens=cache_creation_tokens,
cache_read_tokens=cache_read_tokens,
)
await UsageService.record_usage(
usage = await UsageService.record_usage(
db=self.db,
user=self.user,
api_key=self.api_key,
@@ -170,6 +161,8 @@ class MessageTelemetry:
metadata=response_metadata,
)
total_cost = float(getattr(usage, "total_cost_usd", 0.0) or 0.0)
if self.user and self.api_key:
audit_service.log_api_request(
db=self.db,
@@ -181,8 +174,8 @@ class MessageTelemetry:
success=True,
ip_address=self.client_ip,
status_code=status_code,
input_tokens=input_tokens,
output_tokens=output_tokens,
input_tokens=getattr(usage, "input_tokens", input_tokens),
output_tokens=getattr(usage, "output_tokens", output_tokens),
cost_usd=total_cost,
)

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

@@ -1903,10 +1903,6 @@ class CliMessageHandlerBase(BaseMessageHandler):
logger.warning(f"[{ctx.request_id}] 流式请求失败，未选中提供商")
return
# Claude API 的 input_tokens 已经是非缓存部分，不需要再减去 cached_tokens
# 实际计费的输入 tokens = input_tokens + cache_creation_tokens（缓存读取免费或折扣）
actual_input_tokens = ctx.input_tokens
# 获取新的 DB session
db_gen = get_db()
bg_db = next(db_gen)
@@ -1961,7 +1957,7 @@ class CliMessageHandlerBase(BaseMessageHandler):
is_stream=True,
api_format=ctx.api_format,
provider_request_headers=ctx.provider_request_headers,
input_tokens=actual_input_tokens,
input_tokens=ctx.input_tokens,
output_tokens=ctx.output_tokens,
cache_creation_tokens=ctx.cache_creation_tokens,
cache_read_tokens=ctx.cached_tokens,
@@ -1975,7 +1971,7 @@ class CliMessageHandlerBase(BaseMessageHandler):
logger.debug(f"{self.FORMAT_ID} 流式响应被客户端取消")
logger.info(
f"[CANCEL] {self.request_id[:8]} | {ctx.model} | {ctx.provider_name} | {response_time_ms}ms | "
f"{ctx.status_code} | in:{actual_input_tokens} out:{ctx.output_tokens} cache:{ctx.cached_tokens}"
f"{ctx.status_code} | in:{ctx.input_tokens} out:{ctx.output_tokens} cache:{ctx.cached_tokens}"
)
else:
# 服务端/上游异常：记录为失败
@@ -1991,7 +1987,7 @@ class CliMessageHandlerBase(BaseMessageHandler):
api_format=ctx.api_format,
provider_request_headers=ctx.provider_request_headers,
# 预估 token 信息（来自 message_start 事件）
input_tokens=actual_input_tokens,
input_tokens=ctx.input_tokens,
output_tokens=ctx.output_tokens,
cache_creation_tokens=ctx.cache_creation_tokens,
cache_read_tokens=ctx.cached_tokens,
@@ -2007,7 +2003,7 @@ class CliMessageHandlerBase(BaseMessageHandler):
logger.debug(f"{self.FORMAT_ID} 流式响应中断")
logger.info(
f"[FAIL] {self.request_id[:8]} | {ctx.model} | {ctx.provider_name} | {response_time_ms}ms | "
f"{ctx.status_code} | in:{actual_input_tokens} out:{ctx.output_tokens} cache:{ctx.cached_tokens}"
f"{ctx.status_code} | in:{ctx.input_tokens} out:{ctx.output_tokens} cache:{ctx.cached_tokens}"
)
else:
# 在记录统计前，允许子类从 parsed_chunks 中提取额外的元数据
@@ -2026,12 +2022,12 @@ class CliMessageHandlerBase(BaseMessageHandler):
logger.debug(
f"[{ctx.request_id}] 开始记录 Usage: "
f"provider={ctx.provider_name}, model={ctx.model}, "
f"in={actual_input_tokens}, out={ctx.output_tokens}"
f"in={ctx.input_tokens}, out={ctx.output_tokens}"
)
total_cost = await bg_telemetry.record_success(
provider=ctx.provider_name,
model=ctx.model,
input_tokens=actual_input_tokens,
input_tokens=ctx.input_tokens,
output_tokens=ctx.output_tokens,
response_time_ms=response_time_ms,
first_byte_time_ms=ctx.first_byte_time_ms, # 传递首字时间
@@ -2492,8 +2488,6 @@ class CliMessageHandlerBase(BaseMessageHandler):
output_tokens = usage.get("output_tokens", 0)
cached_tokens = usage.get("cache_read_tokens", 0)
cache_creation_tokens = usage.get("cache_creation_tokens", 0)
# Claude API 的 input_tokens 已经是非缓存部分，不需要再减去 cached_tokens
actual_input_tokens = input_tokens
output_text = self.parser.extract_text_content(response_json)[:200]
@@ -2507,7 +2501,7 @@ class CliMessageHandlerBase(BaseMessageHandler):
total_cost = await self.telemetry.record_success(
provider=provider_name,
model=model,
input_tokens=actual_input_tokens,
input_tokens=input_tokens,
output_tokens=output_tokens,
response_time_ms=response_time_ms,
status_code=status_code,

47
src/services/billing/token_normalization.py Normal file
View File

@@ -0,0 +1,47 @@
"""
计费相关 token 归一化工具。
"""
from __future__ import annotations
from src.core.api_format.enums import ApiFamily
from src.core.api_format.signature import parse_signature_key
def _get_api_family(api_format: str | None) -> ApiFamily | None:
"""解析 api_format 字符串，返回对应的 ApiFamily 枚举。"""
if not api_format:
return None
text = str(api_format).strip()
if not text:
return None
sig = parse_signature_key(text)
return sig.api_family
def normalize_input_tokens_for_billing(
api_format: str | None,
input_tokens: int,
cache_read_tokens: int,
) -> int:
"""
归一化 `input_tokens`，使其在计费中表示"非缓存输入 token"。
计费口径：`input_tokens`=非缓存输入 token；`cache_read_tokens`=缓存命中 token（折扣/免费维度）。
- Claude 系：保持上游口径（不扣除），因为 Claude API 的 input_tokens 本身就不包含缓存部分。
- OpenAI 系：`input_tokens` 包含缓存命中部分，需要扣除 `cache_read_tokens`。
- Gemini 系：`promptTokenCount` 包含 `cachedContentTokenCount`，需要扣除。
"""
if input_tokens <= 0:
return 0 if input_tokens == 0 else input_tokens
if cache_read_tokens <= 0:
return input_tokens
api_family = _get_api_family(api_format)
if api_family == ApiFamily.CLAUDE:
return input_tokens
if api_family in (ApiFamily.OPENAI, ApiFamily.GEMINI):
return max(input_tokens - cache_read_tokens, 0)
# 未知格式，保守处理，不扣除
return input_tokens

34
src/services/usage/service.py
View File

@@ -25,6 +25,7 @@ from src.models.database import (
User,
UserRole,
)
from src.services.billing.token_normalization import normalize_input_tokens_for_billing
from src.services.model.cost import ModelCostService
from src.services.system.config import SystemConfigService
from src.services.usage.error_classifier import classify_error
@@ -843,9 +844,28 @@ class UsageService:
Returns:
(usage_params 字典, total_cost 总成本)
"""
# 计费口径以 Provider 为准（优先 endpoint_api_format）
billing_api_format: str | None = None
if params.endpoint_api_format:
try:
billing_api_format = normalize_signature_key(str(params.endpoint_api_format))
except Exception:
billing_api_format = None
if billing_api_format is None and params.api_format:
try:
billing_api_format = normalize_signature_key(str(params.api_format))
except Exception:
billing_api_format = None
input_tokens_for_billing = normalize_input_tokens_for_billing(
billing_api_format,
params.input_tokens,
params.cache_read_input_tokens,
)
# 获取费率倍数和是否免费套餐（传递 api_format 支持按格式配置的倍率）
actual_rate_multiplier, is_free_tier = await cls._get_rate_multiplier_and_free_tier(
params.db, params.provider_api_key_id, params.provider_id, params.api_format
params.db, params.provider_api_key_id, params.provider_id, billing_api_format
)
metadata = dict(params.metadata or {})
@@ -884,7 +904,7 @@ class UsageService:
request_count = 0 if is_failed_request else 1
dims: dict[str, Any] = {
"input_tokens": params.input_tokens,
"input_tokens": input_tokens_for_billing,
"output_tokens": params.output_tokens,
"cache_creation_input_tokens": params.cache_creation_input_tokens,
"cache_read_input_tokens": params.cache_read_input_tokens,
@@ -956,11 +976,11 @@ class UsageService:
db=params.db,
provider=params.provider,
model=params.model,
input_tokens=params.input_tokens,
input_tokens=input_tokens_for_billing,
output_tokens=params.output_tokens,
cache_creation_input_tokens=params.cache_creation_input_tokens,
cache_read_input_tokens=params.cache_read_input_tokens,
api_format=params.api_format,
api_format=billing_api_format,
cache_ttl_minutes=params.cache_ttl_minutes,
use_tiered_pricing=params.use_tiered_pricing,
is_failed_request=is_failed_request,
@@ -989,8 +1009,8 @@ class UsageService:
provider_id=params.provider_id,
model=params.model,
task_type=billing_task_type,
api_format=params.api_format,
input_tokens=params.input_tokens,
api_format=billing_api_format,
input_tokens=input_tokens_for_billing,
output_tokens=params.output_tokens,
cache_creation_input_tokens=params.cache_creation_input_tokens,
cache_read_input_tokens=params.cache_read_input_tokens,
@@ -1019,7 +1039,7 @@ class UsageService:
api_key=params.api_key,
provider=params.provider,
model=params.model,
input_tokens=params.input_tokens,
input_tokens=input_tokens_for_billing,
output_tokens=params.output_tokens,
cache_creation_input_tokens=params.cache_creation_input_tokens,
cache_read_input_tokens=params.cache_read_input_tokens,

44
tests/services/billing/test_token_normalization.py Normal file
View File

@@ -0,0 +1,44 @@
from src.services.billing.token_normalization import normalize_input_tokens_for_billing
from src.services.billing.usage_mapper import UsageMapper
class TestNormalizeInputTokensForBilling:
def test_openai_family_subtracts_cached_tokens(self) -> None:
assert normalize_input_tokens_for_billing("openai:cli", 160_070, 81_664) == 78_406
def test_claude_family_does_not_change(self) -> None:
assert normalize_input_tokens_for_billing("claude:cli", 160_070, 81_664) == 160_070
def test_gemini_family_subtracts_cached_tokens(self) -> None:
# Gemini 的 promptTokenCount 包含 cachedContentTokenCount，需要扣除
assert normalize_input_tokens_for_billing("gemini:chat", 323_392, 323_384) == 8
assert normalize_input_tokens_for_billing("gemini:cli", 100, 20) == 80
def test_missing_format_does_not_change(self) -> None:
assert normalize_input_tokens_for_billing(None, 100, 20) == 100
assert normalize_input_tokens_for_billing("", 100, 20) == 100
def test_clamps_when_cached_tokens_exceed_input(self) -> None:
assert normalize_input_tokens_for_billing("openai:cli", 10, 20) == 0
assert normalize_input_tokens_for_billing("gemini:chat", 10, 20) == 0
class TestUsageMapperOpenAICacheTokens:
def test_openai_mapping_maps_cached_tokens_details(self) -> None:
raw_usage = {
"prompt_tokens": 100,
"completion_tokens": 50,
"prompt_tokens_details": {"cached_tokens": 20},
}
usage = UsageMapper.map(raw_usage, api_format="openai:chat")
assert usage.input_tokens == 100
assert usage.output_tokens == 50
assert usage.cache_read_tokens == 20
def test_openai_mapping_without_cached_tokens_is_unchanged(self) -> None:
raw_usage = {"prompt_tokens": 100, "completion_tokens": 50}
usage = UsageMapper.map(raw_usage, api_format="openai:chat")
assert usage.input_tokens == 100
assert usage.output_tokens == 50
assert usage.cache_read_tokens == 0