diff --git a/docs/my-website/docs/observability/mavvrik.md b/docs/my-website/docs/observability/mavvrik.md new file mode 100644 index 00000000000..c1c78cedccd --- /dev/null +++ b/docs/my-website/docs/observability/mavvrik.md @@ -0,0 +1,237 @@ +import Tabs from '@theme/Tabs'; +import TabItem from '@theme/TabItem'; + +# Mavvrik Integration + +LiteLLM provides an integration with Mavvrik, allowing you to export your LLM usage data to Mavvrik for AI cost tracking and analysis. + +## Overview + +| Property | Details | +|----------|---------| +| Description | Export LiteLLM daily usage data to Mavvrik | +| Supported Operations | Automatic daily data export, manual data export, dry run testing, cost and token usage tracking | +| Data Format | CSV (gzip-compressed, streamed page-by-page — no row limit) | +| Export Frequency | Hourly scheduler check — exports complete calendar days (never today's partial data) | + +## Setup + +### Step 1: Set the Mavvrik credentials as environment variables + +```bash +export MAVVRIK_API_KEY="mav_xxxxxxxxxx" +export MAVVRIK_API_ENDPOINT="https://api.mavvrik.dev/" +export MAVVRIK_CONNECTION_ID="litellm-prod" +``` + +### Step 2: Start LiteLLM Proxy + +```bash +litellm --config /path/to/config.yaml +``` + +LiteLLM will schedule hourly exports automatically. Registration with the Mavvrik API (and the initial export window determination) happens when the first scheduled job fires, not immediately at process start. If you need exports to begin immediately, use the API-based initialization flow below. + +## Environment Variables + +| Variable | Required | Description | Example | +|----------|----------|-------------|---------| +| `MAVVRIK_API_KEY` | Yes | Your Mavvrik API key (`x-api-key` header) | `mav_xxxxxxxxxx` | +| `MAVVRIK_API_ENDPOINT` | Yes | Mavvrik API base URL including your tenant path | `https://api.mavvrik.dev/` | +| `MAVVRIK_CONNECTION_ID` | Yes | Connection/instance ID assigned by Mavvrik | `litellm-prod` | + +| `MAVVRIK_EXPORT_INTERVAL_MINUTES` | No | Scheduler check frequency in minutes (default: `60`) | `60` | + +## Alternative Setup: API-Based Initialization + +If you prefer to configure Mavvrik without restarting the proxy, use the `/mavvrik/init` endpoint: + +```bash +curl -X POST "http://localhost:4000/mavvrik/init" \ + -H "Content-Type: application/json" \ + -H "Authorization: Bearer sk-admin-key" \ + -d '{ + "api_key": "mav_xxxxxxxxxx", + "api_endpoint": "https://api.mavvrik.dev/", + "connection_id": "litellm-prod" + }' | jq +``` + +**Expected Response:** +```json +{ + "message": "Mavvrik settings initialized successfully", + "status": "success" +} +``` + +This stores encrypted credentials in the database and registers the background export job immediately — no proxy restart required. + +## Testing Your Setup + +### Dry Run Export + +Preview the CSV records that would be uploaded for a given date without sending any data to Mavvrik: + +```bash +curl -X POST "http://localhost:4000/mavvrik/dry-run" \ + -H "Content-Type: application/json" \ + -H "Authorization: Bearer sk-admin-key" \ + -d '{ + "limit": 10 + }' | jq +``` + +**Expected Response:** +```json +{ + "message": "Mavvrik dry run completed", + "status": "success", + "dry_run_data": { + "usage_data": [...], + "csv_preview": "date,model,team_id,user_id,spend,prompt_tokens,completion_tokens,..." + }, + "summary": { + "total_records": 10, + "total_cost": 0.05, + "total_tokens": 1250, + "unique_models": 3, + "unique_teams": 2 + } +} +``` + +### Manual Export + +Trigger an immediate upload to Mavvrik for a specific date: + +```bash +curl -X POST "http://localhost:4000/mavvrik/export" \ + -H "Content-Type: application/json" \ + -H "Authorization: Bearer sk-admin-key" \ + -d '{ + "date_str": "2024-01-15", + "limit": 100 + }' | jq +``` + +**Expected Response:** +```json +{ + "message": "Mavvrik export completed successfully for 2024-01-15", + "status": "success", + "records_exported": 87 +} +``` + +Omitting `date_str` defaults to yesterday. Re-exporting the same date overwrites the previously uploaded object — exports are idempotent. + +### View Current Settings + +Check the current Mavvrik configuration (API key is masked): + +```bash +curl -X GET "http://localhost:4000/mavvrik/settings" \ + -H "Authorization: Bearer sk-admin-key" | jq +``` + +**Expected Response:** +```json +{ + "api_key_masked": "mav_****xxxx", + "api_endpoint": "https://api.mavvrik.dev/", + "connection_id": "litellm-prod", + "status": "configured" +} +``` + +The export cursor (marker) is owned by the Mavvrik API — it is retrieved from Mavvrik at the start of each scheduled run and is not stored locally. + +## Data Export Details + +### Export Schedule + +- **Frequency**: Scheduler runs every 60 minutes (configurable via `MAVVRIK_EXPORT_INTERVAL_MINUTES`) +- **Scope**: Each run exports all complete calendar days since the last marker — never today's partial data +- **First run**: If no marker exists, LiteLLM starts from the earliest date present in `LiteLLM_DailyUserSpend` (i.e. all available history) +- **Catch-up**: If the proxy was offline for multiple days, the scheduler automatically back-fills all missed days on the next run +- **Idempotency**: Each day's data is uploaded to an object named by date (e.g. `2024-01-15`). Re-exporting the same date safely overwrites the previous upload + +### Data Format + +LiteLLM exports daily spend aggregates from `LiteLLM_DailyUserSpend` as a CSV file. Each row represents one model/team/user combination for a given day and includes: + +| Column | Description | +|--------|-------------| +| `date` | Calendar date (YYYY-MM-DD) | +| `model` | LLM model name | +| `team_id` | LiteLLM team identifier | +| `user_id` | LiteLLM user identifier | +| `spend` | Total cost in USD | +| `prompt_tokens` | Input tokens consumed | +| `completion_tokens` | Output tokens generated | +| `successful_requests` | Count of successful API calls | +| `connection_id` | Your Mavvrik connection ID (added by LiteLLM) | + +All rows are exported, including rows where `successful_requests` is `0` (failed requests). Mavvrik handles filtering on their ingestion side. + +## Advanced Configuration + +### Re-export Historical Data + +The export cursor (marker) is owned exclusively by the Mavvrik API and is not settable via `PUT /mavvrik/settings`. + +If Mavvrik asks you to re-export from an earlier date (e.g. after a data reset), contact Mavvrik support to reset the `metricsMarker` on their side. Once reset, the next scheduled run will retrieve the updated marker via `register()` and automatically back-fill all days from that point onwards. + +### Custom Export Frequency + +Change how often the scheduler checks for new days to export: + +```bash +export MAVVRIK_EXPORT_INTERVAL_MINUTES=120 # Check every 2 hours +``` + +### Remove Mavvrik Integration + +Delete all Mavvrik settings and deregister the background job: + +```bash +curl -X DELETE "http://localhost:4000/mavvrik/delete" \ + -H "Authorization: Bearer sk-admin-key" | jq +``` + +## Troubleshooting + +### Common Issues + +1. **`status: not_configured` in logs** + ``` + Initialized Success Callbacks - ['mavvrik'] + ``` + Ensure `MAVVRIK_API_KEY`, `MAVVRIK_API_ENDPOINT`, and `MAVVRIK_CONNECTION_ID` are all set in the environment. The integration auto-initializes from these env vars when the proxy starts. + +2. **401 on registration** + ``` + Mavvrik registration failed: 401 Unauthorized + ``` + Verify your `MAVVRIK_API_KEY` is valid and that `MAVVRIK_API_ENDPOINT` includes your tenant path (e.g. `https://api.mavvrik.dev/`, not just `https://api.mavvrik.dev`). + +3. **No data appearing in Mavvrik** + - Use the dry-run endpoint to verify data exists for the target date + - Check proxy logs for `Mavvrik Orchestrator: up to date` — if the marker is already at today, there is nothing new to export + - Check proxy logs for `no data in DB, skipped` — the date has no spend rows in `LiteLLM_DailyUserSpend` + - Ensure the proxy has been generating traffic (check `LiteLLM_DailyUserSpend` table) + - Only complete calendar days are exported — today's data will appear after midnight UTC + +4. **Missing credentials error on export** + ``` + ValueError: Mavvrik not configured. Call POST /mavvrik/init first. + ``` + Either set the `MAVVRIK_*` environment variables or call `POST /mavvrik/init` to store credentials in the database. + +5. **Export succeeds but `records_exported: 0`** + There are no rows in `LiteLLM_DailyUserSpend` for that date. Verify the proxy was receiving traffic and that spend tracking is enabled. + +## Related Links + +- [Mavvrik Documentation](https://help.mavvrik.ai/) diff --git a/docs/my-website/docs/proxy/config_settings.md b/docs/my-website/docs/proxy/config_settings.md index a886a754f5e..c9b059ce292 100644 --- a/docs/my-website/docs/proxy/config_settings.md +++ b/docs/my-website/docs/proxy/config_settings.md @@ -607,6 +607,11 @@ router_settings: | MCP_OAUTH2_TOKEN_EXPIRY_BUFFER_SECONDS | Seconds to subtract from token expiry when computing cache TTL. Default is 60 | MCP_PER_USER_TOKEN_DEFAULT_TTL | Default TTL in seconds for per-user MCP OAuth tokens stored in Redis. Default is 43200 (12 hours) | MCP_PER_USER_TOKEN_EXPIRY_BUFFER_SECONDS | Seconds to subtract from per-user MCP OAuth token expiry when computing Redis TTL. Default is 60 +| MAVVRIK_API_KEY | API key for Mavvrik cost analytics integration +| MAVVRIK_API_ENDPOINT | Base URL for Mavvrik API endpoint +| MAVVRIK_CONNECTION_ID | Connection ID for Mavvrik data submission +| MAVVRIK_EXPORT_INTERVAL_MINUTES | Interval in minutes between Mavvrik cost data exports. Default is 60 +| MAVVRIK_MAX_FETCHED_DATA_RECORDS | Maximum number of spend records to fetch per Mavvrik export cycle. Default is 50000 | DEFAULT_MOCK_RESPONSE_COMPLETION_TOKEN_COUNT | Default token count for mock response completions. Default is 20 | DEFAULT_MOCK_RESPONSE_PROMPT_TOKEN_COUNT | Default token count for mock response prompts. Default is 10 | DEFAULT_MODEL_CREATED_AT_TIME | Default creation timestamp for models. Default is 1677610602 diff --git a/litellm/__init__.py b/litellm/__init__.py index 89cef667c6e..9d5028aab70 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 012599ab6ab..ebe4f9ace68 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1423,6 +1423,12 @@ 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" SPEND_LOG_RUN_LOOPS = int(os.getenv("SPEND_LOG_RUN_LOOPS", 500)) diff --git a/litellm/integrations/mavvrik/__init__.py b/litellm/integrations/mavvrik/__init__.py new file mode 100644 index 00000000000..ce39976bcc0 --- /dev/null +++ b/litellm/integrations/mavvrik/__init__.py @@ -0,0 +1,380 @@ +"""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 + +import polars as pl + +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=data.get("api_key") or os.getenv("MAVVRIK_API_KEY", ""), + api_endpoint=data.get("api_endpoint") or os.getenv("MAVVRIK_API_ENDPOINT", ""), + connection_id=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() + 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) + return {"message": "Mavvrik settings updated successfully", "status": "success"} + + # ------------------------------------------------------------------ + # delete → DELETE /mavvrik/delete + # ------------------------------------------------------------------ + + async def delete(self) -> dict: + """Remove all Mavvrik settings and deregister the scheduler job. + + Raises: + LookupError: when no settings exist in the database. + """ + from litellm.constants import MAVVRIK_EXPORT_USAGE_DATA_JOB_NAME + + import litellm.proxy.proxy_server as _pserver + + await self._settings.delete() + + _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 + + 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.") + + 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.") + + 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()) 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 + 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/_http.py b/litellm/integrations/mavvrik/_http.py new file mode 100644 index 00000000000..4fa354d7652 --- /dev/null +++ b/litellm/integrations/mavvrik/_http.py @@ -0,0 +1,96 @@ +"""Shared HTTP transport for the Mavvrik integration. + +Provides a single async function with retry and exponential backoff used +by both Client (Mavvrik API calls) and Uploader (GCS calls). + + http_request(method, url, *, headers, json, params, content, timeout, label) + → httpx.Response + +Retry behaviour: + - 5xx responses and network errors: retry up to MAX_RETRIES times + - 4xx responses: returned immediately (client-side error, no retry) + - Backoff: RETRY_BACKOFF_BASE * 2^attempt seconds between retries +""" + +import asyncio +from typing import Any, Dict, Optional + +import httpx + +from litellm._logging import verbose_proxy_logger + +_MAX_RETRIES = 3 +_RETRY_BACKOFF_BASE = 1.0 # seconds; doubles each retry + + +async def http_request( + 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 an HTTP request with retry and exponential backoff. + + Args: + method: HTTP verb — GET, POST, PUT, PATCH, etc. + url: Full URL. + headers: Request headers. + json: JSON-serialisable body (mutually exclusive with content). + params: URL query parameters. + content: Raw bytes body (mutually exclusive with json). + timeout: Per-request timeout in seconds. + label: Short name used in log and error messages (e.g. "register", + "initiate"). Falls back to method when empty. + + Returns: + httpx.Response. Callers check the status themselves. + + Raises: + RuntimeError: after MAX_RETRIES failed attempts on 5xx or network errors. + """ + tag = label or method + last_exc: Exception = RuntimeError("unknown error") + + async with httpx.AsyncClient() as http: + for attempt in range(_MAX_RETRIES): + try: + resp = await http.request( + method, + url, + headers=headers, + json=json, + params=params, + content=content, + timeout=timeout, + ) + + if resp.status_code < 500: + return resp # success or 4xx — return immediately, no retry + + 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}" + ) diff --git a/litellm/integrations/mavvrik/client.py b/litellm/integrations/mavvrik/client.py new file mode 100644 index 00000000000..387c404d793 --- /dev/null +++ b/litellm/integrations/mavvrik/client.py @@ -0,0 +1,188 @@ +"""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 layer (shared by all four methods): + _request() — single httpx call with retry + exponential backoff + _assert_ok() — raises RuntimeError on unexpected status; fast-fails on 4xx +""" + +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.integrations.mavvrik._http import http_request + + +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 + + # ------------------------------------------------------------------ + # 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. + """ + # Both name and datetime are set to date_str so the GCS object path is: + # {connectionType}/{connectionId}/{type}/{date_str} + # This gives each calendar date its own object — backfills write N objects, + # not one overwritten N times. + 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 layer — delegates to shared http_request + # ------------------------------------------------------------------ + + 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 via the shared retry transport.""" + return await http_request( + method, + url, + headers=headers, + json=json, + params=params, + content=content, + timeout=timeout, + label=label, + ) + + @staticmethod + def _assert_ok( + resp: httpx.Response, + expected: Union[set, List[int]], + ) -> None: + """Raise RuntimeError when the response status is not in expected. + + Called immediately after _request() by each public method. + """ + 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..8b1d3502bb1 --- /dev/null +++ b/litellm/integrations/mavvrik/exporter.py @@ -0,0 +1,210 @@ +"""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) → AsyncIterator[str] + _get_usage_data(date_str, limit) → DataFrame + _to_csv(df, connection_id) → str + +DB not connected: all methods log a warning and return empty/None — never raise. +The scheduler skips the date gracefully; user-triggered endpoints surface the +missing-DB error through Settings._ensure_prisma_client() before reaching here. +""" + +import io +from typing import Any, AsyncIterator, List, Optional, Tuple + +import polars as pl + +from litellm._logging import verbose_proxy_logger + +# 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.model 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, + ) -> AsyncIterator[str]: + """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. + Yields nothing when DB is not connected or no rows exist for the date. + """ + client = self._prisma_client + if client is None: + verbose_proxy_logger.warning( + "Exporter: database not connected, skipping stream for %s", date_str + ) + return + + 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 + + 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. + """ + 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.""" + 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..07822b670e0 --- /dev/null +++ b/litellm/integrations/mavvrik/orchestrator.py @@ -0,0 +1,186 @@ +"""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.warning( + "Orchestrator: up to date (start=%s, end=%s), nothing to export", + start, + end, + ) + return + + verbose_logger.warning("Orchestrator: exporting %s → %s", start, end) + + for export_date in self._date_range(start, end): + await self._export(export_date) + await self._advance(export_date) + + verbose_logger.warning("Orchestrator: export complete, last date=%s", end) + + except Exception as exc: + verbose_logger.error( + "Orchestrator: pipeline failed: %s", exc, exc_info=True + ) + await self._client.report_error(str(exc)[:500]) + + # ------------------------------------------------------------------ + # Pipeline steps + # ------------------------------------------------------------------ + + async def _register(self) -> date: + marker_str = await self._client.register() + verbose_logger.warning("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.warning( + "Orchestrator: %s → streamed %d bytes to GCS ✓", date_str, total_bytes + ) + else: + verbose_logger.warning("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.warning( + "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..a7c28b68d2e --- /dev/null +++ b/litellm/integrations/mavvrik/uploader.py @@ -0,0 +1,227 @@ +"""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): + http_request() from _http.py — single httpx call with retry + exponential + backoff. Used by _initiate_resumable_upload, _finalize_upload, _put_chunk. + +GCS resumable upload protocol reference: + https://cloud.google.com/storage/docs/resumable-uploads +""" + +import gzip +import io +from typing import TYPE_CHECKING, Any, AsyncIterator + +from litellm._logging import verbose_proxy_logger +from litellm.integrations.mavvrik._http import http_request + +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 + + @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 http_request( + "POST", + signed_url, + headers={"Content-Type": "application/gzip", "x-goog-resumable": "start"}, + content=metadata, + timeout=30.0, + label="initiate", + ) + if resp.status_code != 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.""" + resp = await http_request( + "PUT", + session_uri, + headers={ + "Content-Type": "application/gzip", + "Content-Encoding": "gzip", + "x-goog-resumable": "stop", + }, + content=gzip_bytes, + timeout=120.0, + label="finalize", + ) + 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 http_request( + "PUT", + session_uri, + headers={ + "Content-Type": "application/gzip", + "Content-Range": content_range, + }, + content=chunk, + timeout=120.0, + label="chunk", + ) + 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 + session_uri: str = "" + has_data = False + + 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) + + 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 f873bfeece5..19ae28ad5d0 100644 --- a/litellm/litellm_core_utils/custom_logger_registry.py +++ b/litellm/litellm_core_utils/custom_logger_registry.py @@ -18,6 +18,7 @@ from litellm.integrations.azure_storage.azure_storage import AzureBlobStorageLog from litellm.integrations.bitbucket import BitBucketPromptManager from litellm.integrations.braintrust_logging import BraintrustLogger from litellm.integrations.cloudzero.cloudzero import CloudZeroLogger +from litellm.integrations.mavvrik import Logger as MavvrikLogger from litellm.integrations.datadog.datadog import DataDogLogger from litellm.integrations.datadog.datadog_llm_obs import DataDogLLMObsLogger from litellm.integrations.datadog.datadog_metrics import DatadogMetricsLogger @@ -100,6 +101,7 @@ class CustomLoggerRegistry: "gitlab": GitLabPromptManager, "cloudzero": CloudZeroLogger, "focus": FocusLogger, + "mavvrik": MavvrikLogger, "vantage": VantageLogger, "posthog": PostHogLogger, } diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 625cb83724b..d8353a52ef0 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -3973,6 +3973,15 @@ def _init_custom_logger_compatible_class( # noqa: PLR0915 galileo_logger = GalileoObserve() _in_memory_loggers.append(galileo_logger) return galileo_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 == "cloudzero": from litellm.integrations.cloudzero.cloudzero import CloudZeroLogger @@ -4360,6 +4369,12 @@ def get_custom_logger_compatible_class( # noqa: PLR0915 for callback in _in_memory_loggers: if isinstance(callback, GalileoObserve): 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 == "cloudzero": from litellm.integrations.cloudzero.cloudzero import CloudZeroLogger diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index aa8122d8fd9..0f4aa544b4c 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -478,6 +478,7 @@ from litellm.proxy.search_endpoints.search_tool_management import ( router as search_tool_management_router, ) from litellm.proxy.spend_tracking.cloudzero_endpoints import router as cloudzero_router +from litellm.proxy.spend_tracking.mavvrik_endpoints import router as mavvrik_router from litellm.proxy.spend_tracking.spend_management_endpoints import ( router as spend_management_router, ) @@ -6717,6 +6718,49 @@ class ProxyStartupEvent: ) await VantageLogger.init_vantage_background_job(scheduler=scheduler) + ######################################################## + # Mavvrik Background Job + ######################################################## + from litellm.constants import ( + MAVVRIK_EXPORT_INTERVAL_MINUTES, + MAVVRIK_EXPORT_USAGE_DATA_JOB_NAME, + ) + from litellm.integrations.mavvrik import ( + Client as MavvrikClient, + Orchestrator as MavvrikOrchestrator, + Settings as MavvrikSettings, + Uploader as MavvrikUploader, + ) + + settings = MavvrikSettings() + if await settings.is_setup(): + # Skip DB load when credentials come from env vars — avoids raising + # if prisma_client is not yet connected at startup. + data = {} if settings.has_env_vars else await settings.load() + import os as _os + + client = MavvrikClient( + api_key=data.get("api_key") or _os.getenv("MAVVRIK_API_KEY", ""), + api_endpoint=data.get("api_endpoint") + or _os.getenv("MAVVRIK_API_ENDPOINT", ""), + connection_id=data.get("connection_id") + or _os.getenv("MAVVRIK_CONNECTION_ID", ""), + ) + uploader = MavvrikUploader(client=client) + orchestrator = MavvrikOrchestrator(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.warning( + "Mavvrik: background export job scheduled every %d min (connection_id=%s)", + MAVVRIK_EXPORT_INTERVAL_MINUTES, + client.connection_id, + ) + ######################################################## # Prometheus Background Job ######################################################## @@ -14131,6 +14175,7 @@ app.include_router(customer_router) app.include_router(spend_management_router) app.include_router(cloudzero_router) app.include_router(vantage_router) +app.include_router(mavvrik_router) app.include_router(caching_router) app.include_router(analytics_router) app.include_router(guardrails_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..3e727f7aca4 --- /dev/null +++ b/litellm/proxy/spend_tracking/mavvrik_endpoints.py @@ -0,0 +1,224 @@ +"""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.integrations.mavvrik import Service +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() + +# --------------------------------------------------------------------------- +# 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 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 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 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 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 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 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..24e854113e5 --- /dev/null +++ b/litellm/types/proxy/mavvrik_endpoints.py @@ -0,0 +1,128 @@ +""" +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 + + +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, + 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 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..ba03e7e115d --- /dev/null +++ b/tests/test_litellm/integrations/mavvrik/conftest.py @@ -0,0 +1,3 @@ +# Local conftest for mavvrik integration tests. +# Avoids importing the top-level conftest which requires the full litellm package +# with all optional dependencies. diff --git a/tests/test_litellm/integrations/mavvrik/test_e2e_upload.py b/tests/test_litellm/integrations/mavvrik/test_e2e_upload.py new file mode 100644 index 00000000000..467424ca157 --- /dev/null +++ b/tests/test_litellm/integrations/mavvrik/test_e2e_upload.py @@ -0,0 +1,174 @@ +"""End-to-end tests for the Mavvrik upload layer against the real API. + +These tests hit the live Mavvrik API. They are skipped automatically +when the required environment variables are absent so they never break CI. + +Set the following env vars before running: + + MAVVRIK_API_KEY= + MAVVRIK_API_ENDPOINT=https://api.mavvrik.dev/ + MAVVRIK_CONNECTION_ID= + +Run with: + poetry run pytest tests/test_litellm/integrations/mavvrik/test_e2e_upload.py -v -s +""" + +import calendar +import os +import sys +from datetime import date, datetime, timedelta + +import pytest + +sys.path.insert(0, os.path.abspath("../../../..")) + +from litellm.integrations.mavvrik.client import Client +from litellm.integrations.mavvrik.uploader import Uploader + +# --------------------------------------------------------------------------- +# Credentials — populated from env vars; test is skipped if any are absent. +# --------------------------------------------------------------------------- + +API_KEY = os.getenv("MAVVRIK_API_KEY", "") +API_ENDPOINT = os.getenv("MAVVRIK_API_ENDPOINT", "") +CONNECTION_ID = os.getenv("MAVVRIK_CONNECTION_ID", "") + +_CREDS_PRESENT = all([API_KEY, API_ENDPOINT, CONNECTION_ID]) +_skip_if_no_creds = pytest.mark.skipif( + not _CREDS_PRESENT, + reason="Mavvrik credentials not configured — set MAVVRIK_API_KEY, MAVVRIK_API_ENDPOINT, MAVVRIK_CONNECTION_ID", +) + +# Name used for the synthetic GCS object. +# Prefixed with "test-" so it is clearly not real data. +_TEST_DATE = "test-e2e-litellm" + +# Minimal synthetic CSV that matches the Mavvrik schema column order +_TEST_CSV = ( + "date,user_id,api_key,model,model_group,custom_llm_provider," + "prompt_tokens,completion_tokens,spend,api_requests,successful_requests," + "failed_requests,cache_creation_input_tokens,cache_read_input_tokens," + "created_at,updated_at,team_id,api_key_alias,team_alias,user_email\n" + "2026-01-01,user-e2e,sk-test,gpt-4o,gpt-4o,openai," + "100,50,0.0025,1,1,0,0,0," + "2026-01-01T00:00:00Z,2026-01-01T00:01:00Z,team-e2e,e2e-key,e2e-team,e2e@example.com\n" +) + + +# --------------------------------------------------------------------------- +# Fixtures +# --------------------------------------------------------------------------- + + +@pytest.fixture(scope="module") +def client(): + return Client( + api_key=API_KEY, + api_endpoint=API_ENDPOINT, + connection_id=CONNECTION_ID, + ) + + +@pytest.fixture(scope="module") +def uploader(client): + return Uploader(client=client) + + +# --------------------------------------------------------------------------- +# Tests +# --------------------------------------------------------------------------- + + +@_skip_if_no_creds +class TestE2ERegister: + @pytest.mark.asyncio + async def test_register_returns_iso_string_or_none(self, client): + """register() must return an ISO-8601 date string or None (first run).""" + marker = await client.register() + print(f"\n register() returned marker: {marker}") + + if marker is not None: + dt = datetime.fromisoformat(marker) + assert dt.year >= 2020, f"Unexpected marker year: {dt.year}" + + @pytest.mark.asyncio + async def test_register_twice_is_idempotent(self, client): + """Calling register() twice should succeed without error.""" + m1 = await client.register() + m2 = await client.register() + print(f"\n First call: {m1}") + print(f" Second call: {m2}") + if m1 is not None: + datetime.fromisoformat(m1) + if m2 is not None: + datetime.fromisoformat(m2) + + +@_skip_if_no_creds +class TestE2EGetSignedUrl: + @pytest.mark.asyncio + async def test_get_signed_url_returns_url(self, client): + """get_signed_url() must return a GCS URL for a given date name.""" + url = await client.get_signed_url(_TEST_DATE) + print(f"\n signed URL: {url[:80]}...") + assert url.startswith("https://"), f"Expected https URL, got: {url[:40]}" + + +@_skip_if_no_creds +class TestE2EUpload: + @pytest.mark.asyncio + async def test_upload_synthetic_csv(self, uploader): + """Full 3-step GCS upload: get signed URL → initiate → PUT gzip bytes.""" + await uploader.upload(_TEST_CSV, date_str=_TEST_DATE) + print(f"\n Upload for date {_TEST_DATE} succeeded") + + @pytest.mark.asyncio + async def test_upload_same_date_twice_is_idempotent(self, uploader): + """Re-uploading the same date must succeed (GCS object is overwritten).""" + await uploader.upload(_TEST_CSV, date_str=_TEST_DATE) + await uploader.upload(_TEST_CSV, date_str=_TEST_DATE) + print(f"\n Two uploads for date {_TEST_DATE} both succeeded (idempotent)") + + @pytest.mark.asyncio + async def test_upload_empty_payload_is_noop(self, uploader): + """Empty payload must return without making any network calls.""" + await uploader.upload(" ", date_str=_TEST_DATE) + print("\n Empty payload correctly skipped") + + +@_skip_if_no_creds +class TestE2EAdvanceMarker: + @pytest.mark.asyncio + async def test_advance_marker_succeeds(self, client): + """advance_marker() must PATCH Mavvrik without raising.""" + epoch = 1700000000 + await client.advance_marker(epoch) + print(f"\n advance_marker({epoch}) succeeded") + + @pytest.mark.asyncio + async def test_advance_marker_with_recent_date(self, client): + """advance_marker() with a recent epoch must also succeed.""" + yesterday = date.today() - timedelta(days=1) + epoch = int(calendar.timegm(yesterday.timetuple())) + await client.advance_marker(epoch) + print(f"\n advance_marker({epoch}) for {yesterday} succeeded") + + +@_skip_if_no_creds +class TestE2EFullFlow: + @pytest.mark.asyncio + async def test_register_then_upload_then_advance(self, client, uploader): + """Simulate one complete scheduled export cycle end-to-end.""" + marker_iso = await client.register() + if marker_iso is not None: + datetime.fromisoformat(marker_iso) + print(f"\n register() marker: {marker_iso}") + + await uploader.upload(_TEST_CSV, date_str=_TEST_DATE) + print(f" upload() for {_TEST_DATE}: OK") + + export_epoch = 1700000000 + await client.advance_marker(export_epoch) + print(f" advance_marker({export_epoch}): OK") + + print("\n Full cycle PASSED") diff --git a/tests/test_litellm/integrations/mavvrik/test_endpoints.py b/tests/test_litellm/integrations/mavvrik/test_endpoints.py new file mode 100644 index 00000000000..b6a01afa9ad --- /dev/null +++ b/tests/test_litellm/integrations/mavvrik/test_endpoints.py @@ -0,0 +1,457 @@ +"""Unit tests for Mavvrik FastAPI admin endpoints.""" + +import os +import sys +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +sys.path.insert(0, os.path.abspath("../../../..")) + +from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth +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, +) + +# 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", + ) + + with patch.object(Settings, "save", new=AsyncMock()), patch( + "litellm.proxy.proxy_server.scheduler", None + ): + 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 diff --git a/tests/test_litellm/integrations/mavvrik/test_http.py b/tests/test_litellm/integrations/mavvrik/test_http.py new file mode 100644 index 00000000000..d2ebcda0c9f --- /dev/null +++ b/tests/test_litellm/integrations/mavvrik/test_http.py @@ -0,0 +1,182 @@ +"""Unit tests for mavvrik._http.http_request — shared retry transport.""" + +import os +import sys +from unittest.mock import AsyncMock, MagicMock, patch + +import httpx +import pytest + +sys.path.insert(0, os.path.abspath("../../../..")) + +from litellm.integrations.mavvrik._http import http_request + + +def _mock_response(status_code: int, text: str = "") -> MagicMock: + resp = MagicMock(spec=httpx.Response) + resp.status_code = status_code + resp.text = text + return resp + + +def _mock_http(return_value=None, side_effect=None): + """Return a patched httpx.AsyncClient context manager.""" + http = MagicMock() + if side_effect: + http.request = AsyncMock(side_effect=side_effect) + else: + http.request = AsyncMock(return_value=return_value) + ctx = MagicMock() + ctx.__aenter__ = AsyncMock(return_value=http) + ctx.__aexit__ = AsyncMock(return_value=False) + return ctx, http + + +# --------------------------------------------------------------------------- +# Success paths +# --------------------------------------------------------------------------- + + +class TestHttpRequestSuccess: + @pytest.mark.asyncio + async def test_returns_response_on_2xx(self): + ctx, _ = _mock_http(return_value=_mock_response(200)) + with patch("httpx.AsyncClient", return_value=ctx): + resp = await http_request("GET", "https://example.com") + assert resp.status_code == 200 + + @pytest.mark.asyncio + async def test_returns_4xx_without_retry(self): + ctx, http = _mock_http(return_value=_mock_response(401, "Unauthorized")) + with patch("httpx.AsyncClient", return_value=ctx), patch( + "asyncio.sleep", new_callable=AsyncMock + ) as mock_sleep: + resp = await http_request("GET", "https://example.com") + assert resp.status_code == 401 + assert http.request.call_count == 1 + mock_sleep.assert_not_called() + + @pytest.mark.asyncio + async def test_passes_headers(self): + captured = [] + + async def fake_request(method, url, headers=None, **kwargs): + captured.append(headers) + return _mock_response(200) + + ctx = MagicMock() + http = MagicMock() + http.request = fake_request + ctx.__aenter__ = AsyncMock(return_value=http) + ctx.__aexit__ = AsyncMock(return_value=False) + + with patch("httpx.AsyncClient", return_value=ctx): + await http_request( + "POST", "https://example.com", headers={"x-api-key": "secret"} + ) + + assert captured[0]["x-api-key"] == "secret" + + @pytest.mark.asyncio + async def test_passes_json_and_params(self): + captured = [] + + async def fake_request(method, url, json=None, params=None, **kwargs): + captured.append({"json": json, "params": params}) + return _mock_response(200) + + ctx = MagicMock() + http = MagicMock() + http.request = fake_request + ctx.__aenter__ = AsyncMock(return_value=http) + ctx.__aexit__ = AsyncMock(return_value=False) + + with patch("httpx.AsyncClient", return_value=ctx): + await http_request( + "GET", "https://example.com", json={"key": "val"}, params={"q": "1"} + ) + + assert captured[0]["json"] == {"key": "val"} + assert captured[0]["params"] == {"q": "1"} + + @pytest.mark.asyncio + async def test_passes_content(self): + captured = [] + + async def fake_request(method, url, content=None, **kwargs): + captured.append(content) + return _mock_response(201) + + ctx = MagicMock() + http = MagicMock() + http.request = fake_request + ctx.__aenter__ = AsyncMock(return_value=http) + ctx.__aexit__ = AsyncMock(return_value=False) + + with patch("httpx.AsyncClient", return_value=ctx): + await http_request("PUT", "https://example.com", content=b"gzip-data") + + assert captured[0] == b"gzip-data" + + +# --------------------------------------------------------------------------- +# Retry on 5xx +# --------------------------------------------------------------------------- + + +class TestHttpRequestRetry: + @pytest.mark.asyncio + async def test_retries_on_5xx_then_raises(self): + ctx, http = _mock_http(return_value=_mock_response(503, "unavailable")) + with patch("httpx.AsyncClient", return_value=ctx), patch( + "asyncio.sleep", new_callable=AsyncMock + ): + with pytest.raises(RuntimeError, match="failed after"): + await http_request("GET", "https://example.com") + assert http.request.call_count == 3 + + @pytest.mark.asyncio + async def test_retries_on_network_error_then_raises(self): + ctx, http = _mock_http(side_effect=httpx.ConnectError("timeout")) + with patch("httpx.AsyncClient", return_value=ctx), patch( + "asyncio.sleep", new_callable=AsyncMock + ): + with pytest.raises(RuntimeError, match="failed after"): + await http_request("GET", "https://example.com") + assert http.request.call_count == 3 + + @pytest.mark.asyncio + async def test_succeeds_on_second_attempt(self): + fail = _mock_response(503, "err") + ok = _mock_response(200) + ctx, http = _mock_http() + http.request = AsyncMock(side_effect=[fail, ok]) + with patch("httpx.AsyncClient", return_value=ctx), patch( + "asyncio.sleep", new_callable=AsyncMock + ): + resp = await http_request("GET", "https://example.com") + assert resp.status_code == 200 + assert http.request.call_count == 2 + + @pytest.mark.asyncio + async def test_uses_exponential_backoff(self): + ctx, http = _mock_http(return_value=_mock_response(503, "err")) + sleep_calls = [] + with patch("httpx.AsyncClient", return_value=ctx), patch( + "asyncio.sleep", + new_callable=AsyncMock, + side_effect=lambda s: sleep_calls.append(s), + ): + with pytest.raises(RuntimeError): + await http_request("GET", "https://example.com") + # 3 attempts → 2 sleeps: 1.0s and 2.0s + assert sleep_calls == [1.0, 2.0] + + @pytest.mark.asyncio + async def test_error_message_contains_label(self): + ctx, _ = _mock_http(return_value=_mock_response(503, "err")) + with patch("httpx.AsyncClient", return_value=ctx), patch( + "asyncio.sleep", new_callable=AsyncMock + ): + with pytest.raises(RuntimeError, match="initiate"): + await http_request("POST", "https://example.com", label="initiate") diff --git a/tests/test_litellm/integrations/mavvrik/test_logger.py b/tests/test_litellm/integrations/mavvrik/test_logger.py new file mode 100644 index 00000000000..7dee475256b --- /dev/null +++ b/tests/test_litellm/integrations/mavvrik/test_logger.py @@ -0,0 +1,24 @@ +"""Unit tests for Mavvrik Logger marker class.""" + +import os +import sys + +import pytest + +sys.path.insert(0, os.path.abspath("../../../..")) + +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): + """Logger is registered as the 'mavvrik' callback in custom_logger_registry.""" + assert Logger is not None diff --git a/tests/test_litellm/integrations/mavvrik/test_scheduler.py b/tests/test_litellm/integrations/mavvrik/test_scheduler.py new file mode 100644 index 00000000000..a8327d96b0d --- /dev/null +++ b/tests/test_litellm/integrations/mavvrik/test_scheduler.py @@ -0,0 +1,369 @@ +"""Unit tests for Orchestrator — pipeline sequencing and pod lock.""" + +import os +import sys +from datetime import date, datetime, timezone +from unittest.mock import AsyncMock, MagicMock, patch + +import polars as pl +import pytest + +sys.path.insert(0, os.path.abspath("../../../..")) + +from litellm.integrations.mavvrik.orchestrator import Orchestrator +from litellm.integrations.mavvrik.client import Client +from litellm.integrations.mavvrik.uploader import Uploader + + +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_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_skips_upload_when_no_data(self): + """When export returns 0 bytes, advance still called (date was processed).""" + 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 always called — even for empty dates + orc._client.advance_marker.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 "API down" 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 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..3d01c235e46 --- /dev/null +++ b/tests/test_litellm/integrations/mavvrik/test_service.py @@ -0,0 +1,299 @@ +"""Tests for Service facade — verifies constructors and method calls are correct.""" + +import os +import sys +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.uploader import Uploader +from litellm.integrations.mavvrik.orchestrator import Orchestrator + + +_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"] == [] 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_transform.py b/tests/test_litellm/integrations/mavvrik/test_transform.py new file mode 100644 index 00000000000..7e4687fa742 --- /dev/null +++ b/tests/test_litellm/integrations/mavvrik/test_transform.py @@ -0,0 +1,474 @@ +"""Unit tests for the Mavvrik transform layer (CSV output).""" + +import io +import os +import sys +from unittest.mock import AsyncMock, MagicMock, patch + +import polars as pl +import pytest + +sys.path.insert(0, os.path.abspath("../../../..")) + +from litellm.integrations.mavvrik.exporter import Exporter + + +def _make_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_df(successful_requests=[3]) + result = transformer._to_csv(df) + assert result != "" + + def test_output_has_header_row(self): + transformer = Exporter() + df = _make_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_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_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_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_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_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_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_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_yields_nothing_when_no_db(self): + """_stream_pages yields nothing when DB not connected.""" + exporter = Exporter() + with patch.object( + type(exporter), + "_prisma_client", + new_callable=lambda: property(lambda self: None), + ): + chunks = [] + async for chunk in exporter._stream_pages("2026-04-10", connection_id="c"): + chunks.append(chunk) + assert chunks == [] + + @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_upload.py b/tests/test_litellm/integrations/mavvrik/test_upload.py new file mode 100644 index 00000000000..1f00079dee9 --- /dev/null +++ b/tests/test_litellm/integrations/mavvrik/test_upload.py @@ -0,0 +1,393 @@ +"""Unit tests for Mavvrik Client — Mavvrik API HTTP calls.""" + +import sys +import os +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 + + +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 — delegates to http_request (retry behaviour tested in test_http.py) +# --------------------------------------------------------------------------- + + +class TestClientRequest: + @pytest.mark.asyncio + async def test_delegates_to_http_request(self): + """_request() delegates to the shared http_request transport.""" + c = _make_client() + mock_resp = _mock_response(200) + + with patch( + "litellm.integrations.mavvrik.client.http_request", + new_callable=AsyncMock, + return_value=mock_resp, + ) as mock_http: + resp = await c._request( + "GET", + "https://example.com", + headers={"x-api-key": "k"}, + json={"a": 1}, + label="test", + ) + + assert resp.status_code == 200 + mock_http.assert_called_once_with( + "GET", + "https://example.com", + headers={"x-api-key": "k"}, + json={"a": 1}, + params=None, + content=None, + timeout=30.0, + label="test", + ) + + +# --------------------------------------------------------------------------- +# _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] 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..5bb50339693 --- /dev/null +++ b/tests/test_litellm/integrations/mavvrik/test_uploader.py @@ -0,0 +1,562 @@ +"""Unit tests for Mavvrik Uploader — GCS resumable upload protocol.""" + +import gzip +import os +import sys +from unittest.mock import AsyncMock, MagicMock, patch + +import httpx +import pytest + +sys.path.insert(0, os.path.abspath("../../../..")) + +from litellm.integrations.mavvrik.uploader import Uploader +from litellm.integrations.mavvrik.client import Client + + +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 + + +# --------------------------------------------------------------------------- +# 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"} + ) + with patch( + "litellm.integrations.mavvrik.uploader.http_request", + new_callable=AsyncMock, + 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={}) + with patch( + "litellm.integrations.mavvrik.uploader.http_request", + new_callable=AsyncMock, + 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") + with patch( + "litellm.integrations.mavvrik.uploader.http_request", + new_callable=AsyncMock, + 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 + + with patch( + "litellm.integrations.mavvrik.uploader.http_request", + 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() + with patch( + "litellm.integrations.mavvrik.uploader.http_request", + new_callable=AsyncMock, + 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() + with patch( + "litellm.integrations.mavvrik.uploader.http_request", + new_callable=AsyncMock, + 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() + with patch( + "litellm.integrations.mavvrik.uploader.http_request", + new_callable=AsyncMock, + 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() + with patch( + "litellm.integrations.mavvrik.uploader.http_request", + new_callable=AsyncMock, + 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() + with patch( + "litellm.integrations.mavvrik.uploader.http_request", + new_callable=AsyncMock, + 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 from _gcs_request — raises in _finalize.""" + u = _make_uploader() + with patch( + "litellm.integrations.mavvrik.uploader.http_request", + new_callable=AsyncMock, + 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) + + with patch( + "litellm.integrations.mavvrik.uploader.http_request", + 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() + with patch( + "litellm.integrations.mavvrik.uploader.http_request", + new_callable=AsyncMock, + 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() + with patch( + "litellm.integrations.mavvrik.uploader.http_request", + new_callable=AsyncMock, + 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) + + with patch( + "litellm.integrations.mavvrik.uploader.http_request", + 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) + + with patch( + "litellm.integrations.mavvrik.uploader.http_request", + 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() + with patch( + "litellm.integrations.mavvrik.uploader.http_request", + new_callable=AsyncMock, + 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_request causes _put_chunk to raise.""" + u = _make_uploader() + with patch( + "litellm.integrations.mavvrik.uploader.http_request", + new_callable=AsyncMock, + return_value=_mock_http_response(403, text="Forbidden"), + ): + with pytest.raises(RuntimeError): + await u._put_chunk("https://session", b"data", offset=0, final=False) + + +# _gcs_request removed — transport tested in test_http.py via http_request + + +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)