feat(client): add support for aiohttp

stainless-app[bot] · stainless-app[bot] · commit 3aa7ee28445e · 2025-06-21T04:33:33.000Z
diff --git a/README.md b/README.md
@@ -79,6 +79,45 @@ asyncio.run(main())
 
 Functionality between the synchronous and asynchronous clients is otherwise identical.
 
+### With aiohttp
+
+By default, the async client uses `httpx` for HTTP requests. However, for improved concurrency performance you may also use `aiohttp` as the HTTP backend.
+
+You can enable this by installing `aiohttp`:
+
+```sh
+# install from PyPI
+pip install --pre zeroentropy[aiohttp]
+```
+
+Then you can enable it by instantiating the client with `http_client=DefaultAioHttpClient()`:
+
+```python
+import os
+import asyncio
+from zeroentropy import DefaultAioHttpClient
+from zeroentropy import AsyncZeroEntropy
+
+
+async def main() -> None:
+    async with AsyncZeroEntropy(
+        api_key=os.environ.get("ZEROENTROPY_API_KEY"),  # This is the default and can be omitted
+        http_client=DefaultAioHttpClient(),
+    ) as client:
+        response = await client.documents.add(
+            collection_name="example_collection",
+            content={
+                "type": "text",
+                "text": "Example Content",
+            },
+            path="my_document.txt",
+        )
+        print(response.message)
+
+
+asyncio.run(main())
+```
+
 ## Using types
 
 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:
diff --git a/pyproject.toml b/pyproject.toml
@@ -37,6 +37,8 @@ classifiers = [
 Homepage = "https://github.com/zeroentropy-ai/zeroentropy-python"
 Repository = "https://github.com/zeroentropy-ai/zeroentropy-python"
 
+[project.optional-dependencies]
+aiohttp = ["aiohttp", "httpx_aiohttp>=0.1.6"]
 
 [tool.rye]
 managed = true
diff --git a/requirements-dev.lock b/requirements-dev.lock
@@ -10,13 +10,24 @@
 #   universal: false
 
 -e file:.
+aiohappyeyeballs==2.6.1
+    # via aiohttp
+aiohttp==3.12.8
+    # via httpx-aiohttp
+    # via zeroentropy
+aiosignal==1.3.2
+    # via aiohttp
 annotated-types==0.6.0
     # via pydantic
 anyio==4.4.0
     # via httpx
     # via zeroentropy
 argcomplete==3.1.2
     # via nox
+async-timeout==5.0.1
+    # via aiohttp
+attrs==25.3.0
+    # via aiohttp
 certifi==2023.7.22
     # via httpcore
     # via httpx
@@ -34,23 +45,33 @@ execnet==2.1.1
     # via pytest-xdist
 filelock==3.12.4
     # via virtualenv
+frozenlist==1.6.2
+    # via aiohttp
+    # via aiosignal
 h11==0.14.0
     # via httpcore
 httpcore==1.0.2
     # via httpx
 httpx==0.28.1
+    # via httpx-aiohttp
     # via respx
     # via zeroentropy
+httpx-aiohttp==0.1.6
+    # via zeroentropy
 idna==3.4
     # via anyio
     # via httpx
+    # via yarl
 importlib-metadata==7.0.0
 iniconfig==2.0.0
     # via pytest
 markdown-it-py==3.0.0
     # via rich
 mdurl==0.1.2
     # via markdown-it-py
+multidict==6.4.4
+    # via aiohttp
+    # via yarl
 mypy==1.14.1
 mypy-extensions==1.0.0
     # via mypy
@@ -65,6 +86,9 @@ platformdirs==3.11.0
     # via virtualenv
 pluggy==1.5.0
     # via pytest
+propcache==0.3.1
+    # via aiohttp
+    # via yarl
 pydantic==2.10.3
     # via zeroentropy
 pydantic-core==2.27.1
@@ -97,12 +121,15 @@ tomli==2.0.2
     # via pytest
 typing-extensions==4.12.2
     # via anyio
+    # via multidict
     # via mypy
     # via pydantic
     # via pydantic-core
     # via pyright
     # via zeroentropy
 virtualenv==20.24.5
     # via nox
+yarl==1.20.0
+    # via aiohttp
 zipp==3.17.0
     # via importlib-metadata
diff --git a/requirements.lock b/requirements.lock
@@ -10,27 +10,51 @@
 #   universal: false
 
 -e file:.
+aiohappyeyeballs==2.6.1
+    # via aiohttp
+aiohttp==3.12.8
+    # via httpx-aiohttp
+    # via zeroentropy
+aiosignal==1.3.2
+    # via aiohttp
 annotated-types==0.6.0
     # via pydantic
 anyio==4.4.0
     # via httpx
     # via zeroentropy
+async-timeout==5.0.1
+    # via aiohttp
+attrs==25.3.0
+    # via aiohttp
 certifi==2023.7.22
     # via httpcore
     # via httpx
 distro==1.8.0
     # via zeroentropy
 exceptiongroup==1.2.2
     # via anyio
+frozenlist==1.6.2
+    # via aiohttp
+    # via aiosignal
 h11==0.14.0
     # via httpcore
 httpcore==1.0.2
     # via httpx
 httpx==0.28.1
+    # via httpx-aiohttp
+    # via zeroentropy
+httpx-aiohttp==0.1.6
     # via zeroentropy
 idna==3.4
     # via anyio
     # via httpx
+    # via yarl
+multidict==6.4.4
+    # via aiohttp
+    # via yarl
+propcache==0.3.1
+    # via aiohttp
+    # via yarl
 pydantic==2.10.3
     # via zeroentropy
 pydantic-core==2.27.1
@@ -40,6 +64,9 @@ sniffio==1.3.0
     # via zeroentropy
 typing-extensions==4.12.2
     # via anyio
+    # via multidict
     # via pydantic
     # via pydantic-core
     # via zeroentropy
+yarl==1.20.0
+    # via aiohttp
diff --git a/src/zeroentropy/__init__.py b/src/zeroentropy/__init__.py
@@ -36,7 +36,7 @@
     UnprocessableEntityError,
     APIResponseValidationError,
 )
