Skip to content

Commit c802a30

Browse files
feat(client): add support for aiohttp
1 parent d17c345 commit c802a30

21 files changed

+192
-20
lines changed

README.md

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -68,6 +68,42 @@ asyncio.run(main())
6868

6969
Functionality between the synchronous and asynchronous clients is otherwise identical.
7070

71+
### With aiohttp
72+
73+
By default, the async client uses `httpx` for HTTP requests. However, for improved concurrency performance you may also use `aiohttp` as the HTTP backend.
74+
75+
You can enable this by installing `aiohttp`:
76+
77+
```sh
78+
# install from PyPI
79+
pip install --pre replicate[aiohttp]
80+
```
81+
82+
Then you can enable it by instantiating the client with `http_client=DefaultAioHttpClient()`:
83+
84+
```python
85+
import os
86+
import asyncio
87+
from replicate import DefaultAioHttpClient
88+
from replicate import AsyncReplicate
89+
90+
91+
async def main() -> None:
92+
async with AsyncReplicate(
93+
bearer_token=os.environ.get(
94+
"REPLICATE_API_TOKEN"
95+
), # This is the default and can be omitted
96+
http_client=DefaultAioHttpClient(),
97+
) as replicate:
98+
prediction = await replicate.predictions.get(
99+
prediction_id="gm3qorzdhgbfurvjtvhg6dckhu",
100+
)
101+
print(prediction.id)
102+
103+
104+
asyncio.run(main())
105+
```
106+
71107
## Using types
72108

