Document and fix behavior of custom client instances (#318)

In #268, we allowed custom headers to be passed to client instances.
However, that didn't apply any custom `User-Agent` header. This is now
fixed.

This PR also updates the README to document how to create a custom
client instance, incorporating this new behavior. This PR is an
alternative to #304.

---------

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

Author: mattt

README.md

@@ -359,6 +359,33 @@ training = replicate.trainings.create(
 )
 ```
 
+## Customize client behavior
+
+The `replicate` package exports a default shared client.
+This client is initialized with an API token
+set by the `REPLICATE_API_TOKEN` environment variable.
+
+You can create your own client instance to
+pass a different API token value,
+add custom headers to requests,
+or control the behavior of the underlying [HTTPX client](https://www.python-httpx.org/api/#client):
+
+```python
+import os
+from replicate.client import Client
+
+replicate = Client(
+  api_token=os.environ["SOME_OTHER_REPLICATE_API_TOKEN"]
+  headers={
+    "User-Agent": "my-app/1.0
+  }
+)
+```
+
+> [!WARNING]
+> Never hardcode authentication credentials like API tokens into your code.
+> Instead, pass them as environment variables when running your program.
+
 ## Development
 
 See [CONTRIBUTING.md](CONTRIBUTING.md)

pyproject.toml

@@ -80,4 +80,5 @@ ignore = [
     "S101",   # Use of assert
     "S106",   # Possible use of hard-coded password function arguments
     "ANN201", # Missing return type annotation for public function
+    "ANN202", # Missing return type annotation for private function
 ]

replicate/client.py

@@ -330,11 +330,11 @@ def _build_httpx_client(
     **kwargs,
 ) -> Union[httpx.Client, httpx.AsyncClient]:
     headers = kwargs.pop("headers", {})
-    headers["User-Agent"] = f"replicate-python/{__version__}"
-
-    if (
+    if "User-Agent" not in headers:
+        headers["User-Agent"] = f"replicate-python/{__version__}"
+    if "Authorization" not in headers and (
         api_token := api_token or os.environ.get("REPLICATE_API_TOKEN")
-    ) and api_token != "":
+    ):
         headers["Authorization"] = f"Bearer {api_token}"
 
     base_url = (

tests/test_client.py

@@ -91,21 +91,26 @@ def test_custom_headers_are_applied():
     import replicate
     from replicate.exceptions import ReplicateError
 
-    custom_headers = {"Custom-Header": "CustomValue"}
-
-    def mock_send(request: httpx.Request, **kwargs) -> httpx.Response:
-        assert "Custom-Header" in request.headers
-        assert request.headers["Custom-Header"] == "CustomValue"
+    custom_headers = {"User-Agent": "my-custom-user-agent/1.0"}
 
+    def mock_send(request):
+        assert "User-Agent" in request.headers, "Custom header not found in request"
+        assert (
+            request.headers["User-Agent"] == "my-custom-user-agent/1.0"
+        ), "Custom header value is incorrect"
         return httpx.Response(401, json={})
 
+    mock_send_wrapper = mock.Mock(side_effect=mock_send)
+
     client = replicate.Client(
         api_token="dummy_token",
         headers=custom_headers,
-        transport=httpx.MockTransport(mock_send),
+        transport=httpx.MockTransport(mock_send_wrapper),
     )
 
     try:
         client.accounts.current()
     except ReplicateError:
         pass
+
+    mock_send_wrapper.assert_called_once()