From 69a94a873c1a7f286b3e1249348b92046702c410 Mon Sep 17 00:00:00 2001 From: Harshit28j Date: Wed, 11 Mar 2026 13:51:36 +0530 Subject: [PATCH 01/21] Add Vantage integration for FOCUS CSV export Adds a pluggable Vantage destination to the existing FOCUS export pipeline, enabling LiteLLM to export spend data in FOCUS format directly to Vantage's cost-import API. Supports automatic hourly exports via scheduled background job, with admin API endpoints for manual control and configuration. Includes CSV serializer, batching for 10K row / 2MB API limits, and enriched Tags JSON with team/user/key metadata for Vantage Token Allocation feature. - Add CSV serializer (FocusCsvSerializer) for FOCUS data - Add Vantage API destination with automatic batching - Add VantageLogger that wraps FocusLogger with Vantage defaults - Add proxy endpoints: /vantage/{init,settings,export,dry-run,delete} - Register "vantage" callback in logger registry and literal type - Wire up background job in proxy_server.py startup - Populate Tags column with JSON metadata (team_id, user_id, user_email, etc.) - Add 14 unit tests covering serializer, destination, and factory All tests pass (23 focus tests total, no regressions). Co-Authored-By: Claude Haiku 4.5 --- litellm/__init__.py | 1 + .../focus/destinations/__init__.py | 2 + .../focus/destinations/factory.py | 21 + .../focus/destinations/vantage_destination.py | 136 +++++ litellm/integrations/focus/export_engine.py | 12 +- litellm/integrations/focus/schema.py | 2 +- .../focus/serializers/__init__.py | 3 +- litellm/integrations/focus/serializers/csv.py | 22 + litellm/integrations/focus/transformer.py | 40 +- litellm/integrations/vantage/__init__.py | 0 .../integrations/vantage/vantage_logger.py | 105 ++++ .../custom_logger_registry.py | 2 + litellm/proxy/proxy_server.py | 11 + .../proxy/spend_tracking/vantage_endpoints.py | 488 ++++++++++++++++++ litellm/types/proxy/vantage_endpoints.py | 81 +++ .../integrations/focus/test_csv_serializer.py | 34 ++ .../focus/test_destination_factory.py | 48 ++ .../focus/test_vantage_destination.py | 138 +++++ 18 files changed, 1139 insertions(+), 7 deletions(-) create mode 100644 litellm/integrations/focus/destinations/vantage_destination.py create mode 100644 litellm/integrations/focus/serializers/csv.py create mode 100644 litellm/integrations/vantage/__init__.py create mode 100644 litellm/integrations/vantage/vantage_logger.py create mode 100644 litellm/proxy/spend_tracking/vantage_endpoints.py create mode 100644 litellm/types/proxy/vantage_endpoints.py create mode 100644 tests/test_litellm/integrations/focus/test_csv_serializer.py create mode 100644 tests/test_litellm/integrations/focus/test_destination_factory.py create mode 100644 tests/test_litellm/integrations/focus/test_vantage_destination.py diff --git a/litellm/__init__.py b/litellm/__init__.py index 4fc71e1270..2334858a05 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -143,6 +143,7 @@ _custom_logger_compatible_callbacks_literal = Literal[ "gitlab", "cloudzero", "focus", + "vantage", "posthog", "levo", ] diff --git a/litellm/integrations/focus/destinations/__init__.py b/litellm/integrations/focus/destinations/__init__.py index 233f1da0c9..775d3a259d 100644 --- a/litellm/integrations/focus/destinations/__init__.py +++ b/litellm/integrations/focus/destinations/__init__.py @@ -3,10 +3,12 @@ from .base import FocusDestination, FocusTimeWindow from .factory import FocusDestinationFactory from .s3_destination import FocusS3Destination +from .vantage_destination import FocusVantageDestination __all__ = [ "FocusDestination", "FocusDestinationFactory", "FocusTimeWindow", "FocusS3Destination", + "FocusVantageDestination", ] diff --git a/litellm/integrations/focus/destinations/factory.py b/litellm/integrations/focus/destinations/factory.py index cb7696a11d..01ea6ca9cb 100644 --- a/litellm/integrations/focus/destinations/factory.py +++ b/litellm/integrations/focus/destinations/factory.py @@ -7,6 +7,7 @@ from typing import Any, Dict, Optional from .base import FocusDestination from .s3_destination import FocusS3Destination +from .vantage_destination import FocusVantageDestination class FocusDestinationFactory: @@ -26,6 +27,8 @@ class FocusDestinationFactory: ) if provider_lower == "s3": return FocusS3Destination(prefix=prefix, config=normalized_config) + if provider_lower == "vantage": + return FocusVantageDestination(prefix=prefix, config=normalized_config) raise NotImplementedError( f"Provider '{provider}' not supported for Focus export" ) @@ -54,6 +57,24 @@ class FocusDestinationFactory: if not resolved.get("bucket_name"): raise ValueError("FOCUS_S3_BUCKET_NAME must be provided for S3 exports") return {k: v for k, v in resolved.items() if v is not None} + if provider == "vantage": + resolved = { + "api_key": overrides.get("api_key") + or os.getenv("VANTAGE_API_KEY"), + "integration_token": overrides.get("integration_token") + or os.getenv("VANTAGE_INTEGRATION_TOKEN"), + "base_url": overrides.get("base_url") + or os.getenv("VANTAGE_BASE_URL", "https://api.vantage.sh"), + } + if not resolved.get("api_key"): + raise ValueError( + "VANTAGE_API_KEY must be provided for Vantage exports" + ) + if not resolved.get("integration_token"): + raise ValueError( + "VANTAGE_INTEGRATION_TOKEN must be provided for Vantage exports" + ) + return {k: v for k, v in resolved.items() if v is not None} raise NotImplementedError( f"Provider '{provider}' not supported for Focus export configuration" ) diff --git a/litellm/integrations/focus/destinations/vantage_destination.py b/litellm/integrations/focus/destinations/vantage_destination.py new file mode 100644 index 0000000000..5612a513d6 --- /dev/null +++ b/litellm/integrations/focus/destinations/vantage_destination.py @@ -0,0 +1,136 @@ +"""Vantage API destination for Focus export.""" + +from __future__ import annotations + +from typing import Any, Optional + +import httpx + +from litellm._logging import verbose_logger + +from .base import FocusDestination, FocusTimeWindow + +# Vantage enforces a 10,000-row / 2 MB limit per upload. +VANTAGE_MAX_ROWS_PER_UPLOAD = 10_000 +VANTAGE_MAX_BYTES_PER_UPLOAD = 2 * 1024 * 1024 # 2 MB + + +class FocusVantageDestination(FocusDestination): + """Upload FOCUS CSV exports to the Vantage cost-import API.""" + + def __init__( + self, + *, + prefix: str, + config: Optional[dict[str, Any]] = None, + ) -> None: + config = config or {} + api_key = config.get("api_key") + integration_token = config.get("integration_token") + if not api_key: + raise ValueError( + "api_key must be provided for Vantage destination " + "(set VANTAGE_API_KEY env var or pass in destination_config)" + ) + if not integration_token: + raise ValueError( + "integration_token must be provided for Vantage destination " + "(set VANTAGE_INTEGRATION_TOKEN env var or pass in destination_config)" + ) + self.api_key = api_key + self.integration_token = integration_token + self.base_url = config.get( + "base_url", "https://api.vantage.sh" + ) + self.prefix = prefix + + async def deliver( + self, + *, + content: bytes, + time_window: FocusTimeWindow, + filename: str, + ) -> None: + """Upload CSV content to the Vantage API, batching if needed.""" + if not content: + verbose_logger.debug("Vantage destination: empty content, skipping upload") + return + + # If the payload is within limits, send in one shot + if len(content) <= VANTAGE_MAX_BYTES_PER_UPLOAD: + await self._upload_csv(content, filename) + return + + # Otherwise split into chunks by line count + await self._upload_batched(content, filename) + + async def _upload_csv(self, csv_bytes: bytes, filename: str) -> None: + url = ( + f"{self.base_url}/v2/integrations/" + f"{self.integration_token}/costs.csv" + ) + headers = { + "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() + + verbose_logger.debug( + "Vantage destination: uploaded %d bytes (%s)", + len(csv_bytes), + filename, + ) + + async def _upload_batched(self, csv_bytes: bytes, filename: str) -> None: + """Split the CSV into batches and upload each.""" + lines = csv_bytes.split(b"\n") + header = lines[0] + data_lines = [line for line in lines[1:] if line.strip()] + + batch_num = 0 + for start in range(0, len(data_lines), VANTAGE_MAX_ROWS_PER_UPLOAD): + batch_lines = data_lines[start : start + VANTAGE_MAX_ROWS_PER_UPLOAD] + batch_csv = header + b"\n" + b"\n".join(batch_lines) + b"\n" + + # 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) + else: + batch_filename = f"{filename}.part{batch_num}" if batch_num > 0 else filename + await self._upload_csv(batch_csv, batch_filename) + batch_num += 1 + + async def _upload_size_limited( + self, + header: bytes, + data_lines: list[bytes], + filename: str, + batch_offset: int, + ) -> None: + """Upload lines in chunks that stay under the 2 MB size limit.""" + current_chunk: list[bytes] = [] + current_size = len(header) + 1 # header + newline + sub_batch = 0 + + for line in data_lines: + line_size = len(line) + 1 # line + newline + 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) + current_chunk = [] + current_size = len(header) + 1 + sub_batch += 1 + current_chunk.append(line) + current_size += line_size + + 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) diff --git a/litellm/integrations/focus/export_engine.py b/litellm/integrations/focus/export_engine.py index 22ebce2a16..a9361d0e98 100644 --- a/litellm/integrations/focus/export_engine.py +++ b/litellm/integrations/focus/export_engine.py @@ -10,7 +10,7 @@ from litellm._logging import verbose_logger from .database import FocusLiteLLMDatabase from .destinations import FocusDestinationFactory, FocusTimeWindow -from .serializers import FocusParquetSerializer, FocusSerializer +from .serializers import FocusCsvSerializer, FocusParquetSerializer, FocusSerializer from .transformer import FocusTransformer @@ -38,9 +38,13 @@ class FocusExportEngine: self._database = FocusLiteLLMDatabase() def _init_serializer(self) -> FocusSerializer: - if self.export_format != "parquet": - raise NotImplementedError("Only parquet export supported currently") - return FocusParquetSerializer() + if self.export_format == "csv": + return FocusCsvSerializer() + if self.export_format == "parquet": + return FocusParquetSerializer() + raise NotImplementedError( + f"Export format '{self.export_format}' not supported. Use 'parquet' or 'csv'." + ) async def dry_run_export_usage_data(self, limit: Optional[int]) -> Dict[str, Any]: data = await self._database.get_usage_data(limit=limit) diff --git a/litellm/integrations/focus/schema.py b/litellm/integrations/focus/schema.py index ac2f33dad0..6ad725f367 100644 --- a/litellm/integrations/focus/schema.py +++ b/litellm/integrations/focus/schema.py @@ -43,7 +43,7 @@ FOCUS_NORMALIZED_SCHEMA = pl.Schema( ("SubAccountId", pl.String), ("SubAccountName", pl.String), ("SubAccountType", pl.String), - ("Tags", pl.Object), + ("Tags", pl.String), ] ) diff --git a/litellm/integrations/focus/serializers/__init__.py b/litellm/integrations/focus/serializers/__init__.py index 18187bf73e..bdbf520454 100644 --- a/litellm/integrations/focus/serializers/__init__.py +++ b/litellm/integrations/focus/serializers/__init__.py @@ -1,6 +1,7 @@ """Serializer package exports for Focus integration.""" from .base import FocusSerializer +from .csv import FocusCsvSerializer from .parquet import FocusParquetSerializer -__all__ = ["FocusSerializer", "FocusParquetSerializer"] +__all__ = ["FocusSerializer", "FocusCsvSerializer", "FocusParquetSerializer"] diff --git a/litellm/integrations/focus/serializers/csv.py b/litellm/integrations/focus/serializers/csv.py new file mode 100644 index 0000000000..30e7f3283d --- /dev/null +++ b/litellm/integrations/focus/serializers/csv.py @@ -0,0 +1,22 @@ +"""CSV serializer for Focus export.""" + +from __future__ import annotations + +import io + +import polars as pl + +from .base import FocusSerializer + + +class FocusCsvSerializer(FocusSerializer): + """Serialize normalized Focus frames to CSV bytes.""" + + extension = "csv" + + def serialize(self, frame: pl.DataFrame) -> bytes: + """Encode the provided frame as a CSV payload.""" + target = frame if not frame.is_empty() else pl.DataFrame(schema=frame.schema) + buffer = io.BytesIO() + target.write_csv(buffer) + return buffer.getvalue() diff --git a/litellm/integrations/focus/transformer.py b/litellm/integrations/focus/transformer.py index cac12b7be1..1e4a17796d 100644 --- a/litellm/integrations/focus/transformer.py +++ b/litellm/integrations/focus/transformer.py @@ -2,6 +2,7 @@ from __future__ import annotations +import json from datetime import timedelta import polars as pl @@ -9,6 +10,29 @@ import polars as pl from .schema import FOCUS_NORMALIZED_SCHEMA +def _build_tags_json(row: dict) -> str: + """Build a JSON string of metadata tags from a DB row. + + Vantage uses this for Token Allocation — enriching billing data with + team, user, and API key metadata. + """ + tags: dict[str, str] = {} + for key in ( + "team_id", + "team_alias", + "user_id", + "user_email", + "api_key_alias", + "model", + "model_group", + "custom_llm_provider", + ): + val = row.get(key) + if val is not None: + tags[key] = str(val) + return json.dumps(tags) if tags else "{}" + + class FocusTransformer: """Transforms LiteLLM DB rows into Focus-compatible schema.""" @@ -19,6 +43,20 @@ class FocusTransformer: if frame.is_empty(): return pl.DataFrame(schema=self.schema) + # Build Tags JSON from metadata columns + tag_col = "Tags" + tag_keys = [ + "team_id", "team_alias", "user_id", "user_email", + "api_key_alias", "model", "model_group", "custom_llm_provider", + ] + available_keys = [k for k in tag_keys if k in frame.columns] + if available_keys: + tags_series = frame.select(available_keys).to_dicts() + tags_json = [_build_tags_json(row) for row in tags_series] + frame = frame.with_columns(pl.Series(tag_col, tags_json)) + else: + frame = frame.with_columns(pl.lit("{}").alias(tag_col)) + # derive period start/end from usage date frame = frame.with_columns( pl.col("date") @@ -86,5 +124,5 @@ class FocusTransformer: pl.col("team_id").cast(pl.String).alias("SubAccountId"), pl.col("team_alias").cast(pl.String).alias("SubAccountName"), none_str.alias("SubAccountType"), - none_str.alias("Tags"), + pl.col(tag_col).cast(pl.String).alias("Tags"), ) diff --git a/litellm/integrations/vantage/__init__.py b/litellm/integrations/vantage/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/litellm/integrations/vantage/vantage_logger.py b/litellm/integrations/vantage/vantage_logger.py new file mode 100644 index 0000000000..01529d537f --- /dev/null +++ b/litellm/integrations/vantage/vantage_logger.py @@ -0,0 +1,105 @@ +"""Vantage logger — thin wrapper around the Focus export pipeline. + +Configures FocusLogger to use the Vantage API destination with CSV format +so users can simply set ``success_callback: ["vantage"]`` in their proxy config. +""" + +from __future__ import annotations + +import os +from typing import TYPE_CHECKING, Any, Dict, List, Optional, cast + +import litellm +from litellm._logging import verbose_logger +from litellm.integrations.custom_logger import CustomLogger +from litellm.integrations.focus.focus_logger import FocusLogger + +if TYPE_CHECKING: + from apscheduler.schedulers.asyncio import AsyncIOScheduler +else: + AsyncIOScheduler = Any + +VANTAGE_USAGE_DATA_JOB_NAME = "vantage_export_usage_data" + + +class VantageLogger(FocusLogger): + """FocusLogger pre-configured for Vantage (CSV format, Vantage API destination). + + Environment Variables: + VANTAGE_API_KEY: Vantage API key for authentication + VANTAGE_INTEGRATION_TOKEN: Vantage integration token for the cost-import endpoint + VANTAGE_BASE_URL: Optional base URL override (default: https://api.vantage.sh) + VANTAGE_EXPORT_FREQUENCY: Export frequency — "hourly" (default), "daily", or "interval" + VANTAGE_EXPORT_INTERVAL_SECONDS: Interval in seconds when frequency is "interval" + """ + + def __init__( + self, + *, + api_key: Optional[str] = None, + integration_token: Optional[str] = None, + base_url: Optional[str] = None, + frequency: Optional[str] = None, + interval_seconds: Optional[int] = None, + **kwargs: Any, + ) -> None: + resolved_api_key = api_key or os.getenv("VANTAGE_API_KEY") + resolved_token = integration_token or os.getenv("VANTAGE_INTEGRATION_TOKEN") + resolved_base_url = base_url or os.getenv( + "VANTAGE_BASE_URL", "https://api.vantage.sh" + ) + resolved_frequency = ( + frequency or os.getenv("VANTAGE_EXPORT_FREQUENCY") or "hourly" + ).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 + + destination_config: Dict[str, Any] = {} + if resolved_api_key: + destination_config["api_key"] = resolved_api_key + if resolved_token: + destination_config["integration_token"] = resolved_token + if resolved_base_url: + destination_config["base_url"] = resolved_base_url + + super().__init__( + provider="vantage", + export_format="csv", + frequency=resolved_frequency, + interval_seconds=resolved_interval, + prefix="vantage_exports", + destination_config=destination_config, + **kwargs, + ) + + verbose_logger.debug( + "VantageLogger initialized (integration_token=%s)", + resolved_token[:8] + "..." if resolved_token else "None", + ) + + @staticmethod + async def init_vantage_background_job( + scheduler: AsyncIOScheduler, + ) -> None: + """Register the Vantage export job with the provided scheduler.""" + vantage_loggers: List[ + CustomLogger + ] = litellm.logging_callback_manager.get_custom_loggers_for_type( + callback_type=VantageLogger + ) + if not vantage_loggers: + verbose_logger.debug( + "No Vantage logger registered; skipping scheduler" + ) + return + + vantage_logger = cast(VantageLogger, vantage_loggers[0]) + trigger_kwargs = vantage_logger._build_scheduler_trigger() + scheduler.add_job( + vantage_logger.initialize_focus_export_job, + **trigger_kwargs, + ) + + +__all__ = ["VantageLogger"] diff --git a/litellm/litellm_core_utils/custom_logger_registry.py b/litellm/litellm_core_utils/custom_logger_registry.py index 2d483f7861..f873bfeece 100644 --- a/litellm/litellm_core_utils/custom_logger_registry.py +++ b/litellm/litellm_core_utils/custom_logger_registry.py @@ -24,6 +24,7 @@ from litellm.integrations.datadog.datadog_metrics import DatadogMetricsLogger from litellm.integrations.deepeval import DeepEvalLogger from litellm.integrations.dotprompt import DotpromptManager from litellm.integrations.focus.focus_logger import FocusLogger +from litellm.integrations.vantage.vantage_logger import VantageLogger from litellm.integrations.galileo import GalileoObserve from litellm.integrations.gcs_bucket.gcs_bucket import GCSBucketLogger from litellm.integrations.gcs_pubsub.pub_sub import GcsPubSubLogger @@ -99,6 +100,7 @@ class CustomLoggerRegistry: "gitlab": GitLabPromptManager, "cloudzero": CloudZeroLogger, "focus": FocusLogger, + "vantage": VantageLogger, "posthog": PostHogLogger, } diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index f3bc4b0803..91196ce5b9 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -473,6 +473,7 @@ from litellm.proxy.search_endpoints.search_tool_management import ( router as search_tool_management_router, ) from litellm.proxy.spend_tracking.cloudzero_endpoints import router as cloudzero_router +from litellm.proxy.spend_tracking.vantage_endpoints import router as vantage_router from litellm.proxy.spend_tracking.spend_management_endpoints import ( router as spend_management_router, ) @@ -6127,6 +6128,15 @@ class ProxyStartupEvent: ######################################################## await FocusLogger.init_focus_export_background_job(scheduler=scheduler) + ######################################################## + # Vantage Background Job + ######################################################## + from litellm.integrations.vantage.vantage_logger import VantageLogger + from litellm.proxy.spend_tracking.vantage_endpoints import is_vantage_setup + + if await is_vantage_setup(): + await VantageLogger.init_vantage_background_job(scheduler=scheduler) + ######################################################## # Prometheus Background Job ######################################################## @@ -13180,6 +13190,7 @@ app.include_router(project_router) app.include_router(customer_router) app.include_router(spend_management_router) app.include_router(cloudzero_router) +app.include_router(vantage_router) app.include_router(caching_router) app.include_router(analytics_router) app.include_router(guardrails_router) diff --git a/litellm/proxy/spend_tracking/vantage_endpoints.py b/litellm/proxy/spend_tracking/vantage_endpoints.py new file mode 100644 index 0000000000..842ca9a258 --- /dev/null +++ b/litellm/proxy/spend_tracking/vantage_endpoints.py @@ -0,0 +1,488 @@ +import json + +from fastapi import APIRouter, Depends, HTTPException + +from litellm._logging import verbose_proxy_logger +from litellm.litellm_core_utils.sensitive_data_masker import SensitiveDataMasker +from litellm.proxy._types import CommonProxyErrors, LitellmUserRoles, UserAPIKeyAuth +from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.common_utils.encrypt_decrypt_utils import ( + decrypt_value_helper, + encrypt_value_helper, +) +from litellm.types.proxy.vantage_endpoints import ( + VantageExportRequest, + VantageExportResponse, + VantageInitRequest, + VantageInitResponse, + VantageSettingsUpdate, + VantageSettingsView, +) + +router = APIRouter() + +_sensitive_masker = SensitiveDataMasker() + +VANTAGE_SETTINGS_PARAM_NAME = "vantage_settings" + + +async def _set_vantage_settings( + api_key: str, integration_token: str, base_url: str +): + """Store Vantage settings in the database with encrypted API key.""" + from litellm.proxy.proxy_server import prisma_client + + if prisma_client is None: + raise HTTPException( + status_code=500, + detail={"error": CommonProxyErrors.db_not_connected_error.value}, + ) + + encrypted_api_key = encrypt_value_helper(api_key) + + vantage_settings = { + "api_key": encrypted_api_key, + "integration_token": integration_token, + "base_url": base_url, + } + + await prisma_client.db.litellm_config.upsert( + where={"param_name": VANTAGE_SETTINGS_PARAM_NAME}, + data={ + "create": { + "param_name": VANTAGE_SETTINGS_PARAM_NAME, + "param_value": json.dumps(vantage_settings), + }, + "update": {"param_value": json.dumps(vantage_settings)}, + }, + ) + + +async def _get_vantage_settings(): + """Retrieve Vantage settings from the database with decrypted API key.""" + from litellm.proxy.proxy_server import prisma_client + + if prisma_client is None: + raise HTTPException( + status_code=500, + detail={"error": CommonProxyErrors.db_not_connected_error.value}, + ) + + vantage_config = await prisma_client.db.litellm_config.find_first( + where={"param_name": VANTAGE_SETTINGS_PARAM_NAME} + ) + if vantage_config is None or vantage_config.param_value is None: + return {} + + if isinstance(vantage_config.param_value, dict): + settings = vantage_config.param_value + elif isinstance(vantage_config.param_value, str): + settings = json.loads(vantage_config.param_value) + else: + settings = dict(vantage_config.param_value) + + encrypted_api_key = settings.get("api_key") + if encrypted_api_key: + decrypted_api_key = decrypt_value_helper( + encrypted_api_key, key="vantage_api_key", exception_type="error" + ) + if decrypted_api_key is None: + raise HTTPException( + status_code=500, + detail={ + "error": "Failed to decrypt Vantage API key. Check your salt key configuration." + }, + ) + settings["api_key"] = decrypted_api_key + + return settings + + +@router.get( + "/vantage/settings", + tags=["Vantage"], + dependencies=[Depends(user_api_key_auth)], + response_model=VantageSettingsView, +) +async def get_vantage_settings( + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +): + """ + View current Vantage settings. + + Returns the current Vantage configuration with the API key masked for security. + Only admin users can view Vantage settings. + """ + if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN: + raise HTTPException( + status_code=403, + detail={"error": CommonProxyErrors.not_allowed_access.value}, + ) + + try: + settings = await _get_vantage_settings() + + if not settings: + return VantageSettingsView( + api_key_masked=None, + integration_token=None, + base_url=None, + status=None, + ) + + masked_settings = _sensitive_masker.mask_dict(settings) + + return VantageSettingsView( + api_key_masked=masked_settings.get("api_key"), + integration_token=settings.get("integration_token"), + base_url=settings.get("base_url"), + status="configured", + ) + + except HTTPException: + raise + except Exception as e: + verbose_proxy_logger.error(f"Error retrieving Vantage settings: {str(e)}") + raise HTTPException( + status_code=500, + detail={"error": f"Failed to retrieve Vantage settings: {str(e)}"}, + ) + + +@router.put( + "/vantage/settings", + tags=["Vantage"], + dependencies=[Depends(user_api_key_auth)], + response_model=VantageInitResponse, +) +async def update_vantage_settings( + request: VantageSettingsUpdate, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +): + """ + Update existing Vantage settings. + + Allows updating individual Vantage configuration fields without requiring all fields. + Only admin users can update Vantage settings. + """ + if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN: + raise HTTPException( + status_code=403, + detail={"error": CommonProxyErrors.not_allowed_access.value}, + ) + + if not any([request.api_key, request.integration_token, request.base_url]): + raise HTTPException( + status_code=400, + detail={"error": "At least one field must be provided for update"}, + ) + + try: + current_settings = await _get_vantage_settings() + + updated_api_key = ( + request.api_key + if request.api_key is not None + else current_settings["api_key"] + ) + updated_token = ( + request.integration_token + if request.integration_token is not None + else current_settings["integration_token"] + ) + updated_base_url = ( + request.base_url + if request.base_url is not None + else current_settings["base_url"] + ) + + await _set_vantage_settings( + api_key=updated_api_key, + integration_token=updated_token, + base_url=updated_base_url, + ) + + verbose_proxy_logger.info("Vantage settings updated successfully") + + return VantageInitResponse( + message="Vantage settings updated successfully", status="success" + ) + + except HTTPException as e: + if e.status_code == 400: + raise HTTPException( + status_code=404, + detail={ + "error": "Vantage settings not found. Please initialize settings first using /vantage/init" + }, + ) + raise + except Exception as e: + verbose_proxy_logger.error(f"Error updating Vantage settings: {str(e)}") + raise HTTPException( + status_code=500, + detail={"error": f"Failed to update Vantage settings: {str(e)}"}, + ) + + +async def is_vantage_setup_in_db() -> bool: + """Check if Vantage is setup in the database.""" + try: + from litellm.proxy.proxy_server import prisma_client + + if prisma_client is None: + return False + + vantage_config = await prisma_client.db.litellm_config.find_first( + where={"param_name": VANTAGE_SETTINGS_PARAM_NAME} + ) + + return vantage_config is not None and vantage_config.param_value is not None + + except Exception as e: + verbose_proxy_logger.error(f"Error checking Vantage status: {str(e)}") + return False + + +def is_vantage_setup_in_config() -> bool: + """Check if Vantage is setup in config.yaml or environment variables.""" + import litellm + + return "vantage" in litellm.callbacks + + +async def is_vantage_setup() -> bool: + """Check if Vantage is setup in either config or database.""" + try: + if is_vantage_setup_in_config(): + return True + if await is_vantage_setup_in_db(): + return True + return False + except Exception as e: + verbose_proxy_logger.error(f"Error checking Vantage setup: {str(e)}") + return False + + +@router.post( + "/vantage/init", + tags=["Vantage"], + dependencies=[Depends(user_api_key_auth)], + response_model=VantageInitResponse, +) +async def init_vantage_settings( + request: VantageInitRequest, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +): + """ + Initialize Vantage settings and store in the database. + + Parameters: + - api_key: Vantage API key for authentication + - integration_token: Vantage integration token for the cost-import endpoint + - base_url: Vantage API base URL (default: https://api.vantage.sh) + + Only admin users can configure Vantage settings. + """ + if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN: + raise HTTPException( + status_code=403, + detail={"error": CommonProxyErrors.not_allowed_access.value}, + ) + + try: + await _set_vantage_settings( + api_key=request.api_key, + integration_token=request.integration_token, + base_url=request.base_url, + ) + + verbose_proxy_logger.info("Vantage settings initialized successfully") + + return VantageInitResponse( + message="Vantage settings initialized successfully", status="success" + ) + + except Exception as e: + verbose_proxy_logger.error( + f"Error initializing Vantage settings: {str(e)}" + ) + raise HTTPException( + status_code=500, + detail={"error": f"Failed to initialize Vantage settings: {str(e)}"}, + ) + + +@router.post( + "/vantage/dry-run", + tags=["Vantage"], + dependencies=[Depends(user_api_key_auth)], + response_model=VantageExportResponse, +) +async def vantage_dry_run_export( + request: VantageExportRequest, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +): + """ + Perform a dry run export using the Vantage logger. + + 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) + + Only admin users can perform Vantage exports. + """ + if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN: + raise HTTPException( + status_code=403, + detail={"error": CommonProxyErrors.not_allowed_access.value}, + ) + + try: + from litellm.integrations.vantage.vantage_logger import VantageLogger + + logger = VantageLogger() + dry_run_result = await logger.dry_run_export_usage_data( + limit=request.limit + ) + + verbose_proxy_logger.info("Vantage dry run export completed successfully") + + return VantageExportResponse( + message="Vantage dry run export completed successfully.", + status="success", + dry_run_data=dry_run_result, + summary=dry_run_result.get("summary") if dry_run_result else None, + ) + + except Exception as e: + verbose_proxy_logger.error( + f"Error performing Vantage dry run export: {str(e)}" + ) + raise HTTPException( + status_code=500, + detail={ + "error": f"Failed to perform Vantage dry run export: {str(e)}" + }, + ) + + +@router.post( + "/vantage/export", + tags=["Vantage"], + dependencies=[Depends(user_api_key_auth)], + response_model=VantageExportResponse, +) +async def vantage_export( + request: VantageExportRequest, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +): + """ + Perform an actual export using the Vantage logger. + + Exports usage data in FOCUS CSV format to the Vantage API. + + Parameters: + - limit: Optional limit on number of records to export + - start_time_utc: Optional start time for data export + - end_time_utc: Optional end time for data export + + Only admin users can perform Vantage exports. + """ + if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN: + raise HTTPException( + status_code=403, + detail={"error": CommonProxyErrors.not_allowed_access.value}, + ) + + try: + settings = await _get_vantage_settings() + + from litellm.integrations.vantage.vantage_logger import VantageLogger + + logger = VantageLogger( + api_key=settings.get("api_key"), + integration_token=settings.get("integration_token"), + base_url=settings.get("base_url"), + ) + await logger.export_usage_data( + limit=request.limit, + start_time_utc=request.start_time_utc, + end_time_utc=request.end_time_utc, + ) + + verbose_proxy_logger.info("Vantage export completed successfully") + + return VantageExportResponse( + message="Vantage export completed successfully", + status="success", + dry_run_data=None, + summary=None, + ) + + except Exception as e: + verbose_proxy_logger.error(f"Error performing Vantage export: {str(e)}") + raise HTTPException( + status_code=500, + detail={"error": f"Failed to perform Vantage export: {str(e)}"}, + ) + + +@router.delete( + "/vantage/delete", + tags=["Vantage"], + dependencies=[Depends(user_api_key_auth)], + response_model=VantageInitResponse, +) +async def delete_vantage_settings( + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +): + """ + Delete Vantage settings from the database. + + Only admin users can delete Vantage settings. + """ + if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN: + raise HTTPException( + status_code=403, + detail={"error": CommonProxyErrors.not_allowed_access.value}, + ) + + try: + from litellm.proxy.proxy_server import prisma_client + + if prisma_client is None: + raise HTTPException( + status_code=500, + detail={"error": CommonProxyErrors.db_not_connected_error.value}, + ) + + vantage_config = await prisma_client.db.litellm_config.find_first( + where={"param_name": VANTAGE_SETTINGS_PARAM_NAME} + ) + + if vantage_config is None: + raise HTTPException( + status_code=404, + detail={"error": "Vantage settings not found"}, + ) + + await prisma_client.db.litellm_config.delete( + where={"param_name": VANTAGE_SETTINGS_PARAM_NAME} + ) + + verbose_proxy_logger.info("Vantage settings deleted successfully") + + return VantageInitResponse( + message="Vantage settings deleted successfully", status="success" + ) + + except HTTPException: + raise + except Exception as e: + verbose_proxy_logger.error(f"Error deleting Vantage settings: {str(e)}") + raise HTTPException( + status_code=500, + detail={"error": f"Failed to delete Vantage settings: {str(e)}"}, + ) diff --git a/litellm/types/proxy/vantage_endpoints.py b/litellm/types/proxy/vantage_endpoints.py new file mode 100644 index 0000000000..394677d1e8 --- /dev/null +++ b/litellm/types/proxy/vantage_endpoints.py @@ -0,0 +1,81 @@ +""" +Vantage endpoint types for LiteLLM Proxy +""" + +from datetime import datetime +from typing import Any, Dict, Optional + +from pydantic import BaseModel, Field + + +class VantageInitRequest(BaseModel): + """Request model for initializing Vantage settings""" + + api_key: str = Field(..., description="Vantage API key for authentication") + integration_token: str = Field( + ..., description="Vantage integration token for the cost-import endpoint" + ) + base_url: str = Field( + default="https://api.vantage.sh", + description="Vantage API base URL (default: https://api.vantage.sh)", + ) + + +class VantageInitResponse(BaseModel): + """Response model for Vantage initialization""" + + message: str + status: str + + +class VantageExportRequest(BaseModel): + """Request model for Vantage export operations""" + + limit: Optional[int] = Field( + None, description="Optional limit on number of records to export" + ) + start_time_utc: Optional[datetime] = Field( + None, description="Start time for data export in UTC" + ) + end_time_utc: Optional[datetime] = Field( + None, description="End time for data export in UTC" + ) + + +class VantageExportResponse(BaseModel): + """Response model for Vantage export operations""" + + message: str + status: str + dry_run_data: Optional[Dict[str, Any]] = Field( + None, description="Dry run data including usage data and FOCUS transformed data" + ) + summary: Optional[Dict[str, Any]] = Field( + None, description="Summary statistics for dry run" + ) + + +class VantageSettingsView(BaseModel): + """Response model for viewing Vantage settings with masked API key""" + + api_key_masked: Optional[str] = Field( + None, + description="Masked API key showing only first 4 and last 4 characters", + ) + integration_token: Optional[str] = Field( + None, description="Vantage integration token" + ) + base_url: Optional[str] = Field(None, description="Vantage API base URL") + status: Optional[str] = Field(None, description="Configuration status") + + +class VantageSettingsUpdate(BaseModel): + """Request model for updating Vantage settings""" + + api_key: Optional[str] = Field( + None, description="New Vantage API key for authentication" + ) + integration_token: Optional[str] = Field( + None, description="New Vantage integration token" + ) + base_url: Optional[str] = Field(None, description="New Vantage API base URL") diff --git a/tests/test_litellm/integrations/focus/test_csv_serializer.py b/tests/test_litellm/integrations/focus/test_csv_serializer.py new file mode 100644 index 0000000000..f527a9dceb --- /dev/null +++ b/tests/test_litellm/integrations/focus/test_csv_serializer.py @@ -0,0 +1,34 @@ +"""Tests for FocusCsvSerializer.""" + +from __future__ import annotations + +import polars as pl + +from litellm.integrations.focus.serializers.csv import FocusCsvSerializer + + +def test_should_serialize_dataframe_to_csv(): + frame = pl.DataFrame({"BilledCost": [1.5, 2.0], "ServiceName": ["openai", "anthropic"]}) + serializer = FocusCsvSerializer() + result = serializer.serialize(frame) + + assert isinstance(result, bytes) + lines = result.decode("utf-8").strip().split("\n") + assert lines[0] == "BilledCost,ServiceName" + assert len(lines) == 3 # header + 2 data rows + + +def test_should_return_header_only_for_empty_frame(): + frame = pl.DataFrame( + schema={"BilledCost": pl.Float64, "ServiceName": pl.Utf8} + ) + serializer = FocusCsvSerializer() + result = serializer.serialize(frame) + + lines = result.decode("utf-8").strip().split("\n") + assert lines[0] == "BilledCost,ServiceName" + assert len(lines) == 1 # header only + + +def test_extension_should_be_csv(): + assert FocusCsvSerializer.extension == "csv" diff --git a/tests/test_litellm/integrations/focus/test_destination_factory.py b/tests/test_litellm/integrations/focus/test_destination_factory.py new file mode 100644 index 0000000000..4888e8231b --- /dev/null +++ b/tests/test_litellm/integrations/focus/test_destination_factory.py @@ -0,0 +1,48 @@ +"""Tests for FocusDestinationFactory with vantage provider.""" + +from __future__ import annotations + +import pytest + +from litellm.integrations.focus.destinations.factory import FocusDestinationFactory +from litellm.integrations.focus.destinations.vantage_destination import ( + FocusVantageDestination, +) + + +def test_should_create_vantage_destination(): + dest = FocusDestinationFactory.create( + provider="vantage", + prefix="exports", + config={ + "api_key": "test-key", + "integration_token": "test-token", + }, + ) + assert isinstance(dest, FocusVantageDestination) + + +def test_should_raise_when_vantage_missing_api_key(): + with pytest.raises(ValueError, match="VANTAGE_API_KEY"): + FocusDestinationFactory.create( + provider="vantage", + prefix="exports", + config={"integration_token": "tok"}, + ) + + +def test_should_raise_when_vantage_missing_token(): + with pytest.raises(ValueError, match="VANTAGE_INTEGRATION_TOKEN"): + FocusDestinationFactory.create( + provider="vantage", + prefix="exports", + config={"api_key": "key"}, + ) + + +def test_should_raise_for_unsupported_provider(): + with pytest.raises(NotImplementedError): + FocusDestinationFactory.create( + provider="unknown_provider", + prefix="exports", + ) diff --git a/tests/test_litellm/integrations/focus/test_vantage_destination.py b/tests/test_litellm/integrations/focus/test_vantage_destination.py new file mode 100644 index 0000000000..181d3d1381 --- /dev/null +++ b/tests/test_litellm/integrations/focus/test_vantage_destination.py @@ -0,0 +1,138 @@ +"""Tests for FocusVantageDestination behavior.""" + +from __future__ import annotations + +from datetime import datetime, timezone +from typing import Any, Dict, List +from unittest.mock import AsyncMock, patch + +import pytest + +from litellm.integrations.focus.destinations.base import FocusTimeWindow +from litellm.integrations.focus.destinations.vantage_destination import ( + FocusVantageDestination, + VANTAGE_MAX_BYTES_PER_UPLOAD, +) + + +def _window(freq: str = "hourly", hour: int = 5) -> FocusTimeWindow: + start = datetime(2024, 1, 2, hour, tzinfo=timezone.utc) + end = start.replace(hour=hour + 1) + return FocusTimeWindow(start_time=start, end_time=end, frequency=freq) + + +def _config(**overrides: Any) -> dict[str, Any]: + base = { + "api_key": "test-api-key", + "integration_token": "test-token-123", + } + base.update(overrides) + return base + + +def test_should_require_api_key(): + with pytest.raises(ValueError, match="api_key"): + FocusVantageDestination( + prefix="exports", + config={"integration_token": "tok"}, + ) + + +def test_should_require_integration_token(): + with pytest.raises(ValueError, match="integration_token"): + FocusVantageDestination( + prefix="exports", + config={"api_key": "key"}, + ) + + +def test_should_initialize_with_valid_config(): + dest = FocusVantageDestination(prefix="exports", config=_config()) + assert dest.api_key == "test-api-key" + assert dest.integration_token == "test-token-123" + assert dest.base_url == "https://api.vantage.sh" + + +def test_should_use_custom_base_url(): + dest = FocusVantageDestination( + prefix="exports", + config=_config(base_url="https://custom.vantage.sh"), + ) + assert dest.base_url == "https://custom.vantage.sh" + + +@pytest.mark.asyncio +async def test_should_skip_empty_content(): + dest = FocusVantageDestination(prefix="exports", config=_config()) + # Should not raise + await dest.deliver(content=b"", time_window=_window(), filename="usage.csv") + + +@pytest.mark.asyncio +async def test_should_upload_csv_to_correct_url(): + dest = FocusVantageDestination(prefix="exports", config=_config()) + captured: Dict[str, Any] = {} + + mock_response = AsyncMock() + mock_response.raise_for_status = lambda: None + + mock_client = AsyncMock() + mock_client.post = AsyncMock(return_value=mock_response) + mock_client.__aenter__ = AsyncMock(return_value=mock_client) + mock_client.__aexit__ = AsyncMock(return_value=False) + + with patch("litellm.integrations.focus.destinations.vantage_destination.httpx.AsyncClient", return_value=mock_client): + await dest.deliver( + content=b"header\nrow1\n", + time_window=_window(), + filename="usage.csv", + ) + + mock_client.post.assert_called_once() + call_args = mock_client.post.call_args + assert "test-token-123" in call_args[0][0] + assert "costs.csv" in call_args[0][0] + assert call_args[1]["headers"]["Authorization"] == "Bearer test-api-key" + + +@pytest.mark.asyncio +async def test_should_batch_large_content(): + dest = FocusVantageDestination(prefix="exports", config=_config()) + + # Create content larger than 2 MB + header = b"col1,col2,col3" + row = b"a" * 100 + b"," + b"b" * 100 + b"," + b"c" * 100 + num_rows = (VANTAGE_MAX_BYTES_PER_UPLOAD // len(row)) + 100 + large_content = header + b"\n" + b"\n".join([row] * num_rows) + b"\n" + + assert len(large_content) > VANTAGE_MAX_BYTES_PER_UPLOAD + + upload_calls: List[bytes] = [] + + mock_response = AsyncMock() + mock_response.raise_for_status = lambda: None + + mock_client = AsyncMock() + mock_client.__aenter__ = AsyncMock(return_value=mock_client) + mock_client.__aexit__ = AsyncMock(return_value=False) + + async def capture_post(url, **kwargs): + files = kwargs.get("files", {}) + if "file" in files: + upload_calls.append(files["file"][1]) + return mock_response + + mock_client.post = capture_post + + with patch("litellm.integrations.focus.destinations.vantage_destination.httpx.AsyncClient", return_value=mock_client): + await dest.deliver( + content=large_content, + time_window=_window(), + filename="usage.csv", + ) + + # Should have made multiple uploads + assert len(upload_calls) > 1 + # Each upload should be within limits + for chunk in upload_calls: + assert len(chunk) <= VANTAGE_MAX_BYTES_PER_UPLOAD From f683befeed0a01695b3ab4250724e8626433d516 Mon Sep 17 00:00:00 2001 From: Harshit28j Date: Wed, 11 Mar 2026 14:02:05 +0530 Subject: [PATCH 02/21] Fix Greptile review issues in Vantage integration - Enforce 10K row limit in single-shot upload path (not just 2MB size) - Fix KeyError crash in update_vantage_settings when no settings exist - Remove unreachable status_code==400 dead code branch - Make dry-run endpoint work without Vantage credentials by using FOCUS database + transformer directly instead of VantageLogger Co-Authored-By: Claude Opus 4.6 --- .../focus/destinations/vantage_destination.py | 12 +++- .../proxy/spend_tracking/vantage_endpoints.py | 55 +++++++++++++------ 2 files changed, 47 insertions(+), 20 deletions(-) diff --git a/litellm/integrations/focus/destinations/vantage_destination.py b/litellm/integrations/focus/destinations/vantage_destination.py index 5612a513d6..32ba6df8ea 100644 --- a/litellm/integrations/focus/destinations/vantage_destination.py +++ b/litellm/integrations/focus/destinations/vantage_destination.py @@ -56,12 +56,18 @@ class FocusVantageDestination(FocusDestination): verbose_logger.debug("Vantage destination: empty content, skipping upload") return - # If the payload is within limits, send in one shot - if len(content) <= VANTAGE_MAX_BYTES_PER_UPLOAD: + # 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 - # Otherwise split into chunks by line count + # Otherwise split into batches respecting both limits await self._upload_batched(content, filename) async def _upload_csv(self, csv_bytes: bytes, filename: str) -> None: diff --git a/litellm/proxy/spend_tracking/vantage_endpoints.py b/litellm/proxy/spend_tracking/vantage_endpoints.py index 842ca9a258..47838158b9 100644 --- a/litellm/proxy/spend_tracking/vantage_endpoints.py +++ b/litellm/proxy/spend_tracking/vantage_endpoints.py @@ -180,20 +180,28 @@ async def update_vantage_settings( try: current_settings = await _get_vantage_settings() + if not current_settings: + raise HTTPException( + status_code=404, + detail={ + "error": "Vantage settings not found. Please initialize settings first using /vantage/init" + }, + ) + updated_api_key = ( request.api_key if request.api_key is not None - else current_settings["api_key"] + else current_settings.get("api_key", "") ) updated_token = ( request.integration_token if request.integration_token is not None - else current_settings["integration_token"] + else current_settings.get("integration_token", "") ) updated_base_url = ( request.base_url if request.base_url is not None - else current_settings["base_url"] + else current_settings.get("base_url", "https://api.vantage.sh") ) await _set_vantage_settings( @@ -208,14 +216,7 @@ async def update_vantage_settings( message="Vantage settings updated successfully", status="success" ) - except HTTPException as e: - if e.status_code == 400: - raise HTTPException( - status_code=404, - detail={ - "error": "Vantage settings not found. Please initialize settings first using /vantage/init" - }, - ) + except HTTPException: raise except Exception as e: verbose_proxy_logger.error(f"Error updating Vantage settings: {str(e)}") @@ -340,12 +341,32 @@ async def vantage_dry_run_export( ) try: - from litellm.integrations.vantage.vantage_logger import VantageLogger + # Dry-run uses the FOCUS database + transformer directly, + # bypassing the destination so no Vantage credentials are required. + from litellm.integrations.focus.database import FocusLiteLLMDatabase + from litellm.integrations.focus.transformer import FocusTransformer - logger = VantageLogger() - dry_run_result = await logger.dry_run_export_usage_data( - limit=request.limit - ) + database = FocusLiteLLMDatabase() + transformer = FocusTransformer() + + data = await database.get_usage_data(limit=request.limit) + normalized = transformer.transform(data) + + 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 [] + + summary = { + "total_records": len(normalized), + "total_spend": float(normalized.select("BilledCost").sum().item()) if not normalized.is_empty() and "BilledCost" in normalized.columns else 0.0, + "unique_teams": normalized["SubAccountId"].n_unique() if not normalized.is_empty() and "SubAccountId" in normalized.columns else 0, + "unique_models": normalized["ResourceType"].n_unique() if not normalized.is_empty() and "ResourceType" in normalized.columns else 0, + } + + dry_run_result = { + "usage_data": usage_sample, + "normalized_data": normalized_sample, + "summary": summary, + } verbose_proxy_logger.info("Vantage dry run export completed successfully") @@ -353,7 +374,7 @@ async def vantage_dry_run_export( message="Vantage dry run export completed successfully.", status="success", dry_run_data=dry_run_result, - summary=dry_run_result.get("summary") if dry_run_result else None, + summary=summary, ) except Exception as e: From e1c44fe0881ed714d91942d9751c9505e8eaccc4 Mon Sep 17 00:00:00 2001 From: Harshit28j Date: Wed, 11 Mar 2026 14:17:03 +0530 Subject: [PATCH 03/21] Fix Greptile round 2 review issues - DB-only Vantage config now registers VantageLogger at startup so background job actually schedules (was silently skipping) - Add missing `except HTTPException: raise` in /vantage/init endpoint - Use consistent batch filenames (always .partN suffix) - Document Tags schema change from pl.Object to pl.String Co-Authored-By: Claude Opus 4.6 --- .../focus/destinations/vantage_destination.py | 2 +- litellm/integrations/focus/schema.py | 3 +++ litellm/proxy/proxy_server.py | 24 ++++++++++++++++++- .../proxy/spend_tracking/vantage_endpoints.py | 2 ++ 4 files changed, 29 insertions(+), 2 deletions(-) diff --git a/litellm/integrations/focus/destinations/vantage_destination.py b/litellm/integrations/focus/destinations/vantage_destination.py index 32ba6df8ea..e8d9532c41 100644 --- a/litellm/integrations/focus/destinations/vantage_destination.py +++ b/litellm/integrations/focus/destinations/vantage_destination.py @@ -108,7 +108,7 @@ class FocusVantageDestination(FocusDestination): if len(batch_csv) > VANTAGE_MAX_BYTES_PER_UPLOAD: await self._upload_size_limited(header, batch_lines, filename, batch_num) else: - batch_filename = f"{filename}.part{batch_num}" if batch_num > 0 else filename + batch_filename = f"{filename}.part{batch_num}" await self._upload_csv(batch_csv, batch_filename) batch_num += 1 diff --git a/litellm/integrations/focus/schema.py b/litellm/integrations/focus/schema.py index 6ad725f367..127e8ee03a 100644 --- a/litellm/integrations/focus/schema.py +++ b/litellm/integrations/focus/schema.py @@ -43,6 +43,9 @@ FOCUS_NORMALIZED_SCHEMA = pl.Schema( ("SubAccountId", pl.String), ("SubAccountName", pl.String), ("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. ("Tags", pl.String), ] ) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 91196ce5b9..9bd4cd7f14 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -6132,9 +6132,31 @@ class ProxyStartupEvent: # Vantage Background Job ######################################################## from litellm.integrations.vantage.vantage_logger import VantageLogger - from litellm.proxy.spend_tracking.vantage_endpoints import is_vantage_setup + from litellm.proxy.spend_tracking.vantage_endpoints import ( + _get_vantage_settings, + is_vantage_setup, + is_vantage_setup_in_config, + is_vantage_setup_in_db, + ) if await is_vantage_setup(): + # If configured via DB but not in config.yaml callbacks, + # instantiate and register a VantageLogger so the scheduler + # can find it. + if not is_vantage_setup_in_config() and await is_vantage_setup_in_db(): + try: + db_settings = await _get_vantage_settings() + if db_settings: + vantage_logger = VantageLogger( + api_key=db_settings.get("api_key"), + integration_token=db_settings.get("integration_token"), + base_url=db_settings.get("base_url"), + ) + litellm.callbacks.append(vantage_logger) # type: ignore[arg-type] + except Exception as e: + verbose_proxy_logger.warning( + "Failed to register VantageLogger from DB settings: %s", e + ) await VantageLogger.init_vantage_background_job(scheduler=scheduler) ######################################################## diff --git a/litellm/proxy/spend_tracking/vantage_endpoints.py b/litellm/proxy/spend_tracking/vantage_endpoints.py index 47838158b9..d0e01fe325 100644 --- a/litellm/proxy/spend_tracking/vantage_endpoints.py +++ b/litellm/proxy/spend_tracking/vantage_endpoints.py @@ -304,6 +304,8 @@ async def init_vantage_settings( message="Vantage settings initialized successfully", status="success" ) + except HTTPException: + raise except Exception as e: verbose_proxy_logger.error( f"Error initializing Vantage settings: {str(e)}" From e07297fa8731676b0de83ae3443ceb2b037e216a Mon Sep 17 00:00:00 2001 From: Harshit28j Date: Wed, 11 Mar 2026 14:27:24 +0530 Subject: [PATCH 04/21] Address Greptile round 3 feedback for improved security and consistency - Encrypt integration_token alongside api_key in Vantage settings storage - Align dry-run summary with FocusExportEngine helper methods - Vectorize Tags JSON building using pl.struct + map_elements - Reuse registered VantageLogger in /export endpoint instead of creating fresh instances Co-Authored-By: Claude Opus 4.6 --- litellm/integrations/focus/transformer.py | 62 +++++++++---------- .../proxy/spend_tracking/vantage_endpoints.py | 57 +++++++++++++---- 2 files changed, 77 insertions(+), 42 deletions(-) diff --git a/litellm/integrations/focus/transformer.py b/litellm/integrations/focus/transformer.py index 1e4a17796d..adac158651 100644 --- a/litellm/integrations/focus/transformer.py +++ b/litellm/integrations/focus/transformer.py @@ -10,27 +10,34 @@ import polars as pl from .schema import FOCUS_NORMALIZED_SCHEMA -def _build_tags_json(row: dict) -> str: - """Build a JSON string of metadata tags from a DB row. +_TAG_KEYS = ( + "team_id", + "team_alias", + "user_id", + "user_email", + "api_key_alias", + "model", + "model_group", + "custom_llm_provider", +) - Vantage uses this for Token Allocation — enriching billing data with - team, user, and API key metadata. + +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. """ - tags: dict[str, str] = {} - for key in ( - "team_id", - "team_alias", - "user_id", - "user_email", - "api_key_alias", - "model", - "model_group", - "custom_llm_provider", - ): - val = row.get(key) - if val is not None: - tags[key] = str(val) - return json.dumps(tags) if tags else "{}" + + def _struct_to_json(row: dict) -> str: + tags = {k: str(v) for k, v in row.items() if v is not None} + return json.dumps(tags) if tags else "{}" + + return ( + pl.struct(available_keys) + .map_elements(_struct_to_json, return_dtype=pl.String) + .alias("Tags") + ) class FocusTransformer: @@ -43,19 +50,12 @@ class FocusTransformer: if frame.is_empty(): return pl.DataFrame(schema=self.schema) - # Build Tags JSON from metadata columns - tag_col = "Tags" - tag_keys = [ - "team_id", "team_alias", "user_id", "user_email", - "api_key_alias", "model", "model_group", "custom_llm_provider", - ] - available_keys = [k for k in tag_keys if k in frame.columns] + # Build Tags JSON from metadata columns using vectorized Polars expression + available_keys = [k for k in _TAG_KEYS if k in frame.columns] if available_keys: - tags_series = frame.select(available_keys).to_dicts() - tags_json = [_build_tags_json(row) for row in tags_series] - frame = frame.with_columns(pl.Series(tag_col, tags_json)) + frame = frame.with_columns(_build_tags_expr(available_keys)) else: - frame = frame.with_columns(pl.lit("{}").alias(tag_col)) + frame = frame.with_columns(pl.lit("{}").alias("Tags")) # derive period start/end from usage date frame = frame.with_columns( @@ -124,5 +124,5 @@ class FocusTransformer: pl.col("team_id").cast(pl.String).alias("SubAccountId"), pl.col("team_alias").cast(pl.String).alias("SubAccountName"), none_str.alias("SubAccountType"), - pl.col(tag_col).cast(pl.String).alias("Tags"), + pl.col("Tags").cast(pl.String).alias("Tags"), ) diff --git a/litellm/proxy/spend_tracking/vantage_endpoints.py b/litellm/proxy/spend_tracking/vantage_endpoints.py index d0e01fe325..a553d1061e 100644 --- a/litellm/proxy/spend_tracking/vantage_endpoints.py +++ b/litellm/proxy/spend_tracking/vantage_endpoints.py @@ -1,5 +1,6 @@ import json +import litellm from fastapi import APIRouter, Depends, HTTPException from litellm._logging import verbose_proxy_logger @@ -26,6 +27,18 @@ _sensitive_masker = SensitiveDataMasker() VANTAGE_SETTINGS_PARAM_NAME = "vantage_settings" +def _get_registered_vantage_logger(): + """Return the VantageLogger already registered in litellm.callbacks, if any.""" + from litellm.integrations.vantage.vantage_logger import VantageLogger + + vantage_loggers = litellm.logging_callback_manager.get_custom_loggers_for_type( + callback_type=VantageLogger + ) + if vantage_loggers: + return vantage_loggers[0] + return None + + async def _set_vantage_settings( api_key: str, integration_token: str, base_url: str ): @@ -39,10 +52,11 @@ async def _set_vantage_settings( ) encrypted_api_key = encrypt_value_helper(api_key) + encrypted_integration_token = encrypt_value_helper(integration_token) vantage_settings = { "api_key": encrypted_api_key, - "integration_token": integration_token, + "integration_token": encrypted_integration_token, "base_url": base_url, } @@ -95,6 +109,22 @@ async def _get_vantage_settings(): ) settings["api_key"] = decrypted_api_key + encrypted_integration_token = settings.get("integration_token") + if encrypted_integration_token: + decrypted_integration_token = decrypt_value_helper( + encrypted_integration_token, + key="vantage_integration_token", + exception_type="error", + ) + if decrypted_integration_token is None: + raise HTTPException( + status_code=500, + detail={ + "error": "Failed to decrypt Vantage integration token. Check your salt key configuration." + }, + ) + settings["integration_token"] = decrypted_integration_token + return settings @@ -346,6 +376,7 @@ async def vantage_dry_run_export( # Dry-run uses the FOCUS database + transformer directly, # bypassing the destination so no Vantage credentials are required. from litellm.integrations.focus.database import FocusLiteLLMDatabase + from litellm.integrations.focus.export_engine import FocusExportEngine from litellm.integrations.focus.transformer import FocusTransformer database = FocusLiteLLMDatabase() @@ -357,11 +388,12 @@ 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 summary = { "total_records": len(normalized), - "total_spend": float(normalized.select("BilledCost").sum().item()) if not normalized.is_empty() and "BilledCost" in normalized.columns else 0.0, - "unique_teams": normalized["SubAccountId"].n_unique() if not normalized.is_empty() and "SubAccountId" in normalized.columns else 0, - "unique_models": normalized["ResourceType"].n_unique() if not normalized.is_empty() and "ResourceType" in normalized.columns else 0, + "total_spend": FocusExportEngine._sum_column(normalized, "BilledCost"), + "unique_teams": FocusExportEngine._count_unique(normalized, "SubAccountId"), + "unique_models": FocusExportEngine._count_unique(normalized, "ResourceType"), } dry_run_result = { @@ -420,15 +452,18 @@ async def vantage_export( ) try: - settings = await _get_vantage_settings() - from litellm.integrations.vantage.vantage_logger import VantageLogger - logger = VantageLogger( - api_key=settings.get("api_key"), - integration_token=settings.get("integration_token"), - base_url=settings.get("base_url"), - ) + # Prefer the already-registered logger to avoid recreating HTTP clients + # on every export call. + logger = _get_registered_vantage_logger() + if logger is None: + settings = await _get_vantage_settings() + logger = VantageLogger( + api_key=settings.get("api_key"), + integration_token=settings.get("integration_token"), + base_url=settings.get("base_url"), + ) await logger.export_usage_data( limit=request.limit, start_time_utc=request.start_time_utc, From 4583c90194c56c4bbc299589fa56138827bb7823 Mon Sep 17 00:00:00 2001 From: Harshit28j Date: Wed, 11 Mar 2026 14:38:19 +0530 Subject: [PATCH 05/21] 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 e8d9532c41..e2a21c773f 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 a9361d0e98..0cad1b3acd 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 a553d1061e..d3d6787e5f 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 394677d1e8..165c21316b 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") From 9cc6df6b8d1a042ce28d80c1a57d3e2c0c3b3fd4 Mon Sep 17 00:00:00 2001 From: Harshit28j Date: Wed, 11 Mar 2026 14:51:03 +0530 Subject: [PATCH 06/21] 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" From 622d5dabba0dcc676247069f2891f3ef1dc9fb0a Mon Sep 17 00:00:00 2001 From: Harshit28j Date: Wed, 11 Mar 2026 14:59:52 +0530 Subject: [PATCH 07/21] Fix HTTPException swallowed as 500 in /vantage/export endpoint Add missing `except HTTPException: raise` guard so intentional 404 responses (e.g. when settings are not configured) are not caught by the generic Exception handler and re-raised as 500 errors. Co-Authored-By: Claude Opus 4.6 --- litellm/proxy/spend_tracking/vantage_endpoints.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/litellm/proxy/spend_tracking/vantage_endpoints.py b/litellm/proxy/spend_tracking/vantage_endpoints.py index d354ed2962..e497e42e69 100644 --- a/litellm/proxy/spend_tracking/vantage_endpoints.py +++ b/litellm/proxy/spend_tracking/vantage_endpoints.py @@ -487,6 +487,8 @@ async def vantage_export( summary=None, ) + except HTTPException: + raise except Exception as e: verbose_proxy_logger.error(f"Error performing Vantage export: {str(e)}") raise HTTPException( From 230c85313b9a424652041f51263bf53b0c78664b Mon Sep 17 00:00:00 2001 From: Harshit28j Date: Wed, 11 Mar 2026 15:10:53 +0530 Subject: [PATCH 08/21] Fix dry-run HTTPException guard and VantageLogger instance detection - Add missing except HTTPException: raise in vantage_dry_run_export (consistent with all other endpoints in the file) - Fix is_vantage_setup_in_config() to detect both the string "vantage" and VantageLogger instances in litellm.callbacks, preventing duplicate logger registration on startup Co-Authored-By: Claude Opus 4.6 --- litellm/proxy/spend_tracking/vantage_endpoints.py | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/spend_tracking/vantage_endpoints.py b/litellm/proxy/spend_tracking/vantage_endpoints.py index e497e42e69..f478ba5160 100644 --- a/litellm/proxy/spend_tracking/vantage_endpoints.py +++ b/litellm/proxy/spend_tracking/vantage_endpoints.py @@ -276,10 +276,13 @@ async def is_vantage_setup_in_db() -> bool: def is_vantage_setup_in_config() -> bool: - """Check if Vantage is setup in config.yaml or environment variables.""" - import litellm + """Check if Vantage is setup in config.yaml, environment variables, or programmatically.""" + from litellm.integrations.vantage.vantage_logger import VantageLogger - return "vantage" in litellm.callbacks + for cb in litellm.callbacks: + if cb == "vantage" or isinstance(cb, VantageLogger): + return True + return False async def is_vantage_setup() -> bool: @@ -412,6 +415,8 @@ async def vantage_dry_run_export( summary=summary, ) + except HTTPException: + raise except Exception as e: verbose_proxy_logger.error( f"Error performing Vantage dry run export: {str(e)}" From 9e623ceaff0cb30a261e5241cdd356a37932fe58 Mon Sep 17 00:00:00 2001 From: Harshit28j Date: Wed, 11 Mar 2026 15:16:49 +0530 Subject: [PATCH 09/21] Remove redundant empty-frame guard in FocusCsvSerializer Polars write_csv already handles empty DataFrames correctly (outputs header-only CSV), so the conditional branch was a no-op. Co-Authored-By: Claude Opus 4.6 --- litellm/integrations/focus/serializers/csv.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/litellm/integrations/focus/serializers/csv.py b/litellm/integrations/focus/serializers/csv.py index 30e7f3283d..320c99517d 100644 --- a/litellm/integrations/focus/serializers/csv.py +++ b/litellm/integrations/focus/serializers/csv.py @@ -16,7 +16,6 @@ class FocusCsvSerializer(FocusSerializer): def serialize(self, frame: pl.DataFrame) -> bytes: """Encode the provided frame as a CSV payload.""" - target = frame if not frame.is_empty() else pl.DataFrame(schema=frame.schema) buffer = io.BytesIO() - target.write_csv(buffer) + frame.write_csv(buffer) return buffer.getvalue() From 25c865876188959fd141435391ac0c0636a3e50d Mon Sep 17 00:00:00 2001 From: Harshit28j Date: Wed, 11 Mar 2026 15:34:24 +0530 Subject: [PATCH 10/21] Fix oversized-row handling, batch error resilience, and callback registration Vantage destination: - Skip individual CSV rows exceeding 2MB limit with a warning instead of uploading an oversized batch that Vantage would reject - Wrap each batch upload in try/except so remaining batches continue on failure; re-raise the first error after all batches are attempted Callback registration: - Use litellm.logging_callback_manager.add_litellm_callback() instead of raw litellm.callbacks.append() for DB-bootstrapped VantageLogger, ensuring proper dedup and manager visibility - Add "vantage" init handler in litellm_logging.py (both creation and lookup branches) so config.yaml string callbacks are properly resolved Co-Authored-By: Claude Opus 4.6 --- .../focus/destinations/vantage_destination.py | 49 +++++++++++++++---- litellm/litellm_core_utils/litellm_logging.py | 15 ++++++ litellm/proxy/proxy_server.py | 4 +- 3 files changed, 57 insertions(+), 11 deletions(-) diff --git a/litellm/integrations/focus/destinations/vantage_destination.py b/litellm/integrations/focus/destinations/vantage_destination.py index e2a21c773f..59c4e0cdb9 100644 --- a/litellm/integrations/focus/destinations/vantage_destination.py +++ b/litellm/integrations/focus/destinations/vantage_destination.py @@ -99,26 +99,41 @@ class FocusVantageDestination(FocusDestination): async def _upload_batched( self, client: httpx.AsyncClient, csv_bytes: bytes, filename: str ) -> None: - """Split the CSV into batches and upload each.""" + """Split the CSV into batches and upload each. + + Continues uploading remaining batches even if one fails, then raises + the first error encountered so callers know the export was partial. + """ lines = csv_bytes.split(b"\n") header = lines[0] data_lines = [line for line in lines[1:] if line.strip()] + first_error: Optional[Exception] = None batch_num = 0 for start in range(0, len(data_lines), VANTAGE_MAX_ROWS_PER_UPLOAD): batch_lines = data_lines[start : start + VANTAGE_MAX_ROWS_PER_UPLOAD] batch_csv = header + b"\n" + b"\n".join(batch_lines) + b"\n" - # 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( - client, header, batch_lines, filename, batch_num + try: + # 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( + client, header, batch_lines, filename, batch_num + ) + else: + batch_filename = f"{filename}.part{batch_num}" + await self._upload_csv(client, batch_csv, batch_filename) + except Exception as e: + verbose_logger.error( + "Vantage destination: batch %d failed: %s", batch_num, e ) - else: - batch_filename = f"{filename}.part{batch_num}" - await self._upload_csv(client, batch_csv, batch_filename) + if first_error is None: + first_error = e batch_num += 1 + if first_error is not None: + raise first_error + async def _upload_size_limited( self, client: httpx.AsyncClient, @@ -127,19 +142,33 @@ class FocusVantageDestination(FocusDestination): filename: str, batch_offset: int, ) -> None: - """Upload lines in chunks that stay under the 2 MB size limit.""" + """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. + """ current_chunk: list[bytes] = [] current_size = len(header) + 1 # header + newline sub_batch = 0 + header_size = len(header) + 1 for line in data_lines: line_size = len(line) + 1 # line + newline + + # Skip individual rows that exceed the limit on their own + if header_size + line_size > VANTAGE_MAX_BYTES_PER_UPLOAD: + verbose_logger.warning( + "Vantage destination: skipping oversized row (%d bytes)", + line_size, + ) + continue + 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) current_chunk = [] - current_size = len(header) + 1 + current_size = header_size sub_batch += 1 current_chunk.append(line) current_size += line_size diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 6f587abcdf..7157f4cb63 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -3873,6 +3873,15 @@ def _init_custom_logger_compatible_class( # noqa: PLR0915 focus_logger = FocusLogger() _in_memory_loggers.append(focus_logger) return focus_logger # type: ignore + elif logging_integration == "vantage": + from litellm.integrations.vantage.vantage_logger import VantageLogger + + for callback in _in_memory_loggers: + if isinstance(callback, VantageLogger): + return callback # type: ignore + vantage_logger = VantageLogger() + _in_memory_loggers.append(vantage_logger) + return vantage_logger # type: ignore elif logging_integration == "deepeval": for callback in _in_memory_loggers: if isinstance(callback, DeepEvalLogger): @@ -4243,6 +4252,12 @@ def get_custom_logger_compatible_class( # noqa: PLR0915 for callback in _in_memory_loggers: if isinstance(callback, FocusLogger): return callback + elif logging_integration == "vantage": + from litellm.integrations.vantage.vantage_logger import VantageLogger + + for callback in _in_memory_loggers: + if isinstance(callback, VantageLogger): + return callback elif logging_integration == "deepeval": for callback in _in_memory_loggers: if isinstance(callback, DeepEvalLogger): diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 9bd4cd7f14..c71df3a2ce 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -6152,7 +6152,9 @@ class ProxyStartupEvent: integration_token=db_settings.get("integration_token"), base_url=db_settings.get("base_url"), ) - litellm.callbacks.append(vantage_logger) # type: ignore[arg-type] + litellm.logging_callback_manager.add_litellm_callback( + vantage_logger + ) except Exception as e: verbose_proxy_logger.warning( "Failed to register VantageLogger from DB settings: %s", e From d35abfb55f926ae784b864356531923d649434ef Mon Sep 17 00:00:00 2001 From: Harshit28j Date: Wed, 11 Mar 2026 16:06:59 +0530 Subject: [PATCH 11/21] 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 59c4e0cdb9..0e06c8f441 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 01529d537f..df51bb7286 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 f478ba5160..cd78cf7ebd 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 ce2bad84e3..12d9099f81 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""" From 9200b2807899a267e3ce5e64dbaa70fca244478c Mon Sep 17 00:00:00 2001 From: Harshit28j Date: Wed, 11 Mar 2026 16:31:30 +0530 Subject: [PATCH 12/21] Fix dry-run summary alignment, token logging, and upload error handling - Align dry-run summary to use pre-transform columns (spend, total_tokens, team_id, model) matching FocusExportEngine internals - Reduce token exposure in debug logs to first 4 chars - Wrap raise_for_status in try/except to log httpx errors before re-raising Co-Authored-By: Claude Opus 4.6 --- .../focus/destinations/vantage_destination.py | 10 +++++++++- litellm/integrations/vantage/vantage_logger.py | 2 +- litellm/proxy/spend_tracking/vantage_endpoints.py | 11 ++++++----- 3 files changed, 16 insertions(+), 7 deletions(-) diff --git a/litellm/integrations/focus/destinations/vantage_destination.py b/litellm/integrations/focus/destinations/vantage_destination.py index 0e06c8f441..8275c0690b 100644 --- a/litellm/integrations/focus/destinations/vantage_destination.py +++ b/litellm/integrations/focus/destinations/vantage_destination.py @@ -88,7 +88,15 @@ class FocusVantageDestination(FocusDestination): headers=headers, files={"file": (filename, csv_bytes, "text/csv")}, ) - response.raise_for_status() + try: + response.raise_for_status() + except httpx.HTTPStatusError as e: + verbose_logger.error( + "Vantage destination: upload failed for %s — %s", + filename, + e, + ) + raise verbose_logger.debug( "Vantage destination: uploaded %d bytes (%s)", diff --git a/litellm/integrations/vantage/vantage_logger.py b/litellm/integrations/vantage/vantage_logger.py index df51bb7286..5dbcce8c33 100644 --- a/litellm/integrations/vantage/vantage_logger.py +++ b/litellm/integrations/vantage/vantage_logger.py @@ -83,7 +83,7 @@ class VantageLogger(FocusLogger): verbose_logger.debug( "VantageLogger initialized (integration_token=%s)", - resolved_token[:8] + "..." if resolved_token else "None", + resolved_token[:4] + "***" if resolved_token and len(resolved_token) > 4 else "***", ) @staticmethod diff --git a/litellm/proxy/spend_tracking/vantage_endpoints.py b/litellm/proxy/spend_tracking/vantage_endpoints.py index cd78cf7ebd..4a11bb0e9d 100644 --- a/litellm/proxy/spend_tracking/vantage_endpoints.py +++ b/litellm/proxy/spend_tracking/vantage_endpoints.py @@ -392,13 +392,14 @@ 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 [] - # Compute summary from the FOCUS-normalized DataFrame. - # These use post-transform column names specific to this endpoint. + # Use the same pre-transform column names as + # FocusExportEngine.dry_run_export_usage_data for consistency. summary = { "total_records": len(normalized), - "total_spend": FocusExportEngine._sum_column(normalized, "BilledCost"), - "unique_teams": FocusExportEngine._count_unique(normalized, "SubAccountId"), - "unique_models": FocusExportEngine._count_unique(normalized, "ResourceType"), + "total_spend": FocusExportEngine._sum_column(data, "spend"), + "total_tokens": FocusExportEngine._sum_column(data, "total_tokens"), + "unique_teams": FocusExportEngine._count_unique(data, "team_id"), + "unique_models": FocusExportEngine._count_unique(data, "model"), } dry_run_result = { From 3cc2ccec4ae2bd001cb9b30ce00f3c0de52204e3 Mon Sep 17 00:00:00 2001 From: Harshit28j Date: Wed, 11 Mar 2026 16:42:35 +0530 Subject: [PATCH 13/21] Fix double-scheduling bug and Decimal CSV serialization - Guard FocusLogger.init_focus_export_background_job with exact type check (type(cb) is FocusLogger) to exclude VantageLogger subclass, preventing duplicate hourly exports when VantageLogger is registered programmatically before startup - Cast pl.Decimal columns to Float64 in FocusCsvSerializer so CSV output uses standard floating-point notation instead of fixed-point strings that Vantage's parser may reject Co-Authored-By: Claude Opus 4.6 --- litellm/integrations/focus/focus_logger.py | 14 +++++++++----- litellm/integrations/focus/serializers/csv.py | 12 ++++++++++++ 2 files changed, 21 insertions(+), 5 deletions(-) diff --git a/litellm/integrations/focus/focus_logger.py b/litellm/integrations/focus/focus_logger.py index ade1cf861b..a1b75fbfff 100644 --- a/litellm/integrations/focus/focus_logger.py +++ b/litellm/integrations/focus/focus_logger.py @@ -137,11 +137,15 @@ class FocusLogger(CustomLogger): ) -> None: """Register the export cron/interval job with the provided scheduler.""" - focus_loggers: List[ - CustomLogger - ] = litellm.logging_callback_manager.get_custom_loggers_for_type( - callback_type=FocusLogger - ) + # Use exact type match to exclude subclasses like VantageLogger, + # which have their own dedicated scheduling method. + focus_loggers: List[CustomLogger] = [ + cb + for cb in litellm.logging_callback_manager.get_custom_loggers_for_type( + callback_type=FocusLogger + ) + if type(cb) is FocusLogger + ] if not focus_loggers: verbose_logger.debug( "No Focus export logger registered; skipping scheduler" diff --git a/litellm/integrations/focus/serializers/csv.py b/litellm/integrations/focus/serializers/csv.py index 320c99517d..8e33c557be 100644 --- a/litellm/integrations/focus/serializers/csv.py +++ b/litellm/integrations/focus/serializers/csv.py @@ -16,6 +16,18 @@ class FocusCsvSerializer(FocusSerializer): def serialize(self, frame: pl.DataFrame) -> bytes: """Encode the provided frame as a CSV payload.""" + # Cast Decimal columns to Float64 so CSV output uses standard + # floating-point notation (e.g. "1.5") instead of fixed-point + # strings (e.g. "1.500000") that some parsers may reject. + decimal_cols = [ + col + for col, dtype in zip(frame.columns, frame.dtypes) + if isinstance(dtype, pl.Decimal) + ] + if decimal_cols: + frame = frame.with_columns( + [pl.col(c).cast(pl.Float64) for c in decimal_cols] + ) buffer = io.BytesIO() frame.write_csv(buffer) return buffer.getvalue() From 4a9d03f1b11917a82a540c6166cd4d5b4767d185 Mon Sep 17 00:00:00 2001 From: Harshit28j Date: Wed, 11 Mar 2026 17:01:16 +0530 Subject: [PATCH 14/21] Fix FocusLogger dedup to use exact type match in litellm_logging.py Use `type(cb) is FocusLogger` instead of `isinstance(cb, FocusLogger)` in both _init_custom_logger_compatible_class and get_custom_logger_compatible_class so that a VantageLogger already in _in_memory_loggers is not incorrectly returned for a "focus" lookup. This ensures users with both "vantage" and "focus" in success_callbacks get separate loggers for each export destination. Co-Authored-By: Claude Opus 4.6 --- litellm/litellm_core_utils/litellm_logging.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 7157f4cb63..81f85e3d5b 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -3868,7 +3868,7 @@ def _init_custom_logger_compatible_class( # noqa: PLR0915 from litellm.integrations.focus.focus_logger import FocusLogger for callback in _in_memory_loggers: - if isinstance(callback, FocusLogger): + if type(callback) is FocusLogger: # exact match; exclude subclasses like VantageLogger return callback # type: ignore focus_logger = FocusLogger() _in_memory_loggers.append(focus_logger) @@ -4250,7 +4250,7 @@ def get_custom_logger_compatible_class( # noqa: PLR0915 from litellm.integrations.focus.focus_logger import FocusLogger for callback in _in_memory_loggers: - if isinstance(callback, FocusLogger): + if type(callback) is FocusLogger: # exact match; exclude subclasses like VantageLogger return callback elif logging_integration == "vantage": from litellm.integrations.vantage.vantage_logger import VantageLogger From 212cd0e4aadae0522392f9550a0ca4c968061eca Mon Sep 17 00:00:00 2001 From: Harshit28j Date: Wed, 11 Mar 2026 17:28:19 +0530 Subject: [PATCH 15/21] Fix pod lock collision, empty credential validation, and interval parsing - Override initialize_focus_export_job in VantageLogger to use VANTAGE_USAGE_DATA_JOB_NAME as the Redis pod lock key, preventing silent export skips when both Focus and Vantage loggers are configured - Add field_validator to VantageInitRequest rejecting empty-string api_key and integration_token at init time instead of at export time - Guard FocusLogger FOCUS_INTERVAL_SECONDS with try/except matching the pattern already used in VantageLogger Co-Authored-By: Claude Opus 4.6 --- litellm/integrations/focus/focus_logger.py | 10 +++++- .../integrations/vantage/vantage_logger.py | 31 +++++++++++++++++++ litellm/types/proxy/vantage_endpoints.py | 9 +++++- 3 files changed, 48 insertions(+), 2 deletions(-) diff --git a/litellm/integrations/focus/focus_logger.py b/litellm/integrations/focus/focus_logger.py index a1b75fbfff..f461f8c9fa 100644 --- a/litellm/integrations/focus/focus_logger.py +++ b/litellm/integrations/focus/focus_logger.py @@ -53,7 +53,15 @@ class FocusLogger(CustomLogger): if interval_seconds is not None else os.getenv("FOCUS_INTERVAL_SECONDS") ) - self.interval_seconds = int(raw_interval) if raw_interval is not None else None + self.interval_seconds: Optional[int] = None + if raw_interval is not None: + try: + self.interval_seconds = int(raw_interval) + except (ValueError, TypeError): + verbose_logger.warning( + "Invalid FOCUS_INTERVAL_SECONDS value: %s, ignoring", + raw_interval, + ) env_prefix = os.getenv("FOCUS_PREFIX") self.prefix: str = ( prefix if prefix is not None else (env_prefix if env_prefix else "focus_exports") diff --git a/litellm/integrations/vantage/vantage_logger.py b/litellm/integrations/vantage/vantage_logger.py index 5dbcce8c33..6689d93274 100644 --- a/litellm/integrations/vantage/vantage_logger.py +++ b/litellm/integrations/vantage/vantage_logger.py @@ -86,6 +86,37 @@ class VantageLogger(FocusLogger): resolved_token[:4] + "***" if resolved_token and len(resolved_token) > 4 else "***", ) + async def initialize_focus_export_job(self) -> None: + """Override to use the Vantage-specific pod lock key. + + Without this, VantageLogger and FocusLogger would compete for the + same ``FOCUS_USAGE_DATA_JOB_NAME`` lock, causing one to silently + skip its export cycle when both are configured simultaneously. + """ + from litellm.proxy.proxy_server import proxy_logging_obj + + pod_lock_manager = None + if proxy_logging_obj is not None: + writer = getattr(proxy_logging_obj, "db_spend_update_writer", None) + if writer is not None: + pod_lock_manager = getattr(writer, "pod_lock_manager", None) + + if pod_lock_manager and pod_lock_manager.redis_cache: + acquired = await pod_lock_manager.acquire_lock( + cronjob_id=VANTAGE_USAGE_DATA_JOB_NAME + ) + if not acquired: + verbose_logger.debug("Vantage export: unable to acquire pod lock") + return + try: + await self._run_scheduled_export() + finally: + await pod_lock_manager.release_lock( + cronjob_id=VANTAGE_USAGE_DATA_JOB_NAME + ) + else: + await self._run_scheduled_export() + @staticmethod async def init_vantage_background_job( scheduler: AsyncIOScheduler, diff --git a/litellm/types/proxy/vantage_endpoints.py b/litellm/types/proxy/vantage_endpoints.py index 12d9099f81..86b5a26dbc 100644 --- a/litellm/types/proxy/vantage_endpoints.py +++ b/litellm/types/proxy/vantage_endpoints.py @@ -5,7 +5,7 @@ Vantage endpoint types for LiteLLM Proxy from datetime import datetime from typing import Any, Dict, Optional -from pydantic import BaseModel, Field +from pydantic import BaseModel, Field, field_validator class VantageInitRequest(BaseModel): @@ -20,6 +20,13 @@ class VantageInitRequest(BaseModel): description="Vantage API base URL (default: https://api.vantage.sh)", ) + @field_validator("api_key", "integration_token") + @classmethod + def must_be_non_empty(cls, v: str) -> str: + if not v.strip(): + raise ValueError("must be a non-empty string") + return v + class VantageInitResponse(BaseModel): """Response model for Vantage initialization""" From 358c2fd033231c31ebc2416943c240999e80aaab Mon Sep 17 00:00:00 2001 From: Harshit28j Date: Wed, 11 Mar 2026 17:41:23 +0530 Subject: [PATCH 16/21] Add empty-string credential validator to VantageSettingsUpdate Matches the existing validator on VantageInitRequest so that empty or whitespace-only api_key/integration_token values are rejected at update time rather than silently persisted and failing at the next export. Co-Authored-By: Claude Opus 4.6 --- litellm/types/proxy/vantage_endpoints.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/litellm/types/proxy/vantage_endpoints.py b/litellm/types/proxy/vantage_endpoints.py index 86b5a26dbc..cf4f0a6685 100644 --- a/litellm/types/proxy/vantage_endpoints.py +++ b/litellm/types/proxy/vantage_endpoints.py @@ -95,3 +95,10 @@ class VantageSettingsUpdate(BaseModel): None, description="New Vantage integration token" ) base_url: Optional[str] = Field(None, description="New Vantage API base URL") + + @field_validator("api_key", "integration_token") + @classmethod + def must_be_non_empty(cls, v: Optional[str]) -> Optional[str]: + if v is not None and not v.strip(): + raise ValueError("must be a non-empty string") + return v From ce052d07be9093b85042db1c33c6bbac9c4488d2 Mon Sep 17 00:00:00 2001 From: Harshit28j Date: Wed, 11 Mar 2026 18:12:27 +0530 Subject: [PATCH 17/21] Fix test helper hour=23 bug and add row-count batching test - Use timedelta(hours=1) instead of replace(hour=hour+1) in _window() to avoid ValueError when hour=23 - Add test_should_batch_by_row_count covering the >10K rows batching path (previously only the 2 MB size-limit path was tested) Co-Authored-By: Claude Opus 4.6 --- .../focus/test_vantage_destination.py | 53 ++++++++++++++++++- 1 file changed, 51 insertions(+), 2 deletions(-) diff --git a/tests/test_litellm/integrations/focus/test_vantage_destination.py b/tests/test_litellm/integrations/focus/test_vantage_destination.py index 181d3d1381..c3fa9bda63 100644 --- a/tests/test_litellm/integrations/focus/test_vantage_destination.py +++ b/tests/test_litellm/integrations/focus/test_vantage_destination.py @@ -2,7 +2,7 @@ from __future__ import annotations -from datetime import datetime, timezone +from datetime import datetime, timedelta, timezone from typing import Any, Dict, List from unittest.mock import AsyncMock, patch @@ -12,12 +12,13 @@ from litellm.integrations.focus.destinations.base import FocusTimeWindow from litellm.integrations.focus.destinations.vantage_destination import ( FocusVantageDestination, VANTAGE_MAX_BYTES_PER_UPLOAD, + VANTAGE_MAX_ROWS_PER_UPLOAD, ) def _window(freq: str = "hourly", hour: int = 5) -> FocusTimeWindow: start = datetime(2024, 1, 2, hour, tzinfo=timezone.utc) - end = start.replace(hour=hour + 1) + end = start + timedelta(hours=1) return FocusTimeWindow(start_time=start, end_time=end, frequency=freq) @@ -136,3 +137,51 @@ async def test_should_batch_large_content(): # Each upload should be within limits for chunk in upload_calls: assert len(chunk) <= VANTAGE_MAX_BYTES_PER_UPLOAD + + +@pytest.mark.asyncio +async def test_should_batch_by_row_count(): + """Verify batching triggers when row count exceeds 10K even if under 2 MB.""" + dest = FocusVantageDestination(prefix="exports", config=_config()) + + header = b"col1" + # Short rows so total size stays well under 2 MB + row = b"x" + num_rows = VANTAGE_MAX_ROWS_PER_UPLOAD + 500 + content = header + b"\n" + b"\n".join([row] * num_rows) + b"\n" + + # Confirm content is under 2 MB but over 10K rows + assert len(content) < VANTAGE_MAX_BYTES_PER_UPLOAD + assert num_rows > VANTAGE_MAX_ROWS_PER_UPLOAD + + upload_calls: List[bytes] = [] + + mock_response = AsyncMock() + mock_response.raise_for_status = lambda: None + + mock_client = AsyncMock() + mock_client.__aenter__ = AsyncMock(return_value=mock_client) + mock_client.__aexit__ = AsyncMock(return_value=False) + + async def capture_post(url, **kwargs): + files = kwargs.get("files", {}) + if "file" in files: + upload_calls.append(files["file"][1]) + return mock_response + + mock_client.post = capture_post + + with patch("litellm.integrations.focus.destinations.vantage_destination.httpx.AsyncClient", return_value=mock_client): + await dest.deliver( + content=content, + time_window=_window(), + filename="usage.csv", + ) + + # Should have made at least 2 uploads due to row count + assert len(upload_calls) >= 2 + # Each batch should have at most 10K data rows (header + data rows + trailing newline) + for chunk in upload_calls: + lines = chunk.split(b"\n") + data_lines = [line for line in lines[1:] if line.strip()] + assert len(data_lines) <= VANTAGE_MAX_ROWS_PER_UPLOAD From 24d4e5bc60bf54156f7d72ed3b31932e2efddcb7 Mon Sep 17 00:00:00 2001 From: Harshit28j Date: Wed, 11 Mar 2026 19:28:27 +0530 Subject: [PATCH 18/21] Deregister VantageLogger on delete and add Decimal cast test MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - DELETE /vantage/delete now removes the in-memory VantageLogger from litellm.callbacks via remove_callbacks_by_type, preventing the scheduler from continuing to fire exports with stale credentials - Add test_should_cast_decimal_columns_to_float covering the Decimal→Float64 cast in FocusCsvSerializer Co-Authored-By: Claude Opus 4.6 --- litellm/proxy/spend_tracking/vantage_endpoints.py | 7 +++++++ .../integrations/focus/test_csv_serializer.py | 13 +++++++++++++ 2 files changed, 20 insertions(+) diff --git a/litellm/proxy/spend_tracking/vantage_endpoints.py b/litellm/proxy/spend_tracking/vantage_endpoints.py index 4a11bb0e9d..19859bfb83 100644 --- a/litellm/proxy/spend_tracking/vantage_endpoints.py +++ b/litellm/proxy/spend_tracking/vantage_endpoints.py @@ -547,6 +547,13 @@ async def delete_vantage_settings( where={"param_name": VANTAGE_SETTINGS_PARAM_NAME} ) + # Deregister in-memory VantageLogger so the scheduler stops firing + from litellm.integrations.vantage.vantage_logger import VantageLogger + + litellm.logging_callback_manager.remove_callbacks_by_type( + litellm.callbacks, VantageLogger + ) + verbose_proxy_logger.info("Vantage settings deleted successfully") return VantageInitResponse( diff --git a/tests/test_litellm/integrations/focus/test_csv_serializer.py b/tests/test_litellm/integrations/focus/test_csv_serializer.py index f527a9dceb..f3256808e4 100644 --- a/tests/test_litellm/integrations/focus/test_csv_serializer.py +++ b/tests/test_litellm/integrations/focus/test_csv_serializer.py @@ -30,5 +30,18 @@ def test_should_return_header_only_for_empty_frame(): assert len(lines) == 1 # header only +def test_should_cast_decimal_columns_to_float(): + frame = pl.DataFrame( + {"BilledCost": [1, 2], "ServiceName": ["openai", "anthropic"]} + ).cast({"BilledCost": pl.Decimal(18, 6)}) + serializer = FocusCsvSerializer() + result = serializer.serialize(frame) + + lines = result.decode("utf-8").strip().split("\n") + # Should use floating-point notation (e.g. "1.0") not fixed-point ("1.000000") + assert "1.000000" not in lines[1] + assert lines[1].startswith("1.0,") + + def test_extension_should_be_csv(): assert FocusCsvSerializer.extension == "csv" From e878941da5e037dc615e0a63557176a09dc8e31e Mon Sep 17 00:00:00 2001 From: Harshit28j Date: Wed, 11 Mar 2026 19:42:29 +0530 Subject: [PATCH 19/21] Fix Decimal serialization crash in dry-run endpoint Cast Polars Decimal columns to Float64 before calling .to_dicts() in vantage_dry_run_export so the response contains JSON-serializable float values instead of decimal.Decimal objects that FastAPI cannot encode. Also cast summary totals to float for the same reason. Co-Authored-By: Claude Opus 4.6 --- .../proxy/spend_tracking/vantage_endpoints.py | 25 ++++++++++++++++--- 1 file changed, 21 insertions(+), 4 deletions(-) diff --git a/litellm/proxy/spend_tracking/vantage_endpoints.py b/litellm/proxy/spend_tracking/vantage_endpoints.py index 19859bfb83..14c0ebcb54 100644 --- a/litellm/proxy/spend_tracking/vantage_endpoints.py +++ b/litellm/proxy/spend_tracking/vantage_endpoints.py @@ -386,18 +386,35 @@ async def vantage_dry_run_export( database = FocusLiteLLMDatabase() transformer = FocusTransformer() + import polars as pl + data = await database.get_usage_data(limit=request.limit) normalized = transformer.transform(data) - 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 [] + def _to_json_safe_dicts(frame: pl.DataFrame) -> list: + """Cast Decimal columns to Float64 so .to_dicts() produces + JSON-serializable float values instead of decimal.Decimal.""" + decimal_cols = [ + col for col, dtype in zip(frame.columns, frame.dtypes) + if isinstance(dtype, pl.Decimal) + ] + if decimal_cols: + frame = frame.with_columns( + [pl.col(c).cast(pl.Float64) for c in decimal_cols] + ) + return frame.to_dicts() + + usage_sample = _to_json_safe_dicts(data.head(min(50, len(data)))) if not data.is_empty() else [] + normalized_sample = _to_json_safe_dicts(normalized.head(min(50, len(normalized)))) if not normalized.is_empty() else [] # Use the same pre-transform column names as # FocusExportEngine.dry_run_export_usage_data for consistency. + total_spend = FocusExportEngine._sum_column(data, "spend") + total_tokens = FocusExportEngine._sum_column(data, "total_tokens") summary = { "total_records": len(normalized), - "total_spend": FocusExportEngine._sum_column(data, "spend"), - "total_tokens": FocusExportEngine._sum_column(data, "total_tokens"), + "total_spend": float(total_spend) if total_spend is not None else 0, + "total_tokens": float(total_tokens) if total_tokens is not None else 0, "unique_teams": FocusExportEngine._count_unique(data, "team_id"), "unique_models": FocusExportEngine._count_unique(data, "model"), } From 118d8114f8d606faa40f7802765f53ffce37762c Mon Sep 17 00:00:00 2001 From: Harshit28j Date: Fri, 13 Mar 2026 15:37:13 +0530 Subject: [PATCH 20/21] Include time window in export filename to prevent Vantage overwrites MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Vantage deduplicates by filename — uploading the same filename overwrites previous data. Changed _build_filename to include the time window (e.g. usage_20240102T050000Z_20240102T060000Z.csv) so each hourly/daily export has a unique filename. Co-Authored-By: Claude Opus 4.6 --- litellm/integrations/focus/export_engine.py | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/litellm/integrations/focus/export_engine.py b/litellm/integrations/focus/export_engine.py index 1f51e1bac9..8b77fad9e1 100644 --- a/litellm/integrations/focus/export_engine.py +++ b/litellm/integrations/focus/export_engine.py @@ -101,13 +101,17 @@ class FocusExportEngine: await self._destination.deliver( content=payload, time_window=window, - filename=self._build_filename(), + filename=self._build_filename(window), ) - def _build_filename(self) -> str: + def _build_filename(self, window: FocusTimeWindow) -> str: if not self._serializer.extension: raise ValueError("Serializer must declare a file extension") - return f"usage.{self._serializer.extension}" + # Include time window in filename so Vantage (which deduplicates + # by filename) doesn't overwrite previous uploads. + start_str = window.start_time.strftime("%Y%m%dT%H%M%SZ") + end_str = window.end_time.strftime("%Y%m%dT%H%M%SZ") + return f"usage_{start_str}_{end_str}.{self._serializer.extension}" @staticmethod def _sum_column(frame: pl.DataFrame, column: str) -> float: From d72a34dbe2e167d78e7f7baf79438ceff5f4aa17 Mon Sep 17 00:00:00 2001 From: Harshit28j Date: Sat, 14 Mar 2026 22:19:56 +0530 Subject: [PATCH 21/21] add docs and export fixes --- docs/my-website/docs/observability/vantage.md | 148 ++++++++++++++++++ .../focus/destinations/vantage_destination.py | 102 ++++++++++-- litellm/integrations/focus/export_engine.py | 27 ++++ litellm/integrations/focus/focus_logger.py | 27 +++- 4 files changed, 288 insertions(+), 16 deletions(-) create mode 100644 docs/my-website/docs/observability/vantage.md diff --git a/docs/my-website/docs/observability/vantage.md b/docs/my-website/docs/observability/vantage.md new file mode 100644 index 0000000000..31b43a76c3 --- /dev/null +++ b/docs/my-website/docs/observability/vantage.md @@ -0,0 +1,148 @@ +import Tabs from '@theme/Tabs'; +import TabItem from '@theme/TabItem'; + +# Vantage Integration + +LiteLLM can export proxy spend data to [Vantage](https://vantage.sh) as [FOCUS 1.2](https://focus.finops.org/) formatted cost reports. This lets you visualize LLM spend alongside your cloud infrastructure costs in the Vantage dashboard. + +## Overview + +| Property | Details | +|----------|---------| +| Destination | Export LiteLLM usage data to Vantage Custom Provider | +| Data format | FOCUS CSV (automatically transformed from LiteLLM spend data) | +| Supported operations | Manual export, automatic scheduled export (hourly/daily/interval) | +| Authentication | Vantage API key + Custom Provider token | + +## Prerequisites + +You need two credentials from the [Vantage console](https://console.vantage.sh): + +1. **API Key** — Go to **Settings → API Access Tokens** → Create a token with **Write** scope. The token looks like `vntg_tkn_...`. +2. **Custom Provider Token** — Go to **Settings → Integrations** → Create a **Custom Provider** integration → Copy the Provider ID (looks like `accss_crdntl_...`). + +## Setup via API + +The recommended setup uses the proxy admin endpoints. No config file changes needed. + +### 1. Initialize credentials + +```bash +curl -X POST http://localhost:4000/vantage/init \ + -H "Authorization: Bearer $LITELLM_ADMIN_KEY" \ + -H "Content-Type: application/json" \ + -d '{ + "api_key": "vntg_tkn_YOUR_VANTAGE_API_KEY", + "integration_token": "accss_crdntl_YOUR_PROVIDER_TOKEN" + }' +``` + +Credentials are encrypted and stored in the proxy database. + +### 2. Preview data (dry run) + +```bash +curl -X POST http://localhost:4000/vantage/dry-run \ + -H "Authorization: Bearer $LITELLM_ADMIN_KEY" \ + -H "Content-Type: application/json" \ + -d '{"limit": 10}' +``` + +This returns FOCUS-transformed data without sending anything to Vantage. Use it to verify the pipeline works and inspect the data mapping. + +### 3. Export to Vantage + +```bash +curl -X POST http://localhost:4000/vantage/export \ + -H "Authorization: Bearer $LITELLM_ADMIN_KEY" \ + -H "Content-Type: application/json" \ + -d '{}' +``` + +Optional parameters: +- `limit` — Max number of records to export +- `start_time_utc` / `end_time_utc` — Filter by time range (must be provided together) + +### 4. Verify in Vantage + +Go to **Settings → Integrations → your Custom Provider → Import Costs** tab to see uploaded CSVs. Once the status changes from "Importing and Processing" to "Stable", costs appear in **Cost Reporting → All Resources**. + +## Setup via Environment Variables + +For automatic scheduled exports, configure via environment variables and proxy config: + +### Environment variables + +| Variable | Required | Description | +|----------|----------|-------------| +| `VANTAGE_API_KEY` | Yes | Vantage API access token | +| `VANTAGE_INTEGRATION_TOKEN` | Yes | Custom Provider token from Vantage dashboard | +| `VANTAGE_BASE_URL` | No | API URL override (default: `https://api.vantage.sh`) | +| `VANTAGE_EXPORT_FREQUENCY` | No | `hourly` (default), `daily`, or `interval` | +| `VANTAGE_EXPORT_INTERVAL_SECONDS` | No | Seconds between exports when frequency is `interval` | + +### Proxy config + +```yaml +model_list: + - model_name: gpt-4o + litellm_params: + model: openai/gpt-4o + api_key: sk-your-key + +litellm_settings: + callbacks: ["vantage"] +``` + +```bash +export VANTAGE_API_KEY="vntg_tkn_..." +export VANTAGE_INTEGRATION_TOKEN="accss_crdntl_..." +litellm --config /path/to/config.yaml +``` + +The proxy registers a background job that exports data on the configured schedule. + +## API Endpoints + +All endpoints require admin authentication. + +| Method | Endpoint | Description | +|--------|----------|-------------| +| `POST` | `/vantage/init` | Store Vantage credentials (encrypted) | +| `GET` | `/vantage/settings` | View current config (credentials masked) | +| `PUT` | `/vantage/settings` | Update credentials or base URL | +| `POST` | `/vantage/dry-run` | Preview FOCUS data without uploading | +| `POST` | `/vantage/export` | Upload cost data to Vantage | +| `DELETE` | `/vantage/delete` | Remove credentials and stop scheduled exports | + +## FOCUS Field Mapping + +LiteLLM spend data is transformed into the FOCUS 1.2 schema: + +| LiteLLM Field | FOCUS Column | Description | +|---------------|-------------|-------------| +| `spend` | BilledCost, EffectiveCost | Cost of the usage | +| `model` | ChargeDescription, ResourceId | Model identifier | +| `model_group` | ServiceName | Model group / deployment | +| `custom_llm_provider` | ProviderName, PublisherName | Provider (openai, anthropic, etc.) | +| `api_key` | BillingAccountId | Hashed API key | +| `api_key_alias` | BillingAccountName | Human-readable key alias | +| `team_id` | SubAccountId | Team identifier | +| `team_alias` | SubAccountName | Team name | + +Additional metadata (user_id, model_group, etc.) is included in the `Tags` column as JSON. + +## Upload Limits + +Vantage enforces per-upload limits. LiteLLM handles these automatically: + +- **10,000 rows** per upload — large exports are split into batches +- **2 MB** per upload — oversized batches are further split by size +- **Unsupported columns** are stripped before upload + +## Related Links + +- [Vantage](https://vantage.sh) +- [Vantage Custom Providers](https://docs.vantage.sh/connecting_custom_providers) +- [FOCUS Specification](https://focus.finops.org/) +- [Focus Export (S3/Parquet)](./focus.md) diff --git a/litellm/integrations/focus/destinations/vantage_destination.py b/litellm/integrations/focus/destinations/vantage_destination.py index 8275c0690b..9e6028900f 100644 --- a/litellm/integrations/focus/destinations/vantage_destination.py +++ b/litellm/integrations/focus/destinations/vantage_destination.py @@ -2,6 +2,8 @@ from __future__ import annotations +import csv +import io from typing import Any, Optional import httpx @@ -14,6 +16,77 @@ from .base import FocusDestination, FocusTimeWindow VANTAGE_MAX_ROWS_PER_UPLOAD = 10_000 VANTAGE_MAX_BYTES_PER_UPLOAD = 2 * 1024 * 1024 # 2 MB +# Columns that Vantage actually supports for custom provider CSV uploads. +# See: https://docs.vantage.sh/connecting_custom_providers +# Columns not in this set are silently dropped before upload so Vantage +# does not reject the file. +VANTAGE_SUPPORTED_COLUMNS = { + # Required + "ChargeCategory", + "ChargePeriodStart", + "BilledCost", + "ServiceName", + # Optional + "BillingCurrency", + "BillingAccountId", + "BillingAccountName", + "ChargePeriodEnd", + "ChargeDescription", + "ChargeFrequency", + "ConsumedQuantity", + "ConsumedUnit", + "ContractedCost", + "EffectiveCost", + "ListCost", + "RegionId", + "RegionName", + "ResourceId", + "ResourceName", + "ResourceType", + "ServiceCategory", + "ServiceSubcategory", + "SubAccountId", + "SubAccountName", + "Tags", +} + + +def _strip_unsupported_columns(csv_bytes: bytes) -> bytes: + """Remove CSV columns not in VANTAGE_SUPPORTED_COLUMNS. + + Parses the header row, identifies column indices to keep, and + rebuilds the CSV with only those columns. + """ + lines = csv_bytes.split(b"\n") + if not lines: + return csv_bytes + + header_cols = lines[0].decode("utf-8").split(",") + keep_indices = [ + i + for i, col in enumerate(header_cols) + if col.strip('"') in VANTAGE_SUPPORTED_COLUMNS + ] + + # If all columns are supported, return as-is + if len(keep_indices) == len(header_cols): + return csv_bytes + + dropped = [col for i, col in enumerate(header_cols) if i not in keep_indices] + verbose_logger.debug( + "Vantage destination: dropping unsupported columns: %s", dropped + ) + + output = io.StringIO() + writer = csv.writer(output) + reader = csv.reader(io.StringIO(csv_bytes.decode("utf-8"))) + for row in reader: + if not row: + continue + writer.writerow([row[i] for i in keep_indices]) + + return output.getvalue().encode("utf-8") + class FocusVantageDestination(FocusDestination): """Upload FOCUS CSV exports to the Vantage cost-import API.""" @@ -39,9 +112,7 @@ class FocusVantageDestination(FocusDestination): ) self.api_key = api_key self.integration_token = integration_token - self.base_url = config.get( - "base_url", "https://api.vantage.sh" - ) + self.base_url = config.get("base_url", "https://api.vantage.sh") self.prefix = prefix async def deliver( @@ -56,6 +127,10 @@ class FocusVantageDestination(FocusDestination): verbose_logger.debug("Vantage destination: empty content, skipping upload") return + # Strip columns that Vantage does not support to avoid silent + # rejection (e.g. InvoiceIssuerName, ProviderName, PublisherName). + content = _strip_unsupported_columns(content) + # 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 @@ -75,10 +150,7 @@ class FocusVantageDestination(FocusDestination): 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" - ) + url = f"{self.base_url}/v2/integrations/" f"{self.integration_token}/costs.csv" headers = { "Authorization": f"Bearer {self.api_key}", } @@ -86,15 +158,16 @@ class FocusVantageDestination(FocusDestination): response = await client.post( url, headers=headers, - files={"file": (filename, csv_bytes, "text/csv")}, + files={"csv": (filename, csv_bytes, "text/csv")}, ) try: response.raise_for_status() except httpx.HTTPStatusError as e: verbose_logger.error( - "Vantage destination: upload failed for %s — %s", + "Vantage destination: upload failed for %s — %s — response body: %s", filename, e, + response.text, ) raise @@ -174,7 +247,10 @@ class FocusVantageDestination(FocusDestination): ) continue - if current_size + line_size > VANTAGE_MAX_BYTES_PER_UPLOAD and current_chunk: + 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}" try: @@ -182,7 +258,8 @@ class FocusVantageDestination(FocusDestination): except Exception as e: verbose_logger.error( "Vantage destination: sub-batch %s failed: %s", - batch_filename, e, + batch_filename, + e, ) if first_error is None: first_error = e @@ -200,7 +277,8 @@ class FocusVantageDestination(FocusDestination): except Exception as e: verbose_logger.error( "Vantage destination: sub-batch %s failed: %s", - batch_filename, e, + batch_filename, + e, ) if first_error is None: first_error = e diff --git a/litellm/integrations/focus/export_engine.py b/litellm/integrations/focus/export_engine.py index 8b77fad9e1..37da18a0eb 100644 --- a/litellm/integrations/focus/export_engine.py +++ b/litellm/integrations/focus/export_engine.py @@ -67,6 +67,33 @@ class FocusExportEngine: "summary": summary, } + async def export_all( + self, + *, + limit: Optional[int], + ) -> None: + """Export all available data without time-window filtering.""" + data = await self._database.get_usage_data(limit=limit) + if data.is_empty(): + verbose_logger.debug("Focus export: no usage data available") + return + + normalized = self._transformer.transform(data) + if normalized.is_empty(): + verbose_logger.debug("Focus export: normalized data empty") + return + + # Build a window spanning the full data range for the filename + from datetime import datetime, timezone + + now = datetime.now(timezone.utc) + window = FocusTimeWindow( + start_time=now.replace(hour=0, minute=0, second=0, microsecond=0), + end_time=now, + frequency="all", + ) + await self._serialize_and_upload(normalized, window) + async def export_window( self, *, diff --git a/litellm/integrations/focus/focus_logger.py b/litellm/integrations/focus/focus_logger.py index f461f8c9fa..083b0e1463 100644 --- a/litellm/integrations/focus/focus_logger.py +++ b/litellm/integrations/focus/focus_logger.py @@ -64,7 +64,9 @@ class FocusLogger(CustomLogger): ) env_prefix = os.getenv("FOCUS_PREFIX") self.prefix: str = ( - prefix if prefix is not None else (env_prefix if env_prefix else "focus_exports") + prefix + if prefix is not None + else (env_prefix if env_prefix else "focus_exports") ) self._destination_config = destination_config @@ -90,7 +92,13 @@ class FocusLogger(CustomLogger): start_time_utc: Optional[datetime] = None, end_time_utc: Optional[datetime] = None, ) -> None: - """Public hook to trigger export immediately.""" + """Public hook to trigger export immediately. + + When called without time bounds (manual /vantage/export with no + start/end), exports **all** available data instead of the last + scheduled window. The hourly/daily window only applies to + automatic scheduler runs. + """ if bool(start_time_utc) ^ bool(end_time_utc): raise ValueError( "start_time_utc and end_time_utc must be provided together" @@ -102,9 +110,10 @@ class FocusLogger(CustomLogger): end_time=end_time_utc, frequency=self.frequency, ) + await self._export_window(window=window, limit=limit) else: - window = self._compute_time_window(datetime.now(timezone.utc)) - await self._export_window(window=window, limit=limit) + # No time bounds → export all available data + await self._export_all(limit=limit) async def dry_run_export_usage_data( self, limit: Optional[int] = DEFAULT_DRY_RUN_LIMIT @@ -190,6 +199,15 @@ class FocusLogger(CustomLogger): window = self._compute_time_window(datetime.now(timezone.utc)) await self._export_window(window=window, limit=None) + async def _export_all( + self, + *, + limit: Optional[int], + ) -> None: + """Export all available data without a time window filter.""" + engine = self._ensure_engine() + await engine.export_all(limit=limit) + async def _export_window( self, *, @@ -220,4 +238,5 @@ class FocusLogger(CustomLogger): frequency=self.frequency, ) + __all__ = ["FocusLogger"]