Spaces:
Running
Running
| from typing import Any | |
| def extract_token_usage(message: Any) -> dict[str, int]: | |
| usage = getattr(message, "usage_metadata", None) | |
| if isinstance(usage, dict): | |
| return { | |
| "input_tokens": int(usage.get("input_tokens", 0) or 0), | |
| "output_tokens": int(usage.get("output_tokens", 0) or 0), | |
| "total_tokens": int(usage.get("total_tokens", 0) or 0), | |
| } | |
| response_metadata = getattr(message, "response_metadata", None) | |
| if isinstance(response_metadata, dict): | |
| token_usage = response_metadata.get("token_usage") | |
| if isinstance(token_usage, dict): | |
| prompt_tokens = int(token_usage.get("prompt_tokens", 0) or 0) | |
| completion_tokens = int(token_usage.get("completion_tokens", 0) or 0) | |
| total_tokens = int(token_usage.get("total_tokens", prompt_tokens + completion_tokens) or 0) | |
| return { | |
| "input_tokens": prompt_tokens, | |
| "output_tokens": completion_tokens, | |
| "total_tokens": total_tokens, | |
| } | |
| return {"input_tokens": 0, "output_tokens": 0, "total_tokens": 0} | |
| def combine_usage(*items: dict[str, int]) -> dict[str, int]: | |
| result = {"input_tokens": 0, "output_tokens": 0, "total_tokens": 0} | |
| for item in items: | |
| result["input_tokens"] += int(item.get("input_tokens", 0) or 0) | |
| result["output_tokens"] += int(item.get("output_tokens", 0) or 0) | |
| result["total_tokens"] += int(item.get("total_tokens", 0) or 0) | |
| return result | |