fix(mavvrik_focus): carry token counts (prompt/completion/total) in FOCUS Tags

[pghuge-cloudwiz]

Relevant issues

Replaces #33551, retargeted to litellm_internal_staging (same 3 commits, already reviewed by Greptile at 5/5 on litellm_oss_staging).

Mavvrik FOCUS exports do not carry LLM token usage (prompt_tokens, completion_tokens, total_tokens), even though LiteLLM's own DB and UI have this data. FOCUS v1.2 has no standard column for token counts, so the shared FocusTransformer (used by every FOCUS destination: Mavvrik, Vantage, CloudZero) drops prompt_tokens/completion_tokens during transform(), even though database.py's query already selects them from LiteLLM_DailyUserSpend. total_tokens has no stored column at all; it is derived here as the sum of the other two.

Pre-Submission checklist

• I have added meaningful tests
• My PR passes all unit tests on make test-unit
• My PR's scope is as isolated as possible; it only solves 1 specific problem
• I have requested a Greptile review by commenting @greptileai and received a Confidence Score of at least 4/5 (5/5 on the identical commits against fix(mavvrik_focus): carry token counts (prompt/completion/total) in FOCUS Tags #33551)

Type

Bug Fix

Changes

FocusTransformer.transform() (shared core code, not touched by this PR) builds the final FOCUS-shaped frame via an explicit column select(...) that enumerates ~35 output columns. prompt_tokens/completion_tokens are not in that list, so they are silently dropped between the raw query result and the exported CSV. This is true for every destination built on this transformer, not just Mavvrik. total_tokens isn't stored in LiteLLM_DailyUserSpend at all, so there is nothing to drop or select for it.

Rather than changing the shared transformer (which would affect Vantage and CloudZero too), this PR merges the token counts into the existing Tags JSON column, entirely inside MavvrikFocusLogger._export_window(), a Mavvrik-only file. Tags is FOCUS v1.2's own escape hatch for non-standard fields, and core already uses it to carry team_id, model, custom_llm_provider, etc.

_export_window() holds both the pre-transform frame (data, still has the token columns) and the post-transform frame (normalized, tokens already dropped) at the same point, right before serialization. The _with_token_tags() helper zips the two frames by row position (transform() only adds/renames columns and never filters or reorders rows, so a 1:1 row correspondence is guaranteed), adds prompt_tokens/completion_tokens as extra string keys into each row's existing Tags JSON, and additionally adds total_tokens as their sum when both source counts are present for that row (omitted otherwise, to avoid emitting a partial/wrong total).

No changes to litellm/integrations/focus/database.py or litellm/integrations/focus/transformer.py.

Also fixes the _Frame test double in test_mavvrik_focus_logger.py (was missing a columns attribute, which _with_token_tags reads) and adds unit tests for _with_token_tags: merge with total_tokens, partial-column no-total, no-token-columns no-op, and row-count-mismatch no-op.

Note: litellm_internal_staging is currently missing the earlier fix(mavvrik): advance metricsMarker after upload; fix scheduler startup fix (#31068), which only merged into litellm_oss_staging and was never synced forward. This PR's diff applies cleanly regardless, since mavvrik_focus_logger.py was otherwise identical between the two branches at the time of this PR, but the metricsMarker gap is a separate pre-existing issue worth tracking independently.

Screenshots / Proof of Fix

E2E verified on the QA VM using a native LiteLLM proxy (not Docker) against a dedicated Postgres database (litellm_tokentest), with real Azure gpt-4o-mini completions and a real Mavvrik sandbox connection. FOCUS_CRON_OFFSET was used to fire the daily export job a few minutes after startup; test rows were backdated one day in Postgres so the daily window (which only exports strictly-past dates) would pick them up.

Before fix, commit 9076c3334760d4c4d6be4b2555c874e9d49c2733 (unpatched mavvrik_focus_logger.py)

3 real completions produced this Postgres row:

date       | model             | api_requests | prompt_tokens | completion_tokens | spend
2026-07-15 | azure/gpt-4o-mini | 3            | 54             | 461                | 0.00031317

Resulting Tags value in the exported CSV, no token counts despite Postgres having them:

{"user_id": "default_user_id", "model": "azure/gpt-4o-mini", "model_group": "gpt-4o-mini", "custom_llm_provider": "azure"}

After fix (prompt/completion tokens), commit b33978e5cf...

3 new real completions produced this Postgres row:

date       | model             | api_requests | prompt_tokens | completion_tokens | spend
2026-07-15 | azure/gpt-4o-mini | 3            | 57             | 753                | 0.000506385

Resulting Tags value:

{"user_id": "default_user_id", "model": "azure/gpt-4o-mini", "model_group": "gpt-4o-mini", "custom_llm_provider": "azure", "prompt_tokens": "57", "completion_tokens": "753"}

After adding total_tokens, commit f7df2ec14c...

3 new real completions produced this Postgres row:

date       | model             | api_requests | prompt_tokens | completion_tokens | spend
2026-07-16 | azure/gpt-4o-mini | 3            | 48             | 274                | 0.00018876

Export log:

04:16:00 - LiteLLM:DEBUG: mavvrik_destination.py:318 - Mavvrik FOCUS destination: uploading 1287 bytes for date=2026-07-16 (usage_20260716T000000Z_20260717T041600Z.csv)
04:16:01 - LiteLLM:DEBUG: mavvrik_destination.py:198 - Mavvrik FOCUS destination: GCS session started, uploading 533 gzip bytes in 1 chunk(s)
04:16:01 - LiteLLM:DEBUG: mavvrik_destination.py:329 - Mavvrik FOCUS destination: upload complete for date=2026-07-16

Resulting Tags value in the exported CSV, downloaded from GCS and decompressed. total_tokens now present and equal to prompt_tokens + completion_tokens (48 + 274 = 322), matching Postgres exactly:

{"user_id": "default_user_id", "model": "azure/gpt-4o-mini", "model_group": "gpt-4o-mini", "custom_llm_provider": "azure", "prompt_tokens": "48", "completion_tokens": "274", "total_tokens": "322"}

All other FOCUS columns (BilledCost, ConsumedQuantity, ChargePeriodStart/End, etc.) are unchanged across all three runs, confirming the fix is additive to Tags only.

Unit tests (tests/test_litellm/integrations/mavvrik_focus/test_mavvrik_focus_logger.py), run against a real litellm install with polars/pytest-asyncio available:

test_export_window_delivers_empty_payload_for_empty_export[True-False-not-used] PASSED
test_export_window_delivers_empty_payload_for_empty_export[False-True-not-used] PASSED
test_export_window_delivers_empty_payload_for_empty_export[False-False-] PASSED
test_with_token_tags_merges_prompt_and_completion_tokens PASSED
test_with_token_tags_omits_total_when_only_one_token_column_present PASSED
test_with_token_tags_noop_when_token_columns_absent PASSED
test_with_token_tags_noop_on_row_count_mismatch PASSED
7 passed in 4.18s

[pghuge-cloudwiz]

@greptileai please review.

[greptile-apps]

Greptile Summary

This PR fixes missing LLM token counts (prompt_tokens, completion_tokens, total_tokens) in Mavvrik FOCUS exports by merging them into the existing FOCUS Tags JSON column — FOCUS v1.2's own extension point — entirely within the Mavvrik-specific _export_window() method, leaving the shared transformer used by Vantage and CloudZero untouched.

• Adds _with_token_tags(data, normalized) which zips the pre-transform frame (still has token columns) with the post-transform frame by row position, adds all four token tag keys as string values, and derives total_tokens as the sum of prompt and completion tokens only when both are present.
• Hardens all failure modes from the previous review: malformed Tags JSON falls back to {}, non-dict parsed values fall back to {}, a missing Tags column short-circuits cleanly, and row-count mismatches are silently skipped.
• Adds seven new unit tests covering the happy path, cache-token columns, malformed-JSON recovery, non-dict recovery, absent Tags column, partial token columns, and row-count mismatch; also repairs the _Frame test double by adding a columns attribute required by the new code.

Confidence Score: 5/5

Safe to merge — the change is strictly additive to Tags JSON on one Mavvrik-only code path and cannot affect Vantage or CloudZero exports.

All four defensive-coding gaps raised in the prior review have been addressed in this revision, and the new tests exercise every newly added branch. No existing assertions were weakened; the only test-file change is adding a missing columns attribute to a mock object that would otherwise raise AttributeError under the new call path.

No files require special attention.

Important Files Changed

Filename	Overview
litellm/integrations/mavvrik_focus/mavvrik_focus_logger.py	Adds _with_token_tags() helper to zip token counts from the pre-transform frame into the FOCUS Tags JSON column; all previously flagged issues (malformed JSON, non-dict Tags, missing Tags column guard) are addressed in this version
tests/test_litellm/integrations/mavvrik_focus/test_mavvrik_focus_logger.py	Adds columns: list = [] to the _Frame mock (necessary to avoid AttributeError when _with_token_tags reads data.columns) and adds seven new unit tests covering the major branches of _with_token_tags, including malformed-JSON recovery and non-dict Tags recovery

_{Reviews (8): Last reviewed commit: "fix(mavvrik_focus): also carry cache tok..." | Re-trigger Greptile}

[greptile-apps]

Greptile Summary

This PR adds _with_token_tags() to mavvrik_focus_logger.py, merging prompt_tokens, completion_tokens, and a derived total_tokens into each row's existing FOCUS Tags JSON column immediately before serialization. The approach is correctly scoped to the Mavvrik-only file and avoids touching the shared FocusTransformer used by other destinations.

• _with_token_tags() zips the pre-transform frame (which still has token columns) with the post-transform frame by row position, injecting tokens as string keys into the Tags JSON — a semantically correct use of FOCUS v1.2's extension mechanism.
• The _Frame test double gains a columns attribute to prevent AttributeError; four new unit tests cover the merge, partial-column, no-column, and row-count-mismatch branches.
• Three minor defensive gaps exist: missing Tags column existence check, unvalidated json.loads return type, and an implicit row-ordering invariant with no observability if broken.

Confidence Score: 4/5

The change is additive, confined to a single Mavvrik-specific file, and backed by E2E proof against a live Postgres database with real completions.

The core logic is correct and the approach is well-reasoned. The three defensive gaps — missing Tags column guard, unvalidated json.loads return type, and implicit row-ordering reliance — are all export-path-only and require unexpected inputs or future transformer changes to manifest.

litellm/integrations/mavvrik_focus/mavvrik_focus_logger.py — specifically the _with_token_tags helper.

Important Files Changed

Filename	Overview
litellm/integrations/mavvrik_focus/mavvrik_focus_logger.py	Adds _with_token_tags() to merge prompt/completion tokens into the FOCUS Tags column; logic is sound but has minor defensive gaps.
tests/test_litellm/integrations/mavvrik_focus/test_mavvrik_focus_logger.py	Adds four unit tests for _with_token_tags() branches and fixes _Frame mock; existing coverage is preserved.

_{Reviews (2): Last reviewed commit: "fix(mavvrik_focus): also derive total_to..." | Re-trigger Greptile}

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!

[codspeed-hq]

Merging this PR will not alter performance

✅ 31 untouched benchmarks

---

_{Comparing cloudwizio:fix/mavvrik-focus-token-counts-internal (4b3905c) with litellm_internal_staging (6375923)}

[yuneng-berri]

Please see if the Greptile comments are applicable.

I do not believe we have pl as a dependency, and adding one just for the integration will introduce bloat. Can you accomplish this without the pl dependency

[pghuge-cloudwiz]

Addressed the flagged concern: wrapped json.loads in _merge with a try/except so a malformed Tags value degrades to an empty dict for that row instead of aborting the whole export window with an unhandled JSONDecodeError. Added a dedicated test for this case; all 8 tests pass.

On the import polars as pl note: this file already transitively imports polars via FocusLogger/export_engine.py/transformer.py/etc., so the module-level import here doesn't introduce a new failure mode in practice.

@greptileai please re-review.

[pghuge-cloudwiz]

@greptileai please re-review.