From 9cc6df6b8d1a042ce28d80c1a57d3e2c0c3b3fd4 Mon Sep 17 00:00:00 2001 From: Harshit28j Date: Wed, 11 Mar 2026 14:51:03 +0530 Subject: [PATCH] Fix Greptile round 4: preserve backward compat, add guards, fix defaults - Revert FocusExportEngine.dry_run_export_usage_data to use original raw column names (spend, total_tokens, team_id, model) preserving backward compatibility for existing callers - Vantage dry-run endpoint computes its own summary from FOCUS columns independently, avoiding coupling to the engine method - Set VantageExportRequest.limit default to 500 (was None) matching docstring - Add empty-settings guard in /vantage/export returning 404 instead of deferring ValueError to runtime - Fix misleading docstring in _build_tags_expr about Rust-level execution - Clarify Tags schema comment: parquet is self-describing so existing exports are unaffected Co-Authored-By: Claude Opus 4.6 --- litellm/integrations/focus/export_engine.py | 7 ++++--- litellm/integrations/focus/schema.py | 5 ++++- litellm/integrations/focus/transformer.py | 6 ++++-- litellm/proxy/spend_tracking/vantage_endpoints.py | 10 +++++++++- litellm/types/proxy/vantage_endpoints.py | 2 +- 5 files changed, 22 insertions(+), 8 deletions(-) diff --git a/litellm/integrations/focus/export_engine.py b/litellm/integrations/focus/export_engine.py index 0cad1b3acd..1f51e1bac9 100644 --- a/litellm/integrations/focus/export_engine.py +++ b/litellm/integrations/focus/export_engine.py @@ -55,9 +55,10 @@ class FocusExportEngine: summary = { "total_records": len(normalized), - "total_spend": self._sum_column(normalized, "BilledCost"), - "unique_teams": self._count_unique(normalized, "SubAccountId"), - "unique_models": self._count_unique(normalized, "ResourceType"), + "total_spend": self._sum_column(data, "spend"), + "total_tokens": self._sum_column(data, "total_tokens"), + "unique_teams": self._count_unique(data, "team_id"), + "unique_models": self._count_unique(data, "model"), } return { diff --git a/litellm/integrations/focus/schema.py b/litellm/integrations/focus/schema.py index 127e8ee03a..c06ca9982a 100644 --- a/litellm/integrations/focus/schema.py +++ b/litellm/integrations/focus/schema.py @@ -45,7 +45,10 @@ FOCUS_NORMALIZED_SCHEMA = pl.Schema( ("SubAccountType", pl.String), # Changed from pl.Object to pl.String to hold JSON metadata # (team_id, user_id, etc.) needed by Vantage Token Allocation. - # Previously Tags was always None so no existing data is lost. + # This schema is only used for creating empty DataFrames (e.g. + # when transform() receives no rows). Parquet files are + # self-describing and embed their own schema, so existing S3 + # exports are unaffected. Previously Tags was always None. ("Tags", pl.String), ] ) diff --git a/litellm/integrations/focus/transformer.py b/litellm/integrations/focus/transformer.py index adac158651..b7d28e3dbb 100644 --- a/litellm/integrations/focus/transformer.py +++ b/litellm/integrations/focus/transformer.py @@ -25,8 +25,10 @@ _TAG_KEYS = ( def _build_tags_expr(available_keys: list[str]) -> pl.Expr: """Build a Polars expression that produces a JSON Tags string per row. - Uses ``pl.struct`` + ``map_elements`` so the heavy iteration stays inside - Polars rather than materialising every row to a Python dict first. + Uses ``pl.struct`` + ``map_elements`` to avoid materialising the entire + DataFrame to a list of Python dicts. The JSON serialisation callback + still runs in Python (GIL-bound), but struct-packing and loop dispatch + are handled by Polars' Rust engine. """ def _struct_to_json(row: dict) -> str: diff --git a/litellm/proxy/spend_tracking/vantage_endpoints.py b/litellm/proxy/spend_tracking/vantage_endpoints.py index d3d6787e5f..d354ed2962 100644 --- a/litellm/proxy/spend_tracking/vantage_endpoints.py +++ b/litellm/proxy/spend_tracking/vantage_endpoints.py @@ -388,7 +388,8 @@ async def vantage_dry_run_export( usage_sample = data.head(min(50, len(data))).to_dicts() if not data.is_empty() else [] normalized_sample = normalized.head(min(50, len(normalized))).to_dicts() if not normalized.is_empty() else [] - # Use the same column names as FocusExportEngine.dry_run_export_usage_data + # Compute summary from the FOCUS-normalized DataFrame. + # These use post-transform column names specific to this endpoint. summary = { "total_records": len(normalized), "total_spend": FocusExportEngine._sum_column(normalized, "BilledCost"), @@ -459,6 +460,13 @@ async def vantage_export( logger = _get_registered_vantage_logger() if logger is None: settings = await _get_vantage_settings() + if not settings: + raise HTTPException( + status_code=404, + detail={ + "error": "Vantage settings not found. Please initialize settings first using /vantage/init" + }, + ) logger = VantageLogger( api_key=settings.get("api_key"), integration_token=settings.get("integration_token"), diff --git a/litellm/types/proxy/vantage_endpoints.py b/litellm/types/proxy/vantage_endpoints.py index 165c21316b..ce2bad84e3 100644 --- a/litellm/types/proxy/vantage_endpoints.py +++ b/litellm/types/proxy/vantage_endpoints.py @@ -32,7 +32,7 @@ class VantageExportRequest(BaseModel): """Request model for Vantage export operations""" limit: Optional[int] = Field( - None, description="Optional limit on number of records to export" + 500, description="Limit on number of records to export (default: 500)" ) start_time_utc: Optional[datetime] = Field( None, description="Start time for data export in UTC"