diff --git a/litellm/integrations/mavvrik_focus/mavvrik_focus_logger.py b/litellm/integrations/mavvrik_focus/mavvrik_focus_logger.py index 26b2f32f32f0..b9e00e95bcc7 100644 --- a/litellm/integrations/mavvrik_focus/mavvrik_focus_logger.py +++ b/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"", 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 cd21807e8874..377fab8c32ca 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,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