mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-04 02:31:27 +00:00
feat(focus): add Mavvrik destination for FOCUS export
This commit is contained in:
Praveen Ghuge
parent
8dd7ab7f3f
commit
6e9ee80754
11 changed files with 1407 additions and 3 deletions
|
|
@ -154,6 +154,7 @@ _custom_logger_compatible_callbacks_literal = Literal[
|
|||
"gitlab",
|
||||
"cloudzero",
|
||||
"focus",
|
||||
"mavvrik",
|
||||
"vantage",
|
||||
"posthog",
|
||||
"levo",
|
||||
|
|
|
|||
|
|
@ -1480,6 +1480,7 @@ DB_SPEND_UPDATE_JOB_NAME = "db_spend_update_job"
|
|||
DB_DAILY_TAG_SPEND_UPDATE_JOB_NAME = "db_daily_tag_spend_update_job"
|
||||
PROMETHEUS_EMIT_BUDGET_METRICS_JOB_NAME = "prometheus_emit_budget_metrics"
|
||||
CLOUDZERO_EXPORT_USAGE_DATA_JOB_NAME = "cloudzero_export_usage_data"
|
||||
MAVVRIK_FOCUS_EXPORT_JOB_NAME = "mavvrik_focus_export_usage_data"
|
||||
CLOUDZERO_MAX_FETCHED_DATA_RECORDS = int(
|
||||
os.getenv("CLOUDZERO_MAX_FETCHED_DATA_RECORDS", 50000)
|
||||
)
|
||||
|
|
|
|||
|
|
@ -4,6 +4,7 @@ from .base import FocusDestination, FocusTimeWindow
|
|||
from .factory import FocusDestinationFactory
|
||||
from .gcs_destination import FocusGCSDestination
|
||||
from .s3_destination import FocusS3Destination
|
||||
from .mavvrik_destination import FocusMavvrikDestination
|
||||
from .vantage_destination import FocusVantageDestination
|
||||
|
||||
__all__ = [
|
||||
|
|
@ -12,5 +13,6 @@ __all__ = [
|
|||
"FocusGCSDestination",
|
||||
"FocusTimeWindow",
|
||||
"FocusS3Destination",
|
||||
"FocusMavvrikDestination",
|
||||
"FocusVantageDestination",
|
||||
]
|
||||
|
|
|
|||
|
|
@ -8,6 +8,7 @@ from typing import Any, Dict, Optional
|
|||
from .base import FocusDestination
|
||||
from .gcs_destination import FocusGCSDestination
|
||||
from .s3_destination import FocusS3Destination
|
||||
from .mavvrik_destination import FocusMavvrikDestination
|
||||
from .vantage_destination import FocusVantageDestination
|
||||
|
||||
|
||||
|
|
@ -32,6 +33,8 @@ class FocusDestinationFactory:
|
|||
return FocusVantageDestination(prefix=prefix, config=normalized_config)
|
||||
if provider_lower == "gcs":
|
||||
return FocusGCSDestination(prefix=prefix, config=normalized_config)
|
||||
if provider_lower == "mavvrik":
|
||||
return FocusMavvrikDestination(prefix=prefix, config=normalized_config)
|
||||
raise NotImplementedError(
|
||||
f"Provider '{provider}' not supported for Focus export"
|
||||
)
|
||||
|
|
@ -87,6 +90,15 @@ class FocusDestinationFactory:
|
|||
"FOCUS_GCS_BUCKET_NAME must be provided for GCS exports"
|
||||
)
|
||||
return {k: v for k, v in resolved.items() if v is not None}
|
||||
if provider == "mavvrik":
|
||||
resolved = {
|
||||
"api_key": overrides.get("api_key") or os.getenv("MAVVRIK_API_KEY"),
|
||||
"api_endpoint": overrides.get("api_endpoint")
|
||||
or os.getenv("MAVVRIK_API_ENDPOINT"),
|
||||
"connection_id": overrides.get("connection_id")
|
||||
or os.getenv("MAVVRIK_CONNECTION_ID"),
|
||||
}
|
||||
return {k: v for k, v in resolved.items() if v is not None}
|
||||
raise NotImplementedError(
|
||||
f"Provider '{provider}' not supported for Focus export configuration"
|
||||
)
|
||||
|
|
|
|||
345
litellm/integrations/focus/destinations/mavvrik_destination.py
Normal file
345
litellm/integrations/focus/destinations/mavvrik_destination.py
Normal file
|
|
@ -0,0 +1,345 @@
|
|||
"""Mavvrik GCS destination for FOCUS export.
|
||||
|
||||
Flow:
|
||||
1. GET /metrics/agent/ai/{connection_id}/upload-url → GCS signed URL
|
||||
2. PUT <signed_url> with CSV content
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import gzip
|
||||
from typing import Any, Optional
|
||||
from urllib.parse import urlparse
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.llms.custom_httpx.http_handler import (
|
||||
AsyncHTTPHandler,
|
||||
get_async_httpx_client,
|
||||
httpxSpecialProvider,
|
||||
)
|
||||
|
||||
from .base import FocusDestination, FocusTimeWindow
|
||||
|
||||
_MAVVRIK_ALLOWED_SUFFIXES = (".mavvrik.dev", ".mavvrik.ai", ".mavvrik.app")
|
||||
|
||||
# GCS requires intermediate chunks to be a multiple of 256 KB.
|
||||
# 8 MB gives a good balance between round-trips and memory pressure.
|
||||
_GCS_CHUNK_SIZE = 8 * 1024 * 1024 # 8 MB
|
||||
|
||||
|
||||
def _validate_api_endpoint(api_endpoint: str) -> None:
|
||||
if not api_endpoint.startswith("https://"):
|
||||
raise ValueError("MAVVRIK_API_ENDPOINT must be an HTTPS URL")
|
||||
hostname = (urlparse(api_endpoint).hostname or "").lower()
|
||||
if not any(hostname.endswith(suffix) for suffix in _MAVVRIK_ALLOWED_SUFFIXES):
|
||||
raise ValueError(
|
||||
"MAVVRIK_API_ENDPOINT host must be a Mavvrik domain "
|
||||
"(e.g. https://api.mavvrik.dev/<tenant_id>)"
|
||||
)
|
||||
|
||||
|
||||
def _validate_gcs_url(url: str, label: str) -> None:
|
||||
parsed = urlparse(url)
|
||||
if parsed.scheme != "https":
|
||||
raise ValueError(
|
||||
f"Mavvrik FOCUS destination: {label} must be HTTPS, got scheme '{parsed.scheme}'"
|
||||
)
|
||||
hostname = (parsed.hostname or "").lower()
|
||||
if not (
|
||||
hostname == "storage.googleapis.com"
|
||||
or hostname.endswith(".storage.googleapis.com")
|
||||
):
|
||||
raise ValueError(
|
||||
f"Mavvrik FOCUS destination: {label} must be a GCS endpoint "
|
||||
f"(storage.googleapis.com), got '{hostname}'"
|
||||
)
|
||||
|
||||
|
||||
class FocusMavvrikDestination(FocusDestination):
|
||||
"""Upload FOCUS CSV exports to Mavvrik via GCS signed URL."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
prefix: str,
|
||||
config: Optional[dict[str, Any]] = None,
|
||||
) -> None:
|
||||
config = config or {}
|
||||
api_key = config.get("api_key")
|
||||
api_endpoint = config.get("api_endpoint")
|
||||
connection_id = config.get("connection_id")
|
||||
|
||||
if not api_key:
|
||||
raise ValueError(
|
||||
"MAVVRIK_API_KEY must be provided for Mavvrik FOCUS destination "
|
||||
"(set MAVVRIK_API_KEY env var or pass in destination_config)"
|
||||
)
|
||||
if not api_endpoint:
|
||||
raise ValueError(
|
||||
"MAVVRIK_API_ENDPOINT must be provided for Mavvrik FOCUS destination "
|
||||
"(set MAVVRIK_API_ENDPOINT env var or pass in destination_config)"
|
||||
)
|
||||
if not connection_id:
|
||||
raise ValueError(
|
||||
"MAVVRIK_CONNECTION_ID must be provided for Mavvrik FOCUS destination "
|
||||
"(set MAVVRIK_CONNECTION_ID env var or pass in destination_config)"
|
||||
)
|
||||
|
||||
_validate_api_endpoint(api_endpoint)
|
||||
|
||||
self.api_key = api_key
|
||||
self.api_endpoint = api_endpoint.rstrip("/")
|
||||
self.connection_id = connection_id
|
||||
self.prefix = prefix
|
||||
self._http: AsyncHTTPHandler = get_async_httpx_client(
|
||||
llm_provider=httpxSpecialProvider.LoggingCallback
|
||||
)
|
||||
self._registered = False
|
||||
|
||||
@property
|
||||
def _agent_url(self) -> str:
|
||||
return f"{self.api_endpoint}/metrics/agent/ai/{self.connection_id}"
|
||||
|
||||
@property
|
||||
def _upload_url_endpoint(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}
|
||||
|
||||
async def _ensure_registered(self) -> Optional[int]:
|
||||
"""POST agent endpoint to register/initialize the connector (once per instance).
|
||||
|
||||
Returns metricsMarker from the Mavvrik response — the last date index
|
||||
Mavvrik has successfully processed. Used by the logger to catch up any
|
||||
dates that were missed due to previous export failures.
|
||||
|
||||
Returns None if the connector was already registered (cached).
|
||||
"""
|
||||
if self._registered:
|
||||
return None
|
||||
resp = await self._http.client.request(
|
||||
method="POST",
|
||||
url=self._agent_url,
|
||||
headers=self._auth_headers,
|
||||
json={"name": self.connection_id},
|
||||
timeout=30.0,
|
||||
)
|
||||
if resp.status_code == 410:
|
||||
# Connector has been disconnected in Mavvrik — reset flag so next
|
||||
# delivery attempt re-registers after it becomes active again.
|
||||
self._registered = False
|
||||
raise RuntimeError(
|
||||
"Mavvrik FOCUS destination: connector is disconnected (410). "
|
||||
"Re-enable the connection in the Mavvrik dashboard."
|
||||
)
|
||||
if resp.status_code >= 400:
|
||||
raise RuntimeError(
|
||||
f"Mavvrik FOCUS destination: register failed "
|
||||
f"({resp.status_code}): {resp.text[:200]}"
|
||||
)
|
||||
self._registered = True
|
||||
metrics_marker = resp.json().get("metricsMarker", 0)
|
||||
verbose_logger.debug(
|
||||
"Mavvrik FOCUS destination: connector registered (metricsMarker=%s)",
|
||||
metrics_marker,
|
||||
)
|
||||
return metrics_marker
|
||||
|
||||
async def _get_signed_url(self, date_str: str) -> str:
|
||||
"""GET upload-url endpoint → GCS signed URL for the given date."""
|
||||
params = {"name": date_str, "type": "metrics", "datetime": date_str}
|
||||
resp = await self._http.client.request(
|
||||
method="GET",
|
||||
url=self._upload_url_endpoint,
|
||||
headers=self._auth_headers,
|
||||
params=params,
|
||||
timeout=30.0,
|
||||
)
|
||||
if resp.status_code >= 400:
|
||||
raise RuntimeError(
|
||||
f"Mavvrik FOCUS destination: failed to get signed URL "
|
||||
f"({resp.status_code}): {resp.text[:200]}"
|
||||
)
|
||||
signed_url = resp.json().get("url")
|
||||
if not signed_url:
|
||||
raise RuntimeError(
|
||||
f"Mavvrik FOCUS destination: response missing 'url' field: {resp.json()}"
|
||||
)
|
||||
_validate_gcs_url(signed_url, "signed URL")
|
||||
verbose_logger.debug(
|
||||
"Mavvrik FOCUS destination: got signed URL for date %s", date_str
|
||||
)
|
||||
return signed_url
|
||||
|
||||
async def _upload_to_gcs(self, signed_url: str, content: bytes) -> None:
|
||||
"""Upload gzip-compressed CSV to GCS via chunked resumable upload.
|
||||
|
||||
The full CSV is gzip-compressed first, then uploaded in _GCS_CHUNK_SIZE
|
||||
chunks using the GCS resumable upload protocol. GCS assembles the chunks
|
||||
server-side into a single complete object — the bucket receives one file
|
||||
regardless of how many chunks were sent.
|
||||
|
||||
Intermediate chunks: Content-Range: bytes X-Y/* → expect 308
|
||||
Final chunk: Content-Range: bytes X-Y/T → expect 200/201
|
||||
|
||||
This handles exports larger than available memory for a single PUT while
|
||||
keeping the destination code self-contained (no changes to the FOCUS
|
||||
pipeline upstream).
|
||||
"""
|
||||
gzip_bytes = gzip.compress(content)
|
||||
total = len(gzip_bytes)
|
||||
|
||||
# Step 1: initiate resumable upload session
|
||||
metadata = b'{"contentEncoding":"gzip","contentDisposition":"attachment"}'
|
||||
init_resp = await self._http.client.request(
|
||||
method="POST",
|
||||
url=signed_url,
|
||||
headers={
|
||||
"Content-Type": "application/gzip",
|
||||
"x-goog-resumable": "start",
|
||||
},
|
||||
content=metadata,
|
||||
timeout=30.0,
|
||||
)
|
||||
if init_resp.status_code not in (200, 201):
|
||||
raise RuntimeError(
|
||||
f"Mavvrik FOCUS destination: GCS session init failed "
|
||||
f"({init_resp.status_code}): {init_resp.text[:400]}"
|
||||
)
|
||||
|
||||
session_uri = init_resp.headers.get("Location")
|
||||
if not session_uri:
|
||||
raise RuntimeError(
|
||||
"Mavvrik FOCUS destination: GCS session init missing Location header"
|
||||
)
|
||||
_validate_gcs_url(session_uri, "session URI")
|
||||
|
||||
verbose_logger.debug(
|
||||
"Mavvrik FOCUS destination: GCS session started, uploading %d gzip bytes "
|
||||
"in %d chunk(s)",
|
||||
total,
|
||||
max(1, -(-total // _GCS_CHUNK_SIZE)), # ceiling division
|
||||
)
|
||||
|
||||
# Step 2: upload in chunks; cancel session on any failure to avoid
|
||||
# lingering GCS sessions (they stay open for ~1 week otherwise).
|
||||
offset = 0
|
||||
try:
|
||||
while offset < total:
|
||||
chunk = gzip_bytes[offset : offset + _GCS_CHUNK_SIZE]
|
||||
chunk_end = offset + len(chunk) - 1
|
||||
is_final = (offset + len(chunk)) >= total
|
||||
content_range = (
|
||||
f"bytes {offset}-{chunk_end}/{total}"
|
||||
if is_final
|
||||
else f"bytes {offset}-{chunk_end}/*"
|
||||
)
|
||||
expected_statuses = {200, 201} if is_final else {308}
|
||||
|
||||
resp = await self._http.client.request(
|
||||
method="PUT",
|
||||
url=session_uri,
|
||||
headers={
|
||||
"Content-Type": "application/gzip",
|
||||
"Content-Range": content_range,
|
||||
},
|
||||
content=chunk,
|
||||
timeout=120.0,
|
||||
)
|
||||
if resp.status_code not in expected_statuses:
|
||||
raise RuntimeError(
|
||||
f"Mavvrik FOCUS destination: GCS chunk upload failed "
|
||||
f"(chunk offset={offset}, expected={expected_statuses}, "
|
||||
f"got={resp.status_code}): {resp.text[:400]}"
|
||||
)
|
||||
offset += len(chunk)
|
||||
verbose_logger.debug(
|
||||
"Mavvrik FOCUS destination: uploaded chunk offset=%d/%d",
|
||||
offset,
|
||||
total,
|
||||
)
|
||||
except Exception:
|
||||
# Cancel the open GCS session so it doesn't linger for up to 1 week.
|
||||
try:
|
||||
await self._http.client.request(
|
||||
method="DELETE", url=session_uri, timeout=10.0
|
||||
)
|
||||
verbose_logger.debug(
|
||||
"Mavvrik FOCUS destination: cancelled GCS session after error"
|
||||
)
|
||||
except Exception:
|
||||
pass
|
||||
raise
|
||||
|
||||
async def get_metrics_marker(self) -> Optional[int]:
|
||||
"""Register with Mavvrik and return the current metricsMarker.
|
||||
|
||||
The metricsMarker is a Unix timestamp (seconds) representing the last
|
||||
date Mavvrik has successfully ingested. Called on every scheduled run
|
||||
so the logger can detect and catch up any dates missed due to previous
|
||||
export failures.
|
||||
|
||||
Always calls the Mavvrik register API — unlike deliver() which skips
|
||||
registration once _registered is True, catch-up requires a fresh
|
||||
marker value on every run.
|
||||
"""
|
||||
resp = await self._http.client.request(
|
||||
method="POST",
|
||||
url=self._agent_url,
|
||||
headers=self._auth_headers,
|
||||
json={"name": self.connection_id},
|
||||
timeout=30.0,
|
||||
)
|
||||
if resp.status_code == 410:
|
||||
self._registered = False
|
||||
raise RuntimeError(
|
||||
"Mavvrik FOCUS destination: connector is disconnected (410). "
|
||||
"Re-enable the connection in the Mavvrik dashboard."
|
||||
)
|
||||
if resp.status_code >= 400:
|
||||
raise RuntimeError(
|
||||
f"Mavvrik FOCUS destination: register failed "
|
||||
f"({resp.status_code}): {resp.text[:200]}"
|
||||
)
|
||||
self._registered = True
|
||||
metrics_marker = resp.json().get("metricsMarker", 0)
|
||||
verbose_logger.debug(
|
||||
"Mavvrik FOCUS destination: got metricsMarker=%s", metrics_marker
|
||||
)
|
||||
return metrics_marker
|
||||
|
||||
async def deliver(
|
||||
self,
|
||||
*,
|
||||
content: bytes,
|
||||
time_window: FocusTimeWindow,
|
||||
filename: str,
|
||||
) -> None:
|
||||
"""Upload FOCUS CSV to Mavvrik via GCS signed URL.
|
||||
|
||||
Uses the start date of the time window as the object date key.
|
||||
"""
|
||||
if not content:
|
||||
verbose_logger.debug(
|
||||
"Mavvrik FOCUS destination: empty content, skipping upload"
|
||||
)
|
||||
return
|
||||
|
||||
date_str = time_window.start_time.strftime("%Y-%m-%d")
|
||||
|
||||
verbose_logger.debug(
|
||||
"Mavvrik FOCUS destination: uploading %d bytes for date=%s (%s)",
|
||||
len(content),
|
||||
date_str,
|
||||
filename,
|
||||
)
|
||||
|
||||
await self._ensure_registered()
|
||||
signed_url = await self._get_signed_url(date_str)
|
||||
await self._upload_to_gcs(signed_url, content)
|
||||
|
||||
verbose_logger.debug(
|
||||
"Mavvrik FOCUS destination: upload complete for date=%s", date_str
|
||||
)
|
||||
0
litellm/integrations/mavvrik_focus/__init__.py
Normal file
0
litellm/integrations/mavvrik_focus/__init__.py
Normal file
272
litellm/integrations/mavvrik_focus/mavvrik_focus_logger.py
Normal file
272
litellm/integrations/mavvrik_focus/mavvrik_focus_logger.py
Normal file
|
|
@ -0,0 +1,272 @@
|
|||
"""MavvrikFocusLogger — FOCUS-based Mavvrik export logger.
|
||||
|
||||
Usage in config.yaml:
|
||||
litellm_settings:
|
||||
callbacks: ["mavvrik"]
|
||||
|
||||
Required env vars:
|
||||
MAVVRIK_API_KEY
|
||||
MAVVRIK_API_ENDPOINT
|
||||
MAVVRIK_CONNECTION_ID
|
||||
|
||||
Optional env vars:
|
||||
MAVVRIK_FOCUS_MAX_ROWS — row cap per export window (default: 500000)
|
||||
|
||||
Only daily frequency is supported. The Mavvrik ingestion protocol stores one
|
||||
file per calendar date (metrics/YYYY-MM-DD). Hourly or interval exports would
|
||||
overwrite each other within the same day, producing incomplete data.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from typing import TYPE_CHECKING, Any, List, Optional
|
||||
|
||||
import litellm
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.constants import MAVVRIK_FOCUS_EXPORT_JOB_NAME
|
||||
from litellm.integrations.focus.destinations.base import FocusTimeWindow
|
||||
from litellm.integrations.focus.focus_logger import FocusLogger
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from apscheduler.schedulers.asyncio import AsyncIOScheduler
|
||||
else:
|
||||
AsyncIOScheduler = Any
|
||||
|
||||
|
||||
def _parse_metrics_marker(
|
||||
marker: Optional[object],
|
||||
) -> Optional[datetime]:
|
||||
"""Parse metricsMarker from Mavvrik register response into a UTC datetime.
|
||||
|
||||
Handles both formats Mavvrik may return:
|
||||
- Unix timestamp (int/float): e.g. 1749340800
|
||||
- ISO date string: e.g. "2026-06-09" or "2026-06-09T00:00:00Z"
|
||||
|
||||
Returns None for falsy values (0, None, empty string) which indicate
|
||||
no data has been ingested yet.
|
||||
"""
|
||||
if not marker:
|
||||
return None
|
||||
try:
|
||||
if isinstance(marker, (int, float)):
|
||||
return datetime.fromtimestamp(float(marker), tz=timezone.utc).replace(
|
||||
hour=0, minute=0, second=0, microsecond=0
|
||||
)
|
||||
if isinstance(marker, str):
|
||||
marker = marker.strip()
|
||||
if not marker:
|
||||
return None
|
||||
# Try ISO date first (YYYY-MM-DD), then full ISO datetime
|
||||
for fmt in ("%Y-%m-%d", "%Y-%m-%dT%H:%M:%SZ", "%Y-%m-%dT%H:%M:%S"):
|
||||
try:
|
||||
return datetime.strptime(marker, fmt).replace(tzinfo=timezone.utc)
|
||||
except ValueError:
|
||||
continue
|
||||
except Exception:
|
||||
pass
|
||||
verbose_proxy_logger.warning(
|
||||
"Mavvrik FOCUS: could not parse metricsMarker %r — skipping catch-up", marker
|
||||
)
|
||||
return None
|
||||
|
||||
|
||||
class MavvrikFocusLogger(FocusLogger):
|
||||
"""FOCUS-based export logger that routes to the Mavvrik destination."""
|
||||
|
||||
def __init__(self, **kwargs: Any) -> None:
|
||||
frequency = os.getenv("MAVVRIK_FOCUS_FREQUENCY", "daily").lower()
|
||||
if frequency != "daily":
|
||||
raise ValueError(
|
||||
f"MAVVRIK_FOCUS_FREQUENCY='{frequency}' is not supported. "
|
||||
"Only 'daily' is allowed — the Mavvrik ingestion protocol stores one "
|
||||
"file per calendar date (metrics/YYYY-MM-DD). Hourly or interval "
|
||||
"exports would overwrite each other within the same day."
|
||||
)
|
||||
super().__init__(
|
||||
provider="mavvrik",
|
||||
export_format="csv",
|
||||
frequency="daily",
|
||||
prefix="mavvrik_focus_exports",
|
||||
destination_config={
|
||||
"api_key": os.getenv("MAVVRIK_API_KEY"),
|
||||
"api_endpoint": os.getenv("MAVVRIK_API_ENDPOINT"),
|
||||
"connection_id": os.getenv("MAVVRIK_CONNECTION_ID"),
|
||||
},
|
||||
**kwargs,
|
||||
)
|
||||
raw = os.getenv("MAVVRIK_FOCUS_MAX_ROWS")
|
||||
self._max_rows: Optional[int] = int(raw) if raw else 500_000
|
||||
|
||||
async def _export_window(
|
||||
self,
|
||||
*,
|
||||
window: FocusTimeWindow,
|
||||
limit: Optional[int],
|
||||
) -> None:
|
||||
"""Export with Mavvrik row cap applied when no explicit limit is passed."""
|
||||
effective_limit = limit if limit is not None else self._max_rows
|
||||
engine = self._ensure_engine()
|
||||
data = await engine._database.get_usage_data(
|
||||
limit=effective_limit,
|
||||
start_time_utc=window.start_time,
|
||||
end_time_utc=window.end_time,
|
||||
)
|
||||
if effective_limit is not None and len(data) >= effective_limit:
|
||||
verbose_proxy_logger.warning(
|
||||
"Mavvrik FOCUS export: row cap reached (%d rows). "
|
||||
"Some data for window %s→%s may be excluded. "
|
||||
"Increase MAVVRIK_FOCUS_MAX_ROWS to export all rows.",
|
||||
effective_limit,
|
||||
window.start_time.date(),
|
||||
window.end_time.date(),
|
||||
)
|
||||
if data.is_empty():
|
||||
verbose_proxy_logger.debug(
|
||||
"Mavvrik FOCUS export: no usage data for window %s", window
|
||||
)
|
||||
return
|
||||
normalized = engine._transformer.transform(data)
|
||||
if normalized.is_empty():
|
||||
return
|
||||
payload = engine._serializer.serialize(normalized)
|
||||
if not payload:
|
||||
return
|
||||
await engine._destination.deliver(
|
||||
content=payload,
|
||||
time_window=window,
|
||||
filename=engine._build_filename(window),
|
||||
)
|
||||
|
||||
# Maximum number of days to catch up in a single run. Prevents runaway
|
||||
# loops if the connector was disabled for a long time, and avoids querying
|
||||
# data that has likely been cleaned up from LiteLLM_DailyUserSpend.
|
||||
_MAX_CATCHUP_DAYS = 7
|
||||
|
||||
async def _run_scheduled_export(self) -> None:
|
||||
"""Export today's window, catching up any dates Mavvrik has not yet received.
|
||||
|
||||
On each run:
|
||||
1. Register with Mavvrik → get metricsMarker (last successfully ingested date)
|
||||
2. If metricsMarker is behind yesterday, catch up missed dates (capped at
|
||||
_MAX_CATCHUP_DAYS to avoid runaway loops on long outages)
|
||||
3. Export yesterday (today's daily window)
|
||||
|
||||
This ensures a failed export on day N is automatically retried on day N+1
|
||||
without any manual intervention.
|
||||
"""
|
||||
engine = self._ensure_engine()
|
||||
from litellm.integrations.focus.destinations.mavvrik_destination import ( # noqa: PLC0415
|
||||
FocusMavvrikDestination,
|
||||
)
|
||||
|
||||
destination = engine._destination
|
||||
if not isinstance(destination, FocusMavvrikDestination):
|
||||
await super()._run_scheduled_export()
|
||||
return
|
||||
|
||||
# Register and get the last date Mavvrik has processed.
|
||||
# metricsMarker may be a Unix timestamp (int/float) or an ISO date string.
|
||||
marker = await destination.get_metrics_marker()
|
||||
|
||||
now = datetime.now(timezone.utc)
|
||||
yesterday = now.replace(hour=0, minute=0, second=0, microsecond=0) - timedelta(
|
||||
days=1
|
||||
)
|
||||
|
||||
last_ingested = _parse_metrics_marker(marker)
|
||||
|
||||
# Catch up missed dates, capped at _MAX_CATCHUP_DAYS
|
||||
if last_ingested and last_ingested < yesterday:
|
||||
# Never go further back than _MAX_CATCHUP_DAYS from yesterday
|
||||
earliest_catchup = yesterday - timedelta(days=self._MAX_CATCHUP_DAYS - 1)
|
||||
catch_up_date = max(last_ingested + timedelta(days=1), earliest_catchup)
|
||||
|
||||
if last_ingested + timedelta(days=1) < earliest_catchup:
|
||||
verbose_proxy_logger.warning(
|
||||
"Mavvrik FOCUS export: metricsMarker is more than %d days behind "
|
||||
"(%s). Catching up from %s only; earlier data will not be re-exported.",
|
||||
self._MAX_CATCHUP_DAYS,
|
||||
last_ingested.date(),
|
||||
catch_up_date.date(),
|
||||
)
|
||||
|
||||
while catch_up_date < yesterday:
|
||||
verbose_proxy_logger.info(
|
||||
"Mavvrik FOCUS export: catching up missed date %s",
|
||||
catch_up_date.date(),
|
||||
)
|
||||
window = FocusTimeWindow(
|
||||
start_time=catch_up_date,
|
||||
end_time=catch_up_date + timedelta(days=1),
|
||||
frequency="daily",
|
||||
)
|
||||
await self._export_window(window=window, limit=None)
|
||||
catch_up_date += timedelta(days=1)
|
||||
|
||||
# Export yesterday's window (the normal daily run)
|
||||
window = FocusTimeWindow(
|
||||
start_time=yesterday,
|
||||
end_time=yesterday + timedelta(days=1),
|
||||
frequency="daily",
|
||||
)
|
||||
await self._export_window(window=window, limit=None)
|
||||
|
||||
async def initialize_mavvrik_focus_export_job(self) -> None:
|
||||
"""Scheduler entry point — uses Mavvrik-specific pod-lock key."""
|
||||
from litellm.proxy.proxy_server import proxy_logging_obj # noqa: PLC0415
|
||||
|
||||
pod_lock_manager = None
|
||||
if proxy_logging_obj is not None:
|
||||
writer = getattr(proxy_logging_obj, "db_spend_update_writer", None)
|
||||
if writer is not None:
|
||||
pod_lock_manager = getattr(writer, "pod_lock_manager", None)
|
||||
|
||||
if pod_lock_manager and pod_lock_manager.redis_cache:
|
||||
acquired = await pod_lock_manager.acquire_lock(
|
||||
cronjob_id=MAVVRIK_FOCUS_EXPORT_JOB_NAME
|
||||
)
|
||||
if not acquired:
|
||||
verbose_proxy_logger.debug(
|
||||
"Mavvrik FOCUS export: unable to acquire pod lock"
|
||||
)
|
||||
return
|
||||
try:
|
||||
await self._run_scheduled_export()
|
||||
finally:
|
||||
await pod_lock_manager.release_lock(
|
||||
cronjob_id=MAVVRIK_FOCUS_EXPORT_JOB_NAME
|
||||
)
|
||||
else:
|
||||
await self._run_scheduled_export()
|
||||
|
||||
@staticmethod
|
||||
async def init_mavvrik_focus_background_job(
|
||||
scheduler: AsyncIOScheduler,
|
||||
) -> None:
|
||||
"""Register the Mavvrik FOCUS export job on the provided scheduler."""
|
||||
loggers: List[MavvrikFocusLogger] = [
|
||||
cb
|
||||
for cb in litellm.logging_callback_manager.get_custom_loggers_for_type(
|
||||
callback_type=MavvrikFocusLogger
|
||||
)
|
||||
if type(cb) is MavvrikFocusLogger
|
||||
]
|
||||
if not loggers:
|
||||
verbose_proxy_logger.debug(
|
||||
"No MavvrikFocusLogger registered; skipping scheduler"
|
||||
)
|
||||
return
|
||||
|
||||
logger = loggers[0]
|
||||
trigger_kwargs = logger._build_scheduler_trigger()
|
||||
scheduler.add_job( # type: ignore[attr-defined]
|
||||
logger.initialize_mavvrik_focus_export_job,
|
||||
id=MAVVRIK_FOCUS_EXPORT_JOB_NAME,
|
||||
replace_existing=True,
|
||||
**trigger_kwargs,
|
||||
)
|
||||
verbose_proxy_logger.info(
|
||||
"mavvrik_focus: background export job scheduled (%s)", trigger_kwargs
|
||||
)
|
||||
|
|
@ -25,6 +25,7 @@ from litellm.integrations.datadog.datadog_metrics import DatadogMetricsLogger
|
|||
from litellm.integrations.deepeval import DeepEvalLogger
|
||||
from litellm.integrations.dotprompt import DotpromptManager
|
||||
from litellm.integrations.focus.focus_logger import FocusLogger
|
||||
from litellm.integrations.mavvrik_focus.mavvrik_focus_logger import MavvrikFocusLogger
|
||||
from litellm.integrations.vantage.vantage_logger import VantageLogger
|
||||
from litellm.integrations.galileo import GalileoObserve
|
||||
from litellm.integrations.gcs_bucket.gcs_bucket import GCSBucketLogger
|
||||
|
|
@ -102,6 +103,7 @@ class CustomLoggerRegistry:
|
|||
"gitlab": GitLabPromptManager,
|
||||
"cloudzero": CloudZeroLogger,
|
||||
"focus": FocusLogger,
|
||||
"mavvrik": MavvrikFocusLogger,
|
||||
"vantage": VantageLogger,
|
||||
"posthog": PostHogLogger,
|
||||
}
|
||||
|
|
|
|||
|
|
@ -4124,6 +4124,17 @@ def _init_custom_logger_compatible_class( # noqa: PLR0915
|
|||
focus_logger = FocusLogger()
|
||||
_in_memory_loggers.append(focus_logger)
|
||||
return focus_logger # type: ignore
|
||||
elif logging_integration == "mavvrik":
|
||||
from litellm.integrations.mavvrik_focus.mavvrik_focus_logger import (
|
||||
MavvrikFocusLogger,
|
||||
)
|
||||
|
||||
for callback in _in_memory_loggers:
|
||||
if type(callback) is MavvrikFocusLogger:
|
||||
return callback # type: ignore
|
||||
mavvrik_focus_logger = MavvrikFocusLogger()
|
||||
_in_memory_loggers.append(mavvrik_focus_logger)
|
||||
return mavvrik_focus_logger # type: ignore
|
||||
elif logging_integration == "vantage":
|
||||
from litellm.integrations.vantage.vantage_logger import VantageLogger
|
||||
|
||||
|
|
|
|||
|
|
@ -7840,6 +7840,15 @@ class ProxyStartupEvent:
|
|||
)
|
||||
await VantageLogger.init_vantage_background_job(scheduler=scheduler)
|
||||
|
||||
########################################################
|
||||
# Mavvrik FOCUS Background Job
|
||||
########################################################
|
||||
from litellm.integrations.mavvrik_focus.mavvrik_focus_logger import ( # noqa: PLC0415
|
||||
MavvrikFocusLogger,
|
||||
)
|
||||
|
||||
await MavvrikFocusLogger.init_mavvrik_focus_background_job(scheduler=scheduler)
|
||||
|
||||
########################################################
|
||||
# Prometheus Background Job
|
||||
########################################################
|
||||
|
|
@ -9504,9 +9513,7 @@ async def realtime_websocket_endpoint(
|
|||
if intent == "transcription":
|
||||
route_model = "gpt-realtime-whisper"
|
||||
else:
|
||||
await websocket.close(
|
||||
code=1008, reason="model query parameter is required"
|
||||
)
|
||||
await websocket.close(code=1008, reason="model query parameter is required")
|
||||
return
|
||||
assert route_model is not None
|
||||
try:
|
||||
|
|
|
|||
|
|
@ -0,0 +1,751 @@
|
|||
"""Tests for FocusMavvrikDestination."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from datetime import datetime, timezone
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
from litellm.integrations.focus.destinations.base import FocusTimeWindow
|
||||
from litellm.integrations.focus.destinations.mavvrik_destination import (
|
||||
FocusMavvrikDestination,
|
||||
_validate_api_endpoint,
|
||||
)
|
||||
|
||||
VALID_ENDPOINT = "https://api.mavvrik.ai/tenant123"
|
||||
|
||||
|
||||
def _make_window() -> FocusTimeWindow:
|
||||
return FocusTimeWindow(
|
||||
start_time=datetime(2026, 1, 1, 0, 0, 0, tzinfo=timezone.utc),
|
||||
end_time=datetime(2026, 1, 2, 0, 0, 0, tzinfo=timezone.utc),
|
||||
frequency="daily",
|
||||
)
|
||||
|
||||
|
||||
def _dest(**overrides) -> FocusMavvrikDestination:
|
||||
config = {
|
||||
"api_key": "test-key",
|
||||
"api_endpoint": VALID_ENDPOINT,
|
||||
"connection_id": "conn-123",
|
||||
}
|
||||
config.update(overrides)
|
||||
return FocusMavvrikDestination(prefix="mavvrik_focus_exports", config=config)
|
||||
|
||||
|
||||
def test_missing_api_key_raises():
|
||||
with pytest.raises(ValueError, match="MAVVRIK_API_KEY"):
|
||||
FocusMavvrikDestination(
|
||||
prefix="p",
|
||||
config={"api_endpoint": VALID_ENDPOINT, "connection_id": "c"},
|
||||
)
|
||||
|
||||
|
||||
def test_missing_api_endpoint_raises():
|
||||
with pytest.raises(ValueError, match="MAVVRIK_API_ENDPOINT"):
|
||||
FocusMavvrikDestination(
|
||||
prefix="p",
|
||||
config={"api_key": "k", "connection_id": "c"},
|
||||
)
|
||||
|
||||
|
||||
def test_missing_connection_id_raises():
|
||||
with pytest.raises(ValueError, match="MAVVRIK_CONNECTION_ID"):
|
||||
FocusMavvrikDestination(
|
||||
prefix="p",
|
||||
config={"api_key": "k", "api_endpoint": VALID_ENDPOINT},
|
||||
)
|
||||
|
||||
|
||||
def test_non_https_endpoint_raises():
|
||||
with pytest.raises(ValueError, match="HTTPS"):
|
||||
_validate_api_endpoint("http://api.mavvrik.ai/tenant")
|
||||
|
||||
|
||||
def test_non_mavvrik_domain_raises():
|
||||
with pytest.raises(ValueError, match="Mavvrik domain"):
|
||||
_validate_api_endpoint("https://evil.com/tenant")
|
||||
|
||||
|
||||
def test_valid_mavvrik_domains_accepted():
|
||||
for domain in (
|
||||
"https://api.mavvrik.ai/tenant",
|
||||
"https://api.mavvrik.dev/tenant",
|
||||
"https://api.mavvrik.app/tenant",
|
||||
):
|
||||
_validate_api_endpoint(domain) # must not raise
|
||||
|
||||
|
||||
def test_initializes_with_not_registered():
|
||||
dest = _dest()
|
||||
assert dest._registered is False
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_deliver_skips_empty_content():
|
||||
dest = _dest()
|
||||
await dest.deliver(content=b"", time_window=_make_window(), filename="usage.csv")
|
||||
# _registered still False — _ensure_registered was never called
|
||||
assert dest._registered is False
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_large_content_uploads_in_multiple_chunks():
|
||||
"""Content larger than _GCS_CHUNK_SIZE must be uploaded in multiple chunks.
|
||||
|
||||
GCS assembles intermediate chunks (308) + final chunk (200) into one object.
|
||||
The destination must send Content-Range headers for each chunk correctly.
|
||||
"""
|
||||
from litellm.integrations.focus.destinations.mavvrik_destination import (
|
||||
FocusMavvrikDestination,
|
||||
_GCS_CHUNK_SIZE,
|
||||
)
|
||||
|
||||
dest = FocusMavvrikDestination(
|
||||
prefix="p",
|
||||
config={"api_key": "k", "api_endpoint": VALID_ENDPOINT, "connection_id": "c"},
|
||||
)
|
||||
|
||||
register_resp = MagicMock()
|
||||
register_resp.status_code = 200
|
||||
|
||||
signed_url_resp = MagicMock()
|
||||
signed_url_resp.status_code = 200
|
||||
signed_url_resp.json.return_value = {
|
||||
"url": "https://storage.googleapis.com/upload?sig=x"
|
||||
}
|
||||
|
||||
init_resp = MagicMock()
|
||||
init_resp.status_code = 200
|
||||
init_resp.headers = {"Location": "https://storage.googleapis.com/session"}
|
||||
|
||||
# First chunk → 308, second (final) chunk → 200
|
||||
chunk1_resp = MagicMock()
|
||||
chunk1_resp.status_code = 308
|
||||
|
||||
chunk2_resp = MagicMock()
|
||||
chunk2_resp.status_code = 200
|
||||
|
||||
mock_http = MagicMock()
|
||||
mock_http.client = MagicMock()
|
||||
mock_http.client.request = AsyncMock(
|
||||
side_effect=[
|
||||
register_resp,
|
||||
signed_url_resp,
|
||||
init_resp,
|
||||
chunk1_resp,
|
||||
chunk2_resp,
|
||||
]
|
||||
)
|
||||
dest._http = mock_http
|
||||
|
||||
# Build content that when gzipped exceeds one chunk.
|
||||
# Use incompressible random-ish bytes to ensure gzip doesn't shrink it below the chunk size.
|
||||
import os as _os
|
||||
|
||||
raw = b"col1,col2\n" + _os.urandom(_GCS_CHUNK_SIZE + 1024)
|
||||
|
||||
await dest.deliver(
|
||||
content=raw,
|
||||
time_window=_make_window(),
|
||||
filename="usage.csv",
|
||||
)
|
||||
|
||||
# register + get_signed_url + init + 2 chunk PUTs = 5 calls
|
||||
assert mock_http.client.request.call_count == 5
|
||||
|
||||
# Check Content-Range headers
|
||||
put_calls = mock_http.client.request.call_args_list[3:]
|
||||
assert "bytes" in put_calls[0].kwargs["headers"]["Content-Range"]
|
||||
assert "/*" in put_calls[0].kwargs["headers"]["Content-Range"] # intermediate
|
||||
assert "/*" not in put_calls[1].kwargs["headers"]["Content-Range"] # final
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_deliver_calls_register_get_url_and_upload():
|
||||
dest = _dest()
|
||||
|
||||
register_resp = MagicMock()
|
||||
register_resp.status_code = 200
|
||||
|
||||
signed_url_resp = MagicMock()
|
||||
signed_url_resp.status_code = 200
|
||||
signed_url_resp.json.return_value = {"url": "https://storage.googleapis.com/signed"}
|
||||
|
||||
init_resp = MagicMock()
|
||||
init_resp.status_code = 200
|
||||
init_resp.headers = {"Location": "https://storage.googleapis.com/session-uri"}
|
||||
|
||||
upload_resp = MagicMock()
|
||||
upload_resp.status_code = 200
|
||||
|
||||
mock_http = MagicMock()
|
||||
mock_http.client = MagicMock()
|
||||
# All 4 calls go through self._http.client.request:
|
||||
# 1. register, 2. get_signed_url, 3. GCS session init POST, 4. GCS PUT
|
||||
mock_http.client.request = AsyncMock(
|
||||
side_effect=[register_resp, signed_url_resp, init_resp, upload_resp]
|
||||
)
|
||||
dest._http = mock_http
|
||||
|
||||
await dest.deliver(
|
||||
content=b"header\nrow1\n",
|
||||
time_window=_make_window(),
|
||||
filename="usage.csv",
|
||||
)
|
||||
|
||||
assert dest._registered is True
|
||||
assert mock_http.client.request.call_count == 4
|
||||
# Verify Content-Range header was set on the PUT
|
||||
put_call = mock_http.client.request.call_args_list[3]
|
||||
assert "Content-Range" in put_call.kwargs["headers"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_register_called_only_once_across_multiple_deliveries():
|
||||
dest = _dest()
|
||||
|
||||
register_resp = MagicMock()
|
||||
register_resp.status_code = 200
|
||||
|
||||
def _signed_url_resp():
|
||||
r = MagicMock()
|
||||
r.status_code = 200
|
||||
r.json.return_value = {"url": "https://storage.googleapis.com/signed"}
|
||||
return r
|
||||
|
||||
init_resp = MagicMock()
|
||||
init_resp.status_code = 200
|
||||
init_resp.headers = {"Location": "https://storage.googleapis.com/session-uri"}
|
||||
|
||||
upload_resp = MagicMock()
|
||||
upload_resp.status_code = 200
|
||||
|
||||
mock_http = MagicMock()
|
||||
mock_http.client = MagicMock()
|
||||
# First delivery: register, get_signed_url, GCS init, GCS PUT
|
||||
# Second delivery: get_signed_url, GCS init, GCS PUT (register skipped)
|
||||
mock_http.client.request = AsyncMock(
|
||||
side_effect=[
|
||||
register_resp,
|
||||
_signed_url_resp(),
|
||||
init_resp,
|
||||
upload_resp,
|
||||
_signed_url_resp(),
|
||||
init_resp,
|
||||
upload_resp,
|
||||
]
|
||||
)
|
||||
dest._http = mock_http
|
||||
|
||||
window = _make_window()
|
||||
await dest.deliver(content=b"header\nrow1\n", time_window=window, filename="1.csv")
|
||||
await dest.deliver(content=b"header\nrow2\n", time_window=window, filename="2.csv")
|
||||
|
||||
# 7 total: register(1) + [get_url+init+put](2) × 2 deliveries
|
||||
assert mock_http.client.request.call_count == 7
|
||||
# First call was register
|
||||
first_call = mock_http.client.request.call_args_list[0]
|
||||
assert first_call.kwargs["method"] == "POST"
|
||||
assert "/upload-url" not in first_call.kwargs["url"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_deliver_raises_on_register_failure():
|
||||
dest = _dest()
|
||||
|
||||
fail_resp = MagicMock()
|
||||
fail_resp.status_code = 403
|
||||
fail_resp.text = "Forbidden"
|
||||
|
||||
mock_http = MagicMock()
|
||||
mock_http.client = MagicMock()
|
||||
mock_http.client.request = AsyncMock(return_value=fail_resp)
|
||||
dest._http = mock_http
|
||||
|
||||
with pytest.raises(RuntimeError, match="register failed"):
|
||||
await dest.deliver(
|
||||
content=b"data",
|
||||
time_window=_make_window(),
|
||||
filename="usage.csv",
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_deliver_raises_on_signed_url_api_error():
|
||||
"""_get_signed_url must raise RuntimeError when the API returns a 4xx."""
|
||||
dest = _dest()
|
||||
|
||||
register_resp = MagicMock()
|
||||
register_resp.status_code = 200
|
||||
|
||||
fail_resp = MagicMock()
|
||||
fail_resp.status_code = 500
|
||||
fail_resp.text = "Internal Server Error"
|
||||
|
||||
mock_http = MagicMock()
|
||||
mock_http.client = MagicMock()
|
||||
mock_http.client.request = AsyncMock(side_effect=[register_resp, fail_resp])
|
||||
dest._http = mock_http
|
||||
|
||||
with pytest.raises(RuntimeError, match="failed to get signed URL"):
|
||||
await dest.deliver(
|
||||
content=b"data",
|
||||
time_window=_make_window(),
|
||||
filename="usage.csv",
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_deliver_raises_on_missing_signed_url():
|
||||
dest = _dest()
|
||||
|
||||
register_resp = MagicMock()
|
||||
register_resp.status_code = 200
|
||||
|
||||
bad_url_resp = MagicMock()
|
||||
bad_url_resp.status_code = 200
|
||||
bad_url_resp.json.return_value = {} # no 'url' field
|
||||
|
||||
mock_http = MagicMock()
|
||||
mock_http.client = MagicMock()
|
||||
mock_http.client.request = AsyncMock(side_effect=[register_resp, bad_url_resp])
|
||||
dest._http = mock_http
|
||||
|
||||
with pytest.raises(RuntimeError, match="missing 'url' field"):
|
||||
await dest.deliver(
|
||||
content=b"data",
|
||||
time_window=_make_window(),
|
||||
filename="usage.csv",
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_deliver_raises_on_non_gcs_signed_url():
|
||||
"""Signed URL pointing to a non-GCS host must be rejected before any upload."""
|
||||
from litellm.integrations.focus.destinations.mavvrik_destination import (
|
||||
_validate_gcs_url,
|
||||
)
|
||||
|
||||
with pytest.raises(ValueError, match="GCS endpoint"):
|
||||
_validate_gcs_url("https://evil.com/upload?token=abc", "signed URL")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_deliver_raises_on_non_gcs_session_uri():
|
||||
"""Session URI from Location header pointing to a non-GCS host must be rejected."""
|
||||
dest = _dest()
|
||||
|
||||
register_resp = MagicMock()
|
||||
register_resp.status_code = 200
|
||||
|
||||
signed_url_resp = MagicMock()
|
||||
signed_url_resp.status_code = 200
|
||||
# signed URL is valid GCS
|
||||
signed_url_resp.json.return_value = {
|
||||
"url": "https://storage.googleapis.com/upload?sig=abc"
|
||||
}
|
||||
|
||||
# Location header points to a non-GCS host
|
||||
init_resp = MagicMock()
|
||||
init_resp.status_code = 200
|
||||
init_resp.headers = {"Location": "https://evil.com/session-uri"}
|
||||
|
||||
mock_http = MagicMock()
|
||||
mock_http.client = MagicMock()
|
||||
# register, get_signed_url, GCS session init (returns bad Location)
|
||||
mock_http.client.request = AsyncMock(
|
||||
side_effect=[register_resp, signed_url_resp, init_resp]
|
||||
)
|
||||
dest._http = mock_http
|
||||
|
||||
with pytest.raises(ValueError, match="GCS endpoint"):
|
||||
await dest.deliver(
|
||||
content=b"data",
|
||||
time_window=_make_window(),
|
||||
filename="usage.csv",
|
||||
)
|
||||
|
||||
|
||||
def test_factory_creates_mavvrik_destination(monkeypatch):
|
||||
monkeypatch.setenv("MAVVRIK_API_KEY", "k")
|
||||
monkeypatch.setenv("MAVVRIK_API_ENDPOINT", VALID_ENDPOINT)
|
||||
monkeypatch.setenv("MAVVRIK_CONNECTION_ID", "c")
|
||||
|
||||
from litellm.integrations.focus.destinations.factory import FocusDestinationFactory
|
||||
|
||||
dest = FocusDestinationFactory.create(provider="mavvrik", prefix="p")
|
||||
|
||||
assert isinstance(dest, FocusMavvrikDestination)
|
||||
assert dest.api_key == "k"
|
||||
assert dest.connection_id == "c"
|
||||
|
||||
|
||||
def test_only_daily_frequency_is_supported():
|
||||
"""MavvrikFocusLogger must raise ValueError for non-daily frequencies."""
|
||||
import importlib
|
||||
|
||||
for freq in ("hourly", "interval"):
|
||||
|
||||
def _make(f=freq, monkeypatch=None):
|
||||
import os
|
||||
|
||||
old = os.environ.get("MAVVRIK_FOCUS_FREQUENCY")
|
||||
os.environ["MAVVRIK_FOCUS_FREQUENCY"] = f
|
||||
try:
|
||||
from litellm.integrations.mavvrik_focus import mavvrik_focus_logger
|
||||
|
||||
importlib.reload(mavvrik_focus_logger)
|
||||
with pytest.raises(ValueError, match="Only 'daily' is allowed"):
|
||||
mavvrik_focus_logger.MavvrikFocusLogger()
|
||||
finally:
|
||||
if old is None:
|
||||
os.environ.pop("MAVVRIK_FOCUS_FREQUENCY", None)
|
||||
else:
|
||||
os.environ["MAVVRIK_FOCUS_FREQUENCY"] = old
|
||||
|
||||
_make()
|
||||
|
||||
|
||||
def test_max_rows_defaults_to_500k():
|
||||
"""MAVVRIK_FOCUS_MAX_ROWS defaults to 500_000 when not set."""
|
||||
from litellm.integrations.mavvrik_focus.mavvrik_focus_logger import (
|
||||
MavvrikFocusLogger,
|
||||
)
|
||||
|
||||
logger = MavvrikFocusLogger()
|
||||
assert logger._max_rows == 500_000
|
||||
|
||||
|
||||
def test_max_rows_reads_from_env(monkeypatch):
|
||||
"""MAVVRIK_FOCUS_MAX_ROWS env var is respected."""
|
||||
monkeypatch.setenv("MAVVRIK_FOCUS_MAX_ROWS", "100000")
|
||||
|
||||
from litellm.integrations.mavvrik_focus.mavvrik_focus_logger import (
|
||||
MavvrikFocusLogger,
|
||||
)
|
||||
|
||||
logger = MavvrikFocusLogger()
|
||||
assert logger._max_rows == 100_000
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_export_window_passes_max_rows_as_limit(monkeypatch):
|
||||
"""_export_window must pass _max_rows as limit to get_usage_data."""
|
||||
monkeypatch.setenv("MAVVRIK_FOCUS_MAX_ROWS", "1000")
|
||||
|
||||
import polars as pl
|
||||
from litellm.integrations.mavvrik_focus.mavvrik_focus_logger import (
|
||||
MavvrikFocusLogger,
|
||||
)
|
||||
from litellm.integrations.focus.destinations.base import FocusTimeWindow
|
||||
from datetime import datetime, timezone
|
||||
|
||||
logger = MavvrikFocusLogger()
|
||||
assert logger._max_rows == 1000
|
||||
|
||||
# Mock the engine internals so _export_window runs through our new code path
|
||||
db_mock = MagicMock()
|
||||
db_mock.get_usage_data = AsyncMock(return_value=pl.DataFrame()) # empty → no upload
|
||||
|
||||
engine_mock = MagicMock()
|
||||
engine_mock._database = db_mock
|
||||
logger._engine = engine_mock
|
||||
|
||||
window = FocusTimeWindow(
|
||||
start_time=datetime(2026, 1, 1, tzinfo=timezone.utc),
|
||||
end_time=datetime(2026, 1, 2, tzinfo=timezone.utc),
|
||||
frequency="daily",
|
||||
)
|
||||
await logger._export_window(window=window, limit=None)
|
||||
|
||||
db_mock.get_usage_data.assert_called_once_with(
|
||||
limit=1000,
|
||||
start_time_utc=window.start_time,
|
||||
end_time_utc=window.end_time,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_run_scheduled_export_catches_up_missed_dates():
|
||||
"""If metricsMarker is 2 days behind, _run_scheduled_export exports missed dates first."""
|
||||
import polars as pl
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from litellm.integrations.mavvrik_focus.mavvrik_focus_logger import (
|
||||
MavvrikFocusLogger,
|
||||
)
|
||||
from litellm.integrations.focus.destinations.mavvrik_destination import (
|
||||
FocusMavvrikDestination,
|
||||
)
|
||||
|
||||
logger = MavvrikFocusLogger()
|
||||
|
||||
# metricsMarker = 3 days ago → 2 missed dates (day-2 and day-1) + today's run
|
||||
now = datetime.now(timezone.utc).replace(hour=0, minute=0, second=0, microsecond=0)
|
||||
yesterday = now - timedelta(days=1)
|
||||
two_days_ago = now - timedelta(days=2)
|
||||
three_days_ago = now - timedelta(days=3)
|
||||
|
||||
marker_ts = int(three_days_ago.timestamp())
|
||||
|
||||
# Mock destination
|
||||
dest_mock = MagicMock(spec=FocusMavvrikDestination)
|
||||
dest_mock.get_metrics_marker = AsyncMock(return_value=marker_ts)
|
||||
|
||||
# Mock engine
|
||||
db_mock = MagicMock()
|
||||
db_mock.get_usage_data = AsyncMock(return_value=pl.DataFrame())
|
||||
engine_mock = MagicMock()
|
||||
engine_mock._database = db_mock
|
||||
engine_mock._destination = dest_mock
|
||||
logger._engine = engine_mock
|
||||
|
||||
await logger._run_scheduled_export()
|
||||
|
||||
# Should have queried DB 3 times: day-2, day-1 (yesterday), and the normal yesterday window
|
||||
# Actually: catch-up covers [three_days_ago+1 .. yesterday) = [two_days_ago, yesterday)
|
||||
# = two_days_ago only (1 missed), then normal yesterday = 2 total calls
|
||||
calls = db_mock.get_usage_data.call_args_list
|
||||
assert len(calls) == 2
|
||||
# First call is the catch-up (two_days_ago)
|
||||
assert calls[0].kwargs["start_time_utc"].date() == two_days_ago.date()
|
||||
# Second call is yesterday's normal daily run
|
||||
assert calls[1].kwargs["start_time_utc"].date() == yesterday.date()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_run_scheduled_export_no_catchup_when_marker_is_current():
|
||||
"""If metricsMarker = yesterday, no catch-up needed — just export yesterday."""
|
||||
import polars as pl
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from litellm.integrations.mavvrik_focus.mavvrik_focus_logger import (
|
||||
MavvrikFocusLogger,
|
||||
)
|
||||
from litellm.integrations.focus.destinations.mavvrik_destination import (
|
||||
FocusMavvrikDestination,
|
||||
)
|
||||
|
||||
logger = MavvrikFocusLogger()
|
||||
|
||||
now = datetime.now(timezone.utc).replace(hour=0, minute=0, second=0, microsecond=0)
|
||||
yesterday = now - timedelta(days=1)
|
||||
marker_ts = int(yesterday.timestamp())
|
||||
|
||||
dest_mock = MagicMock(spec=FocusMavvrikDestination)
|
||||
dest_mock.get_metrics_marker = AsyncMock(return_value=marker_ts)
|
||||
|
||||
db_mock = MagicMock()
|
||||
db_mock.get_usage_data = AsyncMock(return_value=pl.DataFrame())
|
||||
engine_mock = MagicMock()
|
||||
engine_mock._database = db_mock
|
||||
engine_mock._destination = dest_mock
|
||||
logger._engine = engine_mock
|
||||
|
||||
await logger._run_scheduled_export()
|
||||
|
||||
# Only one call — yesterday's normal run, no catch-up
|
||||
assert db_mock.get_usage_data.call_count == 1
|
||||
assert (
|
||||
db_mock.get_usage_data.call_args.kwargs["start_time_utc"].date()
|
||||
== yesterday.date()
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_metrics_marker_always_calls_api():
|
||||
"""get_metrics_marker must call the register API every time to get a fresh marker.
|
||||
|
||||
This is the key difference from deliver() — catch-up requires the current
|
||||
metricsMarker on every scheduled run, not just the first one.
|
||||
"""
|
||||
dest = _dest()
|
||||
|
||||
register_resp = MagicMock()
|
||||
register_resp.status_code = 200
|
||||
register_resp.json.return_value = {
|
||||
"id": "litellm-conn-123",
|
||||
"metricsMarker": 1749340800,
|
||||
}
|
||||
|
||||
mock_http = MagicMock()
|
||||
mock_http.client = MagicMock()
|
||||
mock_http.client.request = AsyncMock(return_value=register_resp)
|
||||
dest._http = mock_http
|
||||
|
||||
# First call
|
||||
marker = await dest.get_metrics_marker()
|
||||
assert marker == 1749340800
|
||||
assert dest._registered is True
|
||||
|
||||
# Second call — must call API again to get fresh marker (not return None)
|
||||
marker2 = await dest.get_metrics_marker()
|
||||
assert marker2 == 1749340800
|
||||
assert mock_http.client.request.call_count == 2 # API called both times
|
||||
|
||||
|
||||
def test_parse_metrics_marker_handles_unix_timestamp():
|
||||
from litellm.integrations.mavvrik_focus.mavvrik_focus_logger import (
|
||||
_parse_metrics_marker,
|
||||
)
|
||||
from datetime import datetime, timezone
|
||||
|
||||
# Use a known date and compute its timestamp to avoid hardcoding
|
||||
known_date = datetime(2026, 6, 9, 0, 0, 0, tzinfo=timezone.utc)
|
||||
ts = int(known_date.timestamp())
|
||||
|
||||
result = _parse_metrics_marker(ts)
|
||||
assert result is not None
|
||||
assert result.date().isoformat() == "2026-06-09"
|
||||
assert result.tzinfo == timezone.utc
|
||||
|
||||
|
||||
def test_parse_metrics_marker_handles_iso_date_string():
|
||||
from litellm.integrations.mavvrik_focus.mavvrik_focus_logger import (
|
||||
_parse_metrics_marker,
|
||||
)
|
||||
|
||||
result = _parse_metrics_marker("2026-06-09")
|
||||
assert result is not None
|
||||
assert result.date().isoformat() == "2026-06-09"
|
||||
|
||||
|
||||
def test_parse_metrics_marker_handles_iso_datetime_string():
|
||||
from litellm.integrations.mavvrik_focus.mavvrik_focus_logger import (
|
||||
_parse_metrics_marker,
|
||||
)
|
||||
|
||||
result = _parse_metrics_marker("2026-06-09T00:00:00Z")
|
||||
assert result is not None
|
||||
assert result.date().isoformat() == "2026-06-09"
|
||||
|
||||
|
||||
def test_parse_metrics_marker_returns_none_for_zero():
|
||||
from litellm.integrations.mavvrik_focus.mavvrik_focus_logger import (
|
||||
_parse_metrics_marker,
|
||||
)
|
||||
|
||||
assert _parse_metrics_marker(0) is None
|
||||
assert _parse_metrics_marker(None) is None
|
||||
assert _parse_metrics_marker("") is None
|
||||
|
||||
|
||||
def test_parse_metrics_marker_returns_none_for_garbage():
|
||||
from litellm.integrations.mavvrik_focus.mavvrik_focus_logger import (
|
||||
_parse_metrics_marker,
|
||||
)
|
||||
|
||||
# Should not raise — logs warning and returns None
|
||||
assert _parse_metrics_marker("not-a-date") is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_catchup_capped_at_max_catchup_days():
|
||||
"""Catch-up must not go further back than _MAX_CATCHUP_DAYS."""
|
||||
import polars as pl
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from litellm.integrations.mavvrik_focus.mavvrik_focus_logger import (
|
||||
MavvrikFocusLogger,
|
||||
)
|
||||
from litellm.integrations.focus.destinations.mavvrik_destination import (
|
||||
FocusMavvrikDestination,
|
||||
)
|
||||
|
||||
logger = MavvrikFocusLogger()
|
||||
max_days = MavvrikFocusLogger._MAX_CATCHUP_DAYS
|
||||
|
||||
now = datetime.now(timezone.utc).replace(hour=0, minute=0, second=0, microsecond=0)
|
||||
yesterday = now - timedelta(days=1)
|
||||
# Marker is 30 days ago — well beyond the cap
|
||||
thirty_days_ago = now - timedelta(days=30)
|
||||
marker_ts = int(thirty_days_ago.timestamp())
|
||||
|
||||
dest_mock = MagicMock(spec=FocusMavvrikDestination)
|
||||
dest_mock.get_metrics_marker = AsyncMock(return_value=marker_ts)
|
||||
|
||||
db_mock = MagicMock()
|
||||
db_mock.get_usage_data = AsyncMock(return_value=pl.DataFrame())
|
||||
engine_mock = MagicMock()
|
||||
engine_mock._database = db_mock
|
||||
engine_mock._destination = dest_mock
|
||||
logger._engine = engine_mock
|
||||
|
||||
await logger._run_scheduled_export()
|
||||
|
||||
# Should have queried at most _MAX_CATCHUP_DAYS times
|
||||
# (max_days - 1 catch-up dates + 1 yesterday = max_days total)
|
||||
assert db_mock.get_usage_data.call_count <= max_days
|
||||
|
||||
# First catch-up date must not be earlier than (yesterday - max_days + 1)
|
||||
earliest_allowed = yesterday - timedelta(days=max_days - 1)
|
||||
first_call_start = db_mock.get_usage_data.call_args_list[0].kwargs["start_time_utc"]
|
||||
assert first_call_start.date() >= earliest_allowed.date()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_register_resets_on_410():
|
||||
"""_registered flag must be False after a 410 so next run re-registers."""
|
||||
dest = _dest()
|
||||
|
||||
resp_410 = MagicMock()
|
||||
resp_410.status_code = 410
|
||||
resp_410.text = "Gone"
|
||||
|
||||
mock_http = MagicMock()
|
||||
mock_http.client = MagicMock()
|
||||
mock_http.client.request = AsyncMock(return_value=resp_410)
|
||||
dest._http = mock_http
|
||||
dest._registered = False # not yet registered — trigger the call
|
||||
|
||||
with pytest.raises(RuntimeError, match="disconnected"):
|
||||
await dest._ensure_registered()
|
||||
|
||||
assert dest._registered is False
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_gcs_session_cancelled_on_chunk_failure():
|
||||
"""GCS session must be cancelled (DELETE) when a chunk PUT fails."""
|
||||
dest = _dest()
|
||||
|
||||
register_resp = MagicMock()
|
||||
register_resp.status_code = 200
|
||||
|
||||
signed_url_resp = MagicMock()
|
||||
signed_url_resp.status_code = 200
|
||||
signed_url_resp.json.return_value = {
|
||||
"url": "https://storage.googleapis.com/upload?sig=x"
|
||||
}
|
||||
|
||||
init_resp = MagicMock()
|
||||
init_resp.status_code = 200
|
||||
init_resp.headers = {"Location": "https://storage.googleapis.com/session"}
|
||||
|
||||
# Chunk PUT fails with 500
|
||||
fail_resp = MagicMock()
|
||||
fail_resp.status_code = 500
|
||||
fail_resp.text = "Internal Server Error"
|
||||
|
||||
# DELETE (session cancel)
|
||||
delete_resp = MagicMock()
|
||||
delete_resp.status_code = 200
|
||||
|
||||
mock_http = MagicMock()
|
||||
mock_http.client = MagicMock()
|
||||
mock_http.client.request = AsyncMock(
|
||||
side_effect=[register_resp, signed_url_resp, init_resp, fail_resp, delete_resp]
|
||||
)
|
||||
dest._http = mock_http
|
||||
|
||||
with pytest.raises(RuntimeError, match="GCS chunk upload failed"):
|
||||
await dest.deliver(
|
||||
content=b"header\nrow1\n",
|
||||
time_window=_make_window(),
|
||||
filename="usage.csv",
|
||||
)
|
||||
|
||||
# Verify DELETE was called to cancel the session
|
||||
calls = mock_http.client.request.call_args_list
|
||||
delete_call = calls[4]
|
||||
assert delete_call.kwargs["method"] == "DELETE"
|
||||
assert "storage.googleapis.com/session" in delete_call.kwargs["url"]
|
||||
Loading…
Add table
Reference in a new issue