From a58ab8b19a178999564c70ff8ae34f48b4e391de Mon Sep 17 00:00:00 2001 From: Praveen Ghuge Date: Thu, 16 Jul 2026 21:46:01 +0530 Subject: [PATCH 1/6] fix(mavvrik_focus): carry prompt/completion token counts in FOCUS Tags FOCUS v1.2 has no standard column for LLM token counts, and the shared FocusTransformer used by every destination (Mavvrik, Vantage, CloudZero) drops prompt_tokens/completion_tokens even though the source query already selects them. Merge the two counts into the existing Tags JSON column, which is the spec's own escape hatch for non-standard fields, inside the Mavvrik-only export path so no shared transformer changes. --- .../mavvrik_focus/mavvrik_focus_logger.py | 37 +++++++++++++++++++ 1 file changed, 37 insertions(+) diff --git a/litellm/integrations/mavvrik_focus/mavvrik_focus_logger.py b/litellm/integrations/mavvrik_focus/mavvrik_focus_logger.py index 26b2f32f32f..544960aca54 100644 --- a/litellm/integrations/mavvrik_focus/mavvrik_focus_logger.py +++ b/litellm/integrations/mavvrik_focus/mavvrik_focus_logger.py @@ -19,10 +19,13 @@ overwrite each other within the same day, producing incomplete data. 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,39 @@ if TYPE_CHECKING: 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. +_TOKEN_TAG_KEYS = ("prompt_tokens", "completion_tokens") + + +def _with_token_tags(data: pl.DataFrame, normalized: pl.DataFrame) -> pl.DataFrame: + """Merge prompt/completion token counts 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): + return normalized + + token_rows = data.select(available).to_dicts() + + def _merge(tags_json: str, row: dict) -> str: + tags = json.loads(tags_json) if tags_json else {} + for key in available: + value = row.get(key) + if value is not None: + tags[key] = str(value) + return json.dumps(tags) + + 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 +172,7 @@ class MavvrikFocusLogger(FocusLogger): 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"", From d4f059c9b5d36ad831c824dbfc5d4e9bcafd0a35 Mon Sep 17 00:00:00 2001 From: Praveen Ghuge Date: Thu, 16 Jul 2026 22:04:50 +0530 Subject: [PATCH 2/6] test(mavvrik_focus): cover _with_token_tags and fix mock Frame Add columns attribute to the _Frame test double so _with_token_tags does not raise AttributeError on the existing empty-export parametrized case, and add dedicated unit tests for _with_token_tags covering the merge, no-token-columns, and row-count-mismatch paths. --- .../test_mavvrik_focus_logger.py | 40 ++++++++++++++++++- 1 file changed, 39 insertions(+), 1 deletion(-) 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 cd21807e887..df2d76ecbe2 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 @@ -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,35 @@ 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", + } + + +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 From 24dc17f1fa08bed1a2018357e9945bd722ef5b7f Mon Sep 17 00:00:00 2001 From: Praveen Ghuge Date: Fri, 17 Jul 2026 09:47:46 +0530 Subject: [PATCH 3/6] fix(mavvrik_focus): also derive total_tokens in FOCUS Tags total_tokens has no stored column in LiteLLM_DailyUserSpend at all, so it can't be selected like prompt_tokens/completion_tokens. Derive it as their sum in _with_token_tags, only when both source counts are present for a row, and add tests covering the sum and the partial-data case. --- .../mavvrik_focus/mavvrik_focus_logger.py | 17 ++++++++++++----- .../mavvrik_focus/test_mavvrik_focus_logger.py | 12 ++++++++++++ 2 files changed, 24 insertions(+), 5 deletions(-) diff --git a/litellm/integrations/mavvrik_focus/mavvrik_focus_logger.py b/litellm/integrations/mavvrik_focus/mavvrik_focus_logger.py index 544960aca54..d61653c9559 100644 --- a/litellm/integrations/mavvrik_focus/mavvrik_focus_logger.py +++ b/litellm/integrations/mavvrik_focus/mavvrik_focus_logger.py @@ -41,21 +41,23 @@ else: # 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. +# 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") def _with_token_tags(data: pl.DataFrame, normalized: pl.DataFrame) -> pl.DataFrame: - """Merge prompt/completion token counts 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 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. """ available = [k for k in _TOKEN_TAG_KEYS if k in data.columns] if not available or len(data) != len(normalized): 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: tags = json.loads(tags_json) if tags_json else {} @@ -63,6 +65,11 @@ def _with_token_tags(data: pl.DataFrame, normalized: pl.DataFrame) -> pl.DataFra 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) merged_tags = pl.Series( 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 df2d76ecbe2..64890892915 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 @@ -84,9 +84,21 @@ def test_with_token_tags_merges_prompt_and_completion_tokens() -> None: "model": "azure/gpt-4o-mini", "prompt_tokens": "57", "completion_tokens": "753", + "total_tokens": "810", } +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"})]}) From 19de9894c288a1a4c8df6e307980852fc0de6e22 Mon Sep 17 00:00:00 2001 From: Praveen Ghuge Date: Fri, 17 Jul 2026 19:47:50 +0530 Subject: [PATCH 4/6] fix(mavvrik_focus): guard against malformed Tags JSON in token merge json.loads on the existing Tags value had no error handling; a malformed value would raise JSONDecodeError and abort the entire export window instead of just skipping that row's token merge. --- .../mavvrik_focus/mavvrik_focus_logger.py | 5 ++++- .../mavvrik_focus/test_mavvrik_focus_logger.py | 14 ++++++++++++++ 2 files changed, 18 insertions(+), 1 deletion(-) diff --git a/litellm/integrations/mavvrik_focus/mavvrik_focus_logger.py b/litellm/integrations/mavvrik_focus/mavvrik_focus_logger.py index d61653c9559..83532f215e4 100644 --- a/litellm/integrations/mavvrik_focus/mavvrik_focus_logger.py +++ b/litellm/integrations/mavvrik_focus/mavvrik_focus_logger.py @@ -60,7 +60,10 @@ def _with_token_tags(data: pl.DataFrame, normalized: pl.DataFrame) -> pl.DataFra has_both = "prompt_tokens" in available and "completion_tokens" in available def _merge(tags_json: str, row: dict) -> str: - tags = json.loads(tags_json) if tags_json else {} + try: + tags = json.loads(tags_json) if tags_json else {} + except (TypeError, ValueError): + tags = {} for key in available: value = row.get(key) if value is not None: 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 64890892915..c1b3e1fcedd 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,20 @@ def test_with_token_tags_merges_prompt_and_completion_tokens() -> None: } +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_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"})]}) From 03bc06079e906f942f73b19eab7f0005bc817b96 Mon Sep 17 00:00:00 2001 From: Praveen Ghuge Date: Sat, 18 Jul 2026 11:11:55 +0530 Subject: [PATCH 5/6] fix(mavvrik_focus): guard Tags column and non-dict parsed JSON --- .../mavvrik_focus/mavvrik_focus_logger.py | 4 +++- .../test_mavvrik_focus_logger.py | 23 +++++++++++++++++++ 2 files changed, 26 insertions(+), 1 deletion(-) diff --git a/litellm/integrations/mavvrik_focus/mavvrik_focus_logger.py b/litellm/integrations/mavvrik_focus/mavvrik_focus_logger.py index 83532f215e4..33078329e54 100644 --- a/litellm/integrations/mavvrik_focus/mavvrik_focus_logger.py +++ b/litellm/integrations/mavvrik_focus/mavvrik_focus_logger.py @@ -53,7 +53,7 @@ def _with_token_tags(data: pl.DataFrame, normalized: pl.DataFrame) -> pl.DataFra 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): + if not available or len(data) != len(normalized) or "Tags" not in normalized.columns: return normalized token_rows = data.select(available).to_dicts() @@ -64,6 +64,8 @@ def _with_token_tags(data: pl.DataFrame, normalized: pl.DataFrame) -> pl.DataFra 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: 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 c1b3e1fcedd..62c8391c7ff 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 @@ -102,6 +102,29 @@ def test_with_token_tags_recovers_from_malformed_tags_json() -> None: } +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"})]}) From 4b3905cb8b6763fc66ef63c6f240a8e08ca08259 Mon Sep 17 00:00:00 2001 From: Praveen Ghuge Date: Sat, 18 Jul 2026 11:32:11 +0530 Subject: [PATCH 6/6] 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"]})