Async context managers for LLM API clients in Python

Most LLM integrations in Python start as a bare POST and degrade into a tangle of unclosed sockets and orphaned streams. A well-designed python async context managers llm client keeps connection pools warm, cancels in-flight generations on timeout, and guarantees cleanup when your coroutine is cancelled. This guide gives an ordered path to building one that streams tokens and survives real production pressure.

Why async context managers matter for LLM calls

LLM endpoints hold connections open for seconds while tokens trickle in. Under asyncio, a forgotten response body or an unclosed client leaks file descriptors and stalls the event loop. The async with protocol forces you to pair acquisition with release, even when the body raises or gets cancelled.

Synchronous clients block the loop; threading mitigations waste memory per request. An async client with a shared pool multiplies throughput on the same hardware.

Step 1: Choose a transport that speaks async and streams

httpx.AsyncClient is the pragmatic default. It supports HTTP/2, connection pooling, and async for streaming. Avoid wrapping requests in run_in_executor—you lose cancellation granularity.

import httpx

async def raw_stream(prompt: str):
    async with httpx.AsyncClient(timeout=30.0) as client:
        async with client.stream(
            "POST",
            "https://api.openai.com/v1/chat/completions",
            json={"model": "gpt-4o-mini", "messages": [{"role": "user", "content": prompt}], "stream": True},
            headers={"Authorization": "Bearer $KEY"},
        ) as resp:
            async for line in resp.aiter_lines():
                if line.strip():
                    print(line)

That snippet works, but instantiating a client per call throws away the pool. We need a client that lives across requests but still cleans up.

Step 2: Implement the async context manager protocol

Define a class with __aenter__ and __aexit__. Store the httpx.AsyncClient as a member, created on enter, closed on exit. This turns your python async context managers llm client into a first-class resource.

class LLMClient:
    def __init__(self, base_url: str, api_key: str, timeout: float = 30.0):
        self._base = base_url
        self._key = api_key
        self._timeout = timeout
        self._client: httpx.AsyncClient | None = None

    async def __aenter__(self) -> "LLMClient":
        self._client = httpx.AsyncClient(
            base_url=self._base,
            timeout=self._timeout,
            headers={"Authorization": f"Bearer {self._key}"},
        )
        return self

    async def __aexit__(self, exc_type, exc, tb) -> None:
        if self._client:
            await self._client.aclose()
            self._client = None

Use it:

async with LLMClient("https://api.openai.com/v1", key) as llm:
    await llm.stream_chat("Explain asyncio")

Step 3: Stream completions without leaking responses

Streaming methods must themselves use async with on the response. If you return from inside a stream without exiting, the connection stays open. Wrap the stream in a generator that yields parsed chunks and relies on the outer context to close the client.

class LLMClient:
    # ... __aenter__/__aexit__ as above ...

    async def stream_chat(self, prompt: str, model: str = "gpt-4o-mini"):
        if not self._client:
            raise RuntimeError("Client not entered")
        payload = {
            "model": model,
            "messages": [{"role": "user", "content": prompt}],
            "stream": True,
        }
        async with self._client.stream("POST", "/chat/completions", json=payload) as resp:
            resp.raise_for_status()
            async for line in resp.aiter_lines():
                if not line.startswith("data:"):
                    continue
                data = line[len("data:"):].strip()
                if data == "[DONE]":
                    break
                yield data

Callers consume with async for. The async with on resp guarantees the stream closes even if the caller cancels.

Step 4: Handle cancellation and timeouts explicitly

Asyncio cancellation propagates as asyncio.CancelledError. If you catch it and don’t re-raise, the context manager exit still runs but the task is stuck. Use asyncio.timeout to bound generation time, and never swallow cancellation inside the stream loop.

import asyncio

async def bounded_generate(llm: LLMClient, prompt: str, max_seconds: float = 10.0):
    try:
        async with asyncio.timeout(max_seconds):
            async for chunk in llm.stream_chat(prompt):
                yield chunk
    except asyncio.CancelledError:
        # Propagate; cleanup happens in __aexit__
        raise

