mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-14 23:21:35 +00:00
Guard ServiceName against null model fallback and add transformer tests
- Add secondary null/empty check on `model` column, defaulting to "unknown" when both model_group and model are blank - Add 6 transformer tests covering ServiceName fallback chain and token counts in Tags Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
parent
b51201620a
commit
bb7c3afc7b
2 changed files with 84 additions and 1 deletions
|
|
@ -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"),
|
||||
|
|
|
|||
78
tests/test_litellm/integrations/focus/test_transformer.py
Normal file
78
tests/test_litellm/integrations/focus/test_transformer.py
Normal file
|
|
@ -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
|
||||
Loading…
Add table
Reference in a new issue