feat: auto-detect aiohttp transport for async client

pyproject.toml

@@ -47,6 +47,11 @@ requests = "^2.0.0"
 tokenizers = ">=0.15,<1"
 types-requests = "^2.0.0"
 typing_extensions = ">= 4.0.0"
+aiohttp = {version = ">=3.0", optional = true}
+httpx_aiohttp = {version = ">=0.1.8", optional = true}
+
+[tool.poetry.extras]
+aiohttp = ["aiohttp", "httpx_aiohttp"]
 
 [tool.poetry.group.dev.dependencies]
 mypy = "==1.13.0"

src/cohere/__init__.py

@@ -280,6 +280,7 @@
     )
     from .bedrock_client import BedrockClient, BedrockClientV2
     from .client import AsyncClient, Client
+    from ._default_clients import DefaultAioHttpClient, DefaultAsyncHttpxClient
     from .client_v2 import AsyncClientV2, ClientV2
     from .datasets import DatasetsCreateResponse, DatasetsGetResponse, DatasetsGetUsageResponse, DatasetsListResponse
     from .embed_jobs import CreateEmbedJobRequestTruncate
@@ -440,6 +441,8 @@
     "CreateEmbedJobResponse": ".types",
     "Dataset": ".types",
     "DatasetPart": ".types",
+    "DefaultAioHttpClient": "._default_clients",
+    "DefaultAsyncHttpxClient": "._default_clients",
     "DatasetType": ".types",
     "DatasetValidationStatus": ".types",
     "DatasetsCreateResponse": ".datasets",
@@ -779,6 +782,8 @@ def __dir__():
     "DatasetsGetResponse",
     "DatasetsGetUsageResponse",
     "DatasetsListResponse",
+    "DefaultAioHttpClient",
+    "DefaultAsyncHttpxClient",
     "DebugStreamedChatResponse",
     "DebugV2ChatStreamResponse",
     "DeleteConnectorResponse",

src/cohere/_default_clients.py

@@ -0,0 +1,31 @@
+import typing
+
+import httpx
+
+COHERE_DEFAULT_TIMEOUT = 300
+
+try:
+    import httpx_aiohttp
+except ImportError:
+
+    class DefaultAioHttpClient(httpx.AsyncClient):  # type: ignore
+        def __init__(self, **kwargs: typing.Any) -> None:
+            raise RuntimeError(
+                "To use the aiohttp client, install the aiohttp extra: "
+                "pip install cohere[aiohttp]"
+            )
+
+else:
+
+    class DefaultAioHttpClient(httpx_aiohttp.HttpxAiohttpClient):  # type: ignore
+        def __init__(self, **kwargs: typing.Any) -> None:
+            kwargs.setdefault("timeout", COHERE_DEFAULT_TIMEOUT)
+            kwargs.setdefault("follow_redirects", True)
+            super().__init__(**kwargs)
+
+
+class DefaultAsyncHttpxClient(httpx.AsyncClient):
+    def __init__(self, **kwargs: typing.Any) -> None:
+        kwargs.setdefault("timeout", COHERE_DEFAULT_TIMEOUT)
+        kwargs.setdefault("follow_redirects", True)
+        super().__init__(**kwargs)

src/cohere/base_client.py

@@ -1565,6 +1565,24 @@ def finetuning(self):
         return self._finetuning
 
 
+def _make_default_async_client(
+    timeout: float,
+    follow_redirects: typing.Optional[bool],
+) -> httpx.AsyncClient:
+    try:
+        import httpx_aiohttp
+    except ImportError:
+        pass
+    else:
+        if follow_redirects is not None:
+            return httpx_aiohttp.HttpxAiohttpClient(timeout=timeout, follow_redirects=follow_redirects)
+        return httpx_aiohttp.HttpxAiohttpClient(timeout=timeout)
+
+    if follow_redirects is not None:
+        return httpx.AsyncClient(timeout=timeout, follow_redirects=follow_redirects)
+    return httpx.AsyncClient(timeout=timeout)
+
+
 class AsyncBaseCohere:
     """
     Use this class to access the different functions within the SDK. You can instantiate any number of clients with different configuration that will propagate to these functions.
@@ -1631,9 +1649,7 @@ def __init__(
             headers=headers,
             httpx_client=httpx_client
             if httpx_client is not None
-            else httpx.AsyncClient(timeout=_defaulted_timeout, follow_redirects=follow_redirects)
-            if follow_redirects is not None
-            else httpx.AsyncClient(timeout=_defaulted_timeout),
+            else _make_default_async_client(timeout=_defaulted_timeout, follow_redirects=follow_redirects),
             timeout=_defaulted_timeout,
         )
         self._raw_client = AsyncRawBaseCohere(client_wrapper=self._client_wrapper)

tests/test_aiohttp_autodetect.py

@@ -0,0 +1,119 @@
+import sys
+import typing
+import unittest
+from unittest import mock
+
+import httpx
+
+
+class TestMakeDefaultAsyncClient(unittest.TestCase):
+    """Tests for _make_default_async_client in base_client.py."""
+
+    def test_without_httpx_aiohttp_returns_httpx_async_client(self) -> None:
+        """When httpx_aiohttp is not installed, returns plain httpx.AsyncClient."""
+        with mock.patch.dict(sys.modules, {"httpx_aiohttp": None}):
+            # Re-import to pick up the mocked module state
+            from cohere.base_client import _make_default_async_client
+
+            client = _make_default_async_client(timeout=300, follow_redirects=True)
+            self.assertIsInstance(client, httpx.AsyncClient)
+            self.assertEqual(client.timeout.read, 300)
+            self.assertTrue(client.follow_redirects)
+
+    def test_without_httpx_aiohttp_follow_redirects_none(self) -> None:
+        """When follow_redirects is None, omits it from httpx.AsyncClient."""
+        with mock.patch.dict(sys.modules, {"httpx_aiohttp": None}):
+            from cohere.base_client import _make_default_async_client
+
+            client = _make_default_async_client(timeout=300, follow_redirects=None)
+            self.assertIsInstance(client, httpx.AsyncClient)
+            # httpx default is False when not specified
+            self.assertFalse(client.follow_redirects)
+
+    def test_with_httpx_aiohttp_returns_aiohttp_client(self) -> None:
+        """When httpx_aiohttp is installed, returns HttpxAiohttpClient."""
+        try:
+            import httpx_aiohttp
+        except ImportError:
+            self.skipTest("httpx_aiohttp not installed")
+
+        from cohere.base_client import _make_default_async_client
+
+        client = _make_default_async_client(timeout=300, follow_redirects=True)
+        self.assertIsInstance(client, httpx_aiohttp.HttpxAiohttpClient)
+        self.assertEqual(client.timeout.read, 300)
+        self.assertTrue(client.follow_redirects)
+
+    def test_with_httpx_aiohttp_follow_redirects_none(self) -> None:
+        """When httpx_aiohttp is installed and follow_redirects is None, omits it."""
+        try:
+            import httpx_aiohttp
+        except ImportError:
+            self.skipTest("httpx_aiohttp not installed")
+
+        from cohere.base_client import _make_default_async_client
+
+        client = _make_default_async_client(timeout=300, follow_redirects=None)
+        self.assertIsInstance(client, httpx_aiohttp.HttpxAiohttpClient)
+        # httpx default is False when not specified
+        self.assertFalse(client.follow_redirects)
+
+    def test_explicit_httpx_client_bypasses_autodetect(self) -> None:
+        """When user passes httpx_client explicitly, auto-detect is not used."""
+        explicit_client = httpx.AsyncClient(timeout=60)
+        # Simulate what AsyncBaseCohere.__init__ does:
+        # httpx_client if httpx_client is not None else _make_default_async_client(...)
+        result = explicit_client if explicit_client is not None else None
+        self.assertIs(result, explicit_client)
+        self.assertEqual(result.timeout.read, 60)
+
+
+class TestDefaultClients(unittest.TestCase):
+    """Tests for convenience classes in _default_clients.py."""
+
+    def test_default_async_httpx_client_defaults(self) -> None:
+        """DefaultAsyncHttpxClient applies SDK defaults."""
+        from cohere._default_clients import COHERE_DEFAULT_TIMEOUT, DefaultAsyncHttpxClient
+
+        client = DefaultAsyncHttpxClient()
+        self.assertIsInstance(client, httpx.AsyncClient)
+        self.assertEqual(client.timeout.read, COHERE_DEFAULT_TIMEOUT)
+        self.assertTrue(client.follow_redirects)
+
+    def test_default_async_httpx_client_overrides(self) -> None:
+        """DefaultAsyncHttpxClient allows overriding defaults."""
+        from cohere._default_clients import DefaultAsyncHttpxClient
+
+        client = DefaultAsyncHttpxClient(timeout=60, follow_redirects=False)
+        self.assertEqual(client.timeout.read, 60)
+        self.assertFalse(client.follow_redirects)
+
+    def test_default_aiohttp_client_without_package(self) -> None:
+        """DefaultAioHttpClient raises RuntimeError when httpx_aiohttp not installed."""
+        with mock.patch.dict(sys.modules, {"httpx_aiohttp": None}):
+            # Need to reload the module to pick up the mock
+            import importlib
+            import cohere._default_clients
+
+            importlib.reload(cohere._default_clients)
+
+            with self.assertRaises(RuntimeError) as ctx:
+                cohere._default_clients.DefaultAioHttpClient()
+            self.assertIn("pip install cohere[aiohttp]", str(ctx.exception))
+
+            # Reload again to restore original state
+            importlib.reload(cohere._default_clients)
+
+    def test_default_aiohttp_client_with_package(self) -> None:
+        """DefaultAioHttpClient works when httpx_aiohttp is installed."""
+        try:
+            import httpx_aiohttp
+        except ImportError:
+            self.skipTest("httpx_aiohttp not installed")
+
+        from cohere._default_clients import COHERE_DEFAULT_TIMEOUT, DefaultAioHttpClient
+
+        client = DefaultAioHttpClient()
+        self.assertIsInstance(client, httpx_aiohttp.HttpxAiohttpClient)
+        self.assertEqual(client.timeout.read, COHERE_DEFAULT_TIMEOUT)
+        self.assertTrue(client.follow_redirects)