Async Context Managers

async with EXPR as TARGET:
    BODY

manager = EXPR
aexit = type(manager).__aexit__
value = await type(manager).__aenter__(manager)
try:
    TARGET = value
    BODY
except BaseException as exc:
    if not await aexit(manager, type(exc), exc, exc.__traceback__):
        raise
else:
    await aexit(manager, None, None, None)

from types import TracebackType
from typing import Optional, Type

import asyncpg


class VectorTxn:
    """asyncpg transaction around pgvector writes."""

    def __init__(self, pool: asyncpg.Pool) -> None:
        self._pool = pool

    async def __aenter__(self) -> asyncpg.Connection:
        self._conn = await self._pool.acquire()
        try:
            self._tx = self._conn.transaction()
            await self._tx.start()
        except BaseException:
            await self._pool.release(self._conn)   # __aexit__ will not run
            raise
        return self._conn

    async def __aexit__(
        self,
        exc_type: Optional[Type[BaseException]],
        exc: Optional[BaseException],
        tb: Optional[TracebackType],
    ) -> None:
        try:
            if exc_type is None:
                await self._tx.commit()
            else:
                await self._tx.rollback()
        finally:
            await self._pool.release(self._conn)


async def upsert_chunks(pool, chunks) -> None:
    async with VectorTxn(pool) as conn:
        await conn.executemany(
            "INSERT INTO chunks (id, text, embedding) VALUES ($1, $2, $3) "
            "ON CONFLICT (id) DO UPDATE SET text = EXCLUDED.text, embedding = EXCLUDED.embedding",
            [(c.id, c.text, c.embedding) for c in chunks],
        )

from contextlib import asynccontextmanager
from typing import AsyncIterator


@asynccontextmanager
async def leased_connection(pool) -> AsyncIterator["asyncpg.Connection"]:
    conn = await pool.acquire()
    try:
        yield conn
    finally:
        await pool.release(conn)

import asyncio


class Lease:
    def __init__(self, broker, job_id: str) -> None:
        self.broker, self.job_id = broker, job_id

    async def __aenter__(self) -> "Lease":
        await self.broker.claim(self.job_id)
        return self

    async def __aexit__(self, exc_type, exc, tb) -> None:
        # The job must be released even if we are being cancelled; shield the await
        # so a second cancellation does not interrupt the release request itself.
        await asyncio.shield(self.broker.release(self.job_id))

import asyncio

EMBED_LIMIT = asyncio.Semaphore(16)        # create ONCE, share across tasks
                                           # (3.10+: no loop binding at construction)


async def embed_one(client, text: str) -> list[float]:
    async with EMBED_LIMIT:                # acquire; release guaranteed
        resp = await client.embeddings.create(model="text-embedding-3-small", input=text)
        return resp.data[0].embedding


async def embed_all(client, texts: list[str]) -> list[list[float]]:
    return await asyncio.gather(*(embed_one(client, t) for t in texts))

import asyncio
import contextlib
from contextlib import asynccontextmanager
from typing import AsyncIterator, Awaitable, Callable


@asynccontextmanager
async def background(fn: Callable[[], Awaitable[None]], interval: float) -> AsyncIterator[asyncio.Task]:
    async def loop() -> None:
        while True:
            await fn()
            await asyncio.sleep(interval)

    task = asyncio.create_task(loop())
    try:
        yield task
    finally:
        task.cancel()
        with contextlib.suppress(asyncio.CancelledError):
            await task                    # wait for it to actually stop


async def run_long_agent(session) -> None:
    async with background(session.heartbeat, interval=10):
        await session.run_steps()

import asyncio


async def answer_with_tools(question: str) -> dict:
    async with asyncio.timeout(30):                 # 3.11+: cancels the block on expiry
        async with asyncio.TaskGroup() as tg:       # 3.11+
            web = tg.create_task(search_web(question))
            docs = tg.create_task(search_vector_store(question))
            sql = tg.create_task(query_warehouse(question))
    return {"web": web.result(), "docs": docs.result(), "sql": sql.result()}


try:
    result = asyncio.run(answer_with_tools("Q3 revenue by region?"))
except* TimeoutError:                                # 3.11+ except*
    ...

import httpx


class ChatClient:
    """Thin LLM API wrapper that owns its HTTP pool."""

    def __init__(self, base_url: str, api_key: str) -> None:
        self._http = httpx.AsyncClient(
            base_url=base_url,
            headers={"Authorization": f"Bearer {api_key}"},
            timeout=httpx.Timeout(60.0, connect=5.0),
            limits=httpx.Limits(max_connections=50, max_keepalive_connections=20),
        )

    async def complete(self, messages: list[dict], model: str) -> str:
        r = await self._http.post("/v1/chat/completions", json={"model": model, "messages": messages})
        r.raise_for_status()
        return r.json()["choices"][0]["message"]["content"]

    async def aclose(self) -> None:
        await self._http.aclose()

    async def __aenter__(self) -> "ChatClient":
        return self

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

import aiohttp


async def fetch_all(urls: list[str]) -> list[str]:
    timeout = aiohttp.ClientTimeout(total=30)
    async with aiohttp.ClientSession(timeout=timeout) as session:
        async def one(url: str) -> str:
            async with session.get(url) as r:      # releases the connection back to the pool
                r.raise_for_status()
                return await r.text()
        return await asyncio.gather(*(one(u) for u in urls))

async with client.stream("POST", "/v1/chat/completions", json=payload) as r:
    async for line in r.aiter_lines():
        handle(line)

from contextlib import aclosing   # 3.10+


async def read_events(stream):
    async with aclosing(sse_events(stream)) as events:
        async for ev in events:
            if ev.type == "done":
                break

import time


class Span:
    def __init__(self, name: str) -> None:
        self.name = name

    def _start(self) -> "Span":
        self.t0 = time.perf_counter()
        return self

    def _end(self, exc: BaseException | None) -> None:   # 3.10+ union syntax; use Optional on 3.9
        status = "error" if exc else "ok"
        export(self.name, time.perf_counter() - self.t0, status)

    def __enter__(self) -> "Span":
        return self._start()

    def __exit__(self, exc_type, exc, tb) -> None:
        self._end(exc)

    async def __aenter__(self) -> "Span":
        return self._start()

    async def __aexit__(self, exc_type, exc, tb) -> None:
        self._end(exc)