mirror of
https://github.com/NCBM/plyngent.git
synced 2026-11-07 23:47:54 +08:00
core/agent: normalize provider usage structs in token_usage_from_api
This commit is contained in:
@@ -64,16 +64,6 @@ def _encode_tool_arguments(obj: dict[str, Any]) -> str:
|
|||||||
return json.dumps(obj, ensure_ascii=False, separators=(",", ":"))
|
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(
|
def tool_items_to_anthropic_tools(
|
||||||
tools: Sequence[AnyToolItem] | None,
|
tools: Sequence[AnyToolItem] | None,
|
||||||
) -> list[AnthropicToolDefinition]:
|
) -> 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:
|
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."""
|
"""Normalize Anthropic usage Struct/dict to chat-compatible usage fields."""
|
||||||
if usage is None:
|
from .usage import token_usage_from_api
|
||||||
return None
|
|
||||||
if isinstance(usage, AnthropicUsage):
|
parsed = token_usage_from_api(usage)
|
||||||
prompt = _as_nonneg_int(usage.input_tokens)
|
if parsed is None:
|
||||||
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:
|
|
||||||
return None
|
return None
|
||||||
return {
|
return {
|
||||||
"prompt_tokens": prompt,
|
"prompt_tokens": parsed.prompt_tokens,
|
||||||
"completion_tokens": completion,
|
"completion_tokens": parsed.completion_tokens,
|
||||||
"total_tokens": prompt + completion,
|
"total_tokens": parsed.total_tokens,
|
||||||
"input_tokens": prompt,
|
"input_tokens": parsed.prompt_tokens,
|
||||||
"output_tokens": completion,
|
"output_tokens": parsed.completion_tokens,
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
||||||
|
|||||||
@@ -2,6 +2,7 @@ from __future__ import annotations
|
|||||||
|
|
||||||
from typing import TYPE_CHECKING, cast
|
from typing import TYPE_CHECKING, cast
|
||||||
|
|
||||||
|
import msgspec
|
||||||
from msgspec import UNSET, Struct
|
from msgspec import UNSET, Struct
|
||||||
|
|
||||||
from plyngent.agent.budget import estimate_message_chars, estimate_messages_chars
|
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:
|
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``)
|
Accepts chat completions fields (``prompt_tokens`` / ``completion_tokens``),
|
||||||
and Responses fields (``input_tokens`` / ``output_tokens``).
|
Responses fields (``input_tokens`` / ``output_tokens``), and msgspec Structs
|
||||||
|
carrying either pair (e.g. ``AnthropicUsage``).
|
||||||
"""
|
"""
|
||||||
if usage is None or usage is UNSET:
|
if usage is None or usage is UNSET:
|
||||||
return None
|
return None
|
||||||
|
if isinstance(usage, Struct):
|
||||||
|
usage = msgspec.to_builtins(usage)
|
||||||
if not isinstance(usage, dict):
|
if not isinstance(usage, dict):
|
||||||
return None
|
return None
|
||||||
raw = cast("dict[str, object]", usage)
|
raw = cast("dict[str, object]", usage)
|
||||||
|
|||||||
@@ -58,6 +58,23 @@ def test_token_usage_from_api_responses_fields() -> None:
|
|||||||
assert u.total_tokens == 9
|
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:
|
def test_chars_to_tokens() -> None:
|
||||||
assert chars_to_tokens(0) == 0
|
assert chars_to_tokens(0) == 0
|
||||||
assert chars_to_tokens(1) == 1
|
assert chars_to_tokens(1) == 1
|
||||||
|
|||||||
Reference in New Issue
Block a user