Update client to use httpx

Signed-off-by: Mattt Zmuda <[email protected]>

Add respx dependency

Signed-off-by: Mattt Zmuda <[email protected]>

Author: mattt

pyproject.toml

@@ -17,6 +17,7 @@ optional-dependencies = { dev = [
     "pytest",
     "pytest-asyncio",
     "pytest-recording",
+    "respx",
     "ruff",
 ] }
 

replicate/client.py

@@ -1,106 +1,67 @@
 import os
 import re
-from json import JSONDecodeError
-from typing import Any, Dict, Iterator, Optional, Union
+from typing import Any, Iterator, Optional, Union
 
-import requests
-from requests.adapters import HTTPAdapter, Retry
-from requests.cookies import RequestsCookieJar
+import httpx
 
-from replicate.__about__ import __version__
-from replicate.deployment import DeploymentCollection
-from replicate.exceptions import ModelError, ReplicateError
-from replicate.model import ModelCollection
-from replicate.prediction import PredictionCollection
-from replicate.training import TrainingCollection
+from .__about__ import __version__
+from .deployment import DeploymentCollection
+from .exceptions import ModelError, ReplicateError
+from .model import ModelCollection
+from .prediction import PredictionCollection
+from .training import TrainingCollection
 
 
 class Client:
-    def __init__(self, api_token: Optional[str] = None) -> None:
+    """A Replicate API client library"""
+
+    def __init__(
+        self,
+        api_token: Optional[str] = None,
+        *,
+        base_url: Optional[str] = None,
+        timeout: Optional[httpx.Timeout] = None,
+        **kwargs,
+    ) -> None:
         super().__init__()
