This commit is contained in:
agustin18 2026-11-06 19:31:27 -08:00 • committed by GitHub
commit b8880cd8c4
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 69 additions and 6 deletions

View file

@ -65,13 +65,13 @@ def timeout(timeout_duration: float = 0.0, exception_to_raise=Timeout):
@wraps(func)
async def async_wrapper(*args, **kwargs):
local_timeout_duration = timeout_duration
if "force_timeout" in kwargs:
local_timeout_duration = kwargs["force_timeout"]
elif "request_timeout" in kwargs and kwargs["request_timeout"] is not None:
local_timeout_duration = kwargs["request_timeout"]
local_timeout_duration: float | None = timeout_duration
if "force_timeout" in kwargs and isinstance(kwargs["force_timeout"], (int, float)):
local_timeout_duration = float(kwargs["force_timeout"])
elif "request_timeout" in kwargs and isinstance(kwargs["request_timeout"], (int, float)):
local_timeout_duration = float(kwargs["request_timeout"])
try:
value: Final = await asyncio.wait_for(func(*args, **kwargs), timeout=timeout_duration)
value: Final = await asyncio.wait_for(func(*args, **kwargs), timeout=local_timeout_duration)
return value
except asyncio.TimeoutError:
model: Final = args[0] if len(args) > 0 else kwargs["model"]

View file

@ -0,0 +1,63 @@
"""Unit tests for litellm.timeout decorator."""
import asyncio
import time
from typing import Final
import pytest
from litellm.exceptions import Timeout
from litellm.timeout import timeout
@pytest.mark.asyncio
@pytest.mark.parametrize("timeout_arg", ["request_timeout", "force_timeout"])
async def test_async_timeout_decorator_enforces_per_call_timeout(timeout_arg: str) -> None:
@timeout(timeout_duration=60.0)
async def hung_func(**kwargs):
await asyncio.sleep(10.0)
return "never"
with pytest.raises(Timeout) as exc_info:
await hung_func(**{timeout_arg: 0.001, "model": "test-model"})
assert "0.001 second(s)" in str(exc_info.value)
@pytest.mark.asyncio
@pytest.mark.parametrize("timeout_arg", ["request_timeout", "force_timeout"])
async def test_async_timeout_decorator_extends_duration_with_per_call_timeout(
timeout_arg: str,
) -> None:
@timeout(timeout_duration=0.001)
async def delayed_func(**kwargs):
await asyncio.sleep(0.01)
return "completed"
result: Final = await delayed_func(**{timeout_arg: 60.0, "model": "test-model"})
assert result == "completed"
@pytest.mark.asyncio
async def test_async_timeout_decorator_handles_none_force_timeout() -> None:
@timeout(timeout_duration=0.001)
async def hung_func(**kwargs):
await asyncio.sleep(10.0)
return "never"
with pytest.raises(Timeout) as exc_info:
await hung_func(force_timeout=None, model="test-model")
assert "0.001 second(s)" in str(exc_info.value)
def test_sync_timeout_decorator_enforces_per_call_timeout() -> None:
@timeout(timeout_duration=60.0)
def hung_sync_func(**kwargs):
time.sleep(10.0)
return "never"
with pytest.raises(Timeout) as exc_info:
hung_sync_func(request_timeout=0.001, model="test-model")
assert "0.001 second(s)" in str(exc_info.value)