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"",