diff --git a/litellm/integrations/pointfive/upload_client.py b/litellm/integrations/pointfive/upload_client.py new file mode 100644 index 00000000000..686569988c2 --- /dev/null +++ b/litellm/integrations/pointfive/upload_client.py @@ -0,0 +1,211 @@ +""" +Uploads one batch to PointFive through a presigned URL. + +The proxy holds no cloud credentials. For every batch it asks the PointFive API for a +single-use presigned URL and PUTs the bytes there, so the same plugin runs unchanged on +AWS, GCP, Azure or on-prem. The server picks the object key, so the proxy never chooses +where its data lands. +""" + +import asyncio +from collections.abc import Awaitable, Callable, Mapping +from typing import Final + +import httpx +from pydantic import BaseModel, Field, ValidationError + +from litellm._logging import verbose_logger +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler +from litellm.types.integrations.pointfive import ( + RETRYABLE_UPLOAD_STATUS_CODES, + PointFiveUploadFailure, + PointFiveUploadTarget, +) + +UPLOAD_URL_QUERY: Final = ( + "query UploadUrl($kind: UploadKind!, $byteCount: Int!) " + "{ uploadUrl(kind: $kind, byteCount: $byteCount) { uploadUrl objectKey } }" +) +PING_QUERY: Final = "query IntegrationPing($kind: UploadKind!) { integrationPing(kind: $kind) }" +UPLOAD_KIND: Final = "LITELLM" + + +class _PresignVariables(BaseModel): + kind: str = UPLOAD_KIND + byte_count: int = Field(serialization_alias="byteCount") + + +class _PresignRequest(BaseModel): + variables: _PresignVariables + query: str = UPLOAD_URL_QUERY + + +class _PingVariables(BaseModel): + kind: str = UPLOAD_KIND + + +class _PingRequest(BaseModel): + variables: _PingVariables = _PingVariables() + query: str = PING_QUERY + + +class _TargetPayload(BaseModel): + upload_url: str = Field(alias="uploadUrl") + object_key: str = Field(alias="objectKey") + + +class _DataPayload(BaseModel): + upload_url: _TargetPayload = Field(alias="uploadUrl") + + +class _ErrorPayload(BaseModel): + message: str = "" + + +class _ErrorsPayload(BaseModel): + """Just the errors array, so it reads any operation's response, whatever its data shape.""" + + errors: tuple[_ErrorPayload, ...] | None = None + + +class _EnvelopePayload(BaseModel): + data: _DataPayload | None = None + + +class PointFiveUploadError(Exception): + """A batch could not be uploaded and the failure is worth retrying.""" + + +def _failure_for(status_code: int, detail: str) -> PointFiveUploadFailure: + return PointFiveUploadFailure(detail, retryable=status_code in RETRYABLE_UPLOAD_STATUS_CODES) + + +def _graphql_error(body: str) -> PointFiveUploadFailure | None: + """GraphQL reports failures inside a 200, so the body has to be read either way.""" + try: + envelope: Final = _ErrorsPayload.model_validate_json(body) + except ValidationError: + return PointFiveUploadFailure("pointfive api returned an unreadable body", retryable=False) + if not envelope.errors: + return None + joined: Final = ", ".join(error.message for error in envelope.errors) + return PointFiveUploadFailure(f"pointfive api rejected the request: {joined}", retryable=False) + + +def _parse_target(body: str) -> PointFiveUploadTarget | PointFiveUploadFailure: + """Read the presigned target out of a GraphQL response body.""" + error: Final = _graphql_error(body) + if error is not None: + return error + + try: + envelope: Final = _EnvelopePayload.model_validate_json(body) + except ValidationError: + return PointFiveUploadFailure("pointfive api returned an unreadable body", retryable=False) + if envelope.data is None: + return PointFiveUploadFailure("pointfive api returned no upload url", retryable=False) + + target: Final = envelope.data.upload_url + return PointFiveUploadTarget(upload_url=target.upload_url, object_key=target.object_key) + + +class PointFiveUploadClient: + """Presigns and uploads one batch at a time.""" + + def __init__( + self, + api_key: str, + api_url: str, + http_client: AsyncHTTPHandler, + max_retries: int, + sleep: Callable[[float], Awaitable[None]] = asyncio.sleep, + ) -> None: + self.api_key: Final = api_key + self.api_url: Final = api_url + self.http_client: Final = http_client + self.max_retries: Final = max_retries + self.sleep: Final = sleep + + async def upload(self, body: bytes) -> str | PointFiveUploadFailure: + """ + Upload one gzipped batch, returning the object key it landed at. + + Every attempt presigns again, so a retry never reuses a URL that has expired or + has already been consumed. + """ + for attempt in range(self.max_retries): + match await self._upload_once(body): + case PointFiveUploadFailure(retryable=True) as failure: + if attempt + 1 >= self.max_retries: + return PointFiveUploadFailure( + f"{failure.detail}, gave up after {self.max_retries} attempts", retryable=True + ) + await self.sleep(float(1 << attempt)) + case outcome: + return outcome + return PointFiveUploadFailure("max_upload_retries must be at least 1", retryable=False) + + async def _upload_once(self, body: bytes) -> str | PointFiveUploadFailure: + target: Final = await self._presign(len(body)) + if isinstance(target, PointFiveUploadFailure): + return target + + rejection: Final = await self._put(target, body) + if rejection is not None: + return rejection + + verbose_logger.debug("pointfive: uploaded %s gzipped bytes to %s", len(body), target.object_key) + return target.object_key + + async def ping(self) -> PointFiveUploadFailure | None: + """Report that the proxy is alive when it has nothing to upload.""" + body: Final = await self._query(_PingRequest().model_dump(by_alias=True)) + if isinstance(body, PointFiveUploadFailure): + return body + return _graphql_error(body) + + async def _presign(self, byte_count: int) -> PointFiveUploadTarget | PointFiveUploadFailure: + """Ask the PointFive API for a presigned URL sized to this batch.""" + body: Final = await self._query( + _PresignRequest(variables=_PresignVariables(byte_count=byte_count)).model_dump(by_alias=True) + ) + if isinstance(body, PointFiveUploadFailure): + return body + return _parse_target(body) + + async def _query(self, request: Mapping[str, object]) -> str | PointFiveUploadFailure: + """POST one GraphQL request to the PointFive API and return its raw body.""" + try: + response: Final = await self.http_client.post( + self.api_url, + json=dict(request), # mutable-ok: AsyncHTTPHandler.post types json as dict + headers={ # mutable-ok: AsyncHTTPHandler.post types headers as dict + "Authorization": f"Bearer {self.api_key}", + "Content-Type": "application/json", + }, + ) + except httpx.HTTPStatusError as e: + return _failure_for(e.response.status_code, f"pointfive api returned {e.response.status_code}") + except Exception as e: # noqa: BLE001 # a transport fault is worth another attempt + return PointFiveUploadFailure(f"pointfive api unreachable: {type(e).__name__}", retryable=True) + + if response is None: + return PointFiveUploadFailure("pointfive api returned no response", retryable=True) + return response.text + + async def _put(self, target: PointFiveUploadTarget, body: bytes) -> PointFiveUploadFailure | None: + """PUT the batch to the presigned URL, which carries its own authorization.""" + try: + await self.http_client.put( + target.upload_url, + data=body, + headers={ # mutable-ok: AsyncHTTPHandler.put types headers as dict + "Content-Type": "application/x-ndjson", + "Content-Encoding": "gzip", + }, + ) + except httpx.HTTPStatusError as e: + return _failure_for(e.response.status_code, f"presigned upload returned {e.response.status_code}") + except Exception as e: # noqa: BLE001 # a transport fault is worth another attempt + return PointFiveUploadFailure(f"presigned upload unreachable: {type(e).__name__}", retryable=True) + return None diff --git a/litellm/types/integrations/pointfive.py b/litellm/types/integrations/pointfive.py new file mode 100644 index 00000000000..78393c66970 --- /dev/null +++ b/litellm/types/integrations/pointfive.py @@ -0,0 +1,44 @@ +from dataclasses import dataclass +from typing import Final + +from pydantic import Field + +from litellm.types.integrations.custom_logger import StandardCustomLoggerInitParams + +RETRYABLE_UPLOAD_STATUS_CODES: Final = frozenset({429, 500, 502, 503, 504}) + +DEFAULT_API_URL: Final = "https://api.pointfive.co/query" + + +class PointFiveInitParams(StandardCustomLoggerInitParams): + """ + Params for initializing a PointFive logger on litellm. + + Defaults trade freshness for fewer, larger uploads: every flush becomes one object, so + the interval is minutes rather than seconds. ``max_batch_bytes`` bounds how much a + single object may hold, which matters most when message logging is left on, since an + unredacted payload is orders of magnitude larger than a redacted one. + """ + + api_key: str | None = None + api_url: str | None = None + batch_size: int = Field(default=10_000, gt=0) + flush_interval: int = Field(default=300, gt=0) + max_batch_bytes: int = Field(default=8 * 1024 * 1024, gt=0) + max_upload_retries: int = Field(default=3, ge=1) + + +@dataclass(frozen=True, slots=True) +class PointFiveUploadTarget: + """A single-use presigned destination for one batch, issued by the PointFive API.""" + + upload_url: str + object_key: str + + +@dataclass(frozen=True, slots=True) +class PointFiveUploadFailure: + """Why a batch could not be uploaded, and whether a later attempt could still succeed.""" + + detail: str + retryable: bool diff --git a/tests/test_litellm/integrations/pointfive/test_upload_client.py b/tests/test_litellm/integrations/pointfive/test_upload_client.py new file mode 100644 index 00000000000..a0cbbc15e3a --- /dev/null +++ b/tests/test_litellm/integrations/pointfive/test_upload_client.py @@ -0,0 +1,285 @@ +import json +from collections.abc import Sequence + +import httpx +import pytest + +from litellm.integrations.pointfive.upload_client import PointFiveUploadClient +from litellm.types.integrations.pointfive import PointFiveUploadFailure + +API_URL = "https://api.pointfive.co/query" +UPLOAD_URL = "https://uploads.example.invalid/some/object.ndjson.gz?signature=sig" +OBJECT_KEY = "some/object.ndjson.gz" +BODY = b"gzipped-bytes" + + +def _presigned(status_code: int = 200) -> httpx.Response: + return _response(status_code, {"data": {"uploadUrl": {"uploadUrl": UPLOAD_URL, "objectKey": OBJECT_KEY}}}) + + +def _response(status_code: int, payload: object) -> httpx.Response: + return httpx.Response(status_code, text=json.dumps(payload)) + + +def _accepted() -> httpx.Response: + return httpx.Response(200, text="") + + +class FakeHTTPClient: + """ + Stands in for AsyncHTTPHandler, including its habit of raising on error statuses. + + Scripted results are consumed in order, and the last one repeats, so a test that + cares about a single behaviour passes a single result. + """ + + def __init__( + self, + presign: Sequence[httpx.Response | Exception] | None = None, + put: Sequence[httpx.Response | Exception] | None = None, + ) -> None: + self.presign = list(presign) if presign else [_presigned()] # mutable-ok: results are consumed by popping + self.put_results = list(put) if put else [_accepted()] # mutable-ok: results are consumed by popping + self.presign_calls: list[dict] = [] + self.put_calls: list[dict] = [] + + async def post(self, url, json=None, headers=None, **_): + self.presign_calls.append({"url": url, "json": json, "headers": headers or {}}) + return _next_result(self.presign, url) + + async def put(self, url, data=None, headers=None, **_): + self.put_calls.append({"url": url, "data": data, "headers": headers or {}}) + return _next_result(self.put_results, url) + + +def _next_result(results: list, url: str) -> httpx.Response: + result = results.pop(0) if len(results) > 1 else results[0] + if isinstance(result, Exception): + raise result + if result.status_code >= 400: + request = httpx.Request("POST", url) + raise httpx.HTTPStatusError("boom", request=request, response=httpx.Response(result.status_code, text="denied")) + return result + + +async def _no_backoff(_seconds: float) -> None: + return None + + +def _client(http_client: FakeHTTPClient, max_retries: int = 3) -> PointFiveUploadClient: + return PointFiveUploadClient( + api_key="p5tu_testkey", + api_url=API_URL, + http_client=http_client, + max_retries=max_retries, + sleep=_no_backoff, + ) + + +@pytest.mark.asyncio +async def test_uploads_the_body_to_the_url_the_api_returned(): + http_client = FakeHTTPClient() + + outcome = await _client(http_client).upload(BODY) + + assert outcome == OBJECT_KEY + assert http_client.put_calls[0]["url"] == UPLOAD_URL + assert http_client.put_calls[0]["data"] == BODY + + +@pytest.mark.asyncio +async def test_presign_request_is_authenticated_and_sized(): + http_client = FakeHTTPClient() + + await _client(http_client).upload(BODY) + + call = http_client.presign_calls[0] + assert call["headers"]["Authorization"] == "Bearer p5tu_testkey" + assert call["json"]["variables"] == {"kind": "LITELLM", "byteCount": len(BODY)} + + +@pytest.mark.asyncio +async def test_no_bearer_token_is_sent_to_the_presigned_url(): + """The URL carries its own authorization, so the api key must not travel with it.""" + http_client = FakeHTTPClient() + + await _client(http_client).upload(BODY) + + assert "Authorization" not in http_client.put_calls[0]["headers"] + + +@pytest.mark.asyncio +async def test_the_object_is_declared_as_gzipped_ndjson(): + http_client = FakeHTTPClient() + + await _client(http_client).upload(BODY) + + assert http_client.put_calls[0]["headers"]["Content-Encoding"] == "gzip" + assert http_client.put_calls[0]["headers"]["Content-Type"] == "application/x-ndjson" + + +@pytest.mark.asyncio +async def test_each_retry_presigns_again(): + """A retry must never reuse a URL that was consumed or has expired.""" + http_client = FakeHTTPClient(put=[httpx.Response(503), _accepted()]) + + outcome = await _client(http_client).upload(BODY) + + assert outcome == OBJECT_KEY + assert len(http_client.presign_calls) == 2 + assert len(http_client.put_calls) == 2 + + +@pytest.mark.asyncio +async def test_retryable_upload_failure_gives_up_after_max_retries(): + http_client = FakeHTTPClient(put=[httpx.Response(503)]) + + outcome = await _client(http_client, max_retries=2).upload(BODY) + + assert outcome == PointFiveUploadFailure("presigned upload returned 503, gave up after 2 attempts", retryable=True) + assert len(http_client.put_calls) == 2 + + +@pytest.mark.asyncio +async def test_rejected_upload_is_not_retried(): + http_client = FakeHTTPClient(put=[httpx.Response(403)]) + + outcome = await _client(http_client).upload(BODY) + + assert outcome == PointFiveUploadFailure("presigned upload returned 403", retryable=False) + assert len(http_client.put_calls) == 1 + + +@pytest.mark.asyncio +async def test_bad_api_key_is_not_retried(): + http_client = FakeHTTPClient(presign=[httpx.Response(401)]) + + outcome = await _client(http_client).upload(BODY) + + assert outcome == PointFiveUploadFailure("pointfive api returned 401", retryable=False) + assert http_client.put_calls == [] + + +@pytest.mark.asyncio +async def test_api_server_error_is_retried(): + http_client = FakeHTTPClient(presign=[httpx.Response(503), _presigned()]) + + outcome = await _client(http_client).upload(BODY) + + assert outcome == OBJECT_KEY + assert len(http_client.presign_calls) == 2 + + +@pytest.mark.asyncio +async def test_unreachable_api_is_retried_then_reported_as_retryable(): + http_client = FakeHTTPClient(presign=(ConnectionError("down"),)) + + outcome = await _client(http_client, max_retries=2).upload(BODY) + + assert isinstance(outcome, PointFiveUploadFailure) + assert outcome.retryable + assert "unreachable" in outcome.detail + assert len(http_client.presign_calls) == 2 + + +@pytest.mark.asyncio +async def test_graphql_errors_are_not_retried(): + http_client = FakeHTTPClient(presign=[_response(200, {"errors": [{"message": "forbidden"}]})]) + + outcome = await _client(http_client).upload(BODY) + + assert outcome == PointFiveUploadFailure("pointfive api rejected the request: forbidden", retryable=False) + assert http_client.put_calls == [] + + +@pytest.mark.asyncio +async def test_malformed_api_body_is_not_retried(): + http_client = FakeHTTPClient(presign=[_response(200, {"data": {"uploadUrl": {"objectKey": "k"}}})]) + + outcome = await _client(http_client).upload(BODY) + + assert outcome == PointFiveUploadFailure("pointfive api returned an unreadable body", retryable=False) + assert http_client.put_calls == [] + + +@pytest.mark.asyncio +async def test_ping_reports_a_live_shipper(): + http_client = FakeHTTPClient(presign=(_response(200, {"data": {"integrationPing": True}}),)) + + failure = await _client(http_client).ping() + + assert failure is None + assert http_client.presign_calls[0]["json"]["variables"] == {"kind": "LITELLM"} + assert "integrationPing" in http_client.presign_calls[0]["json"]["query"] + + +@pytest.mark.asyncio +async def test_ping_surfaces_a_graphql_error(): + """GraphQL answers 200 with an errors array, so the body has to be read.""" + http_client = FakeHTTPClient(presign=(_response(200, {"errors": [{"message": "revoked"}]}),)) + + failure = await _client(http_client).ping() + + assert failure is not None + assert "revoked" in failure.detail + + +@pytest.mark.asyncio +async def test_ping_surfaces_an_unreachable_api(): + http_client = FakeHTTPClient(presign=(ConnectionError("down"),)) + + failure = await _client(http_client).ping() + + assert failure is not None + assert failure.retryable + + +@pytest.mark.asyncio +async def test_a_body_that_is_not_json_is_reported_as_unreadable(): + http_client = FakeHTTPClient(presign=(httpx.Response(200, text="gateway"),)) + + outcome = await _client(http_client).upload(BODY) + + assert outcome == PointFiveUploadFailure("pointfive api returned an unreadable body", retryable=False) + assert http_client.put_calls == [] + + +@pytest.mark.asyncio +async def test_a_response_carrying_neither_data_nor_errors_is_reported(): + http_client = FakeHTTPClient(presign=(_response(200, {"data": None}),)) + + outcome = await _client(http_client).upload(BODY) + + assert outcome == PointFiveUploadFailure("pointfive api returned no upload url", retryable=False) + + +@pytest.mark.asyncio +async def test_no_response_at_all_is_worth_retrying(): + class SilentHTTPClient(FakeHTTPClient): + async def post(self, url, json=None, headers=None, **_): + return None + + outcome = await _client(SilentHTTPClient()).upload(BODY) + + assert outcome == PointFiveUploadFailure( + "pointfive api returned no response, gave up after 3 attempts", retryable=True + ) + + +@pytest.mark.asyncio +async def test_a_transport_fault_on_the_upload_itself_is_retryable(): + http_client = FakeHTTPClient(put=(ConnectionError("reset"),)) + + outcome = await _client(http_client, max_retries=1).upload(BODY) + + assert isinstance(outcome, PointFiveUploadFailure) + assert outcome.retryable + assert "presigned upload unreachable" in outcome.detail + + +@pytest.mark.asyncio +async def test_a_client_that_may_not_try_at_all_says_so(): + """max_upload_retries is validated as >= 1, so this guards the loop against a future zero.""" + outcome = await _client(FakeHTTPClient(), max_retries=0).upload(BODY) + + assert outcome == PointFiveUploadFailure("max_upload_retries must be at least 1", retryable=False)