diff --git a/litellm/integrations/cloudzero/transform.py b/litellm/integrations/cloudzero/transform.py index ffc8fe1c1f5..12a0ee55fad 100644 --- a/litellm/integrations/cloudzero/transform.py +++ b/litellm/integrations/cloudzero/transform.py @@ -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.""" diff --git a/tests/test_litellm/integrations/cloudzero/test_transform.py b/tests/test_litellm/integrations/cloudzero/test_transform.py index 3ec2fe6779e..cf8d70702f9 100644 --- a/tests/test_litellm/integrations/cloudzero/test_transform.py +++ b/tests/test_litellm/integrations/cloudzero/test_transform.py @@ -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()