feat(pointfive): add presigned upload client

litellm/integrations/pointfive/upload_client.py

@@ -0,0 +1,194 @@
+"""
+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
+from types import MappingProxyType
+from typing import Final
+
+import httpx
+from pydantic import BaseModel, Field, ValidationError
+
+import litellm
+from litellm._logging import verbose_logger
+from litellm.litellm_core_utils.url_utils import SSRFError, validate_url
+from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
+from litellm.types.integrations.pointfive import (
+    RETRYABLE_UPLOAD_STATUS_CODES,
+    PointFiveUploadFailure,
+    PointFiveUploadTarget,
+)
+
+UPLOAD_KIND: Final = "LITELLM"
+UPLOAD_URL_PATH: Final = "/upload-url"
+PING_PATH: Final = "/ping"
+PUT_HEADERS: Final = MappingProxyType({"Content-Type": "application/x-ndjson", "Content-Encoding": "gzip"})
+
+
+class _PresignRequest(BaseModel):
+    kind: str = UPLOAD_KIND
+    byte_count: int = Field(serialization_alias="byteCount")
+
+
+class _PingRequest(BaseModel):
+    kind: str = UPLOAD_KIND
+
+
+class _TargetPayload(BaseModel):
+    upload_url: str = Field(alias="uploadUrl")
+    object_key: str = Field(alias="objectKey")
+
+
+class _ErrorPayload(BaseModel):
+    error: str = ""
+
+
+class PointFiveUploadError(Exception):
+    """A batch could not be uploaded and the failure is worth retrying."""
+
+
+def _failure_for(response: httpx.Response, what: str) -> PointFiveUploadFailure:
+    detail: Final = f"{what} returned {response.status_code}"
+    reason: Final = _refusal_reason(response.text)
+    return PointFiveUploadFailure(
+        f"{detail}, {reason}" if reason else detail,
+        retryable=response.status_code in RETRYABLE_UPLOAD_STATUS_CODES,
+    )
+
+
+def _refusal_reason(body: str) -> str:
+    try:
+        return _ErrorPayload.model_validate_json(body).error
+    except ValidationError:
+        return ""
+
+
+def _parse_target(body: str) -> PointFiveUploadTarget | PointFiveUploadFailure:
+    try:
+        target: Final = _TargetPayload.model_validate_json(body)
+    except ValidationError:
+        return PointFiveUploadFailure("pointfive api returned an unreadable body", retryable=False)
+    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,
+        validate_upload_url: Callable[[str], tuple[str, str]] = validate_url,
+    ) -> None:
+        self.api_key: Final = api_key
+        self.api_url: Final = api_url.rstrip("/")
+        self.http_client: Final = http_client
+        self.max_retries: Final = max_retries
+        self.sleep: Final = sleep
+        self.validate_upload_url: Final = validate_upload_url
+
+    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._post(PING_PATH, _PingRequest())
+        if isinstance(body, PointFiveUploadFailure):
+            return body
+        return None
+
+    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._post(UPLOAD_URL_PATH, _PresignRequest(byte_count=byte_count))
+        if isinstance(body, PointFiveUploadFailure):
+            return body
+        return _parse_target(body)
+
+    async def _post(self, path: str, request: BaseModel) -> str | PointFiveUploadFailure:
+        """POST one JSON request to the PointFive ingestion API and return its raw body."""
+        try:
+            response: Final = await self.http_client.post(
+                self.api_url + path,
+                json=request.model_dump(by_alias=True),
+                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, "pointfive api")
+        except Exception as e:  # noqa: BLE001  # a transport fault is worth another attempt
+            return PointFiveUploadFailure(f"pointfive api unreachable: {type(e).__name__}", 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.
+
+        The server chose that URL, so it is treated like any other externally supplied
+        destination: the host is checked against blocked networks before connecting, and
+        a redirect is refused rather than followed. A presigned URL never legitimately
+        redirects, and following one would let a compromised endpoint point the proxy at
+        an internal service.
+        """
+        destination: Final = self._destination(target.upload_url)
+        if isinstance(destination, PointFiveUploadFailure):
+            return destination
+        url, host = destination
+        headers: Final = dict(PUT_HEADERS, Host=host) if host else dict(PUT_HEADERS)  # mutable-ok: put wants dict
+        try:
+            await self.http_client.put(url, data=body, headers=headers, follow_redirects=False)
+        except httpx.HTTPStatusError as e:
+            if e.response.is_redirect:
+                return PointFiveUploadFailure(
+                    f"presigned upload redirected with {e.response.status_code}, refusing to follow", retryable=False
+                )
+            return _failure_for(e.response, "presigned upload")
+        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
+
+    def _destination(self, upload_url: str) -> tuple[str, str | None] | PointFiveUploadFailure:
+        if not getattr(litellm, "user_url_validation", True):
+            return upload_url, None
+        try:
+            return self.validate_upload_url(upload_url)
+        except SSRFError as e:
+            return PointFiveUploadFailure(f"presigned upload url refused: {e}", retryable=False)

litellm/types/integrations/pointfive.py

@@ -0,0 +1,45 @@
+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/api/v1/ingestion"
+
+
+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. ``batch_size`` also bounds how much a busy
+    proxy holds in memory between flushes, so it stays modest. ``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=1_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,378 @@
+import json
+from collections.abc import Sequence
+
+import httpx
+import pytest
+
+import litellm
+from litellm.integrations.pointfive.upload_client import PointFiveUploadClient
+from litellm.litellm_core_utils.url_utils import validate_url
+from litellm.types.integrations.pointfive import PointFiveUploadFailure
+
+API_URL = "https://api.pointfive.co/api/v1/ingestion"
+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, {"uploadUrl": UPLOAD_URL, "objectKey": OBJECT_KEY, "expiresAt": "2026-08-25T14:35:00Z"}
+    )
+
+
+def _response(status_code: int, payload: object) -> httpx.Response:
+    return httpx.Response(status_code, text=json.dumps(payload))
+
+
+def _refused(status_code: int, error: str) -> httpx.Response:
+    """The body PointFive sends with every refusal."""
+    return _response(status_code, {"success": False, "error": error})
+
+
+def _accepted() -> httpx.Response:
+    return httpx.Response(200, text="")
+
+
+def _no_content() -> httpx.Response:
+    return httpx.Response(204, 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, follow_redirects=None, **_):
+        self.put_calls.append(
+            {"url": url, "data": data, "headers": headers or {}, "follow_redirects": follow_redirects}
+        )
+        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 >= 300:
+        request = httpx.Request("POST", url)
+        raise httpx.HTTPStatusError(
+            "boom",
+            request=request,
+            response=httpx.Response(result.status_code, text=result.text, headers=result.headers),
+        )
+    return result
+
+
+async def _no_backoff(_seconds: float) -> None:
+    return None
+
+
+def _trusting_validator(url: str) -> tuple[str, str]:
+    """Stands in for validate_url so the fixture hosts need no DNS; the SSRF tests use the real one."""
+    return url, httpx.URL(url).host
+
+
+def _client(
+    http_client: FakeHTTPClient,
+    max_retries: int = 3,
+    api_url: str = API_URL,
+    validate_upload_url=_trusting_validator,
+) -> PointFiveUploadClient:
+    return PointFiveUploadClient(
+        api_key="p5tu_testkey",
+        api_url=api_url,
+        http_client=http_client,
+        max_retries=max_retries,
+        sleep=_no_backoff,
+        validate_upload_url=validate_upload_url,
+    )
+
+
+def _presigned_for(upload_url: str) -> httpx.Response:
+    return _response(200, {"uploadUrl": upload_url, "objectKey": OBJECT_KEY, "expiresAt": "2026-08-25T14:35:00Z"})
+
+
+@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["url"] == "https://api.pointfive.co/api/v1/ingestion/upload-url"
+    assert call["headers"]["Authorization"] == "Bearer p5tu_testkey"
+    assert call["json"] == {"kind": "LITELLM", "byteCount": len(BODY)}
+
+
+@pytest.mark.asyncio
+async def test_a_trailing_slash_on_the_api_url_is_tolerated():
+    """A pasted URL often ends in a slash; it must not produce a double slash in the path."""
+    http_client = FakeHTTPClient()
+
+    await _client(http_client, api_url=API_URL + "/").upload(BODY)
+
+    assert http_client.presign_calls[0]["url"] == "https://api.pointfive.co/api/v1/ingestion/upload-url"
+
+
+@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_upload_pins_the_host_and_never_follows_a_redirect():
+    http_client = FakeHTTPClient()
+
+    await _client(http_client).upload(BODY)
+
+    call = http_client.put_calls[0]
+    assert call["headers"]["Host"] == "uploads.example.invalid"
+    assert call["follow_redirects"] is False
+
+
+@pytest.mark.asyncio
+async def test_a_redirected_upload_is_refused_rather_than_followed():
+    """A presigned URL never redirects legitimately; following one is how a bad endpoint reaches inside."""
+    http_client = FakeHTTPClient(put=[httpx.Response(301, headers={"location": "http://169.254.169.254/"})])
+
+    outcome = await _client(http_client).upload(BODY)
+
+    assert outcome == PointFiveUploadFailure(
+        "presigned upload redirected with 301, refusing to follow", retryable=False
+    )
+    assert len(http_client.put_calls) == 1
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "upload_url",
+    [
+        "http://169.254.169.254/latest/meta-data",
+        "https://10.0.0.7/internal/bucket/object",
+        "http://127.0.0.1:9000/bucket/object",
+    ],
+)
+async def test_an_upload_url_inside_the_network_is_refused_before_any_bytes_leave(upload_url):
+    http_client = FakeHTTPClient(presign=[_presigned_for(upload_url)])
+
+    outcome = await _client(http_client, validate_upload_url=validate_url).upload(BODY)
+
+    assert isinstance(outcome, PointFiveUploadFailure)
+    assert not outcome.retryable
+    assert outcome.detail.startswith("presigned upload url refused: ")
+    assert http_client.put_calls == []
+
+
+@pytest.mark.asyncio
+async def test_an_operator_can_switch_destination_validation_off(monkeypatch):
+    """litellm.user_url_validation is the proxy-wide switch every SSRF guard honours."""
+    monkeypatch.setattr(litellm, "user_url_validation", False)
+    http_client = FakeHTTPClient(presign=[_presigned_for("http://10.0.0.7/bucket/object")])
+
+    outcome = await _client(http_client, validate_upload_url=validate_url).upload(BODY)
+
+    assert outcome == OBJECT_KEY
+    assert http_client.put_calls[0]["url"] == "http://10.0.0.7/bucket/object"
+    assert "Host" 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_the_reason_for_a_refusal_is_surfaced():
+    """A 403 means the key no longer maps to an integration; the operator needs to read why."""
+    http_client = FakeHTTPClient(presign=[_refused(403, "no integration accepts uploads from this api key")])
+
+    outcome = await _client(http_client).upload(BODY)
+
+    assert outcome == PointFiveUploadFailure(
+        "pointfive api returned 403, no integration accepts uploads from this api key", 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_too_many_requests_is_retried():
+    http_client = FakeHTTPClient(presign=[httpx.Response(429), _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_malformed_api_body_is_not_retried():
+    http_client = FakeHTTPClient(presign=[_response(200, {"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_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_ping_reports_a_live_shipper():
+    http_client = FakeHTTPClient(presign=(_no_content(),))
+
+    failure = await _client(http_client).ping()
+
+    assert failure is None
+    assert http_client.presign_calls[0]["url"] == "https://api.pointfive.co/api/v1/ingestion/ping"
+    assert http_client.presign_calls[0]["json"] == {"kind": "LITELLM"}
+
+
+@pytest.mark.asyncio
+async def test_ping_surfaces_a_revoked_key():
+    http_client = FakeHTTPClient(presign=(_refused(403, "no integration accepts uploads from this api key"),))
+
+    failure = await _client(http_client).ping()
+
+    assert failure is not None
+    assert not failure.retryable
+    assert "no integration accepts uploads from this api key" 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_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)