73109
Nested request parameters are [TypedDicts](https://docs.python.org/3/library/typing.html#typing.TypedDict). Responses are [Pydantic models](https://docs.pydantic.dev) which also provide helper methods for things like:

pyproject.toml

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,8 @@ classifiers = [
3737
Homepage = "https://github.com/replicate/replicate-python-stainless"
3838
Repository = "https://github.com/replicate/replicate-python-stainless"
3939

40+
[project.optional-dependencies]
41+
aiohttp = ["aiohttp", "httpx_aiohttp>=0.1.6"]
4042

4143
[tool.rye]
4244
managed = true

requirements-dev.lock

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,13 +10,24 @@
1010
# universal: false
1111

1212
-e file:.
13+
aiohappyeyeballs==2.6.1
14+
# via aiohttp
15+
aiohttp==3.12.8
16+
# via httpx-aiohttp
17+
# via replicate
18+
aiosignal==1.3.2
19+
# via aiohttp
1320
annotated-types==0.6.0
1421
# via pydantic
1522
anyio==4.4.0
1623
# via httpx
1724
# via replicate
1825
argcomplete==3.1.2
1926
# via nox
27+
async-timeout==5.0.1
28+
# via aiohttp
29+
attrs==25.3.0
30+
# via aiohttp
2031
certifi==2023.7.22
2132
# via httpcore
2233
# via httpx
@@ -34,23 +45,33 @@ execnet==2.1.1
3445
# via pytest-xdist
3546
filelock==3.12.4
3647
# via virtualenv
48+
frozenlist==1.6.2
49+
# via aiohttp
50+
# via aiosignal
3751
h11==0.14.0
3852
# via httpcore
3953
httpcore==1.0.2
4054
# via httpx
4155
httpx==0.28.1
56+
# via httpx-aiohttp
4257
# via replicate
4358
# via respx
59+
httpx-aiohttp==0.1.6
60+
# via replicate
4461
idna==3.4
4562
# via anyio
4663
# via httpx
64+
# via yarl
4765
importlib-metadata==7.0.0
4866
iniconfig==2.0.0
4967
# via pytest
5068
markdown-it-py==3.0.0
5169
# via rich
5270
mdurl==0.1.2
5371
# via markdown-it-py
72+
multidict==6.4.4
73+
# via aiohttp
74+
# via yarl
5475
mypy==1.14.1
5576
mypy-extensions==1.0.0
5677
# via mypy
@@ -65,6 +86,9 @@ platformdirs==3.11.0
6586
# via virtualenv
6687
pluggy==1.5.0
6788
# via pytest
89+
propcache==0.3.1
90+
# via aiohttp
91+
# via yarl
6892
pydantic==2.10.3
6993
# via replicate
7094
pydantic-core==2.27.1
@@ -97,12 +121,15 @@ tomli==2.0.2
97121
# via pytest
98122
typing-extensions==4.12.2
99123
# via anyio
124+
# via multidict
100125
# via mypy
101126
# via pydantic
102127
# via pydantic-core
103128
# via pyright
104129
# via replicate
105130
virtualenv==20.24.5
106131
# via nox
132+
yarl==1.20.0
133+
# via aiohttp
107134
zipp==3.17.0
108135
# via importlib-metadata

requirements.lock

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,27 +10,51 @@
1010
# universal: false
1111

1212
-e file:.
13+
aiohappyeyeballs==2.6.1
14+
# via aiohttp
15+
aiohttp==3.12.8
16+
# via httpx-aiohttp
17+
# via replicate
18+
aiosignal==1.3.2
19+
# via aiohttp
1320
annotated-types==0.6.0
1421
# via pydantic
1522
anyio==4.4.0
1623
# via httpx
1724
# via replicate
25+
async-timeout==5.0.1
26+
# via aiohttp
27+
attrs==25.3.0
28+
# via aiohttp
1829
certifi==2023.7.22
1930
# via httpcore
2031
# via httpx
2132
distro==1.8.0
2233
# via replicate
2334
exceptiongroup==1.2.2
2435
# via anyio
36+
frozenlist==1.6.2
37+
# via aiohttp
38+
# via aiosignal
2539
h11==0.14.0
2640
# via httpcore
2741
httpcore==1.0.2
2842
# via httpx
2943
httpx==0.28.1
44+
# via httpx-aiohttp
45+
# via replicate
46+
httpx-aiohttp==0.1.6
3047
# via replicate
3148
idna==3.4
3249
# via anyio
3350
# via httpx
51+
# via yarl
52+
multidict==6.4.4
53+
# via aiohttp
54+
# via yarl
55+
propcache==0.3.1
56+
# via aiohttp
57+
# via yarl
3458
pydantic==2.10.3
3559
# via replicate
3660
pydantic-core==2.27.1
@@ -40,6 +64,9 @@ sniffio==1.3.0
4064
# via replicate
4165
typing-extensions==4.12.2
4266
# via anyio
67+
# via multidict
4368
# via pydantic
4469
# via pydantic-core
4570
# via replicate
71+
yarl==1.20.0
72+
# via aiohttp

src/replicate/__init__.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -41,7 +41,7 @@
4141
APIResponseValidationError,
4242
)
4343
from .lib._models import Model as Model, Version as Version, ModelVersionIdentifier as ModelVersionIdentifier
44-
from ._base_client import DefaultHttpxClient, DefaultAsyncHttpxClient
44+
from ._base_client import DefaultHttpxClient, DefaultAioHttpClient, DefaultAsyncHttpxClient
4545
from ._utils._logs import setup_logging as _setup_logging
4646

4747
__all__ = [
@@ -83,6 +83,7 @@
8383
"DEFAULT_CONNECTION_LIMITS",
8484
"DefaultHttpxClient",
8585
"DefaultAsyncHttpxClient",
86+
"DefaultAioHttpClient",
8687
"FileOutput",
8788
"AsyncFileOutput",
8889
"Model",

src/replicate/_base_client.py

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1303,6 +1303,24 @@ def __init__(self, **kwargs: Any) -> None:
13031303
super().__init__(**kwargs)
13041304

13051305

1306+
try:
1307+
import httpx_aiohttp
1308+
except ImportError:
1309+
1310+
class _DefaultAioHttpClient(httpx.AsyncClient):
1311+
def __init__(self, **_kwargs: Any) -> None:
1312+
raise RuntimeError("To use the aiohttp client you must have installed the package with the `aiohttp` extra")
1313+
else:
1314+
1315+
class _DefaultAioHttpClient(httpx_aiohttp.HttpxAiohttpClient): # type: ignore
1316+
def __init__(self, **kwargs: Any) -> None:
1317+
kwargs.setdefault("timeout", DEFAULT_TIMEOUT)
1318+
kwargs.setdefault("limits", DEFAULT_CONNECTION_LIMITS)
1319+
kwargs.setdefault("follow_redirects", True)
1320+
1321+
super().__init__(**kwargs)
1322+
1323+
13061324
if TYPE_CHECKING:
13071325
DefaultAsyncHttpxClient = httpx.AsyncClient
13081326
"""An alias to `httpx.AsyncClient` that provides the same defaults that this SDK
@@ -1311,8 +1329,12 @@ def __init__(self, **kwargs: Any) -> None:
13111329
This is useful because overriding the `http_client` with your own instance of
13121330
`httpx.AsyncClient` will result in httpx's defaults being used, not ours.
13131331
"""
1332+
1333+
DefaultAioHttpClient = httpx.AsyncClient
1334+
"""An alias to `httpx.AsyncClient` that changes the default HTTP transport to `aiohttp`."""
13141335
else:
13151336
DefaultAsyncHttpxClient = _DefaultAsyncHttpxClient
1337+
DefaultAioHttpClient = _DefaultAioHttpClient
13161338

13171339

13181340
class AsyncHttpxClientWrapper(DefaultAsyncHttpxClient):

tests/api_resources/deployments/test_predictions.py

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -108,7 +108,9 @@ def test_path_params_create(self, client: Replicate) -> None:
108108

109109

110110
class TestAsyncPredictions:
111-
parametrize = pytest.mark.parametrize("async_client", [False, True], indirect=True, ids=["loose", "strict"])
111+
parametrize = pytest.mark.parametrize(
112+
"async_client", [False, True, {"http_client": "aiohttp"}], indirect=True, ids=["loose", "strict", "aiohttp"]
113+
)
112114

113115
@pytest.mark.skip()
114116
@parametrize

tests/api_resources/models/test_examples.py

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -69,7 +69,9 @@ def test_path_params_list(self, client: Replicate) -> None:
6969

7070

7171
class TestAsyncExamples:
72-
parametrize = pytest.mark.parametrize("async_client", [False, True], indirect=True, ids=["loose", "strict"])
72+
parametrize = pytest.mark.parametrize(
73+
"async_client", [False, True, {"http_client": "aiohttp"}], indirect=True, ids=["loose", "strict", "aiohttp"]
74+
)
7375

7476
@pytest.mark.skip()
7577
@parametrize

tests/api_resources/models/test_predictions.py

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -108,7 +108,9 @@ def test_path_params_create(self, client: Replicate) -> None:
108108

109109

110110
class TestAsyncPredictions:
111-
parametrize = pytest.mark.parametrize("async_client", [False, True], indirect=True, ids=["loose", "strict"])
111+
parametrize = pytest.mark.parametrize(
112+
"async_client", [False, True, {"http_client": "aiohttp"}], indirect=True, ids=["loose", "strict", "aiohttp"]
113+
)
112114

113115
@pytest.mark.skip()
114116
@parametrize

tests/api_resources/models/test_readme.py

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -70,7 +70,9 @@ def test_path_params_get(self, client: Replicate) -> None:
7070

7171

7272
class TestAsyncReadme:
73-
parametrize = pytest.mark.parametrize("async_client", [False, True], indirect=True, ids=["loose", "strict"])
73+
parametrize = pytest.mark.parametrize(
74+
"async_client", [False, True, {"http_client": "aiohttp"}], indirect=True, ids=["loose", "strict", "aiohttp"]
75+
)
7476

7577
@pytest.mark.skip()
7678
@parametrize

0 commit comments

Comments
 (0)