-from ._base_client import DefaultHttpxClient, DefaultAsyncHttpxClient
+from ._base_client import DefaultHttpxClient, DefaultAioHttpClient, DefaultAsyncHttpxClient
 from ._utils._logs import setup_logging as _setup_logging
 
 __all__ = [
@@ -78,6 +78,7 @@
     "DEFAULT_CONNECTION_LIMITS",
     "DefaultHttpxClient",
     "DefaultAsyncHttpxClient",
+    "DefaultAioHttpClient",
 ]
 
 if not _t.TYPE_CHECKING:
diff --git a/src/zeroentropy/_base_client.py b/src/zeroentropy/_base_client.py
@@ -1289,6 +1289,24 @@ def __init__(self, **kwargs: Any) -> None:
         super().__init__(**kwargs)
 
 
+try:
+    import httpx_aiohttp
+except ImportError:
+
+    class _DefaultAioHttpClient(httpx.AsyncClient):
+        def __init__(self, **_kwargs: Any) -> None:
+            raise RuntimeError("To use the aiohttp client you must have installed the package with the `aiohttp` extra")
+else:
+
+    class _DefaultAioHttpClient(httpx_aiohttp.HttpxAiohttpClient):  # type: ignore
+        def __init__(self, **kwargs: Any) -> None:
+            kwargs.setdefault("timeout", DEFAULT_TIMEOUT)
+            kwargs.setdefault("limits", DEFAULT_CONNECTION_LIMITS)
+            kwargs.setdefault("follow_redirects", True)
+
+            super().__init__(**kwargs)
+
+
 if TYPE_CHECKING:
     DefaultAsyncHttpxClient = httpx.AsyncClient
     """An alias to `httpx.AsyncClient` that provides the same defaults that this SDK
@@ -1297,8 +1315,12 @@ def __init__(self, **kwargs: Any) -> None:
     This is useful because overriding the `http_client` with your own instance of
     `httpx.AsyncClient` will result in httpx's defaults being used, not ours.
     """
+
+    DefaultAioHttpClient = httpx.AsyncClient
+    """An alias to `httpx.AsyncClient` that changes the default HTTP transport to `aiohttp`."""
 else:
     DefaultAsyncHttpxClient = _DefaultAsyncHttpxClient
+    DefaultAioHttpClient = _DefaultAioHttpClient
 
 
 class AsyncHttpxClientWrapper(DefaultAsyncHttpxClient):
diff --git a/tests/api_resources/test_admin.py b/tests/api_resources/test_admin.py
@@ -50,7 +50,9 @@ def test_streaming_response_create_organization(self, client: ZeroEntropy) -> No
 
 
 class TestAsyncAdmin:
-    parametrize = pytest.mark.parametrize("async_client", [False, True], indirect=True, ids=["loose", "strict"])
+    parametrize = pytest.mark.parametrize(
+        "async_client", [False, True, {"http_client": "aiohttp"}], indirect=True, ids=["loose", "strict", "aiohttp"]
+    )
 
     @parametrize
     async def test_method_create_organization(self, async_client: AsyncZeroEntropy) -> None:
diff --git a/tests/api_resources/test_collections.py b/tests/api_resources/test_collections.py
@@ -110,7 +110,9 @@ def test_streaming_response_get_list(self, client: ZeroEntropy) -> None:
 
 
 class TestAsyncCollections:
-    parametrize = pytest.mark.parametrize("async_client", [False, True], indirect=True, ids=["loose", "strict"])
+    parametrize = pytest.mark.parametrize(
+        "async_client", [False, True, {"http_client": "aiohttp"}], indirect=True, ids=["loose", "strict", "aiohttp"]
+    )
 
     @parametrize
     async def test_method_delete(self, async_client: AsyncZeroEntropy) -> None:
diff --git a/tests/api_resources/test_documents.py b/tests/api_resources/test_documents.py
@@ -295,7 +295,9 @@ def test_streaming_response_get_page_info(self, client: ZeroEntropy) -> None:
 
 
 class TestAsyncDocuments:
-    parametrize = pytest.mark.parametrize("async_client", [False, True], indirect=True, ids=["loose", "strict"])
+    parametrize = pytest.mark.parametrize(
+        "async_client", [False, True, {"http_client": "aiohttp"}], indirect=True, ids=["loose", "strict", "aiohttp"]
+    )
 
     @parametrize
     async def test_method_update(self, async_client: AsyncZeroEntropy) -> None:
diff --git a/tests/api_resources/test_parsers.py b/tests/api_resources/test_parsers.py
@@ -50,7 +50,9 @@ def test_streaming_response_parse_document(self, client: ZeroEntropy) -> None:
 
 
 class TestAsyncParsers:
-    parametrize = pytest.mark.parametrize("async_client", [False, True], indirect=True, ids=["loose", "strict"])
+    parametrize = pytest.mark.parametrize(
+        "async_client", [False, True, {"http_client": "aiohttp"}], indirect=True, ids=["loose", "strict", "aiohttp"]
+    )
 
     @parametrize
     async def test_method_parse_document(self, async_client: AsyncZeroEntropy) -> None:
diff --git a/tests/api_resources/test_queries.py b/tests/api_resources/test_queries.py
@@ -171,7 +171,9 @@ def test_streaming_response_top_snippets(self, client: ZeroEntropy) -> None:
 
 
 class TestAsyncQueries:
-    parametrize = pytest.mark.parametrize("async_client", [False, True], indirect=True, ids=["loose", "strict"])
+    parametrize = pytest.mark.parametrize(
+        "async_client", [False, True, {"http_client": "aiohttp"}], indirect=True, ids=["loose", "strict", "aiohttp"]
+    )
 
     @parametrize
     async def test_method_top_documents(self, async_client: AsyncZeroEntropy) -> None:
diff --git a/tests/api_resources/test_status.py b/tests/api_resources/test_status.py
@@ -51,7 +51,9 @@ def test_streaming_response_get_status(self, client: ZeroEntropy) -> None:
 
 
 class TestAsyncStatus:
