diff --git a/litellm/integrations/focus/transformer.py b/litellm/integrations/focus/transformer.py index c22be2e66fc..d4a29d4ccdc 100644 --- a/litellm/integrations/focus/transformer.py +++ b/litellm/integrations/focus/transformer.py @@ -129,7 +129,12 @@ class FocusTransformer: & (pl.col("model_group").cast(pl.String) != "") ) .then(pl.col("model_group").cast(pl.String)) - .otherwise(pl.col("model").cast(pl.String)) + .when( + pl.col("model").cast(pl.String).is_not_null() + & (pl.col("model").cast(pl.String) != "") + ) + .then(pl.col("model").cast(pl.String)) + .otherwise(pl.lit("unknown")) .alias("ServiceName"), pl.col("team_id").cast(pl.String).alias("SubAccountId"), pl.col("team_alias").cast(pl.String).alias("SubAccountName"), 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..53e54e83799 --- /dev/null +++ b/tests/test_litellm/integrations/focus/test_transformer.py @@ -0,0 +1,78 @@ +"""Tests for FocusTransformer.""" + +from __future__ import annotations + +import json + +import polars as pl + +from litellm.integrations.focus.transformer import FocusTransformer + + +def _make_row(**overrides) -> dict: + """Return a minimal valid spend-log row, with optional overrides.""" + base = { + "date": "2026-03-17", + "spend": 0.05, + "api_key": "sk-test", + "api_key_alias": "my-key", + "model": "claude-sonnet", + "model_group": "sonnet-group", + "custom_llm_provider": "anthropic", + "team_id": "team-1", + "team_alias": "Engineering", + "user_id": "user-1", + "user_email": "test@example.com", + "prompt_tokens": 100, + "completion_tokens": 50, + } + base.update(overrides) + return base + + +def test_service_name_uses_model_group_when_present(): + row = _make_row(model_group="sonnet-group", model="claude-sonnet") + frame = pl.DataFrame([row]) + result = FocusTransformer().transform(frame) + assert result["ServiceName"][0] == "sonnet-group" + + +def test_service_name_falls_back_to_model_when_model_group_blank(): + row = _make_row(model_group="", model="claude-sonnet") + frame = pl.DataFrame([row]) + result = FocusTransformer().transform(frame) + assert result["ServiceName"][0] == "claude-sonnet" + + +def test_service_name_falls_back_to_model_when_model_group_null(): + row = _make_row(model_group=None, model="claude-sonnet") + frame = pl.DataFrame([row]) + result = FocusTransformer().transform(frame) + assert result["ServiceName"][0] == "claude-sonnet" + + +def test_service_name_defaults_to_unknown_when_both_blank(): + row = _make_row(model_group="", model="") + frame = pl.DataFrame([row]) + result = FocusTransformer().transform(frame) + assert result["ServiceName"][0] == "unknown" + + +def test_tags_include_token_counts(): + row = _make_row(prompt_tokens=100, completion_tokens=50) + frame = pl.DataFrame([row]) + result = FocusTransformer().transform(frame) + tags = json.loads(result["Tags"][0]) + assert tags["prompt_tokens"] == "100" + assert tags["completion_tokens"] == "50" + + +def test_tags_include_all_metadata_keys(): + row = _make_row() + frame = pl.DataFrame([row]) + result = FocusTransformer().transform(frame) + tags = json.loads(result["Tags"][0]) + assert "team_id" in tags + assert "model" in tags + assert "prompt_tokens" in tags + assert "completion_tokens" in tags