From 565a622cf9d67147b038b927078a1f3ac01de366 Mon Sep 17 00:00:00 2001 From: Yuta Saito Date: Thu, 8 Jan 2026 12:09:58 +0900 Subject: [PATCH 1/9] feat: scaffold Focus export logging skeleton --- litellm/integrations/focus/__init__.py | 0 litellm/integrations/focus/database.py | 22 +++ .../focus/destinations/__init__.py | 12 ++ .../integrations/focus/destinations/base.py | 30 ++++ .../focus/destinations/factory.py | 21 +++ .../focus/destinations/s3_destination.py | 32 +++++ .../integrations/focus/focus_export_logger.py | 129 ++++++++++++++++++ .../focus/serializers/__init__.py | 6 + .../integrations/focus/serializers/base.py | 18 +++ .../integrations/focus/serializers/parquet.py | 16 +++ litellm/integrations/focus/transformer.py | 13 ++ 11 files changed, 299 insertions(+) create mode 100644 litellm/integrations/focus/__init__.py create mode 100644 litellm/integrations/focus/database.py create mode 100644 litellm/integrations/focus/destinations/__init__.py create mode 100644 litellm/integrations/focus/destinations/base.py create mode 100644 litellm/integrations/focus/destinations/factory.py create mode 100644 litellm/integrations/focus/destinations/s3_destination.py create mode 100644 litellm/integrations/focus/focus_export_logger.py create mode 100644 litellm/integrations/focus/serializers/__init__.py create mode 100644 litellm/integrations/focus/serializers/base.py create mode 100644 litellm/integrations/focus/serializers/parquet.py create mode 100644 litellm/integrations/focus/transformer.py diff --git a/litellm/integrations/focus/__init__.py b/litellm/integrations/focus/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/litellm/integrations/focus/database.py b/litellm/integrations/focus/database.py new file mode 100644 index 00000000000..a23d16278af --- /dev/null +++ b/litellm/integrations/focus/database.py @@ -0,0 +1,22 @@ +"""Database access helpers for Focus export.""" + +from __future__ import annotations + +from datetime import datetime +from typing import Optional + +import polars as pl + + +class FocusLiteLLMDatabase: + """Retrieves LiteLLM usage data for Focus export workflows.""" + + async def get_usage_data( + self, + *, + limit: Optional[int] = None, + start_time_utc: Optional[datetime] = None, + end_time_utc: Optional[datetime] = None, + ) -> pl.DataFrame: + """Return usage data for the requested window.""" + raise NotImplementedError diff --git a/litellm/integrations/focus/destinations/__init__.py b/litellm/integrations/focus/destinations/__init__.py new file mode 100644 index 00000000000..233f1da0c9b --- /dev/null +++ b/litellm/integrations/focus/destinations/__init__.py @@ -0,0 +1,12 @@ +"""Destination implementations for Focus export.""" + +from .base import FocusDestination, FocusTimeWindow +from .factory import FocusDestinationFactory +from .s3_destination import FocusS3Destination + +__all__ = [ + "FocusDestination", + "FocusDestinationFactory", + "FocusTimeWindow", + "FocusS3Destination", +] diff --git a/litellm/integrations/focus/destinations/base.py b/litellm/integrations/focus/destinations/base.py new file mode 100644 index 00000000000..8042a7e23b9 --- /dev/null +++ b/litellm/integrations/focus/destinations/base.py @@ -0,0 +1,30 @@ +"""Abstract destination interfaces for Focus export.""" + +from __future__ import annotations + +from dataclasses import dataclass +from datetime import datetime +from typing import Protocol + + +@dataclass(frozen=True) +class FocusTimeWindow: + """Represents the span of data exported in a single batch.""" + + start_time: datetime + end_time: datetime + frequency: str + + +class FocusDestination(Protocol): + """Protocol for anything that can receive Focus export files.""" + + async def deliver( + self, + *, + content: bytes, + time_window: FocusTimeWindow, + filename: str, + ) -> None: + """Persist the serialized export for the provided time window.""" + ... diff --git a/litellm/integrations/focus/destinations/factory.py b/litellm/integrations/focus/destinations/factory.py new file mode 100644 index 00000000000..36163ece61c --- /dev/null +++ b/litellm/integrations/focus/destinations/factory.py @@ -0,0 +1,21 @@ +"""Factory helpers for Focus export destinations.""" + +from __future__ import annotations + +from typing import Optional + +from .base import FocusDestination + + +class FocusDestinationFactory: + """Builds destination instances based on provider/config settings.""" + + @staticmethod + def create( + *, + provider: str, + prefix: str, + config: Optional[dict] = None, + ) -> FocusDestination: + """Return a destination implementation for the requested provider.""" + raise NotImplementedError diff --git a/litellm/integrations/focus/destinations/s3_destination.py b/litellm/integrations/focus/destinations/s3_destination.py new file mode 100644 index 00000000000..7ead87baa29 --- /dev/null +++ b/litellm/integrations/focus/destinations/s3_destination.py @@ -0,0 +1,32 @@ +"""S3 destination implementation for Focus export.""" + +from __future__ import annotations + +from typing import Any, Optional + +from .base import FocusDestination, FocusTimeWindow + + +class FocusS3Destination(FocusDestination): + """Handles uploading serialized exports to S3 buckets.""" + + def __init__( + self, + *, + prefix: str, + config: Optional[dict[str, Any]] = None, + ) -> None: + self.prefix = prefix.rstrip("/") + self.config = config or {} + + async def deliver( + self, + *, + content: bytes, + time_window: FocusTimeWindow, + filename: str, + ) -> None: + raise NotImplementedError + + def _build_object_key(self, *, time_window: FocusTimeWindow, filename: str) -> str: + raise NotImplementedError diff --git a/litellm/integrations/focus/focus_export_logger.py b/litellm/integrations/focus/focus_export_logger.py new file mode 100644 index 00000000000..723dbc706fb --- /dev/null +++ b/litellm/integrations/focus/focus_export_logger.py @@ -0,0 +1,129 @@ +"""Focus export logger orchestrating DB pull/transform/upload.""" + +from __future__ import annotations + +import os +from datetime import datetime, timedelta, timezone +from typing import TYPE_CHECKING, Any, Optional + +import polars as pl + +from litellm.integrations.custom_logger import CustomLogger + +from .destinations import ( + FocusDestination, + FocusDestinationFactory, + FocusTimeWindow, +) +from .serializers import FocusParquetSerializer, FocusSerializer +from .transformer import FocusTransformer + +if TYPE_CHECKING: + from apscheduler.schedulers.asyncio import AsyncIOScheduler +else: + AsyncIOScheduler = Any + + +class FocusExportLogger(CustomLogger): + """Coordinates Focus export jobs across transformer/serializer/destination layers.""" + + def __init__( + self, + *, + provider: Optional[str] = None, + export_format: Optional[str] = None, + frequency: Optional[str] = None, + cron_offset_minute: Optional[int] = None, + interval_seconds: Optional[int] = None, + prefix: Optional[str] = None, + destination_config: Optional[dict[str, Any]] = None, + **kwargs: Any, + ) -> None: + super().__init__(**kwargs) + self.provider = (provider or os.getenv("FOCUS_EXPORT_PROVIDER") or "s3").lower() + self.export_format = ( + export_format or os.getenv("FOCUS_EXPORT_FORMAT") or "parquet" + ).lower() + self.frequency = ( + frequency or os.getenv("FOCUS_EXPORT_FREQUENCY") or "hourly" + ).lower() + self.cron_offset_minute = ( + cron_offset_minute + if cron_offset_minute is not None + else int(os.getenv("FOCUS_EXPORT_CRON_OFFSET", "5")) + ) + self.interval_seconds = ( + interval_seconds + if interval_seconds is not None + else os.getenv("FOCUS_EXPORT_INTERVAL_SECONDS") + ) + self.prefix = prefix or os.getenv("FOCUS_EXPORT_PREFIX", "focus_exports") + + self._destination = self._init_destination( + destination_config=destination_config, + ) + self._serializer = self._init_serializer() + self._transformer = FocusTransformer() + + 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 _init_destination( + self, + *, + destination_config: Optional[dict[str, Any]], + ) -> FocusDestination: + """Factory for destination implementations.""" + resolved_config = self._resolve_destination_config(destination_config) + return FocusDestinationFactory.create( + provider=self.provider, + prefix=self.prefix, + config=resolved_config, + ) + + def _resolve_destination_config( + self, + destination_config: Optional[dict[str, Any]], + ) -> dict[str, Any]: + """Collect provider-specific configuration for destination creation.""" + raise NotImplementedError + + async def export_usage_data(self) -> None: + """Public hook to trigger export immediately.""" + raise NotImplementedError + + async def dry_run_export_usage_data(self) -> dict: + """Return transformed data without uploading.""" + raise NotImplementedError + + async def initialize_focus_export_job(self) -> None: + """Entry point for scheduler jobs to run export cycle with locking.""" + raise NotImplementedError + + @staticmethod + async def init_focus_export_background_job( + scheduler: AsyncIOScheduler, + ) -> None: + """Register the export cron/interval job with the provided scheduler.""" + raise NotImplementedError + + def _compute_time_window(self, now: datetime) -> FocusTimeWindow: + """Derive the time window to export based on configured frequency.""" + raise NotImplementedError + + def _serialize_and_upload( + self, + frame: pl.DataFrame, + window: FocusTimeWindow, + ) -> None: + """Helper stub for serializing and delegating to destination.""" + raise NotImplementedError + + 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}" diff --git a/litellm/integrations/focus/serializers/__init__.py b/litellm/integrations/focus/serializers/__init__.py new file mode 100644 index 00000000000..18187bf73e5 --- /dev/null +++ b/litellm/integrations/focus/serializers/__init__.py @@ -0,0 +1,6 @@ +"""Serializer package exports for Focus integration.""" + +from .base import FocusSerializer +from .parquet import FocusParquetSerializer + +__all__ = ["FocusSerializer", "FocusParquetSerializer"] diff --git a/litellm/integrations/focus/serializers/base.py b/litellm/integrations/focus/serializers/base.py new file mode 100644 index 00000000000..6da080dae81 --- /dev/null +++ b/litellm/integrations/focus/serializers/base.py @@ -0,0 +1,18 @@ +"""Serializer abstractions for Focus export.""" + +from __future__ import annotations + +from abc import ABC, abstractmethod + +import polars as pl + + +class FocusSerializer(ABC): + """Base serializer turning Focus frames into bytes.""" + + extension: str = "" + + @abstractmethod + def serialize(self, frame: pl.DataFrame) -> bytes: + """Convert the normalized Focus frame into the chosen format.""" + raise NotImplementedError diff --git a/litellm/integrations/focus/serializers/parquet.py b/litellm/integrations/focus/serializers/parquet.py new file mode 100644 index 00000000000..3d42337e388 --- /dev/null +++ b/litellm/integrations/focus/serializers/parquet.py @@ -0,0 +1,16 @@ +"""Parquet serializer for Focus export.""" + +from __future__ import annotations + +import polars as pl + +from .base import FocusSerializer + + +class FocusParquetSerializer(FocusSerializer): + """Placeholder Parquet serializer implementation.""" + + extension = "parquet" + + def serialize(self, frame: pl.DataFrame) -> bytes: + raise NotImplementedError diff --git a/litellm/integrations/focus/transformer.py b/litellm/integrations/focus/transformer.py new file mode 100644 index 00000000000..45a77dd8ac5 --- /dev/null +++ b/litellm/integrations/focus/transformer.py @@ -0,0 +1,13 @@ +"""Focus export data transformer.""" + +from __future__ import annotations + +import polars as pl + + +class FocusTransformer: + """Transforms LiteLLM DB rows into Focus-compatible schema.""" + + def transform(self, frame: pl.DataFrame) -> pl.DataFrame: + """Return a normalized frame expected by downstream serializers.""" + raise NotImplementedError From 3af9bab6a504ccc9e78730d229ebfcea1adddfe3 Mon Sep 17 00:00:00 2001 From: Yuta Saito Date: Thu, 8 Jan 2026 12:17:14 +0900 Subject: [PATCH 2/9] feat: add focus schema --- litellm/integrations/focus/schema.py | 33 +++++++++++++++++++++++ litellm/integrations/focus/transformer.py | 4 +++ 2 files changed, 37 insertions(+) create mode 100644 litellm/integrations/focus/schema.py diff --git a/litellm/integrations/focus/schema.py b/litellm/integrations/focus/schema.py new file mode 100644 index 00000000000..766e29c0856 --- /dev/null +++ b/litellm/integrations/focus/schema.py @@ -0,0 +1,33 @@ +"""Schema definitions for Focus export data.""" + +from __future__ import annotations + +import polars as pl + + +FOCUS_NORMALIZED_SCHEMA = pl.Schema( + { + "usage_date": pl.Datetime(time_unit="us"), + "team_id": pl.String, + "team_alias": pl.String, + "user_id": pl.String, + "user_email": pl.String, + "api_key_alias": pl.String, + "model": pl.String, + "model_group": pl.String, + "custom_llm_provider": pl.String, + "prompt_tokens": pl.Int64, + "completion_tokens": pl.Int64, + "total_tokens": pl.Int64, + "spend": pl.Float64, + "cache_creation_input_tokens": pl.Int64, + "cache_read_input_tokens": pl.Int64, + "api_requests": pl.Int64, + "successful_requests": pl.Int64, + "failed_requests": pl.Int64, + "created_at": pl.Datetime(time_unit="us"), + "updated_at": pl.Datetime(time_unit="us"), + } +) + +__all__ = ["FOCUS_NORMALIZED_SCHEMA"] diff --git a/litellm/integrations/focus/transformer.py b/litellm/integrations/focus/transformer.py index 45a77dd8ac5..612082505fa 100644 --- a/litellm/integrations/focus/transformer.py +++ b/litellm/integrations/focus/transformer.py @@ -4,10 +4,14 @@ from __future__ import annotations import polars as pl +from .schema import FOCUS_NORMALIZED_SCHEMA + class FocusTransformer: """Transforms LiteLLM DB rows into Focus-compatible schema.""" + schema = FOCUS_NORMALIZED_SCHEMA + def transform(self, frame: pl.DataFrame) -> pl.DataFrame: """Return a normalized frame expected by downstream serializers.""" raise NotImplementedError From c00d83fea2338429991a355d59c6d1f548f30a27 Mon Sep 17 00:00:00 2001 From: Yuta Saito Date: Thu, 8 Jan 2026 16:48:06 +0900 Subject: [PATCH 3/9] feat: add support focus export --- litellm/__init__.py | 1 + litellm/integrations/focus/database.py | 95 +++++- .../focus/destinations/factory.py | 44 ++- .../focus/destinations/s3_destination.py | 48 ++- .../integrations/focus/focus_export_logger.py | 129 -------- litellm/integrations/focus/focus_logger.py | 285 ++++++++++++++++++ litellm/integrations/focus/schema.py | 65 ++-- .../integrations/focus/serializers/parquet.py | 11 +- litellm/integrations/focus/transformer.py | 89 +++++- .../custom_logger_registry.py | 2 + litellm/litellm_core_utils/litellm_logging.py | 15 + litellm/proxy/proxy_server.py | 11 +- 12 files changed, 632 insertions(+), 163 deletions(-) delete mode 100644 litellm/integrations/focus/focus_export_logger.py create mode 100644 litellm/integrations/focus/focus_logger.py diff --git a/litellm/__init__.py b/litellm/__init__.py index 8af9a7d10a6..76bdac40c2d 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -134,6 +134,7 @@ _custom_logger_compatible_callbacks_literal = Literal[ "bitbucket", "gitlab", "cloudzero", + "focus", "posthog", "levo", ] diff --git a/litellm/integrations/focus/database.py b/litellm/integrations/focus/database.py index a23d16278af..402a1ec4b4e 100644 --- a/litellm/integrations/focus/database.py +++ b/litellm/integrations/focus/database.py @@ -3,7 +3,7 @@ from __future__ import annotations from datetime import datetime -from typing import Optional +from typing import Any, Dict, Optional import polars as pl @@ -11,6 +11,16 @@ import polars as pl class FocusLiteLLMDatabase: """Retrieves LiteLLM usage data for Focus export workflows.""" + def _ensure_prisma_client(self): + from litellm.proxy.proxy_server import prisma_client + + if prisma_client is None: + raise RuntimeError( + "Database not connected. Connect a database to your proxy - " + "https://docs.litellm.ai/docs/simple_proxy#managing-auth---virtual-keys" + ) + return prisma_client + async def get_usage_data( self, *, @@ -19,4 +29,85 @@ class FocusLiteLLMDatabase: end_time_utc: Optional[datetime] = None, ) -> pl.DataFrame: """Return usage data for the requested window.""" - raise NotImplementedError + client = self._ensure_prisma_client() + + where_clauses = [] + query_params = [] + placeholder_index = 1 + if start_time_utc: + where_clauses.append(f"dus.updated_at >= ${placeholder_index}::timestamptz") + query_params.append(start_time_utc) + placeholder_index += 1 + if end_time_utc: + where_clauses.append(f"dus.updated_at <= ${placeholder_index}::timestamptz") + query_params.append(end_time_utc) + placeholder_index += 1 + + where_clause = "" + if where_clauses: + where_clause = "WHERE " + " AND ".join(where_clauses) + + limit_clause = "" + if limit is not None: + try: + limit_value = int(limit) + except (TypeError, ValueError) as exc: # pragma: no cover - defensive guard + raise ValueError("limit must be an integer") from exc + if limit_value < 0: + raise ValueError("limit must be non-negative") + limit_clause = f" LIMIT ${placeholder_index}" + query_params.append(limit_value) + + query = f""" + SELECT + dus.id, + dus.date, + dus.user_id, + dus.api_key, + dus.model, + dus.model_group, + dus.custom_llm_provider, + dus.prompt_tokens, + dus.completion_tokens, + dus.spend, + dus.api_requests, + dus.successful_requests, + dus.failed_requests, + dus.cache_creation_input_tokens, + dus.cache_read_input_tokens, + dus.created_at, + dus.updated_at, + vt.team_id, + vt.key_alias as api_key_alias, + tt.team_alias, + ut.user_email as user_email + FROM "LiteLLM_DailyUserSpend" dus + LEFT JOIN "LiteLLM_VerificationToken" vt ON dus.api_key = vt.token + LEFT JOIN "LiteLLM_TeamTable" tt ON vt.team_id = tt.team_id + LEFT JOIN "LiteLLM_UserTable" ut ON dus.user_id = ut.user_id + {where_clause} + ORDER BY dus.date DESC, dus.created_at DESC + {limit_clause} + """ + + try: + db_response = await client.db.query_raw(query, *query_params) + return pl.DataFrame(db_response, infer_schema_length=None) + except Exception as exc: + raise RuntimeError(f"Error retrieving usage data: {exc}") from exc + + async def get_table_info(self) -> Dict[str, Any]: + """Return metadata about the spend table for diagnostics.""" + client = self._ensure_prisma_client() + + info_query = """ + SELECT column_name, data_type, is_nullable + FROM information_schema.columns + WHERE table_name = 'LiteLLM_DailyUserSpend' + ORDER BY ordinal_position; + """ + try: + columns_response = await client.db.query_raw(info_query) + return {"columns": columns_response, "table_name": "LiteLLM_DailyUserSpend"} + except Exception as exc: + raise RuntimeError(f"Error getting table info: {exc}") from exc diff --git a/litellm/integrations/focus/destinations/factory.py b/litellm/integrations/focus/destinations/factory.py index 36163ece61c..6f5dc5dc06a 100644 --- a/litellm/integrations/focus/destinations/factory.py +++ b/litellm/integrations/focus/destinations/factory.py @@ -2,9 +2,11 @@ from __future__ import annotations -from typing import Optional +import os +from typing import Any, Dict, Optional from .base import FocusDestination +from .s3_destination import FocusS3Destination class FocusDestinationFactory: @@ -15,7 +17,43 @@ class FocusDestinationFactory: *, provider: str, prefix: str, - config: Optional[dict] = None, + config: Optional[Dict[str, Any]] = None, ) -> FocusDestination: """Return a destination implementation for the requested provider.""" - raise NotImplementedError + provider_lower = provider.lower() + normalized_config = FocusDestinationFactory._resolve_config( + provider=provider_lower, overrides=config or {} + ) + if provider_lower == "s3": + return FocusS3Destination(prefix=prefix, config=normalized_config) + raise NotImplementedError(f"Provider '{provider}' not supported for Focus export") + + @staticmethod + def _resolve_config( + *, + provider: str, + overrides: Dict[str, Any], + ) -> Dict[str, Any]: + if provider == "s3": + resolved = { + "bucket_name": overrides.get("bucket_name") + or os.getenv("FOCUS_S3_BUCKET_NAME"), + "region_name": overrides.get("region_name") + or os.getenv("FOCUS_S3_REGION_NAME"), + "endpoint_url": overrides.get("endpoint_url") + or os.getenv("FOCUS_S3_ENDPOINT_URL"), + "aws_access_key_id": overrides.get("aws_access_key_id") + or os.getenv("FOCUS_S3_ACCESS_KEY"), + "aws_secret_access_key": overrides.get("aws_secret_access_key") + or os.getenv("FOCUS_S3_SECRET_KEY"), + "aws_session_token": overrides.get("aws_session_token") + or os.getenv("FOCUS_S3_SESSION_TOKEN"), + } + if not resolved.get("bucket_name"): + raise ValueError( + "FOCUS_S3_BUCKET_NAME must be provided for S3 exports" + ) + 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" + ) diff --git a/litellm/integrations/focus/destinations/s3_destination.py b/litellm/integrations/focus/destinations/s3_destination.py index 7ead87baa29..c6d5554b438 100644 --- a/litellm/integrations/focus/destinations/s3_destination.py +++ b/litellm/integrations/focus/destinations/s3_destination.py @@ -2,8 +2,12 @@ from __future__ import annotations +import asyncio +from datetime import timezone from typing import Any, Optional +import boto3 + from .base import FocusDestination, FocusTimeWindow @@ -16,8 +20,13 @@ class FocusS3Destination(FocusDestination): prefix: str, config: Optional[dict[str, Any]] = None, ) -> None: + config = config or {} + bucket_name = config.get("bucket_name") + if not bucket_name: + raise ValueError("bucket_name must be provided for S3 destination") + self.bucket_name = bucket_name self.prefix = prefix.rstrip("/") - self.config = config or {} + self.config = config async def deliver( self, @@ -26,7 +35,40 @@ class FocusS3Destination(FocusDestination): time_window: FocusTimeWindow, filename: str, ) -> None: - raise NotImplementedError + object_key = self._build_object_key(time_window=time_window, filename=filename) + await asyncio.to_thread(self._upload, content, object_key) def _build_object_key(self, *, time_window: FocusTimeWindow, filename: str) -> str: - raise NotImplementedError + start_utc = time_window.start_time.astimezone(timezone.utc) + date_component = f"date={start_utc.strftime('%Y-%m-%d')}" + parts = [self.prefix, date_component] + if time_window.frequency == "hourly": + parts.append(f"hour={start_utc.strftime('%H')}") + key_prefix = "/".join(filter(None, parts)) + return f"{key_prefix}/{filename}" if key_prefix else filename + + def _upload(self, content: bytes, object_key: str) -> None: + client_kwargs: dict[str, Any] = {} + region_name = self.config.get("region_name") + if region_name: + client_kwargs["region_name"] = region_name + endpoint_url = self.config.get("endpoint_url") + if endpoint_url: + client_kwargs["endpoint_url"] = endpoint_url + + session_kwargs: dict[str, Any] = {} + for key in ( + "aws_access_key_id", + "aws_secret_access_key", + "aws_session_token", + ): + if self.config.get(key): + session_kwargs[key] = self.config[key] + + s3_client = boto3.client("s3", **client_kwargs, **session_kwargs) + s3_client.put_object( + Bucket=self.bucket_name, + Key=object_key, + Body=content, + ContentType="application/octet-stream", + ) diff --git a/litellm/integrations/focus/focus_export_logger.py b/litellm/integrations/focus/focus_export_logger.py deleted file mode 100644 index 723dbc706fb..00000000000 --- a/litellm/integrations/focus/focus_export_logger.py +++ /dev/null @@ -1,129 +0,0 @@ -"""Focus export logger orchestrating DB pull/transform/upload.""" - -from __future__ import annotations - -import os -from datetime import datetime, timedelta, timezone -from typing import TYPE_CHECKING, Any, Optional - -import polars as pl - -from litellm.integrations.custom_logger import CustomLogger - -from .destinations import ( - FocusDestination, - FocusDestinationFactory, - FocusTimeWindow, -) -from .serializers import FocusParquetSerializer, FocusSerializer -from .transformer import FocusTransformer - -if TYPE_CHECKING: - from apscheduler.schedulers.asyncio import AsyncIOScheduler -else: - AsyncIOScheduler = Any - - -class FocusExportLogger(CustomLogger): - """Coordinates Focus export jobs across transformer/serializer/destination layers.""" - - def __init__( - self, - *, - provider: Optional[str] = None, - export_format: Optional[str] = None, - frequency: Optional[str] = None, - cron_offset_minute: Optional[int] = None, - interval_seconds: Optional[int] = None, - prefix: Optional[str] = None, - destination_config: Optional[dict[str, Any]] = None, - **kwargs: Any, - ) -> None: - super().__init__(**kwargs) - self.provider = (provider or os.getenv("FOCUS_EXPORT_PROVIDER") or "s3").lower() - self.export_format = ( - export_format or os.getenv("FOCUS_EXPORT_FORMAT") or "parquet" - ).lower() - self.frequency = ( - frequency or os.getenv("FOCUS_EXPORT_FREQUENCY") or "hourly" - ).lower() - self.cron_offset_minute = ( - cron_offset_minute - if cron_offset_minute is not None - else int(os.getenv("FOCUS_EXPORT_CRON_OFFSET", "5")) - ) - self.interval_seconds = ( - interval_seconds - if interval_seconds is not None - else os.getenv("FOCUS_EXPORT_INTERVAL_SECONDS") - ) - self.prefix = prefix or os.getenv("FOCUS_EXPORT_PREFIX", "focus_exports") - - self._destination = self._init_destination( - destination_config=destination_config, - ) - self._serializer = self._init_serializer() - self._transformer = FocusTransformer() - - 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 _init_destination( - self, - *, - destination_config: Optional[dict[str, Any]], - ) -> FocusDestination: - """Factory for destination implementations.""" - resolved_config = self._resolve_destination_config(destination_config) - return FocusDestinationFactory.create( - provider=self.provider, - prefix=self.prefix, - config=resolved_config, - ) - - def _resolve_destination_config( - self, - destination_config: Optional[dict[str, Any]], - ) -> dict[str, Any]: - """Collect provider-specific configuration for destination creation.""" - raise NotImplementedError - - async def export_usage_data(self) -> None: - """Public hook to trigger export immediately.""" - raise NotImplementedError - - async def dry_run_export_usage_data(self) -> dict: - """Return transformed data without uploading.""" - raise NotImplementedError - - async def initialize_focus_export_job(self) -> None: - """Entry point for scheduler jobs to run export cycle with locking.""" - raise NotImplementedError - - @staticmethod - async def init_focus_export_background_job( - scheduler: AsyncIOScheduler, - ) -> None: - """Register the export cron/interval job with the provided scheduler.""" - raise NotImplementedError - - def _compute_time_window(self, now: datetime) -> FocusTimeWindow: - """Derive the time window to export based on configured frequency.""" - raise NotImplementedError - - def _serialize_and_upload( - self, - frame: pl.DataFrame, - window: FocusTimeWindow, - ) -> None: - """Helper stub for serializing and delegating to destination.""" - raise NotImplementedError - - 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}" diff --git a/litellm/integrations/focus/focus_logger.py b/litellm/integrations/focus/focus_logger.py new file mode 100644 index 00000000000..5eccf441bea --- /dev/null +++ b/litellm/integrations/focus/focus_logger.py @@ -0,0 +1,285 @@ +"""Focus export logger orchestrating DB pull/transform/upload.""" + +from __future__ import annotations + +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 ( + FocusDestination, + FocusDestinationFactory, + FocusTimeWindow, +) +from .serializers import FocusParquetSerializer, FocusSerializer +from .transformer import FocusTransformer + +if TYPE_CHECKING: + from apscheduler.schedulers.asyncio import AsyncIOScheduler +else: + AsyncIOScheduler = Any + +FOCUS_USAGE_DATA_JOB_NAME = "focus_export_usage_data" +DEFAULT_DRY_RUN_LIMIT = 500 + + +class FocusLogger(CustomLogger): + """Coordinates Focus export jobs across transformer/serializer/destination layers.""" + + def __init__( + self, + *, + provider: Optional[str] = None, + export_format: Optional[str] = None, + frequency: Optional[str] = None, + cron_offset_minute: Optional[int] = None, + interval_seconds: Optional[int] = None, + prefix: Optional[str] = None, + destination_config: Optional[dict[str, Any]] = None, + **kwargs: Any, + ) -> None: + super().__init__(**kwargs) + self.provider = (provider or os.getenv("FOCUS_PROVIDER") or "s3").lower() + self.export_format = ( + export_format or os.getenv("FOCUS_FORMAT") or "parquet" + ).lower() + self.frequency = ( + frequency or os.getenv("FOCUS_FREQUENCY") or "hourly" + ).lower() + self.cron_offset_minute = ( + cron_offset_minute + if cron_offset_minute is not None + else int(os.getenv("FOCUS_CRON_OFFSET", "5")) + ) + raw_interval = ( + interval_seconds + if interval_seconds is not None + else os.getenv("FOCUS_INTERVAL_SECONDS") + ) + self.interval_seconds = int(raw_interval) if raw_interval is not None else None + self.prefix = prefix or os.getenv("FOCUS_PREFIX", "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() + + 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() + + async def export_usage_data( + self, + *, + limit: Optional[int] = None, + start_time_utc: Optional[datetime] = None, + end_time_utc: Optional[datetime] = None, + ) -> None: + """Public hook to trigger export immediately.""" + if bool(start_time_utc) ^ bool(end_time_utc): + raise ValueError("start_time_utc and end_time_utc must be provided together") + + if start_time_utc and end_time_utc: + window = FocusTimeWindow( + start_time=start_time_utc, + end_time=end_time_utc, + frequency=self.frequency, + ) + else: + window = self._compute_time_window(datetime.now(timezone.utc)) + await self._export_window(window=window, limit=limit) + + async def dry_run_export_usage_data( + 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, + } + + async def initialize_focus_export_job(self) -> None: + """Entry point for scheduler jobs to run export cycle with locking.""" + try: + from litellm.proxy.proxy_server import proxy_logging_obj + except ImportError: + proxy_logging_obj = None + + 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=FOCUS_USAGE_DATA_JOB_NAME + ) + if not acquired: + verbose_logger.debug("Focus export: unable to acquire pod lock") + return + try: + await self._run_scheduled_export() + finally: + await pod_lock_manager.release_lock( + cronjob_id=FOCUS_USAGE_DATA_JOB_NAME + ) + else: + await self._run_scheduled_export() + + @staticmethod + async def init_focus_export_background_job( + scheduler: AsyncIOScheduler, + ) -> None: + """Register the export cron/interval job with the provided scheduler.""" + from litellm.integrations.custom_logger import CustomLogger + + focus_loggers: List[ + CustomLogger + ] = litellm.logging_callback_manager.get_custom_loggers_for_type( + callback_type=FocusLogger + ) + if not focus_loggers: + verbose_logger.debug("No Focus export logger registered; skipping scheduler") + return + + focus_logger = cast(FocusLogger, focus_loggers[0]) + trigger_kwargs = focus_logger._build_scheduler_trigger() + scheduler.add_job( + focus_logger.initialize_focus_export_job, + **trigger_kwargs, + ) + + def _build_scheduler_trigger(self) -> Dict[str, Any]: + """Return scheduler configuration for the selected frequency.""" + if self.frequency == "interval": + seconds = self.interval_seconds or 60 + return {"trigger": "interval", "seconds": seconds} + + if self.frequency == "hourly": + minute = max(0, min(59, self.cron_offset_minute)) + return {"trigger": "cron", "minute": minute, "second": 0} + + if self.frequency == "daily": + total_minutes = max(0, self.cron_offset_minute) + hour = min(23, total_minutes // 60) + minute = min(59, total_minutes % 60) + return {"trigger": "cron", "hour": hour, "minute": minute, "second": 0} + + raise ValueError(f"Unsupported frequency: {self.frequency}") + + async def _run_scheduled_export(self) -> None: + """Execute the scheduled export for the configured window.""" + window = self._compute_time_window(datetime.now(timezone.utc)) + await self._export_window(window=window, limit=None) + + 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) + + def _compute_time_window(self, now: datetime) -> FocusTimeWindow: + """Derive the time window to export based on configured frequency.""" + now_utc = now.astimezone(timezone.utc) + if self.frequency == "hourly": + end_time = now_utc.replace(minute=0, second=0, microsecond=0) + # start_time = end_time - timedelta(hours=1) + # Temporary override: export data since start of day instead of last hour + start_time = end_time.replace(hour=0) + elif self.frequency == "daily": + end_time = now_utc.replace(hour=0, minute=0, second=0, microsecond=0) + start_time = end_time - timedelta(days=1) + elif self.frequency == "interval": + interval = timedelta(seconds=self.interval_seconds or 60) + end_time = now_utc + start_time = end_time - interval + else: + raise ValueError(f"Unsupported frequency: {self.frequency}") + return FocusTimeWindow( + start_time=start_time, + end_time=end_time, + 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) diff --git a/litellm/integrations/focus/schema.py b/litellm/integrations/focus/schema.py index 766e29c0856..61ebbf2b9de 100644 --- a/litellm/integrations/focus/schema.py +++ b/litellm/integrations/focus/schema.py @@ -4,29 +4,52 @@ from __future__ import annotations import polars as pl - +# see: https://focus.finops.org/focus-specification/v1-2/ FOCUS_NORMALIZED_SCHEMA = pl.Schema( { - "usage_date": pl.Datetime(time_unit="us"), - "team_id": pl.String, - "team_alias": pl.String, - "user_id": pl.String, - "user_email": pl.String, - "api_key_alias": pl.String, - "model": pl.String, - "model_group": pl.String, - "custom_llm_provider": pl.String, - "prompt_tokens": pl.Int64, - "completion_tokens": pl.Int64, - "total_tokens": pl.Int64, - "spend": pl.Float64, - "cache_creation_input_tokens": pl.Int64, - "cache_read_input_tokens": pl.Int64, - "api_requests": pl.Int64, - "successful_requests": pl.Int64, - "failed_requests": pl.Int64, - "created_at": pl.Datetime(time_unit="us"), - "updated_at": pl.Datetime(time_unit="us"), + "BilledCost": pl.Float64, + "BillingAccountId": pl.String, + "BillingAccountName": pl.String, + "BillingCurrency": pl.String, + "BillingPeriodStart": pl.Datetime(time_unit="us"), + "BillingPeriodEnd": pl.Datetime(time_unit="us"), + "ChargeCategory": pl.String, + "ChargeClass": pl.String, + "ChargeDescription": pl.String, + "ChargeFrequency": pl.String, + "ChargePeriodStart": pl.Datetime(time_unit="us"), + "ChargePeriodEnd": pl.Datetime(time_unit="us"), + "CommitmentDiscountCategory": pl.String, + "CommitmentDiscountId": pl.String, + "CommitmentDiscountName": pl.String, + "CommitmentDiscountStatus": pl.String, + "CommitmentDiscountType": pl.String, + "ConsumedQuantity": pl.Float64, + "ConsumedUnit": pl.Float64, + "ContractedCost": pl.Float64, + "ContractedUnitPrice": pl.Float64, + "EffectiveCost": pl.Float64, + "InvoiceIssuerName": pl.String, + "ListCost": pl.Float64, + "ListUnitPrice": pl.Float64, + "PricingCategory": pl.String, + "PricingQuantity": pl.Float64, + "PricingUnit": pl.String, + "ProviderName": pl.String, + "PublisherName": pl.String, + "RegionId": pl.String, + "RegionName": pl.String, + "ResourceId": pl.String, + "ResourceName": pl.String, + "ResourceType": pl.String, + "ServiceCategory": pl.String, + "ServiceName": pl.String, + "SkuId": pl.String, + "SkuPriceId": pl.String, + "SubAccountId": pl.String, + "SubAccountName": pl.String, + "SubAccountType": pl.String, + "Tags": pl.Object, } ) diff --git a/litellm/integrations/focus/serializers/parquet.py b/litellm/integrations/focus/serializers/parquet.py index 3d42337e388..deac0e45390 100644 --- a/litellm/integrations/focus/serializers/parquet.py +++ b/litellm/integrations/focus/serializers/parquet.py @@ -2,15 +2,22 @@ from __future__ import annotations +import io + import polars as pl from .base import FocusSerializer class FocusParquetSerializer(FocusSerializer): - """Placeholder Parquet serializer implementation.""" + """Serialize normalized Focus frames to Parquet bytes.""" extension = "parquet" def serialize(self, frame: pl.DataFrame) -> bytes: - raise NotImplementedError + """Encode the provided frame as a parquet payload.""" + target = frame if not frame.is_empty() else pl.DataFrame(schema=frame.schema) + buffer = io.BytesIO() + target.write_parquet(buffer, compression="snappy") + print(target.head(5)) # debug + return buffer.getvalue() diff --git a/litellm/integrations/focus/transformer.py b/litellm/integrations/focus/transformer.py index 612082505fa..8957172da86 100644 --- a/litellm/integrations/focus/transformer.py +++ b/litellm/integrations/focus/transformer.py @@ -2,6 +2,8 @@ from __future__ import annotations +from datetime import timedelta + import polars as pl from .schema import FOCUS_NORMALIZED_SCHEMA @@ -14,4 +16,89 @@ class FocusTransformer: def transform(self, frame: pl.DataFrame) -> pl.DataFrame: """Return a normalized frame expected by downstream serializers.""" - raise NotImplementedError + if frame.is_empty(): + return pl.DataFrame(schema=self.schema) + + # derive period start/end from usage date + frame = frame.with_columns( + pl.col("date") + .cast(pl.Utf8) + .str.strptime(pl.Datetime(time_unit="us"), format="%Y-%m-%d", strict=False) + .alias("usage_date"), + ) + frame = frame.with_columns( + pl.col("usage_date").alias("ChargePeriodStart"), + (pl.col("usage_date") + timedelta(days=1)).alias("ChargePeriodEnd"), + ) + + def fmt(col): + return col.dt.strftime("%Y-%m-%dT%H:%M:%SZ") + + DEC = pl.Decimal(18, 6) + def dec(col): + return col.cast(DEC) + + none_str = pl.lit(None, dtype=pl.Utf8) + none_dec = pl.lit(None, dtype=pl.Decimal(18, 6)) + # zero_float = pl.lit(0.0, dtype=pl.Float64) + + return frame.select( + dec(pl.col("spend").fill_null(0.0)).alias("BilledCost"), + pl.col("api_key").cast(pl.String).alias("BillingAccountId"), + pl.col("api_key_alias").cast(pl.String).alias("BillingAccountName"), + pl.lit("API Key").alias("BillingAccountType"), + pl.lit("USD").alias("BillingCurrency"), + fmt(pl.col("ChargePeriodEnd")).alias("BillingPeriodEnd"), + fmt(pl.col("ChargePeriodStart")).alias("BillingPeriodStart"), + pl.lit("Usage").alias("ChargeCategory"), + none_str.alias("ChargeClass"), + pl.col("model").cast(pl.String).alias("ChargeDescription"), + pl.lit("Usage-Based").alias("ChargeFrequency"), + fmt(pl.col("ChargePeriodEnd")).alias("ChargePeriodEnd"), + fmt(pl.col("ChargePeriodStart")).alias("ChargePeriodStart"), + # pl.lit(None).alias("CommitmentDiscountCategory"), + # none_str.alias("CommitmentDiscountId"), + # none_str.alias("CommitmentDiscountName"), + # none_dec.alias("CommitmentDiscountQuantity"), + # none_str.alias("CommitmentDiscountUnit"), + # none_str.alias("CommitmentDiscountStatus"), + # none_str.alias("CommitmentDiscountType"), + dec(pl.lit(1.0)).alias("ConsumedQuantity"), + pl.lit("Requests").alias("ConsumedUnit"), + dec(pl.col("spend").fill_null(0.0)).alias("ContractedCost"), + none_str.alias("ContractedUnitPrice"), + dec(pl.col("spend").fill_null(0.0)).alias("EffectiveCost"), + pl.col("custom_llm_provider").cast(pl.String).alias("InvoiceIssuerName"), + # pl.lit("INVOICE-NOT-ISSUED").alias("InvoiceId"), + none_str.alias("InvoiceId"), + dec(pl.col("spend").fill_null(0.0)).alias("ListCost"), + none_dec.alias("ListUnitPrice"), + none_str.alias("AvailabilityZone"), + # none_str.alias("CapacityReservationId"), + # none_str.alias("CapacityReservationStatus"), + pl.lit("USD").alias("PricingCurrency"), + none_str.alias("PricingCategory"), + dec(pl.lit(1.0)).alias("PricingQuantity"), + none_dec.alias("PricingCurrencyContractedUnitPrice"), + dec(pl.col("spend").fill_null(0.0)).alias("PricingCurrencyEffectiveCost"), + none_dec.alias("PricingCurrencyListUnitPrice"), + pl.lit("Requests").alias("PricingUnit"), + pl.col("custom_llm_provider").cast(pl.String).alias("ProviderName"), + pl.col("custom_llm_provider").cast(pl.String).alias("PublisherName"), + none_str.alias("RegionId"), + none_str.alias("RegionName"), + pl.col("model").cast(pl.String).alias("ResourceId"), + pl.col("model").cast(pl.String).alias("ResourceName"), + pl.col("model").cast(pl.String).alias("ResourceType"), + pl.lit("AI and Machine Learning").alias("ServiceCategory"), + pl.lit("Generative AI").alias("ServiceSubcategory"), + pl.col("model_group").cast(pl.String).alias("ServiceName"), + # none_str.alias("SkuId"), + # none_str.alias("SkuPriceId"), + # none_str.alias("SkuMeter"), + # none_str.alias("SkuPriceDetails"), + pl.col("team_id").cast(pl.String).alias("SubAccountId"), + pl.col("team_alias").cast(pl.String).alias("SubAccountName"), + none_str.alias("SubAccountType"), + none_str.alias("Tags"), + ) diff --git a/litellm/litellm_core_utils/custom_logger_registry.py b/litellm/litellm_core_utils/custom_logger_registry.py index 47cbcb8aec9..a3c25ab65e9 100644 --- a/litellm/litellm_core_utils/custom_logger_registry.py +++ b/litellm/litellm_core_utils/custom_logger_registry.py @@ -18,6 +18,7 @@ from litellm.integrations.azure_storage.azure_storage import AzureBlobStorageLog from litellm.integrations.bitbucket import BitBucketPromptManager from litellm.integrations.braintrust_logging import BraintrustLogger from litellm.integrations.cloudzero.cloudzero import CloudZeroLogger +from litellm.integrations.focus.focus_logger import FocusLogger from litellm.integrations.datadog.datadog import DataDogLogger from litellm.integrations.datadog.datadog_llm_obs import DataDogLLMObsLogger from litellm.integrations.deepeval import DeepEvalLogger @@ -93,6 +94,7 @@ class CustomLoggerRegistry: "bitbucket": BitBucketPromptManager, "gitlab": GitLabPromptManager, "cloudzero": CloudZeroLogger, + "focus": FocusLogger, "posthog": PostHogLogger, } diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 5448fe7c771..b5847722a66 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -3756,6 +3756,15 @@ def _init_custom_logger_compatible_class( # noqa: PLR0915 cloudzero_logger = CloudZeroLogger() _in_memory_loggers.append(cloudzero_logger) return cloudzero_logger # type: ignore + elif logging_integration == "focus": + from litellm.integrations.focus.focus_logger import FocusLogger + + for callback in _in_memory_loggers: + if isinstance(callback, FocusLogger): + return callback # type: ignore + focus_logger = FocusLogger() + _in_memory_loggers.append(focus_logger) + return focus_logger # type: ignore elif logging_integration == "deepeval": for callback in _in_memory_loggers: if isinstance(callback, DeepEvalLogger): @@ -4076,6 +4085,12 @@ def get_custom_logger_compatible_class( # noqa: PLR0915 for callback in _in_memory_loggers: if isinstance(callback, CloudZeroLogger): return callback + elif logging_integration == "focus": + from litellm.integrations.focus.focus_logger import FocusLogger + + for callback in _in_memory_loggers: + if isinstance(callback, FocusLogger): + return callback elif logging_integration == "deepeval": for callback in _in_memory_loggers: if isinstance(callback, DeepEvalLogger): diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 06525e39133..6c5a2397204 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -4681,8 +4681,9 @@ class ProxyStartupEvent: """ Initialize the spend tracking and other background jobs 1. CloudZero Background Job - 2. Prometheus Background Job - 3. Key Rotation Background Job + 2. Focus Background Job + 3. Prometheus Background Job + 4. Key Rotation Background Job Args: scheduler: The scheduler to add the background jobs to @@ -4691,11 +4692,17 @@ class ProxyStartupEvent: # CloudZero Background Job ######################################################## from litellm.integrations.cloudzero.cloudzero import CloudZeroLogger + from litellm.integrations.focus.focus_logger import FocusLogger from litellm.proxy.spend_tracking.cloudzero_endpoints import is_cloudzero_setup if await is_cloudzero_setup(): await CloudZeroLogger.init_cloudzero_background_job(scheduler=scheduler) + ######################################################## + # Focus Background Job + ######################################################## + await FocusLogger.init_focus_export_background_job(scheduler=scheduler) + ######################################################## # Prometheus Background Job ######################################################## From 271ee0959b5c3d64e2d2cb2da92234749b0eefa7 Mon Sep 17 00:00:00 2001 From: Yuta Saito Date: Thu, 8 Jan 2026 16:53:09 +0900 Subject: [PATCH 4/9] test: focus --- litellm/integrations/focus/focus_logger.py | 4 +- litellm/integrations/focus/schema.py | 7 -- .../integrations/focus/serializers/parquet.py | 1 - litellm/integrations/focus/transformer.py | 15 --- .../integrations/focus/test_database.py | 74 +++++++++++++ .../integrations/focus/test_s3_destination.py | 100 ++++++++++++++++++ 6 files changed, 175 insertions(+), 26 deletions(-) create mode 100644 tests/test_litellm/integrations/focus/test_database.py create mode 100644 tests/test_litellm/integrations/focus/test_s3_destination.py diff --git a/litellm/integrations/focus/focus_logger.py b/litellm/integrations/focus/focus_logger.py index 5eccf441bea..28629f64d38 100644 --- a/litellm/integrations/focus/focus_logger.py +++ b/litellm/integrations/focus/focus_logger.py @@ -228,9 +228,7 @@ class FocusLogger(CustomLogger): now_utc = now.astimezone(timezone.utc) if self.frequency == "hourly": end_time = now_utc.replace(minute=0, second=0, microsecond=0) - # start_time = end_time - timedelta(hours=1) - # Temporary override: export data since start of day instead of last hour - start_time = end_time.replace(hour=0) + start_time = end_time - timedelta(hours=1) elif self.frequency == "daily": end_time = now_utc.replace(hour=0, minute=0, second=0, microsecond=0) start_time = end_time - timedelta(days=1) diff --git a/litellm/integrations/focus/schema.py b/litellm/integrations/focus/schema.py index 61ebbf2b9de..6d2e1dc83c1 100644 --- a/litellm/integrations/focus/schema.py +++ b/litellm/integrations/focus/schema.py @@ -19,11 +19,6 @@ FOCUS_NORMALIZED_SCHEMA = pl.Schema( "ChargeFrequency": pl.String, "ChargePeriodStart": pl.Datetime(time_unit="us"), "ChargePeriodEnd": pl.Datetime(time_unit="us"), - "CommitmentDiscountCategory": pl.String, - "CommitmentDiscountId": pl.String, - "CommitmentDiscountName": pl.String, - "CommitmentDiscountStatus": pl.String, - "CommitmentDiscountType": pl.String, "ConsumedQuantity": pl.Float64, "ConsumedUnit": pl.Float64, "ContractedCost": pl.Float64, @@ -44,8 +39,6 @@ FOCUS_NORMALIZED_SCHEMA = pl.Schema( "ResourceType": pl.String, "ServiceCategory": pl.String, "ServiceName": pl.String, - "SkuId": pl.String, - "SkuPriceId": pl.String, "SubAccountId": pl.String, "SubAccountName": pl.String, "SubAccountType": pl.String, diff --git a/litellm/integrations/focus/serializers/parquet.py b/litellm/integrations/focus/serializers/parquet.py index deac0e45390..6b3dde5903d 100644 --- a/litellm/integrations/focus/serializers/parquet.py +++ b/litellm/integrations/focus/serializers/parquet.py @@ -19,5 +19,4 @@ class FocusParquetSerializer(FocusSerializer): target = frame if not frame.is_empty() else pl.DataFrame(schema=frame.schema) buffer = io.BytesIO() target.write_parquet(buffer, compression="snappy") - print(target.head(5)) # debug return buffer.getvalue() diff --git a/litellm/integrations/focus/transformer.py b/litellm/integrations/focus/transformer.py index 8957172da86..a98ea21b1bb 100644 --- a/litellm/integrations/focus/transformer.py +++ b/litellm/integrations/focus/transformer.py @@ -40,7 +40,6 @@ class FocusTransformer: none_str = pl.lit(None, dtype=pl.Utf8) none_dec = pl.lit(None, dtype=pl.Decimal(18, 6)) - # zero_float = pl.lit(0.0, dtype=pl.Float64) return frame.select( dec(pl.col("spend").fill_null(0.0)).alias("BilledCost"), @@ -56,26 +55,16 @@ class FocusTransformer: pl.lit("Usage-Based").alias("ChargeFrequency"), fmt(pl.col("ChargePeriodEnd")).alias("ChargePeriodEnd"), fmt(pl.col("ChargePeriodStart")).alias("ChargePeriodStart"), - # pl.lit(None).alias("CommitmentDiscountCategory"), - # none_str.alias("CommitmentDiscountId"), - # none_str.alias("CommitmentDiscountName"), - # none_dec.alias("CommitmentDiscountQuantity"), - # none_str.alias("CommitmentDiscountUnit"), - # none_str.alias("CommitmentDiscountStatus"), - # none_str.alias("CommitmentDiscountType"), dec(pl.lit(1.0)).alias("ConsumedQuantity"), pl.lit("Requests").alias("ConsumedUnit"), dec(pl.col("spend").fill_null(0.0)).alias("ContractedCost"), none_str.alias("ContractedUnitPrice"), dec(pl.col("spend").fill_null(0.0)).alias("EffectiveCost"), pl.col("custom_llm_provider").cast(pl.String).alias("InvoiceIssuerName"), - # pl.lit("INVOICE-NOT-ISSUED").alias("InvoiceId"), none_str.alias("InvoiceId"), dec(pl.col("spend").fill_null(0.0)).alias("ListCost"), none_dec.alias("ListUnitPrice"), none_str.alias("AvailabilityZone"), - # none_str.alias("CapacityReservationId"), - # none_str.alias("CapacityReservationStatus"), pl.lit("USD").alias("PricingCurrency"), none_str.alias("PricingCategory"), dec(pl.lit(1.0)).alias("PricingQuantity"), @@ -93,10 +82,6 @@ class FocusTransformer: pl.lit("AI and Machine Learning").alias("ServiceCategory"), pl.lit("Generative AI").alias("ServiceSubcategory"), pl.col("model_group").cast(pl.String).alias("ServiceName"), - # none_str.alias("SkuId"), - # none_str.alias("SkuPriceId"), - # none_str.alias("SkuMeter"), - # none_str.alias("SkuPriceDetails"), pl.col("team_id").cast(pl.String).alias("SubAccountId"), pl.col("team_alias").cast(pl.String).alias("SubAccountName"), none_str.alias("SubAccountType"), diff --git a/tests/test_litellm/integrations/focus/test_database.py b/tests/test_litellm/integrations/focus/test_database.py new file mode 100644 index 00000000000..5ee98cc9dd0 --- /dev/null +++ b/tests/test_litellm/integrations/focus/test_database.py @@ -0,0 +1,74 @@ +"""Tests for FocusLiteLLMDatabase query construction.""" + +from datetime import datetime, timezone +from types import SimpleNamespace +from unittest.mock import AsyncMock + +import pytest + +from litellm.integrations.focus.database import FocusLiteLLMDatabase + + +def _setup_db(monkeypatch: pytest.MonkeyPatch, query_return): + """Create a database instance with a stubbed prisma client.""" + query_mock = AsyncMock(return_value=query_return) + mock_client = SimpleNamespace(db=SimpleNamespace(query_raw=query_mock)) + db = FocusLiteLLMDatabase() + monkeypatch.setattr(db, "_ensure_prisma_client", lambda: mock_client) + return db, query_mock + + +@pytest.mark.asyncio +async def test_should_parameterize_filters_and_limit(monkeypatch: pytest.MonkeyPatch): + start = datetime(2024, 1, 1, tzinfo=timezone.utc) + end = datetime(2024, 1, 2, tzinfo=timezone.utc) + db, query_mock = _setup_db(monkeypatch, []) + + await db.get_usage_data(limit=25, start_time_utc=start, end_time_utc=end) + + query_text, *params = query_mock.await_args.args + assert "dus.updated_at >= $1::timestamptz" in query_text + assert "dus.updated_at <= $2::timestamptz" in query_text + assert "LIMIT $3" in query_text + assert params == [start, end, 25] + + +@pytest.mark.asyncio +async def test_should_execute_without_filters(monkeypatch: pytest.MonkeyPatch): + row = { + "id": 1, + "user_id": "user", + "date": datetime(2024, 1, 1, tzinfo=timezone.utc), + } + db, query_mock = _setup_db(monkeypatch, [row]) + + result = await db.get_usage_data() + + query_text, *params = query_mock.await_args.args + assert "WHERE" not in query_text + assert "LIMIT $" not in query_text + assert params == [] + assert result.height == 1 + assert result["id"][0] == 1 + + +@pytest.mark.asyncio +async def test_should_accept_string_timestamps(monkeypatch: pytest.MonkeyPatch): + db, query_mock = _setup_db(monkeypatch, []) + + start = "2024-02-01T00:00:00+00:00" + end = "2024-02-02T00:00:00+00:00" + await db.get_usage_data(start_time_utc=start, end_time_utc=end) + + _, *params = query_mock.await_args.args + assert params == [start, end] + + +@pytest.mark.asyncio +async def test_should_reject_invalid_limit(monkeypatch: pytest.MonkeyPatch): + db, query_mock = _setup_db(monkeypatch, []) + + with pytest.raises(ValueError): + await db.get_usage_data(limit="invalid") + + assert query_mock.await_count == 0 diff --git a/tests/test_litellm/integrations/focus/test_s3_destination.py b/tests/test_litellm/integrations/focus/test_s3_destination.py new file mode 100644 index 00000000000..f915b2c56a3 --- /dev/null +++ b/tests/test_litellm/integrations/focus/test_s3_destination.py @@ -0,0 +1,100 @@ +"""Tests for FocusS3Destination behavior.""" + +from __future__ import annotations + +from datetime import datetime, timezone +from types import SimpleNamespace +from typing import Any, Dict + +import pytest + +import litellm.integrations.focus.destinations.s3_destination as s3_module +from litellm.integrations.focus.destinations.base import FocusTimeWindow +from litellm.integrations.focus.destinations.s3_destination import FocusS3Destination + + +def _window(freq: str = "hourly", hour: int = 5) -> FocusTimeWindow: + start = datetime(2024, 1, 2, hour, tzinfo=timezone.utc) + end = start.replace(hour=hour + 1) + return FocusTimeWindow(start_time=start, end_time=end, frequency=freq) + + +def test_should_require_bucket_name(): + with pytest.raises(ValueError): + FocusS3Destination(prefix="focus", config={}) + + +def test_should_build_hourly_object_key(): + dest = FocusS3Destination(prefix="exports/", config={"bucket_name": "bucket"}) + key = dest._build_object_key( + time_window=_window(freq="hourly", hour=3), filename="data.snappy" + ) + assert key == "exports/date=2024-01-02/hour=03/data.snappy" + + +def test_should_build_daily_key_without_hour_segment(): + dest = FocusS3Destination(prefix="", config={"bucket_name": "bucket"}) + key = dest._build_object_key( + time_window=_window(freq="daily", hour=0), filename="daily.parquet" + ) + assert key == "date=2024-01-02/daily.parquet" + + +@pytest.mark.asyncio +async def test_should_dispatch_upload_via_thread(monkeypatch: pytest.MonkeyPatch): + dest = FocusS3Destination(prefix="focus", config={"bucket_name": "bucket"}) + captured: Dict[str, Any] = {} + + async def fake_to_thread(func, *args, **kwargs): # type: ignore[override] + captured["func"] = func + captured["args"] = args + captured["kwargs"] = kwargs + + monkeypatch.setattr(s3_module.asyncio, "to_thread", fake_to_thread) + + window = _window(freq="hourly", hour=1) + await dest.deliver(content=b"payload", time_window=window, filename="file.bin") + + assert captured["func"] == dest._upload + assert captured["args"][0] == b"payload" + assert captured["args"][1].endswith("/file.bin") + + +def test_should_upload_with_configured_client(monkeypatch: pytest.MonkeyPatch): + config = { + "bucket_name": "bucket", + "region_name": "us-east-2", + "endpoint_url": "http://localhost:4566", + "aws_access_key_id": "key", + "aws_secret_access_key": "secret", + "aws_session_token": "token", + } + dest = FocusS3Destination(prefix="focus", config=config) + captured: Dict[str, Any] = {} + + def fake_client(service: str, **kwargs): + assert service == "s3" + captured["client_kwargs"] = kwargs + + def put_object(**put_kwargs): + captured["put_kwargs"] = put_kwargs + + return SimpleNamespace(put_object=put_object) + + monkeypatch.setattr(s3_module.boto3, "client", fake_client) + + dest._upload(content=b"payload", object_key="path/file.bin") + + assert captured["client_kwargs"] == { + "region_name": "us-east-2", + "endpoint_url": "http://localhost:4566", + "aws_access_key_id": "key", + "aws_secret_access_key": "secret", + "aws_session_token": "token", + } + assert captured["put_kwargs"] == { + "Bucket": "bucket", + "Key": "path/file.bin", + "Body": b"payload", + "ContentType": "application/octet-stream", + } From fb00b38fcdd81677895893251fc023ff7d91b8e3 Mon Sep 17 00:00:00 2001 From: Yuta Saito Date: Thu, 8 Jan 2026 16:56:51 +0900 Subject: [PATCH 5/9] chore: lint --- .../focus/destinations/factory.py | 8 +- litellm/integrations/focus/focus_logger.py | 18 +++-- litellm/integrations/focus/transformer.py | 1 + litellm/proxy/proxy_server.py | 76 +++++++++---------- 4 files changed, 53 insertions(+), 50 deletions(-) diff --git a/litellm/integrations/focus/destinations/factory.py b/litellm/integrations/focus/destinations/factory.py index 6f5dc5dc06a..cb7696a11de 100644 --- a/litellm/integrations/focus/destinations/factory.py +++ b/litellm/integrations/focus/destinations/factory.py @@ -26,7 +26,9 @@ class FocusDestinationFactory: ) if provider_lower == "s3": return FocusS3Destination(prefix=prefix, config=normalized_config) - raise NotImplementedError(f"Provider '{provider}' not supported for Focus export") + raise NotImplementedError( + f"Provider '{provider}' not supported for Focus export" + ) @staticmethod def _resolve_config( @@ -50,9 +52,7 @@ class FocusDestinationFactory: or os.getenv("FOCUS_S3_SESSION_TOKEN"), } if not resolved.get("bucket_name"): - raise ValueError( - "FOCUS_S3_BUCKET_NAME must be provided for S3 exports" - ) + raise ValueError("FOCUS_S3_BUCKET_NAME must be provided for S3 exports") 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" diff --git a/litellm/integrations/focus/focus_logger.py b/litellm/integrations/focus/focus_logger.py index 28629f64d38..5f47590f0b8 100644 --- a/litellm/integrations/focus/focus_logger.py +++ b/litellm/integrations/focus/focus_logger.py @@ -14,7 +14,6 @@ from litellm.integrations.custom_logger import CustomLogger from .database import FocusLiteLLMDatabase from .destinations import ( - FocusDestination, FocusDestinationFactory, FocusTimeWindow, ) @@ -50,9 +49,7 @@ class FocusLogger(CustomLogger): self.export_format = ( export_format or os.getenv("FOCUS_FORMAT") or "parquet" ).lower() - self.frequency = ( - frequency or os.getenv("FOCUS_FREQUENCY") or "hourly" - ).lower() + self.frequency = (frequency or os.getenv("FOCUS_FREQUENCY") or "hourly").lower() self.cron_offset_minute = ( cron_offset_minute if cron_offset_minute is not None @@ -90,7 +87,9 @@ class FocusLogger(CustomLogger): ) -> None: """Public hook to trigger export immediately.""" if bool(start_time_utc) ^ bool(end_time_utc): - raise ValueError("start_time_utc and end_time_utc must be provided together") + raise ValueError( + "start_time_utc and end_time_utc must be provided together" + ) if start_time_utc and end_time_utc: window = FocusTimeWindow( @@ -160,7 +159,6 @@ class FocusLogger(CustomLogger): scheduler: AsyncIOScheduler, ) -> None: """Register the export cron/interval job with the provided scheduler.""" - from litellm.integrations.custom_logger import CustomLogger focus_loggers: List[ CustomLogger @@ -168,7 +166,9 @@ class FocusLogger(CustomLogger): callback_type=FocusLogger ) if not focus_loggers: - verbose_logger.debug("No Focus export logger registered; skipping scheduler") + verbose_logger.debug( + "No Focus export logger registered; skipping scheduler" + ) return focus_logger = cast(FocusLogger, focus_loggers[0]) @@ -218,7 +218,9 @@ class FocusLogger(CustomLogger): normalized = self._transformer.transform(data) if normalized.is_empty(): - verbose_logger.debug("Focus export: normalized data empty for window %s", window) + verbose_logger.debug( + "Focus export: normalized data empty for window %s", window + ) return await self._serialize_and_upload(normalized, window) diff --git a/litellm/integrations/focus/transformer.py b/litellm/integrations/focus/transformer.py index a98ea21b1bb..cac12b7be14 100644 --- a/litellm/integrations/focus/transformer.py +++ b/litellm/integrations/focus/transformer.py @@ -35,6 +35,7 @@ class FocusTransformer: return col.dt.strftime("%Y-%m-%dT%H:%M:%SZ") DEC = pl.Decimal(18, 6) + def dec(col): return col.cast(DEC) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 6c5a2397204..7535a79dfac 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -533,9 +533,9 @@ except ImportError: server_root_path = os.getenv("SERVER_ROOT_PATH", "") _license_check = LicenseCheck() premium_user: bool = _license_check.is_premium() -premium_user_data: Optional["EnterpriseLicenseData"] = ( - _license_check.airgapped_license_data -) +premium_user_data: Optional[ + "EnterpriseLicenseData" +] = _license_check.airgapped_license_data global_max_parallel_request_retries_env: Optional[str] = os.getenv( "LITELLM_GLOBAL_MAX_PARALLEL_REQUEST_RETRIES" ) @@ -1083,9 +1083,7 @@ try: # In non-root Docker, we restructure in /var/lib/litellm/ui. try: _restructure_ui_html_files(ui_path) - verbose_proxy_logger.info( - f"Restructured UI directory: {ui_path}" - ) + verbose_proxy_logger.info(f"Restructured UI directory: {ui_path}") except PermissionError as e: verbose_proxy_logger.exception( f"Permission error while restructuring UI directory {ui_path}: {e}" @@ -1171,9 +1169,9 @@ master_key: Optional[str] = None config_agents: Optional[List[AgentConfig]] = None otel_logging = False prisma_client: Optional[PrismaClient] = None -shared_aiohttp_session: Optional["ClientSession"] = ( - None # Global shared session for connection reuse -) +shared_aiohttp_session: Optional[ + "ClientSession" +] = None # Global shared session for connection reuse user_api_key_cache = DualCache( default_in_memory_ttl=UserAPIKeyCacheTTLEnum.in_memory_cache_ttl.value ) @@ -1181,9 +1179,9 @@ model_max_budget_limiter = _PROXY_VirtualKeyModelMaxBudgetLimiter( dual_cache=user_api_key_cache ) litellm.logging_callback_manager.add_litellm_callback(model_max_budget_limiter) -redis_usage_cache: Optional[RedisCache] = ( - None # redis cache used for tracking spend, tpm/rpm limits -) +redis_usage_cache: Optional[ + RedisCache +] = None # redis cache used for tracking spend, tpm/rpm limits polling_via_cache_enabled: Union[Literal["all"], List[str], bool] = False polling_cache_ttl: int = 3600 # Default 1 hour TTL for polling cache user_custom_auth = None @@ -1522,9 +1520,9 @@ async def update_cache( # noqa: PLR0915 _id = "team_id:{}".format(team_id) try: # Fetch the existing cost for the given user - existing_spend_obj: Optional[LiteLLM_TeamTable] = ( - await user_api_key_cache.async_get_cache(key=_id) - ) + existing_spend_obj: Optional[ + LiteLLM_TeamTable + ] = await user_api_key_cache.async_get_cache(key=_id) if existing_spend_obj is None: # do nothing if team not in api key cache return @@ -1876,7 +1874,6 @@ class ProxyConfig: "environment_variables" in config_to_save and config_to_save["environment_variables"] ): - # decrypt the environment_variables - in case a caller function has already encrypted the environment_variables decrypted_env_vars = self._decrypt_and_set_db_env_variables( environment_variables=config_to_save["environment_variables"], @@ -2794,21 +2791,21 @@ class ProxyConfig: verbose_proxy_logger.debug(f"_alerting_callbacks: {general_settings}") if _alerting_callbacks is None: return - + # Ensure proxy_logging_obj.alerting is set for all alerting types _alerting_value = general_settings.get("alerting", None) - verbose_proxy_logger.debug(f"_load_alerting_settings: Calling update_values with alerting={_alerting_value}") + verbose_proxy_logger.debug( + f"_load_alerting_settings: Calling update_values with alerting={_alerting_value}" + ) proxy_logging_obj.update_values( alerting=_alerting_value, alerting_threshold=general_settings.get("alerting_threshold", 600), alert_types=general_settings.get("alert_types", None), - alert_to_webhook_url=general_settings.get( - "alert_to_webhook_url", None - ), + alert_to_webhook_url=general_settings.get("alert_to_webhook_url", None), alerting_args=general_settings.get("alerting_args", None), redis_cache=redis_usage_cache, ) - + for _alert in _alerting_callbacks: if _alert == "slack": # [OLD] v0 implementation - already handled by update_values above @@ -3279,7 +3276,7 @@ class ProxyConfig: proxy_logging_obj: ProxyLogging """ _general_settings = config_data.get("general_settings", {}) - + if _general_settings is not None and "alerting" in _general_settings: if ( general_settings is not None @@ -3294,7 +3291,8 @@ class ProxyConfig: _merged_alerting = list(_yaml_alerting.union(_db_alerting)) # Preserve order: YAML values first, then DB values _merged_alerting = list(general_settings["alerting"]) + [ - item for item in _general_settings["alerting"] + item + for item in _general_settings["alerting"] if item not in general_settings["alerting"] ] verbose_proxy_logger.debug( @@ -3605,7 +3603,6 @@ class ProxyConfig: await self._init_vector_stores_in_db(prisma_client=prisma_client) if self._should_load_db_object(object_type="vector_store_indexes"): - await self._init_vector_store_indexes_in_db(prisma_client=prisma_client) if self._should_load_db_object(object_type="mcp"): @@ -3804,10 +3801,10 @@ class ProxyConfig: ) try: - guardrails_in_db: List[Guardrail] = ( - await GuardrailRegistry.get_all_guardrails_from_db( - prisma_client=prisma_client - ) + guardrails_in_db: List[ + Guardrail + ] = await GuardrailRegistry.get_all_guardrails_from_db( + prisma_client=prisma_client ) verbose_proxy_logger.debug( "guardrails from the DB %s", str(guardrails_in_db) @@ -4134,9 +4131,9 @@ async def initialize( # noqa: PLR0915 user_api_base = api_base dynamic_config[user_model]["api_base"] = api_base if api_version: - os.environ["AZURE_API_VERSION"] = ( - api_version # set this for azure - litellm can read this from the env - ) + os.environ[ + "AZURE_API_VERSION" + ] = api_version # set this for azure - litellm can read this from the env if max_tokens: # model-specific param dynamic_config[user_model]["max_tokens"] = max_tokens if temperature: # model-specific param @@ -4654,10 +4651,14 @@ class ProxyStartupEvent: replace_existing=True, misfire_grace_time=APSCHEDULER_MISFIRE_GRACE_TIME, ) - verbose_proxy_logger.info("Responses cost check job scheduled successfully") + verbose_proxy_logger.info( + "Responses cost check job scheduled successfully" + ) except Exception as e: - verbose_proxy_logger.debug(f"Failed to setup responses cost checking: {e}") + verbose_proxy_logger.debug( + f"Failed to setup responses cost checking: {e}" + ) verbose_proxy_logger.debug( "Checking responses cost for LiteLLM Managed Files is an Enterprise Feature. Skipping..." ) @@ -5944,7 +5945,6 @@ async def realtime_websocket_endpoint( ), user_api_key_dict=Depends(user_api_key_auth_websocket), ): - await websocket.accept() # Only use explicit parameters, not all query params @@ -9532,9 +9532,9 @@ async def get_config_list( hasattr(sub_field_info, "description") and sub_field_info.description is not None ): - nested_fields[idx].field_description = ( - sub_field_info.description - ) + nested_fields[ + idx + ].field_description = sub_field_info.description idx += 1 _stored_in_db = None From 93cf2d4848867b5ec85549132946693d261a24f4 Mon Sep 17 00:00:00 2001 From: Yuta Saito Date: Thu, 8 Jan 2026 17:19:02 +0900 Subject: [PATCH 6/9] fix: mypy --- litellm/integrations/focus/database.py | 4 +- litellm/integrations/focus/focus_logger.py | 5 +- litellm/integrations/focus/schema.py | 77 +++++++++++----------- 3 files changed, 45 insertions(+), 41 deletions(-) diff --git a/litellm/integrations/focus/database.py b/litellm/integrations/focus/database.py index 402a1ec4b4e..298254670eb 100644 --- a/litellm/integrations/focus/database.py +++ b/litellm/integrations/focus/database.py @@ -31,8 +31,8 @@ class FocusLiteLLMDatabase: """Return usage data for the requested window.""" client = self._ensure_prisma_client() - where_clauses = [] - query_params = [] + where_clauses: list[str] = [] + query_params: list[Any] = [] placeholder_index = 1 if start_time_utc: where_clauses.append(f"dus.updated_at >= ${placeholder_index}::timestamptz") diff --git a/litellm/integrations/focus/focus_logger.py b/litellm/integrations/focus/focus_logger.py index 5f47590f0b8..1589f030fa5 100644 --- a/litellm/integrations/focus/focus_logger.py +++ b/litellm/integrations/focus/focus_logger.py @@ -61,7 +61,10 @@ class FocusLogger(CustomLogger): else os.getenv("FOCUS_INTERVAL_SECONDS") ) self.interval_seconds = int(raw_interval) if raw_interval is not None else None - self.prefix = prefix or os.getenv("FOCUS_PREFIX", "focus_exports") + env_prefix = os.getenv("FOCUS_PREFIX") + self.prefix: str = ( + prefix if prefix is not None else (env_prefix if env_prefix else "focus_exports") + ) self._destination = FocusDestinationFactory.create( provider=self.provider, diff --git a/litellm/integrations/focus/schema.py b/litellm/integrations/focus/schema.py index 6d2e1dc83c1..ac2f33dad0a 100644 --- a/litellm/integrations/focus/schema.py +++ b/litellm/integrations/focus/schema.py @@ -6,44 +6,45 @@ import polars as pl # see: https://focus.finops.org/focus-specification/v1-2/ FOCUS_NORMALIZED_SCHEMA = pl.Schema( - { - "BilledCost": pl.Float64, - "BillingAccountId": pl.String, - "BillingAccountName": pl.String, - "BillingCurrency": pl.String, - "BillingPeriodStart": pl.Datetime(time_unit="us"), - "BillingPeriodEnd": pl.Datetime(time_unit="us"), - "ChargeCategory": pl.String, - "ChargeClass": pl.String, - "ChargeDescription": pl.String, - "ChargeFrequency": pl.String, - "ChargePeriodStart": pl.Datetime(time_unit="us"), - "ChargePeriodEnd": pl.Datetime(time_unit="us"), - "ConsumedQuantity": pl.Float64, - "ConsumedUnit": pl.Float64, - "ContractedCost": pl.Float64, - "ContractedUnitPrice": pl.Float64, - "EffectiveCost": pl.Float64, - "InvoiceIssuerName": pl.String, - "ListCost": pl.Float64, - "ListUnitPrice": pl.Float64, - "PricingCategory": pl.String, - "PricingQuantity": pl.Float64, - "PricingUnit": pl.String, - "ProviderName": pl.String, - "PublisherName": pl.String, - "RegionId": pl.String, - "RegionName": pl.String, - "ResourceId": pl.String, - "ResourceName": pl.String, - "ResourceType": pl.String, - "ServiceCategory": pl.String, - "ServiceName": pl.String, - "SubAccountId": pl.String, - "SubAccountName": pl.String, - "SubAccountType": pl.String, - "Tags": pl.Object, - } + [ + ("BilledCost", pl.Decimal(18, 6)), + ("BillingAccountId", pl.String), + ("BillingAccountName", pl.String), + ("BillingCurrency", pl.String), + ("BillingPeriodStart", pl.Datetime(time_unit="us")), + ("BillingPeriodEnd", pl.Datetime(time_unit="us")), + ("ChargeCategory", pl.String), + ("ChargeClass", pl.String), + ("ChargeDescription", pl.String), + ("ChargeFrequency", pl.String), + ("ChargePeriodStart", pl.Datetime(time_unit="us")), + ("ChargePeriodEnd", pl.Datetime(time_unit="us")), + ("ConsumedQuantity", pl.Decimal(18, 6)), + ("ConsumedUnit", pl.String), + ("ContractedCost", pl.Decimal(18, 6)), + ("ContractedUnitPrice", pl.Decimal(18, 6)), + ("EffectiveCost", pl.Decimal(18, 6)), + ("InvoiceIssuerName", pl.String), + ("ListCost", pl.Decimal(18, 6)), + ("ListUnitPrice", pl.Decimal(18, 6)), + ("PricingCategory", pl.String), + ("PricingQuantity", pl.Decimal(18, 6)), + ("PricingUnit", pl.String), + ("ProviderName", pl.String), + ("PublisherName", pl.String), + ("RegionId", pl.String), + ("RegionName", pl.String), + ("ResourceId", pl.String), + ("ResourceName", pl.String), + ("ResourceType", pl.String), + ("ServiceCategory", pl.String), + ("ServiceSubcategory", pl.String), + ("ServiceName", pl.String), + ("SubAccountId", pl.String), + ("SubAccountName", pl.String), + ("SubAccountType", pl.String), + ("Tags", pl.Object), + ] ) __all__ = ["FOCUS_NORMALIZED_SCHEMA"] From 790c80070f8e0ebec8ce156ac64a1cecf8d68582 Mon Sep 17 00:00:00 2001 From: Yuta Saito Date: Thu, 8 Jan 2026 17:24:51 +0900 Subject: [PATCH 7/9] docs: add FOCUS env --- docs/my-website/docs/proxy/config_settings.md | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/docs/my-website/docs/proxy/config_settings.md b/docs/my-website/docs/proxy/config_settings.md index 5772fbaa487..12995e4a5f8 100644 --- a/docs/my-website/docs/proxy/config_settings.md +++ b/docs/my-website/docs/proxy/config_settings.md @@ -577,6 +577,18 @@ router_settings: | FIREWORKS_AI_56_B_MOE | Size parameter for Fireworks AI 56B MOE model. Default is 56 | FIREWORKS_AI_80_B | Size parameter for Fireworks AI 80B model. Default is 80 | FIREWORKS_AI_176_B_MOE | Size parameter for Fireworks AI 176B MOE model. Default is 176 +| FOCUS_PROVIDER | Destination provider for Focus exports (e.g., `s3`). Defaults to `s3`. +| FOCUS_FORMAT | Output format for Focus exports. Defaults to `parquet`. +| FOCUS_FREQUENCY | Frequency for scheduled Focus exports (`hourly`, `daily`, or `interval`). Defaults to `hourly`. +| FOCUS_CRON_OFFSET | Minute offset used when scheduling hourly/daily Focus exports. Defaults to `5` minutes. +| FOCUS_INTERVAL_SECONDS | Interval (in seconds) for Focus exports when `frequency` is `interval`. +| FOCUS_PREFIX | Object key prefix (or folder) used when uploading Focus export files. Defaults to `focus_exports`. +| FOCUS_S3_BUCKET_NAME | S3 bucket to upload Focus export files when using the S3 destination. +| FOCUS_S3_REGION_NAME | AWS region for the Focus export S3 bucket. +| FOCUS_S3_ENDPOINT_URL | Custom endpoint for the Focus export S3 client (optional; useful for S3-compatible storage). +| FOCUS_S3_ACCESS_KEY | AWS access key ID used by the Focus export S3 client. +| FOCUS_S3_SECRET_KEY | AWS secret access key used by the Focus export S3 client. +| FOCUS_S3_SESSION_TOKEN | AWS session token used by the Focus export S3 client (optional). | FUNCTION_DEFINITION_TOKEN_COUNT | Token count for function definitions. Default is 9 | GALILEO_BASE_URL | Base URL for Galileo platform | GALILEO_PASSWORD | Password for Galileo authentication From cc3b068f2240f4996984a3c28b4af2cd9d9643f9 Mon Sep 17 00:00:00 2001 From: Yuta Saito Date: Fri, 9 Jan 2026 05:35:15 +0900 Subject: [PATCH 8/9] fix: mypy error --- litellm/integrations/focus/focus_logger.py | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/litellm/integrations/focus/focus_logger.py b/litellm/integrations/focus/focus_logger.py index 1589f030fa5..505337ff1c8 100644 --- a/litellm/integrations/focus/focus_logger.py +++ b/litellm/integrations/focus/focus_logger.py @@ -130,10 +130,7 @@ class FocusLogger(CustomLogger): async def initialize_focus_export_job(self) -> None: """Entry point for scheduler jobs to run export cycle with locking.""" - try: - from litellm.proxy.proxy_server import proxy_logging_obj - except ImportError: - proxy_logging_obj = None + from litellm.proxy.proxy_server import proxy_logging_obj pod_lock_manager = None if proxy_logging_obj is not None: From 844577f2d46c5e64c7075630d0ce02981dd2e717 Mon Sep 17 00:00:00 2001 From: Yuta Saito Date: Fri, 9 Jan 2026 06:39:24 +0900 Subject: [PATCH 9/9] refactor: lazy-load focus export engine and isolate polars deps --- litellm/integrations/focus/export_engine.py | 124 ++++++++++++++++++++ litellm/integrations/focus/focus_logger.py | 116 ++++-------------- 2 files changed, 145 insertions(+), 95 deletions(-) create mode 100644 litellm/integrations/focus/export_engine.py diff --git a/litellm/integrations/focus/export_engine.py b/litellm/integrations/focus/export_engine.py new file mode 100644 index 00000000000..22ebce2a168 --- /dev/null +++ b/litellm/integrations/focus/export_engine.py @@ -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) diff --git a/litellm/integrations/focus/focus_logger.py b/litellm/integrations/focus/focus_logger.py index 505337ff1c8..ade1cf861b1 100644 --- a/litellm/integrations/focus/focus_logger.py +++ b/litellm/integrations/focus/focus_logger.py @@ -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"]