150 lines
4.8 KiB
Python
150 lines
4.8 KiB
Python
"""Token usage extraction helpers.
|
|
|
|
The project intentionally avoids model/provider-name special cases. These helpers
|
|
normalize common OpenAI-compatible and provider-specific response shapes by
|
|
looking for usage-like payloads in known response locations and field aliases.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from typing import Any, Dict, Iterable, Optional
|
|
|
|
|
|
INPUT_TOKEN_KEYS = (
|
|
"prompt_tokens",
|
|
"input_tokens",
|
|
"inputTokens",
|
|
"promptTokens",
|
|
"prefill_tokens",
|
|
)
|
|
OUTPUT_TOKEN_KEYS = (
|
|
"completion_tokens",
|
|
"output_tokens",
|
|
"outputTokens",
|
|
"completionTokens",
|
|
"generated_tokens",
|
|
"generatedTokens",
|
|
)
|
|
TOTAL_TOKEN_KEYS = (
|
|
"total_tokens",
|
|
"totalTokens",
|
|
"total_token_count",
|
|
"totalTokenCount",
|
|
)
|
|
CURRENT_CONTEXT_KEYS = (
|
|
"current_context_tokens",
|
|
"currentContextTokens",
|
|
"context_tokens",
|
|
"contextTokens",
|
|
)
|
|
KNOWN_CONTAINER_KEYS = {
|
|
"usage",
|
|
"token_usage",
|
|
"tokenUsage",
|
|
"token_usages",
|
|
"response_metadata",
|
|
"responseMetadata",
|
|
"metadata",
|
|
"meta",
|
|
}
|
|
|
|
|
|
def _to_int(value: Any) -> Optional[int]:
|
|
if value is None or isinstance(value, bool):
|
|
return None
|
|
try:
|
|
number = int(value)
|
|
except (TypeError, ValueError):
|
|
return None
|
|
return number if number >= 0 else None
|
|
|
|
|
|
def _first_int(payload: Dict[str, Any], keys: Iterable[str]) -> Optional[int]:
|
|
for key in keys:
|
|
if key in payload:
|
|
value = _to_int(payload.get(key))
|
|
if value is not None:
|
|
return value
|
|
return None
|
|
|
|
|
|
def normalize_usage_payload(raw: Any) -> Optional[Dict[str, int]]:
|
|
if not isinstance(raw, dict):
|
|
return None
|
|
|
|
prompt_tokens = _first_int(raw, INPUT_TOKEN_KEYS)
|
|
completion_tokens = _first_int(raw, OUTPUT_TOKEN_KEYS)
|
|
total_tokens = _first_int(raw, TOTAL_TOKEN_KEYS)
|
|
current_context_tokens = _first_int(raw, CURRENT_CONTEXT_KEYS)
|
|
|
|
prompt_details = raw.get("prompt_tokens_details") or raw.get("input_tokens_details")
|
|
if isinstance(prompt_details, dict):
|
|
cached = _first_int(prompt_details, ("cached_tokens", "cachedTokens"))
|
|
# cached tokens are still part of prompt tokens in most APIs. Keep the
|
|
# detail accessible for callers that need it, but do not add it again.
|
|
completion_details = raw.get("completion_tokens_details") or raw.get("output_tokens_details")
|
|
if isinstance(completion_details, dict):
|
|
reasoning = _first_int(completion_details, ("reasoning_tokens", "reasoningTokens"))
|
|
if completion_tokens is None and reasoning is not None:
|
|
completion_tokens = reasoning
|
|
|
|
if prompt_tokens is None and completion_tokens is None and total_tokens is None:
|
|
return None
|
|
|
|
if prompt_tokens is None:
|
|
prompt_tokens = max(0, (total_tokens or 0) - (completion_tokens or 0)) if total_tokens is not None else 0
|
|
if completion_tokens is None:
|
|
completion_tokens = max(0, (total_tokens or 0) - prompt_tokens) if total_tokens is not None else 0
|
|
if total_tokens is None:
|
|
total_tokens = prompt_tokens + completion_tokens
|
|
if current_context_tokens is None:
|
|
current_context_tokens = prompt_tokens
|
|
|
|
return {
|
|
"prompt_tokens": int(prompt_tokens),
|
|
"completion_tokens": int(completion_tokens),
|
|
"total_tokens": int(total_tokens),
|
|
"current_context_tokens": int(current_context_tokens),
|
|
}
|
|
|
|
|
|
def _usage_score(payload: Dict[str, int]) -> int:
|
|
return int(payload.get("total_tokens", 0)) + int(payload.get("prompt_tokens", 0)) + int(payload.get("completion_tokens", 0))
|
|
|
|
|
|
def extract_usage_payload(obj: Any) -> Optional[Dict[str, int]]:
|
|
"""Find and normalize the best token usage payload in a response chunk/object."""
|
|
best: Optional[Dict[str, int]] = None
|
|
|
|
def consider(value: Any) -> None:
|
|
nonlocal best
|
|
normalized = normalize_usage_payload(value)
|
|
if not normalized:
|
|
return
|
|
if best is None or _usage_score(normalized) >= _usage_score(best):
|
|
best = normalized
|
|
|
|
def walk(value: Any, *, depth: int = 0, in_known_container: bool = False) -> None:
|
|
if depth > 8:
|
|
return
|
|
if isinstance(value, dict):
|
|
if in_known_container:
|
|
consider(value)
|
|
else:
|
|
# Also accept dicts that directly look like usage payloads.
|
|
consider(value)
|
|
for key, child in value.items():
|
|
child_known = in_known_container or key in KNOWN_CONTAINER_KEYS
|
|
if key in KNOWN_CONTAINER_KEYS:
|
|
consider(child)
|
|
walk(child, depth=depth + 1, in_known_container=child_known)
|
|
elif isinstance(value, list):
|
|
for child in value:
|
|
walk(child, depth=depth + 1, in_known_container=in_known_container)
|
|
|
|
walk(obj)
|
|
return best
|
|
|
|
|
|
__all__ = ["extract_usage_payload", "normalize_usage_payload"]
|