fix(mavvrik_focus): carry token counts (prompt/completion/total) in FOCUS Tags

litellm/integrations/mavvrik_focus/mavvrik_focus_logger.py

@@ -19,10 +19,13 @@
 
 from __future__ import annotations
 
+import json
 import os
 from datetime import datetime, timedelta, timezone
 from typing import TYPE_CHECKING, Any, List, Optional
 
+import polars as pl
+
 import litellm
 from litellm._logging import verbose_proxy_logger
 from litellm.constants import MAVVRIK_FOCUS_EXPORT_JOB_NAME
@@ -34,6 +37,63 @@
 else:
     AsyncIOScheduler = Any
 
+# FOCUS v1.2 has no standard column for token counts; core's transformer
+# 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 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:
+        return normalized
+
+    token_rows = data.select(available).to_dicts()
+    has_both = "prompt_tokens" in available and "completion_tokens" in available
+
+    def _merge(tags_json: str, row: dict) -> str:
+        try:
+            tags = json.loads(tags_json) if tags_json else {}
+        except (TypeError, ValueError):
+            tags = {}
+        if not isinstance(tags, dict):
+            tags = {}
+        for key in available:
+            value = row.get(key)
+            if value is not None:
+                tags[key] = str(value)
+        if has_both:
+            prompt = row.get("prompt_tokens")
+            completion = row.get("completion_tokens")
+            if prompt is not None and completion is not None:
+                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)]
+    )
+    return normalized.with_columns(merged_tags.alias("Tags"))
+
 
 def _parse_metrics_marker(
     marker: Optional[object],
@@ -136,6 +196,7 @@ async def _export_window(
         else:
             normalized = engine._transformer.transform(data)
             if not normalized.is_empty():
+                normalized = _with_token_tags(data, normalized)
                 payload = engine._serializer.serialize(normalized)
         await engine._destination.deliver(
             content=payload or b"",

tests/test_litellm/integrations/mavvrik_focus/test_mavvrik_focus_logger.py

@@ -1,15 +1,21 @@
+import json
 from datetime import datetime, timezone
 from unittest.mock import AsyncMock, MagicMock
 
+import polars as pl
 import pytest
 
 from litellm.integrations.focus.destinations.base import FocusTimeWindow
-from litellm.integrations.mavvrik_focus.mavvrik_focus_logger import MavvrikFocusLogger
+from litellm.integrations.mavvrik_focus.mavvrik_focus_logger import (
+    MavvrikFocusLogger,
+    _with_token_tags,
+)
 
 
 class _Frame:
     def __init__(self, *, empty: bool) -> None:
         self._empty = empty
+        self.columns: list = []
 
     def __len__(self) -> int:
         return 0 if self._empty else 1
@@ -65,3 +71,108 @@ async def test_export_window_delivers_empty_payload_for_empty_export(
         time_window=window,
         filename="metrics.csv",
     )
+
+
+def test_with_token_tags_merges_prompt_and_completion_tokens() -> None:
+    data = pl.DataFrame({"prompt_tokens": [57], "completion_tokens": [753]})
+    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",
+        "total_tokens": "810",
+    }
+
+
+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"]})
+
+    result = _with_token_tags(data, normalized)
+
+    tags = json.loads(result["Tags"][0])
+    assert tags == {
+        "prompt_tokens": "57",
+        "completion_tokens": "753",
+        "total_tokens": "810",
+    }
+
+
+def test_with_token_tags_recovers_from_non_dict_tags_json() -> None:
+    data = pl.DataFrame({"prompt_tokens": [57], "completion_tokens": [753]})
+    normalized = pl.DataFrame({"Tags": ["null"]})
+
+    result = _with_token_tags(data, normalized)
+
+    tags = json.loads(result["Tags"][0])
+    assert tags == {
+        "prompt_tokens": "57",
+        "completion_tokens": "753",
+        "total_tokens": "810",
+    }
+
+
+def test_with_token_tags_noop_when_tags_column_absent() -> None:
+    data = pl.DataFrame({"prompt_tokens": [57], "completion_tokens": [753]})
+    normalized = pl.DataFrame({"OtherColumn": ["x"]})
+
+    result = _with_token_tags(data, normalized)
+
+    assert result is normalized
+
+
+def test_with_token_tags_omits_total_when_only_one_token_column_present() -> None:
+    data = pl.DataFrame({"prompt_tokens": [57]})
+    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"}
+    assert "total_tokens" not in tags
+
+
+def test_with_token_tags_noop_when_token_columns_absent() -> None:
+    data = pl.DataFrame({"model": ["azure/gpt-4o-mini"]})
+    normalized = pl.DataFrame({"Tags": [json.dumps({"model": "azure/gpt-4o-mini"})]})
+
+    result = _with_token_tags(data, normalized)
+
+    assert result is normalized
+
+
+def test_with_token_tags_noop_on_row_count_mismatch() -> None:
+    data = pl.DataFrame({"prompt_tokens": [57, 12], "completion_tokens": [753, 40]})
+    normalized = pl.DataFrame({"Tags": [json.dumps({"model": "azure/gpt-4o-mini"})]})
+
+    result = _with_token_tags(data, normalized)
+
+    assert result is normalized