diff --git a/litellm/__init__.py b/litellm/__init__.py index c8df4394a06..dd260e34cfd 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -166,6 +166,7 @@ _custom_logger_compatible_callbacks_literal = Literal[ "focus", "mavvrik", "vantage", + "ternary", "posthog", "levo", "compression_interception", diff --git a/litellm/integrations/focus/destinations/__init__.py b/litellm/integrations/focus/destinations/__init__.py index 932184cf485..05baf5d06b4 100644 --- a/litellm/integrations/focus/destinations/__init__.py +++ b/litellm/integrations/focus/destinations/__init__.py @@ -5,6 +5,7 @@ from .factory import FocusDestinationFactory from .gcs_destination import FocusGCSDestination from .mavvrik_destination import FocusMavvrikDestination from .s3_destination import FocusS3Destination +from .ternary_destination import FocusTernaryDestination from .vantage_destination import FocusVantageDestination __all__ = [ @@ -13,6 +14,7 @@ __all__ = [ "FocusGCSDestination", "FocusMavvrikDestination", "FocusS3Destination", + "FocusTernaryDestination", "FocusTimeWindow", "FocusVantageDestination", ] diff --git a/litellm/integrations/focus/destinations/factory.py b/litellm/integrations/focus/destinations/factory.py index a46b1ea2dcb..394472e0ef2 100644 --- a/litellm/integrations/focus/destinations/factory.py +++ b/litellm/integrations/focus/destinations/factory.py @@ -9,6 +9,7 @@ from .base import FocusDestination from .gcs_destination import FocusGCSDestination from .mavvrik_destination import FocusMavvrikDestination from .s3_destination import FocusS3Destination +from .ternary_destination import FocusTernaryDestination from .vantage_destination import FocusVantageDestination @@ -33,6 +34,8 @@ class FocusDestinationFactory: return FocusGCSDestination(prefix=prefix, config=normalized_config) if provider_lower == "mavvrik": return FocusMavvrikDestination(prefix=prefix, config=normalized_config) + if provider_lower == "ternary": + return FocusTernaryDestination(prefix=prefix, config=normalized_config) raise NotImplementedError(f"Provider '{provider}' not supported for Focus export") @staticmethod @@ -80,4 +83,17 @@ class FocusDestinationFactory: "connection_id": overrides.get("connection_id") or os.getenv("MAVVRIK_CONNECTION_ID"), } return {k: v for k, v in resolved.items() if v is not None} + if provider == "ternary": + resolved = { + "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 00000000000..80ed3761917 --- /dev/null +++ b/litellm/integrations/focus/destinations/ternary_destination.py @@ -0,0 +1,188 @@ +"""Ternary API destination for FOCUS export: uploads FOCUS CSV to Ternary's cost-ingestion endpoint.""" + +from __future__ import annotations + +import csv +import io +from collections.abc import Sequence +from typing import TYPE_CHECKING, Final +from urllib.parse import quote, urlparse +from uuid import uuid4 + +from litellm._logging import verbose_logger +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler + +from .base import FocusDestination, FocusTimeWindow + +if TYPE_CHECKING: + import httpx + +TERNARY_MAX_ROWS_PER_UPLOAD: Final = 100_000 +TERNARY_MAX_BYTES_PER_UPLOAD: Final = 30 * 1024 * 1024 +TERNARY_UPLOAD_TIMEOUT_SECONDS: Final = 120.0 + +_LOOPBACK_HOSTS: Final = frozenset({"localhost", "127.0.0.1", "::1"}) + + +def _encode_csv(rows: Sequence[Sequence[str]]) -> bytes: + """Serialize CSV rows to UTF-8 bytes via the csv module (quoted fields stay intact).""" + buffer: Final = io.StringIO() + csv.writer(buffer).writerows(rows) + return buffer.getvalue().encode("utf-8") + + +def _require_secure_base_url(base_url: str) -> None: + parsed: Final = urlparse(base_url) + if parsed.scheme == "https": + return + if parsed.scheme == "http" and (parsed.hostname or "").lower() in _LOOPBACK_HOSTS: + return + raise ValueError(f"base_url must be an HTTPS URL (got {base_url!r}); http is allowed only for loopback") + + +class _NoRedirectHTTPHandler(AsyncHTTPHandler): + """AsyncHTTPHandler that disables redirects on every client it creates (SSRF guard).""" + + def create_client(self, *args: object, **kwargs: object) -> httpx.AsyncClient: + client: Final = super().create_client(*args, **kwargs) # pyright: ignore[reportArgumentType] # forwarded verbatim to the base + client.follow_redirects = False + return client + + +class FocusTernaryDestination(FocusDestination): + """Upload FOCUS CSV exports to the Ternary cost-ingestion API.""" + + def __init__( + self, + *, + prefix: str, + config: dict[str, str] | None = None, # mutable-ok: FocusDestination(config) factory contract + ) -> None: + resolved_config: Final = config or {} # mutable-ok: read-only local; empty fallback for absent config + api_key: Final = resolved_config.get("api_key") + connection_id: Final = resolved_config.get("connection_id") + base_url: Final = resolved_config.get("base_url") + if not api_key: + raise ValueError( + "api_key must be provided for Ternary destination " + "(set TERNARY_API_KEY env var or pass in destination_config)" + ) + if not connection_id: + raise ValueError( + "connection_id must be provided for Ternary destination " + "(set TERNARY_CONNECTION_ID env var or pass in destination_config)" + ) + if "/" in connection_id or ".." in connection_id or any(c.isspace() for c in connection_id): + raise ValueError(f"connection_id must not contain '/', '..', or whitespace (got {connection_id!r})") + if not base_url: + raise ValueError( + "base_url must be provided for Ternary destination " + "(set TERNARY_BASE_URL env var or pass in destination_config)" + ) + _require_secure_base_url(str(base_url)) + self.api_key = api_key + self.connection_id = connection_id + self.base_url = str(base_url).rstrip("/") + self.prefix = prefix + + async def deliver( + self, + *, + content: bytes, + time_window: FocusTimeWindow, + filename: str, + ) -> None: + """Upload CSV to the Ternary API, chunking an oversized backfill; aborts on the first failed chunk.""" + if not content: + verbose_logger.debug("Ternary destination: empty content, skipping upload") + return + + client: Final = _NoRedirectHTTPHandler(timeout=TERNARY_UPLOAD_TIMEOUT_SECONDS) + + chunks: Final = self._split_into_chunks(content) + upload_id: Final = uuid4().hex + total: Final = len(chunks) + + try: + for index, chunk in enumerate(chunks): + await self._upload_csv( + client, + chunk, + filename if total == 1 else f"{filename}.part{index + 1}", + upload_id=upload_id, + chunk_index=index, + chunk_total=total, + ) + finally: + await client.close() + + def _split_into_chunks(self, content: bytes) -> Sequence[bytes]: + """Split CSV bytes into chunks within the row/byte limits (csv-aware; a single oversized row errors).""" + 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() + 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 00000000000..e69de29bb2d diff --git a/litellm/integrations/ternary/ternary_logger.py b/litellm/integrations/ternary/ternary_logger.py new file mode 100644 index 00000000000..b2b2c5f1137 --- /dev/null +++ b/litellm/integrations/ternary/ternary_logger.py @@ -0,0 +1,255 @@ +"""Ternary FocusLogger preset (CSV, Ternary destination) that adds per-row token counts to the FOCUS Tags.""" + +from __future__ import annotations + +import json +import os +from datetime import datetime, timedelta, timezone +from typing import TYPE_CHECKING, Final, TypeAlias + +import litellm +from litellm._logging import verbose_logger +from litellm.integrations.custom_logger import CustomLogger +from litellm.integrations.focus.destinations.base import FocusTimeWindow +from litellm.integrations.focus.export_engine import FocusExportEngine +from litellm.integrations.focus.focus_logger import FocusLogger + +if TYPE_CHECKING: + import polars as pl + from apscheduler.schedulers.asyncio import AsyncIOScheduler +else: + AsyncIOScheduler: TypeAlias = object + +TERNARY_USAGE_DATA_JOB_NAME: Final = "ternary_export_usage_data" + +_SUPPORTED_FREQUENCIES: Final = frozenset({"daily", "interval"}) + +_TOKEN_TAG_KEYS: Final = ( + "prompt_tokens", + "completion_tokens", + "cache_read_input_tokens", + "cache_creation_input_tokens", +) + + +def _merge_token_tags(normalized: pl.DataFrame, source: pl.DataFrame) -> pl.DataFrame: + """Zip transformed rows with the row-aligned source and add token keys to each row's Tags; degrades to unchanged on any mismatch.""" + import polars as pl + + if "Tags" not in normalized.columns or normalized.height != source.height: + return normalized + token_cols: Final = tuple(c for c in _TOKEN_TAG_KEYS if c in source.columns) + if not token_cols: + return normalized + + token_rows: Final = source.select(token_cols).to_dicts() + tags_values: Final = normalized["Tags"].to_list() + + merged: Final[list[str]] = [] # mutable-ok: per-row accumulator built in the zip loop + for tags_json, tokens in zip(tags_values, token_rows): + try: + parsed = json.loads(tags_json) if tags_json else {} # mutable-ok: fresh per-row tag dict + if not isinstance(parsed, dict): + parsed = {} # mutable-ok: non-object Tags degrades to empty + except (json.JSONDecodeError, TypeError): + parsed = {} # mutable-ok: malformed Tags degrades to empty + for key, value in tokens.items(): + if value is not None: + parsed[key] = str(value) + merged.append(json.dumps(parsed)) + + return normalized.with_columns(pl.Series("Tags", merged)) + + +def _drop_days_before(data: pl.DataFrame, floor: datetime) -> pl.DataFrame: + """Drop source rows for days older than the window start, so whole-day replace-by-day never truncates a partial older day.""" + import polars as pl + + if "date" not in data.columns: + return data + floor_date: Final = floor.astimezone(timezone.utc).date() + date_col: Final = pl.col("date").cast(pl.Utf8) # cast-ok: polars dtype cast, not typing.cast + parsed: Final = date_col.str.strptime(pl.Date, "%Y-%m-%d", strict=False) + kept: Final = data.filter(parsed >= floor_date) + if data.height > 0 and kept.height == 0: + verbose_logger.warning( + "Ternary export: day-window floor %s dropped all %d rows (unparseable or older `date`?)", + floor_date, + data.height, + ) + return kept + + +def _parse_interval(raw: str | int | None) -> int | None: + """Parse the export interval-seconds override; a non-numeric value is ignored.""" + if raw is None: + return None + try: + return int(raw) + except (ValueError, TypeError): + verbose_logger.warning("Invalid TERNARY_EXPORT_INTERVAL_SECONDS value: %s, ignoring", raw) + return None + + +class _TernaryExportEngine(FocusExportEngine): + """FocusExportEngine that drops pre-window days and merges token counts into Tags before delivery.""" + + async def export_window(self, *, window: FocusTimeWindow, limit: int | None) -> None: + data: Final = await self._database.get_usage_data( + limit=limit, + start_time_utc=window.start_time, + end_time_utc=window.end_time, + ) + await self._deliver_enriched(_drop_days_before(data, window.start_time), window) + + async def export_all(self, *, limit: int | None) -> None: + data: Final = await self._database.get_usage_data(limit=limit) + now: Final = datetime.now(timezone.utc) + window: Final = FocusTimeWindow( + start_time=now.replace(hour=0, minute=0, second=0, microsecond=0), + end_time=now, + frequency="all", + ) + await self._deliver_enriched(data, window) + + async def _deliver_enriched(self, data: pl.DataFrame, window: FocusTimeWindow) -> None: + if data.is_empty(): + verbose_logger.debug("Ternary export: no usage data for window %s", window) + return + transformed: Final = self._transformer.transform(data) + if transformed.is_empty(): + verbose_logger.debug("Ternary export: normalized data empty for window %s", window) + return + enriched: Final = _merge_token_tags(transformed, data) + payload: Final = self._serializer.serialize(enriched) + if not payload: + verbose_logger.debug("Ternary export: serializer returned empty payload") + return + await self._destination.deliver( + content=payload, + time_window=window, + filename=self._build_filename(window), + ) + + +class TernaryLogger(FocusLogger): + """FocusLogger preset for Ternary (CSV, Ternary cost-import API). + + Env: TERNARY_API_KEY, TERNARY_CONNECTION_ID, TERNARY_BASE_URL (required), + TERNARY_EXPORT_FREQUENCY ("daily" default / "interval"), TERNARY_EXPORT_INTERVAL_SECONDS. + """ + + def __init__( + self, + *, + api_key: str | None = None, + connection_id: str | None = None, + base_url: str | None = None, + frequency: str | None = None, + interval_seconds: int | None = None, + ) -> None: + resolved_api_key: Final = api_key or os.getenv("TERNARY_API_KEY") + resolved_connection_id: Final = connection_id or os.getenv("TERNARY_CONNECTION_ID") + resolved_base_url: Final = base_url or os.getenv("TERNARY_BASE_URL") + resolved_frequency: Final = (frequency or os.getenv("TERNARY_EXPORT_FREQUENCY") or "daily").lower() + if resolved_frequency not in _SUPPORTED_FREQUENCIES: + raise ValueError( + f"Unsupported TERNARY_EXPORT_FREQUENCY {resolved_frequency!r}; " + f"Ternary supports {sorted(_SUPPORTED_FREQUENCIES)}. LiteLLM spend is a daily " + "aggregate landed with whole-day replace, so 'hourly' adds no grain -- use " + "'interval' only for short setup-validation loops." + ) + + resolved_interval: Final = _parse_interval( + interval_seconds if interval_seconds is not None else os.getenv("TERNARY_EXPORT_INTERVAL_SECONDS") + ) + if resolved_frequency == "interval" and (resolved_interval is None or resolved_interval <= 0): + raise ValueError( + "TERNARY_EXPORT_INTERVAL_SECONDS must be a positive integer when TERNARY_EXPORT_FREQUENCY is 'interval'" + ) + + destination_config: Final[dict[str, str]] = {} # mutable-ok: built from the config values present below + if resolved_api_key: + destination_config["api_key"] = resolved_api_key + if resolved_connection_id: + destination_config["connection_id"] = resolved_connection_id + if resolved_base_url: + destination_config["base_url"] = resolved_base_url + + super().__init__( + provider="ternary", + export_format="csv", + frequency=resolved_frequency, + interval_seconds=resolved_interval, + prefix="ternary_exports", + destination_config=destination_config, + ) + + verbose_logger.debug( + "TernaryLogger initialized (connection_id=%s)", + ( + resolved_connection_id[:4] + "***" + if resolved_connection_id and len(resolved_connection_id) > 4 + else "***" + ), + ) + + def _compute_time_window(self, now: datetime) -> FocusTimeWindow: + """Snap the window start to the previous UTC midnight so each push carries whole days.""" + now_utc: Final = now.astimezone(timezone.utc) + start_time: Final = (now_utc - timedelta(days=1)).replace(hour=0, minute=0, second=0, microsecond=0) + return FocusTimeWindow(start_time=start_time, end_time=now_utc, frequency=self.frequency) + + def _ensure_engine(self) -> FocusExportEngine: + """Use the Ternary export engine (day-window filtering + Tags token enrichment).""" + if self._engine is None: + self._engine = _TernaryExportEngine( + provider=self.provider, + export_format=self.export_format, + prefix=self.prefix, + destination_config=self._destination_config, + ) + return self._engine + + async def initialize_focus_export_job(self) -> None: + """Run the scheduled export under a Ternary-specific pod lock (distinct from FocusLogger's).""" + from litellm.proxy.proxy_server import proxy_logging_obj + + writer: Final = getattr(proxy_logging_obj, "db_spend_update_writer", None) if proxy_logging_obj else None + pod_lock_manager: Final = getattr(writer, "pod_lock_manager", None) if writer is not None else None + + if pod_lock_manager and pod_lock_manager.redis_cache: + acquired: Final = await pod_lock_manager.acquire_lock(cronjob_id=TERNARY_USAGE_DATA_JOB_NAME) + if not acquired: + verbose_logger.debug("Ternary export: unable to acquire pod lock") + return + try: + await self._run_scheduled_export() + finally: + await pod_lock_manager.release_lock(cronjob_id=TERNARY_USAGE_DATA_JOB_NAME) + else: + await self._run_scheduled_export() + + @staticmethod + async def init_ternary_background_job( + scheduler: AsyncIOScheduler, + ) -> None: + """Register the Ternary export job with the provided scheduler.""" + ternary_loggers: Final[list[CustomLogger]] = ( # mutable-ok: list returned by the shared callback manager + litellm.logging_callback_manager.get_custom_loggers_for_type(callback_type=TernaryLogger) + ) + if not ternary_loggers: + verbose_logger.debug("No Ternary logger registered; skipping scheduler") + return + + ternary_logger: Final = ternary_loggers[0] + if not isinstance(ternary_logger, TernaryLogger): + return + trigger_kwargs: Final = ternary_logger._build_scheduler_trigger() + scheduler.add_job( + ternary_logger.initialize_focus_export_job, + **trigger_kwargs, + ) + + +__all__ = ("TernaryLogger",) diff --git a/litellm/litellm_core_utils/custom_logger_registry.py b/litellm/litellm_core_utils/custom_logger_registry.py index 6294f3bc577..101225c6cb5 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, @@ -108,6 +109,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 8e28a0d543d..e114bd3ec6b 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -4858,6 +4858,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): @@ -5307,6 +5316,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 92a75bf953a..f7e31523235 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -10669,6 +10669,12 @@ class ProxyStartupEvent: await MavvrikFocusLogger.init_mavvrik_focus_background_job(scheduler=scheduler) + from litellm.integrations.ternary.ternary_logger import ( # noqa: PLC0415 # lazy import avoids a circular import at module load + TernaryLogger, + ) + + await TernaryLogger.init_ternary_background_job(scheduler=scheduler) + ######################################################## # Prometheus Background Job ######################################################## diff --git a/tests/test_litellm/integrations/focus/test_ternary_destination.py b/tests/test_litellm/integrations/focus/test_ternary_destination.py new file mode 100644 index 00000000000..dfb2bb444ce --- /dev/null +++ b/tests/test_litellm/integrations/focus/test_ternary_destination.py @@ -0,0 +1,321 @@ +"""Tests for FocusTernaryDestination behavior.""" + +from __future__ import annotations + +import csv +import io +from datetime import datetime, timedelta, timezone +from typing import Any +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +import litellm.integrations.focus.destinations.ternary_destination as td +from litellm.integrations.focus.destinations.base import FocusTimeWindow +from litellm.integrations.focus.destinations.ternary_destination import ( + TERNARY_UPLOAD_TIMEOUT_SECONDS, + FocusTernaryDestination, + _NoRedirectHTTPHandler, +) + +MOCK_TARGET = "litellm.integrations.focus.destinations.ternary_destination._NoRedirectHTTPHandler" + + +def _window(freq: str = "daily", hour: int = 5) -> FocusTimeWindow: + start = datetime(2024, 1, 2, hour, tzinfo=timezone.utc) + end = start + timedelta(hours=1) + return FocusTimeWindow(start_time=start, end_time=end, frequency=freq) + + +def _config(**overrides: Any) -> dict[str, Any]: + base = { + "api_key": "test-api-key", + "connection_id": "conn-1234", + "base_url": "https://ternary.test", + } + base.update(overrides) + return base + + +def _capturing_client() -> tuple[MagicMock, list[dict[str, Any]]]: + """Return a mock client plus a list that records each post() call.""" + calls: list[dict[str, Any]] = [] + + mock_response = AsyncMock() + mock_response.raise_for_status = lambda: None + + mock_client = MagicMock() + mock_client.close = AsyncMock() + + async def capture_post(url, **kwargs): + calls.append({"url": url, **kwargs}) + return mock_response + + mock_client.post = capture_post + return mock_client, calls + + +def _uploaded_part(call: dict[str, Any]) -> tuple[str, bytes]: + field = call["files"]["csv"] + return field[0], field[1] + + +def _rows(content: bytes) -> list[list[str]]: + return list(csv.reader(io.StringIO(content.decode("utf-8")))) + + +def test_should_require_api_key(): + with pytest.raises(ValueError, match="api_key"): + FocusTernaryDestination(prefix="exports", config={"connection_id": "c", "base_url": "u"}) + + +def test_should_require_connection_id(): + with pytest.raises(ValueError, match="connection_id"): + FocusTernaryDestination(prefix="exports", config={"api_key": "k", "base_url": "u"}) + + +def test_should_require_base_url(): + with pytest.raises(ValueError, match="base_url"): + FocusTernaryDestination(prefix="exports", config={"api_key": "k", "connection_id": "c"}) + + +@pytest.mark.parametrize("bad_id", ["a/b", "..", "has space", "tab\tid"]) +def test_should_reject_connection_id_that_could_reroute_the_path(bad_id): + with pytest.raises(ValueError, match="connection_id"): + FocusTernaryDestination(prefix="exports", config=_config(connection_id=bad_id)) + + +def test_should_initialize_with_valid_config(): + dest = FocusTernaryDestination(prefix="exports", config=_config()) + assert dest.api_key == "test-api-key" + assert dest.connection_id == "conn-1234" + assert dest.base_url == "https://ternary.test" + + +def test_should_use_custom_base_url_and_strip_trailing_slash(): + dest = FocusTernaryDestination(prefix="exports", config=_config(base_url="http://localhost:8080/")) + assert dest.base_url == "http://localhost:8080" + + +@pytest.mark.asyncio +async def test_should_skip_empty_content(): + dest = FocusTernaryDestination(prefix="exports", config=_config()) + with patch(MOCK_TARGET, side_effect=AssertionError("client initialized on empty content")): + assert await dest.deliver(content=b"", time_window=_window(), filename="usage.csv") is None + + +@pytest.mark.asyncio +async def test_should_upload_to_correct_url_with_auth_and_upload_headers(): + dest = FocusTernaryDestination(prefix="exports", config=_config()) + mock_client, calls = _capturing_client() + + with patch(MOCK_TARGET, return_value=mock_client): + await dest.deliver(content=b"header\nrow1\n", time_window=_window(), filename="usage.csv") + + assert len(calls) == 1 + call = calls[0] + assert call["url"] == "https://ternary.test/external-cost-sources/v1/conn-1234/focus" + assert call["headers"]["Authorization"] == "Bearer test-api-key" + assert call["headers"]["X-Ternary-Chunk-Index"] == "0" + assert call["headers"]["X-Ternary-Chunk-Total"] == "1" + assert call["headers"]["X-Ternary-Upload-Id"] + filename, body = _uploaded_part(call) + assert filename == "usage.csv" + assert body == b"header\nrow1\n" + assert call["files"]["csv"][2] == "text/csv" + assert call["timeout"] == TERNARY_UPLOAD_TIMEOUT_SECONDS + + +@pytest.mark.asyncio +async def test_should_close_the_client_after_upload(): + dest = FocusTernaryDestination(prefix="exports", config=_config()) + mock_client, _ = _capturing_client() + with patch(MOCK_TARGET, return_value=mock_client): + await dest.deliver(content=b"h\nr\n", time_window=_window(), filename="usage.csv") + mock_client.close.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_no_redirect_handler_disables_redirects_including_retry_client(): + handler = _NoRedirectHTTPHandler() + retry_client = handler.create_client(timeout=1.0, event_hooks=None) + try: + assert handler.client.follow_redirects is False + assert retry_client.follow_redirects is False + finally: + await retry_client.aclose() + await handler.close() + + +@pytest.mark.asyncio +async def test_should_url_encode_the_connection_id(): + dest = FocusTernaryDestination(prefix="exports", config=_config(connection_id="conn+id~ok")) + mock_client, calls = _capturing_client() + with patch(MOCK_TARGET, return_value=mock_client): + await dest.deliver(content=b"h\nr\n", time_window=_window(), filename="usage.csv") + assert calls[0]["url"].endswith("/external-cost-sources/v1/conn%2Bid~ok/focus") + + +@pytest.mark.asyncio +async def test_should_pass_tags_column_through_unstripped(): + """The Ternary sink must not drop/strip any columns (forwards Tags as-is).""" + dest = FocusTernaryDestination(prefix="exports", config=_config()) + mock_client, calls = _capturing_client() + + content = b'ServiceName,Tags,x_unknown\nfoo,"{""team_id"": ""t1""}",keepme\n' + with patch(MOCK_TARGET, return_value=mock_client): + await dest.deliver(content=content, time_window=_window(), filename="usage.csv") + + _, body = _uploaded_part(calls[0]) + assert body == content + + +@pytest.mark.asyncio +async def test_should_chunk_by_row_count_with_stable_upload_id(monkeypatch): + monkeypatch.setattr(td, "TERNARY_MAX_ROWS_PER_UPLOAD", 2) + dest = FocusTernaryDestination(prefix="exports", config=_config()) + + content = b"ServiceName\n" + b"\n".join([b"x"] * 5) + b"\n" + mock_client, calls = _capturing_client() + with patch(MOCK_TARGET, return_value=mock_client): + await dest.deliver(content=content, time_window=_window(), filename="usage.csv") + + assert len(calls) == 3 + total = "3" + assert {c["headers"]["X-Ternary-Upload-Id"] for c in calls} == {calls[0]["headers"]["X-Ternary-Upload-Id"]} + for i, call in enumerate(calls): + assert call["headers"]["X-Ternary-Chunk-Index"] == str(i) + assert call["headers"]["X-Ternary-Chunk-Total"] == total + filename, body = _uploaded_part(call) + assert filename == f"usage.csv.part{i + 1}" + assert len(_rows(body)) - 1 <= 2 + + +@pytest.mark.asyncio +async def test_should_chunk_by_bytes(monkeypatch): + monkeypatch.setattr(td, "TERNARY_MAX_BYTES_PER_UPLOAD", 200) + dest = FocusTernaryDestination(prefix="exports", config=_config()) + + row = b"a" * 40 + b"," + b"b" * 40 + content = b"ServiceName,BilledCost\n" + b"\n".join([row] * 20) + b"\n" + mock_client, calls = _capturing_client() + with patch(MOCK_TARGET, return_value=mock_client): + await dest.deliver(content=content, time_window=_window(), filename="usage.csv") + + assert len(calls) > 1 + for call in calls: + _, body = _uploaded_part(call) + assert len(body) <= 200 + + +@pytest.mark.asyncio +async def test_should_not_mangle_a_quoted_field_containing_a_newline(monkeypatch): + monkeypatch.setattr(td, "TERNARY_MAX_ROWS_PER_UPLOAD", 1) + dest = FocusTernaryDestination(prefix="exports", config=_config()) + + content = b'ServiceName,Tags\nsvc1,"line1\nline2"\nsvc2,"{""k"":""v""}"\n' + mock_client, calls = _capturing_client() + with patch(MOCK_TARGET, return_value=mock_client): + await dest.deliver(content=content, time_window=_window(), filename="usage.csv") + + assert len(calls) == 2 + _, first = _uploaded_part(calls[0]) + rows = _rows(first) + assert rows[0] == ["ServiceName", "Tags"] + assert rows[1] == ["svc1", "line1\nline2"] + + +@pytest.mark.asyncio +async def test_should_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() + mock_client.close = AsyncMock() + + async def failing_post(url, **kwargs): + attempts.append(kwargs["headers"]["X-Ternary-Chunk-Index"]) + raise RuntimeError("boom") + + mock_client.post = failing_post + + with patch(MOCK_TARGET, return_value=mock_client): + with pytest.raises(RuntimeError, match="boom"): + await dest.deliver(content=content, time_window=_window(), filename="usage.csv") + + assert attempts == ["0"] + + +@pytest.mark.parametrize("url", ["http://api.ternary.app", "http://evil.example.com:8080", "ftp://ternary.test"]) +def test_should_reject_non_https_base_url(url): + with pytest.raises(ValueError, match="HTTPS"): + FocusTernaryDestination(prefix="exports", config=_config(base_url=url)) + + +@pytest.mark.parametrize("url", ["https://ternary.test", "http://localhost:8080", "http://127.0.0.1:8080"]) +def test_should_accept_https_or_loopback_base_url(url): + dest = FocusTernaryDestination(prefix="exports", config=_config(base_url=url)) + assert dest.base_url == url.rstrip("/") + + +@pytest.mark.asyncio +async def test_should_return_header_only_content_untouched(monkeypatch): + monkeypatch.setattr(td, "TERNARY_MAX_BYTES_PER_UPLOAD", 10) + dest = FocusTernaryDestination(prefix="exports", config=_config()) + content = b"HeaderOnlyLineWithNoDataRows" + mock_client, calls = _capturing_client() + with patch(MOCK_TARGET, return_value=mock_client): + await dest.deliver(content=content, time_window=_window(), filename="usage.csv") + assert len(calls) == 1 + _, body = _uploaded_part(calls[0]) + assert body == content + + +def test_factory_creates_ternary_destination_from_config(): + from litellm.integrations.focus.destinations.factory import FocusDestinationFactory + + dest = FocusDestinationFactory.create( + provider="ternary", + prefix="exports", + config={"api_key": "k", "connection_id": "c", "base_url": "https://ternary.test"}, + ) + assert isinstance(dest, FocusTernaryDestination) + assert dest.connection_id == "c" + + +def test_factory_resolves_ternary_config_from_env(monkeypatch): + monkeypatch.setenv("TERNARY_API_KEY", "envk") + monkeypatch.setenv("TERNARY_CONNECTION_ID", "envc") + monkeypatch.setenv("TERNARY_BASE_URL", "https://ternary.test") + from litellm.integrations.focus.destinations.factory import FocusDestinationFactory + + dest = FocusDestinationFactory.create(provider="ternary", prefix="exports") + assert isinstance(dest, FocusTernaryDestination) + assert dest.api_key == "envk" + assert dest.connection_id == "envc" + + +@pytest.mark.parametrize("missing", ["TERNARY_API_KEY", "TERNARY_CONNECTION_ID", "TERNARY_BASE_URL"]) +def test_factory_requires_each_ternary_env_var(monkeypatch, missing): + for var in ("TERNARY_API_KEY", "TERNARY_CONNECTION_ID", "TERNARY_BASE_URL"): + monkeypatch.setenv(var, "https://ternary.test" if var == "TERNARY_BASE_URL" else "x") + monkeypatch.delenv(missing, raising=False) + from litellm.integrations.focus.destinations.factory import FocusDestinationFactory + + with pytest.raises(ValueError, match=missing): + FocusDestinationFactory.create(provider="ternary", prefix="exports") diff --git a/tests/test_litellm/integrations/ternary/test_ternary_logger.py b/tests/test_litellm/integrations/ternary/test_ternary_logger.py new file mode 100644 index 00000000000..44f79cf33a0 --- /dev/null +++ b/tests/test_litellm/integrations/ternary/test_ternary_logger.py @@ -0,0 +1,358 @@ +"""Tests for TernaryLogger configuration and Tags token enrichment.""" + +from __future__ import annotations + +import json +from unittest.mock import AsyncMock, MagicMock + +import polars as pl +import pytest + +from litellm.integrations.focus.destinations.base import FocusTimeWindow +from litellm.integrations.focus.focus_logger import FocusLogger +from litellm.integrations.ternary.ternary_logger import ( + TERNARY_USAGE_DATA_JOB_NAME, + TernaryLogger, + _drop_days_before, + _merge_token_tags, + _TernaryExportEngine, +) + + +def _logger() -> TernaryLogger: + return TernaryLogger(api_key="k", connection_id="c", base_url="http://localhost:8080") + + +def _window() -> FocusTimeWindow: + from datetime import datetime, timezone + + now = datetime(2026, 9, 4, tzinfo=timezone.utc) + return FocusTimeWindow(start_time=now, end_time=now, frequency="daily") + + +def test_should_preset_focus_config_for_ternary(): + logger = TernaryLogger(api_key="k", connection_id="c", base_url="http://localhost:8080") + assert isinstance(logger, FocusLogger) + assert logger.provider == "ternary" + assert logger.export_format == "csv" + assert logger.frequency == "daily" + assert logger.prefix == "ternary_exports" + assert logger._destination_config == { + "api_key": "k", + "connection_id": "c", + "base_url": "http://localhost:8080", + } + + +def test_should_read_frequency_from_env(monkeypatch): + monkeypatch.setenv("TERNARY_API_KEY", "k") + monkeypatch.setenv("TERNARY_CONNECTION_ID", "c") + monkeypatch.setenv("TERNARY_EXPORT_FREQUENCY", "daily") + logger = TernaryLogger() + assert logger.frequency == "daily" + assert logger._destination_config["api_key"] == "k" + assert logger._destination_config["connection_id"] == "c" + + +def test_should_default_to_daily(): + logger = TernaryLogger(api_key="k", connection_id="c", base_url="http://localhost:8080") + assert logger.frequency == "daily" + + +def test_should_accept_interval_with_explicit_seconds(): + logger = TernaryLogger( + api_key="k", + connection_id="c", + base_url="http://localhost:8080", + frequency="interval", + interval_seconds=30, + ) + assert logger.frequency == "interval" + assert logger.interval_seconds == 30 + + +def test_should_reject_interval_without_seconds(): + with pytest.raises(ValueError, match="TERNARY_EXPORT_INTERVAL_SECONDS"): + TernaryLogger(api_key="k", connection_id="c", base_url="http://localhost:8080", frequency="interval") + + +@pytest.mark.parametrize("bad_frequency", ["hourly", "weekly", "minutely"]) +def test_should_reject_unsupported_frequency(bad_frequency): + with pytest.raises(ValueError, match="TERNARY_EXPORT_FREQUENCY"): + TernaryLogger(api_key="k", connection_id="c", base_url="http://localhost:8080", frequency=bad_frequency) + + +def test_compute_time_window_is_day_aligned_and_carries_whole_days(): + from datetime import datetime, timezone + + logger = TernaryLogger(api_key="k", connection_id="c", base_url="http://localhost:8080") + now = datetime(2026, 9, 4, 13, 47, 5, tzinfo=timezone.utc) + window = logger._compute_time_window(now) + assert window.start_time == datetime(2026, 9, 3, 0, 0, 0, tzinfo=timezone.utc) + assert window.end_time == now + assert window.start_time.hour == 0 and window.start_time.minute == 0 + + +def test_drop_days_before_removes_older_days(): + from datetime import datetime, timezone + + data = pl.DataFrame({"date": ["2026-09-02", "2026-09-03", "2026-09-04"], "spend": [1.0, 2.0, 3.0]}) + floor = datetime(2026, 9, 3, 0, 0, tzinfo=timezone.utc) + out = _drop_days_before(data, floor) + assert out["date"].to_list() == ["2026-09-03", "2026-09-04"] + + +def test_drop_days_before_noop_without_date_column(): + from datetime import datetime, timezone + + data = pl.DataFrame({"spend": [1.0]}) + out = _drop_days_before(data, datetime(2026, 9, 3, tzinfo=timezone.utc)) + assert out.equals(data) + + +@pytest.mark.asyncio +async def test_export_window_interposes_drop_days_before(): + from datetime import datetime, timezone + + multi_day = pl.DataFrame({"date": ["2026-09-01", "2026-09-03", "2026-09-04"], "spend": [1.0, 2.0, 3.0]}) + + engine = MagicMock() + + async def fake_get_usage_data(**_kwargs): + return multi_day + + engine._database.get_usage_data = fake_get_usage_data + + captured = {} + + async def fake_deliver(data, window): + captured["data"] = data + + engine._deliver_enriched = fake_deliver + + window = FocusTimeWindow( + start_time=datetime(2026, 9, 3, 0, 0, tzinfo=timezone.utc), + end_time=datetime(2026, 9, 4, 12, 0, tzinfo=timezone.utc), + frequency="daily", + ) + await _TernaryExportEngine.export_window(engine, window=window, limit=None) + + assert captured["data"]["date"].to_list() == ["2026-09-03", "2026-09-04"] + + +def test_merge_token_tags_adds_raw_db_keys_and_preserves_existing(): + normalized = pl.DataFrame( + { + "Tags": [json.dumps({"team_id": "t1", "model": "gpt-4o"})], + "BilledCost": [1.23], + } + ) + source = pl.DataFrame( + { + "prompt_tokens": [48], + "completion_tokens": [274], + "cache_read_input_tokens": [12], + "cache_creation_input_tokens": [0], + } + ) + + tags = json.loads(_merge_token_tags(normalized, source)["Tags"][0]) + + assert tags["prompt_tokens"] == "48" + assert tags["completion_tokens"] == "274" + assert tags["cache_read_input_tokens"] == "12" + assert tags["cache_creation_input_tokens"] == "0" + assert not any(k.startswith("x_") for k in tags) + assert tags["team_id"] == "t1" + assert tags["model"] == "gpt-4o" + + +def test_merge_token_tags_aligns_per_row_through_the_real_transformer(): + from litellm.integrations.focus.transformer import FocusTransformer + + data = pl.DataFrame( + { + "date": ["2026-09-03", "2026-09-03", "2026-09-03"], + "spend": [1.0, 2.0, 3.0], + "api_key": ["k1", "k2", "k3"], + "api_key_alias": ["a1", "a2", "a3"], + "model": ["m1", "m2", "m3"], + "model_group": ["g1", "g2", "g3"], + "custom_llm_provider": ["openai", "openai", "anthropic"], + "team_id": ["t1", "t2", "t3"], + "team_alias": ["T1", "T2", "T3"], + "api_requests": [1, 1, 1], + "prompt_tokens": [10, 20, 30], + "completion_tokens": [11, 22, 33], + "cache_read_input_tokens": [0, 0, 0], + "cache_creation_input_tokens": [0, 0, 0], + } + ) + + transformed = FocusTransformer().transform(data) + enriched = _merge_token_tags(transformed, data) + + assert enriched.height == 3 + for i, expected_tokens in enumerate(["10", "20", "30"]): + tags = json.loads(enriched["Tags"][i]) + assert tags["prompt_tokens"] == expected_tokens + assert tags["team_id"] == f"t{i + 1}" + + +def test_merge_token_tags_skips_none_values(): + normalized = pl.DataFrame({"Tags": [json.dumps({"team_id": "t1"})]}) + source = pl.DataFrame({"prompt_tokens": [None], "completion_tokens": [10]}) + tags = json.loads(_merge_token_tags(normalized, source)["Tags"][0]) + assert "prompt_tokens" not in tags + assert tags["completion_tokens"] == "10" + + +def test_merge_token_tags_graceful_on_row_mismatch(): + normalized = pl.DataFrame({"Tags": [json.dumps({"team_id": "t1"})]}) + source = pl.DataFrame({"prompt_tokens": [1, 2]}) + out = _merge_token_tags(normalized, source) + assert out["Tags"][0] == normalized["Tags"][0] + + +def test_merge_token_tags_non_object_tags_degrades_to_empty(): + normalized = pl.DataFrame({"Tags": ["[1, 2, 3]"]}) + source = pl.DataFrame({"prompt_tokens": [5]}) + tags = json.loads(_merge_token_tags(normalized, source)["Tags"][0]) + assert tags == {"prompt_tokens": "5"} + + +def test_drop_days_before_warns_when_everything_dropped(): + from datetime import datetime, timezone + + data = pl.DataFrame({"date": ["2026-09-01", "2026-09-02"], "spend": [1.0, 2.0]}) + out = _drop_days_before(data, datetime(2026, 9, 5, tzinfo=timezone.utc)) + assert out.height == 0 + + +def test_merge_token_tags_graceful_on_malformed_tags(): + normalized = pl.DataFrame({"Tags": ["not-json"]}) + source = pl.DataFrame({"prompt_tokens": [5]}) + tags = json.loads(_merge_token_tags(normalized, source)["Tags"][0]) + assert tags == {"prompt_tokens": "5"} + + +def test_merge_token_tags_noop_without_tags_column(): + normalized = pl.DataFrame({"BilledCost": [1.0]}) + source = pl.DataFrame({"prompt_tokens": [5]}) + out = _merge_token_tags(normalized, source) + assert "Tags" not in out.columns + + +def test_merge_token_tags_noop_without_token_columns(): + normalized = pl.DataFrame({"Tags": [json.dumps({"team_id": "t1"})]}) + source = pl.DataFrame({"spend": [1.0]}) + tags = json.loads(_merge_token_tags(normalized, source)["Tags"][0]) + assert tags == {"team_id": "t1"} + + +def _fake_engine(*, transformed: pl.DataFrame, payload: bytes) -> MagicMock: + engine = MagicMock() + engine._transformer.transform = MagicMock(return_value=transformed) + engine._serializer.serialize = MagicMock(return_value=payload) + engine._destination.deliver = AsyncMock() + engine._build_filename = MagicMock(return_value="usage.csv") + return engine + + +@pytest.mark.asyncio +async def test_transform_enrich_deliver_enriches_tags_then_delivers(): + data = pl.DataFrame({"Tags": [json.dumps({"team_id": "t1"})], "prompt_tokens": [7]}) + transformed = pl.DataFrame({"Tags": [json.dumps({"team_id": "t1"})]}) + engine = _fake_engine(transformed=transformed, payload=b"csv-bytes") + + await _TernaryExportEngine._deliver_enriched(engine, data, _window()) + + serialized_frame = engine._serializer.serialize.call_args.args[0] + assert json.loads(serialized_frame["Tags"][0])["prompt_tokens"] == "7" + engine._destination.deliver.assert_awaited_once() + kwargs = engine._destination.deliver.await_args.kwargs + assert kwargs["content"] == b"csv-bytes" + assert kwargs["filename"] == "usage.csv" + + +@pytest.mark.asyncio +async def test_transform_enrich_deliver_skips_empty_data(): + engine = _fake_engine(transformed=pl.DataFrame(), payload=b"") + await _TernaryExportEngine._deliver_enriched(engine, pl.DataFrame(), _window()) + engine._destination.deliver.assert_not_awaited() + + +def test_pod_lock_key_is_ternary_specific(): + assert TERNARY_USAGE_DATA_JOB_NAME == "ternary_export_usage_data" + + +def test_should_reject_explicit_zero_interval_without_env_fallthrough(monkeypatch): + monkeypatch.setenv("TERNARY_EXPORT_INTERVAL_SECONDS", "300") + with pytest.raises(ValueError, match="TERNARY_EXPORT_INTERVAL_SECONDS"): + TernaryLogger( + api_key="k", connection_id="c", base_url="http://localhost:8080", frequency="interval", interval_seconds=0 + ) + + +def test_should_reject_non_numeric_interval(monkeypatch): + monkeypatch.setenv("TERNARY_EXPORT_INTERVAL_SECONDS", "5m") + with pytest.raises(ValueError, match="TERNARY_EXPORT_INTERVAL_SECONDS"): + TernaryLogger(api_key="k", connection_id="c", base_url="http://localhost:8080", frequency="interval") + + +@pytest.mark.asyncio +async def test_export_all_delivers_with_all_window(): + engine = MagicMock() + frame = pl.DataFrame({"Tags": [json.dumps({"team_id": "t"})], "prompt_tokens": [5]}) + + async def fake_get(**_kwargs): + return frame + + engine._database.get_usage_data = fake_get + + captured = {} + + async def fake_deliver(data, window): + captured["window"] = window + captured["rows"] = data.height + + engine._deliver_enriched = fake_deliver + await _TernaryExportEngine.export_all(engine, limit=None) + assert captured["rows"] == 1 + assert captured["window"].frequency == "all" + + +@pytest.mark.asyncio +async def test_transform_enrich_deliver_skips_when_transform_empty(): + data = pl.DataFrame({"Tags": [json.dumps({"team_id": "t"})], "prompt_tokens": [1]}) + engine = _fake_engine(transformed=pl.DataFrame(), payload=b"") + await _TernaryExportEngine._deliver_enriched(engine, data, _window()) + engine._destination.deliver.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_transform_enrich_deliver_skips_when_payload_empty(): + data = pl.DataFrame({"Tags": [json.dumps({"team_id": "t"})]}) + transformed = pl.DataFrame({"Tags": [json.dumps({"team_id": "t"})]}) + engine = _fake_engine(transformed=transformed, payload=b"") + await _TernaryExportEngine._deliver_enriched(engine, data, _window()) + engine._destination.deliver.assert_not_awaited() + + +def test_ensure_engine_builds_ternary_export_engine(): + logger = _logger() + engine = logger._ensure_engine() + assert isinstance(engine, _TernaryExportEngine) + assert logger._ensure_engine() is engine + + +def test_init_custom_logger_compatible_class_returns_cached_ternary_logger(monkeypatch): + monkeypatch.setenv("TERNARY_API_KEY", "k") + monkeypatch.setenv("TERNARY_CONNECTION_ID", "c") + monkeypatch.setenv("TERNARY_BASE_URL", "https://ternary.test") + from litellm.litellm_core_utils.litellm_logging import _init_custom_logger_compatible_class + + first = _init_custom_logger_compatible_class("ternary", internal_usage_cache=None, llm_router=None) + assert isinstance(first, TernaryLogger) + second = _init_custom_logger_compatible_class("ternary", internal_usage_cache=None, llm_router=None) + assert second is first