diff --git a/litellm/integrations/mavvrik/__init__.py b/litellm/integrations/mavvrik/__init__.py index 97f626e22a8..4fdb858f70c 100644 --- a/litellm/integrations/mavvrik/__init__.py +++ b/litellm/integrations/mavvrik/__init__.py @@ -17,7 +17,11 @@ from datetime import timezone as _tz from typing import Optional from litellm._logging import verbose_proxy_logger -from litellm.constants import MAVVRIK_MAX_FETCHED_DATA_RECORDS +from litellm.constants import ( + MAVVRIK_EXPORT_INTERVAL_MINUTES, + MAVVRIK_EXPORT_USAGE_DATA_JOB_NAME, + 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 @@ -49,6 +53,28 @@ def _build_client(data: dict) -> Client: ) +def _schedule_job( + scheduler: object, api_key: str, api_endpoint: str, connection_id: str +) -> None: + """Register (or replace) the Mavvrik background export job on the given scheduler.""" + client = Client( + api_key=api_key, api_endpoint=api_endpoint, connection_id=connection_id + ) + uploader = Uploader(client=client) + orchestrator = Orchestrator(client=client, uploader=uploader) + scheduler.add_job( # type: ignore[union-attr] + 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, + ) + + class Service: """Public facade that mediates between the REST endpoints and the Mavvrik modules. @@ -105,11 +131,6 @@ class Service: ) # 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) @@ -122,25 +143,12 @@ class Service: "status": "success", } - client = Client( + _schedule_job( + _scheduler, 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", @@ -237,29 +245,16 @@ class Service: # Reschedule the background job with the new credentials so the # running Orchestrator uses the merged values immediately — # without this, the in-memory Client keeps old credentials until restart. - from litellm.constants import ( - MAVVRIK_EXPORT_INTERVAL_MINUTES, - MAVVRIK_EXPORT_USAGE_DATA_JOB_NAME, - ) - import litellm.proxy.proxy_server as _pserver _scheduler = getattr(_pserver, "scheduler", None) if _scheduler is not None: - client = Client( + _schedule_job( + _scheduler, api_key=merged["api_key"], api_endpoint=merged["api_endpoint"], connection_id=merged["connection_id"], ) - uploader = Uploader(client=client) - orchestrator = Orchestrator(client=client, uploader=uploader) - _scheduler.add_job( - orchestrator.run, - "interval", - minutes=MAVVRIK_EXPORT_INTERVAL_MINUTES, - id=MAVVRIK_EXPORT_USAGE_DATA_JOB_NAME, - replace_existing=True, - ) verbose_proxy_logger.info( "mavvrik: background job rescheduled with updated credentials" ) @@ -281,8 +276,6 @@ class Service: Raises: LookupError: only when DB is connected and no settings row exists. """ - from litellm.constants import MAVVRIK_EXPORT_USAGE_DATA_JOB_NAME - import litellm.proxy.proxy_server as _pserver # Deregister scheduler first — independent of whether creds are in DB or env vars. @@ -293,12 +286,18 @@ class Service: except Exception: pass # job may not exist if scheduler was restarted - # Always attempt DB deletion; silently ignore if no row exists or no DB connected - # (env-var-only deployments without a database have no row to remove). + # Attempt DB deletion; skip gracefully when no row exists or DB not connected + # (env-var-only deployments have no row to remove). try: await self._settings.delete() - except Exception: - pass + except LookupError: + verbose_proxy_logger.info( + "mavvrik: no DB settings row found — scheduler deregistered only" + ) + except Exception as exc: + verbose_proxy_logger.warning( + "mavvrik: could not delete DB settings row: %s", exc + ) verbose_proxy_logger.info("mavvrik settings deleted") return {"message": "Mavvrik settings deleted successfully", "status": "success"} diff --git a/litellm/integrations/mavvrik/exporter.py b/litellm/integrations/mavvrik/exporter.py index d5de29835d7..c089c1626ca 100644 --- a/litellm/integrations/mavvrik/exporter.py +++ b/litellm/integrations/mavvrik/exporter.py @@ -10,13 +10,13 @@ Public interface: Returns MIN(date) for first-run start date resolution. Internal methods: - _stream_pages(date_str, connection_id, page_size) → AsyncGenerator[str, None] + stream_pages(date_str, connection_id, page_size) → AsyncGenerator[str, None] _get_usage_data(date_str, limit) → DataFrame _to_csv(df, connection_id) → str DB not connected: _get_usage_data / _to_csv — log warning and return empty/None (scheduler path). - _stream_pages — raises RuntimeError (propagates to Orchestrator try/except). + stream_pages — raises RuntimeError (propagates to Orchestrator try/except). Service.export / dry_run — call Settings._ensure_prisma_client() before reaching here, so they raise before the exporter is called. @@ -104,7 +104,7 @@ class Exporter: csv = self._to_csv(df, connection_id=connection_id) return df, csv - async def _stream_pages( + async def stream_pages( self, date_str: str, connection_id: Optional[str] = None, diff --git a/litellm/integrations/mavvrik/logger.py b/litellm/integrations/mavvrik/logger.py index 97a0ec09c78..25efb21b79f 100644 --- a/litellm/integrations/mavvrik/logger.py +++ b/litellm/integrations/mavvrik/logger.py @@ -7,10 +7,21 @@ 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. """ +import os + +from litellm._logging import verbose_proxy_logger from litellm.integrations.custom_logger import CustomLogger class Logger(CustomLogger): """Mavvrik integration marker — registered via callbacks: ["mavvrik"].""" - pass + def __init__(self) -> None: + super().__init__() + _required = ("MAVVRIK_API_KEY", "MAVVRIK_API_ENDPOINT", "MAVVRIK_CONNECTION_ID") + if not all(os.getenv(v) for v in _required): + verbose_proxy_logger.warning( + "mavvrik: callbacks: ['mavvrik'] is set but credentials are not configured. " + "Call POST /mavvrik/init or set MAVVRIK_API_KEY, MAVVRIK_API_ENDPOINT, " + "and MAVVRIK_CONNECTION_ID to enable exports." + ) diff --git a/litellm/integrations/mavvrik/orchestrator.py b/litellm/integrations/mavvrik/orchestrator.py index a96297c829e..2d1ce415ae8 100644 --- a/litellm/integrations/mavvrik/orchestrator.py +++ b/litellm/integrations/mavvrik/orchestrator.py @@ -134,17 +134,17 @@ class Orchestrator: 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 + 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( + 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) + total_bytes = await self._uploader.stream_upload(pages, date_str=date_str) if total_bytes > 0: verbose_logger.info( "Orchestrator: %s → streamed %d bytes to GCS ✓", date_str, total_bytes diff --git a/litellm/integrations/mavvrik/uploader.py b/litellm/integrations/mavvrik/uploader.py index 27a042190ac..c2148393518 100644 --- a/litellm/integrations/mavvrik/uploader.py +++ b/litellm/integrations/mavvrik/uploader.py @@ -167,7 +167,7 @@ class Uploader: # Streaming upload # ------------------------------------------------------------------ - async def _stream_upload( + async def stream_upload( self, pages: AsyncIterator[str], date_str: str, diff --git a/tests/test_litellm/integrations/mavvrik/test_service.py b/tests/test_litellm/integrations/mavvrik/test_service.py index 590dc9851c0..148abe81080 100644 --- a/tests/test_litellm/integrations/mavvrik/test_service.py +++ b/tests/test_litellm/integrations/mavvrik/test_service.py @@ -658,14 +658,14 @@ class TestResolveFirstRunInvalidDate: 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().""" + """_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 + orc._exporter.stream_pages = fake_stream_pages stream_upload_called_with = [] @@ -676,7 +676,7 @@ class TestStreamingExport: pass return 1024 - orc._uploader._stream_upload = fake_stream_upload + orc._uploader.stream_upload = fake_stream_upload result = await orc._export(date(2026, 4, 9)) @@ -692,12 +692,12 @@ class TestStreamingExport: return yield - orc._exporter._stream_pages = empty_pages + orc._exporter.stream_pages = empty_pages async def fake_stream_upload(pages, date_str): return 0 - orc._uploader._stream_upload = fake_stream_upload + orc._uploader.stream_upload = fake_stream_upload result = await orc._export(date(2026, 4, 9)) assert result == 0 @@ -1497,9 +1497,7 @@ class TestStreamPages: new_callable=lambda: property(lambda self: mock_client), ): chunks = [] - async for chunk in exporter._stream_pages( - "2026-04-10", connection_id="c-1" - ): + async for chunk in exporter.stream_pages("2026-04-10", connection_id="c-1"): chunks.append(chunk) assert len(chunks) >= 1 @@ -1520,9 +1518,7 @@ class TestStreamPages: new_callable=lambda: property(lambda self: mock_client), ): chunks = [] - async for chunk in exporter._stream_pages( - "2026-04-10", connection_id="c-1" - ): + async for chunk in exporter.stream_pages("2026-04-10", connection_id="c-1"): chunks.append(chunk) assert chunks == [] @@ -1554,7 +1550,7 @@ class TestStreamPages: "_prisma_client", new_callable=lambda: property(lambda self: mock_client), ): - async for _ in exporter._stream_pages( + async for _ in exporter.stream_pages( "2026-04-10", connection_id="c", page_size=3 ): pass @@ -1609,7 +1605,7 @@ class TestExporterNoDb: new_callable=lambda: property(lambda self: None), ): with pytest.raises(RuntimeError, match="database not connected"): - async for _ in exporter._stream_pages("2026-04-10", connection_id="c"): + async for _ in exporter.stream_pages("2026-04-10", connection_id="c"): pass @pytest.mark.asyncio @@ -1624,6 +1620,6 @@ class TestExporterNoDb: new_callable=lambda: property(lambda self: mock_client), ): chunks = [] - async for chunk in exporter._stream_pages("2026-04-10", connection_id="c"): + async for chunk in exporter.stream_pages("2026-04-10", connection_id="c"): chunks.append(chunk) assert chunks == [] diff --git a/tests/test_litellm/integrations/mavvrik/test_settings.py b/tests/test_litellm/integrations/mavvrik/test_settings.py index 9fb0ad5a0d2..d1d2f7c79db 100644 --- a/tests/test_litellm/integrations/mavvrik/test_settings.py +++ b/tests/test_litellm/integrations/mavvrik/test_settings.py @@ -71,10 +71,13 @@ class TestIsSetup: "MAVVRIK_CONNECTION_ID", ) } - with patch.dict("os.environ", env), patch.object( - type(s), - "_prisma_client", - new_callable=lambda: property(lambda self: mock_client), + 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() @@ -94,10 +97,13 @@ class TestIsSetup: "MAVVRIK_CONNECTION_ID", ) } - with patch.dict("os.environ", env), patch.object( - type(s), - "_prisma_client", - new_callable=lambda: property(lambda self: mock_client), + 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() @@ -115,8 +121,13 @@ class TestIsSetup: "MAVVRIK_CONNECTION_ID", ) } - with patch.dict("os.environ", env), patch.object( - type(s), "_prisma_client", new_callable=lambda: property(lambda self: None) + with ( + patch.dict("os.environ", env), + patch.object( + type(s), + "_prisma_client", + new_callable=lambda: property(lambda self: None), + ), ): result = await s.is_setup() @@ -135,13 +146,16 @@ class TestSave: 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: + 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", @@ -177,11 +191,14 @@ class TestLoad: ) 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"): + 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" @@ -312,10 +329,13 @@ class TestIsSetupDbException: "MAVVRIK_CONNECTION_ID", ) } - with patch.dict("os.environ", env), patch.object( - type(s), - "_prisma_client", - new_callable=lambda: property(lambda self: mock_client), + 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 @@ -380,11 +400,14 @@ class TestLoadEdgeCases: {"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 ( + 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() diff --git a/tests/test_litellm/integrations/mavvrik/test_uploader.py b/tests/test_litellm/integrations/mavvrik/test_uploader.py index 95f12f2bd85..c30e2de91e5 100644 --- a/tests/test_litellm/integrations/mavvrik/test_uploader.py +++ b/tests/test_litellm/integrations/mavvrik/test_uploader.py @@ -389,7 +389,7 @@ class TestStreamUpload: 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") + count = await u.stream_upload(pages(), date_str="2026-01-01") assert count > 0 assert len(put_calls) >= 1 @@ -417,7 +417,7 @@ class TestStreamUpload: return yield # make it a generator - count = await u._stream_upload(empty_pages(), date_str="2026-01-01") + count = await u.stream_upload(empty_pages(), date_str="2026-01-01") assert count == 0 mock_url.assert_not_called() @@ -508,7 +508,7 @@ class TestStreamUploadEdgeCases: ), patch.object(u, "_put_chunk", new_callable=AsyncMock), ): - result = await u._stream_upload(pages_with_empty(), date_str="2026-04-01") + result = await u.stream_upload(pages_with_empty(), date_str="2026-04-01") assert result > 0 # data was uploaded despite the empty first chunk @@ -551,7 +551,7 @@ class TestStreamUploadEdgeCases: ), patch.object(u, "_put_chunk", side_effect=fake_put), ): - await u._stream_upload(large_pages(), date_str="2026-04-01") + 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"]]