A common bug: wrapping async for in try/except Exception that catches CancelledError (it inherits from BaseException in modern Python). Be explicit and let it propagate.

Step 5: Retries and provider fallback

Transient 429s and 5xx are normal. Implement a small retry loop with backoff, but keep it outside the stream—retry the whole request, not individual chunks. If you point your python async context managers llm client at a gateway that already performs automatic fallback when a provider is rate-limited, such as n4n.ai’s OpenAI-compatible endpoint covering 240+ models, you can skip manual provider selection and focus on transport retries.

async def stream_with_retry(llm: LLMClient, prompt: str, attempts: int = 3):
    for i in range(attempts):
        try:
            async for chunk in llm.stream_chat(prompt):
                yield chunk
            return
        except (httpx.HTTPStatusError, httpx.TransportError) as e:
            if i == attempts - 1:
                raise
            await asyncio.sleep(2 ** i)

The gateway forwards cache-control hints and honors routing directives, so repeated identical prompts may hit provider-side caches without extra code.

Common pitfalls

Creating a new client per request. Defeats pooling. Use one LLMClient per process or per request batch via async with.

Holding the client open forever without supervising. If your app runs indefinitely, close the client on shutdown. Use a lifespan hook in FastAPI or an atexit async handler.

Mixing sync and async. Calling a sync client inside async def blocks the loop. Use AsyncOpenAI or your own httpx wrapper.

Ignoring aiter_lines buffering. SSE lines can split. Parse with json.loads after stripping data:; don’t assume one line equals one JSON object.

Not forwarding cache headers. If your gateway or provider supports prompt caching, pass "cache_control" in the body or relevant headers. The python async context managers llm client should expose a way to set request headers per call.

Tradeoffs

A long-lived client shares a pool (good), but a misconfigured pool size (limits=httpx.Limits(max_connections=100)) can exhaust file descriptors under fan-out. A per-request client isolates failures but adds TLS handshake latency per call.

Streaming saves memory on long outputs but binds a connection for the full generation. If you need high concurrency with long outputs, size the pool for peak simultaneous streams, not peak QPS.

Context managers add boilerplate. For scripts, a module-level singleton with explicit await client.aclose() at exit is fine. For libraries, force async with so callers can’t leak.

Minimal working example

import asyncio
import httpx

class LLMClient:
    def __init__(self, base_url: str, api_key: str, timeout: float = 30.0,
                 limits: httpx.Limits | None = None):
        self._base = base_url
        self._key = api_key
        self._timeout = timeout
        self._limits = limits or httpx.Limits(max_connections=50)
        self._client: httpx.AsyncClient | None = None

    async def __aenter__(self) -> "LLMClient":
        self._client = httpx.AsyncClient(
            base_url=self._base,
            timeout=self._timeout,
            limits=self._limits,
            headers={"Authorization": f"Bearer {self._key}"},
        )
        return self

    async def __aexit__(self, *args) -> None:
        if self._client:
            await self._client.aclose()
            self._client = None

    async def stream_chat(self, prompt: str, model: str = "gpt-4o-mini"):
        if not self._client:
            raise RuntimeError("Enter context first")
        payload = {
            "model": model,
            "messages": [{"role": "user", "content": prompt}],
            "stream": True,
        }
        async with self._client.stream("POST", "/chat/completions", json=payload) as resp:
            resp.raise_for_status()
            async for line in resp.aiter_lines():
                if line.startswith("data:"):
                    data = line[5:].strip()
                    if data == "[DONE]":
                        break
                    yield data

async def main():
    async with LLMClient("https://api.openai.com/v1", "YOUR_KEY") as llm:
        async for chunk in llm.stream_chat("What is asyncio?"):
            print(chunk)

if __name__ == "__main__":
    asyncio.run(main())

That is a complete, production-shaped starting point. Extend it with retry, metrics, and routing as your system demands. The discipline of python async context managers llm client design pays off the first time a request cancels and nothing leaks.