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