Merge pull request #23333 from Harshit28j/litellm_FOCUS_preserve-summary

Add Vantage integration for FOCUS CSV export
This commit is contained in:
yuneng-jiang
2026-03-14 10:15:05 -07:00
committed by GitHub
21 changed files with 1796 additions and 25 deletions
@@ -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)
+1
View File
@@ -144,6 +144,7 @@ _custom_logger_compatible_callbacks_literal = Literal[
"gitlab",
"cloudzero",
"focus",
"vantage",
"posthog",
"levo",
]
@@ -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",
]
@@ -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"
)
@@ -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
+46 -11
View File
@@ -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:
+37 -9
View File
@@ -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,
*,
+7 -1
View File
@@ -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),
]
)
@@ -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"]
@@ -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()
+41 -1
View File
@@ -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"),
)
@@ -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"]
@@ -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,
}
+17 -2
View File
@@ -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:
+35
View File
@@ -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)
@@ -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)}"},
)
+104
View File
@@ -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
@@ -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"
@@ -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",
)
@@ -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