From 24dc17f1fa08bed1a2018357e9945bd722ef5b7f Mon Sep 17 00:00:00 2001 From: Praveen Ghuge Date: Fri, 17 Jul 2026 09:47:46 +0530 Subject: [PATCH] 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"})]})