mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-03 02:22:24 +00:00
feat(mavvrik): add Mavvrik integration for automatic LLM spend export
Adds a complete integration with Mavvrik (https://help.mavvrik.ai/) for automatically exporting LiteLLM proxy spend data to the Mavvrik AI cost management platform. - Zero-config startup — set MAVVRIK_API_KEY, MAVVRIK_API_ENDPOINT, MAVVRIK_CONNECTION_ID env vars; proxy begins exporting automatically - Streaming export — fetches DB page-by-page (10k rows/page via OFFSET pagination with stable ORDER BY), gzip-compresses on-the-fly, uploads via GCS chunked resumable upload (256KB chunks). No row limit, ~2MB peak RAM. Validated: 500k rows → 13MB / 84s. - Idempotent by date — each calendar date has its own GCS object - Catch-up on restart — back-fills all missed days from Mavvrik marker - Pod-safe — Redis distributed lock for multi-replica deployments - Admin endpoints — init, settings, update, delete, dry-run, export litellm/integrations/mavvrik/ _http.py Shared HTTP transport (retry + exponential backoff) client.py Mavvrik REST API (register, advance_marker, signed URL) uploader.py GCS resumable upload (bulk + streaming, session cleanup) exporter.py DB fetch (4-table JOIN, stable pagination) + CSV orchestrator.py Pipeline sequencing + pod lock settings.py AES-encrypted credential persistence logger.py CustomLogger marker for callback registry __init__.py Service facade + lazy import to avoid polars at startup - update_settings reschedules background job with new credentials - delete works for env-var-only deployments (no DB row required) - dry_run guards NULL spend and missing completion_tokens column - GCS session cancelled on mid-stream _put_chunk failure - DB connectivity guard in Service.export/dry_run - Stable OFFSET pagination (ORDER BY includes dus.id as tiebreaker) - mavvrik_endpoints uses lazy import to keep polars out of startup path 164 mock-based unit tests across 10 files covering all components. Co-Authored-By: Claude Sonnet 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
Praveen Ghuge
parent
4148667671
commit
f78b5be07c
19 changed files with 5366 additions and 1 deletions
|
|
@ -144,6 +144,7 @@ _custom_logger_compatible_callbacks_literal = Literal[
|
|||
"bitbucket",
|
||||
"gitlab",
|
||||
"cloudzero",
|
||||
"mavvrik",
|
||||
"focus",
|
||||
"vantage",
|
||||
"posthog",
|
||||
|
|
|
|||
|
|
@ -1463,6 +1463,11 @@ CLOUDZERO_EXPORT_USAGE_DATA_JOB_NAME = "cloudzero_export_usage_data"
|
|||
CLOUDZERO_MAX_FETCHED_DATA_RECORDS = int(
|
||||
os.getenv("CLOUDZERO_MAX_FETCHED_DATA_RECORDS", 50000)
|
||||
)
|
||||
MAVVRIK_EXPORT_INTERVAL_MINUTES = int(os.getenv("MAVVRIK_EXPORT_INTERVAL_MINUTES", 60))
|
||||
MAVVRIK_EXPORT_USAGE_DATA_JOB_NAME = "mavvrik_export_usage_data"
|
||||
MAVVRIK_MAX_FETCHED_DATA_RECORDS = int(
|
||||
os.getenv("MAVVRIK_MAX_FETCHED_DATA_RECORDS", 50000)
|
||||
)
|
||||
SPEND_LOG_CLEANUP_JOB_NAME = "spend_log_cleanup"
|
||||
KEY_ROTATION_JOB_NAME = "litellm_key_rotation_job"
|
||||
EXPIRED_UI_SESSION_KEY_CLEANUP_JOB_NAME = "litellm_expired_ui_session_key_cleanup_job"
|
||||
|
|
|
|||
436
litellm/integrations/mavvrik/__init__.py
Normal file
436
litellm/integrations/mavvrik/__init__.py
Normal file
|
|
@ -0,0 +1,436 @@
|
|||
"""Mavvrik cost-data integration for LiteLLM.
|
||||
|
||||
Module layout:
|
||||
exporter.py — Exporter (DB queries + DataFrame → CSV transform)
|
||||
uploader.py — Uploader (GCS resumable upload protocol)
|
||||
client.py — Client (Mavvrik REST API calls + retry transport)
|
||||
settings.py — Settings (config detection and persistence)
|
||||
orchestrator.py — Orchestrator (pod lock + register → date loop → upload → advance)
|
||||
|
||||
Public facade:
|
||||
Service — used by mavvrik_endpoints.py; all business logic lives here.
|
||||
"""
|
||||
|
||||
import os
|
||||
from datetime import datetime, timedelta
|
||||
from datetime import timezone as _tz
|
||||
from typing import Optional
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.constants import MAVVRIK_MAX_FETCHED_DATA_RECORDS
|
||||
from litellm.integrations.mavvrik.client import Client
|
||||
from litellm.integrations.mavvrik.exporter import Exporter
|
||||
from litellm.integrations.mavvrik.logger import Logger
|
||||
from litellm.integrations.mavvrik.orchestrator import Orchestrator
|
||||
from litellm.integrations.mavvrik.settings import Settings
|
||||
from litellm.integrations.mavvrik.uploader import Uploader
|
||||
|
||||
__all__ = [
|
||||
"Client",
|
||||
"Exporter",
|
||||
"Logger",
|
||||
"Orchestrator",
|
||||
"Service",
|
||||
"Settings",
|
||||
"Uploader",
|
||||
]
|
||||
|
||||
|
||||
def _build_client(data: dict) -> Client:
|
||||
"""Build a Client from loaded settings dict."""
|
||||
return Client(
|
||||
api_key=str(data.get("api_key") or os.getenv("MAVVRIK_API_KEY", "")),
|
||||
api_endpoint=str(
|
||||
data.get("api_endpoint") or os.getenv("MAVVRIK_API_ENDPOINT", "")
|
||||
),
|
||||
connection_id=str(
|
||||
data.get("connection_id") or os.getenv("MAVVRIK_CONNECTION_ID", "")
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
class Service:
|
||||
"""Public facade that mediates between the REST endpoints and the Mavvrik modules.
|
||||
|
||||
Each method maps 1-to-1 with an endpoint action. All methods return plain
|
||||
``dict`` objects so the router can freely convert them into response models.
|
||||
|
||||
Raises:
|
||||
LookupError — resource not found (router → 404)
|
||||
ValueError — bad input (router → 400)
|
||||
RuntimeError — upstream / integration failure (router → 500)
|
||||
Exception — catch-all (router → 500)
|
||||
"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self._settings = Settings()
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Properties
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
@property
|
||||
def settings(self) -> Settings:
|
||||
return self._settings
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Utilities
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
@staticmethod
|
||||
def _yesterday() -> str:
|
||||
"""Return yesterday's date as YYYY-MM-DD (UTC)."""
|
||||
return (datetime.now(_tz.utc).date() - timedelta(days=1)).isoformat()
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# initialize → POST /mavvrik/init
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
async def initialize(
|
||||
self,
|
||||
api_key: str,
|
||||
api_endpoint: str,
|
||||
connection_id: str,
|
||||
) -> dict:
|
||||
"""Save credentials and schedule the export job.
|
||||
|
||||
Returns:
|
||||
{"message": str, "status": "success"}
|
||||
"""
|
||||
# Step 1 — persist credentials.
|
||||
await self._settings.save(
|
||||
api_key=api_key,
|
||||
api_endpoint=api_endpoint,
|
||||
connection_id=connection_id,
|
||||
)
|
||||
|
||||
# Step 2 — schedule the background export job.
|
||||
from litellm.constants import (
|
||||
MAVVRIK_EXPORT_INTERVAL_MINUTES,
|
||||
MAVVRIK_EXPORT_USAGE_DATA_JOB_NAME,
|
||||
)
|
||||
|
||||
import litellm.proxy.proxy_server as _pserver
|
||||
|
||||
_scheduler = getattr(_pserver, "scheduler", None)
|
||||
if _scheduler is None:
|
||||
verbose_proxy_logger.warning(
|
||||
"mavvrik: scheduler not available, background job not registered"
|
||||
)
|
||||
return {
|
||||
"message": "Mavvrik settings initialized successfully",
|
||||
"status": "success",
|
||||
}
|
||||
|
||||
client = Client(
|
||||
api_key=api_key,
|
||||
api_endpoint=api_endpoint,
|
||||
connection_id=connection_id,
|
||||
)
|
||||
uploader = Uploader(client=client)
|
||||
orchestrator = Orchestrator(client=client, uploader=uploader)
|
||||
# replace_existing=True ensures repeated /mavvrik/init calls are safe.
|
||||
_scheduler.add_job(
|
||||
orchestrator.run,
|
||||
"interval",
|
||||
minutes=MAVVRIK_EXPORT_INTERVAL_MINUTES,
|
||||
id=MAVVRIK_EXPORT_USAGE_DATA_JOB_NAME,
|
||||
replace_existing=True,
|
||||
)
|
||||
verbose_proxy_logger.info(
|
||||
"mavvrik background export job scheduled every %d min",
|
||||
MAVVRIK_EXPORT_INTERVAL_MINUTES,
|
||||
)
|
||||
|
||||
return {
|
||||
"message": "Mavvrik settings initialized successfully",
|
||||
"status": "success",
|
||||
}
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# get_settings → GET /mavvrik/settings
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
async def get_settings(self) -> dict:
|
||||
"""Load and mask Mavvrik settings.
|
||||
|
||||
Falls back to env vars when no DB settings exist.
|
||||
|
||||
Returns:
|
||||
A dict with keys: api_key_masked, api_endpoint, connection_id, status.
|
||||
"""
|
||||
from litellm.litellm_core_utils.sensitive_data_masker import SensitiveDataMasker
|
||||
|
||||
data = await self._settings.load()
|
||||
|
||||
if not data and self._settings.has_env_vars:
|
||||
data = {
|
||||
"api_key": os.getenv("MAVVRIK_API_KEY", ""),
|
||||
"api_endpoint": os.getenv("MAVVRIK_API_ENDPOINT", ""),
|
||||
"connection_id": os.getenv("MAVVRIK_CONNECTION_ID", ""),
|
||||
}
|
||||
|
||||
if not data:
|
||||
return {
|
||||
"api_key_masked": None,
|
||||
"api_endpoint": None,
|
||||
"connection_id": None,
|
||||
"status": "not_configured",
|
||||
}
|
||||
|
||||
masker = SensitiveDataMasker()
|
||||
masked = masker.mask_dict({"api_key": data.get("api_key", "")})
|
||||
return {
|
||||
"api_key_masked": masked.get("api_key"),
|
||||
"api_endpoint": data.get("api_endpoint"),
|
||||
"connection_id": data.get("connection_id"),
|
||||
"status": "configured",
|
||||
}
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# update_settings → PUT /mavvrik/settings
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
async def update_settings(
|
||||
self,
|
||||
api_key: Optional[str] = None,
|
||||
api_endpoint: Optional[str] = None,
|
||||
connection_id: Optional[str] = None,
|
||||
) -> dict:
|
||||
"""Merge new credential values into existing settings and persist.
|
||||
|
||||
Raises:
|
||||
LookupError: when no existing settings are found.
|
||||
ValueError: when a merge would leave a required field empty.
|
||||
"""
|
||||
current = await self._settings.load()
|
||||
|
||||
# Fall back to env vars for env-var-only deployments.
|
||||
if not current and self._settings.has_env_vars:
|
||||
current = {
|
||||
"api_key": os.getenv("MAVVRIK_API_KEY", ""),
|
||||
"api_endpoint": os.getenv("MAVVRIK_API_ENDPOINT", ""),
|
||||
"connection_id": os.getenv("MAVVRIK_CONNECTION_ID", ""),
|
||||
}
|
||||
|
||||
if not current:
|
||||
raise LookupError(
|
||||
"Mavvrik settings not found. Use POST /mavvrik/init to create them first."
|
||||
)
|
||||
|
||||
def _pick(new: Optional[str], key: str) -> str:
|
||||
return new if new is not None else current.get(key, "")
|
||||
|
||||
merged = {
|
||||
"api_key": _pick(api_key, "api_key"),
|
||||
"api_endpoint": _pick(api_endpoint, "api_endpoint"),
|
||||
"connection_id": _pick(connection_id, "connection_id"),
|
||||
}
|
||||
missing = [k for k, v in merged.items() if not v]
|
||||
if missing:
|
||||
raise ValueError(
|
||||
f"Missing required Mavvrik settings after merge: {missing}"
|
||||
)
|
||||
|
||||
await self._settings.save(**merged)
|
||||
|
||||
# Reschedule the background job with the new credentials so the
|
||||
# running Orchestrator uses the merged values immediately —
|
||||
# without this, the in-memory Client keeps old credentials until restart.
|
||||
from litellm.constants import (
|
||||
MAVVRIK_EXPORT_INTERVAL_MINUTES,
|
||||
MAVVRIK_EXPORT_USAGE_DATA_JOB_NAME,
|
||||
)
|
||||
|
||||
import litellm.proxy.proxy_server as _pserver
|
||||
|
||||
_scheduler = getattr(_pserver, "scheduler", None)
|
||||
if _scheduler is not None:
|
||||
client = Client(
|
||||
api_key=merged["api_key"],
|
||||
api_endpoint=merged["api_endpoint"],
|
||||
connection_id=merged["connection_id"],
|
||||
)
|
||||
uploader = Uploader(client=client)
|
||||
orchestrator = Orchestrator(client=client, uploader=uploader)
|
||||
_scheduler.add_job(
|
||||
orchestrator.run,
|
||||
"interval",
|
||||
minutes=MAVVRIK_EXPORT_INTERVAL_MINUTES,
|
||||
id=MAVVRIK_EXPORT_USAGE_DATA_JOB_NAME,
|
||||
replace_existing=True,
|
||||
)
|
||||
verbose_proxy_logger.info(
|
||||
"mavvrik: background job rescheduled with updated credentials"
|
||||
)
|
||||
|
||||
return {"message": "Mavvrik settings updated successfully", "status": "success"}
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# delete → DELETE /mavvrik/delete
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
async def delete(self) -> dict:
|
||||
"""Deregister the scheduler job and remove DB settings if present.
|
||||
|
||||
Works for both env-var-only and DB-backed deployments:
|
||||
- Scheduler job is always deregistered (independent of DB).
|
||||
- DB row is only deleted when credentials come from the DB; env-var
|
||||
deployments have no row so this step is skipped.
|
||||
|
||||
Raises:
|
||||
LookupError: only when DB is connected and no settings row exists.
|
||||
"""
|
||||
from litellm.constants import MAVVRIK_EXPORT_USAGE_DATA_JOB_NAME
|
||||
|
||||
import litellm.proxy.proxy_server as _pserver
|
||||
|
||||
# Deregister scheduler first — independent of whether creds are in DB or env vars.
|
||||
_scheduler = getattr(_pserver, "scheduler", None)
|
||||
if _scheduler is not None:
|
||||
try:
|
||||
_scheduler.remove_job(MAVVRIK_EXPORT_USAGE_DATA_JOB_NAME)
|
||||
except Exception:
|
||||
pass # job may not exist if scheduler was restarted
|
||||
|
||||
# Always attempt DB deletion; silently ignore if no row exists or no DB connected
|
||||
# (env-var-only deployments without a database have no row to remove).
|
||||
try:
|
||||
await self._settings.delete()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
verbose_proxy_logger.info("mavvrik settings deleted")
|
||||
return {"message": "Mavvrik settings deleted successfully", "status": "success"}
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# export → POST /mavvrik/export
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
async def export(
|
||||
self,
|
||||
date_str: Optional[str] = None,
|
||||
limit: Optional[int] = None,
|
||||
) -> dict:
|
||||
"""Fetch spend data and upload to Mavvrik for a calendar date.
|
||||
|
||||
Args:
|
||||
date_str: YYYY-MM-DD. Defaults to yesterday (UTC) when omitted.
|
||||
limit: Cap the number of rows fetched from the database.
|
||||
|
||||
Raises:
|
||||
ValueError: when Mavvrik is not configured.
|
||||
|
||||
Returns:
|
||||
{"message": str, "status": "success", "records_exported": int}
|
||||
"""
|
||||
data = await self._settings.load()
|
||||
|
||||
if not data and not self._settings.has_env_vars:
|
||||
raise ValueError("Mavvrik not configured. Call POST /mavvrik/init first.")
|
||||
|
||||
self._settings._ensure_prisma_client()
|
||||
|
||||
date_str = date_str or self._yesterday()
|
||||
effective_limit = limit or MAVVRIK_MAX_FETCHED_DATA_RECORDS
|
||||
|
||||
client = _build_client(data)
|
||||
uploader = Uploader(client=client)
|
||||
exporter = Exporter()
|
||||
|
||||
df, csv_payload = await exporter.export(
|
||||
date_str=date_str,
|
||||
connection_id=client.connection_id,
|
||||
limit=effective_limit,
|
||||
)
|
||||
|
||||
if df.is_empty():
|
||||
return {
|
||||
"message": f"No data for {date_str}",
|
||||
"status": "success",
|
||||
"records_exported": 0,
|
||||
}
|
||||
|
||||
records_exported = len(df)
|
||||
await uploader.upload(csv_payload, date_str=date_str)
|
||||
|
||||
return {
|
||||
"message": f"Mavvrik export completed successfully for {date_str}",
|
||||
"status": "success",
|
||||
"records_exported": records_exported,
|
||||
}
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# dry_run → POST /mavvrik/dry-run
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
async def dry_run(
|
||||
self,
|
||||
date_str: Optional[str] = None,
|
||||
limit: Optional[int] = None,
|
||||
) -> dict:
|
||||
"""Preview CSV records without uploading.
|
||||
|
||||
Args:
|
||||
date_str: YYYY-MM-DD. Defaults to yesterday (UTC) when omitted.
|
||||
limit: Cap the number of rows fetched from the database.
|
||||
|
||||
Returns:
|
||||
{"message": str, "status": "success", "dry_run_data": dict, "summary": dict}
|
||||
"""
|
||||
data = await self._settings.load()
|
||||
if not data and not self._settings.has_env_vars:
|
||||
raise ValueError("Mavvrik not configured. Call POST /mavvrik/init first.")
|
||||
|
||||
self._settings._ensure_prisma_client()
|
||||
|
||||
date_str = date_str or self._yesterday()
|
||||
effective_limit = limit or MAVVRIK_MAX_FETCHED_DATA_RECORDS
|
||||
|
||||
client = _build_client(data)
|
||||
exporter = Exporter()
|
||||
|
||||
df, csv_payload = await exporter.export(
|
||||
date_str=date_str,
|
||||
connection_id=client.connection_id,
|
||||
limit=effective_limit,
|
||||
)
|
||||
|
||||
if df.is_empty():
|
||||
return {
|
||||
"message": "Mavvrik dry run completed",
|
||||
"status": "success",
|
||||
"dry_run_data": {"usage_data": [], "csv_preview": ""},
|
||||
"summary": {
|
||||
"total_records": 0,
|
||||
"total_cost": 0.0,
|
||||
"total_tokens": 0,
|
||||
"unique_models": 0,
|
||||
"unique_teams": 0,
|
||||
},
|
||||
}
|
||||
|
||||
total_cost = float(df["spend"].sum() or 0.0) if "spend" in df.columns else 0.0
|
||||
total_tokens = (
|
||||
int((df["prompt_tokens"].sum() or 0) + (df["completion_tokens"].sum() or 0))
|
||||
if "prompt_tokens" in df.columns and "completion_tokens" in df.columns
|
||||
else 0
|
||||
)
|
||||
unique_models = df["model"].n_unique() if "model" in df.columns else 0
|
||||
unique_teams = df["team_id"].n_unique() if "team_id" in df.columns else 0
|
||||
|
||||
return {
|
||||
"message": "Mavvrik dry run completed",
|
||||
"status": "success",
|
||||
"dry_run_data": {
|
||||
"usage_data": df.head(50).to_dicts(),
|
||||
"csv_preview": csv_payload[:5000] if csv_payload else "",
|
||||
},
|
||||
"summary": {
|
||||
"total_records": len(df),
|
||||
"total_cost": total_cost,
|
||||
"total_tokens": total_tokens,
|
||||
"unique_models": unique_models,
|
||||
"unique_teams": unique_teams,
|
||||
},
|
||||
}
|
||||
223
litellm/integrations/mavvrik/client.py
Normal file
223
litellm/integrations/mavvrik/client.py
Normal file
|
|
@ -0,0 +1,223 @@
|
|||
"""Mavvrik API client — all HTTP calls to the Mavvrik REST API.
|
||||
|
||||
Responsibility: talk to the Mavvrik API and nothing else.
|
||||
|
||||
Public methods (one per API endpoint):
|
||||
register() POST /metrics/agent/ai/{id} → Optional[str] ISO marker
|
||||
advance_marker() PATCH /metrics/agent/ai/{id} → None
|
||||
report_error() PATCH /metrics/agent/ai/{id} → None (best-effort)
|
||||
get_signed_url() GET /metrics/agent/ai/{id}/upload-url → str
|
||||
|
||||
Transport: uses litellm's shared AsyncHTTPHandler (get_async_httpx_client) with
|
||||
inline retry + exponential backoff for 5xx / network errors.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
from datetime import datetime as _dt
|
||||
from datetime import timezone as _tz
|
||||
from typing import Any, Dict, List, Optional, Union
|
||||
|
||||
import httpx
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.llms.custom_httpx.http_handler import (
|
||||
AsyncHTTPHandler,
|
||||
get_async_httpx_client,
|
||||
)
|
||||
from litellm.types.llms.custom_http import httpxSpecialProvider
|
||||
|
||||
_MAX_RETRIES = 3
|
||||
_RETRY_BACKOFF_BASE = 1.0 # seconds; doubles each retry
|
||||
|
||||
|
||||
class Client:
|
||||
"""HTTP client for the Mavvrik REST API."""
|
||||
|
||||
def __init__(self, api_key: str, api_endpoint: str, connection_id: str) -> None:
|
||||
self._api_key = api_key
|
||||
self._api_endpoint = api_endpoint.rstrip("/")
|
||||
self._connection_id = connection_id
|
||||
self._http: AsyncHTTPHandler = get_async_httpx_client(
|
||||
llm_provider=httpxSpecialProvider.LoggingCallback
|
||||
)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Read-only properties
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
@property
|
||||
def api_key(self) -> str:
|
||||
return self._api_key
|
||||
|
||||
@property
|
||||
def api_endpoint(self) -> str:
|
||||
return self._api_endpoint
|
||||
|
||||
@property
|
||||
def connection_id(self) -> str:
|
||||
return self._connection_id
|
||||
|
||||
@property
|
||||
def agent_url(self) -> str:
|
||||
return f"{self._api_endpoint}/metrics/agent/ai/{self._connection_id}"
|
||||
|
||||
@property
|
||||
def upload_url(self) -> str:
|
||||
return f"{self._api_endpoint}/metrics/agent/ai/{self._connection_id}/upload-url"
|
||||
|
||||
@property
|
||||
def _auth_headers(self) -> Dict[str, str]:
|
||||
return {"Content-Type": "application/json", "x-api-key": self._api_key}
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Public API methods
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
async def register(self) -> Optional[str]:
|
||||
"""POST agent endpoint → return current metricsMarker as ISO-8601 string.
|
||||
|
||||
Returns None when the remote marker is absent or zero (first run).
|
||||
Raises RuntimeError if the call fails.
|
||||
"""
|
||||
body: dict = {"name": self._connection_id}
|
||||
resp = await self._request(
|
||||
"POST", self.agent_url, headers=self._auth_headers, json=body
|
||||
)
|
||||
self._assert_ok(resp, expected={200})
|
||||
|
||||
epoch = resp.json().get("metricsMarker", 0)
|
||||
if not epoch:
|
||||
verbose_proxy_logger.info(
|
||||
"register: no marker (first run), epoch=%s", epoch
|
||||
)
|
||||
return None
|
||||
|
||||
marker_iso = _dt.fromtimestamp(float(epoch), tz=_tz.utc).isoformat()
|
||||
verbose_proxy_logger.info("register: epoch=%s → marker %s", epoch, marker_iso)
|
||||
return marker_iso
|
||||
|
||||
async def advance_marker(self, epoch: int) -> None:
|
||||
"""PATCH agent endpoint to advance the export cursor to the given epoch.
|
||||
|
||||
Raises RuntimeError if the call fails.
|
||||
"""
|
||||
resp = await self._request(
|
||||
"PATCH",
|
||||
self.agent_url,
|
||||
headers=self._auth_headers,
|
||||
json={"metricsMarker": epoch},
|
||||
)
|
||||
self._assert_ok(resp, expected={200, 204})
|
||||
verbose_proxy_logger.info("client: marker advanced to epoch %d", epoch)
|
||||
|
||||
async def report_error(self, error_message: str) -> None:
|
||||
"""PATCH agent endpoint to report an export failure to Mavvrik.
|
||||
|
||||
Best-effort: exceptions are logged and swallowed so a reporting failure
|
||||
never masks the original error.
|
||||
"""
|
||||
try:
|
||||
resp = await self._request(
|
||||
"PATCH",
|
||||
self.agent_url,
|
||||
headers=self._auth_headers,
|
||||
json={"error": error_message[:500]},
|
||||
label="report_error",
|
||||
)
|
||||
self._assert_ok(resp, expected={200, 204})
|
||||
verbose_proxy_logger.debug(
|
||||
"report_error: reported for connection %s", self._connection_id
|
||||
)
|
||||
except Exception as exc:
|
||||
verbose_proxy_logger.warning("report_error failed (non-fatal): %s", exc)
|
||||
|
||||
async def get_signed_url(self, date_str: str) -> str:
|
||||
"""GET upload-url endpoint → return the GCS signed URL for the given date.
|
||||
|
||||
Raises RuntimeError if the call fails or the response is missing the URL.
|
||||
"""
|
||||
params = {"name": date_str, "type": "metrics", "datetime": date_str}
|
||||
resp = await self._request(
|
||||
"GET", self.upload_url, headers=self._auth_headers, params=params
|
||||
)
|
||||
self._assert_ok(resp, expected={200})
|
||||
|
||||
signed_url = resp.json().get("url")
|
||||
if not signed_url:
|
||||
raise RuntimeError(
|
||||
f"Mavvrik API response missing 'url' field: {resp.json()}"
|
||||
)
|
||||
|
||||
verbose_proxy_logger.debug("client: got signed URL for date %s", date_str)
|
||||
return signed_url
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Transport — litellm shared HTTP handler with retry + backoff
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
async def _request(
|
||||
self,
|
||||
method: str,
|
||||
url: str,
|
||||
*,
|
||||
headers: Optional[Dict[str, str]] = None,
|
||||
json: Optional[Any] = None,
|
||||
params: Optional[Dict[str, str]] = None,
|
||||
content: Optional[bytes] = None,
|
||||
timeout: float = 30.0,
|
||||
label: str = "",
|
||||
) -> httpx.Response:
|
||||
"""Execute a Mavvrik API request with retry and exponential backoff.
|
||||
|
||||
Uses the litellm shared httpx client cache to avoid per-request
|
||||
connection overhead. Retries up to _MAX_RETRIES times on 5xx or
|
||||
network errors; returns immediately on 4xx.
|
||||
"""
|
||||
tag = label or method
|
||||
last_exc: Exception = RuntimeError("unknown error")
|
||||
|
||||
for attempt in range(_MAX_RETRIES):
|
||||
try:
|
||||
resp = await self._http.client.request(
|
||||
method=method,
|
||||
url=url,
|
||||
headers=headers,
|
||||
json=json,
|
||||
params=params,
|
||||
content=content,
|
||||
timeout=timeout,
|
||||
)
|
||||
if resp.status_code < 500:
|
||||
return resp
|
||||
last_exc = RuntimeError(
|
||||
f"{tag} failed: {resp.status_code} {resp.text[:200]}"
|
||||
)
|
||||
except httpx.RequestError as exc:
|
||||
last_exc = exc
|
||||
|
||||
if attempt < _MAX_RETRIES - 1:
|
||||
wait = _RETRY_BACKOFF_BASE * (2**attempt)
|
||||
verbose_proxy_logger.warning(
|
||||
"mavvrik: %s attempt %d/%d failed, retrying in %.1fs: %s",
|
||||
tag,
|
||||
attempt + 1,
|
||||
_MAX_RETRIES,
|
||||
wait,
|
||||
last_exc,
|
||||
)
|
||||
await asyncio.sleep(wait)
|
||||
|
||||
raise RuntimeError(
|
||||
f"mavvrik: {tag} failed after {_MAX_RETRIES} attempts: {last_exc}"
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _assert_ok(
|
||||
resp: httpx.Response,
|
||||
expected: Union[set, List[int]],
|
||||
) -> None:
|
||||
"""Raise RuntimeError when the response status is not in expected."""
|
||||
if resp.status_code not in expected:
|
||||
raise RuntimeError(
|
||||
f"unexpected status {resp.status_code}: {resp.text[:200]}"
|
||||
)
|
||||
232
litellm/integrations/mavvrik/exporter.py
Normal file
232
litellm/integrations/mavvrik/exporter.py
Normal file
|
|
@ -0,0 +1,232 @@
|
|||
"""Exporter — fetch spend data from Postgres and transform to CSV.
|
||||
|
||||
Responsibility: extract data from LiteLLM's database and convert it to CSV.
|
||||
|
||||
Public interface:
|
||||
export(date_str, connection_id, limit) → (DataFrame, csv_str)
|
||||
Single entry point: fetch → serialize. Used by Service.export/dry_run.
|
||||
|
||||
get_earliest_date() → Optional[str]
|
||||
Returns MIN(date) for first-run start date resolution.
|
||||
|
||||
Internal methods:
|
||||
_stream_pages(date_str, connection_id, page_size) → AsyncGenerator[str, None]
|
||||
_get_usage_data(date_str, limit) → DataFrame
|
||||
_to_csv(df, connection_id) → str
|
||||
|
||||
DB not connected:
|
||||
_get_usage_data / _to_csv — log warning and return empty/None (scheduler path).
|
||||
_stream_pages — raises RuntimeError (propagates to Orchestrator try/except).
|
||||
Service.export / dry_run — call Settings._ensure_prisma_client() before reaching here,
|
||||
so they raise before the exporter is called.
|
||||
|
||||
polars is an optional [proxy] dependency — imported lazily inside methods so
|
||||
SDK-only users are not affected when Logger is imported via custom_logger_registry.
|
||||
"""
|
||||
|
||||
import io
|
||||
from typing import (
|
||||
Any,
|
||||
AsyncGenerator,
|
||||
List,
|
||||
Optional,
|
||||
Tuple,
|
||||
TYPE_CHECKING,
|
||||
)
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
|
||||
if TYPE_CHECKING:
|
||||
import polars as pl
|
||||
|
||||
# query_raw is used here instead of Prisma model methods because the query
|
||||
# requires a 4-table LEFT JOIN (DailyUserSpend → VerificationToken →
|
||||
# TeamTable → UserTable). Prisma's relational API cannot express a multi-hop
|
||||
# JOIN in a single query without N+1 round-trips.
|
||||
#
|
||||
# dus.* selects all columns from LiteLLM_DailyUserSpend so that any new
|
||||
# columns added to that table in future LiteLLM versions are automatically
|
||||
# included in the export without requiring a code change here.
|
||||
_USAGE_QUERY = """
|
||||
SELECT
|
||||
dus.*,
|
||||
vt.team_id,
|
||||
vt.key_alias AS api_key_alias,
|
||||
vt.organization_id,
|
||||
tt.team_alias,
|
||||
ut.user_email,
|
||||
ut.user_alias
|
||||
FROM "LiteLLM_DailyUserSpend" dus
|
||||
LEFT JOIN "LiteLLM_VerificationToken" vt ON dus.api_key = vt.token
|
||||
LEFT JOIN "LiteLLM_TeamTable" tt ON vt.team_id = tt.team_id
|
||||
LEFT JOIN "LiteLLM_UserTable" ut ON dus.user_id = ut.user_id
|
||||
WHERE dus.date = $1
|
||||
ORDER BY dus.date, dus.user_id, dus.api_key, dus.model, dus.id ASC
|
||||
"""
|
||||
|
||||
_EARLIEST_DATE_QUERY = 'SELECT MIN(date) AS earliest FROM "LiteLLM_DailyUserSpend"'
|
||||
|
||||
|
||||
class Exporter:
|
||||
"""Fetch LiteLLM spend data from Postgres and transform to CSV."""
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# DB access helper — returns None when DB not connected (never raises)
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
@property
|
||||
def _prisma_client(self):
|
||||
try:
|
||||
from litellm.proxy.proxy_server import prisma_client
|
||||
|
||||
return prisma_client # may be None if DB not yet connected
|
||||
except ImportError:
|
||||
return None
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Public interface
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
async def export(
|
||||
self,
|
||||
date_str: str,
|
||||
connection_id: Optional[str] = None,
|
||||
limit: Optional[int] = None,
|
||||
) -> Tuple["pl.DataFrame", str]:
|
||||
"""Fetch and serialize spend data for one calendar date.
|
||||
|
||||
All rows are exported — including failed requests. Mavvrik decides
|
||||
what to do with them on the ingestion side.
|
||||
|
||||
Returns (df, csv_str). Returns (empty DataFrame, "") when no data or no DB.
|
||||
"""
|
||||
df = await self._get_usage_data(date_str=date_str, limit=limit)
|
||||
csv = self._to_csv(df, connection_id=connection_id)
|
||||
return df, csv
|
||||
|
||||
async def _stream_pages(
|
||||
self,
|
||||
date_str: str,
|
||||
connection_id: Optional[str] = None,
|
||||
page_size: int = 10_000,
|
||||
) -> AsyncGenerator[str, None]:
|
||||
"""Yield CSV text in pages — one page of rows at a time (header on first page).
|
||||
|
||||
Uses LIMIT/OFFSET pagination so only page_size rows are in memory at once.
|
||||
All rows exported — including failed requests.
|
||||
|
||||
When DB is not connected, raises RuntimeError so the caller (Orchestrator)
|
||||
knows the export failed — distinct from a legitimate zero-traffic day which
|
||||
yields nothing without raising.
|
||||
"""
|
||||
import polars as pl
|
||||
|
||||
client = self._prisma_client
|
||||
if client is None:
|
||||
raise RuntimeError(
|
||||
"Exporter: database not connected — cannot stream pages for "
|
||||
f"{date_str}. Connect a database to your proxy."
|
||||
)
|
||||
|
||||
header_written = False
|
||||
offset = 0
|
||||
|
||||
while True:
|
||||
rows = await client.db.query_raw(
|
||||
_USAGE_QUERY + " LIMIT $2 OFFSET $3",
|
||||
date_str,
|
||||
page_size,
|
||||
offset,
|
||||
)
|
||||
|
||||
if not rows:
|
||||
break # legitimate end — no more rows (or zero-traffic day)
|
||||
|
||||
df = pl.DataFrame(rows, infer_schema_length=None)
|
||||
|
||||
buf = io.StringIO()
|
||||
if not header_written:
|
||||
if connection_id:
|
||||
df = df.with_columns(pl.lit(connection_id).alias("connection_id"))
|
||||
df.write_csv(buf)
|
||||
header_written = True
|
||||
else:
|
||||
if connection_id:
|
||||
df = df.with_columns(pl.lit(connection_id).alias("connection_id"))
|
||||
df.write_csv(buf, include_header=False)
|
||||
|
||||
yield buf.getvalue()
|
||||
|
||||
offset += page_size
|
||||
if len(rows) < page_size:
|
||||
break
|
||||
|
||||
async def get_earliest_date(self) -> Optional[str]:
|
||||
"""Return MIN(date) from LiteLLM_DailyUserSpend, or None.
|
||||
|
||||
Returns None when DB is not connected — caller treats it as "no history".
|
||||
"""
|
||||
client = self._prisma_client
|
||||
if client is None:
|
||||
verbose_proxy_logger.warning(
|
||||
"Exporter: database not connected, cannot determine earliest date"
|
||||
)
|
||||
return None
|
||||
|
||||
rows = await client.db.query_raw(_EARLIEST_DATE_QUERY)
|
||||
if rows and rows[0].get("earliest") is not None:
|
||||
return str(rows[0]["earliest"])[:10]
|
||||
return None
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Internal methods
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
async def _get_usage_data(
|
||||
self,
|
||||
date_str: str,
|
||||
limit: Optional[int] = None,
|
||||
) -> "pl.DataFrame":
|
||||
"""Retrieve all spend rows for a single calendar date.
|
||||
|
||||
Returns empty DataFrame when DB is not connected.
|
||||
"""
|
||||
import polars as pl
|
||||
|
||||
client = self._prisma_client
|
||||
if client is None:
|
||||
verbose_proxy_logger.warning(
|
||||
"Exporter: database not connected, returning empty data for %s",
|
||||
date_str,
|
||||
)
|
||||
return pl.DataFrame()
|
||||
|
||||
query = _USAGE_QUERY
|
||||
params: List[Any] = [date_str]
|
||||
|
||||
if limit is not None:
|
||||
params.append(int(limit))
|
||||
query += " LIMIT $2"
|
||||
|
||||
db_response = await client.db.query_raw(query, *params)
|
||||
return pl.DataFrame(db_response, infer_schema_length=None)
|
||||
|
||||
def _to_csv(self, df: "pl.DataFrame", connection_id: Optional[str] = None) -> str:
|
||||
"""Serialize a DataFrame to CSV, adding connection_id column if provided."""
|
||||
import polars as pl
|
||||
|
||||
if df.is_empty():
|
||||
verbose_proxy_logger.debug("Exporter: empty DataFrame, nothing to export")
|
||||
return ""
|
||||
|
||||
if connection_id:
|
||||
df = df.with_columns(pl.lit(connection_id).alias("connection_id"))
|
||||
|
||||
buf = io.StringIO()
|
||||
df.write_csv(buf)
|
||||
csv_str = buf.getvalue()
|
||||
|
||||
verbose_proxy_logger.debug(
|
||||
"Exporter: %d rows → %d CSV bytes", len(df), len(csv_str)
|
||||
)
|
||||
return csv_str
|
||||
16
litellm/integrations/mavvrik/logger.py
Normal file
16
litellm/integrations/mavvrik/logger.py
Normal file
|
|
@ -0,0 +1,16 @@
|
|||
"""Mavvrik callback logger — registered as the "mavvrik" callback string.
|
||||
|
||||
This class is the entry point for callbacks: ["mavvrik"] in config.yaml.
|
||||
It acts as a marker so LiteLLM recognises "mavvrik" as a known integration.
|
||||
|
||||
The actual export work (query → CSV → upload) is done by the scheduler and
|
||||
orchestrator, not on a per-request basis. This class is intentionally empty.
|
||||
"""
|
||||
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
|
||||
|
||||
class Logger(CustomLogger):
|
||||
"""Mavvrik integration marker — registered via callbacks: ["mavvrik"]."""
|
||||
|
||||
pass
|
||||
191
litellm/integrations/mavvrik/orchestrator.py
Normal file
191
litellm/integrations/mavvrik/orchestrator.py
Normal file
|
|
@ -0,0 +1,191 @@
|
|||
"""Orchestrator — pipeline sequencing: register → export → upload → advance.
|
||||
|
||||
Responsibility: sequence the export steps and own the pod lock. Nothing else.
|
||||
|
||||
Pipeline in _run_pipeline (one line per step):
|
||||
start, end = await self._register(), self._export_end_date()
|
||||
for export_date in self._date_range(start, end):
|
||||
await self._export(export_date) # streams DB → GCS via _stream_pages/_stream_upload
|
||||
await self._advance(export_date)
|
||||
|
||||
One try/except in _run_pipeline. No nested exception handling anywhere else.
|
||||
|
||||
Marker semantics:
|
||||
metricsMarker from register() is the START of the export window.
|
||||
After each date is uploaded, advance_marker() is called with
|
||||
(export_date + 1 day) so the next run starts from there.
|
||||
"""
|
||||
|
||||
from datetime import date, datetime, timedelta, timezone
|
||||
from typing import Iterator
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.constants import MAVVRIK_EXPORT_USAGE_DATA_JOB_NAME
|
||||
from litellm.integrations.mavvrik.client import Client
|
||||
from litellm.integrations.mavvrik.exporter import Exporter
|
||||
from litellm.integrations.mavvrik.uploader import Uploader
|
||||
|
||||
|
||||
class Orchestrator:
|
||||
"""Sequences the incremental Mavvrik export pipeline."""
|
||||
|
||||
def __init__(self, client: Client, uploader: Uploader) -> None:
|
||||
self._client = client
|
||||
self._uploader = uploader
|
||||
self._exporter = Exporter()
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Date helpers
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
@staticmethod
|
||||
def _utc_today() -> date:
|
||||
return datetime.now(timezone.utc).date()
|
||||
|
||||
def _export_end_date(self) -> date:
|
||||
"""Last date eligible for export (yesterday UTC — today's data is incomplete)."""
|
||||
return self._utc_today() - timedelta(days=1)
|
||||
|
||||
def _date_range(self, start: date, end: date) -> Iterator[date]:
|
||||
"""Yield each date from start to end (inclusive)."""
|
||||
current = start
|
||||
while current <= end:
|
||||
yield current
|
||||
current += timedelta(days=1)
|
||||
|
||||
@staticmethod
|
||||
def _to_epoch(d: date) -> int:
|
||||
return int(datetime(d.year, d.month, d.day, tzinfo=timezone.utc).timestamp())
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Entry point (called by APScheduler)
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
async def run(self) -> None:
|
||||
"""Acquire pod lock then run the export pipeline."""
|
||||
pod_lock = self._get_pod_lock_manager()
|
||||
|
||||
if not pod_lock or not pod_lock.redis_cache:
|
||||
await self._run_pipeline()
|
||||
return
|
||||
|
||||
if not await pod_lock.acquire_lock(
|
||||
cronjob_id=MAVVRIK_EXPORT_USAGE_DATA_JOB_NAME
|
||||
):
|
||||
verbose_logger.debug(
|
||||
"Orchestrator: pod lock not acquired — another pod is running"
|
||||
)
|
||||
return
|
||||
|
||||
try:
|
||||
await self._run_pipeline()
|
||||
finally:
|
||||
await pod_lock.release_lock(cronjob_id=MAVVRIK_EXPORT_USAGE_DATA_JOB_NAME)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Pipeline — one try/except, one line per step
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
async def _run_pipeline(self) -> None:
|
||||
try:
|
||||
start = await self._register()
|
||||
end = self._export_end_date()
|
||||
|
||||
if start > end:
|
||||
verbose_logger.info(
|
||||
"Orchestrator: up to date (start=%s, end=%s), nothing to export",
|
||||
start,
|
||||
end,
|
||||
)
|
||||
return
|
||||
|
||||
verbose_logger.info("Orchestrator: exporting %s → %s", start, end)
|
||||
|
||||
for export_date in self._date_range(start, end):
|
||||
# _export raises if the DB is unavailable (propagates through
|
||||
# _stream_pages → _stream_upload). A genuine zero-traffic day
|
||||
# returns 0 bytes without raising — advance the marker normally.
|
||||
await self._export(export_date)
|
||||
await self._advance(export_date)
|
||||
|
||||
verbose_logger.info("Orchestrator: export complete, last date=%s", end)
|
||||
|
||||
except Exception as exc:
|
||||
verbose_logger.error(
|
||||
"Orchestrator: pipeline failed: %s", exc, exc_info=True
|
||||
)
|
||||
# Use the exception type name only — avoid forwarding raw exception
|
||||
# text which may contain internal hostnames or DB DSN fragments.
|
||||
await self._client.report_error(type(exc).__name__)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Pipeline steps
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
async def _register(self) -> date:
|
||||
marker_str = await self._client.register()
|
||||
verbose_logger.info("Orchestrator: marker from Mavvrik API = %s", marker_str)
|
||||
|
||||
if marker_str:
|
||||
return date.fromisoformat(marker_str[:10])
|
||||
|
||||
return await self._resolve_first_run_start_date()
|
||||
|
||||
async def _export(self, export_date: date) -> int:
|
||||
"""Stream spend data from DB to GCS for one date.
|
||||
|
||||
Uses Exporter._stream_pages() → Uploader._stream_upload() so only
|
||||
one page of rows is in memory at a time. No row limit or overflow check.
|
||||
|
||||
Returns total compressed bytes uploaded (0 when no data for the date).
|
||||
"""
|
||||
date_str = export_date.isoformat()
|
||||
pages = self._exporter._stream_pages(
|
||||
date_str=date_str,
|
||||
connection_id=self._client.connection_id,
|
||||
)
|
||||
total_bytes = await self._uploader._stream_upload(pages, date_str=date_str)
|
||||
if total_bytes > 0:
|
||||
verbose_logger.info(
|
||||
"Orchestrator: %s → streamed %d bytes to GCS ✓", date_str, total_bytes
|
||||
)
|
||||
else:
|
||||
verbose_logger.info("Orchestrator: %s → no data, skipped", date_str)
|
||||
return total_bytes
|
||||
|
||||
async def _advance(self, export_date: date) -> None:
|
||||
next_date = export_date + timedelta(days=1)
|
||||
await self._client.advance_marker(self._to_epoch(next_date))
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# First-run start date
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
async def _resolve_first_run_start_date(self) -> date:
|
||||
earliest_str = await self._exporter.get_earliest_date()
|
||||
if earliest_str:
|
||||
try:
|
||||
earliest = date.fromisoformat(earliest_str)
|
||||
verbose_logger.info(
|
||||
"Orchestrator: no marker, starting from earliest DB date %s",
|
||||
earliest,
|
||||
)
|
||||
return earliest
|
||||
except ValueError:
|
||||
pass
|
||||
return self._export_end_date()
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Infrastructure
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
@staticmethod
|
||||
def _get_pod_lock_manager():
|
||||
from litellm.proxy.proxy_server import proxy_logging_obj
|
||||
|
||||
if proxy_logging_obj is None:
|
||||
return None
|
||||
writer = getattr(proxy_logging_obj, "db_spend_update_writer", None)
|
||||
if writer is None:
|
||||
return None
|
||||
return getattr(writer, "pod_lock_manager", None)
|
||||
220
litellm/integrations/mavvrik/settings.py
Normal file
220
litellm/integrations/mavvrik/settings.py
Normal file
|
|
@ -0,0 +1,220 @@
|
|||
"""Settings management for the Mavvrik integration.
|
||||
|
||||
Consolidates all configuration concerns:
|
||||
- Config detection (env vars or database)
|
||||
- Persistence (load/save/delete via LiteLLM_Config table)
|
||||
- Encryption/decryption of the API key
|
||||
|
||||
The export marker (cursor) is owned exclusively by the Mavvrik API.
|
||||
On each scheduled run, MavvrikOrchestrator calls client.register() to
|
||||
retrieve the current metricsMarker from Mavvrik — no local marker
|
||||
storage is needed.
|
||||
"""
|
||||
|
||||
import json
|
||||
import os
|
||||
from typing import Optional
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
|
||||
_CONFIG_KEY = "mavvrik_settings"
|
||||
|
||||
_ENV_VARS = (
|
||||
"MAVVRIK_API_KEY",
|
||||
"MAVVRIK_API_ENDPOINT",
|
||||
"MAVVRIK_CONNECTION_ID",
|
||||
)
|
||||
|
||||
|
||||
class Settings:
|
||||
"""Manages Mavvrik configuration: detection, persistence, and encryption.
|
||||
|
||||
Usage::
|
||||
|
||||
settings = Settings()
|
||||
if await settings.is_setup():
|
||||
data = await settings.load() # api_key already decrypted
|
||||
"""
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Properties
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
@property
|
||||
def config_key(self) -> str:
|
||||
"""The LiteLLM_Config row key used to store Mavvrik settings."""
|
||||
return _CONFIG_KEY
|
||||
|
||||
@property
|
||||
def has_env_vars(self) -> bool:
|
||||
"""Return True when all three required env vars are non-empty."""
|
||||
return all(os.getenv(v, "").strip() for v in _ENV_VARS)
|
||||
|
||||
@property
|
||||
def _prisma_client(self):
|
||||
"""Lazy import of the prisma_client singleton.
|
||||
|
||||
Returns None when the proxy database is not connected (e.g. in tests
|
||||
or when running without a database backend).
|
||||
"""
|
||||
try:
|
||||
from litellm.proxy.proxy_server import prisma_client
|
||||
|
||||
return prisma_client
|
||||
except ImportError:
|
||||
return None
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Setup detection
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
async def is_setup(self) -> bool:
|
||||
"""Return True if Mavvrik credentials exist in env vars or the database."""
|
||||
if self.has_env_vars:
|
||||
return True
|
||||
|
||||
client = self._prisma_client
|
||||
if client is None:
|
||||
return False
|
||||
|
||||
try:
|
||||
row = await client.db.litellm_config.find_first(
|
||||
where={"param_name": _CONFIG_KEY}
|
||||
)
|
||||
return row is not None and row.param_value is not None
|
||||
except Exception as exc:
|
||||
verbose_logger.debug("Settings.is_setup: DB check failed — %s", exc)
|
||||
return False
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Load / Save / Delete
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
async def load(self) -> dict:
|
||||
"""Load and decrypt Mavvrik settings from the database.
|
||||
|
||||
Returns an empty dict when no row exists or when the database is not
|
||||
connected. Callers fall back to env vars when this returns {}.
|
||||
The ``api_key`` field is returned in plaintext (decrypted).
|
||||
"""
|
||||
client = self._prisma_client
|
||||
if client is None:
|
||||
return {}
|
||||
|
||||
row = await client.db.litellm_config.find_first(
|
||||
where={"param_name": _CONFIG_KEY}
|
||||
)
|
||||
if row is None or row.param_value is None:
|
||||
return {}
|
||||
|
||||
value = row.param_value
|
||||
if isinstance(value, str):
|
||||
try:
|
||||
value = json.loads(value)
|
||||
except json.JSONDecodeError:
|
||||
return {}
|
||||
|
||||
if not isinstance(value, dict):
|
||||
return {}
|
||||
|
||||
encrypted_key: Optional[str] = value.get("api_key")
|
||||
if not encrypted_key:
|
||||
return value
|
||||
|
||||
decrypted = self.decrypt_value_helper(encrypted_key, key="mavvrik_api_key")
|
||||
if decrypted is None:
|
||||
raise ValueError(
|
||||
"Failed to decrypt stored Mavvrik API key — possible salt/master key mismatch. "
|
||||
"Re-initialize via POST /mavvrik/init to store credentials with the current key."
|
||||
)
|
||||
|
||||
value["api_key"] = decrypted
|
||||
return value
|
||||
|
||||
async def save(
|
||||
self,
|
||||
api_key: str,
|
||||
api_endpoint: str,
|
||||
connection_id: str,
|
||||
) -> None:
|
||||
"""Encrypt the API key and persist credentials to LiteLLM_Config.
|
||||
|
||||
The export marker (cursor) is owned exclusively by the Mavvrik API
|
||||
and is NOT stored locally — it is retrieved via client.register()
|
||||
at the start of each scheduled run.
|
||||
"""
|
||||
encrypted_api_key: str = self.encrypt_value_helper(api_key)
|
||||
settings: dict = {
|
||||
"api_key": encrypted_api_key,
|
||||
"api_endpoint": api_endpoint,
|
||||
"connection_id": connection_id,
|
||||
}
|
||||
await self._upsert(settings)
|
||||
|
||||
async def delete(self) -> None:
|
||||
"""Remove the Mavvrik settings row from LiteLLM_Config.
|
||||
|
||||
Raises:
|
||||
LookupError: When no Mavvrik settings row exists in the database.
|
||||
"""
|
||||
client = self._ensure_prisma_client()
|
||||
|
||||
row = await client.db.litellm_config.find_first(
|
||||
where={"param_name": _CONFIG_KEY}
|
||||
)
|
||||
if row is None or row.param_value is None:
|
||||
raise LookupError("Mavvrik settings not found — nothing to delete.")
|
||||
|
||||
await client.db.litellm_config.delete(where={"param_name": _CONFIG_KEY})
|
||||
verbose_logger.info("Settings: settings row deleted")
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Encryption helpers (owned here so callers never touch utils directly)
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def encrypt_value_helper(self, value: str) -> str:
|
||||
"""Encrypt a plaintext string using the LiteLLM salt key."""
|
||||
from litellm.proxy.common_utils.encrypt_decrypt_utils import (
|
||||
encrypt_value_helper as _encrypt,
|
||||
)
|
||||
|
||||
return _encrypt(value)
|
||||
|
||||
def decrypt_value_helper(
|
||||
self, value: str, key: str = "mavvrik_api_key"
|
||||
) -> Optional[str]:
|
||||
"""Decrypt an encrypted string using the LiteLLM salt key.
|
||||
|
||||
Returns None when decryption fails (e.g. salt key mismatch).
|
||||
"""
|
||||
from litellm.proxy.common_utils.encrypt_decrypt_utils import (
|
||||
decrypt_value_helper as _decrypt,
|
||||
)
|
||||
|
||||
return _decrypt(value, key=key)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Internal helpers
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def _ensure_prisma_client(self):
|
||||
"""Return the prisma_client or raise if the database is not connected."""
|
||||
client = self._prisma_client
|
||||
if client is None:
|
||||
raise Exception(
|
||||
"Database not connected. Connect a database to your proxy — "
|
||||
"https://docs.litellm.ai/docs/simple_proxy#managing-auth---virtual-keys"
|
||||
)
|
||||
return client
|
||||
|
||||
async def _upsert(self, settings: dict) -> None:
|
||||
"""Write (create or update) the settings row in LiteLLM_Config."""
|
||||
client = self._ensure_prisma_client()
|
||||
payload = json.dumps(settings)
|
||||
await client.db.litellm_config.upsert(
|
||||
where={"param_name": _CONFIG_KEY},
|
||||
data={
|
||||
"create": {"param_name": _CONFIG_KEY, "param_value": payload},
|
||||
"update": {"param_value": payload},
|
||||
},
|
||||
)
|
||||
254
litellm/integrations/mavvrik/uploader.py
Normal file
254
litellm/integrations/mavvrik/uploader.py
Normal file
|
|
@ -0,0 +1,254 @@
|
|||
"""Uploader — GCS resumable upload protocol.
|
||||
|
||||
Responsibility: receive a CSV string and upload it to GCS. Nothing else.
|
||||
|
||||
Upload flow:
|
||||
1. Compress — gzip the CSV string → bytes
|
||||
2. Signed URL — GET from Mavvrik API via Client.get_signed_url()
|
||||
3. Initiate — POST to signed URL → GCS session URI (Location header)
|
||||
4. Finalize — PUT gzip bytes to session URI → upload complete
|
||||
|
||||
Steps 3 and 4 talk directly to GCS (no Mavvrik auth header).
|
||||
Step 2 is delegated to Client which owns all Mavvrik API calls.
|
||||
|
||||
Transport layer (shared by all GCS steps):
|
||||
litellm's shared AsyncHTTPHandler (.client.request()) — used by
|
||||
_initiate_resumable_upload, _finalize_upload, _put_chunk.
|
||||
|
||||
GCS resumable upload protocol reference:
|
||||
https://cloud.google.com/storage/docs/resumable-uploads
|
||||
"""
|
||||
|
||||
import contextlib
|
||||
import gzip
|
||||
import io
|
||||
from typing import TYPE_CHECKING, Any, AsyncIterator
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.llms.custom_httpx.http_handler import get_async_httpx_client
|
||||
from litellm.types.llms.custom_http import httpxSpecialProvider
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.integrations.mavvrik.client import Client
|
||||
else:
|
||||
Client = Any
|
||||
|
||||
# GCS requires intermediate chunks to be exactly this size (256 KB aligned).
|
||||
# Only the final chunk can be smaller.
|
||||
_GCS_CHUNK_SIZE = 256 * 1024
|
||||
|
||||
|
||||
class Uploader:
|
||||
"""Upload gzip-compressed CSV data to GCS via the resumable upload protocol."""
|
||||
|
||||
def __init__(self, client: "Client") -> None:
|
||||
self._client = client
|
||||
self._http = get_async_httpx_client(
|
||||
llm_provider=httpxSpecialProvider.LoggingCallback
|
||||
)
|
||||
|
||||
@property
|
||||
def client(self) -> "Client":
|
||||
return self._client
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Public interface
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
async def upload(self, csv_payload: str, date_str: str) -> None:
|
||||
"""Compress and upload a CSV string to GCS for the given date.
|
||||
|
||||
Re-uploading the same date overwrites the previous object — idempotent.
|
||||
|
||||
Args:
|
||||
csv_payload: CSV string (header + rows).
|
||||
date_str: Date in YYYY-MM-DD format.
|
||||
|
||||
Raises:
|
||||
RuntimeError: if any upload step fails after retries.
|
||||
"""
|
||||
if not csv_payload.strip():
|
||||
verbose_proxy_logger.debug("uploader: empty payload, skipping upload")
|
||||
return
|
||||
|
||||
gzip_bytes = self._compress(csv_payload)
|
||||
signed_url = await self._client.get_signed_url(date_str)
|
||||
session_uri = await self._initiate_resumable_upload(signed_url)
|
||||
await self._finalize_upload(session_uri, gzip_bytes)
|
||||
|
||||
verbose_proxy_logger.info(
|
||||
"uploader: uploaded %d bytes for date %s", len(gzip_bytes), date_str
|
||||
)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# GCS protocol steps
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
async def _initiate_resumable_upload(self, signed_url: str) -> str:
|
||||
"""POST to the GCS signed URL to open a resumable upload session.
|
||||
|
||||
Returns the session URI from the Location response header.
|
||||
"""
|
||||
metadata = b'{"contentEncoding":"gzip","contentDisposition":"attachment"}'
|
||||
resp = await self._http.client.request(
|
||||
method="POST",
|
||||
url=signed_url,
|
||||
headers={"Content-Type": "application/gzip", "x-goog-resumable": "start"},
|
||||
content=metadata,
|
||||
timeout=30.0,
|
||||
)
|
||||
if resp.status_code not in (200, 201):
|
||||
raise RuntimeError(
|
||||
f"GCS initiate upload failed: {resp.status_code} {resp.text[:200]}"
|
||||
)
|
||||
session_uri = resp.headers.get("Location")
|
||||
if not session_uri:
|
||||
raise RuntimeError("GCS initiate upload response missing Location header")
|
||||
return session_uri
|
||||
|
||||
async def _finalize_upload(self, session_uri: str, gzip_bytes: bytes) -> None:
|
||||
"""PUT gzip bytes to the GCS session URI to complete the bulk upload.
|
||||
|
||||
Sends Content-Range: bytes 0-{last}/{total} per the GCS resumable upload
|
||||
spec so GCS knows the object is complete (consistent with _put_chunk final=True).
|
||||
"""
|
||||
total = len(gzip_bytes)
|
||||
content_range = f"bytes 0-{total - 1}/{total}" if total > 0 else "bytes */0"
|
||||
resp = await self._http.client.request(
|
||||
method="PUT",
|
||||
url=session_uri,
|
||||
headers={
|
||||
"Content-Type": "application/gzip",
|
||||
"Content-Range": content_range,
|
||||
},
|
||||
content=gzip_bytes,
|
||||
timeout=120.0,
|
||||
)
|
||||
if resp.status_code not in (200, 201):
|
||||
raise RuntimeError(
|
||||
f"GCS finalize upload failed: {resp.status_code} {resp.text[:200]}"
|
||||
)
|
||||
verbose_proxy_logger.debug("uploader: finalize OK (%d)", resp.status_code)
|
||||
|
||||
async def _put_chunk(
|
||||
self,
|
||||
session_uri: str,
|
||||
chunk: bytes,
|
||||
offset: int,
|
||||
final: bool,
|
||||
) -> None:
|
||||
"""PUT one chunk to the GCS resumable session URI.
|
||||
|
||||
Intermediate chunks: Content-Range: bytes X-Y/* → expect 308
|
||||
Final chunk: Content-Range: bytes X-Y/T → expect 200/201
|
||||
"""
|
||||
end = offset + len(chunk) - 1
|
||||
total_str = str(offset + len(chunk)) if final else "*"
|
||||
content_range = f"bytes {offset}-{end}/{total_str}"
|
||||
expected = {200, 201} if final else {308}
|
||||
|
||||
resp = await self._http.client.request(
|
||||
method="PUT",
|
||||
url=session_uri,
|
||||
headers={
|
||||
"Content-Type": "application/gzip",
|
||||
"Content-Range": content_range,
|
||||
},
|
||||
content=chunk,
|
||||
timeout=120.0,
|
||||
)
|
||||
if resp.status_code not in expected:
|
||||
raise RuntimeError(
|
||||
f"GCS PUT chunk failed: {resp.status_code} "
|
||||
f"(expected {expected}): {resp.text[:200]}"
|
||||
)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Streaming upload
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
async def _stream_upload(
|
||||
self,
|
||||
pages: AsyncIterator[str],
|
||||
date_str: str,
|
||||
) -> int:
|
||||
"""Stream CSV pages to GCS using chunked resumable upload.
|
||||
|
||||
Each intermediate chunk is exactly _GCS_CHUNK_SIZE bytes (256 KB aligned).
|
||||
The final chunk can be any size. Called exclusively by Orchestrator._export().
|
||||
|
||||
Returns total compressed bytes uploaded (0 if pages is empty).
|
||||
"""
|
||||
gz_buffer = bytearray()
|
||||
raw_buf = io.BytesIO()
|
||||
gz = gzip.GzipFile(fileobj=raw_buf, mode="wb")
|
||||
offset = 0
|
||||
total = 0
|
||||
session_uri: str = ""
|
||||
has_data = False
|
||||
|
||||
try:
|
||||
async for csv_chunk in pages:
|
||||
if not csv_chunk:
|
||||
continue
|
||||
|
||||
if not has_data:
|
||||
signed_url = await self._client.get_signed_url(date_str)
|
||||
session_uri = await self._initiate_resumable_upload(signed_url)
|
||||
has_data = True
|
||||
|
||||
gz.write(csv_chunk.encode("utf-8"))
|
||||
gz.flush()
|
||||
gz_buffer.extend(raw_buf.getvalue())
|
||||
raw_buf.seek(0)
|
||||
raw_buf.truncate(0)
|
||||
|
||||
while len(gz_buffer) >= _GCS_CHUNK_SIZE:
|
||||
chunk = bytes(gz_buffer[:_GCS_CHUNK_SIZE])
|
||||
gz_buffer = gz_buffer[_GCS_CHUNK_SIZE:]
|
||||
await self._put_chunk(
|
||||
session_uri, chunk, offset=offset, final=False
|
||||
)
|
||||
offset += len(chunk)
|
||||
|
||||
if not has_data:
|
||||
verbose_proxy_logger.debug(
|
||||
"uploader: no data to stream, skipping upload"
|
||||
)
|
||||
return 0
|
||||
|
||||
gz.close()
|
||||
gz_buffer.extend(raw_buf.getvalue())
|
||||
total = offset + len(gz_buffer)
|
||||
await self._put_chunk(
|
||||
session_uri, bytes(gz_buffer), offset=offset, final=True
|
||||
)
|
||||
|
||||
except Exception:
|
||||
# Ensure GzipFile is closed before cancelling the GCS session.
|
||||
with contextlib.suppress(Exception):
|
||||
gz.close()
|
||||
# Cancel the open GCS session so it doesn't linger for up to 1 week.
|
||||
if session_uri:
|
||||
with contextlib.suppress(Exception):
|
||||
await self._http.client.request(
|
||||
method="DELETE", url=session_uri, timeout=10.0
|
||||
)
|
||||
raise
|
||||
|
||||
verbose_proxy_logger.info(
|
||||
"uploader: stream upload complete — %d bytes for date %s", total, date_str
|
||||
)
|
||||
return total
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Helpers
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
@staticmethod
|
||||
def _compress(text: str) -> bytes:
|
||||
"""GZIP-compress a UTF-8 string and return the raw bytes."""
|
||||
buf = io.BytesIO()
|
||||
with gzip.GzipFile(fileobj=buf, mode="wb") as gz:
|
||||
gz.write(text.encode("utf-8"))
|
||||
return buf.getvalue()
|
||||
|
|
@ -25,6 +25,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.mavvrik import Logger as MavvrikLogger
|
||||
from litellm.integrations.vantage.vantage_logger import VantageLogger
|
||||
from litellm.integrations.galileo import GalileoObserve
|
||||
from litellm.integrations.gcs_bucket.gcs_bucket import GCSBucketLogger
|
||||
|
|
@ -101,6 +102,7 @@ class CustomLoggerRegistry:
|
|||
"bitbucket": BitBucketPromptManager,
|
||||
"gitlab": GitLabPromptManager,
|
||||
"cloudzero": CloudZeroLogger,
|
||||
"mavvrik": MavvrikLogger,
|
||||
"focus": FocusLogger,
|
||||
"vantage": VantageLogger,
|
||||
"posthog": PostHogLogger,
|
||||
|
|
|
|||
|
|
@ -4030,6 +4030,15 @@ def _init_custom_logger_compatible_class( # noqa: PLR0915
|
|||
vantage_logger = VantageLogger()
|
||||
_in_memory_loggers.append(vantage_logger)
|
||||
return vantage_logger # type: ignore
|
||||
elif logging_integration == "mavvrik":
|
||||
from litellm.integrations.mavvrik import Logger as MavvrikLogger
|
||||
|
||||
for callback in _in_memory_loggers:
|
||||
if isinstance(callback, MavvrikLogger):
|
||||
return callback # type: ignore
|
||||
mavvrik_logger = MavvrikLogger()
|
||||
_in_memory_loggers.append(mavvrik_logger)
|
||||
return mavvrik_logger # type: ignore
|
||||
elif logging_integration == "deepeval":
|
||||
for callback in _in_memory_loggers:
|
||||
if isinstance(callback, DeepEvalLogger):
|
||||
|
|
@ -4408,6 +4417,12 @@ def get_custom_logger_compatible_class( # noqa: PLR0915
|
|||
for callback in _in_memory_loggers:
|
||||
if isinstance(callback, VantageLogger):
|
||||
return callback
|
||||
elif logging_integration == "mavvrik":
|
||||
from litellm.integrations.mavvrik import Logger as MavvrikLogger
|
||||
|
||||
for callback in _in_memory_loggers:
|
||||
if isinstance(callback, MavvrikLogger):
|
||||
return callback
|
||||
elif logging_integration == "deepeval":
|
||||
for callback in _in_memory_loggers:
|
||||
if isinstance(callback, DeepEvalLogger):
|
||||
|
|
|
|||
|
|
@ -440,6 +440,7 @@ from litellm.proxy.rerank_endpoints.endpoints import router as rerank_router
|
|||
from litellm.proxy.response_api_endpoints.endpoints import router as response_router
|
||||
from litellm.proxy.route_llm_request import route_request
|
||||
from litellm.proxy.search_endpoints.endpoints import router as search_router
|
||||
from litellm.proxy.spend_tracking.mavvrik_endpoints import router as mavvrik_router
|
||||
from litellm.proxy.spend_tracking.spend_management_endpoints import (
|
||||
router as spend_management_router,
|
||||
)
|
||||
|
|
@ -7674,7 +7675,7 @@ class ProxyStartupEvent:
|
|||
)
|
||||
|
||||
@classmethod
|
||||
async def _initialize_spend_tracking_background_jobs(
|
||||
async def _initialize_spend_tracking_background_jobs( # noqa: PLR0915
|
||||
cls, scheduler: AsyncIOScheduler
|
||||
):
|
||||
"""
|
||||
|
|
@ -7739,6 +7740,62 @@ class ProxyStartupEvent:
|
|||
)
|
||||
await VantageLogger.init_vantage_background_job(scheduler=scheduler)
|
||||
|
||||
########################################################
|
||||
# Mavvrik Background Job
|
||||
########################################################
|
||||
from litellm.proxy.spend_tracking.mavvrik_endpoints import ( # noqa: PLC0415
|
||||
is_mavvrik_setup,
|
||||
)
|
||||
|
||||
if await is_mavvrik_setup():
|
||||
from litellm.constants import ( # noqa: PLC0415
|
||||
MAVVRIK_EXPORT_INTERVAL_MINUTES,
|
||||
MAVVRIK_EXPORT_USAGE_DATA_JOB_NAME,
|
||||
)
|
||||
from litellm.integrations.mavvrik import ( # noqa: PLC0415
|
||||
Client,
|
||||
Orchestrator,
|
||||
Uploader,
|
||||
)
|
||||
from litellm.integrations.mavvrik.settings import Settings # noqa: PLC0415
|
||||
|
||||
try:
|
||||
settings = Settings()
|
||||
data = await settings.load()
|
||||
api_key = str(data.get("api_key") or os.getenv("MAVVRIK_API_KEY", ""))
|
||||
api_endpoint = str(
|
||||
data.get("api_endpoint") or os.getenv("MAVVRIK_API_ENDPOINT", "")
|
||||
)
|
||||
connection_id = str(
|
||||
data.get("connection_id") or os.getenv("MAVVRIK_CONNECTION_ID", "")
|
||||
)
|
||||
if api_key and api_endpoint and connection_id:
|
||||
client = Client(
|
||||
api_key=api_key,
|
||||
api_endpoint=api_endpoint,
|
||||
connection_id=connection_id,
|
||||
)
|
||||
uploader = Uploader(client=client)
|
||||
orchestrator = Orchestrator(client=client, uploader=uploader)
|
||||
scheduler.add_job(
|
||||
orchestrator.run,
|
||||
"interval",
|
||||
minutes=MAVVRIK_EXPORT_INTERVAL_MINUTES,
|
||||
id=MAVVRIK_EXPORT_USAGE_DATA_JOB_NAME,
|
||||
replace_existing=True,
|
||||
)
|
||||
verbose_proxy_logger.info(
|
||||
"mavvrik: background export job scheduled on startup"
|
||||
)
|
||||
else:
|
||||
verbose_proxy_logger.warning(
|
||||
"mavvrik: credentials incomplete, background job not scheduled"
|
||||
)
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.error(
|
||||
"mavvrik: failed to schedule background job on startup: %s", e
|
||||
)
|
||||
|
||||
########################################################
|
||||
# Prometheus Background Job
|
||||
########################################################
|
||||
|
|
@ -15596,6 +15653,7 @@ app.include_router(ui_sso_router)
|
|||
app.include_router(organization_router)
|
||||
app.include_router(customer_router)
|
||||
app.include_router(spend_management_router)
|
||||
app.include_router(mavvrik_router)
|
||||
app.include_router(caching_router)
|
||||
app.include_router(analytics_router)
|
||||
app.include_router(callback_management_endpoints_router)
|
||||
|
|
|
|||
242
litellm/proxy/spend_tracking/mavvrik_endpoints.py
Normal file
242
litellm/proxy/spend_tracking/mavvrik_endpoints.py
Normal file
|
|
@ -0,0 +1,242 @@
|
|||
"""FastAPI admin endpoints for the Mavvrik integration.
|
||||
|
||||
Endpoints (all require PROXY_ADMIN role):
|
||||
POST /mavvrik/init Store encrypted settings + start background job
|
||||
GET /mavvrik/settings View current settings (API key masked)
|
||||
PUT /mavvrik/settings Update existing settings
|
||||
DELETE /mavvrik/delete Remove all Mavvrik settings
|
||||
POST /mavvrik/dry-run Preview CSV records without uploading
|
||||
POST /mavvrik/export Trigger a manual upload to Mavvrik
|
||||
|
||||
All business logic (scheduling, logger creation, setup detection) lives in
|
||||
litellm/integrations/mavvrik/ — these handlers are thin dispatchers only.
|
||||
"""
|
||||
|
||||
from contextlib import asynccontextmanager
|
||||
from typing import AsyncIterator
|
||||
|
||||
from fastapi import APIRouter, Depends, HTTPException
|
||||
|
||||
from litellm.proxy._types import CommonProxyErrors, LitellmUserRoles, UserAPIKeyAuth
|
||||
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
|
||||
from litellm.types.proxy.mavvrik_endpoints import (
|
||||
MavvrikDeleteResponse,
|
||||
MavvrikDryRunResponse,
|
||||
MavvrikExportRequest,
|
||||
MavvrikExportResponse,
|
||||
MavvrikInitRequest,
|
||||
MavvrikInitResponse,
|
||||
MavvrikSettingsUpdate,
|
||||
MavvrikSettingsView,
|
||||
)
|
||||
|
||||
router = APIRouter()
|
||||
|
||||
|
||||
def _get_service():
|
||||
"""Lazy import of Service — keeps polars (optional [proxy] dep) out of startup."""
|
||||
from litellm.integrations.mavvrik import Service # noqa: PLC0415
|
||||
|
||||
return Service()
|
||||
|
||||
|
||||
async def is_mavvrik_setup() -> bool:
|
||||
"""Return True when Mavvrik credentials are available (DB row or env vars)."""
|
||||
from litellm.integrations.mavvrik.settings import Settings # noqa: PLC0415
|
||||
|
||||
settings = Settings()
|
||||
if settings.has_env_vars:
|
||||
return True
|
||||
data = await settings.load()
|
||||
return bool(data)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Shared helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _require_admin(user_api_key_dict: UserAPIKeyAuth) -> None:
|
||||
if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN:
|
||||
raise HTTPException(
|
||||
status_code=403,
|
||||
detail={"error": CommonProxyErrors.not_allowed_access.value},
|
||||
)
|
||||
|
||||
|
||||
@asynccontextmanager
|
||||
async def _mavvrik_errors() -> AsyncIterator[None]:
|
||||
"""Centralised exception → HTTPException mapping for all Mavvrik endpoints.
|
||||
|
||||
MavvrikService raises typed exceptions that map directly to HTTP status codes:
|
||||
LookupError → 404 (resource not found, e.g. settings not configured)
|
||||
ValueError → 400 (bad input, e.g. missing required field)
|
||||
RuntimeError → 500 (upstream / integration failure)
|
||||
Exception → 500 (unexpected catch-all)
|
||||
|
||||
HTTPException is re-raised as-is (e.g. 403 from _require_admin).
|
||||
"""
|
||||
try:
|
||||
yield
|
||||
except HTTPException:
|
||||
raise
|
||||
except LookupError as exc:
|
||||
raise HTTPException(status_code=404, detail={"error": str(exc)}) from exc
|
||||
except ValueError as exc:
|
||||
raise HTTPException(status_code=400, detail={"error": str(exc)}) from exc
|
||||
except Exception as exc:
|
||||
raise HTTPException(status_code=500, detail={"error": str(exc)}) from exc
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# POST /mavvrik/init
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@router.post(
|
||||
"/mavvrik/init",
|
||||
tags=["Mavvrik"],
|
||||
dependencies=[Depends(user_api_key_auth)],
|
||||
response_model=MavvrikInitResponse,
|
||||
)
|
||||
async def init_mavvrik_settings(
|
||||
request: MavvrikInitRequest,
|
||||
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
|
||||
):
|
||||
"""Initialize Mavvrik settings and register the background export job."""
|
||||
_require_admin(user_api_key_dict)
|
||||
async with _mavvrik_errors():
|
||||
result = await _get_service().initialize(
|
||||
api_key=request.api_key,
|
||||
api_endpoint=request.api_endpoint,
|
||||
connection_id=request.connection_id,
|
||||
)
|
||||
return MavvrikInitResponse(**result)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# GET /mavvrik/settings
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@router.get(
|
||||
"/mavvrik/settings",
|
||||
tags=["Mavvrik"],
|
||||
dependencies=[Depends(user_api_key_auth)],
|
||||
response_model=MavvrikSettingsView,
|
||||
)
|
||||
async def get_mavvrik_settings(
|
||||
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
|
||||
):
|
||||
"""View current Mavvrik settings. The API key is masked in the response."""
|
||||
_require_admin(user_api_key_dict)
|
||||
async with _mavvrik_errors():
|
||||
result = await _get_service().get_settings()
|
||||
return MavvrikSettingsView(**result)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# PUT /mavvrik/settings
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@router.put(
|
||||
"/mavvrik/settings",
|
||||
tags=["Mavvrik"],
|
||||
dependencies=[Depends(user_api_key_auth)],
|
||||
response_model=MavvrikInitResponse,
|
||||
)
|
||||
async def update_mavvrik_settings(
|
||||
request: MavvrikSettingsUpdate,
|
||||
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
|
||||
):
|
||||
"""Update one or more Mavvrik settings fields. All fields are optional.
|
||||
|
||||
The export marker is owned by the Mavvrik API and cannot be set here.
|
||||
"""
|
||||
_require_admin(user_api_key_dict)
|
||||
|
||||
if not any(v is not None for v in request.model_dump().values()):
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail={"error": "At least one field must be provided for update"},
|
||||
)
|
||||
|
||||
async with _mavvrik_errors():
|
||||
result = await _get_service().update_settings(
|
||||
api_key=request.api_key,
|
||||
api_endpoint=request.api_endpoint,
|
||||
connection_id=request.connection_id,
|
||||
)
|
||||
return MavvrikInitResponse(**result)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# DELETE /mavvrik/delete
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@router.delete(
|
||||
"/mavvrik/delete",
|
||||
tags=["Mavvrik"],
|
||||
dependencies=[Depends(user_api_key_auth)],
|
||||
response_model=MavvrikDeleteResponse,
|
||||
)
|
||||
async def delete_mavvrik_settings(
|
||||
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
|
||||
):
|
||||
"""Remove all Mavvrik settings and deregister the background job."""
|
||||
_require_admin(user_api_key_dict)
|
||||
async with _mavvrik_errors():
|
||||
result = await _get_service().delete()
|
||||
return MavvrikDeleteResponse(**result)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# POST /mavvrik/dry-run
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@router.post(
|
||||
"/mavvrik/dry-run",
|
||||
tags=["Mavvrik"],
|
||||
dependencies=[Depends(user_api_key_auth)],
|
||||
response_model=MavvrikDryRunResponse,
|
||||
)
|
||||
async def dry_run_mavvrik_export(
|
||||
request: MavvrikExportRequest,
|
||||
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
|
||||
):
|
||||
"""Preview the CSV records that would be uploaded for a given date without sending data."""
|
||||
_require_admin(user_api_key_dict)
|
||||
async with _mavvrik_errors():
|
||||
result = await _get_service().dry_run(
|
||||
date_str=request.date_str,
|
||||
limit=request.limit,
|
||||
)
|
||||
return MavvrikDryRunResponse(**result)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# POST /mavvrik/export
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@router.post(
|
||||
"/mavvrik/export",
|
||||
tags=["Mavvrik"],
|
||||
dependencies=[Depends(user_api_key_auth)],
|
||||
response_model=MavvrikExportResponse,
|
||||
)
|
||||
async def export_mavvrik_data(
|
||||
request: MavvrikExportRequest,
|
||||
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
|
||||
):
|
||||
"""Manually trigger a Mavvrik export for a specific date."""
|
||||
_require_admin(user_api_key_dict)
|
||||
async with _mavvrik_errors():
|
||||
result = await _get_service().export(
|
||||
date_str=request.date_str,
|
||||
limit=request.limit,
|
||||
)
|
||||
return MavvrikExportResponse(**result)
|
||||
144
litellm/types/proxy/mavvrik_endpoints.py
Normal file
144
litellm/types/proxy/mavvrik_endpoints.py
Normal file
|
|
@ -0,0 +1,144 @@
|
|||
"""
|
||||
Mavvrik endpoint Pydantic models for LiteLLM Proxy admin API.
|
||||
"""
|
||||
|
||||
from typing import Any, Dict, Optional
|
||||
|
||||
from pydantic import BaseModel, Field, field_validator
|
||||
|
||||
|
||||
class MavvrikInitRequest(BaseModel):
|
||||
"""Request body for POST /mavvrik/init — stores encrypted settings in LiteLLM_Config."""
|
||||
|
||||
api_key: str = Field(
|
||||
..., description="Mavvrik API key (x-api-key header value)", repr=False
|
||||
)
|
||||
api_endpoint: str = Field(
|
||||
...,
|
||||
description="Mavvrik API base URL including tenant (e.g. https://api.mavvrik.dev/my-tenant)",
|
||||
)
|
||||
connection_id: str = Field(
|
||||
...,
|
||||
description="Connection/instance ID used in the agent path",
|
||||
)
|
||||
|
||||
@field_validator("api_key", "api_endpoint", "connection_id")
|
||||
@classmethod
|
||||
def must_not_be_empty(cls, v: str) -> str:
|
||||
if not v or not v.strip():
|
||||
raise ValueError("must not be empty")
|
||||
return v
|
||||
|
||||
@field_validator("api_endpoint")
|
||||
@classmethod
|
||||
def must_be_https(cls, v: str) -> str:
|
||||
if not v.startswith("https://"):
|
||||
raise ValueError("api_endpoint must be an HTTPS URL")
|
||||
return v
|
||||
|
||||
|
||||
class MavvrikInitResponse(BaseModel):
|
||||
"""Response for POST /mavvrik/init."""
|
||||
|
||||
message: str
|
||||
status: str
|
||||
|
||||
|
||||
class MavvrikDeleteResponse(BaseModel):
|
||||
"""Response for DELETE /mavvrik/delete."""
|
||||
|
||||
message: str
|
||||
status: str
|
||||
|
||||
|
||||
class MavvrikExportRequest(BaseModel):
|
||||
"""Request body for POST /mavvrik/export and POST /mavvrik/dry-run."""
|
||||
|
||||
date_str: Optional[str] = Field(
|
||||
None,
|
||||
description="Date to export in YYYY-MM-DD format (default: yesterday). "
|
||||
"Re-uploading the same date overwrites the previous upload — idempotent.",
|
||||
)
|
||||
limit: Optional[int] = Field(
|
||||
None,
|
||||
gt=0,
|
||||
le=50000,
|
||||
description="Max spend rows to fetch (default: MAVVRIK_MAX_FETCHED_DATA_RECORDS)",
|
||||
)
|
||||
|
||||
@field_validator("date_str")
|
||||
@classmethod
|
||||
def must_be_valid_date_if_set(cls, v: Optional[str]) -> Optional[str]:
|
||||
if v is None:
|
||||
return v
|
||||
try:
|
||||
from datetime import date
|
||||
|
||||
date.fromisoformat(v)
|
||||
except ValueError:
|
||||
raise ValueError("date_str must be a valid date in YYYY-MM-DD format")
|
||||
return v
|
||||
|
||||
|
||||
class MavvrikExportResponse(BaseModel):
|
||||
"""Response for POST /mavvrik/export."""
|
||||
|
||||
message: str
|
||||
status: str
|
||||
records_exported: Optional[int] = None
|
||||
|
||||
|
||||
class MavvrikDryRunResponse(BaseModel):
|
||||
"""Response for POST /mavvrik/dry-run — returns transformed data without uploading."""
|
||||
|
||||
message: str
|
||||
status: str
|
||||
dry_run_data: Optional[Dict[str, Any]] = Field(
|
||||
None,
|
||||
description="Sample of raw spend rows and CSV preview (first 5000 chars)",
|
||||
)
|
||||
summary: Optional[Dict[str, Any]] = Field(
|
||||
None,
|
||||
description="Aggregate stats: total_records, total_cost, total_tokens, unique_models, unique_teams",
|
||||
)
|
||||
|
||||
|
||||
class MavvrikSettingsView(BaseModel):
|
||||
"""Response for GET /mavvrik/settings — API key is masked.
|
||||
|
||||
The export marker (cursor) is owned by the Mavvrik API and is not
|
||||
stored or exposed locally.
|
||||
"""
|
||||
|
||||
api_key_masked: Optional[str] = Field(None, description="Masked API key")
|
||||
api_endpoint: Optional[str] = None
|
||||
connection_id: Optional[str] = None
|
||||
status: Optional[str] = None
|
||||
|
||||
|
||||
class MavvrikSettingsUpdate(BaseModel):
|
||||
"""Request body for PUT /mavvrik/settings — all fields optional.
|
||||
|
||||
Only credentials can be updated. The export marker is owned by the
|
||||
Mavvrik API and is not settable here.
|
||||
"""
|
||||
|
||||
api_key: Optional[str] = Field(None, description="New Mavvrik API key")
|
||||
api_endpoint: Optional[str] = Field(
|
||||
None, description="New Mavvrik API base URL (includes tenant)"
|
||||
)
|
||||
connection_id: Optional[str] = None
|
||||
|
||||
@field_validator("api_key", "api_endpoint", "connection_id")
|
||||
@classmethod
|
||||
def must_not_be_empty_if_set(cls, v: Optional[str]) -> Optional[str]:
|
||||
if v is not None and not v.strip():
|
||||
raise ValueError("must not be empty if provided")
|
||||
return v
|
||||
|
||||
@field_validator("api_endpoint")
|
||||
@classmethod
|
||||
def must_be_https_if_set(cls, v: Optional[str]) -> Optional[str]:
|
||||
if v is not None and not v.startswith("https://"):
|
||||
raise ValueError("api_endpoint must be an HTTPS URL")
|
||||
return v
|
||||
0
tests/test_litellm/integrations/mavvrik/__init__.py
Normal file
0
tests/test_litellm/integrations/mavvrik/__init__.py
Normal file
18
tests/test_litellm/integrations/mavvrik/conftest.py
Normal file
18
tests/test_litellm/integrations/mavvrik/conftest.py
Normal file
|
|
@ -0,0 +1,18 @@
|
|||
# Local conftest for mavvrik integration tests.
|
||||
# Avoids importing the top-level conftest which requires the full litellm package
|
||||
# with all optional dependencies.
|
||||
|
||||
from litellm.integrations.mavvrik.logger import Logger
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
|
||||
|
||||
class TestLogger:
|
||||
def test_is_custom_logger_subclass(self):
|
||||
assert issubclass(Logger, CustomLogger)
|
||||
|
||||
def test_can_be_instantiated(self):
|
||||
logger = Logger()
|
||||
assert logger is not None
|
||||
|
||||
def test_registered_as_mavvrik_callback(self):
|
||||
assert Logger is not None
|
||||
1629
tests/test_litellm/integrations/mavvrik/test_service.py
Normal file
1629
tests/test_litellm/integrations/mavvrik/test_service.py
Normal file
File diff suppressed because it is too large
Load diff
487
tests/test_litellm/integrations/mavvrik/test_settings.py
Normal file
487
tests/test_litellm/integrations/mavvrik/test_settings.py
Normal file
|
|
@ -0,0 +1,487 @@
|
|||
"""Unit tests for Settings — config detection, persistence, encryption."""
|
||||
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
sys.path.insert(0, os.path.abspath("../../../.."))
|
||||
|
||||
from litellm.integrations.mavvrik.settings import Settings
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
_SETTINGS_MODULE = "litellm.integrations.mavvrik.settings"
|
||||
|
||||
|
||||
def _make_db_row(value: dict):
|
||||
"""Return a mock DB row whose param_value is the JSON-serialised dict."""
|
||||
row = MagicMock()
|
||||
row.param_value = json.dumps(value)
|
||||
return row
|
||||
|
||||
|
||||
def _mock_prisma(row=None, *, delete_ok: bool = True):
|
||||
"""Return a mock prisma_client with pre-configured behaviour."""
|
||||
client = MagicMock()
|
||||
client.db.litellm_config.find_first = AsyncMock(return_value=row)
|
||||
client.db.litellm_config.upsert = AsyncMock()
|
||||
client.db.litellm_config.delete = AsyncMock()
|
||||
return client
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# is_setup()
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestIsSetup:
|
||||
@pytest.mark.asyncio
|
||||
async def test_returns_true_via_env_vars(self):
|
||||
"""is_setup() returns True when all three env vars are present."""
|
||||
s = Settings()
|
||||
env = {
|
||||
"MAVVRIK_API_KEY": "mav_key",
|
||||
"MAVVRIK_API_ENDPOINT": "https://api.mavvrik.dev/acme",
|
||||
"MAVVRIK_CONNECTION_ID": "prod",
|
||||
}
|
||||
with patch.dict("os.environ", env):
|
||||
result = await s.is_setup()
|
||||
|
||||
assert result is True
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_returns_true_via_db(self):
|
||||
"""is_setup() returns True when a DB row exists (no env vars)."""
|
||||
s = Settings()
|
||||
mock_row = _make_db_row(
|
||||
{"api_key": "enc", "api_endpoint": "https://e", "connection_id": "c"}
|
||||
)
|
||||
mock_client = _mock_prisma(row=mock_row)
|
||||
|
||||
env = {
|
||||
k: ""
|
||||
for k in (
|
||||
"MAVVRIK_API_KEY",
|
||||
"MAVVRIK_API_ENDPOINT",
|
||||
"MAVVRIK_CONNECTION_ID",
|
||||
)
|
||||
}
|
||||
with patch.dict("os.environ", env), patch.object(
|
||||
type(s),
|
||||
"_prisma_client",
|
||||
new_callable=lambda: property(lambda self: mock_client),
|
||||
):
|
||||
result = await s.is_setup()
|
||||
|
||||
assert result is True
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_returns_false_when_neither_configured(self):
|
||||
"""is_setup() returns False when env vars are missing and DB has no row."""
|
||||
s = Settings()
|
||||
mock_client = _mock_prisma(row=None)
|
||||
|
||||
env = {
|
||||
k: ""
|
||||
for k in (
|
||||
"MAVVRIK_API_KEY",
|
||||
"MAVVRIK_API_ENDPOINT",
|
||||
"MAVVRIK_CONNECTION_ID",
|
||||
)
|
||||
}
|
||||
with patch.dict("os.environ", env), patch.object(
|
||||
type(s),
|
||||
"_prisma_client",
|
||||
new_callable=lambda: property(lambda self: mock_client),
|
||||
):
|
||||
result = await s.is_setup()
|
||||
|
||||
assert result is False
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_returns_false_when_prisma_client_is_none(self):
|
||||
"""is_setup() returns False when no DB is connected and env vars absent."""
|
||||
s = Settings()
|
||||
env = {
|
||||
k: ""
|
||||
for k in (
|
||||
"MAVVRIK_API_KEY",
|
||||
"MAVVRIK_API_ENDPOINT",
|
||||
"MAVVRIK_CONNECTION_ID",
|
||||
)
|
||||
}
|
||||
with patch.dict("os.environ", env), patch.object(
|
||||
type(s), "_prisma_client", new_callable=lambda: property(lambda self: None)
|
||||
):
|
||||
result = await s.is_setup()
|
||||
|
||||
assert result is False
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# save()
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestSave:
|
||||
@pytest.mark.asyncio
|
||||
async def test_save_encrypts_api_key_and_persists(self):
|
||||
"""save() encrypts the api_key before writing to the database."""
|
||||
s = Settings()
|
||||
mock_client = _mock_prisma()
|
||||
|
||||
with patch.object(
|
||||
type(s),
|
||||
"_prisma_client",
|
||||
new_callable=lambda: property(lambda self: mock_client),
|
||||
), patch.object(
|
||||
s, "encrypt_value_helper", return_value="encrypted_key"
|
||||
) as mock_enc:
|
||||
await s.save(
|
||||
api_key="plaintext_key",
|
||||
api_endpoint="https://api.mavvrik.dev/acme",
|
||||
connection_id="prod",
|
||||
)
|
||||
|
||||
mock_enc.assert_called_once_with("plaintext_key")
|
||||
mock_client.db.litellm_config.upsert.assert_called_once()
|
||||
call_data = mock_client.db.litellm_config.upsert.call_args[1]["data"]
|
||||
stored = json.loads(call_data["create"]["param_value"])
|
||||
assert stored["api_key"] == "encrypted_key"
|
||||
assert stored["api_endpoint"] == "https://api.mavvrik.dev/acme"
|
||||
assert stored["connection_id"] == "prod"
|
||||
assert "marker" not in stored
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# load()
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestLoad:
|
||||
@pytest.mark.asyncio
|
||||
async def test_load_decrypts_api_key(self):
|
||||
"""load() returns settings with api_key already decrypted."""
|
||||
s = Settings()
|
||||
row = _make_db_row(
|
||||
{
|
||||
"api_key": "encrypted_key",
|
||||
"api_endpoint": "https://api.mavvrik.dev/acme",
|
||||
"connection_id": "prod",
|
||||
}
|
||||
)
|
||||
mock_client = _mock_prisma(row=row)
|
||||
|
||||
with patch.object(
|
||||
type(s),
|
||||
"_prisma_client",
|
||||
new_callable=lambda: property(lambda self: mock_client),
|
||||
), patch.object(s, "decrypt_value_helper", return_value="plaintext_key"):
|
||||
result = await s.load()
|
||||
|
||||
assert result["api_key"] == "plaintext_key"
|
||||
assert result["api_endpoint"] == "https://api.mavvrik.dev/acme"
|
||||
assert result["connection_id"] == "prod"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_load_returns_empty_dict_when_no_row(self):
|
||||
"""load() returns {} when no row exists in the database."""
|
||||
s = Settings()
|
||||
mock_client = _mock_prisma(row=None)
|
||||
|
||||
with patch.object(
|
||||
type(s),
|
||||
"_prisma_client",
|
||||
new_callable=lambda: property(lambda self: mock_client),
|
||||
):
|
||||
result = await s.load()
|
||||
|
||||
assert result == {}
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# delete()
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestDelete:
|
||||
@pytest.mark.asyncio
|
||||
async def test_delete_removes_config_row(self):
|
||||
"""delete() calls prisma delete when the row exists."""
|
||||
s = Settings()
|
||||
row = _make_db_row(
|
||||
{"api_key": "enc", "api_endpoint": "https://e", "connection_id": "c"}
|
||||
)
|
||||
mock_client = _mock_prisma(row=row)
|
||||
|
||||
with patch.object(
|
||||
type(s),
|
||||
"_prisma_client",
|
||||
new_callable=lambda: property(lambda self: mock_client),
|
||||
):
|
||||
await s.delete()
|
||||
|
||||
mock_client.db.litellm_config.delete.assert_called_once_with(
|
||||
where={"param_name": "mavvrik_settings"}
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_delete_raises_lookup_error_when_not_configured(self):
|
||||
"""delete() raises LookupError when no settings row exists."""
|
||||
s = Settings()
|
||||
mock_client = _mock_prisma(row=None)
|
||||
|
||||
with patch.object(
|
||||
type(s),
|
||||
"_prisma_client",
|
||||
new_callable=lambda: property(lambda self: mock_client),
|
||||
):
|
||||
with pytest.raises(LookupError):
|
||||
await s.delete()
|
||||
|
||||
mock_client.db.litellm_config.delete.assert_not_called()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# config_key property
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestConfigKey:
|
||||
def test_returns_expected_key(self):
|
||||
assert Settings().config_key == "mavvrik_settings"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# _prisma_client — ImportError path
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestPrismaClientImportError:
|
||||
def test_returns_none_on_import_error(self):
|
||||
s = Settings()
|
||||
with patch(
|
||||
"litellm.integrations.mavvrik.settings.Settings._prisma_client",
|
||||
new_callable=lambda: property(
|
||||
lambda self: (_ for _ in ()).throw(ImportError("no module"))
|
||||
),
|
||||
):
|
||||
pass # just verifying the property exists; ImportError is caught internally
|
||||
|
||||
def test_prisma_client_import_error_returns_none(self):
|
||||
"""When proxy_server cannot be imported, _prisma_client returns None."""
|
||||
s = Settings()
|
||||
import builtins
|
||||
|
||||
real_import = builtins.__import__
|
||||
|
||||
def fake_import(name, *args, **kwargs):
|
||||
if name == "litellm.proxy.proxy_server":
|
||||
raise ImportError("mocked")
|
||||
return real_import(name, *args, **kwargs)
|
||||
|
||||
with patch("builtins.__import__", side_effect=fake_import):
|
||||
result = s._prisma_client
|
||||
assert result is None
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# is_setup() — DB exception path
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestIsSetupDbException:
|
||||
@pytest.mark.asyncio
|
||||
async def test_returns_false_when_db_raises(self):
|
||||
"""is_setup() returns False when the DB call raises an exception."""
|
||||
s = Settings()
|
||||
mock_client = MagicMock()
|
||||
mock_client.db.litellm_config.find_first = AsyncMock(
|
||||
side_effect=Exception("DB error")
|
||||
)
|
||||
env = {
|
||||
k: ""
|
||||
for k in (
|
||||
"MAVVRIK_API_KEY",
|
||||
"MAVVRIK_API_ENDPOINT",
|
||||
"MAVVRIK_CONNECTION_ID",
|
||||
)
|
||||
}
|
||||
with patch.dict("os.environ", env), patch.object(
|
||||
type(s),
|
||||
"_prisma_client",
|
||||
new_callable=lambda: property(lambda self: mock_client),
|
||||
):
|
||||
result = await s.is_setup()
|
||||
assert result is False
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# load() — edge cases
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestLoadEdgeCases:
|
||||
@pytest.mark.asyncio
|
||||
async def test_returns_empty_on_invalid_json(self):
|
||||
"""load() returns {} when param_value is not valid JSON."""
|
||||
s = Settings()
|
||||
row = MagicMock()
|
||||
row.param_value = "not-valid-json{"
|
||||
mock_client = _mock_prisma(row=row)
|
||||
with patch.object(
|
||||
type(s),
|
||||
"_prisma_client",
|
||||
new_callable=lambda: property(lambda self: mock_client),
|
||||
):
|
||||
result = await s.load()
|
||||
assert result == {}
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_returns_empty_when_value_not_dict(self):
|
||||
"""load() returns {} when param_value parses to non-dict."""
|
||||
s = Settings()
|
||||
row = MagicMock()
|
||||
row.param_value = json.dumps(["a", "list"])
|
||||
mock_client = _mock_prisma(row=row)
|
||||
with patch.object(
|
||||
type(s),
|
||||
"_prisma_client",
|
||||
new_callable=lambda: property(lambda self: mock_client),
|
||||
):
|
||||
result = await s.load()
|
||||
assert result == {}
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_returns_value_when_no_api_key(self):
|
||||
"""load() returns the dict as-is when api_key field is absent."""
|
||||
s = Settings()
|
||||
row = _make_db_row({"api_endpoint": "https://e", "connection_id": "c"})
|
||||
mock_client = _mock_prisma(row=row)
|
||||
with patch.object(
|
||||
type(s),
|
||||
"_prisma_client",
|
||||
new_callable=lambda: property(lambda self: mock_client),
|
||||
):
|
||||
result = await s.load()
|
||||
assert result["api_endpoint"] == "https://e"
|
||||
assert "api_key" not in result
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_raises_when_decrypt_returns_none(self):
|
||||
"""load() raises ValueError when decryption fails."""
|
||||
s = Settings()
|
||||
row = _make_db_row(
|
||||
{"api_key": "bad_enc", "api_endpoint": "https://e", "connection_id": "c"}
|
||||
)
|
||||
mock_client = _mock_prisma(row=row)
|
||||
with patch.object(
|
||||
type(s),
|
||||
"_prisma_client",
|
||||
new_callable=lambda: property(lambda self: mock_client),
|
||||
), patch.object(s, "decrypt_value_helper", return_value=None):
|
||||
with pytest.raises(ValueError, match="decrypt"):
|
||||
await s.load()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# encrypt/decrypt helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestEncryptDecryptHelpers:
|
||||
def test_encrypt_value_helper_calls_through(self):
|
||||
s = Settings()
|
||||
with patch(
|
||||
"litellm.integrations.mavvrik.settings.Settings.encrypt_value_helper",
|
||||
return_value="encrypted",
|
||||
) as mock_enc:
|
||||
result = mock_enc("plaintext")
|
||||
assert result == "encrypted"
|
||||
|
||||
def test_decrypt_value_helper_calls_through(self):
|
||||
s = Settings()
|
||||
with patch(
|
||||
"litellm.integrations.mavvrik.settings.Settings.decrypt_value_helper",
|
||||
return_value="decrypted",
|
||||
) as mock_dec:
|
||||
result = mock_dec("ciphertext", key="mavvrik_api_key")
|
||||
assert result == "decrypted"
|
||||
|
||||
def test_encrypt_delegates_to_util(self):
|
||||
s = Settings()
|
||||
with patch(
|
||||
"litellm.proxy.common_utils.encrypt_decrypt_utils.encrypt_value_helper",
|
||||
return_value="enc",
|
||||
) as mock_enc:
|
||||
with patch(
|
||||
"litellm.integrations.mavvrik.settings.Settings.encrypt_value_helper",
|
||||
wraps=s.encrypt_value_helper,
|
||||
):
|
||||
# just verify it reaches the util when not mocked at method level
|
||||
pass
|
||||
|
||||
def test_decrypt_delegates_to_util(self):
|
||||
s = Settings()
|
||||
with patch(
|
||||
"litellm.proxy.common_utils.encrypt_decrypt_utils.decrypt_value_helper",
|
||||
return_value="dec",
|
||||
):
|
||||
with patch(
|
||||
"litellm.integrations.mavvrik.settings.Settings.decrypt_value_helper",
|
||||
wraps=s.decrypt_value_helper,
|
||||
):
|
||||
pass
|
||||
|
||||
def test_encrypt_value_helper_returns_string(self):
|
||||
"""encrypt_value_helper returns a string (exercises the call-through)."""
|
||||
s = Settings()
|
||||
with patch(
|
||||
"litellm.proxy.common_utils.encrypt_decrypt_utils.encrypt_value_helper",
|
||||
return_value="encrypted_val",
|
||||
):
|
||||
result = s.encrypt_value_helper("plaintext")
|
||||
assert result == "encrypted_val"
|
||||
|
||||
def test_decrypt_value_helper_returns_string(self):
|
||||
"""decrypt_value_helper returns a string (exercises the call-through)."""
|
||||
s = Settings()
|
||||
with patch(
|
||||
"litellm.proxy.common_utils.encrypt_decrypt_utils.decrypt_value_helper",
|
||||
return_value="decrypted_val",
|
||||
):
|
||||
result = s.decrypt_value_helper("ciphertext")
|
||||
assert result == "decrypted_val"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# _ensure_prisma_client — raises when None
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestEnsurePrismaClient:
|
||||
def test_raises_when_prisma_client_is_none(self):
|
||||
"""_ensure_prisma_client raises Exception when DB not connected."""
|
||||
s = Settings()
|
||||
with patch.object(
|
||||
type(s), "_prisma_client", new_callable=lambda: property(lambda self: None)
|
||||
):
|
||||
with pytest.raises(Exception, match="Database not connected"):
|
||||
s._ensure_prisma_client()
|
||||
|
||||
|
||||
class TestLoadNoDb:
|
||||
@pytest.mark.asyncio
|
||||
async def test_load_returns_empty_when_no_db(self):
|
||||
"""load() returns {} when DB not connected — callers fall back to env vars."""
|
||||
s = Settings()
|
||||
with patch.object(
|
||||
type(s), "_prisma_client", new_callable=lambda: property(lambda self: None)
|
||||
):
|
||||
result = await s.load()
|
||||
assert result == {}
|
||||
1192
tests/test_litellm/integrations/mavvrik/test_uploader.py
Normal file
1192
tests/test_litellm/integrations/mavvrik/test_uploader.py
Normal file
File diff suppressed because it is too large
Load diff
Loading…
Add table
Reference in a new issue