Skip to content
Merged
Show file tree
Hide file tree
Changes from 2 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 9 additions & 0 deletions pydantic_ai_slim/pydantic_ai/concurrency.py
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,11 @@
)


def _validate_max_running(max_running: int) -> None:
if max_running < 1:
raise ValueError(f'max_running must be >= 1, got {max_running}')


class AbstractConcurrencyLimiter(ABC):
"""Abstract base class for concurrency limiters.

Expand Down Expand Up @@ -75,6 +80,9 @@ class ConcurrencyLimit:
max_running: int
max_queued: int | None = None

def __post_init__(self) -> None:
_validate_max_running(self.max_running)

Comment thread
coderabbitai[bot] marked this conversation as resolved.

class ConcurrencyLimiter(AbstractConcurrencyLimiter):
"""A concurrency limiter that tracks waiting operations for observability.
Expand All @@ -101,6 +109,7 @@ def __init__(
a limiter across multiple models or agents.
tracer: OpenTelemetry tracer for span creation.
"""
_validate_max_running(max_running)
self._limiter = anyio.CapacityLimiter(max_running)
self._max_queued = max_queued
self._name = name
Expand Down
29 changes: 28 additions & 1 deletion tests/test_concurrency.py
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,7 @@
import pytest

from pydantic_ai import Agent, ConcurrencyLimit, ConcurrencyLimiter, ConcurrencyLimitExceeded
from pydantic_ai.concurrency import get_concurrency_context
from pydantic_ai.concurrency import get_concurrency_context, normalize_to_limiter
from pydantic_ai.models.test import TestModel

if TYPE_CHECKING:
Expand Down Expand Up @@ -66,6 +66,18 @@ async def test_nowait_acquisition(self):
async with get_concurrency_context(limiter, 'test'):
pass # No waiting

@pytest.mark.parametrize('max_running', [0, -1])
async def test_invalid_max_running(self, max_running: int):
"""Test that invalid concurrency limits are rejected before creating the limiter."""
with pytest.raises(ValueError, match=f'max_running must be >= 1, got {max_running}'):
ConcurrencyLimiter(max_running=max_running)

@pytest.mark.parametrize('max_running', [0, -1])
async def test_invalid_concurrency_limit_config(self, max_running: int):
"""Test that invalid concurrency limit config values fail eagerly."""
with pytest.raises(ValueError, match=f'max_running must be >= 1, got {max_running}'):
ConcurrencyLimit(max_running=max_running)

async def test_waiting_count_tracking(self):
"""Test that waiting_count is accurately tracked."""
limiter = ConcurrencyLimiter(max_running=1)
Expand Down Expand Up @@ -193,6 +205,16 @@ async def test_from_limiter_config(self):
assert limiter.max_running == 5
assert limiter._max_queued == 10

async def test_from_invalid_limit(self):
"""Test that invalid normalized concurrency limits fail eagerly."""
with pytest.raises(ValueError, match='max_running must be >= 1, got 0'):
ConcurrencyLimiter.from_limit(0)

async def test_normalize_to_limiter_rejects_invalid_limit(self):
"""Test that invalid normalized concurrency limits fail before use."""
with pytest.raises(ValueError, match='max_running must be >= 1, got 0'):
normalize_to_limiter(0)

async def test_properties(self):
"""Test the various properties of ConcurrencyLimiter."""
limiter = ConcurrencyLimiter(max_running=5, name='test-limiter')
Expand Down Expand Up @@ -304,6 +326,11 @@ async def test_agent_with_limiter_concurrency(self):
assert agent._concurrency_limiter.max_running == 5
assert agent._concurrency_limiter._max_queued == 10

async def test_agent_invalid_int_concurrency_limit(self):
"""Test that invalid agent concurrency limits fail before a run can hang."""
with pytest.raises(ValueError, match='max_running must be >= 1, got 0'):
Agent(TestModel(), max_concurrency=0)


class TestConcurrencyLimitedModel:
"""Tests for the ConcurrencyLimitedModel wrapper."""
Expand Down
Loading