diff --git a/litellm/integrations/focus/transformer.py b/litellm/integrations/focus/transformer.py index d4a29d4ccdc..b14b1e60ccd 100644 --- a/litellm/integrations/focus/transformer.py +++ b/litellm/integrations/focus/transformer.py @@ -122,13 +122,19 @@ class FocusTransformer: pl.col("model").cast(pl.String).alias("ResourceId"), pl.col("model").cast(pl.String).alias("ResourceName"), pl.col("model").cast(pl.String).alias("ResourceType"), - pl.lit("AI and Machine Learning").alias("ServiceCategory"), - pl.lit("Generative AI").alias("ServiceSubcategory"), pl.when( pl.col("model_group").cast(pl.String).is_not_null() & (pl.col("model_group").cast(pl.String) != "") ) .then(pl.col("model_group").cast(pl.String)) + .otherwise(none_str) + .alias("ServiceCategory"), + pl.lit("Generative AI").alias("ServiceSubcategory"), + pl.when( + pl.col("custom_llm_provider").cast(pl.String).is_not_null() + & (pl.col("custom_llm_provider").cast(pl.String) != "") + ) + .then(pl.col("custom_llm_provider").cast(pl.String)) .when( pl.col("model").cast(pl.String).is_not_null() & (pl.col("model").cast(pl.String) != "") diff --git a/tests/test_litellm/integrations/focus/test_transformer.py b/tests/test_litellm/integrations/focus/test_transformer.py index 53e54e83799..782f77307f1 100644 --- a/tests/test_litellm/integrations/focus/test_transformer.py +++ b/tests/test_litellm/integrations/focus/test_transformer.py @@ -30,34 +30,55 @@ def _make_row(**overrides) -> dict: return base -def test_service_name_uses_model_group_when_present(): - row = _make_row(model_group="sonnet-group", model="claude-sonnet") +def test_service_name_uses_custom_llm_provider_when_present(): + row = _make_row(custom_llm_provider="anthropic", model="claude-sonnet") frame = pl.DataFrame([row]) result = FocusTransformer().transform(frame) - assert result["ServiceName"][0] == "sonnet-group" + assert result["ServiceName"][0] == "anthropic" -def test_service_name_falls_back_to_model_when_model_group_blank(): - row = _make_row(model_group="", model="claude-sonnet") +def test_service_name_falls_back_to_model_when_provider_blank(): + row = _make_row(custom_llm_provider="", 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") +def test_service_name_falls_back_to_model_when_provider_null(): + row = _make_row(custom_llm_provider=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="") + row = _make_row(custom_llm_provider="", model="") frame = pl.DataFrame([row]) result = FocusTransformer().transform(frame) assert result["ServiceName"][0] == "unknown" +def test_service_category_uses_model_group_when_present(): + row = _make_row(model_group="sonnet-group") + frame = pl.DataFrame([row]) + result = FocusTransformer().transform(frame) + assert result["ServiceCategory"][0] == "sonnet-group" + + +def test_service_category_null_when_model_group_blank(): + row = _make_row(model_group="") + frame = pl.DataFrame([row]) + result = FocusTransformer().transform(frame) + assert result["ServiceCategory"][0] is None + + +def test_service_category_null_when_model_group_null(): + row = _make_row(model_group=None) + frame = pl.DataFrame([row]) + result = FocusTransformer().transform(frame) + assert result["ServiceCategory"][0] is None + + def test_tags_include_token_counts(): row = _make_row(prompt_tokens=100, completion_tokens=50) frame = pl.DataFrame([row])