From 03bc06079e906f942f73b19eab7f0005bc817b96 Mon Sep 17 00:00:00 2001 From: Praveen Ghuge Date: Sat, 18 Jul 2026 11:11:55 +0530 Subject: [PATCH] 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"})]})