mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-15 23:31:29 +00:00
877197918b
* 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
312 lines
10 KiB
Python
312 lines
10 KiB
Python
from datetime import datetime
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
import polars as pl
|
|
import pytest
|
|
|
|
|
|
from litellm.integrations.cloudzero.transform import CBFTransformer
|
|
from litellm.types.integrations.cloudzero import CBFRecord
|
|
|
|
|
|
class TestCBFTransformer:
|
|
"""Test suite for CBFTransformer class."""
|
|
|
|
def test_init(self):
|
|
"""Test CBFTransformer initialization."""
|
|
transformer = CBFTransformer()
|
|
assert hasattr(transformer, "czrn_generator")
|
|
assert transformer.czrn_generator is not None
|
|
|
|
def test_transform_empty_dataframe(self):
|
|
"""Test transform method with empty DataFrame."""
|
|
transformer = CBFTransformer()
|
|
empty_df = pl.DataFrame()
|
|
|
|
result = transformer.transform(empty_df)
|
|
|
|
assert result.is_empty()
|
|
assert isinstance(result, pl.DataFrame)
|
|
|
|
def test_transform_with_zero_successful_requests(self):
|
|
"""Test transform method filters out records with zero successful_requests."""
|
|
transformer = CBFTransformer()
|
|
data = pl.DataFrame(
|
|
{
|
|
"date": ["2025-01-19"],
|
|
"successful_requests": [0],
|
|
"spend": [10.0],
|
|
"entity_id": ["test_entity"],
|
|
"model": ["gpt-4"],
|
|
}
|
|
)
|
|
|
|
result = transformer.transform(data)
|
|
|
|
assert result.is_empty()
|
|
|
|
def test_transform_with_valid_data(self):
|
|
"""Test transform method with valid data."""
|
|
transformer = CBFTransformer()
|
|
with patch.object(transformer, "_create_cbf_record") as mock_create:
|
|
mock_create.return_value = CBFRecord({"test": "data"})
|
|
|
|
data = pl.DataFrame(
|
|
{
|
|
"date": ["2025-01-19"],
|
|
"successful_requests": [5],
|
|
"spend": [10.0],
|
|
"entity_id": ["test_entity"],
|
|
"model": ["gpt-4"],
|
|
}
|
|
)
|
|
|
|
result = transformer.transform(data)
|
|
|
|
assert len(result) == 1
|
|
mock_create.assert_called_once()
|
|
|
|
def test_transform_handles_czrn_generation_failures(self):
|
|
"""Test transform method handles CZRN generation failures gracefully."""
|
|
transformer = CBFTransformer()
|
|
with patch.object(transformer, "_create_cbf_record") as mock_create:
|
|
mock_create.side_effect = Exception("CZRN generation failed")
|
|
|
|
data = pl.DataFrame(
|
|
{
|
|
"date": ["2025-01-19"],
|
|
"successful_requests": [5],
|
|
"spend": [10.0],
|
|
"entity_id": ["test_entity"],
|
|
"model": ["gpt-4"],
|
|
}
|
|
)
|
|
|
|
result = transformer.transform(data)
|
|
|
|
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()
|
|
with (
|
|
patch.object(
|
|
transformer.czrn_generator, "create_from_litellm_data"
|
|
) as mock_czrn,
|
|
patch.object(
|
|
transformer.czrn_generator, "extract_components"
|
|
) as mock_extract,
|
|
):
|
|
|
|
mock_czrn.return_value = "test-czrn"
|
|
mock_extract.return_value = (
|
|
"service",
|
|
"provider",
|
|
"region",
|
|
"account",
|
|
"resource",
|
|
"local_id",
|
|
)
|
|
|
|
row = {
|
|
"date": "2025-01-19",
|
|
"spend": 10.5,
|
|
"prompt_tokens": 100,
|
|
"completion_tokens": 50,
|
|
"entity_id": "test_entity",
|
|
"model": "gpt-4",
|
|
"entity_type": "user",
|
|
"model_group": "openai",
|
|
"custom_llm_provider": "openai",
|
|
"api_key": "sk-test123",
|
|
"api_requests": 5,
|
|
"successful_requests": 5,
|
|
"failed_requests": 0,
|
|
}
|
|
|
|
result = transformer._create_cbf_record(row)
|
|
|
|
assert isinstance(result, CBFRecord)
|
|
assert result["cost/cost"] == 10.5
|
|
assert result["usage/amount"] == 150 # 100 + 50
|
|
assert result["usage/units"] == "tokens"
|
|
assert result["resource/id"] == "test-czrn"
|
|
|
|
def test_create_cbf_record_adds_user_email_tag(self):
|
|
"""Test that user_email field is emitted as a resource tag when present."""
|
|
transformer = CBFTransformer()
|
|
with (
|
|
patch.object(
|
|
transformer.czrn_generator, "create_from_litellm_data"
|
|
) as mock_czrn,
|
|
patch.object(
|
|
transformer.czrn_generator, "extract_components"
|
|
) as mock_extract,
|
|
):
|
|
|
|
mock_czrn.return_value = "test-czrn"
|
|
mock_extract.return_value = (
|
|
"service",
|
|
"provider",
|
|
"region",
|
|
"account",
|
|
"resource",
|
|
"local_id",
|
|
)
|
|
|
|
row = {
|
|
"date": "2025-01-19",
|
|
"spend": 1.0,
|
|
"prompt_tokens": 10,
|
|
"completion_tokens": 5,
|
|
"model": "gpt-4",
|
|
"api_key": "sk-useremail",
|
|
"team_id": "team-123",
|
|
"team_alias": "Dev Team",
|
|
"user_email": "user@example.com",
|
|
}
|
|
|
|
result = transformer._create_cbf_record(row)
|
|
|
|
assert result["resource/tag:user_email"] == "user@example.com"
|
|
|
|
def test_create_cbf_record_omits_empty_user_email(self):
|
|
"""Test that empty user_email values are not added as resource tags."""
|
|
transformer = CBFTransformer()
|
|
with (
|
|
patch.object(
|
|
transformer.czrn_generator, "create_from_litellm_data"
|
|
) as mock_czrn,
|
|
patch.object(
|
|
transformer.czrn_generator, "extract_components"
|
|
) as mock_extract,
|
|
):
|
|
|
|
mock_czrn.return_value = "test-czrn"
|
|
mock_extract.return_value = (
|
|
"service",
|
|
"provider",
|
|
"region",
|
|
"account",
|
|
"resource",
|
|
"local_id",
|
|
)
|
|
|
|
row = {
|
|
"date": "2025-01-19",
|
|
"spend": 1.0,
|
|
"prompt_tokens": 10,
|
|
"completion_tokens": 5,
|
|
"model": "gpt-4",
|
|
"api_key": "sk-useremail",
|
|
"team_id": "team-123",
|
|
"team_alias": "Dev Team",
|
|
"user_email": None,
|
|
}
|
|
|
|
result = transformer._create_cbf_record(row)
|
|
|
|
assert "resource/tag:user_email" not in result
|
|
|
|
def test_create_cbf_record_minimal_data(self):
|
|
"""Test _create_cbf_record method with minimal row data."""
|
|
transformer = CBFTransformer()
|
|
with (
|
|
patch.object(
|
|
transformer.czrn_generator, "create_from_litellm_data"
|
|
) as mock_czrn,
|
|
patch.object(
|
|
transformer.czrn_generator, "extract_components"
|
|
) as mock_extract,
|
|
):
|
|
|
|
mock_czrn.return_value = "test-czrn"
|
|
mock_extract.return_value = (
|
|
"service",
|
|
"provider",
|
|
"region",
|
|
"account",
|
|
"resource",
|
|
"local_id",
|
|
)
|
|
|
|
row = {"date": "2025-01-19", "spend": 0.0}
|
|
|
|
result = transformer._create_cbf_record(row)
|
|
|
|
assert isinstance(result, CBFRecord)
|
|
assert result["cost/cost"] == 0.0
|
|
assert result["usage/amount"] == 0 # no tokens
|
|
assert result["usage/units"] == "tokens"
|
|
|
|
def test_parse_date_with_valid_string(self):
|
|
"""Test _parse_date method with valid date string."""
|
|
transformer = CBFTransformer()
|
|
|
|
result = transformer._parse_date("2025-01-19")
|
|
|
|
assert isinstance(result, datetime)
|
|
assert result.year == 2025
|
|
assert result.month == 1
|
|
assert result.day == 19
|
|
|
|
def test_parse_date_with_datetime_object(self):
|
|
"""Test _parse_date method with datetime object."""
|
|
transformer = CBFTransformer()
|
|
dt = datetime(2025, 1, 19)
|
|
|
|
result = transformer._parse_date(dt)
|
|
|
|
assert result == dt
|
|
|
|
def test_parse_date_with_none(self):
|
|
"""Test _parse_date method with None."""
|
|
transformer = CBFTransformer()
|
|
|
|
result = transformer._parse_date(None)
|
|
|
|
assert result is None
|
|
|
|
def test_parse_date_with_invalid_string(self):
|
|
"""Test _parse_date method with invalid date string."""
|
|
transformer = CBFTransformer()
|
|
|
|
result = transformer._parse_date("invalid-date")
|
|
|
|
assert result is None
|
|
|
|
def test_parse_date_with_iso_format(self):
|
|
"""Test _parse_date method with ISO format string."""
|
|
transformer = CBFTransformer()
|
|
|
|
result = transformer._parse_date("2025-01-19T10:30:00Z")
|
|
|
|
assert isinstance(result, datetime)
|
|
assert result.year == 2025
|