From fae6ea618d16a96feffa97bbd4a0fad84baa4953 Mon Sep 17 00:00:00 2001 From: Praveen Ghuge <95286176+pghuge-cloudwiz@users.noreply.github.com> Date: Mon, 1 Jun 2026 15:28:30 +0530 Subject: [PATCH] fix(focus): fix ConsumedQuantity/PricingQuantity BigInt cast and null fallback (#29415) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * fix(focus): fix ConsumedQuantity/PricingQuantity BigInt cast and fill_null Two bugs in FocusTransformer.transform(): 1. fill_null(1.0) → fill_null(0.0): NULL api_requests rows (old schema rows before the column existed) were reported as 1 request instead of 0, inflating consumed/pricing quantity counts. 2. .cast(pl.Float64) on a Postgres BigInt column: query_raw returns BigInt as Python int, which Polars may fail to cast directly to Float64 on some driver versions, producing nulls that then hit the fill_null. Adding an intermediate .cast(pl.Int64) step mirrors the actual column type and ensures the value is preserved before widening to Float64. Co-Authored-By: Claude Sonnet 4.6 (1M context) * chore: remove unused pytest import from test_focus_transformer * style(focus): apply Black py312 formatting to transformer.py * style(focus): reformat transformer.py with Black 26.3.1 to pass CI lint --------- Co-authored-by: Claude Sonnet 4.6 (1M context) --- litellm/integrations/focus/transformer.py | 8 ++- .../focus/test_focus_transformer.py | 69 +++++++++++++++++++ 2 files changed, 75 insertions(+), 2 deletions(-) create mode 100644 tests/test_litellm/integrations/focus/test_focus_transformer.py diff --git a/litellm/integrations/focus/transformer.py b/litellm/integrations/focus/transformer.py index 6f4433b4a05..8496b7ec159 100644 --- a/litellm/integrations/focus/transformer.py +++ b/litellm/integrations/focus/transformer.py @@ -95,7 +95,9 @@ class FocusTransformer: pl.lit("Usage-Based").alias("ChargeFrequency"), fmt(pl.col("ChargePeriodEnd")).alias("ChargePeriodEnd"), fmt(pl.col("ChargePeriodStart")).alias("ChargePeriodStart"), - dec(pl.lit(1.0)).alias("ConsumedQuantity"), + dec( + pl.col("api_requests").cast(pl.Int64).cast(pl.Float64).fill_null(0.0) + ).alias("ConsumedQuantity"), pl.lit("Requests").alias("ConsumedUnit"), dec(pl.col("spend").fill_null(0.0)).alias("ContractedCost"), none_str.alias("ContractedUnitPrice"), @@ -107,7 +109,9 @@ class FocusTransformer: none_str.alias("AvailabilityZone"), pl.lit("USD").alias("PricingCurrency"), none_str.alias("PricingCategory"), - dec(pl.lit(1.0)).alias("PricingQuantity"), + dec( + pl.col("api_requests").cast(pl.Int64).cast(pl.Float64).fill_null(0.0) + ).alias("PricingQuantity"), none_dec.alias("PricingCurrencyContractedUnitPrice"), dec(pl.col("spend").fill_null(0.0)).alias("PricingCurrencyEffectiveCost"), none_dec.alias("PricingCurrencyListUnitPrice"), diff --git a/tests/test_litellm/integrations/focus/test_focus_transformer.py b/tests/test_litellm/integrations/focus/test_focus_transformer.py new file mode 100644 index 00000000000..7e90f7d0a2b --- /dev/null +++ b/tests/test_litellm/integrations/focus/test_focus_transformer.py @@ -0,0 +1,69 @@ +"""Tests for FocusTransformer — ConsumedQuantity / PricingQuantity correctness.""" + +from __future__ import annotations + +from decimal import Decimal + +import polars as pl + +from litellm.integrations.focus.transformer import FocusTransformer + + +def _base_row(**overrides) -> dict: + row = { + "date": "2026-05-25", + "user_id": "u1", + "api_key": "sk-test", + "api_key_alias": "my-key", + "model": "gpt-4o", + "model_group": "openai", + "custom_llm_provider": "openai", + "spend": 0.05, + "api_requests": 3, + "team_id": "team1", + "team_alias": "Engineering", + "user_email": "user@example.com", + } + row.update(overrides) + return row + + +def _transform(rows: list[dict]) -> pl.DataFrame: + frame = pl.DataFrame(rows, infer_schema_length=None) + return FocusTransformer().transform(frame) + + +def test_consumed_quantity_reflects_api_requests(): + result = _transform([_base_row(api_requests=7)]) + assert result["ConsumedQuantity"][0] == Decimal("7.000000") + + +def test_pricing_quantity_reflects_api_requests(): + result = _transform([_base_row(api_requests=7)]) + assert result["PricingQuantity"][0] == Decimal("7.000000") + + +def test_null_api_requests_falls_back_to_zero_not_one(): + """Rows with NULL api_requests (old schema rows) must produce 0, not 1.""" + result = _transform([_base_row(api_requests=None)]) + assert result["ConsumedQuantity"][0] == Decimal("0.000000") + assert result["PricingQuantity"][0] == Decimal("0.000000") + + +def test_zero_api_requests_stays_zero(): + result = _transform([_base_row(api_requests=0)]) + assert result["ConsumedQuantity"][0] == Decimal("0.000000") + assert result["PricingQuantity"][0] == Decimal("0.000000") + + +def test_bigint_api_requests_cast_correctly(): + """api_requests comes from Postgres as BigInt — large values must not overflow.""" + result = _transform([_base_row(api_requests=1_000_000)]) + assert result["ConsumedQuantity"][0] == Decimal("1000000.000000") + assert result["PricingQuantity"][0] == Decimal("1000000.000000") + + +def test_consumed_and_pricing_quantity_match(): + """ConsumedQuantity and PricingQuantity must always be equal.""" + result = _transform([_base_row(api_requests=42)]) + assert result["ConsumedQuantity"][0] == result["PricingQuantity"][0]