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/__init__.py b/litellm/__init__.py index dcd4ce2909..41cc2f5d98 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -144,6 +144,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..9e6028900f --- /dev/null +++ b/litellm/integrations/focus/destinations/vantage_destination.py @@ -0,0 +1,287 @@ +"""Vantage API destination for Focus export.""" + +from __future__ import annotations + +import csv +import io +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 + +# 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.""" + + 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 + + # 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 + 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(client, content, filename) + + 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" + headers = { + "Authorization": f"Bearer {self.api_key}", + } + + response = await client.post( + url, + headers=headers, + 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 — response body: %s", + filename, + e, + response.text, + ) + raise + + verbose_logger.debug( + "Vantage destination: uploaded %d bytes (%s)", + len(csv_bytes), + filename, + ) + + async def _upload_batched( + self, client: httpx.AsyncClient, csv_bytes: bytes, filename: str + ) -> None: + """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" + + 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 + ) + 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, + header: bytes, + data_lines: list[bytes], + filename: str, + batch_offset: int, + ) -> None: + """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. 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 + + # 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}" + 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 + 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}" + 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/focus/export_engine.py b/litellm/integrations/focus/export_engine.py index 22ebce2a16..37da18a0eb 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) @@ -51,10 +55,10 @@ 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(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 { @@ -63,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, *, @@ -97,13 +128,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: diff --git a/litellm/integrations/focus/focus_logger.py b/litellm/integrations/focus/focus_logger.py index f493d47e29..083b0e1463 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 @@ -84,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" @@ -96,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 @@ -139,11 +154,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" @@ -180,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, *, diff --git a/litellm/integrations/focus/schema.py b/litellm/integrations/focus/schema.py index ac2f33dad0..c06ca9982a 100644 --- a/litellm/integrations/focus/schema.py +++ b/litellm/integrations/focus/schema.py @@ -43,7 +43,13 @@ FOCUS_NORMALIZED_SCHEMA = pl.Schema( ("SubAccountId", pl.String), ("SubAccountName", pl.String), ("SubAccountType", pl.String), - ("Tags", pl.Object), + # Changed from pl.Object to pl.String to hold JSON metadata + # (team_id, user_id, etc.) needed by Vantage Token Allocation. + # 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/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..8e33c557be --- /dev/null +++ b/litellm/integrations/focus/serializers/csv.py @@ -0,0 +1,33 @@ +"""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.""" + # 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() diff --git a/litellm/integrations/focus/transformer.py b/litellm/integrations/focus/transformer.py index cac12b7be1..b7d28e3dbb 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,38 @@ import polars as pl from .schema import FOCUS_NORMALIZED_SCHEMA +_TAG_KEYS = ( + "team_id", + "team_alias", + "user_id", + "user_email", + "api_key_alias", + "model", + "model_group", + "custom_llm_provider", +) + + +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`` 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: + 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: """Transforms LiteLLM DB rows into Focus-compatible schema.""" @@ -19,6 +52,13 @@ class FocusTransformer: if frame.is_empty(): return pl.DataFrame(schema=self.schema) + # 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: + frame = frame.with_columns(_build_tags_expr(available_keys)) + else: + frame = frame.with_columns(pl.lit("{}").alias("Tags")) + # derive period start/end from usage date frame = frame.with_columns( pl.col("date") @@ -86,5 +126,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("Tags").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..6689d93274 --- /dev/null +++ b/litellm/integrations/vantage/vantage_logger.py @@ -0,0 +1,144 @@ +"""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: 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: + 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[: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, + ) -> 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/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index e22d057bb6..f567363df1 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -3868,11 +3868,20 @@ 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) 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): @@ -4241,7 +4250,13 @@ 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 + + for callback in _in_memory_loggers: + if isinstance(callback, VantageLogger): return callback elif logging_integration == "deepeval": for callback in _in_memory_loggers: diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 2a9be0a67c..9c29927c5c 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -480,6 +480,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, ) @@ -6175,6 +6176,39 @@ 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 ( + _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.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 + ) + await VantageLogger.init_vantage_background_job(scheduler=scheduler) + ######################################################## # Prometheus Background Job ######################################################## @@ -13250,6 +13284,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..14c0ebcb54 --- /dev/null +++ b/litellm/proxy/spend_tracking/vantage_endpoints.py @@ -0,0 +1,587 @@ +import json + +import litellm +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 ( + VantageDryRunRequest, + VantageExportRequest, + VantageExportResponse, + VantageInitRequest, + VantageInitResponse, + VantageSettingsUpdate, + VantageSettingsView, +) + +router = APIRouter() + +_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 +): + """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) + encrypted_integration_token = encrypt_value_helper(integration_token) + + vantage_settings = { + "api_key": encrypted_api_key, + "integration_token": encrypted_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 + + 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 + + +@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_masked=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_masked=masked_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() + + 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.get("api_key", "") + ) + updated_token = ( + request.integration_token + if request.integration_token is not None + else current_settings.get("integration_token", "") + ) + updated_base_url = ( + request.base_url + if request.base_url is not None + else current_settings.get("base_url", "https://api.vantage.sh") + ) + + 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: + 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, environment variables, or programmatically.""" + from litellm.integrations.vantage.vantage_logger import VantageLogger + + for cb in litellm.callbacks: + if cb == "vantage" or isinstance(cb, VantageLogger): + return True + return False + + +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 HTTPException: + raise + 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: VantageDryRunRequest, + 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: Limit on number of records to preview (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: + # 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() + transformer = FocusTransformer() + + import polars as pl + + data = await database.get_usage_data(limit=request.limit) + normalized = transformer.transform(data) + + 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": 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"), + } + + dry_run_result = { + "usage_data": usage_sample, + "normalized_data": normalized_sample, + "summary": summary, + } + + 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=summary, + ) + + except HTTPException: + raise + 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: + from litellm.integrations.vantage.vantage_logger import VantageLogger + + # 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() + 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"), + 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 HTTPException: + raise + 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} + ) + + # 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( + 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..cf4f0a6685 --- /dev/null +++ b/litellm/types/proxy/vantage_endpoints.py @@ -0,0 +1,104 @@ +""" +Vantage endpoint types for LiteLLM Proxy +""" + +from datetime import datetime +from typing import Any, Dict, Optional + +from pydantic import BaseModel, Field, field_validator + + +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)", + ) + + @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""" + + message: str + status: str + + +class VantageExportRequest(BaseModel): + """Request model for Vantage export operations (actual export, no default limit)""" + + limit: Optional[int] = Field( + 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" + ) + end_time_utc: Optional[datetime] = Field( + None, description="End time for data export in UTC" + ) + + +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""" + + 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_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") + + +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") + + @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 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..f3256808e4 --- /dev/null +++ b/tests/test_litellm/integrations/focus/test_csv_serializer.py @@ -0,0 +1,47 @@ +"""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_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" 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..c3fa9bda63 --- /dev/null +++ b/tests/test_litellm/integrations/focus/test_vantage_destination.py @@ -0,0 +1,187 @@ +"""Tests for FocusVantageDestination behavior.""" + +from __future__ import annotations + +from datetime import datetime, timedelta, 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, + 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 + timedelta(hours=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 + + +@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