Merge 7fd7d43973 into 6762d65eb8

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)
 
 

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")