mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-06 08:16:43 +00:00
Merge pull request #38234 from BerriAI/litellm_request_timeouts
fix(proxy): give every `requests` call a timeout so a silent server cannot hang the caller
This commit is contained in:
commit
c62c2afa09
23 changed files with 293 additions and 59 deletions
|
|
@ -8,16 +8,19 @@ from .exceptions import UnauthorizedError
|
|||
|
||||
|
||||
class ChatClient:
|
||||
def __init__(self, base_url: str, api_key: str | None = None):
|
||||
def __init__(self, base_url: str, api_key: str | None = None, timeout: int = 600):
|
||||
"""
|
||||
Initialize the ChatClient.
|
||||
|
||||
Args:
|
||||
base_url (str): The base URL of the LiteLLM proxy server (e.g., "http://localhost:8000")
|
||||
api_key (Optional[str]): API key for authentication. If provided, it will be sent as a Bearer token.
|
||||
timeout (int): Request timeout in seconds (default: 600, the OpenAI SDK default, since a completion
|
||||
can legitimately take minutes)
|
||||
"""
|
||||
self._base_url = base_url.rstrip("/") # Remove trailing slash if present
|
||||
self._api_key = api_key
|
||||
self._timeout = timeout
|
||||
|
||||
def _get_headers(self) -> dict[str, str]:
|
||||
"""
|
||||
|
|
@ -96,7 +99,7 @@ class ChatClient:
|
|||
# Prepare and send the request
|
||||
session: Final = requests.Session()
|
||||
try:
|
||||
response: Final = session.send(request.prepare())
|
||||
response: Final = session.send(request.prepare(), timeout=self._timeout)
|
||||
response.raise_for_status()
|
||||
return response.json()
|
||||
except requests.exceptions.HTTPError as e:
|
||||
|
|
@ -161,7 +164,9 @@ class ChatClient:
|
|||
# Make streaming request
|
||||
session: Final = requests.Session()
|
||||
try:
|
||||
response: Final = session.post(url, headers=self._get_headers(), json=data, stream=True)
|
||||
response: Final = session.post(
|
||||
url, headers=self._get_headers(), json=data, stream=True, timeout=self._timeout
|
||||
)
|
||||
response.raise_for_status()
|
||||
|
||||
# Parse SSE stream
|
||||
|
|
|
|||
|
|
@ -99,11 +99,6 @@ class CliPollData(TypedDict, total=False):
|
|||
team_id: str
|
||||
|
||||
|
||||
class CliPollRequestKwargs(TypedDict, total=False):
|
||||
timeout: int
|
||||
headers: dict[str, str]
|
||||
|
||||
|
||||
class CliSsoStartData(TypedDict):
|
||||
login_id: str
|
||||
poll_secret: str
|
||||
|
|
@ -518,10 +513,7 @@ def _poll_for_ready_data(
|
|||
) -> CliPollData | None:
|
||||
for attempt in range(total_timeout // poll_interval):
|
||||
try:
|
||||
request_kwargs: CliPollRequestKwargs = {"timeout": request_timeout}
|
||||
if headers is not None:
|
||||
request_kwargs["headers"] = headers
|
||||
response = requests.get(url, **request_kwargs)
|
||||
response = requests.get(url, headers=headers, timeout=request_timeout)
|
||||
if response.status_code == 200:
|
||||
data: CliPollData = response.json()
|
||||
status = data.get("status")
|
||||
|
|
|
|||
|
|
@ -24,7 +24,8 @@ class Client:
|
|||
Args:
|
||||
base_url (str): The base URL of the LiteLLM proxy server (e.g., "http://localhost:4000")
|
||||
api_key (Optional[str]): API key for authentication. If provided, it will be sent as a Bearer token.
|
||||
timeout: Request timeout in seconds (default: 30)
|
||||
timeout: Request timeout in seconds for management calls (default: 30). Chat completions keep
|
||||
ChatClient's own 600 second default, since a completion can legitimately take minutes
|
||||
"""
|
||||
self._base_url = base_url.rstrip("/")
|
||||
# Only use the stored CLI key when it was issued for this server.
|
||||
|
|
@ -33,9 +34,9 @@ class Client:
|
|||
# Initialize resource clients
|
||||
|
||||
self.http = HTTPClient(base_url=base_url, api_key=self._api_key, timeout=timeout)
|
||||
self.models = ModelsManagementClient(base_url=self._base_url, api_key=self._api_key)
|
||||
self.model_groups = ModelGroupsManagementClient(base_url=self._base_url, api_key=self._api_key)
|
||||
self.models = ModelsManagementClient(base_url=self._base_url, api_key=self._api_key, timeout=timeout)
|
||||
self.model_groups = ModelGroupsManagementClient(base_url=self._base_url, api_key=self._api_key, timeout=timeout)
|
||||
self.chat = ChatClient(base_url=self._base_url, api_key=self._api_key)
|
||||
self.keys = KeysManagementClient(base_url=self._base_url, api_key=self._api_key)
|
||||
self.credentials = CredentialsManagementClient(base_url=self._base_url, api_key=self._api_key)
|
||||
self.teams = TeamsManagementClient(base_url=self._base_url, api_key=self._api_key)
|
||||
self.keys = KeysManagementClient(base_url=self._base_url, api_key=self._api_key, timeout=timeout)
|
||||
self.credentials = CredentialsManagementClient(base_url=self._base_url, api_key=self._api_key, timeout=timeout)
|
||||
self.teams = TeamsManagementClient(base_url=self._base_url, api_key=self._api_key, timeout=timeout)
|
||||
|
|
|
|||
|
|
@ -6,16 +6,18 @@ from .exceptions import UnauthorizedError
|
|||
|
||||
|
||||
class CredentialsManagementClient:
|
||||
def __init__(self, base_url: str, api_key: str | None = None):
|
||||
def __init__(self, base_url: str, api_key: str | None = None, timeout: int = 30):
|
||||
"""
|
||||
Initialize the CredentialsManagementClient.
|
||||
|
||||
Args:
|
||||
base_url (str): The base URL of the LiteLLM proxy server (e.g., "http://localhost:8000")
|
||||
api_key (Optional[str]): API key for authentication. If provided, it will be sent as a Bearer token.
|
||||
timeout (int): Request timeout in seconds (default: 30)
|
||||
"""
|
||||
self._base_url = base_url.rstrip("/") # Remove trailing slash if present
|
||||
self._api_key = api_key
|
||||
self._timeout = timeout
|
||||
|
||||
def _get_headers(self) -> dict[str, str]:
|
||||
"""
|
||||
|
|
@ -56,7 +58,7 @@ class CredentialsManagementClient:
|
|||
|
||||
session: Final = requests.Session()
|
||||
try:
|
||||
response: Final = session.send(request.prepare())
|
||||
response: Final = session.send(request.prepare(), timeout=self._timeout)
|
||||
response.raise_for_status()
|
||||
return response.json()
|
||||
except requests.exceptions.HTTPError as e:
|
||||
|
|
@ -103,7 +105,7 @@ class CredentialsManagementClient:
|
|||
|
||||
session: Final = requests.Session()
|
||||
try:
|
||||
response: Final = session.send(request.prepare())
|
||||
response: Final = session.send(request.prepare(), timeout=self._timeout)
|
||||
response.raise_for_status()
|
||||
return response.json()
|
||||
except requests.exceptions.HTTPError as e:
|
||||
|
|
@ -140,7 +142,7 @@ class CredentialsManagementClient:
|
|||
|
||||
session: Final = requests.Session()
|
||||
try:
|
||||
response: Final = session.send(request.prepare())
|
||||
response: Final = session.send(request.prepare(), timeout=self._timeout)
|
||||
response.raise_for_status()
|
||||
return response.json()
|
||||
except requests.exceptions.HTTPError as e:
|
||||
|
|
@ -177,7 +179,7 @@ class CredentialsManagementClient:
|
|||
|
||||
session: Final = requests.Session()
|
||||
try:
|
||||
response: Final = session.send(request.prepare())
|
||||
response: Final = session.send(request.prepare(), timeout=self._timeout)
|
||||
response.raise_for_status()
|
||||
return response.json()
|
||||
except requests.exceptions.HTTPError as e:
|
||||
|
|
|
|||
|
|
@ -9,16 +9,18 @@ from .exceptions import UnauthorizedError
|
|||
|
||||
|
||||
class KeysManagementClient:
|
||||
def __init__(self, base_url: str, api_key: str | None = None):
|
||||
def __init__(self, base_url: str, api_key: str | None = None, timeout: int = 30):
|
||||
"""
|
||||
Initialize the KeysManagementClient.
|
||||
|
||||
Args:
|
||||
base_url (str): The base URL of the LiteLLM proxy server (e.g., "http://localhost:8000")
|
||||
api_key (Optional[str]): API key for authentication. If provided, it will be sent as a Bearer token.
|
||||
timeout (int): Request timeout in seconds (default: 30)
|
||||
"""
|
||||
self._base_url = base_url.rstrip("/") # Remove trailing slash if present
|
||||
self._api_key = api_key
|
||||
self._timeout = timeout
|
||||
|
||||
def _get_headers(self) -> dict[str, str]:
|
||||
"""
|
||||
|
|
@ -99,7 +101,7 @@ class KeysManagementClient:
|
|||
|
||||
session: Final = requests.Session()
|
||||
try:
|
||||
response: Final = session.send(request.prepare())
|
||||
response: Final = session.send(request.prepare(), timeout=self._timeout)
|
||||
response.raise_for_status()
|
||||
return response.json()
|
||||
except requests.exceptions.HTTPError as e:
|
||||
|
|
@ -174,7 +176,7 @@ class KeysManagementClient:
|
|||
|
||||
session: Final = requests.Session()
|
||||
try:
|
||||
response: Final = session.send(request.prepare())
|
||||
response: Final = session.send(request.prepare(), timeout=self._timeout)
|
||||
response.raise_for_status()
|
||||
return response.json()
|
||||
except requests.exceptions.HTTPError as e:
|
||||
|
|
@ -218,7 +220,7 @@ class KeysManagementClient:
|
|||
|
||||
session: Final = requests.Session()
|
||||
try:
|
||||
response: Final = session.send(request.prepare())
|
||||
response: Final = session.send(request.prepare(), timeout=self._timeout)
|
||||
response.raise_for_status()
|
||||
return response.json()
|
||||
except requests.exceptions.HTTPError as e:
|
||||
|
|
@ -279,7 +281,7 @@ class KeysManagementClient:
|
|||
session: Final = requests.Session()
|
||||
response_text: str | None = None
|
||||
try:
|
||||
response: Final = session.send(request.prepare())
|
||||
response: Final = session.send(request.prepare(), timeout=self._timeout)
|
||||
response_text = response.text
|
||||
response.raise_for_status()
|
||||
return response.json()
|
||||
|
|
@ -309,7 +311,7 @@ class KeysManagementClient:
|
|||
|
||||
session: Final = requests.Session()
|
||||
try:
|
||||
response: Final = session.send(request.prepare())
|
||||
response: Final = session.send(request.prepare(), timeout=self._timeout)
|
||||
response.raise_for_status()
|
||||
return response.json()
|
||||
except requests.exceptions.HTTPError as e:
|
||||
|
|
|
|||
|
|
@ -6,16 +6,18 @@ from .exceptions import UnauthorizedError
|
|||
|
||||
|
||||
class ModelGroupsManagementClient:
|
||||
def __init__(self, base_url: str, api_key: str | None = None):
|
||||
def __init__(self, base_url: str, api_key: str | None = None, timeout: int = 30):
|
||||
"""
|
||||
Initialize the ModelGroupsManagementClient.
|
||||
|
||||
Args:
|
||||
base_url (str): The base URL of the LiteLLM proxy server (e.g., "http://localhost:8000")
|
||||
api_key (Optional[str]): API key for authentication. If provided, it will be sent as a Bearer token.
|
||||
timeout (int): Request timeout in seconds (default: 30)
|
||||
"""
|
||||
self._base_url = base_url.rstrip("/") # Remove trailing slash if present
|
||||
self._api_key = api_key
|
||||
self._timeout = timeout
|
||||
|
||||
def _get_headers(self) -> dict[str, str]:
|
||||
"""
|
||||
|
|
@ -53,7 +55,7 @@ class ModelGroupsManagementClient:
|
|||
# Prepare and send the request
|
||||
session: Final = requests.Session()
|
||||
try:
|
||||
response: Final = session.send(request.prepare())
|
||||
response: Final = session.send(request.prepare(), timeout=self._timeout)
|
||||
response.raise_for_status()
|
||||
return response.json()["data"]
|
||||
except requests.exceptions.HTTPError as e:
|
||||
|
|
|
|||
|
|
@ -7,16 +7,18 @@ from .exceptions import NotFoundError, UnauthorizedError
|
|||
|
||||
|
||||
class ModelsManagementClient:
|
||||
def __init__(self, base_url: str, api_key: str | None = None):
|
||||
def __init__(self, base_url: str, api_key: str | None = None, timeout: int = 30):
|
||||
"""
|
||||
Initialize the ModelsManagementClient.
|
||||
|
||||
Args:
|
||||
base_url (str): The base URL of the LiteLLM proxy server (e.g., "http://localhost:8000")
|
||||
api_key (Optional[str]): API key for authentication. If provided, it will be sent as a Bearer token.
|
||||
timeout (int): Request timeout in seconds (default: 30)
|
||||
"""
|
||||
self._base_url = base_url.rstrip("/") # Remove trailing slash if present
|
||||
self._api_key = api_key
|
||||
self._timeout = timeout
|
||||
|
||||
def _get_headers(self) -> dict[str, str]:
|
||||
"""
|
||||
|
|
@ -55,7 +57,7 @@ class ModelsManagementClient:
|
|||
# Prepare and send the request
|
||||
session: Final = requests.Session()
|
||||
try:
|
||||
response: Final = session.send(request.prepare())
|
||||
response: Final = session.send(request.prepare(), timeout=self._timeout)
|
||||
response.raise_for_status()
|
||||
return response.json()["data"]
|
||||
except requests.exceptions.HTTPError as e:
|
||||
|
|
@ -104,7 +106,7 @@ class ModelsManagementClient:
|
|||
# Prepare and send the request
|
||||
session: Final = requests.Session()
|
||||
try:
|
||||
response: Final = session.send(request.prepare())
|
||||
response: Final = session.send(request.prepare(), timeout=self._timeout)
|
||||
response.raise_for_status()
|
||||
return response.json()
|
||||
except requests.exceptions.HTTPError as e:
|
||||
|
|
@ -140,7 +142,7 @@ class ModelsManagementClient:
|
|||
# Prepare and send the request
|
||||
session: Final = requests.Session()
|
||||
try:
|
||||
response: Final = session.send(request.prepare())
|
||||
response: Final = session.send(request.prepare(), timeout=self._timeout)
|
||||
response.raise_for_status()
|
||||
return response.json()
|
||||
except requests.exceptions.HTTPError as e:
|
||||
|
|
@ -232,7 +234,7 @@ class ModelsManagementClient:
|
|||
# Prepare and send the request
|
||||
session: Final = requests.Session()
|
||||
try:
|
||||
response: Final = session.send(request.prepare())
|
||||
response: Final = session.send(request.prepare(), timeout=self._timeout)
|
||||
response.raise_for_status()
|
||||
return response.json()["data"]
|
||||
except requests.exceptions.HTTPError as e:
|
||||
|
|
@ -282,7 +284,7 @@ class ModelsManagementClient:
|
|||
# Prepare and send the request
|
||||
session: Final = requests.Session()
|
||||
try:
|
||||
response: Final = session.send(request.prepare())
|
||||
response: Final = session.send(request.prepare(), timeout=self._timeout)
|
||||
response.raise_for_status()
|
||||
return response.json()
|
||||
except requests.exceptions.HTTPError as e:
|
||||
|
|
|
|||
|
|
@ -11,16 +11,18 @@ from .exceptions import UnauthorizedError
|
|||
class TeamsManagementClient:
|
||||
"""Client for managing teams in LiteLLM proxy."""
|
||||
|
||||
def __init__(self, base_url: str, api_key: str | None = None):
|
||||
def __init__(self, base_url: str, api_key: str | None = None, timeout: int = 30):
|
||||
"""
|
||||
Initialize the TeamsManagementClient.
|
||||
|
||||
Args:
|
||||
base_url (str): The base URL of the LiteLLM proxy server (e.g., "http://localhost:4000")
|
||||
api_key (Optional[str]): API key for authentication. If provided, it will be sent as a Bearer token.
|
||||
timeout (int): Request timeout in seconds (default: 30)
|
||||
"""
|
||||
self._base_url = base_url.rstrip("/") # Remove trailing slash if present
|
||||
self._api_key = api_key
|
||||
self._timeout = timeout
|
||||
|
||||
def _get_headers(self) -> dict[str, str]:
|
||||
"""
|
||||
|
|
@ -60,7 +62,7 @@ class TeamsManagementClient:
|
|||
if organization_id:
|
||||
params["organization_id"] = organization_id
|
||||
|
||||
response: Final = requests.get(url, headers=self._get_headers(), params=params)
|
||||
response: Final = requests.get(url, headers=self._get_headers(), params=params, timeout=self._timeout)
|
||||
|
||||
if response.status_code == 401:
|
||||
raise UnauthorizedError("Authentication failed. Check your API key.")
|
||||
|
|
@ -117,7 +119,7 @@ class TeamsManagementClient:
|
|||
if sort_by:
|
||||
params["sort_by"] = sort_by
|
||||
|
||||
response: Final = requests.get(url, headers=self._get_headers(), params=params)
|
||||
response: Final = requests.get(url, headers=self._get_headers(), params=params, timeout=self._timeout)
|
||||
|
||||
if response.status_code == 401:
|
||||
raise UnauthorizedError("Authentication failed. Check your API key.")
|
||||
|
|
@ -138,7 +140,7 @@ class TeamsManagementClient:
|
|||
"""
|
||||
url: Final = f"{self._base_url}/team/available"
|
||||
|
||||
response: Final = requests.get(url, headers=self._get_headers())
|
||||
response: Final = requests.get(url, headers=self._get_headers(), timeout=self._timeout)
|
||||
|
||||
if response.status_code == 401:
|
||||
raise UnauthorizedError("Authentication failed. Check your API key.")
|
||||
|
|
|
|||
|
|
@ -6,9 +6,10 @@ from .exceptions import NotFoundError, UnauthorizedError
|
|||
|
||||
|
||||
class UsersManagementClient:
|
||||
def __init__(self, base_url: str, api_key: str | None = None):
|
||||
def __init__(self, base_url: str, api_key: str | None = None, timeout: int = 30):
|
||||
self.base_url = base_url.rstrip("/")
|
||||
self.api_key = api_key
|
||||
self.timeout = timeout
|
||||
|
||||
def _get_headers(self) -> dict[str, str]:
|
||||
headers: Final = {"Content-Type": "application/json"}
|
||||
|
|
@ -19,7 +20,7 @@ class UsersManagementClient:
|
|||
def list_users(self, params: dict[str, Any] | None = None) -> list[dict[str, Any]]:
|
||||
"""List users (GET /user/list)"""
|
||||
url: Final = f"{self.base_url}/user/list"
|
||||
response: Final = requests.get(url, headers=self._get_headers(), params=params)
|
||||
response: Final = requests.get(url, headers=self._get_headers(), params=params, timeout=self.timeout)
|
||||
if response.status_code == 401:
|
||||
raise UnauthorizedError(response.text)
|
||||
response.raise_for_status()
|
||||
|
|
@ -29,7 +30,7 @@ class UsersManagementClient:
|
|||
"""Get user info (GET /user/info)"""
|
||||
url: Final = f"{self.base_url}/user/info"
|
||||
params: Final = {"user_id": user_id} if user_id else {}
|
||||
response: Final = requests.get(url, headers=self._get_headers(), params=params)
|
||||
response: Final = requests.get(url, headers=self._get_headers(), params=params, timeout=self.timeout)
|
||||
if response.status_code == 401:
|
||||
raise UnauthorizedError(response.text)
|
||||
if response.status_code == 404:
|
||||
|
|
@ -41,7 +42,7 @@ class UsersManagementClient:
|
|||
"""Get user info v2 - lightweight, returns only user object (GET /v2/user/info)"""
|
||||
url: Final = f"{self.base_url}/v2/user/info"
|
||||
params: Final = {"user_id": user_id} if user_id else {}
|
||||
response: Final = requests.get(url, headers=self._get_headers(), params=params)
|
||||
response: Final = requests.get(url, headers=self._get_headers(), params=params, timeout=self.timeout)
|
||||
if response.status_code == 401:
|
||||
raise UnauthorizedError(response.text)
|
||||
if response.status_code == 404:
|
||||
|
|
@ -52,7 +53,7 @@ class UsersManagementClient:
|
|||
def create_user(self, user_data: dict[str, Any]) -> dict[str, Any]:
|
||||
"""Create a new user (POST /user/new)"""
|
||||
url: Final = f"{self.base_url}/user/new"
|
||||
response: Final = requests.post(url, headers=self._get_headers(), json=user_data)
|
||||
response: Final = requests.post(url, headers=self._get_headers(), json=user_data, timeout=self.timeout)
|
||||
if response.status_code == 401:
|
||||
raise UnauthorizedError(response.text)
|
||||
response.raise_for_status()
|
||||
|
|
@ -61,7 +62,9 @@ class UsersManagementClient:
|
|||
def delete_user(self, user_ids: list[str]) -> dict[str, Any]:
|
||||
"""Delete users (POST /user/delete)"""
|
||||
url: Final = f"{self.base_url}/user/delete"
|
||||
response: Final = requests.post(url, headers=self._get_headers(), json={"user_ids": user_ids})
|
||||
response: Final = requests.post(
|
||||
url, headers=self._get_headers(), json={"user_ids": user_ids}, timeout=self.timeout
|
||||
)
|
||||
if response.status_code == 401:
|
||||
raise UnauthorizedError(response.text)
|
||||
response.raise_for_status()
|
||||
|
|
|
|||
|
|
@ -37,6 +37,9 @@ if TYPE_CHECKING:
|
|||
from litellm.types.proxy.guardrails.guardrail_hooks.base import GuardrailConfigModel
|
||||
|
||||
|
||||
_AUTH_TIMEOUT_SECONDS: Final[float] = 30.0
|
||||
|
||||
|
||||
class _HiddenlayerEvaluation(TypedDict, total=False):
|
||||
action: str
|
||||
threat_level: str
|
||||
|
|
@ -157,10 +160,10 @@ def is_saas(host: str) -> bool:
|
|||
return False
|
||||
|
||||
|
||||
def _get_jwt(auth_url, api_id, api_key) -> str:
|
||||
def _get_jwt(auth_url, api_id, api_key, timeout: float = _AUTH_TIMEOUT_SECONDS) -> str:
|
||||
token_url: Final = f"{auth_url}/oauth2/token?grant_type=client_credentials"
|
||||
|
||||
resp: Final = requests.post(token_url, auth=HTTPBasicAuth(api_id, api_key))
|
||||
resp: Final = requests.post(token_url, auth=HTTPBasicAuth(api_id, api_key), timeout=timeout)
|
||||
|
||||
if not resp.ok:
|
||||
raise RuntimeError(
|
||||
|
|
|
|||
|
|
@ -6,7 +6,8 @@ lint.extend-select = [
|
|||
"T20", "PGH004", "RUF008", "RUF009", "RUF100",
|
||||
"B033", "FURB136", "FURB168", "FURB188", "I001", "PERF402", "PIE790", "PIE800", "PLC0208",
|
||||
"PLR0402", "PLR1711", "PLR1730", "PLR2044", "PLW0133", "PYI030", "PYI041", "PYI064", "RET501",
|
||||
"RUF010", "RUF022", "RUF023", "RUF051", "SIM114", "SIM118", "TC005", "UP006", "UP007", "UP008",
|
||||
"RUF010", "RUF022", "RUF023", "RUF051", "S113", "SIM114", "SIM118", "TC005", "UP006", "UP007",
|
||||
"UP008",
|
||||
"UP012", "UP018", "UP024", "UP032", "UP034", "UP035", "UP037", "UP045",
|
||||
]
|
||||
# RUF100 (unused-noqa) only knows the rules enabled in THIS config, so it would strip
|
||||
|
|
|
|||
|
|
@ -82,7 +82,7 @@ async def test_poll_for_ready_404(sleep_mock, request_mock):
|
|||
_poll_for_ready_data(
|
||||
"https://litellm.com", poll_interval=1, total_timeout=1, request_timeout=42
|
||||
)
|
||||
request_mock.assert_called_once_with("https://litellm.com", timeout=42)
|
||||
request_mock.assert_called_once_with("https://litellm.com", headers=None, timeout=42)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -103,7 +103,7 @@ async def test_poll_for_ready_200_ready(sleep_mock, click_mock, request_mock):
|
|||
)
|
||||
assert actual == {"status": "ready", "json": "data"}
|
||||
click_mock.assert_not_called()
|
||||
request_mock.assert_called_once_with("https://litellm.com", timeout=42)
|
||||
request_mock.assert_called_once_with("https://litellm.com", headers=None, timeout=42)
|
||||
sleep_mock.assert_not_called()
|
||||
|
||||
|
||||
|
|
@ -131,8 +131,8 @@ async def test_poll_for_ready_single_pending(sleep_mock, click_mock, request_moc
|
|||
click_mock.assert_not_called()
|
||||
request_mock.assert_has_calls(
|
||||
[
|
||||
call("https://litellm.com", timeout=42),
|
||||
call("https://litellm.com", timeout=42),
|
||||
call("https://litellm.com", headers=None, timeout=42),
|
||||
call("https://litellm.com", headers=None, timeout=42),
|
||||
]
|
||||
)
|
||||
sleep_mock.assert_called_once_with(1)
|
||||
|
|
@ -168,8 +168,8 @@ async def test_poll_for_ready_pending(sleep_mock, click_mock, request_mock):
|
|||
click_mock.assert_has_calls([call("Pending message"), call("Pending message")])
|
||||
request_mock.assert_has_calls(
|
||||
[
|
||||
call("https://litellm.com", timeout=42),
|
||||
call("https://litellm.com", timeout=42),
|
||||
call("https://litellm.com", headers=None, timeout=42),
|
||||
call("https://litellm.com", headers=None, timeout=42),
|
||||
]
|
||||
)
|
||||
sleep_mock.assert_has_calls([call(1), call(1)])
|
||||
|
|
@ -194,7 +194,7 @@ async def test_poll_for_ready_connection_failure(sleep_mock, click_mock, request
|
|||
click_mock.assert_called_once_with("Connection error (will retry): ERROR")
|
||||
request_mock.assert_has_calls(
|
||||
[
|
||||
call("https://litellm.com", timeout=42),
|
||||
call("https://litellm.com", headers=None, timeout=42),
|
||||
]
|
||||
)
|
||||
sleep_mock.assert_has_calls([call(1), call(1)])
|
||||
|
|
|
|||
38
tests/test_litellm/proxy/client/conftest.py
Normal file
38
tests/test_litellm/proxy/client/conftest.py
Normal file
|
|
@ -0,0 +1,38 @@
|
|||
import threading
|
||||
|
||||
import pytest
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def hanging_server():
|
||||
"""A server that accepts the connection and never answers, so only a timeout ends the call."""
|
||||
from http.server import BaseHTTPRequestHandler, HTTPServer
|
||||
from socketserver import ThreadingMixIn
|
||||
|
||||
stop: threading.Event = threading.Event()
|
||||
|
||||
class SilentRequestHandler(BaseHTTPRequestHandler):
|
||||
protocol_version = "HTTP/1.1"
|
||||
|
||||
def _hang(self):
|
||||
stop.wait(timeout=30)
|
||||
|
||||
do_GET = _hang
|
||||
do_POST = _hang
|
||||
|
||||
def log_message(self, format, *args):
|
||||
pass
|
||||
|
||||
class ThreadedServer(ThreadingMixIn, HTTPServer):
|
||||
daemon_threads = True
|
||||
|
||||
server = ThreadedServer(("127.0.0.1", 0), SilentRequestHandler)
|
||||
thread = threading.Thread(target=server.serve_forever, daemon=True)
|
||||
thread.start()
|
||||
try:
|
||||
yield f"http://127.0.0.1:{server.server_port}"
|
||||
finally:
|
||||
stop.set()
|
||||
server.shutdown()
|
||||
server.server_close()
|
||||
thread.join(timeout=5)
|
||||
|
|
@ -1,6 +1,7 @@
|
|||
import importlib
|
||||
import importlib.util
|
||||
from importlib.machinery import PathFinder
|
||||
import time
|
||||
import site
|
||||
import sys
|
||||
|
||||
|
|
@ -227,3 +228,31 @@ def test_completions_other_errors(client, sample_messages):
|
|||
with pytest.raises(requests.exceptions.HTTPError) as exc_info:
|
||||
client.completions(model="gpt-4", messages=sample_messages)
|
||||
assert exc_info.value.response.status_code == 500
|
||||
|
||||
|
||||
def test_completions_gives_up_at_the_timeout_instead_of_hanging(hanging_server):
|
||||
"""
|
||||
A proxy that accepts the connection but never answers used to pin the caller's
|
||||
process forever, since the request carried no timeout at all.
|
||||
"""
|
||||
client = ChatClient(base_url=hanging_server, api_key="sk-test", timeout=1)
|
||||
|
||||
started = time.monotonic()
|
||||
with pytest.raises(requests.exceptions.Timeout):
|
||||
client.completions(model="gpt-5.4", messages=[{"role": "user", "content": "hi"}])
|
||||
|
||||
assert time.monotonic() - started < 10
|
||||
|
||||
|
||||
def test_completions_stream_gives_up_at_the_timeout_instead_of_hanging(hanging_server):
|
||||
"""
|
||||
The streaming call opens the response before reading chunks, so a proxy that never
|
||||
sends its headers used to hang here forever too.
|
||||
"""
|
||||
client = ChatClient(base_url=hanging_server, api_key="sk-test", timeout=1)
|
||||
|
||||
started = time.monotonic()
|
||||
with pytest.raises(requests.exceptions.Timeout):
|
||||
next(client.completions_stream(model="gpt-5.4", messages=[{"role": "user", "content": "hi"}]))
|
||||
|
||||
assert time.monotonic() - started < 10
|
||||
|
|
|
|||
|
|
@ -82,6 +82,12 @@ def test_client_initialization():
|
|||
assert client.http._base_url == "http://localhost:4000"
|
||||
assert client.http._api_key == "test-key"
|
||||
assert client.http._timeout == 60
|
||||
assert client.teams._timeout == 60
|
||||
assert client.keys._timeout == 60
|
||||
assert client.credentials._timeout == 60
|
||||
assert client.models._timeout == 60
|
||||
assert client.model_groups._timeout == 60
|
||||
assert client.chat._timeout == 600
|
||||
|
||||
|
||||
def test_client_default_timeout():
|
||||
|
|
@ -92,6 +98,8 @@ def test_client_default_timeout():
|
|||
)
|
||||
|
||||
assert client.http._timeout == 30
|
||||
assert client.keys._timeout == 30
|
||||
assert client.chat._timeout == 600
|
||||
|
||||
|
||||
def test_client_without_api_key():
|
||||
|
|
|
|||
|
|
@ -1,4 +1,5 @@
|
|||
|
||||
import time
|
||||
import pytest
|
||||
import requests
|
||||
|
||||
|
|
@ -276,3 +277,17 @@ def test_encrypt_credential_values_does_not_mutate_original(monkeypatch):
|
|||
assert encrypted.credential_values["api_key"] != "sk-123"
|
||||
assert credential.credential_values["api_key"] == "sk-123"
|
||||
assert encrypted.credential_name == credential.credential_name
|
||||
|
||||
|
||||
def test_list_gives_up_at_the_timeout_instead_of_hanging(hanging_server):
|
||||
"""
|
||||
A proxy that accepts the connection but never answers used to pin the caller's
|
||||
process forever, since the request carried no timeout at all.
|
||||
"""
|
||||
client = CredentialsManagementClient(base_url=hanging_server, api_key="sk-test", timeout=1)
|
||||
|
||||
started = time.monotonic()
|
||||
with pytest.raises(requests.exceptions.Timeout):
|
||||
client.list()
|
||||
|
||||
assert time.monotonic() - started < 10
|
||||
|
|
|
|||
|
|
@ -1,3 +1,4 @@
|
|||
import time
|
||||
import traceback
|
||||
|
||||
import pytest
|
||||
|
|
@ -509,3 +510,17 @@ def test_not_found_error_redacts_wrapped_key():
|
|||
assert "REDACTED" in str(wrapped)
|
||||
assert LEAKY_KEY not in str(wrapped.orig_exception)
|
||||
assert wrapped.orig_exception.response.status_code == 404
|
||||
|
||||
|
||||
def test_list_gives_up_at_the_timeout_instead_of_hanging(hanging_server):
|
||||
"""
|
||||
A proxy that accepts the connection but never answers used to pin the caller's
|
||||
process forever, since the request carried no timeout at all.
|
||||
"""
|
||||
client = KeysManagementClient(base_url=hanging_server, api_key="sk-test", timeout=1)
|
||||
|
||||
started = time.monotonic()
|
||||
with pytest.raises(requests.exceptions.Timeout):
|
||||
client.list()
|
||||
|
||||
assert time.monotonic() - started < 10
|
||||
|
|
|
|||
|
|
@ -1,4 +1,5 @@
|
|||
|
||||
import time
|
||||
import pytest
|
||||
import requests
|
||||
|
||||
|
|
@ -172,3 +173,17 @@ def test_client_initialization_without_api_key(base_url):
|
|||
|
||||
assert client._api_key is None
|
||||
assert client.model_groups._api_key is None
|
||||
|
||||
|
||||
def test_info_gives_up_at_the_timeout_instead_of_hanging(hanging_server):
|
||||
"""
|
||||
A proxy that accepts the connection but never answers used to pin the caller's
|
||||
process forever, since the request carried no timeout at all.
|
||||
"""
|
||||
client = ModelGroupsManagementClient(base_url=hanging_server, api_key="sk-test", timeout=1)
|
||||
|
||||
started = time.monotonic()
|
||||
with pytest.raises(requests.exceptions.Timeout):
|
||||
client.info()
|
||||
|
||||
assert time.monotonic() - started < 10
|
||||
|
|
|
|||
|
|
@ -1,4 +1,5 @@
|
|||
|
||||
import time
|
||||
import pytest
|
||||
import requests
|
||||
|
||||
|
|
@ -732,3 +733,17 @@ def test_update_other_errors(client):
|
|||
with pytest.raises(requests.exceptions.HTTPError) as exc_info:
|
||||
client.update(model_id=model_id, model_params=model_params)
|
||||
assert exc_info.value.response.status_code == 500
|
||||
|
||||
|
||||
def test_list_gives_up_at_the_timeout_instead_of_hanging(hanging_server):
|
||||
"""
|
||||
A proxy that accepts the connection but never answers used to pin the caller's
|
||||
process forever, since the request carried no timeout at all.
|
||||
"""
|
||||
client = ModelsManagementClient(base_url=hanging_server, api_key="sk-test", timeout=1)
|
||||
|
||||
started = time.monotonic()
|
||||
with pytest.raises(requests.exceptions.Timeout):
|
||||
client.list()
|
||||
|
||||
assert time.monotonic() - started < 10
|
||||
|
|
|
|||
20
tests/test_litellm/proxy/client/test_teams.py
Normal file
20
tests/test_litellm/proxy/client/test_teams.py
Normal file
|
|
@ -0,0 +1,20 @@
|
|||
import time
|
||||
|
||||
import pytest
|
||||
import requests
|
||||
|
||||
from litellm.proxy.client.teams import TeamsManagementClient
|
||||
|
||||
|
||||
def test_list_gives_up_at_the_timeout_instead_of_hanging(hanging_server):
|
||||
"""
|
||||
A proxy that accepts the connection but never answers used to pin the caller's
|
||||
process forever, since the request carried no timeout at all.
|
||||
"""
|
||||
client = TeamsManagementClient(base_url=hanging_server, api_key="sk-test", timeout=1)
|
||||
|
||||
started = time.monotonic()
|
||||
with pytest.raises(requests.exceptions.Timeout):
|
||||
client.list()
|
||||
|
||||
assert time.monotonic() - started < 10
|
||||
|
|
@ -1,6 +1,8 @@
|
|||
import time
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
import requests
|
||||
|
||||
|
||||
|
||||
|
|
@ -82,3 +84,17 @@ def test_delete_user_unauthorized(mock_post, client):
|
|||
mock_post.return_value.text = "unauthorized"
|
||||
with pytest.raises(UnauthorizedError):
|
||||
client.delete_user(["u1"])
|
||||
|
||||
|
||||
def test_delete_user_gives_up_at_the_timeout_instead_of_hanging(hanging_server):
|
||||
"""
|
||||
A proxy that accepts the connection but never answers used to pin the caller's
|
||||
process forever, since the request carried no timeout at all.
|
||||
"""
|
||||
client = UsersManagementClient(base_url=hanging_server, api_key="sk-test", timeout=1)
|
||||
|
||||
started = time.monotonic()
|
||||
with pytest.raises(requests.exceptions.Timeout):
|
||||
client.delete_user(["u1"])
|
||||
|
||||
assert time.monotonic() - started < 10
|
||||
|
|
|
|||
|
|
@ -1,4 +1,6 @@
|
|||
import os
|
||||
import threading
|
||||
import time
|
||||
import uuid
|
||||
from typing import List, cast
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
|
@ -6,6 +8,7 @@ from unittest.mock import AsyncMock, MagicMock, patch
|
|||
import pytest
|
||||
from fastapi import HTTPException
|
||||
from httpx import Request, Response
|
||||
import requests
|
||||
|
||||
|
||||
import litellm
|
||||
|
|
@ -14,6 +17,7 @@ from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging
|
|||
from litellm.proxy.guardrails.guardrail_hooks.hiddenlayer.hiddenlayer import (
|
||||
HiddenlayerGuardrail,
|
||||
HiddenlayerGuardrailV2,
|
||||
_get_jwt,
|
||||
)
|
||||
from litellm.proxy.guardrails.init_guardrails import init_guardrails_v2
|
||||
from litellm.types.utils import (
|
||||
|
|
@ -1088,3 +1092,47 @@ class TestHiddenlayerGuardrailV2:
|
|||
config_model = HiddenlayerGuardrailV2.get_config_model()
|
||||
assert config_model is not None
|
||||
assert config_model.__name__ == "HiddenlayerGuardrailConfigModel"
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def hanging_auth_server():
|
||||
"""A server that accepts the connection and never answers, so only a timeout ends the call."""
|
||||
from http.server import BaseHTTPRequestHandler, HTTPServer
|
||||
from socketserver import ThreadingMixIn
|
||||
|
||||
stop: threading.Event = threading.Event()
|
||||
|
||||
class SilentRequestHandler(BaseHTTPRequestHandler):
|
||||
protocol_version = "HTTP/1.1"
|
||||
|
||||
def do_POST(self):
|
||||
stop.wait(timeout=30)
|
||||
|
||||
def log_message(self, format, *args):
|
||||
pass
|
||||
|
||||
class ThreadedServer(ThreadingMixIn, HTTPServer):
|
||||
daemon_threads = True
|
||||
|
||||
server = ThreadedServer(("127.0.0.1", 0), SilentRequestHandler)
|
||||
thread = threading.Thread(target=server.serve_forever, daemon=True)
|
||||
thread.start()
|
||||
try:
|
||||
yield f"http://127.0.0.1:{server.server_port}"
|
||||
finally:
|
||||
stop.set()
|
||||
server.shutdown()
|
||||
server.server_close()
|
||||
thread.join(timeout=5)
|
||||
|
||||
|
||||
def test_get_jwt_gives_up_at_the_timeout_instead_of_blocking_the_event_loop(hanging_auth_server):
|
||||
"""
|
||||
`_get_jwt` runs synchronously inside `_call_hiddenlayer`, so an auth host that
|
||||
accepts and never answers used to park the whole worker's event loop.
|
||||
"""
|
||||
started = time.monotonic()
|
||||
with pytest.raises(requests.exceptions.Timeout):
|
||||
_get_jwt(auth_url=hanging_auth_server, api_id="id", api_key="secret", timeout=1)
|
||||
|
||||
assert time.monotonic() - started < 10
|
||||
|
|
|
|||
|
|
@ -1,6 +1,6 @@
|
|||
{
|
||||
"LIT001": {
|
||||
"limit": 22705
|
||||
"limit": 22704
|
||||
},
|
||||
"LIT002": {
|
||||
"limit": 26854
|
||||
|
|
@ -33,6 +33,6 @@
|
|||
"limit": 5577
|
||||
},
|
||||
"LIT012": {
|
||||
"limit": 4508
|
||||
"limit": 4506
|
||||
}
|
||||
}
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue