diff --git a/litellm/__init__.py b/litellm/__init__.py index 42c0ea881fda..34611831edbc 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -162,6 +162,7 @@ def _dev_env_hot_reload_enabled() -> bool: "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 932184cf485b..05baf5d06b41 100644 --- a/litellm/integrations/focus/destinations/__init__.py +++ b/litellm/integrations/focus/destinations/__init__.py @@ -5,6 +5,7 @@ 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 @@ "FocusGCSDestination", "FocusMavvrikDestination", "FocusS3Destination", + "FocusTernaryDestination", "FocusTimeWindow", "FocusVantageDestination", ] diff --git a/litellm/integrations/focus/destinations/factory.py b/litellm/integrations/focus/destinations/factory.py index a46b1ea2dcb4..394472e0ef2b 100644 --- a/litellm/integrations/focus/destinations/factory.py +++ b/litellm/integrations/focus/destinations/factory.py @@ -9,6 +9,7 @@ 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 @@ def create( 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 @@ def _resolve_config( "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 = { + "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} 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 000000000000..c7595d39d851 --- /dev/null +++ b/litellm/integrations/focus/destinations/ternary_destination.py @@ -0,0 +1,203 @@ +"""Ternary API destination for Focus export. + +Uploads FOCUS CSV exports to a Ternary cost-ingestion endpoint so LiteLLM +spend can be allocated in Ternary. A thin sink over the shared FOCUS +transformer/serializer: it adds no columns and forwards the FOCUS CSV as-is. + +A large backfill may exceed the per-request limits and be split into several +chunks. Every chunk of one export carries a stable ``X-Ternary-Upload-Id`` plus +its 0-based index and total count, so the receiver stages the parts and swaps +the whole export window once, after all parts arrive. +""" + +from __future__ import annotations + +import csv +import io +from collections.abc import Sequence +from typing import 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, + get_async_httpx_client, + httpxSpecialProvider, +) + +from .base import FocusDestination, FocusTimeWindow + +# Chunk oversized backfills; the byte cap stays under the receiver's 32 MB limit. +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 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 content to the Ternary API, chunking an oversized backfill. + + Aborts on the first failed chunk (raising) rather than continuing: under + the receiver's stage-then-swap an incomplete upload is never committed, so + stopping early leaves nothing partial landed and the next scheduled run + retries the whole window. All chunks share one upload id. + """ + if not content: + verbose_logger.debug("Ternary destination: empty content, skipping upload") + return + + client: Final = get_async_httpx_client( + llm_provider=httpxSpecialProvider.LoggingCallback, + ) + + chunks: Final = self._split_into_chunks(content) + upload_id: Final = uuid4().hex + total: Final = len(chunks) + + 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, + ) + + def _split_into_chunks(self, content: bytes) -> Sequence[bytes]: + """Split CSV bytes into chunks within the row and byte limits. + + A steady-state export fits in one chunk and is returned untouched (no + re-encoding). Only an oversized backfill is parsed and repartitioned; + parsing goes through the csv module so a quoted field containing a + newline is not mis-split. A single row larger than the byte limit cannot + be split and is a hard error rather than silently dropped. + """ + 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: + raise ValueError( + f"Ternary destination: a single CSV row is {row_bytes} bytes, exceeding the " + f"{TERNARY_MAX_BYTES_PER_UPLOAD}-byte upload limit and cannot be split" + ) + 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() # reset in place after flushing (no rebinding) + current_size = header_bytes # rebind-ok: reset after flushing a chunk + current.append(row) + current_size += row_bytes + + if current: + chunks.append(_encode_csv((header, *current))) + return tuple(chunks) or (content,) + + 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 000000000000..e69de29bb2d1 diff --git a/litellm/integrations/ternary/ternary_logger.py b/litellm/integrations/ternary/ternary_logger.py new file mode 100644 index 000000000000..b4cc9b42e94e --- /dev/null +++ b/litellm/integrations/ternary/ternary_logger.py @@ -0,0 +1,292 @@ +"""Ternary logger — thin wrapper around the Focus export pipeline. + +Configures FocusLogger to use the Ternary API destination with CSV format so +users can simply set ``callbacks: ["ternary"]`` in their proxy config. + +Beyond presetting the destination, this +logger enriches the FOCUS ``Tags`` column with per-row token counts that the +shared FocusTransformer drops. Ternary weights cost allocation by token +consumption, so the token breakdown must survive the export. The enrichment is +confined to this Ternary-only code path and only *adds* keys to ``Tags`` +(FOCUS v1.2's escape hatch for non-standard fields) — it never modifies the +shared transformer, so other FOCUS exports are unaffected. +""" + +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.focus_logger import FocusLogger + +if TYPE_CHECKING: + import polars as pl + from apscheduler.schedulers.asyncio import AsyncIOScheduler + + from litellm.integrations.focus.export_engine import FocusExportEngine +else: + AsyncIOScheduler: TypeAlias = object + +TERNARY_USAGE_DATA_JOB_NAME: Final = "ternary_export_usage_data" + +# No "hourly": spend is a daily aggregate under whole-day replace, so sub-daily adds no grain. +_SUPPORTED_FREQUENCIES: Final = frozenset({"daily", "interval"}) + +# Raw LiteLLM daily-spend token columns carried in the FOCUS Tags JSON. +_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: + """Merge raw token counts from the source rows into the FOCUS Tags JSON. + + The transformer emits rows 1:1 and row-aligned with the source, so we zip + by position and add the token keys to each row's existing Tags dict. + Degrades gracefully: on a row-count mismatch, missing token columns, or a + malformed Tags value, the data is left unchanged rather than failing the + export. + """ + import polars as pl # local import: polars is a heavy optional dependency + + 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. + + The receiver replaces cost by whole UTC day (``ChargePeriodStart``), but + ``get_usage_data`` filters by ``updated_at`` -- so a row for an *older* day that + was merely touched inside the window would partially replace that day and + truncate it. Keeping only days at or after the window start means every day we + send is sent in full, so replace-by-day never truncates. (A backlog draining for + older days is consequently not delivered by the scheduled path; that needs a + wider/backfill window, by design.) + """ + import polars as pl # local import: polars is a heavy optional dependency + + 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 TernaryLogger(FocusLogger): + """FocusLogger pre-configured for Ternary (CSV format, Ternary cost-import API). + + Environment Variables: + TERNARY_API_KEY: per-connection shared secret for the Ternary receiver + TERNARY_CONNECTION_ID: external-cost-source connection id (used in the URL path) + TERNARY_BASE_URL: required — the Ternary API host (region-specific; no default) + TERNARY_EXPORT_FREQUENCY: export cadence — "daily" (default) or "interval" + (short setup-validation / test loops only). "hourly" is unsupported: LiteLLM + spend is a daily aggregate landed with whole-day replace, so it adds no grain. + TERNARY_EXPORT_INTERVAL_SECONDS: interval in seconds when frequency is "interval" + """ + + 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 + complete days -- the receiver replaces landed cost by whole day.""" + 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) + + async def _export_window(self, *, window: FocusTimeWindow, limit: int | None) -> None: + engine: Final = self._ensure_engine() + data: Final = await engine._database.get_usage_data( + limit=limit, + start_time_utc=window.start_time, + end_time_utc=window.end_time, + ) + windowed: Final = _drop_days_before(data, window.start_time) + await self._transform_enrich_deliver(engine=engine, data=windowed, window=window) + + async def _export_all(self, *, limit: int | None) -> None: + engine: Final = self._ensure_engine() + data: Final = await engine._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._transform_enrich_deliver(engine=engine, data=data, window=window) + + async def _transform_enrich_deliver( + self, + *, + engine: FocusExportEngine, + 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 = engine._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 = engine._serializer.serialize(enriched) + if not payload: + verbose_logger.debug("Ternary export: serializer returned empty payload") + return + await engine._destination.deliver( + content=payload, + time_window=window, + filename=engine._build_filename(window), + ) + + async def initialize_focus_export_job(self) -> None: + """Override to use a Ternary-specific pod lock key. + + Without this, TernaryLogger and FocusLogger would compete for the same + ``FOCUS_USAGE_DATA_JOB_NAME`` lock, causing one to silently skip its + export cycle when both are configured simultaneously. + """ + 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 6449aa4d46ee..a86b41d6c8d1 100644 --- a/litellm/litellm_core_utils/custom_logger_registry.py +++ b/litellm/litellm_core_utils/custom_logger_registry.py @@ -47,6 +47,7 @@ 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, @@ -106,6 +107,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 17a19f05fa32..343163638cdf 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -4495,6 +4495,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): @@ -4892,6 +4901,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 5a39b8c610a1..164e0d1e4fa9 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -9780,6 +9780,15 @@ async def _initialize_spend_tracking_background_jobs(cls, scheduler: AsyncIOSche await MavvrikFocusLogger.init_mavvrik_focus_background_job(scheduler=scheduler) + ######################################################## + # Ternary Background Job + ######################################################## + 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/test_litellm/integrations/focus/test_ternary_destination.py b/tests/test_litellm/integrations/focus/test_ternary_destination.py new file mode 100644 index 000000000000..3516d176e31e --- /dev/null +++ b/tests/test_litellm/integrations/focus/test_ternary_destination.py @@ -0,0 +1,262 @@ +"""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, +) + +MOCK_TARGET = "litellm.integrations.focus.destinations.ternary_destination.get_async_httpx_client" + + +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() + + 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 + + +@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_raise_on_a_single_row_larger_than_the_byte_limit(monkeypatch): + monkeypatch.setattr(td, "TERNARY_MAX_BYTES_PER_UPLOAD", 50) + dest = FocusTernaryDestination(prefix="exports", config=_config()) + + content = b"ServiceName\n" + b"z" * 200 + b"\n" + mock_client, _ = _capturing_client() + with patch(MOCK_TARGET, return_value=mock_client): + with pytest.raises(ValueError, match="cannot be split"): + 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() + + 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 diff --git a/tests/test_litellm/integrations/ternary/test_ternary_logger.py b/tests/test_litellm/integrations/ternary/test_ternary_logger.py new file mode 100644 index 000000000000..460d6d3889c4 --- /dev/null +++ b/tests/test_litellm/integrations/ternary/test_ternary_logger.py @@ -0,0 +1,353 @@ +"""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, +) + + +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(): + # Guards the base FocusLogger's silent 60s fallback (would re-upload the whole window every minute). + 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): + # Sub-daily cadence + whole-day replace would silently drop the rest of the day. + 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") + # Start snaps to 00:00 UTC of the prior day, never a mid-day slice the day-replace receiver would apply. + 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(monkeypatch): + # The scheduled path must drop days older than the window start before transform/deliver. + from datetime import datetime, timezone + + logger = _logger() + 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 + monkeypatch.setattr(logger, "_ensure_engine", lambda: engine) + + captured = {} + + async def fake_transform_enrich_deliver(*, engine, data, window): + captured["data"] = data + + monkeypatch.setattr(logger, "_transform_enrich_deliver", fake_transform_enrich_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 logger._export_window(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]) + + # Raw DB column names verbatim — no Ternary-specific x_* names — and existing tags preserved. + 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(): + # Enrichment rests on transformed[i] matching source[i]; a transformer reorder would fail this. + 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]}) # 2 rows vs 1 + 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]"]}) # valid JSON, not an object + 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)) # all older than floor + 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(): + logger = _logger() + 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 logger._transform_enrich_deliver(engine=engine, data=data, window=_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(): + logger = _logger() + engine = _fake_engine(transformed=pl.DataFrame(), payload=b"") + await logger._transform_enrich_deliver(engine=engine, data=pl.DataFrame(), window=_window()) + engine._destination.deliver.assert_not_awaited() + + +def test_pod_lock_key_is_ternary_specific(): + # Distinct from FocusLogger's key so the two don't evict each other's lock. + 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") # must not be used when 0 is passed + 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(monkeypatch): + logger = _logger() + 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 + monkeypatch.setattr(logger, "_ensure_engine", lambda: engine) + + captured = {} + + async def fake_ted(*, engine, data, window): + captured["window"] = window + captured["rows"] = data.height + + monkeypatch.setattr(logger, "_transform_enrich_deliver", fake_ted) + await logger._export_all(limit=None) + assert captured["rows"] == 1 + assert captured["window"].frequency == "all" + + +@pytest.mark.asyncio +async def test_transform_enrich_deliver_skips_when_transform_empty(): + logger = _logger() + data = pl.DataFrame({"Tags": [json.dumps({"team_id": "t"})], "prompt_tokens": [1]}) + engine = _fake_engine(transformed=pl.DataFrame(), payload=b"") + await logger._transform_enrich_deliver(engine=engine, data=data, window=_window()) + engine._destination.deliver.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_transform_enrich_deliver_skips_when_payload_empty(): + logger = _logger() + 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 logger._transform_enrich_deliver(engine=engine, data=data, window=_window()) + engine._destination.deliver.assert_not_awaited()