mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-12 03:48:56 +00:00
* refactor: expose core private symbols with compatibility aliases Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * types: narrow core migration diagnostics Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix: preserve runtime behavior in core symbol migration Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix: preserve core private usage migration behavior Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix: preserve private value rebinding compatibility Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix: restore optional imports and cover public helpers Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test: exempt router property from call coverage Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test: align recursive detector ignore names Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: mateo <mateo@berri.ai> Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
70 lines
2.7 KiB
Python
70 lines
2.7 KiB
Python
from datetime import datetime, timezone
|
|
|
|
import polars as pl
|
|
|
|
from litellm.integrations.focus.database import FocusLiteLLMDatabase
|
|
from litellm.integrations.focus.destinations import FocusTimeWindow
|
|
from litellm.integrations.focus.destinations.factory import FocusDestinationFactory
|
|
from litellm.integrations.focus.export_engine import FocusExportEngine
|
|
from litellm.integrations.focus.serializers import FocusParquetSerializer
|
|
from litellm.integrations.focus.transformer import FocusTransformer
|
|
|
|
|
|
def _vantage_csv_engine() -> FocusExportEngine:
|
|
return FocusExportEngine(
|
|
provider="vantage",
|
|
export_format="csv",
|
|
prefix="exports",
|
|
destination_config={"api_key": "test-key", "integration_token": "test-token"},
|
|
)
|
|
|
|
|
|
def test_legacy_engine_properties_forward_to_public_storage():
|
|
engine = _vantage_csv_engine()
|
|
|
|
replacement_destination = FocusDestinationFactory.create(
|
|
provider="vantage",
|
|
prefix="replacement",
|
|
config={"api_key": "test-key", "integration_token": "test-token"},
|
|
)
|
|
engine._destination = replacement_destination
|
|
assert engine._destination is engine.destination is replacement_destination
|
|
|
|
replacement_serializer = FocusParquetSerializer()
|
|
engine._serializer = replacement_serializer
|
|
assert engine._serializer is engine.serializer is replacement_serializer
|
|
|
|
replacement_transformer = FocusTransformer()
|
|
engine._transformer = replacement_transformer
|
|
assert engine._transformer is engine.transformer is replacement_transformer
|
|
|
|
replacement_database = FocusLiteLLMDatabase()
|
|
engine._database = replacement_database
|
|
assert engine._database is engine.database is replacement_database
|
|
|
|
|
|
def test_build_filename_uses_the_window_and_serializer_format():
|
|
engine = _vantage_csv_engine()
|
|
window = FocusTimeWindow(
|
|
start_time=datetime(2024, 1, 2, 3, 4, 5, tzinfo=timezone.utc),
|
|
end_time=datetime(2024, 1, 2, 4, 5, 6, tzinfo=timezone.utc),
|
|
frequency="hourly",
|
|
)
|
|
|
|
assert engine.build_filename(window) == "usage_20240102T030405Z_20240102T040506Z.csv"
|
|
|
|
|
|
def test_export_aggregates_compute_values_and_handle_missing_columns():
|
|
frame = pl.DataFrame(
|
|
{
|
|
"spend": [1.5, 2.5, 1.0],
|
|
"team_id": ["team-a", "team-a", "team-b"],
|
|
}
|
|
)
|
|
|
|
assert FocusExportEngine.sum_column(frame, "spend") == 5.0
|
|
assert FocusExportEngine.count_unique(frame, "team_id") == 2
|
|
assert FocusExportEngine.sum_column(frame, "missing") == 0.0
|
|
assert FocusExportEngine.count_unique(frame, "missing") == 0
|
|
assert FocusExportEngine.sum_column(pl.DataFrame(), "spend") == 0.0
|
|
assert FocusExportEngine.count_unique(pl.DataFrame(), "team_id") == 0
|