feat(pointfive): add presigned upload client

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

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

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="<html>gateway</html>"),))
+
+    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)