codekingpro/portable-devtools

1"""Token usage extraction for Google ADK."""2 3from __future__ import annotations4 5from typing import Any, Optional6 7 8def extract_usage_from_response(llm_response: Any) -> dict[str, Any]:9    """Extract token usage from LlmResponse."""10    usage: dict[str, Any] = {}11    usage_metadata = getattr(llm_response, "usage_metadata", None)12    if not usage_metadata:13        return usage14 15    if (v := getattr(usage_metadata, "prompt_token_count", None)) is not None:16        usage["input_tokens"] = int(v)17    if (v := getattr(usage_metadata, "candidates_token_count", None)) is not None:18        usage["output_tokens"] = int(v)19    if (v := getattr(usage_metadata, "total_token_count", None)) is not None:20        usage["total_tokens"] = int(v)21    if (v := getattr(usage_metadata, "cached_content_token_count", None)) is not None:22        usage.setdefault("input_token_details", {})["cache_read"] = int(v)23    if (v := getattr(usage_metadata, "thoughts_token_count", None)) is not None:24        usage.setdefault("output_token_details", {})["reasoning"] = int(v)25 26    return usage27 28 29def extract_model_name(llm_request: Any) -> Optional[str]:30    """Extract the model name from an LlmRequest."""31    if config := getattr(llm_request, "config", None):32        if model := getattr(config, "model", None):33            return str(model)34    if model := getattr(llm_request, "model", None):35        return str(model)36    return None37