mirror of
https://github.com/NCBM/plyngent.git
synced 2026-07-23 14:14:57 +08:00
NCBM
2ae7aa0e05
Parse OpenAI-compatible usage dicts into prompt/completion/total counts for session and turn aggregation.
40 lines
1.3 KiB
Python
40 lines
1.3 KiB
Python
from __future__ import annotations
|
|
|
|
from msgspec import UNSET
|
|
|
|
from plyngent.agent.usage import TokenUsage, token_usage_from_api
|
|
|
|
|
|
def test_token_usage_add() -> None:
|
|
a = TokenUsage(prompt_tokens=10, completion_tokens=5, total_tokens=15)
|
|
b = TokenUsage(prompt_tokens=3, completion_tokens=2, total_tokens=5)
|
|
c = a.add(b)
|
|
assert c.prompt_tokens == 13
|
|
assert c.completion_tokens == 7
|
|
assert c.total_tokens == 20
|
|
assert a.prompt_tokens == 10 # immutable-ish via new struct
|
|
|
|
|
|
def test_token_usage_from_api() -> None:
|
|
assert token_usage_from_api(None) is None
|
|
assert token_usage_from_api(UNSET) is None
|
|
assert token_usage_from_api({}) is None
|
|
u = token_usage_from_api({"prompt_tokens": 11, "completion_tokens": 4, "total_tokens": 15})
|
|
assert u is not None
|
|
assert u.prompt_tokens == 11
|
|
assert u.completion_tokens == 4
|
|
assert u.total_tokens == 15
|
|
|
|
|
|
def test_token_usage_from_api_infers_total() -> None:
|
|
u = token_usage_from_api({"prompt_tokens": 2, "completion_tokens": 3})
|
|
assert u is not None
|
|
assert u.total_tokens == 5
|
|
|
|
|
|
def test_format_line() -> None:
|
|
line = TokenUsage(prompt_tokens=1, completion_tokens=2, total_tokens=3).format_line()
|
|
assert "prompt=1" in line
|
|
assert "completion=2" in line
|
|
assert "total=3" in line
|