diff --git a/litellm/integrations/focus/transformer.py b/litellm/integrations/focus/transformer.py index b7d28e3dbb9..b3f321f9773 100644 --- a/litellm/integrations/focus/transformer.py +++ b/litellm/integrations/focus/transformer.py @@ -42,6 +42,12 @@ def _build_tags_expr(available_keys: list[str]) -> pl.Expr: ) +def _non_blank_string_expr(column_name: str) -> pl.Expr: + """Return a stripped string column, treating null/blank values as null.""" + value = pl.col(column_name).cast(pl.String).str.strip_chars() + return pl.when(value.is_not_null() & (value != "")).then(value).otherwise(None) + + class FocusTransformer: """Transforms LiteLLM DB rows into Focus-compatible schema.""" @@ -81,6 +87,12 @@ class FocusTransformer: none_str = pl.lit(None, dtype=pl.Utf8) none_dec = pl.lit(None, dtype=pl.Decimal(18, 6)) + service_name = pl.coalesce( + _non_blank_string_expr("model_group"), + _non_blank_string_expr("model"), + _non_blank_string_expr("custom_llm_provider"), + pl.lit("litellm-proxy"), + ) return frame.select( dec(pl.col("spend").fill_null(0.0)).alias("BilledCost"), @@ -122,7 +134,7 @@ class FocusTransformer: pl.col("model").cast(pl.String).alias("ResourceType"), pl.lit("AI and Machine Learning").alias("ServiceCategory"), pl.lit("Generative AI").alias("ServiceSubcategory"), - pl.col("model_group").cast(pl.String).alias("ServiceName"), + service_name.alias("ServiceName"), pl.col("team_id").cast(pl.String).alias("SubAccountId"), pl.col("team_alias").cast(pl.String).alias("SubAccountName"), none_str.alias("SubAccountType"), diff --git a/tests/test_litellm/integrations/focus/test_transformer.py b/tests/test_litellm/integrations/focus/test_transformer.py new file mode 100644 index 00000000000..c7b14b39ede --- /dev/null +++ b/tests/test_litellm/integrations/focus/test_transformer.py @@ -0,0 +1,51 @@ +"""Tests for FocusTransformer normalization.""" + +from __future__ import annotations + +import polars as pl + +from litellm.integrations.focus.transformer import FocusTransformer + + +def _usage_frame(rows: list[dict]) -> pl.DataFrame: + defaults = { + "date": "2026-05-17", + "user_id": "user-1", + "api_key": "virtual-key-id", + "api_key_alias": "test-key", + "model": "gpt-4.1", + "model_group": "gpt-4.1", + "custom_llm_provider": "openai", + "spend": 0.1, + "team_id": "team-1", + "team_alias": "Test Team", + "user_email": "user@example.com", + } + return pl.DataFrame([{**defaults, **row} for row in rows]) + + +def test_should_use_model_group_as_service_name_when_present(): + result = FocusTransformer().transform( + _usage_frame([{"model": "gpt-4.1", "model_group": "production-gpt"}]) + ) + + assert result["ServiceName"].to_list() == ["production-gpt"] + + +def test_should_fallback_service_name_when_model_group_is_blank(): + result = FocusTransformer().transform( + _usage_frame( + [ + {"model": "mcp-vector-store", "model_group": ""}, + {"model": "", "model_group": " ", "custom_llm_provider": "proxy"}, + {"model": None, "model_group": None, "custom_llm_provider": None}, + ] + ) + ) + + assert result["ServiceName"].to_list() == [ + "mcp-vector-store", + "proxy", + "litellm-proxy", + ] + assert all(service_name.strip() for service_name in result["ServiceName"])