diff --git a/litellm/__init__.py b/litellm/__init__.py index a7ea9ce4583..290d15bd07a 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -170,6 +170,7 @@ _custom_logger_compatible_callbacks_literal = Literal[ "focus", "mavvrik", "vantage", + "ternary", "posthog", "levo", "compression_interception", diff --git a/litellm/integrations/focus/destinations/__init__.py b/litellm/integrations/focus/destinations/__init__.py index 932184cf485..05baf5d06b4 100644 --- a/litellm/integrations/focus/destinations/__init__.py +++ b/litellm/integrations/focus/destinations/__init__.py @@ -5,6 +5,7 @@ from .factory import FocusDestinationFactory from .gcs_destination import FocusGCSDestination from .mavvrik_destination import FocusMavvrikDestination from .s3_destination import FocusS3Destination +from .ternary_destination import FocusTernaryDestination from .vantage_destination import FocusVantageDestination __all__ = [ @@ -13,6 +14,7 @@ __all__ = [ "FocusGCSDestination", "FocusMavvrikDestination", "FocusS3Destination", + "FocusTernaryDestination", "FocusTimeWindow", "FocusVantageDestination", ] diff --git a/litellm/integrations/focus/destinations/ternary_destination.py b/litellm/integrations/focus/destinations/ternary_destination.py new file mode 100644 index 00000000000..b2a8727ea61 --- /dev/null +++ b/litellm/integrations/focus/destinations/ternary_destination.py @@ -0,0 +1,152 @@ +"""Ternary API destination for FOCUS export: uploads each day's FOCUS CSV via a Ternary-signed GCS URL.""" + +from __future__ import annotations + +from collections.abc import Mapping +from typing import TYPE_CHECKING, Final, Protocol +from urllib.parse import quote, urlparse + +from typing_extensions import NotRequired, ReadOnly, TypedDict + +from litellm._logging import verbose_logger +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler + +from .base import FocusDestination, FocusTimeWindow + +if TYPE_CHECKING: + import httpx + +TERNARY_UPLOAD_TIMEOUT_SECONDS: Final = 120.0 + +_LOOPBACK_HOSTS: Final = frozenset({"localhost", "127.0.0.1", "::1"}) + + +class TernaryUploadUrlBody(TypedDict): + url: ReadOnly[NotRequired[str]] + method: ReadOnly[NotRequired[str]] + headers: ReadOnly[NotRequired[Mapping[str, str]]] + + +class _UploadUrlResponse(Protocol): + def json(self) -> TernaryUploadUrlBody: ... + + +def _upload_url_body(response: _UploadUrlResponse) -> TernaryUploadUrlBody: + return response.json() + + +def _require_secure_base_url(base_url: str) -> None: + parsed: Final = urlparse(base_url) + if parsed.scheme == "https": + return + if parsed.scheme == "http" and (parsed.hostname or "").lower() in _LOOPBACK_HOSTS: + return + raise ValueError(f"base_url must be an HTTPS URL (got {base_url!r}); http is allowed only for loopback") + + +def _require_gcs_url(url: str) -> None: + parsed: Final = urlparse(url) + hostname: Final = (parsed.hostname or "").lower() + is_gcs: Final = hostname == "storage.googleapis.com" or hostname.endswith(".storage.googleapis.com") + if parsed.scheme != "https" or not is_gcs: + raise ValueError(f"Ternary destination: upload URL must be an HTTPS GCS URL (got host {hostname!r})") + + +class _NoRedirectHTTPHandler(AsyncHTTPHandler): + """AsyncHTTPHandler that disables redirects on every client it creates (SSRF guard).""" + + def create_client( + self, + *args: object, + **kwargs: object, # kwargs-ok: passthrough to AsyncHTTPHandler.create_client + ) -> httpx.AsyncClient: + client: Final = super().create_client(*args, **kwargs) # pyright: ignore[reportArgumentType] # passthrough + client.follow_redirects = False + return client + + +class FocusTernaryDestination(FocusDestination): + """Upload one UTC day of FOCUS CSV per delivery to Ternary through a signed GCS URL.""" + + def __init__( + self, + *, + prefix: str, + config: dict[str, str] | None = None, # mutable-ok: FocusDestination(config) factory contract + ) -> None: + resolved_config: Final = config or {} + api_key: Final = resolved_config.get("api_key") + connection_id: Final = resolved_config.get("connection_id") + base_url: Final = resolved_config.get("base_url") + if not api_key: + raise ValueError( + "api_key must be provided for Ternary destination " + "(set TERNARY_API_KEY env var or pass in destination_config)" + ) + if not connection_id: + raise ValueError( + "connection_id must be provided for Ternary destination " + "(set TERNARY_CONNECTION_ID env var or pass in destination_config)" + ) + if "/" in connection_id or ".." in connection_id or any(c.isspace() for c in connection_id): + raise ValueError(f"connection_id must not contain '/', '..', or whitespace (got {connection_id!r})") + if not base_url: + raise ValueError( + "base_url must be provided for Ternary destination " + "(set TERNARY_BASE_URL env var or pass in destination_config)" + ) + _require_secure_base_url(str(base_url)) + self.api_key = api_key + self.connection_id = connection_id + self.base_url = str(base_url).rstrip("/") + self.prefix = prefix + + async def deliver( + self, + *, + content: bytes, + time_window: FocusTimeWindow, + filename: str, + ) -> None: + """Upload one UTC day: request a signed URL for the window's start day, then PUT the CSV to it.""" + if not content: + verbose_logger.debug("Ternary destination: empty content, skipping upload") + return + + day: Final = time_window.start_time.strftime("%Y-%m-%d") + handler: Final = _NoRedirectHTTPHandler(timeout=TERNARY_UPLOAD_TIMEOUT_SECONDS) + try: + body: Final = await self._request_upload_url(handler.client, day) + url: Final = body.get("url") + headers: Final = body.get("headers") + if not url or headers is None: + raise RuntimeError(f"Ternary destination: upload-url response for {day} is missing 'url' or 'headers'") + if body.get("method") != "PUT": + raise RuntimeError(f"Ternary destination: unsupported upload method {body.get('method')!r} for {day}") + _require_gcs_url(url) + uploaded: Final = await handler.client.put( + url, content=content, headers=headers, timeout=TERNARY_UPLOAD_TIMEOUT_SECONDS + ) + if uploaded.status_code >= 400: + raise RuntimeError( + f"Ternary destination: upload for {day} failed ({uploaded.status_code}): {uploaded.text[:200]}" + ) + finally: + await handler.close() + + verbose_logger.debug("Ternary destination: uploaded %d bytes for %s (%s)", len(content), day, filename) + + async def _request_upload_url(self, client: httpx.AsyncClient, day: str) -> TernaryUploadUrlBody: + endpoint: Final = f"{self.base_url}/external-cost-sources/v1/{quote(self.connection_id, safe='')}/upload-url" + response: Final = await client.post( + endpoint, + headers={"Authorization": f"Bearer {self.api_key}"}, + json={"day": day}, + timeout=TERNARY_UPLOAD_TIMEOUT_SECONDS, + ) + if response.status_code >= 400: + raise RuntimeError( + f"Ternary destination: upload-url request for {day} failed ({response.status_code}): " + f"{response.text[:200]}" + ) + return _upload_url_body(response) diff --git a/litellm/integrations/ternary/__init__.py b/litellm/integrations/ternary/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/litellm/integrations/ternary/export_engine.py b/litellm/integrations/ternary/export_engine.py new file mode 100644 index 00000000000..57f7a04b31e --- /dev/null +++ b/litellm/integrations/ternary/export_engine.py @@ -0,0 +1,247 @@ +"""Ternary FOCUS export engine (polars-backed): one UTC day per upload, with token counts merged into Tags.""" + +from __future__ import annotations + +import json +from collections.abc import Callable, Mapping, Sequence +from datetime import date, datetime, timedelta, timezone +from functools import partial +from typing import Final, Literal, Protocol + +from pydantic import TypeAdapter, ValidationError + +from litellm._logging import verbose_logger +from litellm.integrations.focus.database import FocusLiteLLMDatabase +from litellm.integrations.focus.destinations.base import FocusTimeWindow +from litellm.integrations.focus.destinations.ternary_destination import FocusTernaryDestination +from litellm.integrations.focus.export_engine import FocusExportEngine +from litellm.integrations.focus.transformer import FocusTransformer + +_TOKEN_TAG_KEYS: Final = ( + "prompt_tokens", + "completion_tokens", + "cache_read_input_tokens", + "cache_creation_input_tokens", +) +_SESSION_TIMEZONE_MARGIN: Final = timedelta(days=1) +_TAGS: Final = TypeAdapter(Mapping[str, object]) + + +class _Series(Protocol): + def to_list(self) -> Sequence[object]: ... + + def alias(self, name: str) -> _Series: ... + + +class _Frame(Protocol): + """The polars DataFrame surface this module uses.""" + + @property + def columns(self) -> Sequence[str]: ... + + @property + def height(self) -> int: ... + + def is_empty(self) -> bool: ... + + def get_column(self, name: str) -> _Series: ... + + def select(self, columns: Sequence[str], /) -> _Frame: ... + + def hstack(self, columns: _Frame, /) -> _Frame: ... + + def filter(self, mask: Sequence[bool], /) -> _Frame: ... + + def map_rows(self, function: Callable[[tuple[object, ...]], str], /) -> _Frame: ... + + def to_series(self) -> _Series: ... + + def with_columns(self, column: _Series, /) -> _Frame: ... + + def partition_by( + self, by: str, /, *, maintain_order: bool, as_dict: Literal[True] + ) -> Mapping[tuple[object, ...], _Frame]: ... + + +class _UsageSource(Protocol): + async def get_usage_data(self, *, limit: int | None, start_time_utc: datetime | None = None) -> _Frame: ... + + +class _Transformer(Protocol): + def transform(self, frame: _Frame) -> _Frame: ... + + +class _Serializer(Protocol): + def serialize(self, frame: _Frame) -> bytes: ... + + +def _existing_tags(raw: object) -> Mapping[str, object]: + if not isinstance(raw, str) or not raw: + return {} + try: + return _TAGS.validate_json(raw) + except ValidationError: + return {} + + +def _tags_with_tokens(token_cols: Sequence[str], row: tuple[object, ...]) -> str: + """Row is (Tags, *token values in `token_cols` order); a missing or non-object Tags starts from empty.""" + tokens: Final = {key: str(value) for key, value in zip(token_cols, row[1:]) if value is not None} + return json.dumps({**_existing_tags(row[0]), **tokens}) + + +def _merge_token_tags(normalized: _Frame, source: _Frame) -> _Frame: + """Zip transformed rows with the row-aligned source and add token keys to each row's Tags; degrades to unchanged on any mismatch.""" + if "Tags" not in normalized.columns: + verbose_logger.warning( + "Ternary export: transformed data has no `Tags` column; token counts not merged into %d row(s)", + normalized.height, + ) + return normalized + if normalized.height != source.height: + verbose_logger.warning( + "Ternary export: row-count mismatch (normalized=%d, source=%d); token counts not merged", + normalized.height, + source.height, + ) + return normalized + token_cols: Final = tuple(c for c in _TOKEN_TAG_KEYS if c in source.columns) + if normalized.is_empty() or not token_cols: + return normalized + merged: Final = ( + normalized.select(("Tags",)).hstack(source.select(token_cols)).map_rows(partial(_tags_with_tokens, token_cols)) + ) + return normalized.with_columns(merged.to_series().alias("Tags")) + + +def _canonical_day(raw: object) -> str | None: + """Return the zero-padded YYYY-MM-DD for a `date` value, or None when it is missing or unparseable.""" + if isinstance(raw, date): + return raw.isoformat() + if not isinstance(raw, str): + return None + try: + parsed: Final = datetime.strptime(raw, "%Y-%m-%d").replace(tzinfo=timezone.utc) + except ValueError: + return None + return parsed.date().isoformat() + + +def _keep_window_days(data: _Frame, window: FocusTimeWindow) -> _Frame: + """Keep source rows whose UTC `date` falls on a day the window touches; each kept day is then uploaded whole.""" + if "date" not in data.columns: + return data + first_day: Final = window.start_time.astimezone(timezone.utc).date().isoformat() + last_day: Final = (window.end_time.astimezone(timezone.utc) - timedelta(microseconds=1)).date().isoformat() + kept: Final = data.filter( + [ + day is not None and first_day <= day <= last_day + for day in map(_canonical_day, data.get_column("date").to_list()) + ] + ) + if data.height > 0 and kept.height == 0: + verbose_logger.warning( + "Ternary export: day window %s..%s dropped all %d rows (unparseable or out-of-window `date`?)", + first_day, + last_day, + data.height, + ) + return kept + + +def _partition_by_day(data: _Frame) -> Sequence[tuple[str, _Frame]]: + """Split source rows by their UTC `date` (one upload per day); rows with no parseable date are dropped.""" + if "date" not in data.columns: + verbose_logger.warning("Ternary export: source data has no `date` column; %d row(s) not exported", data.height) + return () + parts: Final = tuple( + (_canonical_day(key[0]), part) + for key, part in data.partition_by("date", maintain_order=True, as_dict=True).items() + ) + undated: Final = sum(part.height for day, part in parts if day is None) + if undated: + verbose_logger.warning("Ternary export: dropped %d row(s) with a missing or unparseable `date`", undated) + return tuple(sorted(((day, part) for day, part in parts if day is not None), key=lambda entry: entry[0])) + + +def _day_window(day: str, frequency: str) -> FocusTimeWindow: + """The single UTC day a per-day upload covers.""" + start: Final = datetime.strptime(day, "%Y-%m-%d").replace(tzinfo=timezone.utc) + return FocusTimeWindow(start_time=start, end_time=start + timedelta(days=1), frequency=frequency) + + +async def _read_usage(database: _UsageSource, *, limit: int | None, window: FocusTimeWindow | None) -> _Frame: + """Read every row, or the rows dated on the days `window` touches. + + `updated_at` is naive UTC filtered as `timestamptz` in the database session timezone, so the read starts one + day before the window with no upper bound, and each row's own UTC `date` decides what is kept. + """ + if window is None: + return await database.get_usage_data(limit=limit) + data: Final = await database.get_usage_data( + limit=limit, start_time_utc=window.start_time - _SESSION_TIMEZONE_MARGIN + ) + return _keep_window_days(data, window) + + +def _day_payload(data: _Frame, *, transformer: _Transformer, serializer: _Serializer) -> bytes: + transformed: Final = transformer.transform(data) + if transformed.is_empty(): + return b"" + return serializer.serialize(_merge_token_tags(transformed, data)) + + +class TernaryExportEngine(FocusExportEngine): + """FocusExportEngine that uploads one whole UTC day per delivery, with token counts merged into Tags.""" + + def __init__(self, *, prefix: str, destination: FocusTernaryDestination) -> None: + """Set the base engine's fields directly, since the generic destination factory has no Ternary destination.""" + self.provider = "ternary" + self.export_format = "csv" + self.prefix = prefix + self.destination = destination + self.serializer = self._init_serializer() + self.transformer = FocusTransformer() + self.database = FocusLiteLLMDatabase() + + async def export_window(self, *, window: FocusTimeWindow, limit: int | None) -> None: + data: Final = await _read_usage(self._database, limit=limit, window=window) + await self._deliver_enriched(data, window) + + async def export_all(self, *, limit: int | None) -> None: + data: Final = await _read_usage(self._database, limit=limit, window=None) + now: Final = datetime.now(timezone.utc) + window: Final = FocusTimeWindow( + start_time=now.replace(hour=0, minute=0, second=0, microsecond=0), + end_time=now, + frequency="all", + ) + await self._deliver_enriched(data, window) + + async def _deliver_enriched(self, data: _Frame, window: FocusTimeWindow) -> None: + if data.is_empty(): + verbose_logger.debug("Ternary export: no usage data for window %s", window) + return + failed: Final[list[str]] = [] # mutable-ok: days whose upload failed, reported once after all days run + for day, day_data in _partition_by_day(data): + try: + await self._deliver_day(day_data, _day_window(day, window.frequency)) + except Exception as e: # noqa: BLE001 # one day's failure must not block the rest; logged and re-raised + verbose_logger.warning("Ternary export: upload for %s failed: %s", day, e) + failed.append(day) + if failed: + raise RuntimeError(f"Ternary export: {len(failed)} day(s) failed to upload: {', '.join(failed)}") + + async def _deliver_day(self, data: _Frame, window: FocusTimeWindow) -> None: + payload: Final = _day_payload(data, transformer=self._transformer, serializer=self._serializer) + if not payload: + verbose_logger.debug("Ternary export: nothing to upload for %s", window.start_time.date()) + return + await self._destination.deliver( + content=payload, + time_window=window, + filename=self._build_filename(window), + ) + + +__all__ = ("TernaryExportEngine",) diff --git a/litellm/integrations/ternary/ternary_logger.py b/litellm/integrations/ternary/ternary_logger.py new file mode 100644 index 00000000000..b7a664fe3b6 --- /dev/null +++ b/litellm/integrations/ternary/ternary_logger.py @@ -0,0 +1,167 @@ +"""Ternary FocusLogger preset (CSV, Ternary destination) that adds per-row token counts to the FOCUS Tags.""" + +from __future__ import annotations + +import os +from collections.abc import Awaitable, Callable +from datetime import datetime, timedelta, timezone +from typing import TYPE_CHECKING, Final, Protocol + +import litellm +from litellm._logging import verbose_logger +from litellm.integrations.custom_logger import CustomLogger +from litellm.integrations.focus.destinations.base import FocusTimeWindow +from litellm.integrations.focus.destinations.ternary_destination import FocusTernaryDestination +from litellm.integrations.focus.focus_logger import FocusLogger + +if TYPE_CHECKING: + from litellm.integrations.focus.export_engine import FocusExportEngine + +TERNARY_USAGE_DATA_JOB_NAME: Final = "ternary_export_usage_data" + +_SUPPORTED_FREQUENCIES: Final = frozenset({"daily", "interval"}) + + +class _JobScheduler(Protocol): + """The slice of ``AsyncIOScheduler`` this module uses, which ships no type information""" + + def add_job( + self, + func: Callable[[], Awaitable[None]], + /, + **trigger: str | int, # kwargs-ok: APScheduler add_job trigger options + ) -> object: ... + + +def _parse_interval(raw: str | int | None) -> int | None: + """Parse the export interval-seconds override; a non-numeric value is ignored.""" + if raw is None: + return None + try: + return int(raw) + except (ValueError, TypeError): + verbose_logger.warning("Invalid TERNARY_EXPORT_INTERVAL_SECONDS value: %s, ignoring", raw) + return None + + +class TernaryLogger(FocusLogger): + """FocusLogger preset for Ternary (CSV, Ternary cost-import API). + + Env: TERNARY_API_KEY, TERNARY_CONNECTION_ID, TERNARY_BASE_URL (required), + TERNARY_EXPORT_FREQUENCY ("daily" default / "interval"), TERNARY_EXPORT_INTERVAL_SECONDS. + """ + + def __init__( + self, + *, + api_key: str | None = None, + connection_id: str | None = None, + base_url: str | None = None, + frequency: str | None = None, + interval_seconds: int | None = None, + ) -> None: + resolved_api_key: Final = api_key or os.getenv("TERNARY_API_KEY") + resolved_connection_id: Final = connection_id or os.getenv("TERNARY_CONNECTION_ID") + resolved_base_url: Final = base_url or os.getenv("TERNARY_BASE_URL") + resolved_frequency: Final = (frequency or os.getenv("TERNARY_EXPORT_FREQUENCY") or "daily").lower() + if resolved_frequency not in _SUPPORTED_FREQUENCIES: + raise ValueError( + f"Unsupported TERNARY_EXPORT_FREQUENCY {resolved_frequency!r}; " + f"Ternary supports {sorted(_SUPPORTED_FREQUENCIES)}. LiteLLM spend is a daily " + "aggregate landed with whole-day replace, so 'hourly' adds no grain -- use " + "'interval' only for short setup-validation loops." + ) + + resolved_interval: Final = _parse_interval( + interval_seconds if interval_seconds is not None else os.getenv("TERNARY_EXPORT_INTERVAL_SECONDS") + ) + if resolved_frequency == "interval" and (resolved_interval is None or resolved_interval <= 0): + raise ValueError( + "TERNARY_EXPORT_INTERVAL_SECONDS must be a positive integer when TERNARY_EXPORT_FREQUENCY is 'interval'" + ) + + destination_config: Final[dict[str, str]] = {} # mutable-ok: built from the config values present below + if resolved_api_key: + destination_config["api_key"] = resolved_api_key + if resolved_connection_id: + destination_config["connection_id"] = resolved_connection_id + if resolved_base_url: + destination_config["base_url"] = resolved_base_url + + super().__init__( + provider="ternary", + export_format="csv", + frequency=resolved_frequency, + interval_seconds=resolved_interval, + prefix="ternary_exports", + destination_config=destination_config, + ) + + verbose_logger.debug( + "TernaryLogger initialized (connection_id=%s)", + ( + resolved_connection_id[:4] + "***" + if resolved_connection_id and len(resolved_connection_id) > 4 + else "***" + ), + ) + + def _compute_time_window(self, now: datetime) -> FocusTimeWindow: + """Start at UTC midnight two days back so each push re-sends whole days, including spend flushed late.""" + now_utc: Final = now.astimezone(timezone.utc) + start_time: Final = (now_utc - timedelta(days=2)).replace(hour=0, minute=0, second=0, microsecond=0) + return FocusTimeWindow(start_time=start_time, end_time=now_utc, frequency=self.frequency) + + def _ensure_engine(self) -> FocusExportEngine: + """Instantiate the polars-backed Ternary export engine lazily.""" + if self._engine is None: + from .export_engine import TernaryExportEngine # noqa: PLC0415 # polars is a proxy extra + + self._engine = TernaryExportEngine( + prefix=self.prefix, + destination=FocusTernaryDestination(prefix=self.prefix, config=self._destination_config), + ) + return self._engine + + async def initialize_focus_export_job(self) -> None: + """Run the scheduled export under a Ternary-specific pod lock (distinct from FocusLogger's).""" + from litellm.proxy.proxy_server import proxy_logging_obj + + writer: Final = getattr(proxy_logging_obj, "db_spend_update_writer", None) if proxy_logging_obj else None + pod_lock_manager: Final = getattr(writer, "pod_lock_manager", None) if writer is not None else None + + if pod_lock_manager and pod_lock_manager.redis_cache: + acquired: Final = await pod_lock_manager.acquire_lock(cronjob_id=TERNARY_USAGE_DATA_JOB_NAME) + if not acquired: + verbose_logger.debug("Ternary export: unable to acquire pod lock") + return + try: + await self._run_scheduled_export() + finally: + await pod_lock_manager.release_lock(cronjob_id=TERNARY_USAGE_DATA_JOB_NAME) + else: + await self._run_scheduled_export() + + @staticmethod + async def init_ternary_background_job( + scheduler: _JobScheduler, + ) -> None: + """Register the Ternary export job with the provided scheduler.""" + ternary_loggers: Final[list[CustomLogger]] = ( # mutable-ok: list returned by the shared callback manager + litellm.logging_callback_manager.get_custom_loggers_for_type(callback_type=TernaryLogger) + ) + if not ternary_loggers: + verbose_logger.debug("No Ternary logger registered; skipping scheduler") + return + + ternary_logger: Final = ternary_loggers[0] + if not isinstance(ternary_logger, TernaryLogger): + return + trigger_kwargs: Final = ternary_logger._build_scheduler_trigger() + scheduler.add_job( + ternary_logger.initialize_focus_export_job, + **trigger_kwargs, + ) + + +__all__ = ("TernaryLogger",) diff --git a/litellm/litellm_core_utils/custom_logger_registry.py b/litellm/litellm_core_utils/custom_logger_registry.py index 8367c9ee00e..74badc2c8c1 100644 --- a/litellm/litellm_core_utils/custom_logger_registry.py +++ b/litellm/litellm_core_utils/custom_logger_registry.py @@ -48,6 +48,7 @@ from litellm.integrations.posthog import PostHogLogger from litellm.integrations.prometheus import PrometheusLogger from litellm.integrations.s3_v2 import S3Logger from litellm.integrations.sqs import SQSLogger +from litellm.integrations.ternary.ternary_logger import TernaryLogger from litellm.integrations.vantage.vantage_logger import VantageLogger from litellm.integrations.vector_store_integrations.vector_store_pre_call_hook import ( VectorStorePreCallHook, @@ -117,6 +118,7 @@ class CustomLoggerRegistry: "focus": FocusLogger, "mavvrik": MavvrikFocusLogger, "vantage": VantageLogger, + "ternary": TernaryLogger, "posthog": PostHogLogger, "newrelic": NewRelicLogger, } diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index d73d8813308..7b965dea28a 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -4951,6 +4951,15 @@ def _init_custom_logger_compatible_class( vantage_logger: Final = VantageLogger() _in_memory_loggers.append(vantage_logger) return vantage_logger + elif logging_integration == "ternary": + from litellm.integrations.ternary.ternary_logger import TernaryLogger + + for callback in _in_memory_loggers: + if isinstance(callback, TernaryLogger): + return callback + ternary_logger: Final = TernaryLogger() + _in_memory_loggers.append(ternary_logger) + return ternary_logger elif logging_integration == "deepeval": for callback in _in_memory_loggers: if isinstance(callback, DeepEvalLogger): @@ -5456,6 +5465,12 @@ def get_custom_logger_compatible_class( for callback in _in_memory_loggers: if isinstance(callback, VantageLogger): return callback + elif logging_integration == "ternary": + from litellm.integrations.ternary.ternary_logger import TernaryLogger + + for callback in _in_memory_loggers: + if isinstance(callback, TernaryLogger): + 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 7fe175131d8..c2ca045236c 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -11339,6 +11339,14 @@ class ProxyStartupEvent: await MavvrikFocusLogger.init_mavvrik_focus_background_job(scheduler=scheduler) + from litellm.integrations.ternary.ternary_logger import ( # noqa: PLC0415 # lazy import avoids a circular import at module load + TernaryLogger, + ) + + await TernaryLogger.init_ternary_background_job( + scheduler=scheduler, # pyright: ignore[reportUnknownArgumentType] # APScheduler ships no type information + ) + ######################################################## # Prometheus Background Job ######################################################## diff --git a/tests/integration/observability/_ternary_support.py b/tests/integration/observability/_ternary_support.py new file mode 100644 index 00000000000..99f1c760c7e --- /dev/null +++ b/tests/integration/observability/_ternary_support.py @@ -0,0 +1,229 @@ +"""Owned Ternary peers for the FOCUS export contracts: the cost-import API and the signed-URL storage host. + +The API answers `upload-url` with a signed `https://storage.googleapis.com/...` PUT target, the way Ternary +does. The proxy only PUTs to an HTTPS GCS host, so the storage peer serves TLS under that name and the +proxy reaches it through an owned tunnelling forward proxy that trusts the peer's own certificate. +""" + +from __future__ import annotations + +import csv +import io +import json +import threading +import uuid +from collections.abc import Iterator, Mapping +from contextlib import contextmanager +from dataclasses import dataclass, field +from datetime import date, datetime, timedelta, timezone +from pathlib import Path +from types import MappingProxyType +from typing import Final +from urllib.parse import urlsplit + +import psycopg +import yaml +from integration._support.forward_proxy import tunnelling_forward_proxy +from integration._support.tls import server_context, write_self_signed_cert +from integration._support.wire import Reply, Request, Wire, wire_server +from psycopg import sql + +TERNARY_API_KEY: Final = "synthetic-ternary-api-key" +CONNECTION_ID: Final = "ternary-integration-connection" +GCS_HOST: Final = "storage.googleapis.com" +BUCKET: Final = "ternary-integration-bucket" +UPLOAD_URL_PATH: Final = f"/external-cost-sources/v1/{CONNECTION_ID}/upload-url" +SIGNED_HEADERS: Final = MappingProxyType( + { + "Content-Encoding": "identity", + "Content-Type": "text/csv", + "x-goog-content-length-range": "0,67108864", + "x-goog-meta-ecs-connection-id": CONNECTION_ID, + "x-goog-meta-ecs-provider": "litellm", + } +) +_SCRUBBED: Final = frozenset( + {"HTTPS_PROXY", "HTTP_PROXY", "ALL_PROXY", "NO_PROXY", "SSL_CERT_FILE", "SSL_CERT_DIR", "REQUESTS_CA_BUNDLE"} +) + + +def utc_today() -> date: + return datetime.now(timezone.utc).date() + + +def signed_target(day: str) -> str: + return f"/{BUCKET}/{CONNECTION_ID}/{day}.csv?X-Goog-Signature=synthetic-{day}" + + +@dataclass(frozen=True, slots=True) +class Upload: + day: str + headers: Mapping[str, str] + body: bytes + + def rows(self) -> tuple[Mapping[str, str], ...]: + return tuple(MappingProxyType(row) for row in csv.DictReader(io.StringIO(self.body.decode("utf-8")))) + + def cost_by_service(self) -> Mapping[str, float]: + return MappingProxyType({row["ServiceName"]: float(row["BilledCost"]) for row in self.rows()}) + + +@dataclass(slots=True) +class TernaryPeers: + api: Wire + storage: Wire + environment: Mapping[str, str] + _upload_url_requests: list[Request] = field(default_factory=list) # mutable-ok: drain() consumes the wire queue + _puts: list[Request] = field(default_factory=list) # mutable-ok: drain() consumes the wire queue + + def upload_url_requests(self) -> tuple[Request, ...]: + self._upload_url_requests.extend(self.api.drain()) + return tuple(self._upload_url_requests) + + def puts(self) -> tuple[Request, ...]: + self._puts.extend(self.storage.drain()) + return tuple(self._puts) + + def uploads(self) -> tuple[Upload, ...]: + return tuple( + Upload(day=urlsplit(put.target).path.rsplit("/", 1)[-1].removesuffix(".csv"), headers=put.headers, body=put.body) + for put in self.puts() + if put.method == "PUT" + ) + + def latest_by_day(self) -> Mapping[str, Upload]: + return MappingProxyType({upload.day: upload for upload in self.uploads()}) + + def outbound(self) -> int: + return len(self.upload_url_requests()) + len(self.puts()) + + +def _api(request: Request) -> Reply: + if (request.method, request.target) != ("POST", UPLOAD_URL_PATH): + return Reply(status=404, body=b'{"message": "unexpected Ternary request"}') + if request.headers.get("authorization") != f"Bearer {TERNARY_API_KEY}": + return Reply(status=401, body=b'{"message": "invalid secret"}') + day: Final = str(json.loads(request.body)["day"]) + return Reply( + body=json.dumps( + { + "url": f"https://{GCS_HOST}{signed_target(day)}", + "method": "PUT", + "headers": dict(SIGNED_HEADERS), + "expiresAt": (datetime.now(timezone.utc) + timedelta(seconds=899)).isoformat(), + "day": day, + } + ).encode() + ) + + +def _storage(request: Request) -> Reply: + if request.method != "PUT" or not request.target.startswith(f"/{BUCKET}/{CONNECTION_ID}/"): + return Reply(status=404, body=b"NoSuchKey", content_type="application/xml") + return Reply(body=b"", content_type="text/plain") + + +@contextmanager +def ternary_peers(directory: Path) -> Iterator[TernaryPeers]: + cert, key = write_self_signed_cert(directory, (GCS_HOST,)) + with ( + wire_server(_storage, tls=server_context(cert, key)) as storage, + tunnelling_forward_proxy(int(storage.url.rsplit(":", 1)[1])) as tunnel, + wire_server(_api) as api, + ): + yield TernaryPeers( + api=api, + storage=storage, + environment=MappingProxyType( + { + "TERNARY_API_KEY": TERNARY_API_KEY, + "TERNARY_CONNECTION_ID": CONNECTION_ID, + "TERNARY_BASE_URL": api.url, + "HTTPS_PROXY": tunnel.url, + "NO_PROXY": "127.0.0.1,localhost", + "SSL_CERT_FILE": str(cert), + } + ), + ) + + +def scrubbed_environment(inherited: Mapping[str, str]) -> tuple[str, ...]: + return tuple(name for name in inherited if name.upper() in _SCRUBBED) + + +def callbacks_config(directory: Path, callbacks: tuple[str, ...]) -> Path: + config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + config["litellm_settings"]["callbacks"] = list(callbacks) + path: Final = directory / "ternary-callbacks.yaml" + path.write_text(yaml.safe_dump(config)) + return path + + +def set_database_timezone(database_url: str, zone: str) -> None: + name: Final = urlsplit(database_url).path.lstrip("/") + with psycopg.connect(database_url, autocommit=True) as connection: + connection.execute(sql.SQL("ALTER DATABASE {} SET timezone TO {}").format(sql.Identifier(name), sql.Literal(zone))) + + +@dataclass(frozen=True, slots=True) +class SpendRow: + """One `LiteLLM_DailyUserSpend` row; `service` becomes the FOCUS `ServiceName` (the model group).""" + + day: date + service: str + spend: float + updated_at: datetime + prompt_tokens: int = 7 + completion_tokens: int = 3 + + +def seed_daily_spend(database_url: str, rows: tuple[SpendRow, ...]) -> None: + """Insert every row in one transaction so no export run can read a partial seed.""" + with psycopg.connect(database_url) as connection, connection.transaction(): + for row in rows: + naive_utc: Final = row.updated_at.astimezone(timezone.utc).replace(tzinfo=None) + connection.execute( + """ + INSERT INTO "LiteLLM_DailyUserSpend" + (id, user_id, date, api_key, model, model_group, custom_llm_provider, + prompt_tokens, completion_tokens, spend, api_requests, successful_requests, + failed_requests, created_at, updated_at) + VALUES (%s, %s, %s, %s, %s, %s, 'openai', %s, %s, %s, 1, 1, 0, %s, %s) + """, + ( + str(uuid.uuid4()), + "ternary-integration-user", + row.day.isoformat(), + f"ternary-integration-key-{row.service}", + "gpt-4o-mini", + row.service, + row.prompt_tokens, + row.completion_tokens, + row.spend, + naive_utc, + naive_utc, + ), + ) + + +@contextmanager +def refreshing_updated_at(database_url: str, service: str) -> Iterator[None]: + """Keep a seeded row's `updated_at` at the current UTC instant, as ongoing traffic would.""" + stop: Final = threading.Event() + + def refresh() -> None: + with psycopg.connect(database_url, autocommit=True) as connection: + while not stop.wait(0.2): + connection.execute( + """UPDATE "LiteLLM_DailyUserSpend" SET updated_at = (now() AT TIME ZONE 'UTC') WHERE model_group = %s""", + (service,), + ) + + worker: Final = threading.Thread(target=refresh, daemon=True) + worker.start() + try: + yield + finally: + stop.set() + worker.join(timeout=5) + assert not worker.is_alive(), "updated_at refresher survived cleanup" diff --git a/tests/integration/observability/test_ternary_focus_export.py b/tests/integration/observability/test_ternary_focus_export.py new file mode 100644 index 00000000000..e077b70a3b3 --- /dev/null +++ b/tests/integration/observability/test_ternary_focus_export.py @@ -0,0 +1,293 @@ +"""Ternary FOCUS export: one whole UTC day per signed-URL PUT, whatever the database session timezone.""" + +from __future__ import annotations + +import json +import os +import subprocess +import sys +import uuid +from collections.abc import Iterator, Mapping +from contextlib import contextmanager +from datetime import date, datetime, time, timedelta, timezone +from pathlib import Path +from types import MappingProxyType +from typing import Final + +import pytest +from _ternary_support import ( + SIGNED_HEADERS, + TERNARY_API_KEY, + UPLOAD_URL_PATH, + SpendRow, + TernaryPeers, + callbacks_config, + refreshing_updated_at, + scrubbed_environment, + seed_daily_spend, + set_database_timezone, + signed_target, + ternary_peers, + utc_today, +) +from integration._support.client import Gateway, eventually +from integration._support.database import scratch_database +from integration._support.process import OwnedProxy, owned_proxy_process +from integration._support.wire import Reply, Request, wire_server + +_INTERVAL_EXPORT: Final = MappingProxyType( + {"TERNARY_EXPORT_FREQUENCY": "interval", "TERNARY_EXPORT_INTERVAL_SECONDS": "1"} +) +_PROVIDER_KEY: Final = "synthetic-provider-key" +_FOCUS_REFUSAL: Final = "Provider 'ternary' not supported for Focus export" + + +def _midnight(day: date) -> datetime: + return datetime.combine(day, time(), tzinfo=timezone.utc) + + +def _upstream(request: Request) -> Reply: + body: Final = json.loads(request.body) + return Reply( + body=json.dumps( + { + "id": f"chatcmpl-{uuid.uuid4().hex}", + "object": "chat.completion", + "created": 1, + "model": body["model"], + "choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 11, "completion_tokens": 4, "total_tokens": 15}, + } + ).encode() + ) + + +@contextmanager +def _ternary_proxy( + gateway: Gateway, + directory: Path, + peers: TernaryPeers, + *, + callbacks: tuple[str, ...], + database_url: str | None = None, + extra: Mapping[str, str] = MappingProxyType({}), +) -> Iterator[OwnedProxy]: + overrides: Final = { + **peers.environment, + **({"DATABASE_URL": database_url} if database_url else {}), + **extra, + } + with owned_proxy_process( + gateway, + directory, + overrides, + config=callbacks_config(directory, callbacks), + remove_environment=scrubbed_environment(os.environ), + ) as owned: + yield owned + + +def test_ternary_export_uploads_one_csv_per_utc_day_with_the_signed_headers(gateway: Gateway, tmp_path: Path) -> None: + marker: Final = uuid.uuid4().hex[:10] + today: Final = utc_today() + midnight: Final = _midnight(today) + seeds: Final = ( + SpendRow(today - timedelta(days=3), f"before-window-{marker}", 0.5, midnight - timedelta(days=3, hours=-12)), + SpendRow(today - timedelta(days=2), f"oldest-{marker}", 0.0123, midnight - timedelta(days=2, hours=-12)), + SpendRow(today - timedelta(days=2), f"late-flush-{marker}", 0.0045, midnight - timedelta(hours=23, minutes=30)), + SpendRow(today - timedelta(days=1), f"yesterday-{marker}", 0.0067, midnight - timedelta(hours=12)), + SpendRow(today, f"today-{marker}", 0.0089, datetime.now(timezone.utc) - timedelta(minutes=1)), + ) + + with ( + scratch_database("ternary_happy") as database_url, + ternary_peers(tmp_path) as peers, + wire_server(_upstream) as provider, + _ternary_proxy( + gateway, tmp_path, peers, callbacks=("ternary",), database_url=database_url, extra=_INTERVAL_EXPORT + ) as owned, + owned.gateway.scenario() as scenario, + ): + model: Final = scenario.model( + api_base=provider.url + "/v1", + api_key=_PROVIDER_KEY, + input_cost_per_token=0.001, + output_cost_per_token=0.002, + ) + key: Final = scenario.key(models=[model]) + response: Final = owned.gateway.request( + "POST", + "/v1/chat/completions", + {"model": model, "messages": [{"role": "user", "content": marker}], "cache": {"no-cache": True}}, + key=key, + ) + assert response.status_code == 200, response.text + seed_daily_spend(database_url, seeds) + window: Final = tuple(str(today - timedelta(days=offset)) for offset in (2, 1, 0)) + by_day: Final = eventually( + peers.latest_by_day, + lambda latest: set(window) <= set(latest) and model in latest[str(today)].cost_by_service(), + seconds=90, + ) + requests: Final = peers.upload_url_requests() + uploads: Final = peers.uploads() + + assert {upload.day for upload in uploads} == set(window), "a day outside the two-day window was uploaded" + assert dict(by_day[window[0]].cost_by_service()) == pytest.approx( + {f"oldest-{marker}": 0.0123, f"late-flush-{marker}": 0.0045} + ) + assert dict(by_day[window[1]].cost_by_service()) == pytest.approx({f"yesterday-{marker}": 0.0067}) + assert dict(by_day[window[2]].cost_by_service()) == pytest.approx({f"today-{marker}": 0.0089, model: 0.019}) + for day, upload in by_day.items(): + assert {row["ChargePeriodStart"][:10] for row in upload.rows()} == {day} + oldest: Final = next(row for row in by_day[window[0]].rows() if row["ServiceName"] == f"oldest-{marker}") + assert json.loads(oldest["Tags"])["prompt_tokens"] == "7" + assert json.loads(oldest["Tags"])["completion_tokens"] == "3" + for request in requests: + assert (request.method, request.target) == ("POST", UPLOAD_URL_PATH) + assert request.headers["authorization"] == f"Bearer {TERNARY_API_KEY}" + assert json.loads(request.body)["day"] in window + for put in peers.puts(): + assert put.method == "PUT" + assert "authorization" not in put.headers + assert {name.lower(): value for name, value in SIGNED_HEADERS.items()}.items() <= put.headers.items() + assert put.target == signed_target(put.target.rsplit("/", 1)[-1].split(".csv", 1)[0]) + assert TERNARY_API_KEY.encode() not in put.body + + +@pytest.mark.parametrize( + ("zone", "label"), + [("Asia/Tokyo", "east"), ("Pacific/Kiritimati", "east-most")], +) +def test_ternary_export_keeps_the_whole_oldest_day_when_the_database_session_is_east_of_utc( + gateway: Gateway, tmp_path: Path, zone: str, label: str +) -> None: + marker: Final = uuid.uuid4().hex[:10] + today: Final = utc_today() + oldest: Final = today - timedelta(days=2) + seeds: Final = ( + SpendRow(oldest - timedelta(days=1), f"before-window-{marker}", 0.5, _midnight(oldest) - timedelta(hours=12)), + SpendRow(oldest, f"early-{label}-{marker}", 0.0021, _midnight(oldest) + timedelta(hours=3)), + SpendRow(oldest, f"midday-{label}-{marker}", 0.0034, _midnight(oldest) + timedelta(hours=15)), + ) + + with scratch_database("ternary_east") as database_url: + set_database_timezone(database_url, zone) + with ( + ternary_peers(tmp_path) as peers, + _ternary_proxy( + gateway, tmp_path, peers, callbacks=("ternary",), database_url=database_url, extra=_INTERVAL_EXPORT + ), + ): + seed_daily_spend(database_url, seeds) + by_day: Final = eventually(peers.latest_by_day, lambda latest: str(oldest) in latest, seconds=90) + uploads: Final = peers.uploads() + + assert str(oldest - timedelta(days=1)) not in {upload.day for upload in uploads} + assert dict(by_day[str(oldest)].cost_by_service()) == pytest.approx( + {f"early-{label}-{marker}": 0.0021, f"midday-{label}-{marker}": 0.0034} + ), f"the {zone} session dropped spend updated in the first hours of the oldest UTC day" + + +def test_ternary_export_includes_spend_updated_in_the_last_hours_when_the_database_session_is_west_of_utc( + gateway: Gateway, tmp_path: Path +) -> None: + marker: Final = uuid.uuid4().hex[:10] + today: Final = utc_today() + yesterday: Final = today - timedelta(days=1) + seeds: Final = ( + SpendRow(yesterday, f"midday-{marker}", 0.0056, _midnight(yesterday) + timedelta(hours=12)), + SpendRow(yesterday, f"just-flushed-{marker}", 0.0078, datetime.now(timezone.utc) - timedelta(minutes=30)), + ) + + with scratch_database("ternary_west") as database_url: + set_database_timezone(database_url, "America/Los_Angeles") + with ( + ternary_peers(tmp_path) as peers, + _ternary_proxy( + gateway, tmp_path, peers, callbacks=("ternary",), database_url=database_url, extra=_INTERVAL_EXPORT + ), + ): + seed_daily_spend(database_url, seeds) + by_day: Final = eventually(peers.latest_by_day, lambda latest: str(yesterday) in latest, seconds=90) + + assert dict(by_day[str(yesterday)].cost_by_service()) == pytest.approx( + {f"midday-{marker}": 0.0056, f"just-flushed-{marker}": 0.0078} + ), "the America/Los_Angeles session dropped spend flushed in the last hours" + + +@pytest.mark.parametrize("export_format", ["parquet", "csv"]) +def test_the_generic_focus_logger_does_not_upload_to_ternary(gateway: Gateway, tmp_path: Path, export_format: str) -> None: + service: Final = f"generic-focus-{uuid.uuid4().hex[:10]}" + today: Final = utc_today() + + with ( + scratch_database("ternary_generic") as database_url, + ternary_peers(tmp_path) as peers, + _ternary_proxy( + gateway, + tmp_path, + peers, + callbacks=("focus",), + database_url=database_url, + extra={ + "FOCUS_PROVIDER": "ternary", + "FOCUS_FORMAT": export_format, + "FOCUS_FREQUENCY": "interval", + "FOCUS_INTERVAL_SECONDS": "1", + }, + ) as owned, + ): + seed_daily_spend(database_url, (SpendRow(today, service, 0.0042, datetime.now(timezone.utc)),)) + with refreshing_updated_at(database_url, service): + outbound, refused = eventually( + lambda: (peers.outbound(), _FOCUS_REFUSAL in owned.log.read_text(errors="replace")), + lambda observed: observed[0] > 0 or observed[1], + seconds=60, + ) + + assert outbound == 0, f"callbacks: ['focus'] with FOCUS_PROVIDER=ternary uploaded {export_format} to Ternary" + assert refused + + +@pytest.mark.parametrize("callbacks", [(), ("ternary",)], ids=["no-focus-callback", "ternary-callback"]) +def test_a_proxy_installed_without_polars_boots_and_serves( + gateway: Gateway, tmp_path: Path, callbacks: tuple[str, ...] +) -> None: + shadow: Final = tmp_path / "without-polars" + (shadow / "polars").mkdir(parents=True) + (shadow / "polars" / "__init__.py").write_text( + "raise ModuleNotFoundError(\"No module named 'polars'\", name='polars')\n" + ) + python_path: Final = os.pathsep.join(filter(None, (str(shadow), os.environ.get("PYTHONPATH", "")))) + probe: Final = subprocess.run( # noqa: S603 # fixed argv: proves the shadow hides polars from a child + (sys.executable, "-P", "-c", "import polars"), + env={**os.environ, "PYTHONPATH": python_path}, + capture_output=True, + text=True, + check=False, + ) + assert probe.returncode != 0 and "No module named 'polars'" in probe.stderr, probe.stderr + + with ( + ternary_peers(tmp_path) as peers, + wire_server(_upstream) as provider, + _ternary_proxy( + gateway, tmp_path, peers, callbacks=callbacks, extra={"PYTHONPATH": python_path, **_INTERVAL_EXPORT} + ) as owned, + owned.gateway.scenario() as scenario, + ): + model: Final = scenario.model(api_base=provider.url + "/v1", api_key=_PROVIDER_KEY) + key: Final = scenario.key(models=[model]) + response: Final = owned.gateway.request( + "POST", "/v1/chat/completions", {"model": model, "messages": [{"role": "user", "content": "hi"}]}, key=key + ) + if callbacks: + eventually( + lambda: "No module named 'polars'" in owned.log.read_text(errors="replace"), + lambda failed: failed, + seconds=60, + ) + + assert response.status_code == 200, response.text + assert response.json()["choices"][0]["message"]["content"] == "ok" diff --git a/tests/unit/integrations/focus/test_ternary_destination.py b/tests/unit/integrations/focus/test_ternary_destination.py new file mode 100644 index 00000000000..6fe8a2b2931 --- /dev/null +++ b/tests/unit/integrations/focus/test_ternary_destination.py @@ -0,0 +1,274 @@ +"""Tests for FocusTernaryDestination behavior.""" + +from __future__ import annotations + +from datetime import datetime, timedelta, timezone +from typing import Any +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +from litellm.integrations.focus.destinations.base import FocusTimeWindow +from litellm.integrations.focus.destinations.ternary_destination import ( + TERNARY_UPLOAD_TIMEOUT_SECONDS, + FocusTernaryDestination, + _NoRedirectHTTPHandler, +) + +MOCK_TARGET = "litellm.integrations.focus.destinations.ternary_destination._NoRedirectHTTPHandler" + + +def _window(freq: str = "daily", hour: int = 5) -> FocusTimeWindow: + start = datetime(2024, 1, 2, hour, tzinfo=timezone.utc) + end = start + timedelta(hours=1) + return FocusTimeWindow(start_time=start, end_time=end, frequency=freq) + + +def _config(**overrides: Any) -> dict[str, Any]: + base = { + "api_key": "test-api-key", + "connection_id": "conn-1234", + "base_url": "https://ternary.test", + } + base.update(overrides) + return base + + +_SIGNED_URL = "https://storage.googleapis.com/ter-ecs-t-us/days/cloud-1/2024-01-02.csv?X-Goog-Signature=abc" +_SIGNED_HEADERS = { + "Content-Type": "text/csv", + "x-goog-meta-ecs-provider": "litellm", + "x-goog-meta-ecs-connection-id": "conn-1234", + "x-goog-content-length-range": "0,1073741824", +} + + +def _response(status: int = 200, body: Any = None) -> MagicMock: + response = MagicMock() + response.status_code = status + response.json = MagicMock(return_value=body) + response.text = "" if body is None else str(body) + return response + + +def _upload_url_body(**overrides: Any) -> dict[str, Any]: + body = { + "url": _SIGNED_URL, + "method": "PUT", + "headers": dict(_SIGNED_HEADERS), + "expiresAt": "2024-01-02T05:15:00Z", + "day": "2024-01-02", + } + body.update(overrides) + return body + + +def _handler( + *, url_response: MagicMock | None = None, put_response: MagicMock | None = None +) -> tuple[MagicMock, list[dict[str, Any]], list[dict[str, Any]]]: + """Return a mock handler whose client records each upload-url POST and signed-URL PUT.""" + posts: list[dict[str, Any]] = [] + puts: list[dict[str, Any]] = [] + + async def capture_post(url, **kwargs): + posts.append({"url": url, **kwargs}) + return url_response if url_response is not None else _response(body=_upload_url_body()) + + async def capture_put(url, **kwargs): + puts.append({"url": url, **kwargs}) + return put_response if put_response is not None else _response() + + handler = MagicMock() + handler.close = AsyncMock() + handler.client.post = capture_post + handler.client.put = capture_put + return handler, posts, puts + + +@pytest.mark.asyncio +async def test_should_skip_empty_content(): + dest = FocusTernaryDestination(prefix="exports", config=_config()) + with patch(MOCK_TARGET, side_effect=AssertionError("client initialized on empty content")): + assert await dest.deliver(content=b"", time_window=_window(), filename="usage.csv") is None + + +@pytest.mark.asyncio +async def test_should_request_a_signed_url_for_the_window_day_then_put_the_csv(): + dest = FocusTernaryDestination(prefix="exports", config=_config()) + handler, posts, puts = _handler() + + with patch(MOCK_TARGET, return_value=handler): + await dest.deliver(content=b"header\nrow1\n", time_window=_window(), filename="usage.csv") + + assert len(posts) == 1 + assert posts[0]["url"] == "https://ternary.test/external-cost-sources/v1/conn-1234/upload-url" + assert posts[0]["headers"] == {"Authorization": "Bearer test-api-key"} + assert posts[0]["json"] == {"day": "2024-01-02"} + assert posts[0]["timeout"] == TERNARY_UPLOAD_TIMEOUT_SECONDS + assert len(puts) == 1 + assert puts[0]["url"] == _SIGNED_URL + assert puts[0]["content"] == b"header\nrow1\n" + assert puts[0]["headers"] == _SIGNED_HEADERS + assert "Authorization" not in puts[0]["headers"] + handler.close.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_should_key_the_upload_on_the_window_start_day_not_its_end(): + dest = FocusTernaryDestination(prefix="exports", config=_config()) + start = datetime(2024, 1, 2, tzinfo=timezone.utc) + one_day = FocusTimeWindow(start_time=start, end_time=start + timedelta(days=1), frequency="daily") + handler, posts, _ = _handler() + with patch(MOCK_TARGET, return_value=handler): + await dest.deliver(content=b"h\nr\n", time_window=one_day, filename="usage.csv") + assert posts[0]["json"] == {"day": "2024-01-02"} + + +@pytest.mark.asyncio +async def test_no_redirect_handler_disables_redirects_including_retry_client(): + handler = _NoRedirectHTTPHandler() + retry_client = handler.create_client(timeout=1.0, event_hooks=None) + try: + assert handler.client.follow_redirects is False + assert retry_client.follow_redirects is False + finally: + await retry_client.aclose() + await handler.close() + + +@pytest.mark.asyncio +async def test_should_url_encode_the_connection_id(): + dest = FocusTernaryDestination(prefix="exports", config=_config(connection_id="conn+id~ok")) + handler, posts, _ = _handler() + with patch(MOCK_TARGET, return_value=handler): + await dest.deliver(content=b"h\nr\n", time_window=_window(), filename="usage.csv") + assert posts[0]["url"].endswith("/external-cost-sources/v1/conn%2Bid~ok/upload-url") + + +@pytest.mark.asyncio +async def test_should_pass_tags_column_through_unstripped(): + """The Ternary sink must not drop/strip any columns (forwards Tags as-is).""" + dest = FocusTernaryDestination(prefix="exports", config=_config()) + handler, _, puts = _handler() + + content = b'ServiceName,Tags,x_unknown\nfoo,"{""team_id"": ""t1""}",keepme\n' + with patch(MOCK_TARGET, return_value=handler): + await dest.deliver(content=content, time_window=_window(), filename="usage.csv") + + assert puts[0]["content"] == content + + +@pytest.mark.asyncio +@pytest.mark.parametrize("status", [400, 401, 404, 422, 429, 500]) +async def test_should_raise_and_skip_the_put_when_the_upload_url_request_fails(status): + dest = FocusTernaryDestination(prefix="exports", config=_config()) + handler, _, puts = _handler(url_response=_response(status=status, body={"message": "nope"})) + with patch(MOCK_TARGET, return_value=handler): + with pytest.raises(RuntimeError, match=f"upload-url request for 2024-01-02 failed \\({status}\\)"): + await dest.deliver(content=b"h\nr\n", time_window=_window(), filename="usage.csv") + assert puts == [] + handler.close.assert_awaited_once() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("status", [400, 403, 503]) +async def test_should_raise_when_the_signed_put_fails(status): + dest = FocusTernaryDestination(prefix="exports", config=_config()) + handler, _, _ = _handler(put_response=_response(status=status, body="SignatureDoesNotMatch")) + with patch(MOCK_TARGET, return_value=handler): + with pytest.raises(RuntimeError, match=f"upload for 2024-01-02 failed \\({status}\\)"): + await dest.deliver(content=b"h\nr\n", time_window=_window(), filename="usage.csv") + handler.close.assert_awaited_once() + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "url", + [ + "http://storage.googleapis.com/b/days/c/2024-01-02.csv", + "https://evil.example.com/b/days/c/2024-01-02.csv", + "https://storage.googleapis.com.evil.example.com/b/o", + "https://169.254.169.254/computeMetadata/v1/", + ], +) +async def test_should_refuse_to_put_to_a_non_gcs_url(url): + dest = FocusTernaryDestination(prefix="exports", config=_config()) + handler, _, puts = _handler(url_response=_response(body=_upload_url_body(url=url))) + with patch(MOCK_TARGET, return_value=handler): + with pytest.raises(ValueError, match="HTTPS GCS URL"): + await dest.deliver(content=b"h\nr\n", time_window=_window(), filename="usage.csv") + assert puts == [] + + +@pytest.mark.asyncio +async def test_should_accept_a_virtual_hosted_gcs_url(): + dest = FocusTernaryDestination(prefix="exports", config=_config()) + url = "https://ter-ecs-t-us.storage.googleapis.com/days/c/2024-01-02.csv?X-Goog-Signature=abc" + handler, _, puts = _handler(url_response=_response(body=_upload_url_body(url=url))) + with patch(MOCK_TARGET, return_value=handler): + await dest.deliver(content=b"h\nr\n", time_window=_window(), filename="usage.csv") + assert puts[0]["url"] == url + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("overrides", "match"), + [ + ({"url": ""}, "missing 'url' or 'headers'"), + ({"headers": None}, "missing 'url' or 'headers'"), + ({"method": "POST"}, "unsupported upload method 'POST'"), + ], +) +async def test_should_reject_an_incomplete_upload_url_response(overrides, match): + dest = FocusTernaryDestination(prefix="exports", config=_config()) + body = _upload_url_body(**overrides) + handler, _, puts = _handler(url_response=_response(body=body)) + with patch(MOCK_TARGET, return_value=handler): + with pytest.raises(RuntimeError, match=match): + await dest.deliver(content=b"h\nr\n", time_window=_window(), filename="usage.csv") + assert puts == [] + + +@pytest.mark.parametrize("url", ["http://api.ternary.app", "http://evil.example.com:8080", "ftp://ternary.test"]) +def test_should_reject_non_https_base_url(url): + with pytest.raises(ValueError, match="HTTPS"): + FocusTernaryDestination(prefix="exports", config=_config(base_url=url)) + + +@pytest.mark.parametrize("url", ["https://ternary.test", "http://localhost:8080", "http://127.0.0.1:8080"]) +def test_should_accept_https_or_loopback_base_url(url): + dest = FocusTernaryDestination(prefix="exports", config=_config(base_url=url)) + assert dest.base_url == url.rstrip("/") + + +def test_generic_focus_factory_does_not_build_a_ternary_destination(): + from litellm.integrations.focus.destinations.factory import FocusDestinationFactory + + with pytest.raises(NotImplementedError, match="'ternary' not supported"): + FocusDestinationFactory.create( + provider="ternary", + prefix="exports", + config={"api_key": "k", "connection_id": "c", "base_url": "https://ternary.test"}, + ) + + +def test_ternary_logger_resolves_its_destination_from_env(monkeypatch): + monkeypatch.setenv("TERNARY_API_KEY", "envk") + monkeypatch.setenv("TERNARY_CONNECTION_ID", "envc") + monkeypatch.setenv("TERNARY_BASE_URL", "https://ternary.test") + from litellm.integrations.ternary.ternary_logger import TernaryLogger + + dest = TernaryLogger()._ensure_engine().destination + assert isinstance(dest, FocusTernaryDestination) + assert (dest.api_key, dest.connection_id, dest.base_url) == ("envk", "envc", "https://ternary.test") + + +@pytest.mark.parametrize("missing", ["TERNARY_API_KEY", "TERNARY_CONNECTION_ID", "TERNARY_BASE_URL"]) +def test_ternary_logger_requires_each_env_var(monkeypatch, missing): + for var in ("TERNARY_API_KEY", "TERNARY_CONNECTION_ID", "TERNARY_BASE_URL"): + monkeypatch.setenv(var, "https://ternary.test" if var == "TERNARY_BASE_URL" else "x") + monkeypatch.delenv(missing, raising=False) + from litellm.integrations.ternary.ternary_logger import TernaryLogger + + with pytest.raises(ValueError, match=missing): + TernaryLogger()._ensure_engine() diff --git a/tests/unit/integrations/ternary/__init__.py b/tests/unit/integrations/ternary/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/integrations/ternary/test_ternary_logger.py b/tests/unit/integrations/ternary/test_ternary_logger.py new file mode 100644 index 00000000000..457389a063e --- /dev/null +++ b/tests/unit/integrations/ternary/test_ternary_logger.py @@ -0,0 +1,668 @@ +"""Tests for TernaryLogger configuration and Tags token enrichment.""" + +from __future__ import annotations + +import csv +import io +import json +from datetime import datetime, timedelta, timezone +from unittest.mock import AsyncMock, MagicMock + +import polars as pl +import pytest + +from litellm.integrations.focus.destinations.base import FocusTimeWindow +from litellm.integrations.focus.focus_logger import FocusLogger +from litellm.integrations.focus.destinations.ternary_destination import FocusTernaryDestination +from litellm.integrations.focus.serializers import FocusCsvSerializer +from litellm.integrations.ternary.export_engine import ( + TernaryExportEngine, + _canonical_day, + _day_window, + _keep_window_days, + _merge_token_tags, + _partition_by_day, +) +from litellm.integrations.ternary.ternary_logger import ( + TERNARY_USAGE_DATA_JOB_NAME, + TernaryLogger, +) +from tests.test_litellm_rust.support.child_interpreter import run_child_interpreter + + +def _logger() -> TernaryLogger: + return TernaryLogger(api_key="k", connection_id="c", base_url="http://localhost:8080") + + +def _window() -> FocusTimeWindow: + from datetime import datetime, timezone + + now = datetime(2026, 9, 4, tzinfo=timezone.utc) + return FocusTimeWindow(start_time=now, end_time=now, frequency="daily") + + +def test_should_preset_focus_config_for_ternary(): + logger = TernaryLogger(api_key="k", connection_id="c", base_url="http://localhost:8080") + assert isinstance(logger, FocusLogger) + assert logger.provider == "ternary" + assert logger.export_format == "csv" + assert logger.frequency == "daily" + assert logger.prefix == "ternary_exports" + assert logger._destination_config == { + "api_key": "k", + "connection_id": "c", + "base_url": "http://localhost:8080", + } + + +def test_should_read_frequency_from_env(monkeypatch): + monkeypatch.setenv("TERNARY_API_KEY", "k") + monkeypatch.setenv("TERNARY_CONNECTION_ID", "c") + monkeypatch.setenv("TERNARY_EXPORT_FREQUENCY", "daily") + logger = TernaryLogger() + assert logger.frequency == "daily" + assert logger._destination_config["api_key"] == "k" + assert logger._destination_config["connection_id"] == "c" + + +def test_should_default_to_daily(): + logger = TernaryLogger(api_key="k", connection_id="c", base_url="http://localhost:8080") + assert logger.frequency == "daily" + + +def test_should_accept_interval_with_explicit_seconds(): + logger = TernaryLogger( + api_key="k", + connection_id="c", + base_url="http://localhost:8080", + frequency="interval", + interval_seconds=30, + ) + assert logger.frequency == "interval" + assert logger.interval_seconds == 30 + + +def test_should_reject_interval_without_seconds(): + with pytest.raises(ValueError, match="TERNARY_EXPORT_INTERVAL_SECONDS"): + TernaryLogger(api_key="k", connection_id="c", base_url="http://localhost:8080", frequency="interval") + + +@pytest.mark.parametrize("bad_frequency", ["hourly", "weekly", "minutely"]) +def test_should_reject_unsupported_frequency(bad_frequency): + with pytest.raises(ValueError, match="TERNARY_EXPORT_FREQUENCY"): + TernaryLogger(api_key="k", connection_id="c", base_url="http://localhost:8080", frequency=bad_frequency) + + +def test_compute_time_window_is_day_aligned_and_carries_whole_days(): + from datetime import datetime, timezone + + logger = TernaryLogger(api_key="k", connection_id="c", base_url="http://localhost:8080") + now = datetime(2026, 9, 4, 13, 47, 5, tzinfo=timezone.utc) + window = logger._compute_time_window(now) + assert window.start_time == datetime(2026, 9, 2, 0, 0, 0, tzinfo=timezone.utc) + assert window.end_time == now + + +@pytest.mark.asyncio +async def test_daily_export_resends_yesterday_whole_including_spend_flushed_after_midnight(): + """LiteLLM flushes spend in batches, so a 2026-10-07 row can be written at 2026-10-08 00:10, after that + day's 00:05 export. The next daily run (2026-10-09 00:05) must still upload 2026-10-07 complete.""" + from datetime import datetime, timezone + + rows = pl.DataFrame( + { + "date": ["2026-10-07", "2026-10-07", "2026-10-08"], + "updated_at": [ + datetime(2026, 10, 7, 12, 0, tzinfo=timezone.utc), + datetime(2026, 10, 8, 0, 10, tzinfo=timezone.utc), + datetime(2026, 10, 8, 9, 0, tzinfo=timezone.utc), + ], + "spend": [1.0, 2.0, 3.0], + } + ) + logger = TernaryLogger(api_key="k", connection_id="c", base_url="http://localhost:8080") + window = logger._compute_time_window(datetime(2026, 10, 9, 0, 5, tzinfo=timezone.utc)) + + engine = MagicMock() + + async def fake_get_usage_data(*, limit, start_time_utc, end_time_utc=None): + kept = rows.filter(pl.col("updated_at") >= start_time_utc) + return kept if end_time_utc is None else kept.filter(pl.col("updated_at") <= end_time_utc) + + engine._database.get_usage_data = fake_get_usage_data + captured = {} + + async def fake_deliver(data, _window): + captured["data"] = data + + engine._deliver_enriched = fake_deliver + await TernaryExportEngine.export_window(engine, window=window, limit=None) + + sent = captured["data"] + assert sent.filter(pl.col("date") == "2026-10-07")["spend"].sum() == 3.0 + + +def _day_window_from(start: datetime, end: datetime) -> FocusTimeWindow: + return FocusTimeWindow(start_time=start, end_time=end, frequency="interval") + + +def test_keep_window_days_removes_days_outside_the_window(): + data = pl.DataFrame( + {"date": ["2026-09-02", "2026-09-03", "2026-09-04", "2026-09-05"], "spend": [1.0, 2.0, 3.0, 4.0]} + ) + window = _day_window_from( + datetime(2026, 9, 3, tzinfo=timezone.utc), datetime(2026, 9, 4, 13, 0, tzinfo=timezone.utc) + ) + out = _keep_window_days(data, window) + assert out["date"].to_list() == ["2026-09-03", "2026-09-04"] + + +def test_keep_window_days_excludes_the_day_a_window_ends_at_midnight(): + data = pl.DataFrame({"date": ["2026-09-01", "2026-09-02"], "spend": [1.0, 2.0]}) + window = _day_window_from(datetime(2026, 9, 1, tzinfo=timezone.utc), datetime(2026, 9, 2, tzinfo=timezone.utc)) + assert _keep_window_days(data, window)["date"].to_list() == ["2026-09-01"] + + +def test_keep_window_days_noop_without_date_column(): + data = pl.DataFrame({"spend": [1.0]}) + window = _day_window_from(datetime(2026, 9, 3, tzinfo=timezone.utc), datetime(2026, 9, 4, tzinfo=timezone.utc)) + out = _keep_window_days(data, window) + assert out.equals(data) + + +@pytest.mark.asyncio +async def test_export_window_keeps_only_the_window_days(): + from datetime import datetime, timezone + + multi_day = pl.DataFrame({"date": ["2026-09-01", "2026-09-03", "2026-09-04"], "spend": [1.0, 2.0, 3.0]}) + + engine = MagicMock() + + async def fake_get_usage_data(**_kwargs): + return multi_day + + engine._database.get_usage_data = fake_get_usage_data + + captured = {} + + async def fake_deliver(data, window): + captured["data"] = data + + engine._deliver_enriched = fake_deliver + + window = FocusTimeWindow( + start_time=datetime(2026, 9, 3, 0, 0, tzinfo=timezone.utc), + end_time=datetime(2026, 9, 4, 12, 0, tzinfo=timezone.utc), + frequency="daily", + ) + await TernaryExportEngine.export_window(engine, window=window, limit=None) + + assert captured["data"]["date"].to_list() == ["2026-09-03", "2026-09-04"] + + +def test_merge_token_tags_adds_raw_db_keys_and_preserves_existing(): + normalized = pl.DataFrame( + { + "Tags": [json.dumps({"team_id": "t1", "model": "gpt-4o"})], + "BilledCost": [1.23], + } + ) + source = pl.DataFrame( + { + "prompt_tokens": [48], + "completion_tokens": [274], + "cache_read_input_tokens": [12], + "cache_creation_input_tokens": [0], + } + ) + + tags = json.loads(_merge_token_tags(normalized, source)["Tags"][0]) + + assert tags["prompt_tokens"] == "48" + assert tags["completion_tokens"] == "274" + assert tags["cache_read_input_tokens"] == "12" + assert tags["cache_creation_input_tokens"] == "0" + assert not any(k.startswith("x_") for k in tags) + assert tags["team_id"] == "t1" + assert tags["model"] == "gpt-4o" + + +def test_merge_token_tags_aligns_per_row_through_the_real_transformer(): + from litellm.integrations.focus.transformer import FocusTransformer + + data = pl.DataFrame( + { + "date": ["2026-09-03", "2026-09-03", "2026-09-03"], + "spend": [1.0, 2.0, 3.0], + "api_key": ["k1", "k2", "k3"], + "api_key_alias": ["a1", "a2", "a3"], + "model": ["m1", "m2", "m3"], + "model_group": ["g1", "g2", "g3"], + "custom_llm_provider": ["openai", "openai", "anthropic"], + "team_id": ["t1", "t2", "t3"], + "team_alias": ["T1", "T2", "T3"], + "api_requests": [1, 1, 1], + "prompt_tokens": [10, 20, 30], + "completion_tokens": [11, 22, 33], + "cache_read_input_tokens": [0, 0, 0], + "cache_creation_input_tokens": [0, 0, 0], + } + ) + + transformed = FocusTransformer().transform(data) + enriched = _merge_token_tags(transformed, data) + + assert enriched.height == 3 + for i, expected_tokens in enumerate(["10", "20", "30"]): + tags = json.loads(enriched["Tags"][i]) + assert tags["prompt_tokens"] == expected_tokens + assert tags["team_id"] == f"t{i + 1}" + + +def test_merge_token_tags_skips_none_values(): + normalized = pl.DataFrame({"Tags": [json.dumps({"team_id": "t1"})]}) + source = pl.DataFrame({"prompt_tokens": [None], "completion_tokens": [10]}) + tags = json.loads(_merge_token_tags(normalized, source)["Tags"][0]) + assert "prompt_tokens" not in tags + assert tags["completion_tokens"] == "10" + + +def test_merge_token_tags_graceful_on_row_mismatch(): + normalized = pl.DataFrame({"Tags": [json.dumps({"team_id": "t1"})]}) + source = pl.DataFrame({"prompt_tokens": [1, 2]}) + out = _merge_token_tags(normalized, source) + assert out["Tags"][0] == normalized["Tags"][0] + + +def test_merge_token_tags_non_object_tags_degrades_to_empty(): + normalized = pl.DataFrame({"Tags": ["[1, 2, 3]"]}) + source = pl.DataFrame({"prompt_tokens": [5]}) + tags = json.loads(_merge_token_tags(normalized, source)["Tags"][0]) + assert tags == {"prompt_tokens": "5"} + + +def test_keep_window_days_warns_when_everything_dropped(caplog): + import logging + + from litellm._logging import verbose_logger + + data = pl.DataFrame({"date": ["2026-09-01", "2026-09-02"], "spend": [1.0, 2.0]}) + window = _day_window_from(datetime(2026, 9, 5, tzinfo=timezone.utc), datetime(2026, 9, 6, tzinfo=timezone.utc)) + with caplog.at_level(logging.WARNING, logger=verbose_logger.name): + out = _keep_window_days(data, window) + assert out.height == 0 + assert "dropped all 2 rows" in caplog.text + + +def test_merge_token_tags_graceful_on_malformed_tags(): + normalized = pl.DataFrame({"Tags": ["not-json"]}) + source = pl.DataFrame({"prompt_tokens": [5]}) + tags = json.loads(_merge_token_tags(normalized, source)["Tags"][0]) + assert tags == {"prompt_tokens": "5"} + + +def test_merge_token_tags_noop_without_tags_column(): + normalized = pl.DataFrame({"BilledCost": [1.0]}) + source = pl.DataFrame({"prompt_tokens": [5]}) + out = _merge_token_tags(normalized, source) + assert "Tags" not in out.columns + + +def test_merge_token_tags_noop_without_token_columns(): + normalized = pl.DataFrame({"Tags": [json.dumps({"team_id": "t1"})]}) + source = pl.DataFrame({"spend": [1.0]}) + tags = json.loads(_merge_token_tags(normalized, source)["Tags"][0]) + assert tags == {"team_id": "t1"} + + +def test_merge_token_tags_warns_without_tags_column(caplog): + import logging + + from litellm._logging import verbose_logger + + normalized = pl.DataFrame({"BilledCost": [1.0]}) + source = pl.DataFrame({"prompt_tokens": [5]}) + with caplog.at_level(logging.WARNING, logger=verbose_logger.name): + out = _merge_token_tags(normalized, source) + assert "Tags" not in out.columns + assert "no `Tags` column" in caplog.text + + +def test_merge_token_tags_warns_on_row_mismatch(caplog): + import logging + + from litellm._logging import verbose_logger + + normalized = pl.DataFrame({"Tags": [json.dumps({"team_id": "t1"})]}) + source = pl.DataFrame({"prompt_tokens": [1, 2]}) + with caplog.at_level(logging.WARNING, logger=verbose_logger.name): + out = _merge_token_tags(normalized, source) + assert out["Tags"][0] == normalized["Tags"][0] + assert "row-count mismatch" in caplog.text + + +def _fake_engine(*, transformed: pl.DataFrame, payload: bytes) -> MagicMock: + engine = MagicMock() + engine._transformer.transform = MagicMock(return_value=transformed) + engine._serializer.serialize = MagicMock(return_value=payload) + engine._destination.deliver = AsyncMock() + engine._build_filename = MagicMock(return_value="usage.csv") + engine._deliver_day = lambda data, window: TernaryExportEngine._deliver_day(engine, data, window) + return engine + + +@pytest.mark.asyncio +async def test_transform_enrich_deliver_enriches_tags_then_delivers(): + data = pl.DataFrame({"date": ["2026-09-03"], "Tags": [json.dumps({"team_id": "t1"})], "prompt_tokens": [7]}) + transformed = pl.DataFrame({"Tags": [json.dumps({"team_id": "t1"})]}) + engine = _fake_engine(transformed=transformed, payload=b"csv-bytes") + + await TernaryExportEngine._deliver_enriched(engine, data, _window()) + + serialized_frame = engine._serializer.serialize.call_args.args[0] + assert json.loads(serialized_frame["Tags"][0])["prompt_tokens"] == "7" + engine._destination.deliver.assert_awaited_once() + kwargs = engine._destination.deliver.await_args.kwargs + assert kwargs["content"] == b"csv-bytes" + assert kwargs["filename"] == "usage.csv" + assert kwargs["time_window"].start_time.isoformat() == "2026-09-03T00:00:00+00:00" + + +@pytest.mark.asyncio +async def test_transform_enrich_deliver_skips_empty_data(): + engine = _fake_engine(transformed=pl.DataFrame(), payload=b"") + await TernaryExportEngine._deliver_enriched(engine, pl.DataFrame(), _window()) + engine._destination.deliver.assert_not_awaited() + + +def _real_engine() -> MagicMock: + from litellm.integrations.focus.serializers import FocusCsvSerializer + from litellm.integrations.focus.transformer import FocusTransformer + + engine = MagicMock() + engine._transformer = FocusTransformer() + engine._serializer = FocusCsvSerializer() + engine._destination.deliver = AsyncMock() + engine._build_filename = MagicMock(side_effect=lambda w: f"usage_{w.start_time:%Y%m%d}.csv") + engine._deliver_day = lambda data, window: TernaryExportEngine._deliver_day(engine, data, window) + return engine + + +def _source_rows(dates: list[str]) -> pl.DataFrame: + n = len(dates) + return pl.DataFrame( + { + "date": dates, + "spend": [float(i + 1) for i in range(n)], + "api_key": [f"k{i}" for i in range(n)], + "api_key_alias": [f"a{i}" for i in range(n)], + "model": [f"m{i}" for i in range(n)], + "model_group": [f"g{i}" for i in range(n)], + "custom_llm_provider": ["anthropic"] * n, + "team_id": [f"t{i}" for i in range(n)], + "team_alias": [f"T{i}" for i in range(n)], + "api_requests": [1] * n, + "prompt_tokens": [10 * (i + 1) for i in range(n)], + "completion_tokens": [i for i in range(n)], + "cache_read_input_tokens": [0] * n, + "cache_creation_input_tokens": [0] * n, + } + ) + + +@pytest.mark.asyncio +async def test_delivers_one_upload_per_utc_day_through_the_real_transformer(): + engine = _real_engine() + data = _source_rows(["2026-10-06", "2026-10-05", "2026-10-06"]) + + await TernaryExportEngine._deliver_enriched(engine, data, _window()) + + calls = engine._destination.deliver.await_args_list + assert [c.kwargs["time_window"].start_time.date().isoformat() for c in calls] == ["2026-10-05", "2026-10-06"] + for call, expected_rows, expected_day in ((calls[0], 1, "2026-10-05"), (calls[1], 2, "2026-10-06")): + window = call.kwargs["time_window"] + assert window.end_time - window.start_time == timedelta(days=1) + assert window.frequency == "daily" + rows = list(csv.DictReader(io.StringIO(call.kwargs["content"].decode("utf-8")))) + assert len(rows) == expected_rows + assert {row["ChargePeriodStart"] for row in rows} == {f"{expected_day}T00:00:00Z"} + assert all(json.loads(row["Tags"])["prompt_tokens"] for row in rows) + + +@pytest.mark.asyncio +async def test_a_failed_day_does_not_block_the_remaining_days(): + engine = _real_engine() + data = _source_rows(["2026-10-04", "2026-10-05", "2026-10-06"]) + attempted: list[str] = [] + + async def deliver(*, content, time_window, filename): + day = time_window.start_time.date().isoformat() + attempted.append(day) + if day == "2026-10-05": + raise RuntimeError("upload-url request for 2026-10-05 failed (500)") + + engine._destination.deliver = deliver + + with pytest.raises(RuntimeError, match=r"1 day\(s\) failed to upload: 2026-10-05"): + await TernaryExportEngine._deliver_enriched(engine, data, _window()) + + assert attempted == ["2026-10-04", "2026-10-05", "2026-10-06"] + + +def test_partition_by_day_sorts_canonicalizes_and_drops_undated_rows(caplog): + import logging + from datetime import date + + from litellm._logging import verbose_logger + + data = pl.DataFrame( + {"date": ["2026-10-06", None, "2026-10-05", "not-a-date", "2026-10-06"], "spend": [1.0, 2.0, 3.0, 4.0, 5.0]} + ) + with caplog.at_level(logging.WARNING, logger=verbose_logger.name): + parts = _partition_by_day(data) + + assert [(day, part["spend"].to_list()) for day, part in parts] == [ + ("2026-10-05", [3.0]), + ("2026-10-06", [1.0, 5.0]), + ] + assert "dropped 2 row(s)" in caplog.text + assert _canonical_day(date(2026, 1, 2)) == "2026-01-02" + assert _canonical_day("2026-1-2") == "2026-01-02" + assert _canonical_day(20260102) is None + + +def test_partition_by_day_warns_and_exports_nothing_without_a_date_column(caplog): + import logging + + from litellm._logging import verbose_logger + + with caplog.at_level(logging.WARNING, logger=verbose_logger.name): + parts = _partition_by_day(pl.DataFrame({"spend": [1.0, 2.0]})) + assert parts == () + assert "no `date` column; 2 row(s) not exported" in caplog.text + + +def test_day_window_covers_exactly_one_utc_day(): + window = _day_window("2026-10-06", "interval") + assert window.start_time.isoformat() == "2026-10-06T00:00:00+00:00" + assert window.end_time.isoformat() == "2026-10-07T00:00:00+00:00" + assert window.frequency == "interval" + + +def test_pod_lock_key_is_ternary_specific(): + assert TERNARY_USAGE_DATA_JOB_NAME == "ternary_export_usage_data" + + +def test_should_reject_explicit_zero_interval_without_env_fallthrough(monkeypatch): + monkeypatch.setenv("TERNARY_EXPORT_INTERVAL_SECONDS", "300") + with pytest.raises(ValueError, match="TERNARY_EXPORT_INTERVAL_SECONDS"): + TernaryLogger( + api_key="k", connection_id="c", base_url="http://localhost:8080", frequency="interval", interval_seconds=0 + ) + + +def test_should_reject_non_numeric_interval(monkeypatch): + monkeypatch.setenv("TERNARY_EXPORT_INTERVAL_SECONDS", "5m") + with pytest.raises(ValueError, match="TERNARY_EXPORT_INTERVAL_SECONDS"): + TernaryLogger(api_key="k", connection_id="c", base_url="http://localhost:8080", frequency="interval") + + +@pytest.mark.asyncio +async def test_export_all_delivers_with_all_window(): + engine = MagicMock() + frame = pl.DataFrame({"Tags": [json.dumps({"team_id": "t"})], "prompt_tokens": [5]}) + + async def fake_get(**_kwargs): + return frame + + engine._database.get_usage_data = fake_get + + captured = {} + + async def fake_deliver(data, window): + captured["window"] = window + captured["rows"] = data.height + + engine._deliver_enriched = fake_deliver + await TernaryExportEngine.export_all(engine, limit=None) + assert captured["rows"] == 1 + assert captured["window"].frequency == "all" + + +@pytest.mark.asyncio +async def test_transform_enrich_deliver_skips_when_transform_empty(): + data = pl.DataFrame({"date": ["2026-09-03"], "Tags": [json.dumps({"team_id": "t"})], "prompt_tokens": [1]}) + engine = _fake_engine(transformed=pl.DataFrame(), payload=b"") + await TernaryExportEngine._deliver_enriched(engine, data, _window()) + engine._destination.deliver.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_transform_enrich_deliver_skips_when_payload_empty(): + data = pl.DataFrame({"date": ["2026-09-03"], "Tags": [json.dumps({"team_id": "t"})]}) + transformed = pl.DataFrame({"Tags": [json.dumps({"team_id": "t"})]}) + engine = _fake_engine(transformed=transformed, payload=b"") + await TernaryExportEngine._deliver_enriched(engine, data, _window()) + engine._destination.deliver.assert_not_awaited() + + +def test_ensure_engine_builds_ternary_export_engine(): + logger = _logger() + engine = logger._ensure_engine() + assert isinstance(engine, TernaryExportEngine) + assert isinstance(engine.destination, FocusTernaryDestination) + assert isinstance(engine.serializer, FocusCsvSerializer) + assert logger._ensure_engine() is engine + + +def test_logger_registry_and_ternary_logger_import_without_polars(): + result = run_child_interpreter( + "import sys\n" + "sys.modules['polars'] = None\n" + "from litellm.litellm_core_utils.custom_logger_registry import CustomLoggerRegistry\n" + "from litellm.integrations.ternary.ternary_logger import TernaryLogger\n" + "assert CustomLoggerRegistry.CALLBACK_CLASS_STR_TO_CLASS_TYPE['ternary'] is TernaryLogger\n" + "TernaryLogger(api_key='k', connection_id='c', base_url='https://ternary.test')\n", + timeout=120, + ) + assert result.returncode == 0, result.stderr + + +@pytest.mark.asyncio +async def test_export_window_reads_a_day_before_the_floor_with_no_upper_bound(): + from datetime import datetime, timezone + + engine = MagicMock() + captured = {} + + async def fake_get_usage_data(**kwargs): + captured.update(kwargs) + return pl.DataFrame({"date": ["2026-09-03"], "spend": [1.0]}) + + engine._database.get_usage_data = fake_get_usage_data + engine._deliver_enriched = AsyncMock() + window = FocusTimeWindow( + start_time=datetime(2026, 9, 3, 0, 0, tzinfo=timezone.utc), + end_time=datetime(2026, 9, 5, 13, 0, tzinfo=timezone.utc), + frequency="interval", + ) + await TernaryExportEngine.export_window(engine, window=window, limit=None) + + assert captured == {"limit": None, "start_time_utc": datetime(2026, 9, 2, 0, 0, tzinfo=timezone.utc)} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("session_offset_hours", [-12, -7, 0, 9, 14]) +async def test_export_window_keeps_every_row_of_a_day_whatever_the_session_offset(session_offset_hours): + """Emulate Postgres reading naive-UTC `updated_at` in a session `session_offset_hours` east of UTC.""" + from datetime import datetime, timezone + + rows = pl.DataFrame( + { + "date": ["2026-09-02", "2026-09-03", "2026-09-03", "2026-09-05"], + "updated_at": [ + datetime(2026, 9, 2, 12, 0, tzinfo=timezone.utc), + datetime(2026, 9, 3, 1, 0, tzinfo=timezone.utc), + datetime(2026, 9, 3, 20, 0, tzinfo=timezone.utc), + datetime(2026, 9, 5, 12, 50, tzinfo=timezone.utc), + ], + "spend": [8.0, 1.0, 2.0, 4.0], + } + ) + session_offset = timedelta(hours=session_offset_hours) + + async def session_filtered(*, limit, start_time_utc, end_time_utc=None): + as_read = pl.col("updated_at") - session_offset + kept = rows.filter(as_read >= start_time_utc) + return kept if end_time_utc is None else kept.filter(as_read <= end_time_utc) + + engine = MagicMock() + engine._database.get_usage_data = session_filtered + captured = {} + + async def fake_deliver(data, _window): + captured["data"] = data + + engine._deliver_enriched = fake_deliver + window = FocusTimeWindow( + start_time=datetime(2026, 9, 3, 0, 0, tzinfo=timezone.utc), + end_time=datetime(2026, 9, 5, 13, 0, tzinfo=timezone.utc), + frequency="interval", + ) + await TernaryExportEngine.export_window(engine, window=window, limit=None) + + sent = captured["data"] + assert sent["date"].to_list() == ["2026-09-03", "2026-09-03", "2026-09-05"] + assert sent["spend"].sum() == 7.0 + + +def test_init_custom_logger_compatible_class_returns_cached_ternary_logger(monkeypatch): + monkeypatch.setenv("TERNARY_API_KEY", "k") + monkeypatch.setenv("TERNARY_CONNECTION_ID", "c") + monkeypatch.setenv("TERNARY_BASE_URL", "https://ternary.test") + from litellm.litellm_core_utils.litellm_logging import _init_custom_logger_compatible_class + + first = _init_custom_logger_compatible_class("ternary", internal_usage_cache=None, llm_router=None) + assert isinstance(first, TernaryLogger) + second = _init_custom_logger_compatible_class("ternary", internal_usage_cache=None, llm_router=None) + assert second is first + + +@pytest.mark.asyncio +async def test_a_bounded_export_window_uploads_only_the_days_it_covers(): + engine = MagicMock() + + async def fake_get_usage_data(**_kwargs): + return pl.DataFrame({"date": ["2026-09-01", "2026-09-02", "2026-09-03"], "spend": [1.0, 2.0, 3.0]}) + + engine._database.get_usage_data = fake_get_usage_data + captured = {} + + async def fake_deliver(data, _window): + captured["data"] = data + + engine._deliver_enriched = fake_deliver + window = _day_window_from(datetime(2026, 9, 1, tzinfo=timezone.utc), datetime(2026, 9, 2, tzinfo=timezone.utc)) + await TernaryExportEngine.export_window(engine, window=window, limit=None) + + assert captured["data"]["date"].to_list() == ["2026-09-01"]