scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-03 02:22:24 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 81

refactor(mavvrik): address review comments #3 #4 #5 #6 #7

Browse source 
- Move MAVVRIK_EXPORT_INTERVAL_MINUTES/JOB_NAME to top-level imports (#3)
- Extract _schedule_job() helper to eliminate 3× duplicated scheduler
  registration blocks in initialize(), update_settings(), proxy_server.py (#4)
- delete() now logs informatively on LookupError instead of bare pass (#5)
- Logger.__init__ warns at startup if credentials are not configured (#6)
- Rename _stream_pages/_stream_upload to stream_pages/stream_upload (#7)

Co-Authored-By: Claude Sonnet 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
Praveen Ghuge 2026-05-12 22:18:55 +05:30
parent 28ddef18ee
commit 388e2a0c08
8 changed files with 128 additions and 99 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

83
litellm/integrations/mavvrik/__init__.py
View file

@ -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"}

6
litellm/integrations/mavvrik/exporter.py
View file

@ -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,

13
litellm/integrations/mavvrik/logger.py
View file

@ -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."
)

6
litellm/integrations/mavvrik/orchestrator.py
View file

@ -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

2
litellm/integrations/mavvrik/uploader.py
View file

@ -167,7 +167,7 @@ class Uploader:
# Streaming upload
# ------------------------------------------------------------------
async def _stream_upload(
async def stream_upload(
self,
pages: AsyncIterator[str],
date_str: str,

24
tests/test_litellm/integrations/mavvrik/test_service.py
View file

@ -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 == []

85
tests/test_litellm/integrations/mavvrik/test_settings.py
View file

@ -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()

8
tests/test_litellm/integrations/mavvrik/test_uploader.py
View file

@ -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"]]