From d35abfb55f926ae784b864356531923d649434ef Mon Sep 17 00:00:00 2001 From: Harshit28j Date: Wed, 11 Mar 2026 16:06:59 +0530 Subject: [PATCH] Fix data truncation, sub-batch resilience, and env var safety - Split VantageExportRequest (limit=None) and VantageDryRunRequest (limit=500) so actual exports don't silently truncate large datasets - Add try/except around each sub-batch upload in _upload_size_limited, consistent with _upload_batched's continue-on-failure guarantee - Guard VANTAGE_EXPORT_INTERVAL_SECONDS against non-numeric values with try/except instead of bare int() cast Co-Authored-By: Claude Opus 4.6 --- .../focus/destinations/vantage_destination.py | 28 +++++++++++++++++-- .../integrations/vantage/vantage_logger.py | 10 ++++++- .../proxy/spend_tracking/vantage_endpoints.py | 5 ++-- litellm/types/proxy/vantage_endpoints.py | 12 ++++++-- 4 files changed, 47 insertions(+), 8 deletions(-) diff --git a/litellm/integrations/focus/destinations/vantage_destination.py b/litellm/integrations/focus/destinations/vantage_destination.py index 59c4e0cdb9b..0e06c8f4413 100644 --- a/litellm/integrations/focus/destinations/vantage_destination.py +++ b/litellm/integrations/focus/destinations/vantage_destination.py @@ -145,12 +145,15 @@ class FocusVantageDestination(FocusDestination): """Upload lines in chunks that stay under the 2 MB size limit. Individual rows that exceed the limit on their own are skipped with - a warning — they cannot be split further. + a warning — they cannot be split further. Sub-batch failures are + recorded and the first error is re-raised after all sub-batches have + been attempted, consistent with ``_upload_batched``. """ current_chunk: list[bytes] = [] current_size = len(header) + 1 # header + newline sub_batch = 0 header_size = len(header) + 1 + first_error: Optional[Exception] = None for line in data_lines: line_size = len(line) + 1 # line + newline @@ -166,7 +169,15 @@ class FocusVantageDestination(FocusDestination): if current_size + line_size > VANTAGE_MAX_BYTES_PER_UPLOAD and current_chunk: batch_csv = header + b"\n" + b"\n".join(current_chunk) + b"\n" batch_filename = f"{filename}.part{batch_offset}_{sub_batch}" - await self._upload_csv(client, batch_csv, batch_filename) + try: + await self._upload_csv(client, batch_csv, batch_filename) + except Exception as e: + verbose_logger.error( + "Vantage destination: sub-batch %s failed: %s", + batch_filename, e, + ) + if first_error is None: + first_error = e current_chunk = [] current_size = header_size sub_batch += 1 @@ -176,4 +187,15 @@ class FocusVantageDestination(FocusDestination): if current_chunk: batch_csv = header + b"\n" + b"\n".join(current_chunk) + b"\n" batch_filename = f"{filename}.part{batch_offset}_{sub_batch}" - await self._upload_csv(client, batch_csv, batch_filename) + try: + await self._upload_csv(client, batch_csv, batch_filename) + except Exception as e: + verbose_logger.error( + "Vantage destination: sub-batch %s failed: %s", + batch_filename, e, + ) + if first_error is None: + first_error = e + + if first_error is not None: + raise first_error diff --git a/litellm/integrations/vantage/vantage_logger.py b/litellm/integrations/vantage/vantage_logger.py index 01529d537f7..df51bb72866 100644 --- a/litellm/integrations/vantage/vantage_logger.py +++ b/litellm/integrations/vantage/vantage_logger.py @@ -53,7 +53,15 @@ class VantageLogger(FocusLogger): ).lower() raw_interval = interval_seconds or os.getenv("VANTAGE_EXPORT_INTERVAL_SECONDS") - resolved_interval = int(raw_interval) if raw_interval is not None else None + resolved_interval: Optional[int] = None + if raw_interval is not None: + try: + resolved_interval = int(raw_interval) + except (ValueError, TypeError): + verbose_logger.warning( + "Invalid VANTAGE_EXPORT_INTERVAL_SECONDS value: %s, ignoring", + raw_interval, + ) destination_config: Dict[str, Any] = {} if resolved_api_key: diff --git a/litellm/proxy/spend_tracking/vantage_endpoints.py b/litellm/proxy/spend_tracking/vantage_endpoints.py index f478ba51604..cd78cf7ebd9 100644 --- a/litellm/proxy/spend_tracking/vantage_endpoints.py +++ b/litellm/proxy/spend_tracking/vantage_endpoints.py @@ -12,6 +12,7 @@ from litellm.proxy.common_utils.encrypt_decrypt_utils import ( encrypt_value_helper, ) from litellm.types.proxy.vantage_endpoints import ( + VantageDryRunRequest, VantageExportRequest, VantageExportResponse, VantageInitRequest, @@ -356,7 +357,7 @@ async def init_vantage_settings( response_model=VantageExportResponse, ) async def vantage_dry_run_export( - request: VantageExportRequest, + request: VantageDryRunRequest, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ): """ @@ -365,7 +366,7 @@ async def vantage_dry_run_export( Returns the data that would be exported without actually sending it to Vantage. Parameters: - - limit: Optional limit on number of records to process (default: 500) + - limit: Limit on number of records to preview (default: 500) Only admin users can perform Vantage exports. """ diff --git a/litellm/types/proxy/vantage_endpoints.py b/litellm/types/proxy/vantage_endpoints.py index ce2bad84e31..12d9099f819 100644 --- a/litellm/types/proxy/vantage_endpoints.py +++ b/litellm/types/proxy/vantage_endpoints.py @@ -29,10 +29,10 @@ class VantageInitResponse(BaseModel): class VantageExportRequest(BaseModel): - """Request model for Vantage export operations""" + """Request model for Vantage export operations (actual export, no default limit)""" limit: Optional[int] = Field( - 500, description="Limit on number of records to export (default: 500)" + None, description="Optional limit on number of records to export (default: no limit)" ) start_time_utc: Optional[datetime] = Field( None, description="Start time for data export in UTC" @@ -42,6 +42,14 @@ class VantageExportRequest(BaseModel): ) +class VantageDryRunRequest(BaseModel): + """Request model for Vantage dry-run operations (capped for preview)""" + + limit: Optional[int] = Field( + 500, description="Limit on number of records to preview (default: 500)" + ) + + class VantageExportResponse(BaseModel): """Response model for Vantage export operations"""