-        # Client is instantiated at import time, so do as little as possible.
-        # This includes resolving environment variables -- they might be set programmatically.
-        self.api_token = api_token
-        self.base_url = os.environ.get(
+
+        api_token = api_token or os.environ.get("REPLICATE_API_TOKEN")
+
+        base_url = base_url or os.environ.get(
             "REPLICATE_API_BASE_URL", "https://api.replicate.com"
         )
-        self.poll_interval = float(os.environ.get("REPLICATE_POLL_INTERVAL", "0.5"))
 
-        # TODO: make thread safe
-        self.read_session = _create_session()
-        read_retries = Retry(
-            total=5,
-            backoff_factor=2,
-            # Only retry 500s on GET so we don't unintionally mutute data
-            allowed_methods=["GET"],
-            # https://support.cloudflare.com/hc/en-us/articles/115003011431-Troubleshooting-Cloudflare-5XX-errors
-            status_forcelist=[
-                429,
-                500,
-                502,
-                503,
-                504,
-                520,
-                521,
-                522,
-                523,
-                524,
-                526,
-                527,
-            ],
+        timeout = timeout or httpx.Timeout(
+            5.0, read=30.0, write=30.0, connect=5.0, pool=10.0
         )
-        self.read_session.mount("http://", HTTPAdapter(max_retries=read_retries))
-        self.read_session.mount("https://", HTTPAdapter(max_retries=read_retries))
-
-        self.write_session = _create_session()
-        write_retries = Retry(
-            total=5,
-            backoff_factor=2,
-            allowed_methods=["POST", "PUT"],
-            # Only retry POST/PUT requests on rate limits, so we don't unintionally mutute data
-            status_forcelist=[429],
-        )
-        self.write_session.mount("http://", HTTPAdapter(max_retries=write_retries))
-        self.write_session.mount("https://", HTTPAdapter(max_retries=write_retries))
-
-    def _request(self, method: str, path: str, **kwargs) -> requests.Response:
-        # from requests.Session
-        if method in ["GET", "OPTIONS"]:
-            kwargs.setdefault("allow_redirects", True)
-        if method in ["HEAD"]:
-            kwargs.setdefault("allow_redirects", False)
-        kwargs.setdefault("headers", {})
-        kwargs["headers"].update(self._headers())
-        session = self.read_session
-        if method in ["POST", "PUT", "DELETE", "PATCH"]:
-            session = self.write_session
-        resp = session.request(method, self.base_url + path, **kwargs)
-        if 400 <= resp.status_code < 600:
-            try:
-                raise ReplicateError(resp.json()["detail"])
-            except (JSONDecodeError, KeyError):
-                pass
-            raise ReplicateError(f"HTTP error: {resp.status_code, resp.reason}")
-        return resp
 
-    def _headers(self) -> Dict[str, str]:
-        return {
-            "Authorization": f"Token {self._api_token()}",
+        self.poll_interval = float(os.environ.get("REPLICATE_POLL_INTERVAL", "0.5"))
+
+        headers = {
+            "Authorization": f"Token {api_token}",
             "User-Agent": f"replicate-python/{__version__}",
         }
 
-    def _api_token(self) -> str:
-        token = self.api_token
-        # Evaluate lazily in case environment variable is set with dotenv, or something
-        if token is None:
-            token = os.environ.get("REPLICATE_API_TOKEN")
-        if not token:
-            raise ReplicateError(
-                """No API token provided. You need to set the REPLICATE_API_TOKEN environment variable or create a client with `replicate.Client(api_token=...)`.
+        transport = kwargs.pop("transport", httpx.HTTPTransport())
 
-You can find your API key on https://replicate.com"""
-            )
-        return token
+        self._client = self._build_client(
+            **kwargs,
+            base_url=base_url,
+            headers=headers,
+            timeout=timeout,
+            transport=transport,
+        )
+
+    def _build_client(self, **kwargs) -> httpx.Client:
+        return httpx.Client(**kwargs)
+
+    def _request(self, method: str, path: str, **kwargs) -> httpx.Response:
+        resp = self._client.request(method, path, **kwargs)
+
+        if 400 <= resp.status_code < 600:
+            raise ReplicateError(resp.json()["detail"])
+
+        return resp
 
     @property
     def models(self) -> ModelCollection:
@@ -150,21 +111,3 @@ def run(self, model_version: str, **kwargs) -> Union[Any, Iterator[Any]]:
         if prediction.status == "failed":
             raise ModelError(prediction.error)
         return prediction.output
-
-
-class _NonpersistentCookieJar(RequestsCookieJar):
-    """
-    A cookie jar that doesn't persist cookies between requests.
-    """
-
-    def set(self, name, value, **kwargs) -> None:
-        return
-
-    def set_cookie(self, cookie, *args, **kwargs) -> None:
-        return
-
-
-def _create_session() -> requests.Session:
-    s = requests.Session()
-    s.cookies = _NonpersistentCookieJar()
-    return s

replicate/files.py

@@ -4,7 +4,7 @@
 import os
 from typing import Optional
 
-import requests
+import httpx
 
 
 def upload_file(fh: io.IOBase, output_file_prefix: Optional[str] = None) -> str:
@@ -24,7 +24,7 @@ def upload_file(fh: io.IOBase, output_file_prefix: Optional[str] = None) -> str:
     if output_file_prefix is not None:
         name = getattr(fh, "name", "output")
         url = output_file_prefix + os.path.basename(name)
-        resp = requests.put(url, files={"file": fh}, timeout=None)
+        resp = httpx.put(url, files={"file": fh}, timeout=None)  # type: ignore
         resp.raise_for_status()
         return url
 

requirements-dev.txt

@@ -21,7 +21,9 @@ h11==0.14.0
 httpcore==0.17.3
     # via httpx
 httpx==0.24.1
-    # via replicate (pyproject.toml)
+    # via
+    #   replicate (pyproject.toml)
+    #   respx
 idna==3.4
     # via
     #   anyio
@@ -63,6 +65,8 @@ pytest-recording==0.13.0
     # via replicate (pyproject.toml)
 pyyaml==6.0.1
     # via vcrpy
+respx==0.20.2
+    # via replicate (pyproject.toml)
 ruff==0.0.278
     # via replicate (pyproject.toml)
 sniffio==1.3.0

tests/conftest.py

@@ -9,7 +9,10 @@ def mock_replicate_api_token(scope="class"):
     if os.environ.get("REPLICATE_API_TOKEN", "") != "":
         yield
     else:
-        with mock.patch.dict(os.environ, {"REPLICATE_API_TOKEN": "test-token"}):
+        with mock.patch.dict(
+            os.environ,
+            {"REPLICATE_API_TOKEN": "test-token", "REPLICATE_POLL_INTERVAL": "0.0"},
+        ):
             yield
 
 

tests/test_deployment.py

@@ -1,26 +1,16 @@
-import responses
-from responses import matchers
+import httpx
+import respx
 
 from replicate.client import Client
 
-
-@responses.activate
-def test_deployment_predictions_create():
-    client = Client(api_token="abc123")
-
-    deployment = client.deployments.get("test/model")
-
-    rsp = responses.post(
-        "https://api.replicate.com/v1/deployments/test/model/predictions",
-        match=[
-            matchers.json_params_matcher(
-                {
-                    "input": {"text": "world"},
-                    "webhook": "https://example.com/webhook",
-                    "webhook_events_filter": ["completed"],
-                }
-            ),
-        ],
+router = respx.Router(base_url="https://api.replicate.com/v1")
+router.route(
+    method="POST",
+    path="/deployments/test/model/predictions",
+    name="deployments.predictions.create",
+).mock(
+    return_value=httpx.Response(
+        201,
         json={
             "id": "p1",
             "version": "v1",
@@ -31,17 +21,28 @@ def test_deployment_predictions_create():
             "created_at": "2022-04-26T20:00:40.658234Z",
             "source": "api",
             "status": "processing",
-            "input": {"text": "hello"},
+            "input": {"text": "world"},
             "output": None,
             "error": None,
             "logs": "",
         },
     )
+)
+router.route(host="api.replicate.com").pass_through()
+
+
+def test_deployment_predictions_create():
+    client = Client(
+        api_token="test-token", transport=httpx.MockTransport(router.handler)
+    )
+    deployment = client.deployments.get("test/model")
 
-    deployment.predictions.create(
+    prediction = deployment.predictions.create(
         input={"text": "world"},
         webhook="https://example.com/webhook",
         webhook_events_filter=["completed"],
     )
 
-    assert rsp.call_count == 1
+    assert router["deployments.predictions.create"].called
+    assert prediction.id == "p1"
+    assert prediction.input == {"text": "world"}