64 lines
2.2 KiB
Python
64 lines
2.2 KiB
Python
from decimal import Decimal
|
|
import unittest
|
|
|
|
from core.loop import _extract_usage_details
|
|
from core.storage.usage import _fallback_chat_cost_cny
|
|
|
|
|
|
class UsageAccountingTests(unittest.TestCase):
|
|
def test_extract_usage_details_includes_cache_tokens(self) -> None:
|
|
usage = {
|
|
"prompt_tokens": 1200,
|
|
"completion_tokens": 80,
|
|
"prompt_cache_hit_tokens": 900,
|
|
"prompt_cache_miss_tokens": 300,
|
|
"completion_tokens_details": {"reasoning_tokens": 12},
|
|
}
|
|
|
|
details = _extract_usage_details(usage)
|
|
|
|
self.assertEqual(details["tokens_in"], 1200)
|
|
self.assertEqual(details["tokens_out"], 80)
|
|
self.assertEqual(details["cache_hit_tokens"], 900)
|
|
self.assertEqual(details["cache_miss_tokens"], 300)
|
|
self.assertEqual(details["reasoning_tokens"], 12)
|
|
|
|
def test_fallback_chat_cost_uses_price_snapshots(self) -> None:
|
|
cost = _fallback_chat_cost_cny(
|
|
prompt_tokens=1_000_000,
|
|
completion_tokens=500_000,
|
|
input_cny_per_mtoken=1.0,
|
|
output_cny_per_mtoken=10.0,
|
|
)
|
|
|
|
self.assertEqual(cost, Decimal("6.000000"))
|
|
|
|
def test_fallback_chat_cost_discounts_cache_hits(self) -> None:
|
|
# 100 万输入里 80 万命中缓存(0.1 价),20 万未命中(1.0 价),50 万输出(10 价)
|
|
cost = _fallback_chat_cost_cny(
|
|
prompt_tokens=1_000_000,
|
|
completion_tokens=500_000,
|
|
input_cny_per_mtoken=1.0,
|
|
output_cny_per_mtoken=10.0,
|
|
cache_hit_tokens=800_000,
|
|
cache_hit_cny_per_mtoken=0.1,
|
|
)
|
|
# 0.2(miss) + 0.08(hit) + 5.0(out) = 5.28
|
|
self.assertEqual(cost, Decimal("5.280000"))
|
|
|
|
def test_fallback_chat_cost_no_cache_price_charges_full(self) -> None:
|
|
# 未配缓存价(0)→ 命中段不打折,按 input 全价(老行为,绝不少记)
|
|
cost = _fallback_chat_cost_cny(
|
|
prompt_tokens=1_000_000,
|
|
completion_tokens=0,
|
|
input_cny_per_mtoken=1.0,
|
|
output_cny_per_mtoken=10.0,
|
|
cache_hit_tokens=900_000,
|
|
cache_hit_cny_per_mtoken=0.0,
|
|
)
|
|
self.assertEqual(cost, Decimal("1.000000"))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|