diff --git a/litellm/__init__.py b/litellm/__init__.py index 3f93166ecc8..401b513172d 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -169,6 +169,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/factory.py b/litellm/integrations/focus/destinations/factory.py index 0ecf04326b3..e5b4ac6b833 100644 --- a/litellm/integrations/focus/destinations/factory.py +++ b/litellm/integrations/focus/destinations/factory.py @@ -9,6 +9,7 @@ from .base import FocusDestination 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 @@ -33,6 +34,8 @@ class FocusDestinationFactory: return FocusGCSDestination(prefix=prefix, config=normalized_config) if provider_lower == "mavvrik": return FocusMavvrikDestination(prefix=prefix, config=normalized_config) + if provider_lower == "ternary": + return FocusTernaryDestination(prefix=prefix, config=normalized_config) raise NotImplementedError(f"Provider '{provider}' not supported for Focus export") @staticmethod @@ -80,4 +83,17 @@ class FocusDestinationFactory: "connection_id": overrides.get("connection_id") or os.getenv("MAVVRIK_CONNECTION_ID"), } return {k: v for k, v in resolved.items() if v is not None} + if provider == "ternary": + resolved = { # mutable-ok: read-only config dict, returned filtered below + "api_key": overrides.get("api_key") or os.getenv("TERNARY_API_KEY"), + "connection_id": overrides.get("connection_id") or os.getenv("TERNARY_CONNECTION_ID"), + "base_url": overrides.get("base_url") or os.getenv("TERNARY_BASE_URL"), + } + if not resolved.get("api_key"): + raise ValueError("TERNARY_API_KEY must be provided for Ternary exports") + if not resolved.get("connection_id"): + raise ValueError("TERNARY_CONNECTION_ID must be provided for Ternary exports") + if not resolved.get("base_url"): + raise ValueError("TERNARY_BASE_URL must be provided for Ternary exports") + return {k: v for k, v in resolved.items() if v is not None} # mutable-ok: drops None-valued config keys raise NotImplementedError(f"Provider '{provider}' not supported for Focus export configuration") diff --git a/litellm/integrations/focus/destinations/ternary_destination.py b/litellm/integrations/focus/destinations/ternary_destination.py new file mode 100644 index 00000000000..b1be6d0dcc5 --- /dev/null +++ b/litellm/integrations/focus/destinations/ternary_destination.py @@ -0,0 +1,147 @@ +"""Ternary API destination for FOCUS export: uploads each day's FOCUS CSV via a Ternary-signed GCS URL.""" + +from __future__ import annotations + +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[dict[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) -> 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 {} # mutable-ok: read-only local; empty fallback for absent config + 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}"}, # mutable-ok: request headers for the client + json={"day": day}, # mutable-ok: JSON request body for the client + 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/ternary_logger.py b/litellm/integrations/ternary/ternary_logger.py new file mode 100644 index 00000000000..4d9037636df --- /dev/null +++ b/litellm/integrations/ternary/ternary_logger.py @@ -0,0 +1,316 @@ +"""Ternary FocusLogger preset (CSV, Ternary destination) that adds per-row token counts to the FOCUS Tags.""" + +from __future__ import annotations + +import json +import os +from collections.abc import Sequence +from datetime import date, datetime, timedelta, timezone +from typing import TYPE_CHECKING, Final, TypeAlias + +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.export_engine import FocusExportEngine +from litellm.integrations.focus.focus_logger import FocusLogger + +if TYPE_CHECKING: + import polars as pl + from apscheduler.schedulers.asyncio import AsyncIOScheduler +else: + AsyncIOScheduler: TypeAlias = object + +TERNARY_USAGE_DATA_JOB_NAME: Final = "ternary_export_usage_data" + +_SUPPORTED_FREQUENCIES: Final = frozenset({"daily", "interval"}) + +_TOKEN_TAG_KEYS: Final = ( + "prompt_tokens", + "completion_tokens", + "cache_read_input_tokens", + "cache_creation_input_tokens", +) + + +def _merge_token_tags(normalized: pl.DataFrame, source: pl.DataFrame) -> pl.DataFrame: + """Zip transformed rows with the row-aligned source and add token keys to each row's Tags; degrades to unchanged on any mismatch.""" + import polars as pl + + 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 not token_cols: + return normalized + + token_rows: Final = source.select(token_cols).to_dicts() + tags_values: Final = normalized["Tags"].to_list() + + merged: Final[list[str]] = [] # mutable-ok: per-row accumulator built in the zip loop + for tags_json, tokens in zip(tags_values, token_rows): + try: + parsed = json.loads(tags_json) if tags_json else {} # mutable-ok: fresh per-row tag dict + if not isinstance(parsed, dict): + parsed = {} # mutable-ok: non-object Tags degrades to empty + except (json.JSONDecodeError, TypeError): + parsed = {} # mutable-ok: malformed Tags degrades to empty + for key, value in tokens.items(): + if value is not None: + parsed[key] = str(value) + merged.append(json.dumps(parsed)) + + return normalized.with_columns(pl.Series("Tags", merged)) + + +def _drop_days_before(data: pl.DataFrame, floor: datetime) -> pl.DataFrame: + """Drop source rows for days older than the window start, so whole-day replace-by-day never truncates a partial older day.""" + import polars as pl + + if "date" not in data.columns: + return data + floor_date: Final = floor.astimezone(timezone.utc).date() + date_col: Final = pl.col("date").cast(pl.Utf8) # cast-ok: polars dtype cast, not typing.cast + parsed: Final = date_col.str.strptime(pl.Date, "%Y-%m-%d", strict=False) + kept: Final = data.filter(parsed >= floor_date) + if data.height > 0 and kept.height == 0: + verbose_logger.warning( + "Ternary export: day-window floor %s dropped all %d rows (unparseable or older `date`?)", + floor_date, + data.height, + ) + return kept + + +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 _partition_by_day(data: pl.DataFrame) -> Sequence[tuple[str, pl.DataFrame]]: + """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 () + days: Final[list[tuple[str, pl.DataFrame]]] = [] # mutable-ok: per-day parts, sorted and returned as a tuple + undated = 0 # rebind-ok: running count of rows dropped for an unusable `date` + for key, part in data.partition_by("date", as_dict=True, maintain_order=True).items(): + partition_key: object = key[0] # pyright: ignore[reportAny] # polars partition keys are stubbed as Any + day = _canonical_day(partition_key) + if day is None: + undated += part.height + continue + days.append((day, part)) + if undated: + verbose_logger.warning("Ternary export: dropped %d row(s) with a missing or unparseable `date`", undated) + return tuple(sorted(days, 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) + + +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 _TernaryExportEngine(FocusExportEngine): + """FocusExportEngine that drops pre-window days and merges token counts into Tags before delivery.""" + + async def export_window(self, *, window: FocusTimeWindow, limit: int | None) -> None: + data: Final = await self._database.get_usage_data( + limit=limit, + start_time_utc=window.start_time, + end_time_utc=window.end_time, + ) + await self._deliver_enriched(_drop_days_before(data, window.start_time), window) + + async def export_all(self, *, limit: int | None) -> None: + data: Final = await self._database.get_usage_data(limit=limit) + 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: pl.DataFrame, 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: pl.DataFrame, window: FocusTimeWindow) -> None: + transformed: Final = self._transformer.transform(data) + if transformed.is_empty(): + verbose_logger.debug("Ternary export: normalized data empty for %s", window.start_time.date()) + return + enriched: Final = _merge_token_tags(transformed, data) + payload: Final = self._serializer.serialize(enriched) + if not payload: + verbose_logger.debug("Ternary export: serializer returned empty payload for %s", window.start_time.date()) + return + await self._destination.deliver( + content=payload, + time_window=window, + filename=self._build_filename(window), + ) + + +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: + """Use the Ternary export engine (day-window filtering + Tags token enrichment).""" + if self._engine is None: + self._engine = _TernaryExportEngine( + provider=self.provider, + export_format=self.export_format, + prefix=self.prefix, + destination_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: AsyncIOScheduler, + ) -> 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 715e492b159..0a239a92059 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -4939,6 +4939,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): @@ -5444,6 +5453,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 e035f061db9..7743c076ca4 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -11351,6 +11351,12 @@ 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) + ######################################################## # Prometheus Background Job ######################################################## 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..289834ab560 --- /dev/null +++ b/tests/unit/integrations/focus/test_ternary_destination.py @@ -0,0 +1,276 @@ +"""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_factory_creates_ternary_destination_from_config(): + from litellm.integrations.focus.destinations.factory import FocusDestinationFactory + + dest = FocusDestinationFactory.create( + provider="ternary", + prefix="exports", + config={"api_key": "k", "connection_id": "c", "base_url": "https://ternary.test"}, + ) + assert isinstance(dest, FocusTernaryDestination) + assert dest.connection_id == "c" + + +def test_factory_resolves_ternary_config_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.focus.destinations.factory import FocusDestinationFactory + + dest = FocusDestinationFactory.create(provider="ternary", prefix="exports") + assert isinstance(dest, FocusTernaryDestination) + assert dest.api_key == "envk" + assert dest.connection_id == "envc" + + +@pytest.mark.parametrize("missing", ["TERNARY_API_KEY", "TERNARY_CONNECTION_ID", "TERNARY_BASE_URL"]) +def test_factory_requires_each_ternary_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.focus.destinations.factory import FocusDestinationFactory + + with pytest.raises(ValueError, match=missing): + FocusDestinationFactory.create(provider="ternary", prefix="exports") 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..5a96e5bd679 --- /dev/null +++ b/tests/unit/integrations/ternary/test_ternary_logger.py @@ -0,0 +1,543 @@ +"""Tests for TernaryLogger configuration and Tags token enrichment.""" + +from __future__ import annotations + +import csv +import io +import json +from datetime import timedelta +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.ternary.ternary_logger import ( + TERNARY_USAGE_DATA_JOB_NAME, + TernaryLogger, + _canonical_day, + _day_window, + _drop_days_before, + _merge_token_tags, + _partition_by_day, + _TernaryExportEngine, +) + + +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): + return rows.filter((pl.col("updated_at") >= start_time_utc) & (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 test_drop_days_before_removes_older_days(): + from datetime import datetime, timezone + + data = pl.DataFrame({"date": ["2026-09-02", "2026-09-03", "2026-09-04"], "spend": [1.0, 2.0, 3.0]}) + floor = datetime(2026, 9, 3, 0, 0, tzinfo=timezone.utc) + out = _drop_days_before(data, floor) + assert out["date"].to_list() == ["2026-09-03", "2026-09-04"] + + +def test_drop_days_before_noop_without_date_column(): + from datetime import datetime, timezone + + data = pl.DataFrame({"spend": [1.0]}) + out = _drop_days_before(data, datetime(2026, 9, 3, tzinfo=timezone.utc)) + assert out.equals(data) + + +@pytest.mark.asyncio +async def test_export_window_interposes_drop_days_before(): + 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_drop_days_before_warns_when_everything_dropped(): + from datetime import datetime, timezone + + data = pl.DataFrame({"date": ["2026-09-01", "2026-09-02"], "spend": [1.0, 2.0]}) + out = _drop_days_before(data, datetime(2026, 9, 5, tzinfo=timezone.utc)) + assert out.height == 0 + + +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 logger._ensure_engine() is engine + + +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