Python Context Managers for AI Resources: Managing API Clients, DB Connections, and Sessions

The Resource Leak Problem in AI Applications

AI agent applications manage many external resources simultaneously: API client sessions, database connections, file handles for vector stores, WebSocket connections for streaming, and temporary files for processing. If any exception occurs mid-pipeline, these resources must still be properly closed. Context managers guarantee cleanup happens regardless of how the block exits.

The with statement is Python's solution to the resource acquisition and release pattern. For AI engineers building long-running agent processes, getting this right is the difference between a stable system and one that leaks connections until it crashes.

Building an API Client Manager

The most common resource in AI applications is the HTTP client session. Creating a new session per request is wasteful. Sharing one session without proper lifecycle management leads to leaks.

import httpx
from typing import AsyncIterator

class LLMClient:
    def __init__(self, api_key: str, base_url: str):
        self.api_key = api_key
        self.base_url = base_url
        self._client: httpx.AsyncClient | None = None

    async def __aenter__(self) -> "LLMClient":
        self._client = httpx.AsyncClient(
            base_url=self.base_url,
            headers={"Authorization": f"Bearer {self.api_key}"},
            timeout=httpx.Timeout(30.0, connect=5.0),
        )
        return self

    async def __aexit__(self, exc_type, exc_val, exc_tb) -> None:
        if self._client:
            await self._client.aclose()
            self._client = None
        return False  # do not suppress exceptions

    async def complete(self, prompt: str) -> str:
        response = await self._client.post(
            "/v1/chat/completions",
            json={"model": "gpt-4o", "messages": [{"role": "user", "content": prompt}]},
        )
        response.raise_for_status()
        return response.json()["choices"][0]["message"]["content"]

# Usage - client is always properly closed
async def main():
    async with LLMClient("sk-...", "https://api.openai.com") as llm:
        answer = await llm.complete("What is agentic AI?")
        print(answer)

contextlib Shortcuts

For simpler cases, contextlib provides decorator-based context managers that avoid writing full classes.

from contextlib import asynccontextmanager, contextmanager
from typing import AsyncGenerator
import asyncpg

@asynccontextmanager
async def get_db_connection(dsn: str) -> AsyncGenerator[asyncpg.Connection, None]:
    conn = await asyncpg.connect(dsn)
    try:
        yield conn
    finally:
        await conn.close()

@asynccontextmanager
async def db_transaction(dsn: str) -> AsyncGenerator[asyncpg.Connection, None]:
    async with get_db_connection(dsn) as conn:
        tx = conn.transaction()
        await tx.start()
        try:
            yield conn
            await tx.commit()
        except Exception:
            await tx.rollback()
            raise

# Usage
async def save_agent_memory(dsn: str, agent_id: str, memory: dict):
    async with db_transaction(dsn) as conn:
        await conn.execute(
            "INSERT INTO agent_memories (agent_id, data) VALUES ($1, $2)",
            agent_id, memory,
        )

Managing Temporary Files for AI Processing

AI pipelines often need temporary files for audio transcription, image processing, or document parsing.

import tempfile
import os
from contextlib import contextmanager
from pathlib import Path

@contextmanager
def temp_audio_file(suffix: str = ".wav"):
    tmp = tempfile.NamedTemporaryFile(suffix=suffix, delete=False)
    try:
        yield Path(tmp.name)
    finally:
        tmp.close()
        if os.path.exists(tmp.name):
            os.unlink(tmp.name)

# Audio is always cleaned up, even if transcription fails
with temp_audio_file(".mp3") as audio_path:
    download_audio(url, audio_path)
    transcript = transcribe(audio_path)

Combining Multiple Context Managers

Agent pipelines often need several resources open simultaneously. Use contextlib.AsyncExitStack to manage dynamic sets of resources.

from contextlib import AsyncExitStack

async def run_agent_pipeline(config):
    async with AsyncExitStack() as stack:
        llm = await stack.enter_async_context(
            LLMClient(config.api_key, config.base_url)
        )
        db = await stack.enter_async_context(
            get_db_connection(config.db_dsn)
        )
        cache = await stack.enter_async_context(
            RedisConnection(config.redis_url)
        )

        # All three resources are guaranteed cleanup
        result = await llm.complete("Analyze this data")
        await db.execute("INSERT INTO results ...", result)
        await cache.set("latest_result", result)

FAQ

When should I use a class-based context manager versus contextlib?

Use @contextmanager or @asynccontextmanager for simple acquire-yield-release patterns. Use a class with __enter__/__exit__ when you need the context manager to maintain state, offer additional methods on the yielded object, or handle exception types selectively in __exit__.

Can context managers be nested safely?

Yes, and this is the recommended pattern. Nesting ensures resources are released in reverse acquisition order. AsyncExitStack is the cleanest approach when you need to manage a variable number of resources determined at runtime.

How do async context managers differ from sync ones?

Async context managers use __aenter__ and __aexit__ instead of __enter__ and __exit__, and must be used with async with. The key difference is that setup and teardown can perform I/O operations like closing network connections without blocking the event loop.

---

#Python #ContextManagers #ResourceManagement #AIEngineering #AgenticAI #LearnAI