diff --git a/litellm/__init__.py b/litellm/__init__.py index 7c92623358d..5aea5a2c7cb 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -144,6 +144,7 @@ _custom_logger_compatible_callbacks_literal = Literal[ "bitbucket", "gitlab", "cloudzero", + "mavvrik", "focus", "vantage", "posthog", diff --git a/litellm/constants.py b/litellm/constants.py index fb765c0226c..008ad09dadb 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -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" diff --git a/litellm/integrations/mavvrik/__init__.py b/litellm/integrations/mavvrik/__init__.py new file mode 100644 index 00000000000..97f626e22a8 --- /dev/null +++ b/litellm/integrations/mavvrik/__init__.py @@ -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, + }, + } diff --git a/litellm/integrations/mavvrik/client.py b/litellm/integrations/mavvrik/client.py new file mode 100644 index 00000000000..ff293dd57aa --- /dev/null +++ b/litellm/integrations/mavvrik/client.py @@ -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]}" + ) diff --git a/litellm/integrations/mavvrik/exporter.py b/litellm/integrations/mavvrik/exporter.py new file mode 100644 index 00000000000..d5de29835d7 --- /dev/null +++ b/litellm/integrations/mavvrik/exporter.py @@ -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 diff --git a/litellm/integrations/mavvrik/logger.py b/litellm/integrations/mavvrik/logger.py new file mode 100644 index 00000000000..97a0ec09c78 --- /dev/null +++ b/litellm/integrations/mavvrik/logger.py @@ -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 diff --git a/litellm/integrations/mavvrik/orchestrator.py b/litellm/integrations/mavvrik/orchestrator.py new file mode 100644 index 00000000000..a96297c829e --- /dev/null +++ b/litellm/integrations/mavvrik/orchestrator.py @@ -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) diff --git a/litellm/integrations/mavvrik/settings.py b/litellm/integrations/mavvrik/settings.py new file mode 100644 index 00000000000..05e5b6416a5 --- /dev/null +++ b/litellm/integrations/mavvrik/settings.py @@ -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}, + }, + ) diff --git a/litellm/integrations/mavvrik/uploader.py b/litellm/integrations/mavvrik/uploader.py new file mode 100644 index 00000000000..27a042190ac --- /dev/null +++ b/litellm/integrations/mavvrik/uploader.py @@ -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() diff --git a/litellm/litellm_core_utils/custom_logger_registry.py b/litellm/litellm_core_utils/custom_logger_registry.py index fd402b90d88..e2c0a8c00bf 100644 --- a/litellm/litellm_core_utils/custom_logger_registry.py +++ b/litellm/litellm_core_utils/custom_logger_registry.py @@ -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, diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 63fa0e64695..8c3dad88168 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -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): diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 4af7caead0e..28eccdf85e7 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -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) diff --git a/litellm/proxy/spend_tracking/mavvrik_endpoints.py b/litellm/proxy/spend_tracking/mavvrik_endpoints.py new file mode 100644 index 00000000000..71d1b643a8c --- /dev/null +++ b/litellm/proxy/spend_tracking/mavvrik_endpoints.py @@ -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) diff --git a/litellm/types/proxy/mavvrik_endpoints.py b/litellm/types/proxy/mavvrik_endpoints.py new file mode 100644 index 00000000000..c6520d6c3ac --- /dev/null +++ b/litellm/types/proxy/mavvrik_endpoints.py @@ -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 diff --git a/tests/test_litellm/integrations/mavvrik/__init__.py b/tests/test_litellm/integrations/mavvrik/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/integrations/mavvrik/conftest.py b/tests/test_litellm/integrations/mavvrik/conftest.py new file mode 100644 index 00000000000..ef11d356b55 --- /dev/null +++ b/tests/test_litellm/integrations/mavvrik/conftest.py @@ -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 diff --git a/tests/test_litellm/integrations/mavvrik/test_service.py b/tests/test_litellm/integrations/mavvrik/test_service.py new file mode 100644 index 00000000000..590dc9851c0 --- /dev/null +++ b/tests/test_litellm/integrations/mavvrik/test_service.py @@ -0,0 +1,1629 @@ +"""Tests for Service facade — verifies constructors and method calls are correct.""" + +import io +import os +import sys +from datetime import date, datetime, timezone +from typing import Optional +from unittest.mock import AsyncMock, MagicMock, patch + +import polars as pl +import pytest + +sys.path.insert(0, os.path.abspath("../../../..")) + +from litellm.integrations.mavvrik import Service +from litellm.integrations.mavvrik.client import Client +from litellm.integrations.mavvrik.exporter import Exporter +from litellm.integrations.mavvrik.orchestrator import Orchestrator +from litellm.integrations.mavvrik.settings import Settings +from litellm.integrations.mavvrik.uploader import Uploader + + +_CREDS = { + "api_key": "key", + "api_endpoint": "https://api.mavvrik.dev/t", + "connection_id": "c-1", +} + + +def _mock_settings(data=None): + s = MagicMock() + s.load = AsyncMock(return_value=data if data is not None else dict(_CREDS)) + s.save = AsyncMock() + s.delete = AsyncMock() + s.has_env_vars = False + return s + + +def _make_df(rows=3): + return pl.DataFrame( + { + "date": ["2026-04-10"] * rows, + "user_id": ["user-1"] * rows, + "model": ["gpt-4o"] * rows, + "spend": [0.015] * rows, + "successful_requests": [5] * rows, + "prompt_tokens": [100] * rows, + "completion_tokens": [50] * rows, + "team_id": ["team-1"] * rows, + } + ) + + +def _mock_exporter(df): + """Return a mock Exporter instance with stubbed export() method.""" + exporter = MagicMock() + csv = "" if df.is_empty() else "col\nval\n" + exporter.export = AsyncMock(return_value=(df, csv)) + return exporter + + +def _mock_uploader(): + """Return a mock Uploader instance.""" + uploader = MagicMock() + uploader.upload = AsyncMock() + return uploader + + +# --------------------------------------------------------------------------- +# Service.initialize — schedules Orchestrator with correct constructors +# --------------------------------------------------------------------------- + + +def _mock_proxy_server(scheduler=None): + """Return a mock proxy_server module with a stubbed scheduler.""" + mock_pserver = MagicMock() + mock_pserver.scheduler = scheduler + return mock_pserver + + +class TestServiceInitialize: + @pytest.mark.asyncio + async def test_initialize_builds_client_uploader_orchestrator(self): + """initialize() must construct Client, Uploader(client=), Orchestrator(client=, uploader=).""" + svc = Service() + svc._settings = _mock_settings() + + created = {} + + mock_client_inst = MagicMock(spec=Client) + mock_uploader_inst = MagicMock(spec=Uploader) + mock_orchestrator_inst = MagicMock(spec=Orchestrator) + mock_orchestrator_inst.run = AsyncMock() + + MockClient = MagicMock(return_value=mock_client_inst) + MockUploader = MagicMock(return_value=mock_uploader_inst) + + def capture_orchestrator(client, uploader): + created["client"] = client + created["uploader"] = uploader + return mock_orchestrator_inst + + MockOrchestrator = MagicMock(side_effect=capture_orchestrator) + mock_scheduler = MagicMock() + + # Build a mock proxy_server module with scheduler set. + # Use patch.dict to inject it — but also cover the case where + # proxy_server is already loaded in CI by overwriting its scheduler attr. + import sys + + mock_pserver = _mock_proxy_server(scheduler=mock_scheduler) + + # If proxy_server already loaded, patch its scheduler directly too. + real_pserver = sys.modules.get("litellm.proxy.proxy_server") + real_scheduler = ( + getattr(real_pserver, "scheduler", "MISSING") if real_pserver else "MISSING" + ) + if real_pserver: + real_pserver.scheduler = mock_scheduler + + try: + with ( + patch.dict(sys.modules, {"litellm.proxy.proxy_server": mock_pserver}), + patch("litellm.integrations.mavvrik.Client", MockClient), + patch("litellm.integrations.mavvrik.Uploader", MockUploader), + patch("litellm.integrations.mavvrik.Orchestrator", MockOrchestrator), + ): + await svc.initialize( + api_key="key", + api_endpoint="https://api.mavvrik.dev/t", + connection_id="c-1", + ) + finally: + if real_pserver and real_scheduler != "MISSING": + real_pserver.scheduler = real_scheduler + + MockClient.assert_called_once_with( + api_key="key", + api_endpoint="https://api.mavvrik.dev/t", + connection_id="c-1", + ) + MockUploader.assert_called_once_with(client=mock_client_inst) + assert created["client"] is mock_client_inst + assert created["uploader"] is mock_uploader_inst + mock_scheduler.add_job.assert_called_once() + + @pytest.mark.asyncio + async def test_initialize_returns_success_when_no_scheduler(self): + """initialize() returns success even when scheduler is unavailable.""" + svc = Service() + svc._settings = _mock_settings() + + mock_pserver = _mock_proxy_server(scheduler=None) + + import sys + + with patch.dict(sys.modules, {"litellm.proxy.proxy_server": mock_pserver}): + result = await svc.initialize( + api_key="key", + api_endpoint="https://api.mavvrik.dev/t", + connection_id="c-1", + ) + + assert result["status"] == "success" + + +# --------------------------------------------------------------------------- +# Service.export — uses Exporter + Uploader directly +# --------------------------------------------------------------------------- + + +class TestServiceExport: + @pytest.mark.asyncio + async def test_export_returns_record_count(self): + """export() must return records_exported from the pipeline.""" + svc = Service() + svc._settings = _mock_settings() + + mock_exporter_inst = _mock_exporter(_make_df(rows=7)) + mock_uploader_inst = _mock_uploader() + + with ( + patch( + "litellm.integrations.mavvrik.Exporter", return_value=mock_exporter_inst + ), + patch( + "litellm.integrations.mavvrik.Uploader", return_value=mock_uploader_inst + ), + patch("litellm.integrations.mavvrik.Client"), + ): + result = await svc.export(date_str="2026-04-10") + + assert result["status"] == "success" + assert result["records_exported"] == 7 + mock_uploader_inst.upload.assert_called_once() + + @pytest.mark.asyncio + async def test_export_raises_when_not_configured(self): + """export() raises ValueError when settings missing and no env vars.""" + svc = Service() + svc._settings = _mock_settings(data={}) + svc._settings.has_env_vars = False + + with pytest.raises(ValueError, match="not configured"): + await svc.export(date_str="2026-04-10") + + @pytest.mark.asyncio + async def test_export_returns_zero_when_no_data(self): + """export() returns 0 records when DB has no rows for the date.""" + svc = Service() + svc._settings = _mock_settings() + + mock_exporter_inst = _mock_exporter(pl.DataFrame()) + mock_uploader_inst = _mock_uploader() + + with ( + patch( + "litellm.integrations.mavvrik.Exporter", return_value=mock_exporter_inst + ), + patch( + "litellm.integrations.mavvrik.Uploader", return_value=mock_uploader_inst + ), + patch("litellm.integrations.mavvrik.Client"), + ): + result = await svc.export(date_str="2026-04-10") + + assert result["records_exported"] == 0 + mock_uploader_inst.upload.assert_not_called() + + @pytest.mark.asyncio + async def test_export_builds_uploader_with_client(self): + """export() must pass client= to Uploader, not credential kwargs.""" + svc = Service() + svc._settings = _mock_settings() + + mock_client_inst = MagicMock(spec=Client) + mock_client_inst.connection_id = "c-1" + mock_uploader_inst = _mock_uploader() + mock_exporter_inst = _mock_exporter(_make_df(rows=2)) + + MockClient = MagicMock(return_value=mock_client_inst) + MockUploader = MagicMock(return_value=mock_uploader_inst) + + with ( + patch("litellm.integrations.mavvrik.Client", MockClient), + patch("litellm.integrations.mavvrik.Uploader", MockUploader), + patch( + "litellm.integrations.mavvrik.Exporter", return_value=mock_exporter_inst + ), + ): + await svc.export(date_str="2026-04-10") + + MockUploader.assert_called_once_with(client=mock_client_inst) + + +# --------------------------------------------------------------------------- +# Service.dry_run — uses Exporter only, never calls uploader.upload +# --------------------------------------------------------------------------- + + +class TestServiceDryRun: + @pytest.mark.asyncio + async def test_dry_run_returns_preview_without_uploading(self): + """dry_run() must return preview data and never call uploader.upload.""" + svc = Service() + svc._settings = _mock_settings() + + mock_exporter_inst = _mock_exporter(_make_df(rows=5)) + mock_uploader_inst = _mock_uploader() + + with ( + patch( + "litellm.integrations.mavvrik.Exporter", return_value=mock_exporter_inst + ), + patch( + "litellm.integrations.mavvrik.Uploader", return_value=mock_uploader_inst + ), + patch("litellm.integrations.mavvrik.Client"), + ): + result = await svc.dry_run(date_str="2026-04-10") + + assert result["status"] == "success" + assert "dry_run_data" in result + assert "summary" in result + mock_uploader_inst.upload.assert_not_called() + + @pytest.mark.asyncio + async def test_dry_run_raises_when_not_configured(self): + """dry_run() raises ValueError when not configured.""" + svc = Service() + svc._settings = _mock_settings(data={}) + svc._settings.has_env_vars = False + + with pytest.raises(ValueError, match="not configured"): + await svc.dry_run(date_str="2026-04-10") + + @pytest.mark.asyncio + async def test_dry_run_returns_empty_when_no_data(self): + """dry_run() returns zero summary when DB has no rows.""" + svc = Service() + svc._settings = _mock_settings() + + mock_exporter_inst = _mock_exporter(pl.DataFrame()) + + with ( + patch( + "litellm.integrations.mavvrik.Exporter", return_value=mock_exporter_inst + ), + patch("litellm.integrations.mavvrik.Client"), + ): + result = await svc.dry_run(date_str="2026-04-10") + + assert result["summary"]["total_records"] == 0 + assert result["dry_run_data"]["usage_data"] == [] + + +# --------------------------------------------------------------------------- +# Merged from test_scheduler.py +# --------------------------------------------------------------------------- + + +def _make_client(**kwargs) -> Client: + defaults = dict( + api_key="mav_key", + api_endpoint="https://api.mavvrik.dev/acme", + connection_id="litellm-test", + ) + defaults.update(kwargs) + return Client(**defaults) + + +def _make_uploader(client=None) -> Uploader: + return Uploader(client=client or _make_client()) + + +def _make_orchestrator() -> Orchestrator: + client = _make_client() + uploader = _make_uploader(client=client) + return Orchestrator(client=client, uploader=uploader) + + +def _make_scheduler_df(rows=3) -> pl.DataFrame: + return pl.DataFrame( + { + "date": ["2026-04-10"] * rows, + "user_id": ["user-alice"] * rows, + "model": ["gpt-4o"] * rows, + "spend": [0.015] * rows, + "successful_requests": [5] * rows, + "prompt_tokens": [100] * rows, + "completion_tokens": [50] * rows, + } + ) + + +# --------------------------------------------------------------------------- +# _resolve_first_run_start_date +# --------------------------------------------------------------------------- + + +class TestResolveFirstRunStartDate: + @pytest.mark.asyncio + async def test_uses_earliest_db_date(self): + orc = _make_orchestrator() + orc._exporter = MagicMock() + orc._exporter.get_earliest_date = AsyncMock(return_value="2026-02-15") + + result = await orc._resolve_first_run_start_date() + + assert result == date(2026, 2, 15) + + @pytest.mark.asyncio + async def test_falls_back_to_yesterday_when_db_empty(self): + orc = _make_orchestrator() + orc._exporter = MagicMock() + orc._exporter.get_earliest_date = AsyncMock(return_value=None) + + with patch.object(Orchestrator, "_utc_today", return_value=date(2026, 4, 16)): + result = await orc._resolve_first_run_start_date() + + assert result == date(2026, 4, 15) + + +# --------------------------------------------------------------------------- +# run() / _run_pipeline() +# --------------------------------------------------------------------------- + + +class TestRunExportLoop: + @pytest.mark.asyncio + async def test_uploads_all_dates_since_marker(self): + """Exports each day from marker to yesterday (inclusive).""" + orc = _make_orchestrator() + exported_dates = [] + + async def fake_export(export_date): + exported_dates.append(export_date.isoformat()) + return 1024 + + orc._client.register = AsyncMock(return_value="2026-04-09") + orc._client.advance_marker = AsyncMock() + orc._client.report_error = AsyncMock() + + with ( + patch.object(orc, "_export", side_effect=fake_export), + patch.object(Orchestrator, "_utc_today", return_value=date(2026, 4, 11)), + patch.object(Orchestrator, "_get_pod_lock_manager", return_value=None), + ): + await orc.run() + + assert exported_dates == ["2026-04-09", "2026-04-10"] + assert orc._client.advance_marker.call_count == 2 + + @pytest.mark.asyncio + async def test_advance_marker_uses_next_day(self): + """advance_marker is called with (export_date + 1) epoch.""" + orc = _make_orchestrator() + + orc._client.register = AsyncMock(return_value="2026-04-09") + orc._client.advance_marker = AsyncMock() + orc._client.report_error = AsyncMock() + + with ( + patch.object(orc, "_export", new_callable=AsyncMock, return_value=512), + patch.object(Orchestrator, "_utc_today", return_value=date(2026, 4, 10)), + patch.object(Orchestrator, "_get_pod_lock_manager", return_value=None), + ): + await orc.run() + + expected_epoch = int(datetime(2026, 4, 10, tzinfo=timezone.utc).timestamp()) + orc._client.advance_marker.assert_called_once_with(expected_epoch) + + @pytest.mark.asyncio + async def test_does_nothing_when_marker_up_to_date(self): + """No exports when marker is already at today.""" + orc = _make_orchestrator() + + orc._client.register = AsyncMock(return_value="2026-04-11") # = today + orc._client.advance_marker = AsyncMock() + orc._client.report_error = AsyncMock() + + with ( + patch.object(orc, "_export", new_callable=AsyncMock) as mock_export, + patch.object(Orchestrator, "_utc_today", return_value=date(2026, 4, 11)), + patch.object(Orchestrator, "_get_pod_lock_manager", return_value=None), + ): + await orc.run() + + mock_export.assert_not_called() + orc._client.advance_marker.assert_not_called() + + @pytest.mark.asyncio + async def test_first_run_uses_earliest_db_date(self): + """First run: register() returns None → start from MIN(date) in DB.""" + orc = _make_orchestrator() + exported_dates = [] + + async def fake_export(export_date): + exported_dates.append(export_date.isoformat()) + return 1024 + + orc._client.register = AsyncMock(return_value=None) + orc._client.advance_marker = AsyncMock() + orc._client.report_error = AsyncMock() + orc._exporter = MagicMock() + orc._exporter.get_earliest_date = AsyncMock(return_value="2026-04-09") + + with ( + patch.object(orc, "_export", side_effect=fake_export), + patch.object(Orchestrator, "_utc_today", return_value=date(2026, 4, 11)), + patch.object(Orchestrator, "_get_pod_lock_manager", return_value=None), + ): + await orc.run() + + assert exported_dates == ["2026-04-09", "2026-04-10"] + + @pytest.mark.asyncio + async def test_advances_marker_on_zero_traffic_day(self): + """Zero-traffic days (0 bytes, no DB error) still advance the marker. + + A legitimate date with no spend rows returns 0 bytes without raising. + The marker must advance so the pipeline doesn't stall on quiet days. + DB-unavailability is handled differently — _stream_pages raises, which + propagates through _export and aborts before _advance is reached. + """ + orc = _make_orchestrator() + + orc._client.register = AsyncMock(return_value="2026-04-09") + orc._client.advance_marker = AsyncMock() + orc._client.report_error = AsyncMock() + + with ( + patch.object(orc, "_export", new_callable=AsyncMock, return_value=0), + patch.object(Orchestrator, "_utc_today", return_value=date(2026, 4, 10)), + patch.object(Orchestrator, "_get_pod_lock_manager", return_value=None), + ): + await orc.run() + + # advance_marker IS called — zero bytes is a legitimate empty day + orc._client.advance_marker.assert_called_once() + orc._client.report_error.assert_not_called() + + @pytest.mark.asyncio + async def test_does_not_advance_when_db_unavailable(self): + """When DB is unavailable, _export raises — marker NOT advanced.""" + orc = _make_orchestrator() + + orc._client.register = AsyncMock(return_value="2026-04-09") + orc._client.advance_marker = AsyncMock() + orc._client.report_error = AsyncMock() + + with ( + patch.object( + orc, + "_export", + new_callable=AsyncMock, + side_effect=RuntimeError("database not connected"), + ), + patch.object(Orchestrator, "_utc_today", return_value=date(2026, 4, 10)), + patch.object(Orchestrator, "_get_pod_lock_manager", return_value=None), + ): + await orc.run() + + orc._client.advance_marker.assert_not_called() + orc._client.report_error.assert_called_once() + + @pytest.mark.asyncio + async def test_reports_error_on_pipeline_failure(self): + """Any pipeline exception is reported to Mavvrik via report_error.""" + orc = _make_orchestrator() + + orc._client.register = AsyncMock(side_effect=RuntimeError("API down")) + orc._client.report_error = AsyncMock() + + with patch.object(Orchestrator, "_get_pod_lock_manager", return_value=None): + await orc.run() # must not raise + + orc._client.report_error.assert_called_once() + assert "RuntimeError" in orc._client.report_error.call_args.args[0] + + +# --------------------------------------------------------------------------- +# Overflow detection +# --------------------------------------------------------------------------- + + +class TestOrchestratorHelpers: + def test_utc_today_returns_date(self): + result = Orchestrator._utc_today() + assert isinstance(result, date) + + def test_to_epoch_converts_date(self): + d = date(2026, 1, 1) + epoch = Orchestrator._to_epoch(d) + assert isinstance(epoch, int) + assert epoch > 0 + + def test_export_end_date_is_yesterday(self): + orc = _make_orchestrator() + with patch.object(Orchestrator, "_utc_today", return_value=date(2026, 4, 12)): + end = orc._export_end_date() + assert end == date(2026, 4, 11) + + def test_date_range_takes_explicit_end(self): + orc = _make_orchestrator() + end = date(2026, 4, 11) + dates = list(orc._date_range(date(2026, 4, 10), end)) + assert dates == [date(2026, 4, 10), date(2026, 4, 11)] + + def test_date_range_single_date_when_start_equals_end(self): + orc = _make_orchestrator() + dates = list(orc._date_range(date(2026, 4, 10), date(2026, 4, 10))) + assert dates == [date(2026, 4, 10)] + + def test_pipeline_skips_when_start_equals_end_plus_one(self): + """start > end means nothing to export — no dates yielded.""" + orc = _make_orchestrator() + dates = list(orc._date_range(date(2026, 4, 12), date(2026, 4, 11))) + assert dates == [] + + def test_get_pod_lock_manager_returns_none_when_proxy_logging_none(self): + with patch( + "litellm.integrations.mavvrik.orchestrator.proxy_logging_obj", + None, + create=True, + ): + pass # import is lazy; tested via run() path below + + @pytest.mark.asyncio + async def test_get_pod_lock_manager_returns_none_when_logging_obj_none(self): + orc = _make_orchestrator() + with patch( + "litellm.integrations.mavvrik.orchestrator.Orchestrator._get_pod_lock_manager", + return_value=None, + ): + orc._client.register = AsyncMock(return_value="2099-01-01") + orc._client.report_error = AsyncMock() + with patch.object( + Orchestrator, "_utc_today", return_value=date(2026, 4, 10) + ): + await orc.run() # no lock, runs directly + + +class TestPodLockAcquired: + @pytest.mark.asyncio + async def test_runs_pipeline_when_lock_acquired(self): + """When Redis is available and lock is acquired, pipeline runs.""" + orc = _make_orchestrator() + orc._client.register = AsyncMock(return_value="2099-01-01") + orc._client.report_error = AsyncMock() + + mock_lock = MagicMock() + mock_lock.redis_cache = MagicMock() + mock_lock.acquire_lock = AsyncMock(return_value=True) + mock_lock.release_lock = AsyncMock() + + with ( + patch.object(Orchestrator, "_get_pod_lock_manager", return_value=mock_lock), + patch.object(Orchestrator, "_utc_today", return_value=date(2026, 4, 10)), + ): + await orc.run() + + mock_lock.acquire_lock.assert_called_once() + mock_lock.release_lock.assert_called_once() + + @pytest.mark.asyncio + async def test_skips_pipeline_when_lock_not_acquired(self): + """When Redis lock is not acquired, pipeline does not run.""" + orc = _make_orchestrator() + orc._client.register = AsyncMock() + + mock_lock = MagicMock() + mock_lock.redis_cache = MagicMock() + mock_lock.acquire_lock = AsyncMock(return_value=False) + + with patch.object( + Orchestrator, "_get_pod_lock_manager", return_value=mock_lock + ): + await orc.run() + + orc._client.register.assert_not_called() + + +class TestResolveFirstRunInvalidDate: + @pytest.mark.asyncio + async def test_falls_back_to_yesterday_on_invalid_date_string(self): + """When DB returns a non-ISO date string, fall back to yesterday.""" + orc = _make_orchestrator() + orc._exporter = MagicMock() + orc._exporter.get_earliest_date = AsyncMock(return_value="not-a-date") + + with patch.object(Orchestrator, "_utc_today", return_value=date(2026, 4, 16)): + result = await orc._resolve_first_run_start_date() + + assert result == date(2026, 4, 15) + + +class TestStreamingExport: + @pytest.mark.asyncio + async def test_export_calls_stream_pages_and_stream_upload(self): + """_export() wires exporter._stream_pages() into uploader._stream_upload().""" + orc = _make_orchestrator() + + async def fake_stream_pages(**kwargs): + yield "date,model\n" + yield "2026-04-09,gpt-4o\n" + + orc._exporter._stream_pages = fake_stream_pages + + stream_upload_called_with = [] + + async def fake_stream_upload(pages, date_str): + stream_upload_called_with.append(date_str) + # consume the generator + async for _ in pages: + pass + return 1024 + + orc._uploader._stream_upload = fake_stream_upload + + result = await orc._export(date(2026, 4, 9)) + + assert result == 1024 + assert stream_upload_called_with == ["2026-04-09"] + + @pytest.mark.asyncio + async def test_export_returns_zero_when_no_data(self): + """_export() returns 0 when _stream_upload returns 0 (no data).""" + orc = _make_orchestrator() + + async def empty_pages(**kwargs): + return + yield + + orc._exporter._stream_pages = empty_pages + + async def fake_stream_upload(pages, date_str): + return 0 + + orc._uploader._stream_upload = fake_stream_upload + + result = await orc._export(date(2026, 4, 9)) + assert result == 0 + + +# --------------------------------------------------------------------------- +# Merged from test_endpoints.py +# --------------------------------------------------------------------------- + + +from litellm.proxy.spend_tracking.mavvrik_endpoints import ( + delete_mavvrik_settings, + dry_run_mavvrik_export, + get_mavvrik_settings, + init_mavvrik_settings, + update_mavvrik_settings, +) +from litellm.types.proxy.mavvrik_endpoints import ( + MavvrikExportRequest, + MavvrikInitRequest, + MavvrikSettingsUpdate, +) +from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth + +# Patch target prefix — MavvrikService methods live in the integrations package. +_SVC = "litellm.integrations.mavvrik.Service" + + +def _admin_user() -> UserAPIKeyAuth: + return UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN) + + +def _non_admin_user() -> UserAPIKeyAuth: + return UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER) + + +# --------------------------------------------------------------------------- +# Auth gate +# --------------------------------------------------------------------------- + + +class TestAdminGate: + @pytest.mark.asyncio + async def test_init_rejects_non_admin(self): + from fastapi import HTTPException + + req = MavvrikInitRequest( + api_key="k", api_endpoint="https://e.com/t", connection_id="c" + ) + with pytest.raises(HTTPException) as exc_info: + await init_mavvrik_settings(req, user_api_key_dict=_non_admin_user()) + assert exc_info.value.status_code == 403 + + @pytest.mark.asyncio + async def test_get_settings_rejects_non_admin(self): + from fastapi import HTTPException + + with pytest.raises(HTTPException) as exc_info: + await get_mavvrik_settings(user_api_key_dict=_non_admin_user()) + assert exc_info.value.status_code == 403 + + @pytest.mark.asyncio + async def test_delete_rejects_non_admin(self): + from fastapi import HTTPException + + with pytest.raises(HTTPException) as exc_info: + await delete_mavvrik_settings(user_api_key_dict=_non_admin_user()) + assert exc_info.value.status_code == 403 + + +# --------------------------------------------------------------------------- +# POST /mavvrik/init +# --------------------------------------------------------------------------- + + +class TestInitSettings: + @pytest.mark.asyncio + async def test_init_stores_settings_and_returns_success(self): + req = MavvrikInitRequest( + api_key="mav_key", + api_endpoint="https://api.mavvrik.dev/acme", + connection_id="litellm-prod", + ) + + with patch( + f"{_SVC}.initialize", + new=AsyncMock( + return_value={ + "message": "Mavvrik settings initialized successfully", + "status": "success", + } + ), + ): + resp = await init_mavvrik_settings(req, user_api_key_dict=_admin_user()) + + assert resp.status == "success" + + @pytest.mark.asyncio + async def test_init_succeeds_even_when_scheduler_unavailable(self): + """MavvrikService.initialize() succeeds even if scheduler is not available.""" + from litellm.integrations.mavvrik.settings import Settings + + req = MavvrikInitRequest( + api_key="mav_key", + api_endpoint="https://api.mavvrik.dev/acme", + connection_id="litellm-prod", + ) + + mock_pserver = MagicMock() + mock_pserver.scheduler = None + with ( + patch.object(Settings, "save", new=AsyncMock()), + patch("litellm.proxy.proxy_server", mock_pserver, create=True), + ): + resp = await init_mavvrik_settings(req, user_api_key_dict=_admin_user()) + + assert resp.status == "success" + + +# --------------------------------------------------------------------------- +# GET /mavvrik/settings +# --------------------------------------------------------------------------- + + +class TestGetSettings: + @pytest.mark.asyncio + async def test_returns_not_configured_when_no_row(self): + with patch( + f"{_SVC}.get_settings", + new=AsyncMock( + return_value={ + "api_key_masked": None, + "api_endpoint": None, + "connection_id": None, + "status": "not_configured", + } + ), + ): + resp = await get_mavvrik_settings(user_api_key_dict=_admin_user()) + + assert resp.status == "not_configured" + assert resp.api_key_masked is None + + @pytest.mark.asyncio + async def test_returns_masked_key_when_configured(self): + with patch( + f"{_SVC}.get_settings", + new=AsyncMock( + return_value={ + "api_key_masked": "mav_*******", + "api_endpoint": "https://api.mavvrik.dev/acme", + "connection_id": "prod", + "status": "configured", + } + ), + ): + resp = await get_mavvrik_settings(user_api_key_dict=_admin_user()) + + assert resp.status == "configured" + assert resp.api_key_masked is not None + assert "mav_plaintextkey" not in (resp.api_key_masked or "") + assert resp.connection_id == "prod" + + @pytest.mark.asyncio + async def test_raises_500_on_service_error(self): + """Any exception from MavvrikService.get_settings() → 500.""" + from fastapi import HTTPException + + with patch( + f"{_SVC}.get_settings", + new=AsyncMock(side_effect=Exception("DB exploded")), + ): + with pytest.raises(HTTPException) as exc_info: + await get_mavvrik_settings(user_api_key_dict=_admin_user()) + + assert exc_info.value.status_code == 500 + + +# --------------------------------------------------------------------------- +# PUT /mavvrik/settings +# --------------------------------------------------------------------------- + + +class TestUpdateSettings: + @pytest.mark.asyncio + async def test_update_rejects_empty_request(self): + from fastapi import HTTPException + + req = MavvrikSettingsUpdate() # all None + with pytest.raises(HTTPException) as exc_info: + await update_mavvrik_settings(req, user_api_key_dict=_admin_user()) + assert exc_info.value.status_code == 400 + + def test_update_rejects_empty_api_key(self): + with pytest.raises(Exception): + MavvrikSettingsUpdate(api_key="") + + +# --------------------------------------------------------------------------- +# DELETE /mavvrik/delete +# --------------------------------------------------------------------------- + + +class TestDeleteSettings: + @pytest.mark.asyncio + async def test_delete_returns_404_when_not_configured(self): + """Settings.delete() raises LookupError → endpoint returns 404.""" + from fastapi import HTTPException + + with patch( + f"{_SVC}.delete", + new=AsyncMock( + side_effect=LookupError( + "Mavvrik settings not found — nothing to delete." + ) + ), + ): + with pytest.raises(HTTPException) as exc_info: + await delete_mavvrik_settings(user_api_key_dict=_admin_user()) + assert exc_info.value.status_code == 404 + + @pytest.mark.asyncio + async def test_delete_removes_row_and_returns_success(self): + with patch( + f"{_SVC}.delete", + new=AsyncMock( + return_value={ + "message": "Mavvrik settings deleted successfully", + "status": "success", + } + ), + ): + resp = await delete_mavvrik_settings(user_api_key_dict=_admin_user()) + + assert resp.status == "success" + + +# --------------------------------------------------------------------------- +# POST /mavvrik/dry-run +# --------------------------------------------------------------------------- + + +class TestDryRunMavvrikExport: + @pytest.mark.asyncio + async def test_dry_run_returns_preview(self): + req = MavvrikExportRequest(date_str="2024-01-14") + + with patch( + f"{_SVC}.dry_run", + new=AsyncMock( + return_value={ + "message": "Mavvrik dry run completed", + "status": "success", + "dry_run_data": { + "usage_data": [ + {"date": "2024-01-14", "model": "gpt-4o", "spend": 1.5} + ], + "csv_preview": "date,model,spend\n2024-01-14,gpt-4o,1.5", + }, + "summary": { + "total_records": 1, + "total_cost": 1.5, + "total_tokens": 100, + "unique_models": 1, + "unique_teams": 1, + }, + } + ), + ): + resp = await dry_run_mavvrik_export(req, user_api_key_dict=_admin_user()) + + assert resp.status == "success" + assert resp.summary is not None + assert resp.summary["total_records"] == 1 + + @pytest.mark.asyncio + async def test_dry_run_defaults_to_yesterday_when_no_date(self): + req = MavvrikExportRequest() # no date_str + + with patch( + f"{_SVC}.dry_run", + new=AsyncMock( + return_value={ + "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, + }, + } + ), + ) as mock_dry_run: + resp = await dry_run_mavvrik_export(req, user_api_key_dict=_admin_user()) + + assert resp.status == "success" + mock_dry_run.assert_called_once() + # date_str=None is passed through; MavvrikService.dry_run() resolves it to yesterday + _, kwargs = mock_dry_run.call_args + assert "date_str" in kwargs + + +# --------------------------------------------------------------------------- +# POST /mavvrik/export +# --------------------------------------------------------------------------- + + +class TestExportMavvrikData: + @pytest.mark.asyncio + async def test_export_returns_success_and_record_count(self): + from litellm.proxy.spend_tracking.mavvrik_endpoints import export_mavvrik_data + + req = MavvrikExportRequest(date_str="2024-01-15") + + with patch( + f"{_SVC}.export", + new=AsyncMock( + return_value={ + "message": "Mavvrik export completed successfully for 2024-01-15", + "status": "success", + "records_exported": 7, + } + ), + ): + resp = await export_mavvrik_data(req, user_api_key_dict=_admin_user()) + + assert resp.status == "success" + assert resp.records_exported == 7 + assert "2024-01-15" in resp.message + + @pytest.mark.asyncio + async def test_export_returns_400_when_not_configured(self): + from fastapi import HTTPException + + from litellm.proxy.spend_tracking.mavvrik_endpoints import export_mavvrik_data + + req = MavvrikExportRequest(date_str="2024-01-15") + + with patch( + f"{_SVC}.export", + new=AsyncMock( + side_effect=ValueError( + "Mavvrik not configured. Call POST /mavvrik/init first." + ) + ), + ): + with pytest.raises(HTTPException) as exc_info: + await export_mavvrik_data(req, user_api_key_dict=_admin_user()) + + assert exc_info.value.status_code == 400 + + @pytest.mark.asyncio + async def test_export_defaults_to_yesterday_when_no_date(self): + from litellm.proxy.spend_tracking.mavvrik_endpoints import export_mavvrik_data + + req = MavvrikExportRequest() # no date_str + + with patch( + f"{_SVC}.export", + new=AsyncMock( + return_value={ + "message": "Mavvrik export completed successfully for 2024-01-15", + "status": "success", + "records_exported": 3, + } + ), + ) as mock_export: + resp = await export_mavvrik_data(req, user_api_key_dict=_admin_user()) + + assert resp.status == "success" + mock_export.assert_called_once() + _, kwargs = mock_export.call_args + assert "date_str" in kwargs + + +# --------------------------------------------------------------------------- +# Lifecycle flow — init → settings → update marker → delete → export fails +# --------------------------------------------------------------------------- + + +class TestLifecycleFlow: + @pytest.mark.asyncio + async def test_get_settings_returns_not_configured_after_delete(self): + """DELETE succeeds → subsequent GET returns not_configured.""" + with patch( + f"{_SVC}.delete", + new=AsyncMock( + return_value={ + "message": "Mavvrik settings deleted successfully", + "status": "success", + } + ), + ): + del_resp = await delete_mavvrik_settings(user_api_key_dict=_admin_user()) + + assert del_resp.status == "success" + + with patch( + f"{_SVC}.get_settings", + new=AsyncMock( + return_value={ + "api_key_masked": None, + "api_endpoint": None, + "connection_id": None, + "status": "not_configured", + } + ), + ): + get_resp = await get_mavvrik_settings(user_api_key_dict=_admin_user()) + + assert get_resp.status == "not_configured" + + +# --------------------------------------------------------------------------- +# Settings — setup detection +# --------------------------------------------------------------------------- + + +class TestSettingsSetup: + @pytest.mark.asyncio + async def test_is_mavvrik_setup_true_when_env_vars_set(self): + """Settings.is_setup returns True when all env vars are present.""" + from litellm.integrations.mavvrik.settings import Settings + + with patch.dict( + "os.environ", + { + "MAVVRIK_API_KEY": "mav_key", + "MAVVRIK_API_ENDPOINT": "https://api.mavvrik.dev/acme", + "MAVVRIK_CONNECTION_ID": "prod", + }, + ): + result = await Settings().is_setup() + + assert result is True + + @pytest.mark.asyncio + async def test_is_mavvrik_setup_false_when_no_env_and_no_db(self): + """Settings.is_setup returns False when env vars missing and DB not connected.""" + from litellm.integrations.mavvrik.settings import Settings + + env = { + k: "" + for k in ( + "MAVVRIK_API_KEY", + "MAVVRIK_API_ENDPOINT", + "MAVVRIK_CONNECTION_ID", + ) + } + with patch.dict("os.environ", env): + with patch( + "litellm.integrations.mavvrik.settings.Settings._prisma_client", + new_callable=lambda: property(lambda self: None), + ): + result = await Settings().is_setup() + + assert result is False + + +# --------------------------------------------------------------------------- +# Merged from test_exporter.py +# --------------------------------------------------------------------------- + + +def _make_exporter_df(**kwargs) -> pl.DataFrame: + """Helper: build a minimal spend DataFrame, overriding defaults with kwargs.""" + defaults = { + "date": ["2025-01-19"], + "user_id": ["user-1"], + "api_key": ["sk-abc"], + "model": ["gpt-4o"], + "model_group": ["gpt-4o-group"], + "custom_llm_provider": ["openai"], + "prompt_tokens": [100], + "completion_tokens": [50], + "spend": [1.5], + "api_requests": [5], + "successful_requests": [4], + "failed_requests": [1], + "team_id": ["team-1"], + "api_key_alias": ["prod-key"], + "team_alias": ["Alpha"], + "user_email": ["alice@example.com"], + } + defaults.update(kwargs) + return pl.DataFrame(defaults) + + +class TestExporterToCsv: + def test_empty_dataframe_returns_empty_string(self): + transformer = Exporter() + result = transformer._to_csv(pl.DataFrame()) + assert result == "" + + def test_nonzero_successful_requests_in_output(self): + transformer = Exporter() + df = _make_exporter_df(successful_requests=[3]) + result = transformer._to_csv(df) + assert result != "" + + def test_output_has_header_row(self): + transformer = Exporter() + df = _make_exporter_df() + result = transformer._to_csv(df) + header = result.split("\n")[0] + assert "model" in header + assert "spend" in header + + def test_all_db_columns_present_in_header(self): + transformer = Exporter() + df = _make_exporter_df() + header = transformer._to_csv(df).split("\n")[0] + for col in [ + "date", + "user_id", + "api_key", + "model", + "model_group", + "custom_llm_provider", + "prompt_tokens", + "completion_tokens", + "spend", + "api_requests", + "successful_requests", + "failed_requests", + "team_id", + "api_key_alias", + "team_alias", + "user_email", + ]: + assert col in header, f"Expected column '{col}' in CSV header" + + def test_spend_value_in_output(self): + transformer = Exporter() + df = _make_exporter_df(spend=[42.5]) + result = transformer._to_csv(df) + assert "42.5" in result + + def test_model_value_in_output(self): + transformer = Exporter() + df = _make_exporter_df(model=["claude-3-5-sonnet"]) + result = transformer._to_csv(df) + assert "claude-3-5-sonnet" in result + + def test_multiple_rows_all_in_output(self): + transformer = Exporter() + df = pl.DataFrame( + { + "date": ["2025-01-19", "2025-01-20"], + "successful_requests": [2, 3], + "spend": [1.0, 2.0], + "model": ["gpt-4", "claude-3"], + } + ) + result = transformer._to_csv(df) + lines = [l for l in result.strip().split("\n") if l] + assert len(lines) == 3 # header + 2 data rows + + def test_output_is_valid_csv(self): + transformer = Exporter() + df = _make_exporter_df() + result = transformer._to_csv(df) + # Polars can re-read its own CSV output + reloaded = pl.read_csv(io.StringIO(result)) + assert len(reloaded) == 1 + assert "model" in reloaded.columns + + def test_return_type_is_str(self): + transformer = Exporter() + df = _make_exporter_df() + result = transformer._to_csv(df) + assert isinstance(result, str) + + +# --------------------------------------------------------------------------- +# to_csv — connection_id column +# --------------------------------------------------------------------------- + + +class TestToCsvConnectionId: + def test_adds_connection_id_column_when_provided(self): + exporter = Exporter() + df = _make_exporter_df() + result = exporter._to_csv(df, connection_id="conn-123") + assert "connection_id" in result + assert "conn-123" in result + + def test_no_connection_id_column_when_omitted(self): + exporter = Exporter() + df = _make_exporter_df() + result = exporter._to_csv(df) + assert "connection_id" not in result.split("\n")[0] + + +# --------------------------------------------------------------------------- +# _prisma_client — raises when DB not connected +# --------------------------------------------------------------------------- + + +class TestExporterPrismaClient: + def test_raises_runtime_error_when_db_not_connected(self): + """_prisma_client raises RuntimeError when prisma_client is None.""" + exporter = Exporter() + with patch( + "litellm.integrations.mavvrik.exporter.prisma_client", None, create=True + ): + with patch( + "litellm.integrations.mavvrik.exporter.Exporter._prisma_client", + new_callable=lambda: property( + lambda self: (_ for _ in ()).throw( + RuntimeError("Database not connected") + ) + ), + ): + with pytest.raises(RuntimeError, match="Database not connected"): + _ = exporter._prisma_client + + +# --------------------------------------------------------------------------- +# get_usage_data and get_earliest_date — mocked prisma +# --------------------------------------------------------------------------- + + +class TestExporterDbMethods: + @pytest.mark.asyncio + async def test_get_usage_data_returns_dataframe(self): + """get_usage_data() returns a Polars DataFrame from query_raw results.""" + exporter = Exporter() + mock_rows = [ + { + "date": "2026-04-10", + "user_id": "user-1", + "model": "gpt-4o", + "spend": 0.015, + "successful_requests": 5, + "prompt_tokens": 100, + "completion_tokens": 50, + } + ] + mock_client = MagicMock() + mock_client.db.query_raw = AsyncMock(return_value=mock_rows) + + with patch.object( + type(exporter), + "_prisma_client", + new_callable=lambda: property(lambda self: mock_client), + ): + df = await exporter._get_usage_data("2026-04-10") + + assert len(df) == 1 + assert "model" in df.columns + + @pytest.mark.asyncio + async def test_get_usage_data_with_limit(self): + """get_usage_data() appends LIMIT clause when limit is provided.""" + exporter = Exporter() + mock_client = MagicMock() + mock_client.db.query_raw = AsyncMock(return_value=[]) + captured = [] + + async def fake_query_raw(query, *params): + captured.append((query, params)) + return [] + + mock_client.db.query_raw = fake_query_raw + + with patch.object( + type(exporter), + "_prisma_client", + new_callable=lambda: property(lambda self: mock_client), + ): + await exporter._get_usage_data("2026-04-10", limit=100) + + assert "LIMIT" in captured[0][0] + assert 100 in captured[0][1] + + @pytest.mark.asyncio + async def test_get_earliest_date_returns_date_string(self): + """get_earliest_date() returns first 10 chars of the MIN(date) result.""" + exporter = Exporter() + mock_client = MagicMock() + mock_client.db.query_raw = AsyncMock( + return_value=[{"earliest": "2026-01-01T00:00:00"}] + ) + + with patch.object( + type(exporter), + "_prisma_client", + new_callable=lambda: property(lambda self: mock_client), + ): + result = await exporter.get_earliest_date() + + assert result == "2026-01-01" + + @pytest.mark.asyncio + async def test_get_earliest_date_returns_none_when_empty(self): + """get_earliest_date() returns None when table is empty.""" + exporter = Exporter() + mock_client = MagicMock() + mock_client.db.query_raw = AsyncMock(return_value=[{"earliest": None}]) + + with patch.object( + type(exporter), + "_prisma_client", + new_callable=lambda: property(lambda self: mock_client), + ): + result = await exporter.get_earliest_date() + + assert result is None + + +# --------------------------------------------------------------------------- +# Exporter.export() — public method combining _get_usage_data + filter + _to_csv +# --------------------------------------------------------------------------- + + +class TestExporterExport: + @pytest.mark.asyncio + async def test_export_returns_dataframe_and_csv(self): + """export() returns (filtered_df, csv_str) in one call.""" + exporter = Exporter() + mock_client = MagicMock() + mock_rows = [ + { + "date": "2026-04-10", + "user_id": "user-1", + "model": "gpt-4o", + "spend": 0.015, + "successful_requests": 5, + "prompt_tokens": 100, + "completion_tokens": 50, + } + ] + mock_client.db.query_raw = AsyncMock(return_value=mock_rows) + + with patch.object( + type(exporter), + "_prisma_client", + new_callable=lambda: property(lambda self: mock_client), + ): + df, csv = await exporter.export( + date_str="2026-04-10", connection_id="conn-1" + ) + + assert len(df) == 1 + assert "conn-1" in csv + assert isinstance(csv, str) + + @pytest.mark.asyncio + async def test_export_returns_empty_when_no_data(self): + """export() returns empty DataFrame and empty string when no rows.""" + exporter = Exporter() + mock_client = MagicMock() + mock_client.db.query_raw = AsyncMock(return_value=[]) + + with patch.object( + type(exporter), + "_prisma_client", + new_callable=lambda: property(lambda self: mock_client), + ): + df, csv = await exporter.export( + date_str="2026-04-10", connection_id="conn-1" + ) + + assert df.is_empty() + assert csv == "" + + +# --------------------------------------------------------------------------- +# Exporter._stream_pages — async generator for paginated DB fetch +# --------------------------------------------------------------------------- + + +class TestStreamPages: + @pytest.mark.asyncio + async def test_yields_header_then_csv_rows(self): + """_stream_pages() first yields a CSV header, then row data.""" + exporter = Exporter() + mock_rows = [ + { + "date": "2026-04-10", + "model": "gpt-4o", + "spend": 0.01, + "successful_requests": 1, + }, + ] + # page 1 returns 1 row, page 2 returns empty → stop + mock_client = MagicMock() + mock_client.db.query_raw = AsyncMock(side_effect=[mock_rows, []]) + + with patch.object( + type(exporter), + "_prisma_client", + new_callable=lambda: property(lambda self: mock_client), + ): + chunks = [] + async for chunk in exporter._stream_pages( + "2026-04-10", connection_id="c-1" + ): + chunks.append(chunk) + + assert len(chunks) >= 1 + combined = "".join(chunks) + assert "date" in combined and "model" in combined + assert "gpt-4o" in combined + + @pytest.mark.asyncio + async def test_yields_nothing_when_db_empty(self): + """_stream_pages() yields nothing when DB returns no rows.""" + exporter = Exporter() + mock_client = MagicMock() + mock_client.db.query_raw = AsyncMock(return_value=[]) + + with patch.object( + type(exporter), + "_prisma_client", + new_callable=lambda: property(lambda self: mock_client), + ): + chunks = [] + async for chunk in exporter._stream_pages( + "2026-04-10", connection_id="c-1" + ): + chunks.append(chunk) + + assert chunks == [] + + @pytest.mark.asyncio + async def test_paginates_using_offset(self): + """_stream_pages() uses OFFSET to fetch subsequent pages.""" + exporter = Exporter() + page1 = [ + { + "date": "2026-04-10", + "model": "gpt-4o", + "spend": 0.01, + "successful_requests": 1, + } + ] * 3 + page2 = [] + mock_client = MagicMock() + captured_queries = [] + + async def fake_query_raw(query, *params): + captured_queries.append(params) + return page1 if len(captured_queries) == 1 else page2 + + mock_client.db.query_raw = fake_query_raw + + with patch.object( + type(exporter), + "_prisma_client", + new_callable=lambda: property(lambda self: mock_client), + ): + async for _ in exporter._stream_pages( + "2026-04-10", connection_id="c", page_size=3 + ): + pass + + # Two queries: page 1 (offset 0) and page 2 (offset 3 → empty → stop) + assert len(captured_queries) == 2 + assert captured_queries[0][-1] == 0 # first OFFSET is 0 + assert captured_queries[1][-1] == 3 # second OFFSET is page_size + + +# --------------------------------------------------------------------------- +# Exporter — no DB connected: log warning, return gracefully +# --------------------------------------------------------------------------- + + +class TestExporterNoDb: + @pytest.mark.asyncio + async def test_get_usage_data_returns_empty_when_no_db(self): + """_get_usage_data returns empty DataFrame when DB not connected.""" + exporter = Exporter() + with patch.object( + type(exporter), + "_prisma_client", + new_callable=lambda: property(lambda self: None), + ): + df = await exporter._get_usage_data("2026-04-10") + assert df.is_empty() + + @pytest.mark.asyncio + async def test_get_earliest_date_returns_none_when_no_db(self): + """get_earliest_date returns None when DB not connected.""" + exporter = Exporter() + with patch.object( + type(exporter), + "_prisma_client", + new_callable=lambda: property(lambda self: None), + ): + result = await exporter.get_earliest_date() + assert result is None + + @pytest.mark.asyncio + async def test_stream_pages_raises_when_no_db(self): + """_stream_pages raises RuntimeError when DB not connected. + + Raising (not silently returning) ensures the Orchestrator's try/except + catches it and does NOT advance the marker — preventing silent data loss. + """ + exporter = Exporter() + with patch.object( + type(exporter), + "_prisma_client", + new_callable=lambda: property(lambda self: None), + ): + with pytest.raises(RuntimeError, match="database not connected"): + async for _ in exporter._stream_pages("2026-04-10", connection_id="c"): + pass + + @pytest.mark.asyncio + async def test_stream_pages_yields_nothing_when_db_empty(self): + """_stream_pages yields nothing when DB returns no rows for the date.""" + exporter = Exporter() + mock_client = MagicMock() + mock_client.db.query_raw = AsyncMock(return_value=[]) + with patch.object( + type(exporter), + "_prisma_client", + new_callable=lambda: property(lambda self: mock_client), + ): + chunks = [] + async for chunk in exporter._stream_pages("2026-04-10", connection_id="c"): + chunks.append(chunk) + assert chunks == [] diff --git a/tests/test_litellm/integrations/mavvrik/test_settings.py b/tests/test_litellm/integrations/mavvrik/test_settings.py new file mode 100644 index 00000000000..9fb0ad5a0d2 --- /dev/null +++ b/tests/test_litellm/integrations/mavvrik/test_settings.py @@ -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 == {} diff --git a/tests/test_litellm/integrations/mavvrik/test_uploader.py b/tests/test_litellm/integrations/mavvrik/test_uploader.py new file mode 100644 index 00000000000..95f12f2bd85 --- /dev/null +++ b/tests/test_litellm/integrations/mavvrik/test_uploader.py @@ -0,0 +1,1192 @@ +"""Unit tests for Mavvrik Uploader — GCS resumable upload protocol.""" + +import gzip +import os +import sys +from datetime import datetime +from unittest.mock import AsyncMock, MagicMock, patch + +import httpx +import pytest + +sys.path.insert(0, os.path.abspath("../../../..")) + +from litellm.integrations.mavvrik.client import Client +from litellm.integrations.mavvrik.uploader import Uploader + + +def _make_client(**kwargs) -> Client: + defaults = dict( + api_key="test-key", + api_endpoint="https://api.mavvrik.dev/acme", + connection_id="litellm-001", + ) + defaults.update(kwargs) + return Client(**defaults) + + +def _make_uploader(**kwargs) -> Uploader: + return Uploader(client=_make_client(**kwargs)) + + +def _mock_http_response(status_code: int, text="", headers=None) -> MagicMock: + resp = MagicMock(spec=httpx.Response) + resp.status_code = status_code + resp.text = text + resp.headers = headers or {} + return resp + + +def _patch_gcs_request(return_value=None, side_effect=None): + """Patch the shared httpx client used by Uploader for GCS calls. + + Since Uploader stores _http at __init__ time, we patch the instance's + _http.client.request directly after construction. + """ + mock_request = AsyncMock() + if side_effect is not None: + mock_request.side_effect = side_effect + else: + mock_request.return_value = return_value + + class _Ctx: + def __init__(self, uploader=None): + self._uploader = uploader + self._orig = None + + def __enter__(self): + return self + + def __exit__(self, *args): + pass + + # Return a sentinel that test code uses as context manager. + # Usage: with _patch_gcs_request(return_value=resp)[0]: + # But we need the uploader instance to patch. Use patch.object approach below. + return ( + patch("litellm.integrations.mavvrik.uploader.get_async_httpx_client"), + mock_request, + ) + + +def _patch_uploader_http(u: "Uploader", return_value=None, side_effect=None): + """Patch u._http.client.request directly on an Uploader instance.""" + if side_effect is not None: + u._http.client.request = AsyncMock(side_effect=side_effect) + else: + u._http.client.request = AsyncMock(return_value=return_value) + + +# --------------------------------------------------------------------------- +# Init +# --------------------------------------------------------------------------- + + +class TestUploaderInit: + def test_accepts_client(self): + client = _make_client() + u = Uploader(client=client) + assert u.client is client + + def test_raises_without_client(self): + with pytest.raises(TypeError): + Uploader() # client is required + + +# --------------------------------------------------------------------------- +# _compress +# --------------------------------------------------------------------------- + + +class TestCompress: + def test_returns_bytes(self): + u = _make_uploader() + result = u._compress("hello,world\n") + assert isinstance(result, bytes) + + def test_gzip_decompresses_back_to_original(self): + u = _make_uploader() + original = "date,model,spend\n2025-01-15,gpt-4o,1.5\n" + assert gzip.decompress(u._compress(original)) == original.encode("utf-8") + + def test_empty_string_compresses(self): + u = _make_uploader() + assert isinstance(u._compress(""), bytes) + + +# --------------------------------------------------------------------------- +# _initiate_resumable_upload +# --------------------------------------------------------------------------- + + +class TestInitiateResumableUpload: + @pytest.mark.asyncio + async def test_returns_location_header_on_201(self): + u = _make_uploader() + resp = _mock_http_response( + 201, headers={"Location": "https://gcs.example.com/session"} + ) + _patch_uploader_http(u, return_value=resp) + session_uri = await u._initiate_resumable_upload("https://signed-url") + assert session_uri == "https://gcs.example.com/session" + + @pytest.mark.asyncio + async def test_raises_on_missing_location_header(self): + u = _make_uploader() + resp = _mock_http_response(201, headers={}) + _patch_uploader_http(u, return_value=resp) + with pytest.raises(RuntimeError, match="Location"): + await u._initiate_resumable_upload("https://signed-url") + + @pytest.mark.asyncio + async def test_raises_on_non_201(self): + u = _make_uploader() + resp = _mock_http_response(403, text="Forbidden") + _patch_uploader_http(u, return_value=resp) + with pytest.raises(RuntimeError, match="initiate"): + await u._initiate_resumable_upload("https://signed-url") + + @pytest.mark.asyncio + async def test_sends_gzip_content_type(self): + u = _make_uploader() + resp = _mock_http_response( + 201, headers={"Location": "https://gcs.example.com/session"} + ) + captured = [] + + async def fake_gcs_request(method, url, *, headers=None, **kwargs): + captured.append(headers) + return resp + + _patch_uploader_http(u, side_effect=fake_gcs_request) + await u._initiate_resumable_upload("https://signed-url") + + assert captured[0]["Content-Type"] == "application/gzip" + assert captured[0]["x-goog-resumable"] == "start" + + @pytest.mark.asyncio + async def test_retries_on_5xx_then_raises(self): + """Retry behaviour is tested in TestGcsRequest — just confirm 5xx propagates.""" + u = _make_uploader() + _patch_uploader_http(u, side_effect=RuntimeError("failed after 3 attempts")) + with pytest.raises(RuntimeError, match="failed after"): + await u._initiate_resumable_upload("https://signed-url") + + @pytest.mark.asyncio + async def test_retries_on_request_error_then_raises(self): + u = _make_uploader() + _patch_uploader_http(u, side_effect=RuntimeError("failed after 3 attempts")) + with pytest.raises(RuntimeError, match="failed after"): + await u._initiate_resumable_upload("https://signed-url") + + +# --------------------------------------------------------------------------- +# _finalize_upload +# --------------------------------------------------------------------------- + + +class TestFinalizeUpload: + @pytest.mark.asyncio + async def test_accepts_200(self): + u = _make_uploader() + _patch_uploader_http(u, return_value=_mock_http_response(200)) + await u._finalize_upload("https://session-uri", b"gzip-bytes") + + @pytest.mark.asyncio + async def test_accepts_201(self): + u = _make_uploader() + _patch_uploader_http(u, return_value=_mock_http_response(201)) + await u._finalize_upload("https://session-uri", b"gzip-bytes") + + @pytest.mark.asyncio + async def test_raises_on_error_status(self): + u = _make_uploader() + _patch_uploader_http( + u, return_value=_mock_http_response(500, text="Server Error") + ) + with pytest.raises(RuntimeError, match="finalize"): + await u._finalize_upload("https://session-uri", b"gzip-bytes") + + @pytest.mark.asyncio + async def test_raises_immediately_on_4xx(self): + """4xx from GCS returns immediately — raises in _finalize.""" + u = _make_uploader() + _patch_uploader_http(u, return_value=_mock_http_response(403, text="Forbidden")) + with pytest.raises(RuntimeError, match="finalize"): + await u._finalize_upload("https://session-uri", b"gzip-bytes") + + @pytest.mark.asyncio + async def test_sends_gzip_bytes_as_body(self): + u = _make_uploader() + captured = [] + + async def fake_gcs_request(method, url, *, content=None, **kwargs): + captured.append(content) + return _mock_http_response(200) + + _patch_uploader_http(u, side_effect=fake_gcs_request) + await u._finalize_upload("https://session-uri", b"my-gzip-data") + + assert captured[0] == b"my-gzip-data" + + @pytest.mark.asyncio + async def test_retries_on_5xx_then_raises(self): + """Retry behaviour owned by _gcs_request — confirm propagation.""" + u = _make_uploader() + _patch_uploader_http(u, side_effect=RuntimeError("failed after 3 attempts")) + with pytest.raises(RuntimeError, match="failed after"): + await u._finalize_upload("https://session-uri", b"data") + + @pytest.mark.asyncio + async def test_retries_on_request_error_then_raises(self): + u = _make_uploader() + _patch_uploader_http(u, side_effect=RuntimeError("failed after 3 attempts")) + with pytest.raises(RuntimeError, match="failed after"): + await u._finalize_upload("https://session-uri", b"data") + + +# --------------------------------------------------------------------------- +# upload() +# --------------------------------------------------------------------------- + + +class TestUpload: + @pytest.mark.asyncio + async def test_skips_all_steps_on_empty_payload(self): + u = _make_uploader() + with ( + patch.object( + u.client, "get_signed_url", new_callable=AsyncMock + ) as mock_url, + patch.object( + u, "_initiate_resumable_upload", new_callable=AsyncMock + ) as mock_init, + patch.object(u, "_finalize_upload", new_callable=AsyncMock) as mock_fin, + ): + await u.upload(" ", date_str="2025-01-15") + + mock_url.assert_not_called() + mock_init.assert_not_called() + mock_fin.assert_not_called() + + @pytest.mark.asyncio + async def test_calls_all_three_gcs_steps_in_order(self): + u = _make_uploader() + csv = "date,model,spend\n2025-01-15,gpt-4o,1.5" + call_order = [] + + async def fake_get_signed_url(date_str): + call_order.append("get_signed_url") + return "https://signed" + + async def fake_initiate(signed_url): + call_order.append("initiate") + assert signed_url == "https://signed" + return "https://session" + + async def fake_finalize(session_uri, data): + call_order.append("finalize") + assert session_uri == "https://session" + + with ( + patch.object(u.client, "get_signed_url", side_effect=fake_get_signed_url), + patch.object(u, "_initiate_resumable_upload", side_effect=fake_initiate), + patch.object(u, "_finalize_upload", side_effect=fake_finalize), + ): + await u.upload(csv, date_str="2025-01-15") + + assert call_order == ["get_signed_url", "initiate", "finalize"] + + @pytest.mark.asyncio + async def test_uploads_gzip_compressed_bytes(self): + u = _make_uploader() + csv = "date,model,spend\n2025-01-15,gpt-4o,1.5" + captured_bytes = [] + + async def fake_finalize(session_uri, data): + captured_bytes.append(data) + + with ( + patch.object( + u.client, + "get_signed_url", + new_callable=AsyncMock, + return_value="https://signed", + ), + patch.object( + u, + "_initiate_resumable_upload", + new_callable=AsyncMock, + return_value="https://session", + ), + patch.object(u, "_finalize_upload", side_effect=fake_finalize), + ): + await u.upload(csv, date_str="2025-01-15") + + assert len(captured_bytes) == 1 + assert gzip.decompress(captured_bytes[0]) == csv.encode("utf-8") + + @pytest.mark.asyncio + async def test_passes_date_str_to_get_signed_url(self): + u = _make_uploader() + captured_dates = [] + + async def fake_get_signed_url(date_str): + captured_dates.append(date_str) + return "https://signed" + + with ( + patch.object(u.client, "get_signed_url", side_effect=fake_get_signed_url), + patch.object( + u, + "_initiate_resumable_upload", + new_callable=AsyncMock, + return_value="https://session", + ), + patch.object(u, "_finalize_upload", new_callable=AsyncMock), + ): + await u.upload("col\nval", date_str="2025-03-10") + + assert captured_dates[0] == "2025-03-10" + + +# --------------------------------------------------------------------------- +# Uploader._stream_upload — chunked streaming GCS upload +# --------------------------------------------------------------------------- + + +class TestStreamUpload: + @pytest.mark.asyncio + async def test_stream_upload_sends_chunks_and_returns_count(self): + """_stream_upload() sends 256KB chunks and returns total row count.""" + u = _make_uploader() + + # signed URL + session URI from client + with ( + patch.object( + u.client, + "get_signed_url", + new_callable=AsyncMock, + return_value="https://signed", + ), + patch.object( + u, + "_initiate_resumable_upload", + new_callable=AsyncMock, + return_value="https://session", + ), + ): + + put_calls = [] + + async def fake_put(session_uri, chunk, offset, final): + put_calls.append({"size": len(chunk), "final": final, "offset": offset}) + + with patch.object(u, "_put_chunk", side_effect=fake_put): + # Feed 3 pages of CSV text — enough to trigger at least one 256KB chunk + async def pages(): + yield "date,model,spend\n" # header + yield "2026-01-01,gpt-4o,0.01\n" * 5000 # page 1 + yield "2026-01-01,gpt-4o,0.01\n" * 5000 # page 2 + + count = await u._stream_upload(pages(), date_str="2026-01-01") + + assert count > 0 + assert len(put_calls) >= 1 + # final chunk must be marked final=True + assert put_calls[-1]["final"] is True + # all intermediate chunks must be False + for call in put_calls[:-1]: + assert call["final"] is False + + @pytest.mark.asyncio + async def test_stream_upload_empty_pages_skips_upload(self): + """_stream_upload() with empty generator skips all GCS steps.""" + u = _make_uploader() + + with ( + patch.object( + u.client, "get_signed_url", new_callable=AsyncMock + ) as mock_url, + patch.object( + u, "_initiate_resumable_upload", new_callable=AsyncMock + ) as mock_init, + ): + + async def empty_pages(): + return + yield # make it a generator + + count = await u._stream_upload(empty_pages(), date_str="2026-01-01") + + assert count == 0 + mock_url.assert_not_called() + mock_init.assert_not_called() + + @pytest.mark.asyncio + async def test_put_chunk_intermediate_sends_308_content_range(self): + """_put_chunk() sends Content-Range with * total for intermediate chunks.""" + u = _make_uploader() + captured = [] + + async def fake_gcs_request(method, url, *, headers=None, **kwargs): + captured.append(headers) + return _mock_http_response(308) + + _patch_uploader_http(u, side_effect=fake_gcs_request) + await u._put_chunk("https://session", b"x" * 100, offset=0, final=False) + + assert "Content-Range" in captured[0] + assert captured[0]["Content-Range"].endswith("/*") + + @pytest.mark.asyncio + async def test_put_chunk_final_sends_total_in_content_range(self): + """_put_chunk() declares total size in Content-Range for final chunk.""" + u = _make_uploader() + captured = [] + + async def fake_gcs_request(method, url, *, headers=None, **kwargs): + captured.append(headers) + return _mock_http_response(200) + + _patch_uploader_http(u, side_effect=fake_gcs_request) + await u._put_chunk("https://session", b"x" * 100, offset=256, final=True) + + cr = captured[0]["Content-Range"] + assert cr == "bytes 256-355/356" + + +# --------------------------------------------------------------------------- +# _put_chunk — retry on 5xx +# --------------------------------------------------------------------------- + + +class TestPutChunkRetry: + @pytest.mark.asyncio + async def test_put_chunk_retries_on_5xx_then_raises(self): + """_put_chunk propagates retry failure from _gcs_request.""" + u = _make_uploader() + _patch_uploader_http(u, side_effect=RuntimeError("failed after 3 attempts")) + with pytest.raises(RuntimeError, match="failed after"): + await u._put_chunk("https://session", b"data", offset=0, final=False) + + @pytest.mark.asyncio + async def test_put_chunk_raises_immediately_on_4xx(self): + """4xx returned by GCS causes _put_chunk to raise.""" + u = _make_uploader() + _patch_uploader_http(u, return_value=_mock_http_response(403, text="Forbidden")) + with pytest.raises(RuntimeError): + await u._put_chunk("https://session", b"data", offset=0, final=False) + + +# Transport tested via Client._request() which uses litellm shared HTTP handler + + +class TestStreamUploadEdgeCases: + @pytest.mark.asyncio + async def test_skips_empty_string_chunks(self): + """_stream_upload skips empty string chunks without opening GCS session.""" + u = _make_uploader() + + async def pages_with_empty(): + yield "" # empty — should be skipped + yield "date,model\n" + yield "2026-04-01,gpt-4o\n" + + with ( + patch.object( + u.client, + "get_signed_url", + new_callable=AsyncMock, + return_value="https://s", + ), + patch.object( + u, + "_initiate_resumable_upload", + new_callable=AsyncMock, + return_value="https://sess", + ), + patch.object(u, "_put_chunk", new_callable=AsyncMock), + ): + result = await u._stream_upload(pages_with_empty(), date_str="2026-04-01") + + assert result > 0 # data was uploaded despite the empty first chunk + + @pytest.mark.asyncio + @pytest.mark.asyncio + async def test_sends_intermediate_chunks_when_buffer_exceeds_threshold(self): + """_stream_upload sends intermediate PUT chunks when buffer exceeds 256KB. + + Uses import random to generate incompressible data that won't shrink + below _GCS_CHUNK_SIZE after gzip compression. + """ + import random, string + + u = _make_uploader() + put_calls = [] + + async def fake_put(session_uri, chunk, offset, final): + put_calls.append({"size": len(chunk), "final": final}) + + # Use random printable chars — gzip can't compress this below raw size. + # 400KB of random data will produce ~400KB gzipped, exceeding 256KB. + rng = random.Random(42) + random_data = "".join(rng.choices(string.printable, k=400_000)) + + async def large_pages(): + yield random_data + + with ( + patch.object( + u.client, + "get_signed_url", + new_callable=AsyncMock, + return_value="https://s", + ), + patch.object( + u, + "_initiate_resumable_upload", + new_callable=AsyncMock, + return_value="https://sess", + ), + patch.object(u, "_put_chunk", side_effect=fake_put), + ): + await u._stream_upload(large_pages(), date_str="2026-04-01") + + # At least one intermediate 256KB chunk must have been sent + intermediate = [c for c in put_calls if not c["final"]] + assert len(intermediate) >= 1 + assert all(c["size"] == 256 * 1024 for c in intermediate) + + +# --------------------------------------------------------------------------- +# Merged from test_upload.py +# --------------------------------------------------------------------------- + + +def _make_client(**kwargs) -> Client: + defaults = dict( + api_key="test-key", + api_endpoint="https://api.mavvrik.dev/acme", + connection_id="litellm-001", + ) + defaults.update(kwargs) + return Client(**defaults) + + +def _mock_response( + status_code: int, json_body=None, text="", headers=None +) -> MagicMock: + resp = MagicMock(spec=httpx.Response) + resp.status_code = status_code + resp.text = text + resp.json.return_value = json_body or {} + resp.headers = headers or {} + return resp + + +# --------------------------------------------------------------------------- +# Init +# --------------------------------------------------------------------------- + + +class TestClientInit: + def test_strips_trailing_slash(self): + c = Client( + api_key="k", api_endpoint="https://api.mavvrik.dev/acme/", connection_id="x" + ) + assert c.api_endpoint == "https://api.mavvrik.dev/acme" + + def test_stores_attributes(self): + c = Client( + api_key="mvk-key", + api_endpoint="https://api.mavvrik.dev/t", + connection_id="inst-1", + ) + assert c.api_key == "mvk-key" + assert c.connection_id == "inst-1" + + def test_agent_url_includes_connection_id(self): + c = _make_client(connection_id="prod-001") + assert "prod-001" in c.agent_url + + def test_upload_url_includes_connection_id(self): + c = _make_client(connection_id="prod-001") + assert "prod-001" in c.upload_url + assert "upload-url" in c.upload_url + + +# --------------------------------------------------------------------------- +# _request — uses litellm shared HTTP handler with inline retry + backoff +# --------------------------------------------------------------------------- + + +class TestClientRequest: + @pytest.mark.asyncio + async def test_delegates_to_shared_http_handler(self): + """_request() uses the litellm shared httpx client.""" + c = _make_client() + mock_resp = _mock_response(200) + mock_client = MagicMock() + mock_client.request = AsyncMock(return_value=mock_resp) + mock_handler = MagicMock() + mock_handler.client = mock_client + + with patch( + "litellm.integrations.mavvrik.client.get_async_httpx_client", + return_value=mock_handler, + ): + c._http = mock_handler + resp = await c._request( + "GET", + "https://example.com", + headers={"x-api-key": "k"}, + json={"a": 1}, + label="test", + ) + + assert resp.status_code == 200 + mock_client.request.assert_called_once_with( + method="GET", + url="https://example.com", + headers={"x-api-key": "k"}, + json={"a": 1}, + params=None, + content=None, + timeout=30.0, + ) + + +# --------------------------------------------------------------------------- +# _assert_ok — status checker +# --------------------------------------------------------------------------- + + +class TestAssertOk: + def test_passes_on_expected_status(self): + resp = _mock_response(200) + Client._assert_ok(resp, expected={200, 201}) # no raise + + def test_raises_on_unexpected_status(self): + resp = _mock_response(403, text="Forbidden") + with pytest.raises(RuntimeError, match="403"): + Client._assert_ok(resp, expected={200}) + + def test_accepts_any_code_in_set(self): + for code in (200, 201, 204): + Client._assert_ok(_mock_response(code), expected={200, 201, 204}) + + +# --------------------------------------------------------------------------- +# register() +# --------------------------------------------------------------------------- + + +class TestClientRegister: + @pytest.mark.asyncio + async def test_returns_iso_string_from_epoch(self): + c = _make_client() + with patch.object( + c, + "_request", + return_value=_mock_response(200, {"metricsMarker": 1737000000}), + ): + marker = await c.register() + assert "2025-01-16" in marker + assert "+00:00" in marker or "Z" in marker + + @pytest.mark.asyncio + async def test_returns_none_when_marker_zero(self): + c = _make_client() + with patch.object( + c, "_request", return_value=_mock_response(200, {"metricsMarker": 0}) + ): + assert await c.register() is None + + @pytest.mark.asyncio + async def test_returns_none_when_marker_absent(self): + c = _make_client() + with patch.object(c, "_request", return_value=_mock_response(200, {"id": "x"})): + assert await c.register() is None + + @pytest.mark.asyncio + async def test_raises_on_non_200(self): + c = _make_client() + with patch.object( + c, "_request", return_value=_mock_response(401, text="Unauthorized") + ): + with pytest.raises(RuntimeError, match="401"): + await c.register() + + @pytest.mark.asyncio + async def test_posts_to_agent_url(self): + c = _make_client(connection_id="prod-001") + calls = [] + + async def fake_request(method, url, **kwargs): + calls.append((method, url)) + return _mock_response(200, {"metricsMarker": 1700000000}) + + with patch.object(c, "_request", side_effect=fake_request): + await c.register() + + assert calls[0] == ("POST", c.agent_url) + assert "prod-001" in calls[0][1] + + @pytest.mark.asyncio + async def test_sends_auth_header(self): + c = _make_client() + captured = [] + + async def fake_request(method, url, *, headers=None, **kwargs): + captured.append(headers) + return _mock_response(200, {"metricsMarker": 1700000000}) + + with patch.object(c, "_request", side_effect=fake_request): + await c.register() + + assert captured[0].get("x-api-key") == "test-key" + + +# --------------------------------------------------------------------------- +# advance_marker() +# --------------------------------------------------------------------------- + + +class TestClientAdvanceMarker: + @pytest.mark.asyncio + async def test_accepts_204(self): + c = _make_client() + with patch.object(c, "_request", return_value=_mock_response(204)): + await c.advance_marker(1737000000) + + @pytest.mark.asyncio + async def test_accepts_200(self): + c = _make_client() + with patch.object(c, "_request", return_value=_mock_response(200)): + await c.advance_marker(1737000000) + + @pytest.mark.asyncio + async def test_raises_on_error_status(self): + c = _make_client() + with patch.object( + c, "_request", return_value=_mock_response(403, text="Forbidden") + ): + with pytest.raises(RuntimeError, match="403"): + await c.advance_marker(1737000000) + + @pytest.mark.asyncio + async def test_sends_correct_body(self): + c = _make_client() + captured = [] + + async def fake_request(method, url, *, json=None, **kwargs): + captured.append(json) + return _mock_response(204) + + with patch.object(c, "_request", side_effect=fake_request): + await c.advance_marker(1737000000) + + assert captured[0] == {"metricsMarker": 1737000000} + + @pytest.mark.asyncio + async def test_patches_agent_url(self): + c = _make_client(connection_id="prod-001") + calls = [] + + async def fake_request(method, url, **kwargs): + calls.append((method, url)) + return _mock_response(204) + + with patch.object(c, "_request", side_effect=fake_request): + await c.advance_marker(1737000000) + + assert calls[0][0] == "PATCH" + assert "prod-001" in calls[0][1] + + @pytest.mark.asyncio + async def test_sends_auth_header(self): + c = _make_client() + captured = [] + + async def fake_request(method, url, *, headers=None, **kwargs): + captured.append(headers) + return _mock_response(204) + + with patch.object(c, "_request", side_effect=fake_request): + await c.advance_marker(1737000000) + + assert captured[0].get("x-api-key") == "test-key" + + +# --------------------------------------------------------------------------- +# report_error() +# --------------------------------------------------------------------------- + + +class TestClientReportError: + @pytest.mark.asyncio + async def test_swallows_exception(self): + c = _make_client() + with patch.object(c, "_request", side_effect=RuntimeError("network down")): + await c.report_error("something went wrong") # must not raise + + @pytest.mark.asyncio + async def test_truncates_message_to_500_chars(self): + c = _make_client() + captured = [] + + async def fake_request(method, url, *, json=None, **kwargs): + captured.append(json) + return _mock_response(204) + + with patch.object(c, "_request", side_effect=fake_request): + await c.report_error("x" * 600) + + assert len(captured[0]["error"]) == 500 + + @pytest.mark.asyncio + async def test_sends_error_field_in_body(self): + c = _make_client() + captured = [] + + async def fake_request(method, url, *, json=None, **kwargs): + captured.append(json) + return _mock_response(204) + + with patch.object(c, "_request", side_effect=fake_request): + await c.report_error("export failed") + + assert captured[0] == {"error": "export failed"} + + @pytest.mark.asyncio + async def test_logs_warning_on_unexpected_status(self): + """report_error logs a warning when response status is not 200/204.""" + c = _make_client() + with patch.object(c, "_request", return_value=_mock_response(500, text="err")): + await c.report_error("something broke") # must not raise + + +# --------------------------------------------------------------------------- +# get_signed_url() +# --------------------------------------------------------------------------- + + +class TestClientGetSignedUrl: + @pytest.mark.asyncio + async def test_returns_signed_url_on_200(self): + c = _make_client() + with patch.object( + c, + "_request", + return_value=_mock_response( + 200, {"url": "https://storage.example.com/signed"} + ), + ): + url = await c.get_signed_url("2025-01-15") + assert url == "https://storage.example.com/signed" + + @pytest.mark.asyncio + async def test_raises_on_missing_url_field(self): + c = _make_client() + with patch.object(c, "_request", return_value=_mock_response(200, {})): + with pytest.raises(RuntimeError, match="missing 'url' field"): + await c.get_signed_url("2025-01-15") + + @pytest.mark.asyncio + async def test_raises_on_non_200(self): + c = _make_client() + with patch.object( + c, "_request", return_value=_mock_response(403, text="Forbidden") + ): + with pytest.raises(RuntimeError, match="403"): + await c.get_signed_url("2025-01-15") + + @pytest.mark.asyncio + async def test_sends_date_as_name_param(self): + c = _make_client() + captured = [] + + async def fake_request(method, url, *, params=None, **kwargs): + captured.append(params) + return _mock_response(200, {"url": "https://example.com/signed"}) + + with patch.object(c, "_request", side_effect=fake_request): + await c.get_signed_url("2025-01-15") + + assert captured[0]["name"] == "2025-01-15" + assert captured[0]["type"] == "metrics" + assert captured[0]["datetime"] == "2025-01-15" + + @pytest.mark.asyncio + async def test_sends_auth_header(self): + c = _make_client() + captured = [] + + async def fake_request(method, url, *, headers=None, **kwargs): + captured.append(headers) + return _mock_response(200, {"url": "https://example.com/signed"}) + + with patch.object(c, "_request", side_effect=fake_request): + await c.get_signed_url("2025-01-15") + + assert captured[0].get("x-api-key") == "test-key" + + @pytest.mark.asyncio + async def test_gets_upload_url_with_connection_id(self): + c = _make_client(connection_id="prod-001") + calls = [] + + async def fake_request(method, url, **kwargs): + calls.append(url) + return _mock_response(200, {"url": "https://example.com/signed"}) + + with patch.object(c, "_request", side_effect=fake_request): + await c.get_signed_url("2025-01-15") + + assert "prod-001" in calls[0] + assert "upload-url" in calls[0] + + +# --------------------------------------------------------------------------- +# Merged from test_e2e_upload.py +# --------------------------------------------------------------------------- + +# Tests the full Client → Uploader → Orchestrator flow without real network +# calls, verifying that components wire together correctly end-to-end. + + +_TEST_CSV = ( + "date,user_id,api_key,model,spend\n" "2026-01-01,user-1,sk-test,gpt-4o,0.0025\n" +) +_TEST_DATE = "2026-01-01" + + +def _make_e2e_client() -> Client: + return Client( + api_key="test-key", + api_endpoint="https://api.mavvrik.dev/test", + connection_id="litellm-test", + ) + + +def _make_uploader(client=None) -> Uploader: + return Uploader(client=client or _make_e2e_client()) + + +def _mock_response( + status_code: int, json_body=None, headers=None, text="" +) -> MagicMock: + resp = MagicMock(spec=httpx.Response) + resp.status_code = status_code + resp.text = text + resp.json.return_value = json_body or {} + resp.headers = headers or {} + return resp + + +# --------------------------------------------------------------------------- +# Client integration — each public method wired through _request +# --------------------------------------------------------------------------- + + +class TestClientIntegration: + @pytest.mark.asyncio + async def test_register_returns_iso_marker(self): + """register() parses metricsMarker epoch into ISO-8601 string.""" + client = _make_e2e_client() + with patch.object( + client, + "_request", + return_value=_mock_response(200, {"metricsMarker": 1737000000}), + ): + marker = await client.register() + assert marker is not None + dt = datetime.fromisoformat(marker) + assert dt.year == 2025 + + @pytest.mark.asyncio + async def test_register_returns_none_on_first_run(self): + """register() returns None when metricsMarker is 0 (first run).""" + client = _make_e2e_client() + with patch.object( + client, + "_request", + return_value=_mock_response(200, {"metricsMarker": 0}), + ): + marker = await client.register() + assert marker is None + + @pytest.mark.asyncio + async def test_advance_marker_sends_correct_epoch(self): + """advance_marker() sends metricsMarker in body.""" + client = _make_e2e_client() + captured = [] + + async def fake_request(method, url, *, json=None, **kwargs): + captured.append({"method": method, "json": json}) + return _mock_response(204) + + with patch.object(client, "_request", side_effect=fake_request): + await client.advance_marker(1775001600) + + assert captured[0]["method"] == "PATCH" + assert captured[0]["json"] == {"metricsMarker": 1775001600} + + @pytest.mark.asyncio + async def test_get_signed_url_returns_url(self): + """get_signed_url() extracts url field from response.""" + client = _make_e2e_client() + with patch.object( + client, + "_request", + return_value=_mock_response( + 200, {"url": "https://storage.googleapis.com/signed"} + ), + ): + url = await client.get_signed_url(_TEST_DATE) + assert url == "https://storage.googleapis.com/signed" + + @pytest.mark.asyncio + async def test_get_signed_url_sends_name_and_datetime(self): + """get_signed_url() passes name and datetime params.""" + client = _make_e2e_client() + captured = [] + + async def fake_request(method, url, *, params=None, **kwargs): + captured.append(params) + return _mock_response(200, {"url": "https://example.com/signed"}) + + with patch.object(client, "_request", side_effect=fake_request): + await client.get_signed_url("2026-04-01") + + assert captured[0]["name"] == "2026-04-01" + assert captured[0]["datetime"] == "2026-04-01" + assert captured[0]["type"] == "metrics" + + @pytest.mark.asyncio + async def test_report_error_swallows_exceptions(self): + """report_error() never raises even when the request fails.""" + client = _make_e2e_client() + with patch.object(client, "_request", side_effect=RuntimeError("network down")): + await client.report_error("something broke") # must not raise + + +# --------------------------------------------------------------------------- +# Uploader integration — bulk path wires Client + GCS steps +# --------------------------------------------------------------------------- + + +class TestUploaderIntegration: + @pytest.mark.asyncio + async def test_upload_calls_signed_url_then_gcs(self): + """upload() gets signed URL from Client then initiates and finalises GCS session.""" + uploader = _make_uploader() + call_order = [] + + async def fake_get_signed_url(date_str): + call_order.append("get_signed_url") + return "https://signed" + + async def fake_initiate(signed_url): + call_order.append("initiate") + assert signed_url == "https://signed" + return "https://session" + + async def fake_finalize(session_uri, gzip_bytes): + call_order.append("finalize") + assert session_uri == "https://session" + assert isinstance(gzip_bytes, bytes) + assert gzip.decompress(gzip_bytes) == _TEST_CSV.encode("utf-8") + + with ( + patch.object( + uploader.client, "get_signed_url", side_effect=fake_get_signed_url + ), + patch.object( + uploader, "_initiate_resumable_upload", side_effect=fake_initiate + ), + patch.object(uploader, "_finalize_upload", side_effect=fake_finalize), + ): + await uploader.upload(_TEST_CSV, date_str=_TEST_DATE) + + assert call_order == ["get_signed_url", "initiate", "finalize"] + + @pytest.mark.asyncio + async def test_upload_skips_all_gcs_on_empty_payload(self): + """upload() makes no network calls when payload is blank.""" + uploader = _make_uploader() + with patch.object( + uploader.client, "get_signed_url", new_callable=AsyncMock + ) as mock_url: + await uploader.upload(" ", date_str=_TEST_DATE) + mock_url.assert_not_called() + + @pytest.mark.asyncio + async def test_upload_is_idempotent(self): + """Calling upload() twice for the same date succeeds both times.""" + uploader = _make_uploader() + call_count = [0] + + async def fake_get_signed_url(date_str): + call_count[0] += 1 + return "https://signed" + + with ( + patch.object( + uploader.client, "get_signed_url", side_effect=fake_get_signed_url + ), + patch.object( + uploader, + "_initiate_resumable_upload", + new_callable=AsyncMock, + return_value="https://session", + ), + patch.object(uploader, "_finalize_upload", new_callable=AsyncMock), + ): + await uploader.upload(_TEST_CSV, date_str=_TEST_DATE) + await uploader.upload(_TEST_CSV, date_str=_TEST_DATE) + + assert call_count[0] == 2 # both calls went through + + +# --------------------------------------------------------------------------- +# Full pipeline — Client + Uploader wired together +# --------------------------------------------------------------------------- + + +class TestFullPipelineIntegration: + @pytest.mark.asyncio + async def test_register_then_upload_then_advance(self): + """Simulate one complete export cycle: register → upload → advance.""" + client = _make_e2e_client() + uploader = Uploader(client=client) + advance_calls = [] + + with ( + patch.object( + client, + "_request", + side_effect=[ + # register() → metricsMarker + _mock_response(200, {"metricsMarker": 1775001600}), + # get_signed_url() → url + _mock_response(200, {"url": "https://signed"}), + # advance_marker() → 204 + _mock_response(204), + ], + ), + patch.object( + uploader, + "_initiate_resumable_upload", + new_callable=AsyncMock, + return_value="https://session", + ), + patch.object(uploader, "_finalize_upload", new_callable=AsyncMock), + ): + marker = await client.register() + assert marker is not None + + await uploader.upload(_TEST_CSV, date_str=_TEST_DATE) + + await client.advance_marker(1775088000)