From 4583c90194c56c4bbc299589fa56138827bb7823 Mon Sep 17 00:00:00 2001 From: Harshit28j Date: Wed, 11 Mar 2026 14:38:19 +0530 Subject: [PATCH] Address remaining Greptile feedback: mask token, reuse HTTP client, align columns - Mask integration_token in GET /vantage/settings response (renamed field to integration_token_masked) - Reuse single httpx.AsyncClient across all batch uploads in deliver() - Align FocusExportEngine.dry_run_export_usage_data to use post-transform FOCUS columns (BilledCost, SubAccountId, ResourceType) matching the Vantage dry-run endpoint Co-Authored-By: Claude Opus 4.6 --- .../focus/destinations/vantage_destination.py | 58 +++++++++++-------- litellm/integrations/focus/export_engine.py | 7 +-- .../proxy/spend_tracking/vantage_endpoints.py | 4 +- litellm/types/proxy/vantage_endpoints.py | 5 +- 4 files changed, 41 insertions(+), 33 deletions(-) diff --git a/litellm/integrations/focus/destinations/vantage_destination.py b/litellm/integrations/focus/destinations/vantage_destination.py index e8d9532c41c..e2a21c773f4 100644 --- a/litellm/integrations/focus/destinations/vantage_destination.py +++ b/litellm/integrations/focus/destinations/vantage_destination.py @@ -56,21 +56,25 @@ class FocusVantageDestination(FocusDestination): verbose_logger.debug("Vantage destination: empty content, skipping upload") return - # Check both size and row-count limits before single-shot upload - lines = content.split(b"\n") - data_line_count = sum(1 for line in lines[1:] if line.strip()) - within_limits = ( - len(content) <= VANTAGE_MAX_BYTES_PER_UPLOAD - and data_line_count <= VANTAGE_MAX_ROWS_PER_UPLOAD - ) - if within_limits: - await self._upload_csv(content, filename) - return + # Reuse a single HTTP client for the entire deliver() call + async with httpx.AsyncClient(timeout=60.0) as client: + # Check both size and row-count limits before single-shot upload + lines = content.split(b"\n") + data_line_count = sum(1 for line in lines[1:] if line.strip()) + within_limits = ( + len(content) <= VANTAGE_MAX_BYTES_PER_UPLOAD + and data_line_count <= VANTAGE_MAX_ROWS_PER_UPLOAD + ) + if within_limits: + await self._upload_csv(client, content, filename) + return - # Otherwise split into batches respecting both limits - await self._upload_batched(content, filename) + # Otherwise split into batches respecting both limits + await self._upload_batched(client, content, filename) - async def _upload_csv(self, csv_bytes: bytes, filename: str) -> None: + async def _upload_csv( + self, client: httpx.AsyncClient, csv_bytes: bytes, filename: str + ) -> None: url = ( f"{self.base_url}/v2/integrations/" f"{self.integration_token}/costs.csv" @@ -79,13 +83,12 @@ class FocusVantageDestination(FocusDestination): "Authorization": f"Bearer {self.api_key}", } - async with httpx.AsyncClient(timeout=60.0) as client: - response = await client.post( - url, - headers=headers, - files={"file": (filename, csv_bytes, "text/csv")}, - ) - response.raise_for_status() + response = await client.post( + url, + headers=headers, + files={"file": (filename, csv_bytes, "text/csv")}, + ) + response.raise_for_status() verbose_logger.debug( "Vantage destination: uploaded %d bytes (%s)", @@ -93,7 +96,9 @@ class FocusVantageDestination(FocusDestination): filename, ) - async def _upload_batched(self, csv_bytes: bytes, filename: str) -> None: + async def _upload_batched( + self, client: httpx.AsyncClient, csv_bytes: bytes, filename: str + ) -> None: """Split the CSV into batches and upload each.""" lines = csv_bytes.split(b"\n") header = lines[0] @@ -106,14 +111,17 @@ class FocusVantageDestination(FocusDestination): # If a single batch still exceeds 2 MB, split further by size if len(batch_csv) > VANTAGE_MAX_BYTES_PER_UPLOAD: - await self._upload_size_limited(header, batch_lines, filename, batch_num) + await self._upload_size_limited( + client, header, batch_lines, filename, batch_num + ) else: batch_filename = f"{filename}.part{batch_num}" - await self._upload_csv(batch_csv, batch_filename) + await self._upload_csv(client, batch_csv, batch_filename) batch_num += 1 async def _upload_size_limited( self, + client: httpx.AsyncClient, header: bytes, data_lines: list[bytes], filename: str, @@ -129,7 +137,7 @@ 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(batch_csv, batch_filename) + await self._upload_csv(client, batch_csv, batch_filename) current_chunk = [] current_size = len(header) + 1 sub_batch += 1 @@ -139,4 +147,4 @@ 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(batch_csv, batch_filename) + await self._upload_csv(client, batch_csv, batch_filename) diff --git a/litellm/integrations/focus/export_engine.py b/litellm/integrations/focus/export_engine.py index a9361d0e988..0cad1b3acd8 100644 --- a/litellm/integrations/focus/export_engine.py +++ b/litellm/integrations/focus/export_engine.py @@ -55,10 +55,9 @@ class FocusExportEngine: summary = { "total_records": len(normalized), - "total_spend": self._sum_column(normalized, "spend"), - "total_tokens": self._sum_column(normalized, "total_tokens"), - "unique_teams": self._count_unique(normalized, "team_id"), - "unique_models": self._count_unique(normalized, "model"), + "total_spend": self._sum_column(normalized, "BilledCost"), + "unique_teams": self._count_unique(normalized, "SubAccountId"), + "unique_models": self._count_unique(normalized, "ResourceType"), } return { diff --git a/litellm/proxy/spend_tracking/vantage_endpoints.py b/litellm/proxy/spend_tracking/vantage_endpoints.py index a553d1061ec..d3d6787e5fb 100644 --- a/litellm/proxy/spend_tracking/vantage_endpoints.py +++ b/litellm/proxy/spend_tracking/vantage_endpoints.py @@ -155,7 +155,7 @@ async def get_vantage_settings( if not settings: return VantageSettingsView( api_key_masked=None, - integration_token=None, + integration_token_masked=None, base_url=None, status=None, ) @@ -164,7 +164,7 @@ async def get_vantage_settings( return VantageSettingsView( api_key_masked=masked_settings.get("api_key"), - integration_token=settings.get("integration_token"), + integration_token_masked=masked_settings.get("integration_token"), base_url=settings.get("base_url"), status="configured", ) diff --git a/litellm/types/proxy/vantage_endpoints.py b/litellm/types/proxy/vantage_endpoints.py index 394677d1e8f..165c21316bd 100644 --- a/litellm/types/proxy/vantage_endpoints.py +++ b/litellm/types/proxy/vantage_endpoints.py @@ -62,8 +62,9 @@ class VantageSettingsView(BaseModel): None, description="Masked API key showing only first 4 and last 4 characters", ) - integration_token: Optional[str] = Field( - None, description="Vantage integration token" + integration_token_masked: Optional[str] = Field( + None, + description="Masked integration token showing only first 4 and last 4 characters", ) base_url: Optional[str] = Field(None, description="Vantage API base URL") status: Optional[str] = Field(None, description="Configuration status")