"""Concurrency limiting infrastructure with OpenTelemetry observability.""" from __future__ import annotations as _annotations from abc import ABC, abstractmethod from collections.abc import AsyncIterator from contextlib import AbstractAsyncContextManager, asynccontextmanager from dataclasses import dataclass from typing import TypeAlias import anyio from opentelemetry.trace import Tracer, get_tracer from typing_extensions import Self __all__ = ( 'AbstractConcurrencyLimiter', 'ConcurrencyLimiter', 'ConcurrencyLimit', 'AnyConcurrencyLimit', ) class AbstractConcurrencyLimiter(ABC): """Abstract base class for concurrency limiters. Subclass this to create custom concurrency limiters (e.g., Redis-backed distributed limiters). Example: ```python from pydantic_ai.concurrency import AbstractConcurrencyLimiter class RedisConcurrencyLimiter(AbstractConcurrencyLimiter): def __init__(self, redis_client, key: str, max_running: int): self._redis = redis_client self._key = key self._max_running = max_running async def acquire(self, source: str) -> None: # Implement Redis-based distributed locking ... def release(self) -> None: # Release the Redis lock ... ``` """ @abstractmethod async def acquire(self, source: str) -> None: """Acquire a slot, waiting if necessary. Args: source: Identifier for observability (e.g., 'model:gpt-4o'). """ ... @abstractmethod def release(self) -> None: """Release a slot.""" ... @dataclass class ConcurrencyLimit: """Configuration for concurrency limiting with optional backpressure. Args: max_running: Maximum number of concurrent operations allowed. max_queued: Maximum number of operations waiting in the queue. If None, the queue is unlimited. If exceeded, raises `ConcurrencyLimitExceeded`. """ max_running: int max_queued: int | None = None class ConcurrencyLimiter(AbstractConcurrencyLimiter): """A concurrency limiter that tracks waiting operations for observability. This class wraps an anyio.CapacityLimiter and tracks the number of waiting operations. When an operation has to wait to acquire a slot, a span is created for observability purposes. """ def __init__( self, max_running: int, *, max_queued: int | None = None, name: str | None = None, tracer: Tracer | None = None, ): """Initialize the ConcurrencyLimiter. Args: max_running: Maximum number of concurrent operations. max_queued: Maximum queue depth before raising ConcurrencyLimitExceeded. name: Optional name for this limiter, used for observability when sharing a limiter across multiple models or agents. tracer: OpenTelemetry tracer for span creation. """ self._limiter = anyio.CapacityLimiter(max_running) self._max_queued = max_queued self._name = name self._tracer = tracer # Lock and counter to atomically check and track waiting tasks for max_queued enforcement self._queue_lock = anyio.Lock() self._waiting_count = 0 @classmethod def from_limit( cls, limit: int | ConcurrencyLimit, *, name: str | None = None, tracer: Tracer | None = None, ) -> Self: """Create a ConcurrencyLimiter from a ConcurrencyLimit configuration. Args: limit: Either an int for simple limiting or a ConcurrencyLimit for full config. name: Optional name for this limiter, used for observability. tracer: OpenTelemetry tracer for span creation. Returns: A configured ConcurrencyLimiter. """ if isinstance(limit, int): return cls(max_running=limit, name=name, tracer=tracer) else: return cls( max_running=limit.max_running, max_queued=limit.max_queued, name=name, tracer=tracer, ) @property def name(self) -> str | None: """Name of the limiter for observability.""" return self._name @property def waiting_count(self) -> int: """Number of operations currently waiting to acquire a slot.""" return self._waiting_count @property def running_count(self) -> int: """Number of operations currently running.""" return self._limiter.statistics().borrowed_tokens @property def available_count(self) -> int: """Number of slots available.""" return int(self._limiter.available_tokens) @property def max_running(self) -> int: """Maximum concurrent operations allowed.""" return int(self._limiter.total_tokens) def _get_tracer(self) -> Tracer: """Get the tracer, falling back to global tracer if not set.""" if self._tracer is not None: return self._tracer return get_tracer('pydantic-ai') async def acquire(self, source: str) -> None: """Acquire a slot, creating a span if waiting is required. Args: source: Identifier for the source of this acquisition (e.g., 'agent:my-agent' or 'model:gpt-4'). """ from .exceptions import ConcurrencyLimitExceeded # Try to acquire immediately without blocking try: self._limiter.acquire_nowait() return except anyio.WouldBlock: pass # We need to wait - atomically check queue limits and register ourselves as waiting # This prevents a race condition where multiple tasks could pass the check before # any of them actually start waiting on the limiter async with self._queue_lock: if self._max_queued is not None and self._waiting_count >= self._max_queued: # Use limiter name if set, otherwise use source for error messages display_name = self._name or source raise ConcurrencyLimitExceeded( f'Concurrency queue depth ({self._waiting_count + 1}) exceeds max_queued ({self._max_queued})' + (f' for {display_name}' if display_name else '') ) # Register ourselves as waiting before releasing the lock self._waiting_count += 1 # Now we're registered as waiting, proceed to wait on the limiter # Use try/finally to ensure we decrement the counter even on cancellation try: # Create a span for observability while waiting tracer = self._get_tracer() display_name = self._name or source attributes: dict[str, str | int] = { 'source': source, 'waiting_count': self._waiting_count, 'max_running': int(self._limiter.total_tokens), } if self._name is not None: attributes['limiter_name'] = self._name if self._max_queued is not None: attributes['max_queued'] = self._max_queued # Span name uses limiter name if set, otherwise source span_name = f'waiting for {display_name} concurrency' with tracer.start_as_current_span(span_name, attributes=attributes): await self._limiter.acquire() finally: # We're no longer waiting (either we acquired or we were cancelled) self._waiting_count -= 1 def release(self) -> None: """Release a slot.""" self._limiter.release() AnyConcurrencyLimit: TypeAlias = 'int | ConcurrencyLimit | AbstractConcurrencyLimiter | None' """Type alias for concurrency limit configuration. Can be: - An `int`: Simple limit on concurrent operations (unlimited queue). - A `ConcurrencyLimit`: Full configuration with optional backpressure. - An `AbstractConcurrencyLimiter`: A pre-created limiter instance for sharing across multiple models/agents. - `None`: No concurrency limiting (default). """ @asynccontextmanager async def _null_context() -> AsyncIterator[None]: """A no-op async context manager.""" yield @asynccontextmanager async def _limiter_context(limiter: AbstractConcurrencyLimiter, source: str) -> AsyncIterator[None]: """Context manager that acquires and releases a limiter with the given source.""" await limiter.acquire(source) try: yield finally: limiter.release() def get_concurrency_context( limiter: AbstractConcurrencyLimiter | None, source: str = 'unnamed', ) -> AbstractAsyncContextManager[None]: """Get an async context manager for the concurrency limiter. If limiter is None, returns a no-op context manager. Args: limiter: The AbstractConcurrencyLimiter or None. source: Identifier for the source of this acquisition (e.g., 'agent:my-agent' or 'model:gpt-4'). Returns: An async context manager. """ if limiter is None: return _null_context() return _limiter_context(limiter, source) def normalize_to_limiter( limit: AnyConcurrencyLimit, *, name: str | None = None, ) -> AbstractConcurrencyLimiter | None: """Normalize a concurrency limit configuration to an AbstractConcurrencyLimiter. Args: limit: The concurrency limit configuration. name: Optional name for the limiter if one is created. Returns: An AbstractConcurrencyLimiter if limit is not None, otherwise None. """ if limit is None: return None elif isinstance(limit, AbstractConcurrencyLimiter): return limit else: return ConcurrencyLimiter.from_limit(limit, name=name)