mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-13 03:59:12 +00:00
feat(focus): upload one UTC day per Ternary signed-URL PUT
Replace the multipart POST with chunking by a per-day signed-URL upload, matching the Mavvrik destination's pattern. The logger partitions the raw usage rows on `date` (dropping unparseable dates with a warning), then transforms, enriches and delivers each day separately; one day's failure is logged and the remaining days still upload, with a single error raised at the end. The destination requests a signed URL for the day, validates it is an HTTPS GCS URL, and PUTs the CSV with exactly the signed headers. Removes the 30 MiB / 100k-row chunking, the X-Ternary-* upload headers and the .partN filenames; a re-sent day overwrites its object, so retries stay idempotent without them.
This commit is contained in:
Aaron Sproul
parent
284e2cfac1
commit
32bf19bbef
4 changed files with 359 additions and 279 deletions
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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())
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue