don't share httpx.AsyncHTTPTransport between event loops

Author: ewjoachim

pydantic_ai_slim/pydantic_ai/models/__init__.py

@@ -6,12 +6,13 @@
 
 from __future__ import annotations as _annotations
 
+import asyncio
 from abc import ABC, abstractmethod
 from collections.abc import AsyncIterator, Iterator
 from contextlib import asynccontextmanager, contextmanager
 from dataclasses import dataclass, field
 from datetime import datetime
-from functools import cache
+from functools import cache, lru_cache
 
 import httpx
 from typing_extensions import Literal, TypeAliasType
@@ -495,25 +496,35 @@ def cached_async_http_client(*, provider: str | None = None, timeout: int = 600,
     The default timeouts match those of OpenAI,
     see <https://github.com/openai/openai-python/blob/v1.54.4/src/openai/_constants.py#L9>.
     """
-    client = _cached_async_http_client(provider=provider, timeout=timeout, connect=connect)
+    try:
+        loop = asyncio.get_running_loop()
+    except RuntimeError:
+        loop = None
+
+    client = _cached_async_http_client(loop=loop, provider=provider, timeout=timeout, connect=connect)
     if client.is_closed:
         # This happens if the context manager is used, so we need to create a new client.
         _cached_async_http_client.cache_clear()
-        client = _cached_async_http_client(provider=provider, timeout=timeout, connect=connect)
+        client = _cached_async_http_client(loop=loop, provider=provider, timeout=timeout, connect=connect)
     return client
 
 
-@cache
-def _cached_async_http_client(provider: str | None, timeout: int = 600, connect: int = 5) -> httpx.AsyncClient:
+@lru_cache(maxsize=32)
+def _cached_async_http_client(
+    loop: asyncio.AbstractEventLoop, provider: str | None, timeout: int = 600, connect: int = 5
+) -> httpx.AsyncClient:
     return httpx.AsyncClient(
-        transport=_cached_async_http_transport(),
+        transport=_cached_async_http_transport(loop=loop),
         timeout=httpx.Timeout(timeout=timeout, connect=connect),
         headers={'User-Agent': get_user_agent()},
     )
 
 
-@cache
-def _cached_async_http_transport() -> httpx.AsyncHTTPTransport:
+@lru_cache(maxsize=32)
+def _cached_async_http_transport(loop: asyncio.AbstractEventLoop) -> httpx.AsyncHTTPTransport:
+    # The loop argument is unused, but it's here to ensure the cache key is different
+    # for each event loop, because a `httpx.AsyncHTTPTransport instanciated in a loop
+    # cannot be used in another loop.
     return httpx.AsyncHTTPTransport()