diff --git a/litellm/integrations/focus/destinations/ternary_destination.py b/litellm/integrations/focus/destinations/ternary_destination.py index 6d930718cdc..b1be6d0dcc5 100644 --- a/litellm/integrations/focus/destinations/ternary_destination.py +++ b/litellm/integrations/focus/destinations/ternary_destination.py @@ -1,13 +1,11 @@ -"""Ternary API destination for FOCUS export: uploads FOCUS CSV to Ternary's cost-ingestion endpoint.""" +"""Ternary API destination for FOCUS export: uploads each day's FOCUS CSV via a Ternary-signed GCS URL.""" from __future__ import annotations -import csv -import io -from collections.abc import Sequence -from typing import TYPE_CHECKING, Final +from typing import TYPE_CHECKING, Final, Protocol from urllib.parse import quote, urlparse -from uuid import uuid4 + +from typing_extensions import NotRequired, ReadOnly, TypedDict from litellm._logging import verbose_logger from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler @@ -17,18 +15,23 @@ from .base import FocusDestination, FocusTimeWindow if TYPE_CHECKING: import httpx -TERNARY_MAX_ROWS_PER_UPLOAD: Final = 100_000 -TERNARY_MAX_BYTES_PER_UPLOAD: Final = 30 * 1024 * 1024 TERNARY_UPLOAD_TIMEOUT_SECONDS: Final = 120.0 _LOOPBACK_HOSTS: Final = frozenset({"localhost", "127.0.0.1", "::1"}) -def _encode_csv(rows: Sequence[Sequence[str]]) -> bytes: - """Serialize CSV rows to UTF-8 bytes via the csv module (quoted fields stay intact).""" - buffer: Final = io.StringIO() - csv.writer(buffer).writerows(rows) - return buffer.getvalue().encode("utf-8") +class TernaryUploadUrlBody(TypedDict): + url: ReadOnly[NotRequired[str]] + method: ReadOnly[NotRequired[str]] + headers: ReadOnly[NotRequired[dict[str, str]]] + + +class _UploadUrlResponse(Protocol): + def json(self) -> TernaryUploadUrlBody: ... + + +def _upload_url_body(response: _UploadUrlResponse) -> TernaryUploadUrlBody: + return response.json() def _require_secure_base_url(base_url: str) -> None: @@ -40,6 +43,14 @@ def _require_secure_base_url(base_url: str) -> None: raise ValueError(f"base_url must be an HTTPS URL (got {base_url!r}); http is allowed only for loopback") +def _require_gcs_url(url: str) -> None: + parsed: Final = urlparse(url) + hostname: Final = (parsed.hostname or "").lower() + is_gcs: Final = hostname == "storage.googleapis.com" or hostname.endswith(".storage.googleapis.com") + if parsed.scheme != "https" or not is_gcs: + raise ValueError(f"Ternary destination: upload URL must be an HTTPS GCS URL (got host {hostname!r})") + + class _NoRedirectHTTPHandler(AsyncHTTPHandler): """AsyncHTTPHandler that disables redirects on every client it creates (SSRF guard).""" @@ -50,7 +61,7 @@ class _NoRedirectHTTPHandler(AsyncHTTPHandler): class FocusTernaryDestination(FocusDestination): - """Upload FOCUS CSV exports to the Ternary cost-ingestion API.""" + """Upload one UTC day of FOCUS CSV per delivery to Ternary through a signed GCS URL.""" def __init__( self, @@ -92,98 +103,45 @@ class FocusTernaryDestination(FocusDestination): time_window: FocusTimeWindow, filename: str, ) -> None: - """Upload CSV to the Ternary API, chunking an oversized backfill; aborts on the first failed chunk.""" + """Upload one UTC day: request a signed URL for the window's start day, then PUT the CSV to it.""" if not content: verbose_logger.debug("Ternary destination: empty content, skipping upload") return - chunks: Final = self._split_into_chunks(content) - upload_id: Final = uuid4().hex - total: Final = len(chunks) - - client: Final = _NoRedirectHTTPHandler(timeout=TERNARY_UPLOAD_TIMEOUT_SECONDS) + day: Final = time_window.start_time.strftime("%Y-%m-%d") + handler: Final = _NoRedirectHTTPHandler(timeout=TERNARY_UPLOAD_TIMEOUT_SECONDS) try: - for index, chunk in enumerate(chunks): - await self._upload_csv( - client, - chunk, - filename if total == 1 else f"{filename}.part{index + 1}", - upload_id=upload_id, - chunk_index=index, - chunk_total=total, + body: Final = await self._request_upload_url(handler.client, day) + url: Final = body.get("url") + headers: Final = body.get("headers") + if not url or headers is None: + raise RuntimeError(f"Ternary destination: upload-url response for {day} is missing 'url' or 'headers'") + if body.get("method") != "PUT": + raise RuntimeError(f"Ternary destination: unsupported upload method {body.get('method')!r} for {day}") + _require_gcs_url(url) + uploaded: Final = await handler.client.put( + url, content=content, headers=headers, timeout=TERNARY_UPLOAD_TIMEOUT_SECONDS + ) + if uploaded.status_code >= 400: + raise RuntimeError( + f"Ternary destination: upload for {day} failed ({uploaded.status_code}): {uploaded.text[:200]}" ) finally: - await client.close() + await handler.close() - def _split_into_chunks(self, content: bytes) -> Sequence[bytes]: - """Split CSV bytes into row/byte-limited chunks; an unsplittable oversized row is skipped with a warning.""" - newline_count: Final = content.count(b"\n") - if len(content) <= TERNARY_MAX_BYTES_PER_UPLOAD and newline_count <= TERNARY_MAX_ROWS_PER_UPLOAD: - return (content,) + verbose_logger.debug("Ternary destination: uploaded %d bytes for %s (%s)", len(content), day, filename) - rows: Final = tuple(tuple(row) for row in csv.reader(io.StringIO(content.decode("utf-8")))) - if len(rows) <= 1: - return (content,) - header: Final = rows[0] - data_rows: Final = rows[1:] - header_bytes: Final = len(_encode_csv((header,))) - - chunks: Final[list[bytes]] = [] # mutable-ok: local accumulator, appended to and returned frozen - current: Final[list[Sequence[str]]] = [] # mutable-ok: rows buffered for the in-progress chunk - current_size = header_bytes # rebind-ok: running byte tally advanced across the loop - - for row in data_rows: - row_bytes = len(_encode_csv((header, row))) - header_bytes - if header_bytes + row_bytes > TERNARY_MAX_BYTES_PER_UPLOAD: - verbose_logger.warning( - "Ternary destination: skipping a single row of %d bytes, over the %d-byte upload limit", - row_bytes, - TERNARY_MAX_BYTES_PER_UPLOAD, - ) - continue - exceeds_rows = len(current) >= TERNARY_MAX_ROWS_PER_UPLOAD - exceeds_bytes = current_size + row_bytes > TERNARY_MAX_BYTES_PER_UPLOAD - if current and (exceeds_rows or exceeds_bytes): - chunks.append(_encode_csv((header, *current))) - current.clear() - current_size = header_bytes - current.append(row) - current_size += row_bytes - - if current: - chunks.append(_encode_csv((header, *current))) - return tuple(chunks) - - async def _upload_csv( - self, - client: AsyncHTTPHandler, - csv_bytes: bytes, - filename: str, - *, - upload_id: str, - chunk_index: int, - chunk_total: int, - ) -> None: - url: Final = f"{self.base_url}/external-cost-sources/v1/{quote(self.connection_id, safe='')}/focus" - headers: Final = { # mutable-ok: request headers handed to the HTTP client - "Authorization": f"Bearer {self.api_key}", - "X-Ternary-Upload-Id": upload_id, - "X-Ternary-Chunk-Index": str(chunk_index), - "X-Ternary-Chunk-Total": str(chunk_total), - } - - await client.post( - url, - headers=headers, - files={"csv": (filename, csv_bytes, "text/csv")}, # mutable-ok: multipart payload for the client + async def _request_upload_url(self, client: httpx.AsyncClient, day: str) -> TernaryUploadUrlBody: + endpoint: Final = f"{self.base_url}/external-cost-sources/v1/{quote(self.connection_id, safe='')}/upload-url" + response: Final = await client.post( + endpoint, + headers={"Authorization": f"Bearer {self.api_key}"}, # mutable-ok: request headers for the client + json={"day": day}, # mutable-ok: JSON request body for the client timeout=TERNARY_UPLOAD_TIMEOUT_SECONDS, ) - - verbose_logger.debug( - "Ternary destination: uploaded %d bytes (%s, upload_id=%s, chunk %d/%d)", - len(csv_bytes), - filename, - upload_id, - chunk_index, - chunk_total, - ) + if response.status_code >= 400: + raise RuntimeError( + f"Ternary destination: upload-url request for {day} failed ({response.status_code}): " + f"{response.text[:200]}" + ) + return _upload_url_body(response) diff --git a/litellm/integrations/ternary/ternary_logger.py b/litellm/integrations/ternary/ternary_logger.py index f3cb8dd46ce..4a29eafd9ae 100644 --- a/litellm/integrations/ternary/ternary_logger.py +++ b/litellm/integrations/ternary/ternary_logger.py @@ -4,7 +4,8 @@ from __future__ import annotations import json import os -from datetime import datetime, timedelta, timezone +from collections.abc import Sequence +from datetime import date, datetime, timedelta, timezone from typing import TYPE_CHECKING, Final, TypeAlias import litellm @@ -91,6 +92,44 @@ def _drop_days_before(data: pl.DataFrame, floor: datetime) -> pl.DataFrame: return kept +def _canonical_day(raw: object) -> str | None: + """Return the zero-padded YYYY-MM-DD for a `date` value, or None when it is missing or unparseable.""" + if isinstance(raw, date): + return raw.isoformat() + if not isinstance(raw, str): + return None + try: + parsed: Final = datetime.strptime(raw, "%Y-%m-%d").replace(tzinfo=timezone.utc) + except ValueError: + return None + return parsed.date().isoformat() + + +def _partition_by_day(data: pl.DataFrame) -> Sequence[tuple[str, pl.DataFrame]]: + """Split source rows by their UTC `date` (one upload per day); rows with no parseable date are dropped.""" + if "date" not in data.columns: + verbose_logger.warning("Ternary export: source data has no `date` column; %d row(s) not exported", data.height) + return () + days: Final[list[tuple[str, pl.DataFrame]]] = [] # mutable-ok: per-day parts, sorted and returned as a tuple + undated = 0 # rebind-ok: running count of rows dropped for an unusable `date` + for key, part in data.partition_by("date", as_dict=True, maintain_order=True).items(): + partition_key: object = key[0] # pyright: ignore[reportAny] # polars partition keys are stubbed as Any + day = _canonical_day(partition_key) + if day is None: + undated += part.height + continue + days.append((day, part)) + if undated: + verbose_logger.warning("Ternary export: dropped %d row(s) with a missing or unparseable `date`", undated) + return tuple(sorted(days, key=lambda entry: entry[0])) + + +def _day_window(day: str, frequency: str) -> FocusTimeWindow: + """The single UTC day a per-day upload covers.""" + start: Final = datetime.strptime(day, "%Y-%m-%d").replace(tzinfo=timezone.utc) + return FocusTimeWindow(start_time=start, end_time=start + timedelta(days=1), frequency=frequency) + + def _parse_interval(raw: str | int | None) -> int | None: """Parse the export interval-seconds override; a non-numeric value is ignored.""" if raw is None: @@ -127,14 +166,25 @@ class _TernaryExportEngine(FocusExportEngine): if data.is_empty(): verbose_logger.debug("Ternary export: no usage data for window %s", window) return + failed: Final[list[str]] = [] # mutable-ok: days whose upload failed, reported once after all days run + for day, day_data in _partition_by_day(data): + try: + await self._deliver_day(day_data, _day_window(day, window.frequency)) + except Exception as e: # noqa: BLE001 # one day's failure must not block the rest; logged and re-raised + verbose_logger.warning("Ternary export: upload for %s failed: %s", day, e) + failed.append(day) + if failed: + raise RuntimeError(f"Ternary export: {len(failed)} day(s) failed to upload: {', '.join(failed)}") + + async def _deliver_day(self, data: pl.DataFrame, window: FocusTimeWindow) -> None: transformed: Final = self._transformer.transform(data) if transformed.is_empty(): - verbose_logger.debug("Ternary export: normalized data empty for window %s", window) + verbose_logger.debug("Ternary export: normalized data empty for %s", window.start_time.date()) return enriched: Final = _merge_token_tags(transformed, data) payload: Final = self._serializer.serialize(enriched) if not payload: - verbose_logger.debug("Ternary export: serializer returned empty payload") + verbose_logger.debug("Ternary export: serializer returned empty payload for %s", window.start_time.date()) return await self._destination.deliver( content=payload, diff --git a/tests/unit/integrations/focus/test_ternary_destination.py b/tests/unit/integrations/focus/test_ternary_destination.py index d83c98388de..289834ab560 100644 --- a/tests/unit/integrations/focus/test_ternary_destination.py +++ b/tests/unit/integrations/focus/test_ternary_destination.py @@ -2,15 +2,12 @@ from __future__ import annotations -import csv -import io from datetime import datetime, timedelta, timezone from typing import Any from unittest.mock import AsyncMock, MagicMock, patch import pytest -import litellm.integrations.focus.destinations.ternary_destination as td from litellm.integrations.focus.destinations.base import FocusTimeWindow from litellm.integrations.focus.destinations.ternary_destination import ( TERNARY_UPLOAD_TIMEOUT_SECONDS, @@ -37,64 +34,55 @@ def _config(**overrides: Any) -> dict[str, Any]: return base -def _capturing_client() -> tuple[MagicMock, list[dict[str, Any]]]: - """Return a mock client plus a list that records each post() call.""" - calls: list[dict[str, Any]] = [] +_SIGNED_URL = "https://storage.googleapis.com/ter-ecs-t-us/days/cloud-1/2024-01-02.csv?X-Goog-Signature=abc" +_SIGNED_HEADERS = { + "Content-Type": "text/csv", + "x-goog-meta-ecs-provider": "litellm", + "x-goog-meta-ecs-connection-id": "conn-1234", + "x-goog-content-length-range": "0,1073741824", +} - mock_response = AsyncMock() - mock_response.raise_for_status = lambda: None - mock_client = MagicMock() - mock_client.close = AsyncMock() +def _response(status: int = 200, body: Any = None) -> MagicMock: + response = MagicMock() + response.status_code = status + response.json = MagicMock(return_value=body) + response.text = "" if body is None else str(body) + return response + + +def _upload_url_body(**overrides: Any) -> dict[str, Any]: + body = { + "url": _SIGNED_URL, + "method": "PUT", + "headers": dict(_SIGNED_HEADERS), + "expiresAt": "2024-01-02T05:15:00Z", + "day": "2024-01-02", + } + body.update(overrides) + return body + + +def _handler( + *, url_response: MagicMock | None = None, put_response: MagicMock | None = None +) -> tuple[MagicMock, list[dict[str, Any]], list[dict[str, Any]]]: + """Return a mock handler whose client records each upload-url POST and signed-URL PUT.""" + posts: list[dict[str, Any]] = [] + puts: list[dict[str, Any]] = [] async def capture_post(url, **kwargs): - calls.append({"url": url, **kwargs}) - return mock_response + posts.append({"url": url, **kwargs}) + return url_response if url_response is not None else _response(body=_upload_url_body()) - mock_client.post = capture_post - return mock_client, calls + async def capture_put(url, **kwargs): + puts.append({"url": url, **kwargs}) + return put_response if put_response is not None else _response() - -def _uploaded_part(call: dict[str, Any]) -> tuple[str, bytes]: - field = call["files"]["csv"] - return field[0], field[1] - - -def _rows(content: bytes) -> list[list[str]]: - return list(csv.reader(io.StringIO(content.decode("utf-8")))) - - -def test_should_require_api_key(): - with pytest.raises(ValueError, match="api_key"): - FocusTernaryDestination(prefix="exports", config={"connection_id": "c", "base_url": "u"}) - - -def test_should_require_connection_id(): - with pytest.raises(ValueError, match="connection_id"): - FocusTernaryDestination(prefix="exports", config={"api_key": "k", "base_url": "u"}) - - -def test_should_require_base_url(): - with pytest.raises(ValueError, match="base_url"): - FocusTernaryDestination(prefix="exports", config={"api_key": "k", "connection_id": "c"}) - - -@pytest.mark.parametrize("bad_id", ["a/b", "..", "has space", "tab\tid"]) -def test_should_reject_connection_id_that_could_reroute_the_path(bad_id): - with pytest.raises(ValueError, match="connection_id"): - FocusTernaryDestination(prefix="exports", config=_config(connection_id=bad_id)) - - -def test_should_initialize_with_valid_config(): - dest = FocusTernaryDestination(prefix="exports", config=_config()) - assert dest.api_key == "test-api-key" - assert dest.connection_id == "conn-1234" - assert dest.base_url == "https://ternary.test" - - -def test_should_use_custom_base_url_and_strip_trailing_slash(): - dest = FocusTernaryDestination(prefix="exports", config=_config(base_url="http://localhost:8080/")) - assert dest.base_url == "http://localhost:8080" + handler = MagicMock() + handler.close = AsyncMock() + handler.client.post = capture_post + handler.client.put = capture_put + return handler, posts, puts @pytest.mark.asyncio @@ -105,26 +93,35 @@ async def test_should_skip_empty_content(): @pytest.mark.asyncio -async def test_should_upload_to_correct_url_with_auth_and_upload_headers(): +async def test_should_request_a_signed_url_for_the_window_day_then_put_the_csv(): dest = FocusTernaryDestination(prefix="exports", config=_config()) - mock_client, calls = _capturing_client() + handler, posts, puts = _handler() - with patch(MOCK_TARGET, return_value=mock_client): + with patch(MOCK_TARGET, return_value=handler): await dest.deliver(content=b"header\nrow1\n", time_window=_window(), filename="usage.csv") - assert len(calls) == 1 - call = calls[0] - assert call["url"] == "https://ternary.test/external-cost-sources/v1/conn-1234/focus" - assert call["headers"]["Authorization"] == "Bearer test-api-key" - assert call["headers"]["X-Ternary-Chunk-Index"] == "0" - assert call["headers"]["X-Ternary-Chunk-Total"] == "1" - assert call["headers"]["X-Ternary-Upload-Id"] - filename, body = _uploaded_part(call) - assert filename == "usage.csv" - assert body == b"header\nrow1\n" - assert call["files"]["csv"][2] == "text/csv" - assert call["timeout"] == TERNARY_UPLOAD_TIMEOUT_SECONDS - mock_client.close.assert_awaited_once() + assert len(posts) == 1 + assert posts[0]["url"] == "https://ternary.test/external-cost-sources/v1/conn-1234/upload-url" + assert posts[0]["headers"] == {"Authorization": "Bearer test-api-key"} + assert posts[0]["json"] == {"day": "2024-01-02"} + assert posts[0]["timeout"] == TERNARY_UPLOAD_TIMEOUT_SECONDS + assert len(puts) == 1 + assert puts[0]["url"] == _SIGNED_URL + assert puts[0]["content"] == b"header\nrow1\n" + assert puts[0]["headers"] == _SIGNED_HEADERS + assert "Authorization" not in puts[0]["headers"] + handler.close.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_should_key_the_upload_on_the_window_start_day_not_its_end(): + dest = FocusTernaryDestination(prefix="exports", config=_config()) + start = datetime(2024, 1, 2, tzinfo=timezone.utc) + one_day = FocusTimeWindow(start_time=start, end_time=start + timedelta(days=1), frequency="daily") + handler, posts, _ = _handler() + with patch(MOCK_TARGET, return_value=handler): + await dest.deliver(content=b"h\nr\n", time_window=one_day, filename="usage.csv") + assert posts[0]["json"] == {"day": "2024-01-02"} @pytest.mark.asyncio @@ -142,128 +139,94 @@ async def test_no_redirect_handler_disables_redirects_including_retry_client(): @pytest.mark.asyncio async def test_should_url_encode_the_connection_id(): dest = FocusTernaryDestination(prefix="exports", config=_config(connection_id="conn+id~ok")) - mock_client, calls = _capturing_client() - with patch(MOCK_TARGET, return_value=mock_client): + handler, posts, _ = _handler() + with patch(MOCK_TARGET, return_value=handler): await dest.deliver(content=b"h\nr\n", time_window=_window(), filename="usage.csv") - assert calls[0]["url"].endswith("/external-cost-sources/v1/conn%2Bid~ok/focus") + assert posts[0]["url"].endswith("/external-cost-sources/v1/conn%2Bid~ok/upload-url") @pytest.mark.asyncio async def test_should_pass_tags_column_through_unstripped(): """The Ternary sink must not drop/strip any columns (forwards Tags as-is).""" dest = FocusTernaryDestination(prefix="exports", config=_config()) - mock_client, calls = _capturing_client() + handler, _, puts = _handler() content = b'ServiceName,Tags,x_unknown\nfoo,"{""team_id"": ""t1""}",keepme\n' - with patch(MOCK_TARGET, return_value=mock_client): + with patch(MOCK_TARGET, return_value=handler): await dest.deliver(content=content, time_window=_window(), filename="usage.csv") - _, body = _uploaded_part(calls[0]) - assert body == content + assert puts[0]["content"] == content @pytest.mark.asyncio -async def test_should_chunk_by_row_count_with_stable_upload_id(monkeypatch): - monkeypatch.setattr(td, "TERNARY_MAX_ROWS_PER_UPLOAD", 2) +@pytest.mark.parametrize("status", [400, 401, 404, 422, 429, 500]) +async def test_should_raise_and_skip_the_put_when_the_upload_url_request_fails(status): dest = FocusTernaryDestination(prefix="exports", config=_config()) - - content = b"ServiceName\n" + b"\n".join([b"x"] * 5) + b"\n" - mock_client, calls = _capturing_client() - with patch(MOCK_TARGET, return_value=mock_client): - await dest.deliver(content=content, time_window=_window(), filename="usage.csv") - - assert len(calls) == 3 - total = "3" - assert {c["headers"]["X-Ternary-Upload-Id"] for c in calls} == {calls[0]["headers"]["X-Ternary-Upload-Id"]} - for i, call in enumerate(calls): - assert call["headers"]["X-Ternary-Chunk-Index"] == str(i) - assert call["headers"]["X-Ternary-Chunk-Total"] == total - filename, body = _uploaded_part(call) - assert filename == f"usage.csv.part{i + 1}" - assert len(_rows(body)) - 1 <= 2 + handler, _, puts = _handler(url_response=_response(status=status, body={"message": "nope"})) + with patch(MOCK_TARGET, return_value=handler): + with pytest.raises(RuntimeError, match=f"upload-url request for 2024-01-02 failed \\({status}\\)"): + await dest.deliver(content=b"h\nr\n", time_window=_window(), filename="usage.csv") + assert puts == [] + handler.close.assert_awaited_once() @pytest.mark.asyncio -async def test_should_chunk_by_bytes(monkeypatch): - monkeypatch.setattr(td, "TERNARY_MAX_BYTES_PER_UPLOAD", 200) +@pytest.mark.parametrize("status", [400, 403, 503]) +async def test_should_raise_when_the_signed_put_fails(status): dest = FocusTernaryDestination(prefix="exports", config=_config()) - - row = b"a" * 40 + b"," + b"b" * 40 - content = b"ServiceName,BilledCost\n" + b"\n".join([row] * 20) + b"\n" - mock_client, calls = _capturing_client() - with patch(MOCK_TARGET, return_value=mock_client): - await dest.deliver(content=content, time_window=_window(), filename="usage.csv") - - assert len(calls) > 1 - for call in calls: - _, body = _uploaded_part(call) - assert len(body) <= 200 + handler, _, _ = _handler(put_response=_response(status=status, body="SignatureDoesNotMatch")) + with patch(MOCK_TARGET, return_value=handler): + with pytest.raises(RuntimeError, match=f"upload for 2024-01-02 failed \\({status}\\)"): + await dest.deliver(content=b"h\nr\n", time_window=_window(), filename="usage.csv") + handler.close.assert_awaited_once() @pytest.mark.asyncio -async def test_should_not_mangle_a_quoted_field_containing_a_newline(monkeypatch): - monkeypatch.setattr(td, "TERNARY_MAX_ROWS_PER_UPLOAD", 1) +@pytest.mark.parametrize( + "url", + [ + "http://storage.googleapis.com/b/days/c/2024-01-02.csv", + "https://evil.example.com/b/days/c/2024-01-02.csv", + "https://storage.googleapis.com.evil.example.com/b/o", + "https://169.254.169.254/computeMetadata/v1/", + ], +) +async def test_should_refuse_to_put_to_a_non_gcs_url(url): dest = FocusTernaryDestination(prefix="exports", config=_config()) - - content = b'ServiceName,Tags\nsvc1,"line1\nline2"\nsvc2,"{""k"":""v""}"\n' - mock_client, calls = _capturing_client() - with patch(MOCK_TARGET, return_value=mock_client): - await dest.deliver(content=content, time_window=_window(), filename="usage.csv") - - assert len(calls) == 2 - _, first = _uploaded_part(calls[0]) - rows = _rows(first) - assert rows[0] == ["ServiceName", "Tags"] - assert rows[1] == ["svc1", "line1\nline2"] + handler, _, puts = _handler(url_response=_response(body=_upload_url_body(url=url))) + with patch(MOCK_TARGET, return_value=handler): + with pytest.raises(ValueError, match="HTTPS GCS URL"): + await dest.deliver(content=b"h\nr\n", time_window=_window(), filename="usage.csv") + assert puts == [] @pytest.mark.asyncio -async def test_should_skip_an_unsplittable_oversized_row_and_upload_the_rest(monkeypatch): - # An unsplittable row over the byte limit is dropped (warned); the rest of the day still lands. - monkeypatch.setattr(td, "TERNARY_MAX_BYTES_PER_UPLOAD", 80) +async def test_should_accept_a_virtual_hosted_gcs_url(): dest = FocusTernaryDestination(prefix="exports", config=_config()) - content = b"ServiceName\nkeep-me\n" + b"z" * 300 + b"\n" - mock_client, calls = _capturing_client() - with patch(MOCK_TARGET, return_value=mock_client): - await dest.deliver(content=content, time_window=_window(), filename="usage.csv") - body = b"".join(_uploaded_part(call)[1] for call in calls) - assert b"keep-me" in body - assert b"zzz" not in body + url = "https://ter-ecs-t-us.storage.googleapis.com/days/c/2024-01-02.csv?X-Goog-Signature=abc" + handler, _, puts = _handler(url_response=_response(body=_upload_url_body(url=url))) + with patch(MOCK_TARGET, return_value=handler): + await dest.deliver(content=b"h\nr\n", time_window=_window(), filename="usage.csv") + assert puts[0]["url"] == url @pytest.mark.asyncio -async def test_should_not_create_a_client_when_chunking_fails(monkeypatch): - # Chunking runs before the client is created, so a failure there leaks no client. - monkeypatch.setattr(td, "TERNARY_MAX_BYTES_PER_UPLOAD", 10) +@pytest.mark.parametrize( + ("overrides", "match"), + [ + ({"url": ""}, "missing 'url' or 'headers'"), + ({"headers": None}, "missing 'url' or 'headers'"), + ({"method": "POST"}, "unsupported upload method 'POST'"), + ], +) +async def test_should_reject_an_incomplete_upload_url_response(overrides, match): dest = FocusTernaryDestination(prefix="exports", config=_config()) - content = b"ServiceName\n" + b"\xff\xfe" * 20 + b"\n" - with patch(MOCK_TARGET, side_effect=AssertionError("client created before chunking succeeded")): - with pytest.raises(UnicodeDecodeError): - await dest.deliver(content=content, time_window=_window(), filename="usage.csv") - - -@pytest.mark.asyncio -async def test_should_abort_on_first_chunk_failure(monkeypatch): - monkeypatch.setattr(td, "TERNARY_MAX_ROWS_PER_UPLOAD", 1) - dest = FocusTernaryDestination(prefix="exports", config=_config()) - - content = b"ServiceName\n" + b"\n".join([b"x"] * 3) + b"\n" - - attempts: list[str] = [] - mock_client = MagicMock() - mock_client.close = AsyncMock() - - async def failing_post(url, **kwargs): - attempts.append(kwargs["headers"]["X-Ternary-Chunk-Index"]) - raise RuntimeError("boom") - - mock_client.post = failing_post - - with patch(MOCK_TARGET, return_value=mock_client): - with pytest.raises(RuntimeError, match="boom"): - await dest.deliver(content=content, time_window=_window(), filename="usage.csv") - - assert attempts == ["0"] + body = _upload_url_body(**overrides) + handler, _, puts = _handler(url_response=_response(body=body)) + with patch(MOCK_TARGET, return_value=handler): + with pytest.raises(RuntimeError, match=match): + await dest.deliver(content=b"h\nr\n", time_window=_window(), filename="usage.csv") + assert puts == [] @pytest.mark.parametrize("url", ["http://api.ternary.app", "http://evil.example.com:8080", "ftp://ternary.test"]) @@ -278,19 +241,6 @@ def test_should_accept_https_or_loopback_base_url(url): assert dest.base_url == url.rstrip("/") -@pytest.mark.asyncio -async def test_should_return_header_only_content_untouched(monkeypatch): - monkeypatch.setattr(td, "TERNARY_MAX_BYTES_PER_UPLOAD", 10) - dest = FocusTernaryDestination(prefix="exports", config=_config()) - content = b"HeaderOnlyLineWithNoDataRows" - mock_client, calls = _capturing_client() - with patch(MOCK_TARGET, return_value=mock_client): - await dest.deliver(content=content, time_window=_window(), filename="usage.csv") - assert len(calls) == 1 - _, body = _uploaded_part(calls[0]) - assert body == content - - def test_factory_creates_ternary_destination_from_config(): from litellm.integrations.focus.destinations.factory import FocusDestinationFactory diff --git a/tests/unit/integrations/ternary/test_ternary_logger.py b/tests/unit/integrations/ternary/test_ternary_logger.py index 896d97945d4..2e76490f7ec 100644 --- a/tests/unit/integrations/ternary/test_ternary_logger.py +++ b/tests/unit/integrations/ternary/test_ternary_logger.py @@ -2,7 +2,10 @@ from __future__ import annotations +import csv +import io import json +from datetime import timedelta from unittest.mock import AsyncMock, MagicMock import polars as pl @@ -13,8 +16,11 @@ from litellm.integrations.focus.focus_logger import FocusLogger from litellm.integrations.ternary.ternary_logger import ( TERNARY_USAGE_DATA_JOB_NAME, TernaryLogger, + _canonical_day, + _day_window, _drop_days_before, _merge_token_tags, + _partition_by_day, _TernaryExportEngine, ) @@ -282,12 +288,13 @@ def _fake_engine(*, transformed: pl.DataFrame, payload: bytes) -> MagicMock: engine._serializer.serialize = MagicMock(return_value=payload) engine._destination.deliver = AsyncMock() engine._build_filename = MagicMock(return_value="usage.csv") + engine._deliver_day = lambda data, window: _TernaryExportEngine._deliver_day(engine, data, window) return engine @pytest.mark.asyncio async def test_transform_enrich_deliver_enriches_tags_then_delivers(): - data = pl.DataFrame({"Tags": [json.dumps({"team_id": "t1"})], "prompt_tokens": [7]}) + data = pl.DataFrame({"date": ["2026-09-03"], "Tags": [json.dumps({"team_id": "t1"})], "prompt_tokens": [7]}) transformed = pl.DataFrame({"Tags": [json.dumps({"team_id": "t1"})]}) engine = _fake_engine(transformed=transformed, payload=b"csv-bytes") @@ -299,6 +306,7 @@ async def test_transform_enrich_deliver_enriches_tags_then_delivers(): kwargs = engine._destination.deliver.await_args.kwargs assert kwargs["content"] == b"csv-bytes" assert kwargs["filename"] == "usage.csv" + assert kwargs["time_window"].start_time.isoformat() == "2026-09-03T00:00:00+00:00" @pytest.mark.asyncio @@ -308,6 +316,120 @@ async def test_transform_enrich_deliver_skips_empty_data(): engine._destination.deliver.assert_not_awaited() +def _real_engine() -> MagicMock: + from litellm.integrations.focus.serializers import FocusCsvSerializer + from litellm.integrations.focus.transformer import FocusTransformer + + engine = MagicMock() + engine._transformer = FocusTransformer() + engine._serializer = FocusCsvSerializer() + engine._destination.deliver = AsyncMock() + engine._build_filename = MagicMock(side_effect=lambda w: f"usage_{w.start_time:%Y%m%d}.csv") + engine._deliver_day = lambda data, window: _TernaryExportEngine._deliver_day(engine, data, window) + return engine + + +def _source_rows(dates: list[str]) -> pl.DataFrame: + n = len(dates) + return pl.DataFrame( + { + "date": dates, + "spend": [float(i + 1) for i in range(n)], + "api_key": [f"k{i}" for i in range(n)], + "api_key_alias": [f"a{i}" for i in range(n)], + "model": [f"m{i}" for i in range(n)], + "model_group": [f"g{i}" for i in range(n)], + "custom_llm_provider": ["anthropic"] * n, + "team_id": [f"t{i}" for i in range(n)], + "team_alias": [f"T{i}" for i in range(n)], + "api_requests": [1] * n, + "prompt_tokens": [10 * (i + 1) for i in range(n)], + "completion_tokens": [i for i in range(n)], + "cache_read_input_tokens": [0] * n, + "cache_creation_input_tokens": [0] * n, + } + ) + + +@pytest.mark.asyncio +async def test_delivers_one_upload_per_utc_day_through_the_real_transformer(): + engine = _real_engine() + data = _source_rows(["2026-10-06", "2026-10-05", "2026-10-06"]) + + await _TernaryExportEngine._deliver_enriched(engine, data, _window()) + + calls = engine._destination.deliver.await_args_list + assert [c.kwargs["time_window"].start_time.date().isoformat() for c in calls] == ["2026-10-05", "2026-10-06"] + for call, expected_rows, expected_day in ((calls[0], 1, "2026-10-05"), (calls[1], 2, "2026-10-06")): + window = call.kwargs["time_window"] + assert window.end_time - window.start_time == timedelta(days=1) + assert window.frequency == "daily" + rows = list(csv.DictReader(io.StringIO(call.kwargs["content"].decode("utf-8")))) + assert len(rows) == expected_rows + assert {row["ChargePeriodStart"] for row in rows} == {f"{expected_day}T00:00:00Z"} + assert all(json.loads(row["Tags"])["prompt_tokens"] for row in rows) + + +@pytest.mark.asyncio +async def test_a_failed_day_does_not_block_the_remaining_days(): + engine = _real_engine() + data = _source_rows(["2026-10-04", "2026-10-05", "2026-10-06"]) + attempted: list[str] = [] + + async def deliver(*, content, time_window, filename): + day = time_window.start_time.date().isoformat() + attempted.append(day) + if day == "2026-10-05": + raise RuntimeError("upload-url request for 2026-10-05 failed (500)") + + engine._destination.deliver = deliver + + with pytest.raises(RuntimeError, match=r"1 day\(s\) failed to upload: 2026-10-05"): + await _TernaryExportEngine._deliver_enriched(engine, data, _window()) + + assert attempted == ["2026-10-04", "2026-10-05", "2026-10-06"] + + +def test_partition_by_day_sorts_canonicalizes_and_drops_undated_rows(caplog): + import logging + from datetime import date + + from litellm._logging import verbose_logger + + data = pl.DataFrame( + {"date": ["2026-10-06", None, "2026-10-05", "not-a-date", "2026-10-06"], "spend": [1.0, 2.0, 3.0, 4.0, 5.0]} + ) + with caplog.at_level(logging.WARNING, logger=verbose_logger.name): + parts = _partition_by_day(data) + + assert [(day, part["spend"].to_list()) for day, part in parts] == [ + ("2026-10-05", [3.0]), + ("2026-10-06", [1.0, 5.0]), + ] + assert "dropped 2 row(s)" in caplog.text + assert _canonical_day(date(2026, 1, 2)) == "2026-01-02" + assert _canonical_day("2026-1-2") == "2026-01-02" + assert _canonical_day(20260102) is None + + +def test_partition_by_day_warns_and_exports_nothing_without_a_date_column(caplog): + import logging + + from litellm._logging import verbose_logger + + with caplog.at_level(logging.WARNING, logger=verbose_logger.name): + parts = _partition_by_day(pl.DataFrame({"spend": [1.0, 2.0]})) + assert parts == () + assert "no `date` column; 2 row(s) not exported" in caplog.text + + +def test_day_window_covers_exactly_one_utc_day(): + window = _day_window("2026-10-06", "interval") + assert window.start_time.isoformat() == "2026-10-06T00:00:00+00:00" + assert window.end_time.isoformat() == "2026-10-07T00:00:00+00:00" + assert window.frequency == "interval" + + def test_pod_lock_key_is_ternary_specific(): assert TERNARY_USAGE_DATA_JOB_NAME == "ternary_export_usage_data" @@ -350,7 +472,7 @@ async def test_export_all_delivers_with_all_window(): @pytest.mark.asyncio async def test_transform_enrich_deliver_skips_when_transform_empty(): - data = pl.DataFrame({"Tags": [json.dumps({"team_id": "t"})], "prompt_tokens": [1]}) + data = pl.DataFrame({"date": ["2026-09-03"], "Tags": [json.dumps({"team_id": "t"})], "prompt_tokens": [1]}) engine = _fake_engine(transformed=pl.DataFrame(), payload=b"") await _TernaryExportEngine._deliver_enriched(engine, data, _window()) engine._destination.deliver.assert_not_awaited() @@ -358,7 +480,7 @@ async def test_transform_enrich_deliver_skips_when_transform_empty(): @pytest.mark.asyncio async def test_transform_enrich_deliver_skips_when_payload_empty(): - data = pl.DataFrame({"Tags": [json.dumps({"team_id": "t"})]}) + data = pl.DataFrame({"date": ["2026-09-03"], "Tags": [json.dumps({"team_id": "t"})]}) transformed = pl.DataFrame({"Tags": [json.dumps({"team_id": "t"})]}) engine = _fake_engine(transformed=transformed, payload=b"") await _TernaryExportEngine._deliver_enriched(engine, data, _window())