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:
Harshit28j 2026-03-19 18:04:42 +05:30
parent b51201620a
commit bb7c3afc7b
2 changed files with 84 additions and 1 deletions

View file

@ -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"),

View 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