chore: initial public snapshot for github upload
This commit is contained in:
Your Name
@@ -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
|
||||
@@ -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",
|
||||
]
|
||||
@@ -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."""
|
||||
...
|
||||
@@ -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"
|
||||
)
|
||||
@@ -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",
|
||||
)
|
||||
@@ -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)
|
||||
@@ -0,0 +1,214 @@
|
||||
"""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"]
|
||||
@@ -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"]
|
||||
@@ -0,0 +1,6 @@
|
||||
"""Serializer package exports for Focus integration."""
|
||||
|
||||
from .base import FocusSerializer
|
||||
from .parquet import FocusParquetSerializer
|
||||
|
||||
__all__ = ["FocusSerializer", "FocusParquetSerializer"]
|
||||
@@ -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
|
||||
@@ -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()
|
||||
@@ -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"),
|
||||
)
|
||||
Reference in New Issue
Block a user