From be43680949c8b95b92456b99ca869a08a8db8c13 Mon Sep 17 00:00:00 2001 From: Harshit28j Date: Sat, 21 Mar 2026 04:46:48 +0530 Subject: [PATCH] Remap ServiceName to custom_llm_provider and ServiceCategory to model_group Per Vantage team feedback: - ServiceName now maps from custom_llm_provider (fallback: model -> "unknown") - ServiceCategory now maps from model_group (null when blank) instead of hardcoded "AI and Machine Learning" - Updated tests for new mappings Co-Authored-By: Claude Opus 4.6 --- litellm/integrations/focus/transformer.py | 10 ++++- .../integrations/focus/test_transformer.py | 37 +++++++++++++++---- 2 files changed, 37 insertions(+), 10 deletions(-) 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])