mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-28 01:32:17 +00:00
feat(focus): add ternary destination for FOCUS export
Add a `ternary` FOCUS export destination (plus a `ternary` callback) so a LiteLLM proxy can push FOCUS-formatted spend to Ternary with `callbacks: ["ternary"]` -- no separate export job - FocusTernaryDestination: a pure sink modeled on the sibling vantage/mavvrik/gcs destinations; POSTs the FOCUS CSV to the Ternary cost-ingestion endpoint, chunking an oversized backfill (csv-module repartition, quoted-newline safe; oversized single row raises rather than silently dropping) with a stable X-Ternary-Upload-Id + 0-based index/total so the receiver stages-then-swaps. connection_id is validated + URL-encoded; base_url must be HTTPS (loopback exempt for local); the upload has a bounded timeout; it aborts on the first failed chunk so nothing partial lands - TernaryLogger: presets the FOCUS pipeline (CSV) and enriches the FOCUS Tags JSON with the per-request token counts the shared transformer drops (Ternary weights allocation by tokens); the enrichment is confined to this path and only adds keys. Cadence is daily (default) or interval (validation loops, an explicit interval-seconds required); the export window is day-aligned and drops days before the window start so a whole-day-replace receiver never truncates - Registration mirrors vantage/mavvrik: ternary callback literal, logging dispatch, scheduler job, and CustomLoggerRegistry entry -- no shared-code changes - Tests: mocked unit tests only (no live calls), 100% line coverage of both new modules
This commit is contained in:
Aaron Sproul
parent
b3882d8e43
commit
15a5f879b9
11 changed files with 1164 additions and 0 deletions
|
|
@ -166,6 +166,7 @@ _custom_logger_compatible_callbacks_literal = Literal[
|
|||
"focus",
|
||||
"mavvrik",
|
||||
"vantage",
|
||||
"ternary",
|
||||
"posthog",
|
||||
"levo",
|
||||
"compression_interception",
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
]
|
||||
|
|
|
|||
|
|
@ -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")
|
||||
|
|
|
|||
188
litellm/integrations/focus/destinations/ternary_destination.py
Normal file
188
litellm/integrations/focus/destinations/ternary_destination.py
Normal file
|
|
@ -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,
|
||||
)
|
||||
0
litellm/integrations/ternary/__init__.py
Normal file
0
litellm/integrations/ternary/__init__.py
Normal file
255
litellm/integrations/ternary/ternary_logger.py
Normal file
255
litellm/integrations/ternary/ternary_logger.py
Normal file
|
|
@ -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",)
|
||||
|
|
@ -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,
|
||||
}
|
||||
|
|
|
|||
|
|
@ -4587,6 +4587,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):
|
||||
|
|
@ -5036,6 +5045,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):
|
||||
|
|
|
|||
|
|
@ -10198,6 +10198,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
|
||||
########################################################
|
||||
|
|
|
|||
|
|
@ -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")
|
||||
358
tests/test_litellm/integrations/ternary/test_ternary_logger.py
Normal file
358
tests/test_litellm/integrations/ternary/test_ternary_logger.py
Normal file
|
|
@ -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
|
||||
Loading…
Add table
Reference in a new issue