scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-09 03:18:44 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 108

Merge pull request #18802 from BerriAI/litellm_feat_focus_backend

Browse source 
[feat] Focus export support
This commit is contained in:
YutaSaito 2026-01-09 07:21:04 +09:00 • committed by GitHub
commit 661f03058c
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
20 changed files with 1060 additions and 40 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

12
docs/my-website/docs/proxy/config_settings.md
View file

@ -581,6 +581,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

1
litellm/__init__.py
View file

@ -134,6 +134,7 @@ _custom_logger_compatible_callbacks_literal = Literal[
"bitbucket",
"gitlab",
"cloudzero",
"focus",
"posthog",
"levo",
]

0
litellm/integrations/focus/__init__.py Normal file
View file

113
litellm/integrations/focus/database.py Normal file
View file

@ -0,0 +1,113 @@
"""Database access helpers for Focus export."""
from __future__ import annotations
from datetime import datetime
from typing import Any, Dict, Optional
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,
*,
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."""
client = self._ensure_prisma_client()
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")
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

12
litellm/integrations/focus/destinations/__init__.py Normal file
View file

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

30
litellm/integrations/focus/destinations/base.py Normal file
View file

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

59
litellm/integrations/focus/destinations/factory.py Normal file
View file

@ -0,0 +1,59 @@
"""Factory helpers for Focus export destinations."""
from __future__ import annotations
import os
from typing import Any, Dict, Optional
from .base import FocusDestination
from .s3_destination import FocusS3Destination
class FocusDestinationFactory:
"""Builds destination instances based on provider/config settings."""
@staticmethod
def create(
*,
provider: str,
prefix: str,
config: Optional[Dict[str, Any]] = None,
) -> FocusDestination:
"""Return a destination implementation for the requested provider."""
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"
)

74
litellm/integrations/focus/destinations/s3_destination.py Normal file
View file

@ -0,0 +1,74 @@
"""S3 destination implementation for Focus export."""
from __future__ import annotations
import asyncio
from datetime import timezone
from typing import Any, Optional
import boto3
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:
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
async def deliver(
self,
*,
content: bytes,
time_window: FocusTimeWindow,
filename: str,
) -> None:
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:
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",
)

124
litellm/integrations/focus/export_engine.py Normal file
View 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)

211
litellm/integrations/focus/focus_logger.py Normal file
View file

@ -0,0 +1,211 @@
"""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 litellm
from litellm._logging import verbose_logger
from litellm.integrations.custom_logger import CustomLogger
from .destinations import FocusTimeWindow
if TYPE_CHECKING:
from apscheduler.schedulers.asyncio import AsyncIOScheduler
from .export_engine import FocusExportEngine
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
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_config = destination_config
self._engine: Optional["FocusExportEngine"] = None
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,
*,
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."""
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."""
from litellm.proxy.proxy_server import proxy_logging_obj
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."""
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:
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."""
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)
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,
)
__all__ = ["FocusLogger"]

50
litellm/integrations/focus/schema.py Normal file
View file

@ -0,0 +1,50 @@
"""Schema definitions for Focus export data."""
from __future__ import annotations
import polars as pl
# see: https://focus.finops.org/focus-specification/v1-2/
FOCUS_NORMALIZED_SCHEMA = pl.Schema(
[
("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"]

6
litellm/integrations/focus/serializers/__init__.py Normal file
View file

@ -0,0 +1,6 @@
"""Serializer package exports for Focus integration."""
from .base import FocusSerializer
from .parquet import FocusParquetSerializer
__all__ = ["FocusSerializer", "FocusParquetSerializer"]

18
litellm/integrations/focus/serializers/base.py Normal file
View file

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

22
litellm/integrations/focus/serializers/parquet.py Normal file
View file

@ -0,0 +1,22 @@
"""Parquet serializer for Focus export."""
from __future__ import annotations
import io
import polars as pl
from .base import FocusSerializer
class FocusParquetSerializer(FocusSerializer):
"""Serialize normalized Focus frames to Parquet bytes."""
extension = "parquet"
def serialize(self, frame: pl.DataFrame) -> bytes:
"""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")
return buffer.getvalue()

90
litellm/integrations/focus/transformer.py Normal file
View file

@ -0,0 +1,90 @@
"""Focus export data transformer."""
from __future__ import annotations
from datetime import timedelta
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."""
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))
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"),
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"),
none_str.alias("InvoiceId"),
dec(pl.col("spend").fill_null(0.0)).alias("ListCost"),
none_dec.alias("ListUnitPrice"),
none_str.alias("AvailabilityZone"),
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"),
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"),
)

2
litellm/litellm_core_utils/custom_logger_registry.py
View file

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

15
litellm/litellm_core_utils/litellm_logging.py
View file

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

87
litellm/proxy/proxy_server.py
View file

@ -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
@ -3357,7 +3354,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
@ -3372,7 +3369,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(
@ -3686,7 +3684,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"):
@ -3885,10 +3882,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)
@ -4215,9 +4212,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
@ -4735,10 +4732,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..."
)
@ -4762,8 +4763,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
@ -4772,11 +4774,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
########################################################
@ -6018,7 +6026,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
@ -9606,9 +9613,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

74
tests/test_litellm/integrations/focus/test_database.py Normal file
View file

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

100
tests/test_litellm/integrations/focus/test_s3_destination.py Normal file
View file

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