From 7fd7d4397333dea7ad14d68f8d364422689e5058 Mon Sep 17 00:00:00 2001 From: juser0719 Date: Thu, 15 Oct 2026 16:12:37 +0900 Subject: [PATCH] feat(usage): break Ask AI token counts down by type The Usage dashboard's Ask AI tools read the daily activity aggregates, which carry prompt, completion, cache-read and cache-creation token counts per model and per team/tag, but the summaries only handed the model total_tokens. Report the four counters in the header, per-model lines and per-entity lines, and say so in the tool description so the model asks for them. --- .../usage_endpoints/ai_usage_chat.py | 56 +++++++++++++++++-- .../usage_endpoints/test_ai_usage_chat.py | 23 ++++++++ 2 files changed, 73 insertions(+), 6 deletions(-) diff --git a/litellm/proxy/management_endpoints/usage_endpoints/ai_usage_chat.py b/litellm/proxy/management_endpoints/usage_endpoints/ai_usage_chat.py index 7a45d0438de..85bb147e223 100644 --- a/litellm/proxy/management_endpoints/usage_endpoints/ai_usage_chat.py +++ b/litellm/proxy/management_endpoints/usage_endpoints/ai_usage_chat.py @@ -92,6 +92,10 @@ class _EntityTotal(NamedTuple): spend: float requests: float tokens: float + prompt_tokens: float = 0.0 + completion_tokens: float = 0.0 + cache_read_input_tokens: float = 0.0 + cache_creation_input_tokens: float = 0.0 class _UsageDump(Protocol): @@ -133,8 +137,10 @@ _TOOL_USAGE: Final[_ToolDef] = { "name": "get_usage_data", "description": ( "Fetch aggregated global usage/spend data. Returns daily spend, " - "token counts, request counts, and breakdowns by model, provider, " - "and API key. Use for overall spend, top models, top providers." + "token counts split into prompt, completion, cache read and cache " + "creation tokens, request counts, and breakdowns by model, provider, " + "and API key. Use for overall spend, top models, top providers, and " + "per-model token type breakdowns." ), "parameters": { "type": "object", @@ -359,6 +365,23 @@ def _accumulate_breakdown( return totals +TOKEN_TYPE_FIELDS: Final = ( + "prompt_tokens", + "completion_tokens", + "cache_read_input_tokens", + "cache_creation_input_tokens", +) + + +def _token_types(d: Mapping[str, float]) -> str: + """Render the four token-type counters the daily aggregates carry.""" + return ( + f"prompt {int(d.get('prompt_tokens', 0))}, completion {int(d.get('completion_tokens', 0))}, " + f"cache read {int(d.get('cache_read_input_tokens', 0))}, " + f"cache creation {int(d.get('cache_creation_input_tokens', 0))}" + ) + + def _ranked_lines( totals: dict[str, dict[str, float]], fmt: Callable[[str, dict[str, float]], str], @@ -377,15 +400,26 @@ def _summarise_usage_data(data: _UsageDump) -> str: f"Total Requests: {meta.get('total_api_requests', 0)}\n" f"Successful: {meta.get('total_successful_requests', 0)} | " f"Failed: {meta.get('total_failed_requests', 0)}\n" - f"Total Tokens: {meta.get('total_tokens', 0)}" + f"Total Tokens: {meta.get('total_tokens', 0)} " + f"(prompt {meta.get('total_prompt_tokens', 0)}, " + f"completion {meta.get('total_completion_tokens', 0)}, " + f"cache read {meta.get('total_cache_read_input_tokens', 0)}, " + f"cache creation {meta.get('total_cache_creation_input_tokens', 0)})" ) - models: Final = _accumulate_breakdown(results, "models", ["spend", "api_requests", "total_tokens"]) + models: Final = _accumulate_breakdown( + results, + "models", + ["spend", "api_requests", "total_tokens", *TOKEN_TYPE_FIELDS], + ) providers: Final = _accumulate_breakdown(results, "providers", ["spend", "api_requests"]) model_lines: Final = _ranked_lines( models, - lambda n, d: f" - {n}: ${d['spend']:.4f} ({int(d['api_requests'])} reqs, {int(d['total_tokens'])} tokens)", + lambda n, d: ( + f" - {n}: ${d['spend']:.4f} ({int(d['api_requests'])} reqs, {int(d['total_tokens'])} tokens: " + f"{_token_types(d)})" + ), TOP_N_MODELS, ) provider_lines: Final = _ranked_lines( @@ -416,12 +450,22 @@ def _summarise_entity_data(data: _UsageDump, entity_label: str) -> str: spend=(previous.spend if previous is not None else 0.0) + m.get("spend", 0), requests=(previous.requests if previous is not None else 0) + m.get("api_requests", 0), tokens=(previous.tokens if previous is not None else 0) + m.get("total_tokens", 0), + prompt_tokens=(previous.prompt_tokens if previous is not None else 0) + m.get("prompt_tokens", 0), + completion_tokens=(previous.completion_tokens if previous is not None else 0) + + m.get("completion_tokens", 0), + cache_read_input_tokens=(previous.cache_read_input_tokens if previous is not None else 0) + + m.get("cache_read_input_tokens", 0), + cache_creation_input_tokens=(previous.cache_creation_input_tokens if previous is not None else 0) + + m.get("cache_creation_input_tokens", 0), ) lines: Final = [f"{entity_label} Usage ({len(totals)} {entity_label.lower()}s):", ""] for eid, d in sorted(totals.items(), key=lambda x: -x[1].spend): label = d.alias if d.alias != eid else eid - lines.append(f"- {label} (ID: {eid}): ${d.spend:.4f} | {int(d.requests)} reqs | {int(d.tokens)} tokens") + lines.append( + f"- {label} (ID: {eid}): ${d.spend:.4f} | {int(d.requests)} reqs | {int(d.tokens)} tokens " + f"({_token_types(d._asdict())})" + ) return "\n".join(lines) diff --git a/tests/unit/proxy/management_endpoints/usage_endpoints/test_ai_usage_chat.py b/tests/unit/proxy/management_endpoints/usage_endpoints/test_ai_usage_chat.py index 4daddb89350..26457963fdb 100644 --- a/tests/unit/proxy/management_endpoints/usage_endpoints/test_ai_usage_chat.py +++ b/tests/unit/proxy/management_endpoints/usage_endpoints/test_ai_usage_chat.py @@ -43,6 +43,10 @@ SAMPLE_AGGREGATED_RESPONSE = { "spend": 40.0, "api_requests": 300, "total_tokens": 25000, + "prompt_tokens": 15000, + "completion_tokens": 8000, + "cache_read_input_tokens": 1500, + "cache_creation_input_tokens": 500, }, "metadata": {}, "api_key_breakdown": {}, @@ -155,6 +159,21 @@ class TestSummariseUsageData: summary = _summarise_usage_data(SAMPLE_AGGREGATED_RESPONSE) assert "gpt-4" in summary + def test_summarise_breaks_model_tokens_down_by_type(self): + summary = _summarise_usage_data(SAMPLE_AGGREGATED_RESPONSE) + assert "prompt 15000, completion 8000, cache read 1500, cache creation 500" in summary + + def test_summarise_breaks_total_tokens_down_by_type(self): + meta = { + **SAMPLE_AGGREGATED_RESPONSE["metadata"], + "total_prompt_tokens": 20000, + "total_completion_tokens": 10000, + "total_cache_read_input_tokens": 0, + "total_cache_creation_input_tokens": 0, + } + summary = _summarise_usage_data({**SAMPLE_AGGREGATED_RESPONSE, "metadata": meta}) + assert "prompt 20000, completion 10000, cache read 0, cache creation 0" in summary + def test_summarise_includes_providers(self): summary = _summarise_usage_data(SAMPLE_AGGREGATED_RESPONSE) assert "openai" in summary @@ -173,6 +192,10 @@ class TestSummariseEntityData: assert "$60.0" in summary assert "$40.0" in summary + def test_team_summary_breaks_tokens_down_by_type(self): + summary = _summarise_entity_data(SAMPLE_TEAM_RESPONSE, "Team") + assert summary.count("cache creation") == 2 + def test_team_summary_empty(self): empty = {"results": [], "metadata": {}} summary = _summarise_entity_data(empty, "Team")