From 78da522ceb431550415ad38d1487b34bee8aef52 Mon Sep 17 00:00:00 2001 From: Aaron Sproul Date: Fri, 4 Sep 2026 19:17:14 +0000 Subject: [PATCH 1/4] feat(focus): add ternary destination for FOCUS export Add a `ternary` FOCUS export destination (plus a `ternary` callback) so a LiteLLM proxy can push FOCUS-formatted spend to Ternary with `callbacks: ["ternary"]` -- no separate export job - FocusTernaryDestination: a pure sink modeled on the sibling vantage/mavvrik/gcs destinations; POSTs the FOCUS CSV to the Ternary cost-ingestion endpoint, chunking an oversized backfill (csv-module repartition, quoted-newline safe; oversized single row raises rather than silently dropping) with a stable X-Ternary-Upload-Id + 0-based index/total so the receiver stages-then-swaps. connection_id is validated + URL-encoded; base_url must be HTTPS (loopback exempt for local); the upload has a bounded timeout; it aborts on the first failed chunk so nothing partial lands - TernaryLogger: presets the FOCUS pipeline (CSV) and enriches the FOCUS Tags JSON with the per-request token counts the shared transformer drops (Ternary weights allocation by tokens); the enrichment is confined to this path and only adds keys. Cadence is daily (default) or interval (validation loops, an explicit interval-seconds required); the export window is day-aligned and drops days before the window start so a whole-day-replace receiver never truncates - Registration mirrors vantage/mavvrik: ternary callback literal, logging dispatch, scheduler job, and CustomLoggerRegistry entry -- no shared-code changes - Tests: mocked unit tests only (no live calls), 100% line coverage of both new modules --- litellm/__init__.py | 1 + .../focus/destinations/__init__.py | 2 + .../focus/destinations/factory.py | 16 + .../focus/destinations/ternary_destination.py | 189 +++++++++ litellm/integrations/ternary/__init__.py | 0 .../integrations/ternary/ternary_logger.py | 255 +++++++++++++ .../custom_logger_registry.py | 2 + litellm/litellm_core_utils/litellm_logging.py | 15 + litellm/proxy/proxy_server.py | 6 + .../focus/test_ternary_destination.py | 326 ++++++++++++++++ tests/unit/integrations/ternary/__init__.py | 0 .../ternary/test_ternary_logger.py | 358 ++++++++++++++++++ 12 files changed, 1170 insertions(+) create mode 100644 litellm/integrations/focus/destinations/ternary_destination.py create mode 100644 litellm/integrations/ternary/__init__.py create mode 100644 litellm/integrations/ternary/ternary_logger.py create mode 100644 tests/unit/integrations/focus/test_ternary_destination.py create mode 100644 tests/unit/integrations/ternary/__init__.py create mode 100644 tests/unit/integrations/ternary/test_ternary_logger.py diff --git a/litellm/__init__.py b/litellm/__init__.py index e1da202b9ee..59257e107eb 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -168,6 +168,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 a46b1ea2dcb..c5fe19b5bb5 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..6d930718cdc --- /dev/null +++ b/litellm/integrations/focus/destinations/ternary_destination.py @@ -0,0 +1,189 @@ +"""Ternary API destination for FOCUS export: uploads FOCUS CSV to Ternary's cost-ingestion endpoint.""" + +from __future__ import annotations + +import csv +import io +from collections.abc import Sequence +from typing import TYPE_CHECKING, Final +from urllib.parse import quote, urlparse +from uuid import uuid4 + +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_MAX_ROWS_PER_UPLOAD: Final = 100_000 +TERNARY_MAX_BYTES_PER_UPLOAD: Final = 30 * 1024 * 1024 +TERNARY_UPLOAD_TIMEOUT_SECONDS: Final = 120.0 + +_LOOPBACK_HOSTS: Final = frozenset({"localhost", "127.0.0.1", "::1"}) + + +def _encode_csv(rows: Sequence[Sequence[str]]) -> bytes: + """Serialize CSV rows to UTF-8 bytes via the csv module (quoted fields stay intact).""" + buffer: Final = io.StringIO() + csv.writer(buffer).writerows(rows) + return buffer.getvalue().encode("utf-8") + + +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") + + +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 FOCUS CSV exports to the Ternary cost-ingestion API.""" + + 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 CSV to the Ternary API, chunking an oversized backfill; aborts on the first failed chunk.""" + if not content: + verbose_logger.debug("Ternary destination: empty content, skipping upload") + return + + chunks: Final = self._split_into_chunks(content) + upload_id: Final = uuid4().hex + total: Final = len(chunks) + + client: Final = _NoRedirectHTTPHandler(timeout=TERNARY_UPLOAD_TIMEOUT_SECONDS) + try: + for index, chunk in enumerate(chunks): + await self._upload_csv( + client, + chunk, + filename if total == 1 else f"{filename}.part{index + 1}", + upload_id=upload_id, + chunk_index=index, + chunk_total=total, + ) + finally: + await client.close() + + def _split_into_chunks(self, content: bytes) -> Sequence[bytes]: + """Split CSV bytes into row/byte-limited chunks; an unsplittable oversized row is skipped with a warning.""" + newline_count: Final = content.count(b"\n") + if len(content) <= TERNARY_MAX_BYTES_PER_UPLOAD and newline_count <= TERNARY_MAX_ROWS_PER_UPLOAD: + return (content,) + + rows: Final = tuple(tuple(row) for row in csv.reader(io.StringIO(content.decode("utf-8")))) + if len(rows) <= 1: + return (content,) + header: Final = rows[0] + data_rows: Final = rows[1:] + header_bytes: Final = len(_encode_csv((header,))) + + chunks: Final[list[bytes]] = [] # mutable-ok: local accumulator, appended to and returned frozen + current: Final[list[Sequence[str]]] = [] # mutable-ok: rows buffered for the in-progress chunk + current_size = header_bytes # rebind-ok: running byte tally advanced across the loop + + for row in data_rows: + row_bytes = len(_encode_csv((header, row))) - header_bytes + if header_bytes + row_bytes > TERNARY_MAX_BYTES_PER_UPLOAD: + verbose_logger.warning( + "Ternary destination: skipping a single row of %d bytes, over the %d-byte upload limit", + row_bytes, + TERNARY_MAX_BYTES_PER_UPLOAD, + ) + continue + exceeds_rows = len(current) >= TERNARY_MAX_ROWS_PER_UPLOAD + exceeds_bytes = current_size + row_bytes > TERNARY_MAX_BYTES_PER_UPLOAD + if current and (exceeds_rows or exceeds_bytes): + chunks.append(_encode_csv((header, *current))) + current.clear() + current_size = header_bytes + current.append(row) + current_size += row_bytes + + if current: + chunks.append(_encode_csv((header, *current))) + return tuple(chunks) + + async def _upload_csv( + self, + client: AsyncHTTPHandler, + csv_bytes: bytes, + filename: str, + *, + upload_id: str, + chunk_index: int, + chunk_total: int, + ) -> None: + url: Final = f"{self.base_url}/external-cost-sources/v1/{quote(self.connection_id, safe='')}/focus" + headers: Final = { # mutable-ok: request headers handed to the HTTP client + "Authorization": f"Bearer {self.api_key}", + "X-Ternary-Upload-Id": upload_id, + "X-Ternary-Chunk-Index": str(chunk_index), + "X-Ternary-Chunk-Total": str(chunk_total), + } + + await client.post( + url, + headers=headers, + files={"csv": (filename, csv_bytes, "text/csv")}, # mutable-ok: multipart payload for the client + timeout=TERNARY_UPLOAD_TIMEOUT_SECONDS, + ) + + verbose_logger.debug( + "Ternary destination: uploaded %d bytes (%s, upload_id=%s, chunk %d/%d)", + len(csv_bytes), + filename, + upload_id, + chunk_index, + chunk_total, + ) 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..b2b2c5f1137 --- /dev/null +++ b/litellm/integrations/ternary/ternary_logger.py @@ -0,0 +1,255 @@ +"""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 datetime import 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 or 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 _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 + transformed: Final = self._transformer.transform(data) + if transformed.is_empty(): + verbose_logger.debug("Ternary export: normalized data empty for window %s", window) + 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") + 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: + """Snap the window start to the previous UTC midnight so each push carries whole days.""" + now_utc: Final = now.astimezone(timezone.utc) + start_time: Final = (now_utc - timedelta(days=1)).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 1d277995211..e39ece3cd8c 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, @@ -111,6 +112,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 e292ab7b2ec..57fe4642159 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -4844,6 +4844,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): @@ -5325,6 +5334,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 a994293a4d8..9ea8bd7fb0a 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -10969,6 +10969,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..d83c98388de --- /dev/null +++ b/tests/unit/integrations/focus/test_ternary_destination.py @@ -0,0 +1,326 @@ +"""Tests for FocusTernaryDestination behavior.""" + +from __future__ import annotations + +import csv +import io +from datetime import datetime, timedelta, timezone +from typing import Any +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +import litellm.integrations.focus.destinations.ternary_destination as td +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 + + +def _capturing_client() -> tuple[MagicMock, list[dict[str, Any]]]: + """Return a mock client plus a list that records each post() call.""" + calls: list[dict[str, Any]] = [] + + mock_response = AsyncMock() + mock_response.raise_for_status = lambda: None + + mock_client = MagicMock() + mock_client.close = AsyncMock() + + async def capture_post(url, **kwargs): + calls.append({"url": url, **kwargs}) + return mock_response + + mock_client.post = capture_post + return mock_client, calls + + +def _uploaded_part(call: dict[str, Any]) -> tuple[str, bytes]: + field = call["files"]["csv"] + return field[0], field[1] + + +def _rows(content: bytes) -> list[list[str]]: + return list(csv.reader(io.StringIO(content.decode("utf-8")))) + + +def test_should_require_api_key(): + with pytest.raises(ValueError, match="api_key"): + FocusTernaryDestination(prefix="exports", config={"connection_id": "c", "base_url": "u"}) + + +def test_should_require_connection_id(): + with pytest.raises(ValueError, match="connection_id"): + FocusTernaryDestination(prefix="exports", config={"api_key": "k", "base_url": "u"}) + + +def test_should_require_base_url(): + with pytest.raises(ValueError, match="base_url"): + FocusTernaryDestination(prefix="exports", config={"api_key": "k", "connection_id": "c"}) + + +@pytest.mark.parametrize("bad_id", ["a/b", "..", "has space", "tab\tid"]) +def test_should_reject_connection_id_that_could_reroute_the_path(bad_id): + with pytest.raises(ValueError, match="connection_id"): + FocusTernaryDestination(prefix="exports", config=_config(connection_id=bad_id)) + + +def test_should_initialize_with_valid_config(): + dest = FocusTernaryDestination(prefix="exports", config=_config()) + assert dest.api_key == "test-api-key" + assert dest.connection_id == "conn-1234" + assert dest.base_url == "https://ternary.test" + + +def test_should_use_custom_base_url_and_strip_trailing_slash(): + dest = FocusTernaryDestination(prefix="exports", config=_config(base_url="http://localhost:8080/")) + assert dest.base_url == "http://localhost:8080" + + +@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_upload_to_correct_url_with_auth_and_upload_headers(): + dest = FocusTernaryDestination(prefix="exports", config=_config()) + mock_client, calls = _capturing_client() + + with patch(MOCK_TARGET, return_value=mock_client): + await dest.deliver(content=b"header\nrow1\n", time_window=_window(), filename="usage.csv") + + assert len(calls) == 1 + call = calls[0] + assert call["url"] == "https://ternary.test/external-cost-sources/v1/conn-1234/focus" + assert call["headers"]["Authorization"] == "Bearer test-api-key" + assert call["headers"]["X-Ternary-Chunk-Index"] == "0" + assert call["headers"]["X-Ternary-Chunk-Total"] == "1" + assert call["headers"]["X-Ternary-Upload-Id"] + filename, body = _uploaded_part(call) + assert filename == "usage.csv" + assert body == b"header\nrow1\n" + assert call["files"]["csv"][2] == "text/csv" + assert call["timeout"] == TERNARY_UPLOAD_TIMEOUT_SECONDS + mock_client.close.assert_awaited_once() + + +@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")) + mock_client, calls = _capturing_client() + with patch(MOCK_TARGET, return_value=mock_client): + await dest.deliver(content=b"h\nr\n", time_window=_window(), filename="usage.csv") + assert calls[0]["url"].endswith("/external-cost-sources/v1/conn%2Bid~ok/focus") + + +@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()) + mock_client, calls = _capturing_client() + + content = b'ServiceName,Tags,x_unknown\nfoo,"{""team_id"": ""t1""}",keepme\n' + with patch(MOCK_TARGET, return_value=mock_client): + await dest.deliver(content=content, time_window=_window(), filename="usage.csv") + + _, body = _uploaded_part(calls[0]) + assert body == content + + +@pytest.mark.asyncio +async def test_should_chunk_by_row_count_with_stable_upload_id(monkeypatch): + monkeypatch.setattr(td, "TERNARY_MAX_ROWS_PER_UPLOAD", 2) + dest = FocusTernaryDestination(prefix="exports", config=_config()) + + content = b"ServiceName\n" + b"\n".join([b"x"] * 5) + b"\n" + mock_client, calls = _capturing_client() + with patch(MOCK_TARGET, return_value=mock_client): + await dest.deliver(content=content, time_window=_window(), filename="usage.csv") + + assert len(calls) == 3 + total = "3" + assert {c["headers"]["X-Ternary-Upload-Id"] for c in calls} == {calls[0]["headers"]["X-Ternary-Upload-Id"]} + for i, call in enumerate(calls): + assert call["headers"]["X-Ternary-Chunk-Index"] == str(i) + assert call["headers"]["X-Ternary-Chunk-Total"] == total + filename, body = _uploaded_part(call) + assert filename == f"usage.csv.part{i + 1}" + assert len(_rows(body)) - 1 <= 2 + + +@pytest.mark.asyncio +async def test_should_chunk_by_bytes(monkeypatch): + monkeypatch.setattr(td, "TERNARY_MAX_BYTES_PER_UPLOAD", 200) + dest = FocusTernaryDestination(prefix="exports", config=_config()) + + row = b"a" * 40 + b"," + b"b" * 40 + content = b"ServiceName,BilledCost\n" + b"\n".join([row] * 20) + b"\n" + mock_client, calls = _capturing_client() + with patch(MOCK_TARGET, return_value=mock_client): + await dest.deliver(content=content, time_window=_window(), filename="usage.csv") + + assert len(calls) > 1 + for call in calls: + _, body = _uploaded_part(call) + assert len(body) <= 200 + + +@pytest.mark.asyncio +async def test_should_not_mangle_a_quoted_field_containing_a_newline(monkeypatch): + monkeypatch.setattr(td, "TERNARY_MAX_ROWS_PER_UPLOAD", 1) + dest = FocusTernaryDestination(prefix="exports", config=_config()) + + content = b'ServiceName,Tags\nsvc1,"line1\nline2"\nsvc2,"{""k"":""v""}"\n' + mock_client, calls = _capturing_client() + with patch(MOCK_TARGET, return_value=mock_client): + await dest.deliver(content=content, time_window=_window(), filename="usage.csv") + + assert len(calls) == 2 + _, first = _uploaded_part(calls[0]) + rows = _rows(first) + assert rows[0] == ["ServiceName", "Tags"] + assert rows[1] == ["svc1", "line1\nline2"] + + +@pytest.mark.asyncio +async def test_should_skip_an_unsplittable_oversized_row_and_upload_the_rest(monkeypatch): + # An unsplittable row over the byte limit is dropped (warned); the rest of the day still lands. + monkeypatch.setattr(td, "TERNARY_MAX_BYTES_PER_UPLOAD", 80) + dest = FocusTernaryDestination(prefix="exports", config=_config()) + content = b"ServiceName\nkeep-me\n" + b"z" * 300 + b"\n" + mock_client, calls = _capturing_client() + with patch(MOCK_TARGET, return_value=mock_client): + await dest.deliver(content=content, time_window=_window(), filename="usage.csv") + body = b"".join(_uploaded_part(call)[1] for call in calls) + assert b"keep-me" in body + assert b"zzz" not in body + + +@pytest.mark.asyncio +async def test_should_not_create_a_client_when_chunking_fails(monkeypatch): + # Chunking runs before the client is created, so a failure there leaks no client. + monkeypatch.setattr(td, "TERNARY_MAX_BYTES_PER_UPLOAD", 10) + dest = FocusTernaryDestination(prefix="exports", config=_config()) + content = b"ServiceName\n" + b"\xff\xfe" * 20 + b"\n" + with patch(MOCK_TARGET, side_effect=AssertionError("client created before chunking succeeded")): + with pytest.raises(UnicodeDecodeError): + await dest.deliver(content=content, time_window=_window(), filename="usage.csv") + + +@pytest.mark.asyncio +async def test_should_abort_on_first_chunk_failure(monkeypatch): + monkeypatch.setattr(td, "TERNARY_MAX_ROWS_PER_UPLOAD", 1) + dest = FocusTernaryDestination(prefix="exports", config=_config()) + + content = b"ServiceName\n" + b"\n".join([b"x"] * 3) + b"\n" + + attempts: list[str] = [] + mock_client = MagicMock() + mock_client.close = AsyncMock() + + async def failing_post(url, **kwargs): + attempts.append(kwargs["headers"]["X-Ternary-Chunk-Index"]) + raise RuntimeError("boom") + + mock_client.post = failing_post + + with patch(MOCK_TARGET, return_value=mock_client): + with pytest.raises(RuntimeError, match="boom"): + await dest.deliver(content=content, time_window=_window(), filename="usage.csv") + + assert attempts == ["0"] + + +@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("/") + + +@pytest.mark.asyncio +async def test_should_return_header_only_content_untouched(monkeypatch): + monkeypatch.setattr(td, "TERNARY_MAX_BYTES_PER_UPLOAD", 10) + dest = FocusTernaryDestination(prefix="exports", config=_config()) + content = b"HeaderOnlyLineWithNoDataRows" + mock_client, calls = _capturing_client() + with patch(MOCK_TARGET, return_value=mock_client): + await dest.deliver(content=content, time_window=_window(), filename="usage.csv") + assert len(calls) == 1 + _, body = _uploaded_part(calls[0]) + assert body == content + + +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..44f79cf33a0 --- /dev/null +++ b/tests/unit/integrations/ternary/test_ternary_logger.py @@ -0,0 +1,358 @@ +"""Tests for TernaryLogger configuration and Tags token enrichment.""" + +from __future__ import annotations + +import json +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, + _drop_days_before, + _merge_token_tags, + _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, 3, 0, 0, 0, tzinfo=timezone.utc) + assert window.end_time == now + assert window.start_time.hour == 0 and window.start_time.minute == 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 _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") + return engine + + +@pytest.mark.asyncio +async def test_transform_enrich_deliver_enriches_tags_then_delivers(): + data = pl.DataFrame({"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" + + +@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 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({"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({"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 From 284e2cfac1c84aa420c1d1d66f6cf115166e12f1 Mon Sep 17 00:00:00 2001 From: Aaron Sproul Date: Thu, 1 Oct 2026 23:17:12 +0000 Subject: [PATCH 2/4] fix(focus): warn instead of silently dropping token tags on mismatch _merge_token_tags returned the frame unchanged when the transformed data had no `Tags` column or its row count disagreed with the source, so every token count was dropped with no signal -- the receiver then lands NULLs that read as "this gateway reports no tokens". Emit a warning on each degenerate path (mirroring _drop_days_before), and cover both with caplog tests asserting the warning and the unchanged return. --- .../integrations/ternary/ternary_logger.py | 13 +++++++++- .../ternary/test_ternary_logger.py | 26 +++++++++++++++++++ 2 files changed, 38 insertions(+), 1 deletion(-) diff --git a/litellm/integrations/ternary/ternary_logger.py b/litellm/integrations/ternary/ternary_logger.py index b2b2c5f1137..f3cb8dd46ce 100644 --- a/litellm/integrations/ternary/ternary_logger.py +++ b/litellm/integrations/ternary/ternary_logger.py @@ -36,7 +36,18 @@ def _merge_token_tags(normalized: pl.DataFrame, source: pl.DataFrame) -> pl.Data """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 or normalized.height != source.height: + 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: diff --git a/tests/unit/integrations/ternary/test_ternary_logger.py b/tests/unit/integrations/ternary/test_ternary_logger.py index 44f79cf33a0..896d97945d4 100644 --- a/tests/unit/integrations/ternary/test_ternary_logger.py +++ b/tests/unit/integrations/ternary/test_ternary_logger.py @@ -250,6 +250,32 @@ def test_merge_token_tags_noop_without_token_columns(): 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) From 32bf19bbef7104ae8a1910ec48076e451a034567 Mon Sep 17 00:00:00 2001 From: Aaron Sproul Date: Wed, 7 Oct 2026 17:04:29 +0000 Subject: [PATCH 3/4] feat(focus): upload one UTC day per Ternary signed-URL PUT Replace the multipart POST with chunking by a per-day signed-URL upload, matching the Mavvrik destination's pattern. The logger partitions the raw usage rows on `date` (dropping unparseable dates with a warning), then transforms, enriches and delivers each day separately; one day's failure is logged and the remaining days still upload, with a single error raised at the end. The destination requests a signed URL for the day, validates it is an HTTPS GCS URL, and PUTs the CSV with exactly the signed headers. Removes the 30 MiB / 100k-row chunking, the X-Ternary-* upload headers and the .partN filenames; a re-sent day overwrites its object, so retries stay idempotent without them. --- .../focus/destinations/ternary_destination.py | 154 ++++----- .../integrations/ternary/ternary_logger.py | 56 +++- .../focus/test_ternary_destination.py | 300 ++++++++---------- .../ternary/test_ternary_logger.py | 128 +++++++- 4 files changed, 359 insertions(+), 279 deletions(-) diff --git a/litellm/integrations/focus/destinations/ternary_destination.py b/litellm/integrations/focus/destinations/ternary_destination.py index 6d930718cdc..b1be6d0dcc5 100644 --- a/litellm/integrations/focus/destinations/ternary_destination.py +++ b/litellm/integrations/focus/destinations/ternary_destination.py @@ -1,13 +1,11 @@ -"""Ternary API destination for FOCUS export: uploads FOCUS CSV to Ternary's cost-ingestion endpoint.""" +"""Ternary API destination for FOCUS export: uploads each day's FOCUS CSV via a Ternary-signed GCS URL.""" from __future__ import annotations -import csv -import io -from collections.abc import Sequence -from typing import TYPE_CHECKING, Final +from typing import TYPE_CHECKING, Final, Protocol from urllib.parse import quote, urlparse -from uuid import uuid4 + +from typing_extensions import NotRequired, ReadOnly, TypedDict from litellm._logging import verbose_logger from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler @@ -17,18 +15,23 @@ from .base import FocusDestination, FocusTimeWindow if TYPE_CHECKING: import httpx -TERNARY_MAX_ROWS_PER_UPLOAD: Final = 100_000 -TERNARY_MAX_BYTES_PER_UPLOAD: Final = 30 * 1024 * 1024 TERNARY_UPLOAD_TIMEOUT_SECONDS: Final = 120.0 _LOOPBACK_HOSTS: Final = frozenset({"localhost", "127.0.0.1", "::1"}) -def _encode_csv(rows: Sequence[Sequence[str]]) -> bytes: - """Serialize CSV rows to UTF-8 bytes via the csv module (quoted fields stay intact).""" - buffer: Final = io.StringIO() - csv.writer(buffer).writerows(rows) - return buffer.getvalue().encode("utf-8") +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: @@ -40,6 +43,14 @@ def _require_secure_base_url(base_url: str) -> None: 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).""" @@ -50,7 +61,7 @@ class _NoRedirectHTTPHandler(AsyncHTTPHandler): class FocusTernaryDestination(FocusDestination): - """Upload FOCUS CSV exports to the Ternary cost-ingestion API.""" + """Upload one UTC day of FOCUS CSV per delivery to Ternary through a signed GCS URL.""" def __init__( self, @@ -92,98 +103,45 @@ class FocusTernaryDestination(FocusDestination): time_window: FocusTimeWindow, filename: str, ) -> None: - """Upload CSV to the Ternary API, chunking an oversized backfill; aborts on the first failed chunk.""" + """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 - chunks: Final = self._split_into_chunks(content) - upload_id: Final = uuid4().hex - total: Final = len(chunks) - - client: Final = _NoRedirectHTTPHandler(timeout=TERNARY_UPLOAD_TIMEOUT_SECONDS) + day: Final = time_window.start_time.strftime("%Y-%m-%d") + handler: Final = _NoRedirectHTTPHandler(timeout=TERNARY_UPLOAD_TIMEOUT_SECONDS) try: - for index, chunk in enumerate(chunks): - await self._upload_csv( - client, - chunk, - filename if total == 1 else f"{filename}.part{index + 1}", - upload_id=upload_id, - chunk_index=index, - chunk_total=total, + 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 client.close() + await handler.close() - def _split_into_chunks(self, content: bytes) -> Sequence[bytes]: - """Split CSV bytes into row/byte-limited chunks; an unsplittable oversized row is skipped with a warning.""" - newline_count: Final = content.count(b"\n") - if len(content) <= TERNARY_MAX_BYTES_PER_UPLOAD and newline_count <= TERNARY_MAX_ROWS_PER_UPLOAD: - return (content,) + verbose_logger.debug("Ternary destination: uploaded %d bytes for %s (%s)", len(content), day, filename) - rows: Final = tuple(tuple(row) for row in csv.reader(io.StringIO(content.decode("utf-8")))) - if len(rows) <= 1: - return (content,) - header: Final = rows[0] - data_rows: Final = rows[1:] - header_bytes: Final = len(_encode_csv((header,))) - - chunks: Final[list[bytes]] = [] # mutable-ok: local accumulator, appended to and returned frozen - current: Final[list[Sequence[str]]] = [] # mutable-ok: rows buffered for the in-progress chunk - current_size = header_bytes # rebind-ok: running byte tally advanced across the loop - - for row in data_rows: - row_bytes = len(_encode_csv((header, row))) - header_bytes - if header_bytes + row_bytes > TERNARY_MAX_BYTES_PER_UPLOAD: - verbose_logger.warning( - "Ternary destination: skipping a single row of %d bytes, over the %d-byte upload limit", - row_bytes, - TERNARY_MAX_BYTES_PER_UPLOAD, - ) - continue - exceeds_rows = len(current) >= TERNARY_MAX_ROWS_PER_UPLOAD - exceeds_bytes = current_size + row_bytes > TERNARY_MAX_BYTES_PER_UPLOAD - if current and (exceeds_rows or exceeds_bytes): - chunks.append(_encode_csv((header, *current))) - current.clear() - current_size = header_bytes - current.append(row) - current_size += row_bytes - - if current: - chunks.append(_encode_csv((header, *current))) - return tuple(chunks) - - async def _upload_csv( - self, - client: AsyncHTTPHandler, - csv_bytes: bytes, - filename: str, - *, - upload_id: str, - chunk_index: int, - chunk_total: int, - ) -> None: - url: Final = f"{self.base_url}/external-cost-sources/v1/{quote(self.connection_id, safe='')}/focus" - headers: Final = { # mutable-ok: request headers handed to the HTTP client - "Authorization": f"Bearer {self.api_key}", - "X-Ternary-Upload-Id": upload_id, - "X-Ternary-Chunk-Index": str(chunk_index), - "X-Ternary-Chunk-Total": str(chunk_total), - } - - await client.post( - url, - headers=headers, - files={"csv": (filename, csv_bytes, "text/csv")}, # mutable-ok: multipart payload for the client + 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, ) - - verbose_logger.debug( - "Ternary destination: uploaded %d bytes (%s, upload_id=%s, chunk %d/%d)", - len(csv_bytes), - filename, - upload_id, - chunk_index, - chunk_total, - ) + 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/ternary_logger.py b/litellm/integrations/ternary/ternary_logger.py index f3cb8dd46ce..4a29eafd9ae 100644 --- a/litellm/integrations/ternary/ternary_logger.py +++ b/litellm/integrations/ternary/ternary_logger.py @@ -4,7 +4,8 @@ from __future__ import annotations import json import os -from datetime import datetime, timedelta, timezone +from collections.abc import Sequence +from datetime import date, datetime, timedelta, timezone from typing import TYPE_CHECKING, Final, TypeAlias import litellm @@ -91,6 +92,44 @@ def _drop_days_before(data: pl.DataFrame, floor: datetime) -> pl.DataFrame: 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: @@ -127,14 +166,25 @@ class _TernaryExportEngine(FocusExportEngine): 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 window %s", window) + 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") + verbose_logger.debug("Ternary export: serializer returned empty payload for %s", window.start_time.date()) return await self._destination.deliver( content=payload, diff --git a/tests/unit/integrations/focus/test_ternary_destination.py b/tests/unit/integrations/focus/test_ternary_destination.py index d83c98388de..289834ab560 100644 --- a/tests/unit/integrations/focus/test_ternary_destination.py +++ b/tests/unit/integrations/focus/test_ternary_destination.py @@ -2,15 +2,12 @@ from __future__ import annotations -import csv -import io from datetime import datetime, timedelta, timezone from typing import Any from unittest.mock import AsyncMock, MagicMock, patch import pytest -import litellm.integrations.focus.destinations.ternary_destination as td from litellm.integrations.focus.destinations.base import FocusTimeWindow from litellm.integrations.focus.destinations.ternary_destination import ( TERNARY_UPLOAD_TIMEOUT_SECONDS, @@ -37,64 +34,55 @@ def _config(**overrides: Any) -> dict[str, Any]: return base -def _capturing_client() -> tuple[MagicMock, list[dict[str, Any]]]: - """Return a mock client plus a list that records each post() call.""" - calls: list[dict[str, Any]] = [] +_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", +} - mock_response = AsyncMock() - mock_response.raise_for_status = lambda: None - mock_client = MagicMock() - mock_client.close = AsyncMock() +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): - calls.append({"url": url, **kwargs}) - return mock_response + posts.append({"url": url, **kwargs}) + return url_response if url_response is not None else _response(body=_upload_url_body()) - mock_client.post = capture_post - return mock_client, calls + async def capture_put(url, **kwargs): + puts.append({"url": url, **kwargs}) + return put_response if put_response is not None else _response() - -def _uploaded_part(call: dict[str, Any]) -> tuple[str, bytes]: - field = call["files"]["csv"] - return field[0], field[1] - - -def _rows(content: bytes) -> list[list[str]]: - return list(csv.reader(io.StringIO(content.decode("utf-8")))) - - -def test_should_require_api_key(): - with pytest.raises(ValueError, match="api_key"): - FocusTernaryDestination(prefix="exports", config={"connection_id": "c", "base_url": "u"}) - - -def test_should_require_connection_id(): - with pytest.raises(ValueError, match="connection_id"): - FocusTernaryDestination(prefix="exports", config={"api_key": "k", "base_url": "u"}) - - -def test_should_require_base_url(): - with pytest.raises(ValueError, match="base_url"): - FocusTernaryDestination(prefix="exports", config={"api_key": "k", "connection_id": "c"}) - - -@pytest.mark.parametrize("bad_id", ["a/b", "..", "has space", "tab\tid"]) -def test_should_reject_connection_id_that_could_reroute_the_path(bad_id): - with pytest.raises(ValueError, match="connection_id"): - FocusTernaryDestination(prefix="exports", config=_config(connection_id=bad_id)) - - -def test_should_initialize_with_valid_config(): - dest = FocusTernaryDestination(prefix="exports", config=_config()) - assert dest.api_key == "test-api-key" - assert dest.connection_id == "conn-1234" - assert dest.base_url == "https://ternary.test" - - -def test_should_use_custom_base_url_and_strip_trailing_slash(): - dest = FocusTernaryDestination(prefix="exports", config=_config(base_url="http://localhost:8080/")) - assert dest.base_url == "http://localhost:8080" + handler = MagicMock() + handler.close = AsyncMock() + handler.client.post = capture_post + handler.client.put = capture_put + return handler, posts, puts @pytest.mark.asyncio @@ -105,26 +93,35 @@ async def test_should_skip_empty_content(): @pytest.mark.asyncio -async def test_should_upload_to_correct_url_with_auth_and_upload_headers(): +async def test_should_request_a_signed_url_for_the_window_day_then_put_the_csv(): dest = FocusTernaryDestination(prefix="exports", config=_config()) - mock_client, calls = _capturing_client() + handler, posts, puts = _handler() - with patch(MOCK_TARGET, return_value=mock_client): + with patch(MOCK_TARGET, return_value=handler): await dest.deliver(content=b"header\nrow1\n", time_window=_window(), filename="usage.csv") - assert len(calls) == 1 - call = calls[0] - assert call["url"] == "https://ternary.test/external-cost-sources/v1/conn-1234/focus" - assert call["headers"]["Authorization"] == "Bearer test-api-key" - assert call["headers"]["X-Ternary-Chunk-Index"] == "0" - assert call["headers"]["X-Ternary-Chunk-Total"] == "1" - assert call["headers"]["X-Ternary-Upload-Id"] - filename, body = _uploaded_part(call) - assert filename == "usage.csv" - assert body == b"header\nrow1\n" - assert call["files"]["csv"][2] == "text/csv" - assert call["timeout"] == TERNARY_UPLOAD_TIMEOUT_SECONDS - mock_client.close.assert_awaited_once() + 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 @@ -142,128 +139,94 @@ async def test_no_redirect_handler_disables_redirects_including_retry_client(): @pytest.mark.asyncio async def test_should_url_encode_the_connection_id(): dest = FocusTernaryDestination(prefix="exports", config=_config(connection_id="conn+id~ok")) - mock_client, calls = _capturing_client() - with patch(MOCK_TARGET, return_value=mock_client): + 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 calls[0]["url"].endswith("/external-cost-sources/v1/conn%2Bid~ok/focus") + 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()) - mock_client, calls = _capturing_client() + handler, _, puts = _handler() content = b'ServiceName,Tags,x_unknown\nfoo,"{""team_id"": ""t1""}",keepme\n' - with patch(MOCK_TARGET, return_value=mock_client): + with patch(MOCK_TARGET, return_value=handler): await dest.deliver(content=content, time_window=_window(), filename="usage.csv") - _, body = _uploaded_part(calls[0]) - assert body == content + assert puts[0]["content"] == content @pytest.mark.asyncio -async def test_should_chunk_by_row_count_with_stable_upload_id(monkeypatch): - monkeypatch.setattr(td, "TERNARY_MAX_ROWS_PER_UPLOAD", 2) +@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()) - - content = b"ServiceName\n" + b"\n".join([b"x"] * 5) + b"\n" - mock_client, calls = _capturing_client() - with patch(MOCK_TARGET, return_value=mock_client): - await dest.deliver(content=content, time_window=_window(), filename="usage.csv") - - assert len(calls) == 3 - total = "3" - assert {c["headers"]["X-Ternary-Upload-Id"] for c in calls} == {calls[0]["headers"]["X-Ternary-Upload-Id"]} - for i, call in enumerate(calls): - assert call["headers"]["X-Ternary-Chunk-Index"] == str(i) - assert call["headers"]["X-Ternary-Chunk-Total"] == total - filename, body = _uploaded_part(call) - assert filename == f"usage.csv.part{i + 1}" - assert len(_rows(body)) - 1 <= 2 + 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 -async def test_should_chunk_by_bytes(monkeypatch): - monkeypatch.setattr(td, "TERNARY_MAX_BYTES_PER_UPLOAD", 200) +@pytest.mark.parametrize("status", [400, 403, 503]) +async def test_should_raise_when_the_signed_put_fails(status): dest = FocusTernaryDestination(prefix="exports", config=_config()) - - row = b"a" * 40 + b"," + b"b" * 40 - content = b"ServiceName,BilledCost\n" + b"\n".join([row] * 20) + b"\n" - mock_client, calls = _capturing_client() - with patch(MOCK_TARGET, return_value=mock_client): - await dest.deliver(content=content, time_window=_window(), filename="usage.csv") - - assert len(calls) > 1 - for call in calls: - _, body = _uploaded_part(call) - assert len(body) <= 200 + 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 -async def test_should_not_mangle_a_quoted_field_containing_a_newline(monkeypatch): - monkeypatch.setattr(td, "TERNARY_MAX_ROWS_PER_UPLOAD", 1) +@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()) - - content = b'ServiceName,Tags\nsvc1,"line1\nline2"\nsvc2,"{""k"":""v""}"\n' - mock_client, calls = _capturing_client() - with patch(MOCK_TARGET, return_value=mock_client): - await dest.deliver(content=content, time_window=_window(), filename="usage.csv") - - assert len(calls) == 2 - _, first = _uploaded_part(calls[0]) - rows = _rows(first) - assert rows[0] == ["ServiceName", "Tags"] - assert rows[1] == ["svc1", "line1\nline2"] + 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_skip_an_unsplittable_oversized_row_and_upload_the_rest(monkeypatch): - # An unsplittable row over the byte limit is dropped (warned); the rest of the day still lands. - monkeypatch.setattr(td, "TERNARY_MAX_BYTES_PER_UPLOAD", 80) +async def test_should_accept_a_virtual_hosted_gcs_url(): dest = FocusTernaryDestination(prefix="exports", config=_config()) - content = b"ServiceName\nkeep-me\n" + b"z" * 300 + b"\n" - mock_client, calls = _capturing_client() - with patch(MOCK_TARGET, return_value=mock_client): - await dest.deliver(content=content, time_window=_window(), filename="usage.csv") - body = b"".join(_uploaded_part(call)[1] for call in calls) - assert b"keep-me" in body - assert b"zzz" not in body + 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 -async def test_should_not_create_a_client_when_chunking_fails(monkeypatch): - # Chunking runs before the client is created, so a failure there leaks no client. - monkeypatch.setattr(td, "TERNARY_MAX_BYTES_PER_UPLOAD", 10) +@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()) - content = b"ServiceName\n" + b"\xff\xfe" * 20 + b"\n" - with patch(MOCK_TARGET, side_effect=AssertionError("client created before chunking succeeded")): - with pytest.raises(UnicodeDecodeError): - await dest.deliver(content=content, time_window=_window(), filename="usage.csv") - - -@pytest.mark.asyncio -async def test_should_abort_on_first_chunk_failure(monkeypatch): - monkeypatch.setattr(td, "TERNARY_MAX_ROWS_PER_UPLOAD", 1) - dest = FocusTernaryDestination(prefix="exports", config=_config()) - - content = b"ServiceName\n" + b"\n".join([b"x"] * 3) + b"\n" - - attempts: list[str] = [] - mock_client = MagicMock() - mock_client.close = AsyncMock() - - async def failing_post(url, **kwargs): - attempts.append(kwargs["headers"]["X-Ternary-Chunk-Index"]) - raise RuntimeError("boom") - - mock_client.post = failing_post - - with patch(MOCK_TARGET, return_value=mock_client): - with pytest.raises(RuntimeError, match="boom"): - await dest.deliver(content=content, time_window=_window(), filename="usage.csv") - - assert attempts == ["0"] + 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"]) @@ -278,19 +241,6 @@ def test_should_accept_https_or_loopback_base_url(url): assert dest.base_url == url.rstrip("/") -@pytest.mark.asyncio -async def test_should_return_header_only_content_untouched(monkeypatch): - monkeypatch.setattr(td, "TERNARY_MAX_BYTES_PER_UPLOAD", 10) - dest = FocusTernaryDestination(prefix="exports", config=_config()) - content = b"HeaderOnlyLineWithNoDataRows" - mock_client, calls = _capturing_client() - with patch(MOCK_TARGET, return_value=mock_client): - await dest.deliver(content=content, time_window=_window(), filename="usage.csv") - assert len(calls) == 1 - _, body = _uploaded_part(calls[0]) - assert body == content - - def test_factory_creates_ternary_destination_from_config(): from litellm.integrations.focus.destinations.factory import FocusDestinationFactory diff --git a/tests/unit/integrations/ternary/test_ternary_logger.py b/tests/unit/integrations/ternary/test_ternary_logger.py index 896d97945d4..2e76490f7ec 100644 --- a/tests/unit/integrations/ternary/test_ternary_logger.py +++ b/tests/unit/integrations/ternary/test_ternary_logger.py @@ -2,7 +2,10 @@ from __future__ import annotations +import csv +import io import json +from datetime import timedelta from unittest.mock import AsyncMock, MagicMock import polars as pl @@ -13,8 +16,11 @@ 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, ) @@ -282,12 +288,13 @@ def _fake_engine(*, transformed: pl.DataFrame, payload: bytes) -> MagicMock: 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({"Tags": [json.dumps({"team_id": "t1"})], "prompt_tokens": [7]}) + 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") @@ -299,6 +306,7 @@ async def test_transform_enrich_deliver_enriches_tags_then_delivers(): 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 @@ -308,6 +316,120 @@ async def test_transform_enrich_deliver_skips_empty_data(): 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" @@ -350,7 +472,7 @@ async def test_export_all_delivers_with_all_window(): @pytest.mark.asyncio async def test_transform_enrich_deliver_skips_when_transform_empty(): - data = pl.DataFrame({"Tags": [json.dumps({"team_id": "t"})], "prompt_tokens": [1]}) + 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() @@ -358,7 +480,7 @@ async def test_transform_enrich_deliver_skips_when_transform_empty(): @pytest.mark.asyncio async def test_transform_enrich_deliver_skips_when_payload_empty(): - data = pl.DataFrame({"Tags": [json.dumps({"team_id": "t"})]}) + 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()) From b9134f2d738d2105bd98ad7f08059adbdf0dad2e Mon Sep 17 00:00:00 2001 From: Aaron Sproul Date: Wed, 7 Oct 2026 19:03:25 +0000 Subject: [PATCH 4/4] fix(focus): re-send the last two whole days so late-flushed spend lands LiteLLM writes DailyUserSpend in batches, so spend dated yesterday can be flushed after the 00:05 UTC daily export. That export only covered the previous UTC midnight onward, and the next day's export moved past it, so the tail of each day was never sent. Start the window at UTC midnight two days back: the following run re-sends yesterday whole, overwriting its per-day object with the complete day. --- .../integrations/ternary/ternary_logger.py | 4 +- .../ternary/test_ternary_logger.py | 41 ++++++++++++++++++- 2 files changed, 41 insertions(+), 4 deletions(-) diff --git a/litellm/integrations/ternary/ternary_logger.py b/litellm/integrations/ternary/ternary_logger.py index 4a29eafd9ae..4d9037636df 100644 --- a/litellm/integrations/ternary/ternary_logger.py +++ b/litellm/integrations/ternary/ternary_logger.py @@ -256,9 +256,9 @@ class TernaryLogger(FocusLogger): ) def _compute_time_window(self, now: datetime) -> FocusTimeWindow: - """Snap the window start to the previous UTC midnight so each push carries whole days.""" + """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=1)).replace(hour=0, minute=0, second=0, microsecond=0) + 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: diff --git a/tests/unit/integrations/ternary/test_ternary_logger.py b/tests/unit/integrations/ternary/test_ternary_logger.py index 2e76490f7ec..5a96e5bd679 100644 --- a/tests/unit/integrations/ternary/test_ternary_logger.py +++ b/tests/unit/integrations/ternary/test_ternary_logger.py @@ -94,9 +94,46 @@ def test_compute_time_window_is_day_aligned_and_carries_whole_days(): 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, 3, 0, 0, 0, tzinfo=timezone.utc) + assert window.start_time == datetime(2026, 9, 2, 0, 0, 0, tzinfo=timezone.utc) assert window.end_time == now - assert window.start_time.hour == 0 and window.start_time.minute == 0 + + +@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():