fix(cloudzero): preserve late resource tags (#39873)

* fix(cloudzero): infer daily batch schema from every row

pl.DataFrame defaults to inferring column types from the first 100 rows,
so a day whose batch starts with more than 100 rows missing team_alias,
api_key_alias or user_email typed that column as Null and then raised a
ComputeError on the first row that had a value, failing the whole export
with a 500 and sending nothing.

Pass infer_schema_length=None when rebuilding each day's DataFrame, the
same guard the usage query already uses.

* test(cloudzero): cover late tag schema inference

Exercise the CloudZero resource tag field after a long run of missing values so a finite inference window fails the regression test.

* fix(cloudzero): preserve late resource tags

* style(cloudzero): remove redundant test comment
This commit is contained in:
yucheng-berri 2026-09-05 12:10:05 -07:00 committed by GitHub
parent 73e1cfb378
commit 877197918b
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 28 additions and 1 deletions

View file

@ -95,7 +95,7 @@ class CBFTransformer:
if len(cbf_data) > 0:
console.print(f"[green]✓ Successfully transformed {len(cbf_data):,} records[/green]")
return pl.DataFrame(cbf_data)
return pl.DataFrame(cbf_data, infer_schema_length=None)
def _create_cbf_record(self, row: dict[str, object]) -> CBFRecord:
"""Create a single CBF record from LiteLLM daily spend row."""

View file

@ -86,6 +86,33 @@ class TestCBFTransformer:
assert result.is_empty()
def test_transform_keeps_tags_first_seen_after_row_100(self):
transformer = CBFTransformer()
teamless_rows = 101
team_rows = 2
total_rows = teamless_rows + team_rows
data = pl.DataFrame(
{
"date": ["2025-01-19"] * total_rows,
"successful_requests": [1] * total_rows,
"spend": [0.5] * total_rows,
"prompt_tokens": [10] * total_rows,
"completion_tokens": [5] * total_rows,
"model": ["gpt-4"] * total_rows,
"custom_llm_provider": ["openai"] * total_rows,
"api_key": ["sk-late-team"] * total_rows,
"team_id": pl.Series([None] * teamless_rows + ["team-late"] * team_rows, dtype=pl.String),
"team_alias": pl.Series([None] * teamless_rows + ["Late Team"] * team_rows, dtype=pl.String),
}
)
result = transformer.transform(data)
assert len(result) == total_rows
assert "resource/tag:team_alias" in result.columns
assert result["resource/tag:team_alias"].to_list() == [None] * teamless_rows + ["Late Team"] * team_rows
assert result["resource/tag:entity_id"].to_list() == [None] * teamless_rows + ["Late Team"] * team_rows
def test_create_cbf_record(self):
"""Test _create_cbf_record method with valid row data."""
transformer = CBFTransformer()