mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-10 03:28:53 +00:00
refactor: lazy-load focus export engine and isolate polars deps
This commit is contained in:
Yuta Saito
parent
cc3b068f22
commit
844577f2d4
2 changed files with 145 additions and 95 deletions
124
litellm/integrations/focus/export_engine.py
Normal file
124
litellm/integrations/focus/export_engine.py
Normal file
|
|
@ -0,0 +1,124 @@
|
|||
"""Core export engine for Focus integrations (heavy dependencies)."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any, Dict, Optional
|
||||
|
||||
import polars as pl
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
|
||||
from .database import FocusLiteLLMDatabase
|
||||
from .destinations import FocusDestinationFactory, FocusTimeWindow
|
||||
from .serializers import FocusParquetSerializer, FocusSerializer
|
||||
from .transformer import FocusTransformer
|
||||
|
||||
|
||||
class FocusExportEngine:
|
||||
"""Engine that fetches, normalizes, and uploads Focus exports."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
provider: str,
|
||||
export_format: str,
|
||||
prefix: str,
|
||||
destination_config: Optional[dict[str, Any]] = None,
|
||||
) -> None:
|
||||
self.provider = provider
|
||||
self.export_format = export_format
|
||||
self.prefix = prefix
|
||||
self._destination = FocusDestinationFactory.create(
|
||||
provider=self.provider,
|
||||
prefix=self.prefix,
|
||||
config=destination_config,
|
||||
)
|
||||
self._serializer = self._init_serializer()
|
||||
self._transformer = FocusTransformer()
|
||||
self._database = FocusLiteLLMDatabase()
|
||||
|
||||
def _init_serializer(self) -> FocusSerializer:
|
||||
if self.export_format != "parquet":
|
||||
raise NotImplementedError("Only parquet export supported currently")
|
||||
return FocusParquetSerializer()
|
||||
|
||||
async def dry_run_export_usage_data(self, limit: Optional[int]) -> Dict[str, Any]:
|
||||
data = await self._database.get_usage_data(limit=limit)
|
||||
normalized = self._transformer.transform(data)
|
||||
|
||||
usage_sample = data.head(min(50, len(data))).to_dicts()
|
||||
normalized_sample = normalized.head(min(50, len(normalized))).to_dicts()
|
||||
|
||||
summary = {
|
||||
"total_records": len(normalized),
|
||||
"total_spend": self._sum_column(normalized, "spend"),
|
||||
"total_tokens": self._sum_column(normalized, "total_tokens"),
|
||||
"unique_teams": self._count_unique(normalized, "team_id"),
|
||||
"unique_models": self._count_unique(normalized, "model"),
|
||||
}
|
||||
|
||||
return {
|
||||
"usage_data": usage_sample,
|
||||
"normalized_data": normalized_sample,
|
||||
"summary": summary,
|
||||
}
|
||||
|
||||
async def export_window(
|
||||
self,
|
||||
*,
|
||||
window: FocusTimeWindow,
|
||||
limit: Optional[int],
|
||||
) -> None:
|
||||
data = await self._database.get_usage_data(
|
||||
limit=limit,
|
||||
start_time_utc=window.start_time,
|
||||
end_time_utc=window.end_time,
|
||||
)
|
||||
if data.is_empty():
|
||||
verbose_logger.debug("Focus export: no usage data for window %s", window)
|
||||
return
|
||||
|
||||
normalized = self._transformer.transform(data)
|
||||
if normalized.is_empty():
|
||||
verbose_logger.debug(
|
||||
"Focus export: normalized data empty for window %s", window
|
||||
)
|
||||
return
|
||||
|
||||
await self._serialize_and_upload(normalized, window)
|
||||
|
||||
async def _serialize_and_upload(
|
||||
self, frame: pl.DataFrame, window: FocusTimeWindow
|
||||
) -> None:
|
||||
payload = self._serializer.serialize(frame)
|
||||
if not payload:
|
||||
verbose_logger.debug("Focus export: serializer returned empty payload")
|
||||
return
|
||||
await self._destination.deliver(
|
||||
content=payload,
|
||||
time_window=window,
|
||||
filename=self._build_filename(),
|
||||
)
|
||||
|
||||
def _build_filename(self) -> str:
|
||||
if not self._serializer.extension:
|
||||
raise ValueError("Serializer must declare a file extension")
|
||||
return f"usage.{self._serializer.extension}"
|
||||
|
||||
@staticmethod
|
||||
def _sum_column(frame: pl.DataFrame, column: str) -> float:
|
||||
if frame.is_empty() or column not in frame.columns:
|
||||
return 0.0
|
||||
value = frame.select(pl.col(column).sum().alias("sum")).row(0)[0]
|
||||
if value is None:
|
||||
return 0.0
|
||||
return float(value)
|
||||
|
||||
@staticmethod
|
||||
def _count_unique(frame: pl.DataFrame, column: str) -> int:
|
||||
if frame.is_empty() or column not in frame.columns:
|
||||
return 0
|
||||
value = frame.select(pl.col(column).n_unique().alias("unique")).row(0)[0]
|
||||
if value is None:
|
||||
return 0
|
||||
return int(value)
|
||||
|
|
@ -6,22 +6,15 @@ import os
|
|||
from datetime import datetime, timedelta, timezone
|
||||
from typing import TYPE_CHECKING, Any, Dict, List, Optional, cast
|
||||
|
||||
import polars as pl
|
||||
|
||||
import litellm
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
|
||||
from .database import FocusLiteLLMDatabase
|
||||
from .destinations import (
|
||||
FocusDestinationFactory,
|
||||
FocusTimeWindow,
|
||||
)
|
||||
from .serializers import FocusParquetSerializer, FocusSerializer
|
||||
from .transformer import FocusTransformer
|
||||
from .destinations import FocusTimeWindow
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from apscheduler.schedulers.asyncio import AsyncIOScheduler
|
||||
from .export_engine import FocusExportEngine
|
||||
else:
|
||||
AsyncIOScheduler = Any
|
||||
|
||||
|
|
@ -66,20 +59,21 @@ class FocusLogger(CustomLogger):
|
|||
prefix if prefix is not None else (env_prefix if env_prefix else "focus_exports")
|
||||
)
|
||||
|
||||
self._destination = FocusDestinationFactory.create(
|
||||
provider=self.provider,
|
||||
prefix=self.prefix,
|
||||
config=destination_config,
|
||||
)
|
||||
self._serializer = self._init_serializer()
|
||||
self._transformer = FocusTransformer()
|
||||
self._database = FocusLiteLLMDatabase()
|
||||
self._destination_config = destination_config
|
||||
self._engine: Optional["FocusExportEngine"] = None
|
||||
|
||||
def _init_serializer(self) -> FocusSerializer:
|
||||
"""Return serializer implementation for requested format."""
|
||||
if self.export_format != "parquet":
|
||||
raise NotImplementedError("Only parquet export supported currently")
|
||||
return FocusParquetSerializer()
|
||||
def _ensure_engine(self) -> "FocusExportEngine":
|
||||
"""Instantiate the heavy export engine lazily."""
|
||||
if self._engine is None:
|
||||
from .export_engine import FocusExportEngine
|
||||
|
||||
self._engine = FocusExportEngine(
|
||||
provider=self.provider,
|
||||
export_format=self.export_format,
|
||||
prefix=self.prefix,
|
||||
destination_config=self._destination_config,
|
||||
)
|
||||
return self._engine
|
||||
|
||||
async def export_usage_data(
|
||||
self,
|
||||
|
|
@ -108,25 +102,8 @@ class FocusLogger(CustomLogger):
|
|||
self, limit: Optional[int] = DEFAULT_DRY_RUN_LIMIT
|
||||
) -> dict[str, Any]:
|
||||
"""Return transformed data without uploading."""
|
||||
data = await self._database.get_usage_data(limit=limit)
|
||||
normalized = self._transformer.transform(data)
|
||||
|
||||
usage_sample = data.head(min(50, len(data))).to_dicts()
|
||||
normalized_sample = normalized.head(min(50, len(normalized))).to_dicts()
|
||||
|
||||
summary = {
|
||||
"total_records": len(normalized),
|
||||
"total_spend": self._sum_column(normalized, "spend"),
|
||||
"total_tokens": self._sum_column(normalized, "total_tokens"),
|
||||
"unique_teams": self._count_unique(normalized, "team_id"),
|
||||
"unique_models": self._count_unique(normalized, "model"),
|
||||
}
|
||||
|
||||
return {
|
||||
"usage_data": usage_sample,
|
||||
"normalized_data": normalized_sample,
|
||||
"summary": summary,
|
||||
}
|
||||
engine = self._ensure_engine()
|
||||
return await engine.dry_run_export_usage_data(limit=limit)
|
||||
|
||||
async def initialize_focus_export_job(self) -> None:
|
||||
"""Entry point for scheduler jobs to run export cycle with locking."""
|
||||
|
|
@ -207,23 +184,8 @@ class FocusLogger(CustomLogger):
|
|||
window: FocusTimeWindow,
|
||||
limit: Optional[int],
|
||||
) -> None:
|
||||
data = await self._database.get_usage_data(
|
||||
limit=limit,
|
||||
start_time_utc=window.start_time,
|
||||
end_time_utc=window.end_time,
|
||||
)
|
||||
if data.is_empty():
|
||||
verbose_logger.debug("Focus export: no usage data for window %s", window)
|
||||
return
|
||||
|
||||
normalized = self._transformer.transform(data)
|
||||
if normalized.is_empty():
|
||||
verbose_logger.debug(
|
||||
"Focus export: normalized data empty for window %s", window
|
||||
)
|
||||
return
|
||||
|
||||
await self._serialize_and_upload(normalized, window)
|
||||
engine = self._ensure_engine()
|
||||
await engine.export_window(window=window, limit=limit)
|
||||
|
||||
def _compute_time_window(self, now: datetime) -> FocusTimeWindow:
|
||||
"""Derive the time window to export based on configured frequency."""
|
||||
|
|
@ -246,40 +208,4 @@ class FocusLogger(CustomLogger):
|
|||
frequency=self.frequency,
|
||||
)
|
||||
|
||||
async def _serialize_and_upload(
|
||||
self, frame: pl.DataFrame, window: FocusTimeWindow
|
||||
) -> None:
|
||||
"""Serialize the normalized frame and upload via destination."""
|
||||
payload = self._serializer.serialize(frame)
|
||||
if not payload:
|
||||
verbose_logger.debug("Focus export: serializer returned empty payload")
|
||||
return
|
||||
await self._destination.deliver(
|
||||
content=payload,
|
||||
time_window=window,
|
||||
filename=self._build_filename(),
|
||||
)
|
||||
|
||||
def _build_filename(self) -> str:
|
||||
"""Return the canonical file name for exports."""
|
||||
if not self._serializer.extension:
|
||||
raise ValueError("Serializer must declare a file extension")
|
||||
return f"usage.{self._serializer.extension}"
|
||||
|
||||
@staticmethod
|
||||
def _sum_column(frame: pl.DataFrame, column: str) -> float:
|
||||
if frame.is_empty() or column not in frame.columns:
|
||||
return 0.0
|
||||
value = frame.select(pl.col(column).sum().alias("sum")).row(0)[0]
|
||||
if value is None:
|
||||
return 0.0
|
||||
return float(value)
|
||||
|
||||
@staticmethod
|
||||
def _count_unique(frame: pl.DataFrame, column: str) -> int:
|
||||
if frame.is_empty() or column not in frame.columns:
|
||||
return 0
|
||||
value = frame.select(pl.col(column).n_unique().alias("unique")).row(0)[0]
|
||||
if value is None:
|
||||
return 0
|
||||
return int(value)
|
||||
__all__ = ["FocusLogger"]
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue