Migrating httpx.AsyncClient to httpx_aiohttp.HttpxAiohttpClient (#287)

Author: jamesbraza

packages/hotpotqa/pyproject.toml

@@ -22,6 +22,7 @@ dependencies = [
     "datasets>=2.15,<4",  # Lower pin for https://github.com/huggingface/datasets/pull/6404, upper pin for https://huggingface.co/datasets/hotpotqa/hotpot_qa/discussions/8
     "fhaviary",
     "httpx",
+    "httpx-aiohttp",
     "pydantic~=2.0",
     "tenacity",
 ]

packages/hotpotqa/src/aviary/envs/hotpotqa/env.py

@@ -24,6 +24,7 @@
 from uuid import UUID
 
 import httpx
+import httpx_aiohttp
 from bs4 import BeautifulSoup
 from datasets import load_dataset
 from pydantic import BaseModel, ConfigDict, Field, field_validator
@@ -424,7 +425,7 @@ async def search(self, entity: str) -> str:
             response_text = self.wiki_cache[search_entity]
         except KeyError:
             # follow_redirects=True because wikipedia frequently redirects to the correct page
-            async with httpx.AsyncClient(
+            async with httpx_aiohttp.HttpxAiohttpClient(
                 follow_redirects=True, proxy=self.proxy
             ) as client:
                 response = await fetch_with_retry(

pyproject.toml

@@ -24,6 +24,7 @@ classifiers = [
 dependencies = [
     "docstring_parser>=0.16",  # Pin for description addition
     "httpx",
+    "httpx-aiohttp",
     "pydantic~=2.0",
 ]
 description = "Gymnasium framework for training language model agents on constructive tasks"
@@ -46,7 +47,7 @@ dev = [
     "fhaviary[image,llm,server,typing,xml]",
     "ipython>=8",  # Pin to keep recent
     "jupyter>=1.0.0",  # For running notebooks
-    "litellm>=1.65.5,<1.71",  # Lower pin for sending tool schemae title in completions, upper pin for VCR cassette breaks (https://github.com/BerriAI/litellm/issues/11724)
+    "litellm>=1.71",  # Lower pin for aiohttp transport adoption
     "mypy>=1.8",  # Pin for mutable-override
     "numpy>=1",  # Pin to keep recent
     "pre-commit>=3.4",  # Pin to keep recent

src/aviary/env_client.py

@@ -4,6 +4,7 @@
 from typing import Any, Generic, TypeVar, cast
 
 import httpx
+import httpx_aiohttp
 from pydantic import BaseModel, Field
 
 from aviary.env import Environment, TaskDataset
@@ -14,7 +15,9 @@
 
 # Not sure why, but mypy complains if we use the TEnvState in aviary.env, so redefine here
 TEnvState = TypeVar("TEnvState")
-TClient = TypeVar("TClient", httpx.Client, httpx.AsyncClient)
+TClient = TypeVar(
+    "TClient", httpx.Client, httpx.AsyncClient, httpx_aiohttp.HttpxAiohttpClient
+)
 
 
 class EnvironmentClient(Environment[TEnvState], ABC, Generic[TEnvState]):
@@ -35,7 +38,7 @@ def __init__(
         self._api_key = api_key
 
     async def _post(self, url: str, json: Mapping[str, Any]) -> httpx.Response:
-        async with httpx.AsyncClient() as client:
+        async with httpx_aiohttp.HttpxAiohttpClient() as client:
             headers = httpx.Headers(self._request_headers)
             if self._api_key:
                 headers["X-API-Key"] = self._api_key
@@ -142,7 +145,7 @@ def __init__(
 
     def _get_http_client(
         self,
-        client_class: type[TClient] = httpx.AsyncClient,  # type: ignore[assignment]
+        client_class: type[TClient] = httpx_aiohttp.HttpxAiohttpClient,  # type: ignore[assignment]
     ) -> TClient:
         headers = {}
         if self.api_key:

tests/conftest.py

@@ -1,13 +1,16 @@
+from collections.abc import AsyncIterator
 from typing import TYPE_CHECKING, Any
 from urllib.parse import parse_qs, urlencode, urlparse
 
+import httpx_aiohttp
+import litellm.llms.custom_httpx.aiohttp_transport
 import pytest
-import vcr.request
+import vcr.stubs.httpx_stubs
 
 from . import CASSETTES_DIR
 
 if TYPE_CHECKING:
-    import vcr.request
+    import vcr.request  # noqa: TC004
 
     from aviary.core import DummyEnv
 
@@ -56,3 +59,38 @@ def fixture_vcr_config() -> dict[str, Any]:
         "allow_playback_repeats": True,
         "cassette_library_dir": str(CASSETTES_DIR),
     }
+
+
+class PreReadCompatibleAiohttpResponseStream(
+    httpx_aiohttp.transport.AiohttpResponseStream
+):
+    """aiohttp-backed response stream that works if the response was pre-read."""
+
+    async def __aiter__(self) -> AsyncIterator[bytes]:
+        with httpx_aiohttp.transport.map_aiohttp_exceptions():
+            if self._aiohttp_response._body is not None:
+                # Happens if some intermediary called `await _aiohttp_response.read()`
+                # TODO: take into account chunk size
+                yield self._aiohttp_response._body
+            else:
+                async for chunk in self._aiohttp_response.content.iter_chunked(
+                    self.CHUNK_SIZE
+                ):
+                    yield chunk
+
+
+async def _async_vcr_send(cassette, real_send, *args, **kwargs):  # noqa: ARG001
+    """VCR send that only sends, not possibly recording or playing back responses."""
+    return await real_send(*args, **kwargs)
+
+
+# Permanently patch the original response stream,
+# to work around https://github.com/karpetrosyan/httpx-aiohttp/issues/23
+# and https://github.com/BerriAI/litellm/issues/11724
+httpx_aiohttp.transport.AiohttpResponseStream = (  # type: ignore[misc]
+    litellm.llms.custom_httpx.aiohttp_transport.AiohttpResponseStream  # type: ignore[misc]
+) = PreReadCompatibleAiohttpResponseStream  # type: ignore[assignment]
+
+# Permanently patch vcrpy's async VCR recording functionality,
+# to work around https://github.com/kevin1024/vcrpy/issues/944
+vcr.stubs.httpx_stubs._async_vcr_send = _async_vcr_send

tests/test_envs.py

@@ -562,6 +562,9 @@ async def test_dummyenv_using_empty_params(self, dummy_env: DummyEnv) -> None:
 def server_async_client() -> AsyncClient:
     dataset = TaskDataset.from_name("dummy")
     server = TaskDatasetServer[DummyEnv](dataset)
+    # Use httpx.AsyncClient over httpx_aiohttp.HttpxAiohttpClient in tests here,
+    # as httpx_aiohttp.AiohttpTransport doesn't support an app argument
+    # as of httpx-aiohttp==0.1.8
     return AsyncClient(transport=ASGITransport(app=server.app), base_url="http://test")