-    parametrize = pytest.mark.parametrize("async_client", [False, True], indirect=True, ids=["loose", "strict"])
+    parametrize = pytest.mark.parametrize(
+        "async_client", [False, True, {"http_client": "aiohttp"}], indirect=True, ids=["loose", "strict", "aiohttp"]
+    )
 
     @parametrize
     async def test_method_get_status(self, async_client: AsyncZeroEntropy) -> None:
diff --git a/tests/conftest.py b/tests/conftest.py
@@ -6,10 +6,12 @@
 import logging
 from typing import TYPE_CHECKING, Iterator, AsyncIterator
 
+import httpx
 import pytest
 from pytest_asyncio import is_async_test
 
-from zeroentropy import ZeroEntropy, AsyncZeroEntropy
+from zeroentropy import ZeroEntropy, AsyncZeroEntropy, DefaultAioHttpClient
+from zeroentropy._utils import is_dict
 
 if TYPE_CHECKING:
     from _pytest.fixtures import FixtureRequest  # pyright: ignore[reportPrivateImportUsage]
@@ -27,6 +29,19 @@ def pytest_collection_modifyitems(items: list[pytest.Function]) -> None:
     for async_test in pytest_asyncio_tests:
         async_test.add_marker(session_scope_marker, append=False)
 
+    # We skip tests that use both the aiohttp client and respx_mock as respx_mock
+    # doesn't support custom transports.
+    for item in items:
+        if "async_client" not in item.fixturenames or "respx_mock" not in item.fixturenames:
+            continue
+
+        if not hasattr(item, "callspec"):
+            continue
+
+        async_client_param = item.callspec.params.get("async_client")
+        if is_dict(async_client_param) and async_client_param.get("http_client") == "aiohttp":
+            item.add_marker(pytest.mark.skip(reason="aiohttp client is not compatible with respx_mock"))
+
 
 base_url = os.environ.get("TEST_API_BASE_URL", "http://127.0.0.1:4010")
 
@@ -45,9 +60,25 @@ def client(request: FixtureRequest) -> Iterator[ZeroEntropy]:
 
 @pytest.fixture(scope="session")
 async def async_client(request: FixtureRequest) -> AsyncIterator[AsyncZeroEntropy]:
-    strict = getattr(request, "param", True)
-    if not isinstance(strict, bool):
-        raise TypeError(f"Unexpected fixture parameter type {type(strict)}, expected {bool}")
-
-    async with AsyncZeroEntropy(base_url=base_url, api_key=api_key, _strict_response_validation=strict) as client:
+    param = getattr(request, "param", True)
+
+    # defaults
+    strict = True
+    http_client: None | httpx.AsyncClient = None
+
+    if isinstance(param, bool):
+        strict = param
+    elif is_dict(param):
+        strict = param.get("strict", True)
+        assert isinstance(strict, bool)
+
+        http_client_type = param.get("http_client", "httpx")
+        if http_client_type == "aiohttp":
+            http_client = DefaultAioHttpClient()
+    else:
+        raise TypeError(f"Unexpected fixture parameter type {type(param)}, expected bool or dict")
+
+    async with AsyncZeroEntropy(
+        base_url=base_url, api_key=api_key, _strict_response_validation=strict, http_client=http_client
+    ) as client:
         yield client

Original file line number	Diff line number	Diff line change
`@@ -36,7 +36,7 @@`
`36`	`36`	`UnprocessableEntityError,`
`37`	`37`	`APIResponseValidationError,`
`38`	`38`	`)`
`39`		`-from ._base_client import DefaultHttpxClient, DefaultAsyncHttpxClient`
	`39`	`+from ._base_client import DefaultHttpxClient, DefaultAioHttpClient, DefaultAsyncHttpxClient`
`40`	`40`	`from ._utils._logs import setup_logging as _setup_logging`
`41`	`41`
`42`	`42`	`__all__ = [`
`@@ -78,6 +78,7 @@`
`78`	`78`	`"DEFAULT_CONNECTION_LIMITS",`
`79`	`79`	`"DefaultHttpxClient",`
`80`	`80`	`"DefaultAsyncHttpxClient",`
	`81`	`+ "DefaultAioHttpClient",`
`81`	`82`	`]`
`82`	`83`
`83`	`84`	`if not _t.TYPE_CHECKING:`