scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-14 23:21:35 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 102

Remap ServiceName to custom_llm_provider and ServiceCategory to model_group

Browse source 
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 <noreply@anthropic.com>
This commit is contained in:
Harshit28j 2026-03-21 04:46:48 +05:30
parent bb7c3afc7b
commit be43680949
2 changed files with 37 additions and 10 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

10
litellm/integrations/focus/transformer.py
View file

@ -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) != "")

37
tests/test_litellm/integrations/focus/test_transformer.py
View file

@ -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])