From 4b3905cb8b6763fc66ef63c6f240a8e08ca08259 Mon Sep 17 00:00:00 2001
From: Praveen Ghuge <pghuge@digitalex.io>
Date: Sat, 18 Jul 2026 11:32:11 +0530
Subject: [PATCH] fix(mavvrik_focus): also carry cache token counts in FOCUS
 Tags

cache_creation_input_tokens and cache_read_input_tokens are selected
by the same database.py query as prompt_tokens/completion_tokens and
dropped by the same transformer. Add them to _TOKEN_TAG_KEYS.
---
 .../mavvrik_focus/mavvrik_focus_logger.py     | 32 +++++++++++++------
 .../test_mavvrik_focus_logger.py              | 24 ++++++++++++++
 2 files changed, 46 insertions(+), 10 deletions(-)

diff --git a/litellm/integrations/mavvrik_focus/mavvrik_focus_logger.py b/litellm/integrations/mavvrik_focus/mavvrik_focus_logger.py
index 33078329e54..b9e00e95bcc 100644
--- a/litellm/integrations/mavvrik_focus/mavvrik_focus_logger.py
+++ b/litellm/integrations/mavvrik_focus/mavvrik_focus_logger.py
@@ -38,19 +38,26 @@ else:
     AsyncIOScheduler = Any
 
 # FOCUS v1.2 has no standard column for token counts; core's transformer
-# drops prompt_tokens/completion_tokens even though the source query selects
-# them. Mavvrik carries them through as extra keys in the existing Tags JSON
-# column (the spec's own escape hatch for non-standard fields), rather than
-# changing the shared transformer used by every FOCUS destination. total_tokens
-# isn't a stored column at all -- it's derived here as their sum.
-_TOKEN_TAG_KEYS = ("prompt_tokens", "completion_tokens")
+# drops prompt_tokens/completion_tokens/cache_creation_input_tokens/
+# cache_read_input_tokens even though the source query selects them. Mavvrik
+# carries them through as extra keys in the existing Tags JSON column (the
+# spec's own escape hatch for non-standard fields), rather than changing the
+# shared transformer used by every FOCUS destination. total_tokens isn't a
+# stored column at all -- it's derived here as the sum of prompt and
+# completion tokens.
+_TOKEN_TAG_KEYS = (
+    "prompt_tokens",
+    "completion_tokens",
+    "cache_creation_input_tokens",
+    "cache_read_input_tokens",
+)
 
 
 def _with_token_tags(data: pl.DataFrame, normalized: pl.DataFrame) -> pl.DataFrame:
-    """Merge prompt/completion token counts (and their sum, total_tokens) from
-    the pre-transform frame into ``normalized``'s Tags column. Rows correspond
-    1:1 and in the same order across both frames -- transform() only
-    adds/renames columns, it never filters or reorders rows.
+    """Merge token counts (and their sum, total_tokens) from the pre-transform
+    frame into ``normalized``'s Tags column. Rows correspond 1:1 and in the
+    same order across both frames -- transform() only adds/renames columns,
+    it never filters or reorders rows.
     """
     available = [k for k in _TOKEN_TAG_KEYS if k in data.columns]
     if not available or len(data) != len(normalized) or "Tags" not in normalized.columns:
@@ -77,6 +84,11 @@ def _with_token_tags(data: pl.DataFrame, normalized: pl.DataFrame) -> pl.DataFra
                 tags["total_tokens"] = str(prompt + completion)
         return json.dumps(tags)
 
+    verbose_proxy_logger.debug(
+        "Mavvrik FOCUS export: merging token tags for %d row(s) (keys=%s)",
+        len(token_rows),
+        available,
+    )
     merged_tags = pl.Series(
         [_merge(tags_json, row) for tags_json, row in zip(normalized["Tags"].to_list(), token_rows)]
     )
diff --git a/tests/test_litellm/integrations/mavvrik_focus/test_mavvrik_focus_logger.py b/tests/test_litellm/integrations/mavvrik_focus/test_mavvrik_focus_logger.py
index 62c8391c7ff..377fab8c32c 100644
--- a/tests/test_litellm/integrations/mavvrik_focus/test_mavvrik_focus_logger.py
+++ b/tests/test_litellm/integrations/mavvrik_focus/test_mavvrik_focus_logger.py
@@ -88,6 +88,30 @@ def test_with_token_tags_merges_prompt_and_completion_tokens() -> None:
     }
 
 
+def test_with_token_tags_merges_cache_token_columns() -> None:
+    data = pl.DataFrame(
+        {
+            "prompt_tokens": [57],
+            "completion_tokens": [753],
+            "cache_creation_input_tokens": [10],
+            "cache_read_input_tokens": [5],
+        }
+    )
+    normalized = pl.DataFrame({"Tags": [json.dumps({"model": "azure/gpt-4o-mini"})]})
+
+    result = _with_token_tags(data, normalized)
+
+    tags = json.loads(result["Tags"][0])
+    assert tags == {
+        "model": "azure/gpt-4o-mini",
+        "prompt_tokens": "57",
+        "completion_tokens": "753",
+        "cache_creation_input_tokens": "10",
+        "cache_read_input_tokens": "5",
+        "total_tokens": "810",
+    }
+
+
 def test_with_token_tags_recovers_from_malformed_tags_json() -> None:
     data = pl.DataFrame({"prompt_tokens": [57], "completion_tokens": [753]})
     normalized = pl.DataFrame({"Tags": ["not-valid-json"]})