mirror of
https://github.com/tiennm99/litellm.git
synced 2026-07-25 06:21:17 +00:00
Merge pull request #23333 from Harshit28j/litellm_FOCUS_preserve-summary
Add Vantage integration for FOCUS CSV export
This commit is contained in:
@@ -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)
|
||||
@@ -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
|
||||
@@ -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:
|
||||
|
||||
@@ -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,
|
||||
*,
|
||||
|
||||
@@ -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()
|
||||
@@ -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,
|
||||
}
|
||||
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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)}"},
|
||||
)
|
||||
@@ -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
|
||||
Reference in New Issue
Block a user