diff --git a/src/plyngent/agent/messages_bridge.py b/src/plyngent/agent/messages_bridge.py index a6f50c6..d661c15 100644 --- a/src/plyngent/agent/messages_bridge.py +++ b/src/plyngent/agent/messages_bridge.py @@ -64,16 +64,6 @@ def _encode_tool_arguments(obj: dict[str, Any]) -> str: return json.dumps(obj, ensure_ascii=False, separators=(",", ":")) -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 tool_items_to_anthropic_tools( tools: Sequence[AnyToolItem] | None, ) -> list[AnthropicToolDefinition]: @@ -235,23 +225,17 @@ def anthropic_stop_to_finish_reason( def anthropic_usage_to_dict(usage: AnthropicUsage | dict[str, Any] | None) -> dict[str, Any] | None: """Normalize Anthropic usage Struct/dict to chat-compatible usage fields.""" - if usage is None: - return None - if isinstance(usage, AnthropicUsage): - prompt = _as_nonneg_int(usage.input_tokens) - completion = _as_nonneg_int(usage.output_tokens) - else: - raw = cast("dict[str, object]", usage) - prompt = _as_nonneg_int(raw.get("input_tokens") or raw.get("prompt_tokens") or 0) - completion = _as_nonneg_int(raw.get("output_tokens") or raw.get("completion_tokens") or 0) - if prompt == 0 and completion == 0: + from .usage import token_usage_from_api + + parsed = token_usage_from_api(usage) + if parsed is None: return None return { - "prompt_tokens": prompt, - "completion_tokens": completion, - "total_tokens": prompt + completion, - "input_tokens": prompt, - "output_tokens": completion, + "prompt_tokens": parsed.prompt_tokens, + "completion_tokens": parsed.completion_tokens, + "total_tokens": parsed.total_tokens, + "input_tokens": parsed.prompt_tokens, + "output_tokens": parsed.completion_tokens, } diff --git a/src/plyngent/agent/usage.py b/src/plyngent/agent/usage.py index 116c209..6d01f32 100644 --- a/src/plyngent/agent/usage.py +++ b/src/plyngent/agent/usage.py @@ -2,6 +2,7 @@ from __future__ import annotations from typing import TYPE_CHECKING, cast +import msgspec from msgspec import UNSET, Struct from plyngent.agent.budget import estimate_message_chars, estimate_messages_chars @@ -86,13 +87,16 @@ def _as_nonneg_int(value: object) -> int: def token_usage_from_api(usage: object) -> TokenUsage | None: - """Parse OpenAI-style usage dict; return None if missing/empty. + """Parse API usage into :class:`TokenUsage`; return None if missing/empty. - Accepts chat completions fields (``prompt_tokens`` / ``completion_tokens``) - and Responses fields (``input_tokens`` / ``output_tokens``). + Accepts chat completions fields (``prompt_tokens`` / ``completion_tokens``), + Responses fields (``input_tokens`` / ``output_tokens``), and msgspec Structs + carrying either pair (e.g. ``AnthropicUsage``). """ if usage is None or usage is UNSET: return None + if isinstance(usage, Struct): + usage = msgspec.to_builtins(usage) if not isinstance(usage, dict): return None raw = cast("dict[str, object]", usage) diff --git a/tests/test_agent/test_usage.py b/tests/test_agent/test_usage.py index 54bdf82..c05f8eb 100644 --- a/tests/test_agent/test_usage.py +++ b/tests/test_agent/test_usage.py @@ -58,6 +58,23 @@ def test_token_usage_from_api_responses_fields() -> None: assert u.total_tokens == 9 +def test_token_usage_from_api_struct() -> None: + from plyngent.lmproto.anthropic.model import AnthropicUsage + + u = token_usage_from_api(AnthropicUsage(input_tokens=7, output_tokens=2)) + assert u is not None + assert u.prompt_tokens == 7 + assert u.completion_tokens == 2 + assert u.total_tokens == 9 + assert u.source == "api" + + +def test_token_usage_from_api_empty_struct_returns_none() -> None: + from plyngent.lmproto.anthropic.model import AnthropicUsage + + assert token_usage_from_api(AnthropicUsage()) is None + + def test_chars_to_tokens() -> None: assert chars_to_tokens(0) == 0 assert chars_to_tokens(1) == 1