core/agent: add TokenUsage and API usage parser

Parse OpenAI-compatible usage dicts into prompt/completion/total counts
for session and turn aggregation.

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)

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