diff --git a/src/plyngent/agent/usage.py b/src/plyngent/agent/usage.py new file mode 100644 index 0000000..0b8048f --- /dev/null +++ b/src/plyngent/agent/usage.py @@ -0,0 +1,56 @@ +from __future__ import annotations + +from typing import cast + +from msgspec import UNSET, Struct + + +class TokenUsage(Struct, omit_defaults=True): + """Token counts from a provider ``usage`` object (OpenAI-compatible).""" + + prompt_tokens: int = 0 + completion_tokens: int = 0 + total_tokens: int = 0 + + def add(self, other: TokenUsage) -> TokenUsage: + return TokenUsage( + prompt_tokens=self.prompt_tokens + other.prompt_tokens, + completion_tokens=self.completion_tokens + other.completion_tokens, + total_tokens=self.total_tokens + other.total_tokens, + ) + + def is_zero(self) -> bool: + return self.prompt_tokens == 0 and self.completion_tokens == 0 and self.total_tokens == 0 + + def format_line(self) -> str: + return ( + f"tokens prompt={self.prompt_tokens} completion={self.completion_tokens} " + f"total={self.total_tokens}" + ) + + +def _as_nonneg_int(value: object) -> int: + if isinstance(value, bool): + return 0 + if isinstance(value, int): + return max(0, value) + if isinstance(value, float): + return max(0, int(value)) + return 0 + + +def token_usage_from_api(usage: object) -> TokenUsage | None: + """Parse OpenAI-style usage dict; return None if missing/empty.""" + if usage is None or usage is UNSET: + return None + if not isinstance(usage, dict): + return None + raw = cast("dict[str, object]", usage) + prompt = _as_nonneg_int(raw.get("prompt_tokens")) + completion = _as_nonneg_int(raw.get("completion_tokens")) + total = _as_nonneg_int(raw.get("total_tokens")) + if total == 0 and (prompt or completion): + total = prompt + completion + if prompt == 0 and completion == 0 and total == 0: + return None + return TokenUsage(prompt_tokens=prompt, completion_tokens=completion, total_tokens=total) diff --git a/tests/test_agent/test_usage.py b/tests/test_agent/test_usage.py new file mode 100644 index 0000000..06dd42d --- /dev/null +++ b/tests/test_agent/test_usage.py @@ -0,0 +1,39 @@ +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