mirror of
https://github.com/BerriAI/litellm.git
synced 2025-04-26 03:04:13 +00:00
513b1904ab
All checks were successful
Read Version from pyproject.toml / read-version (push) Successful in 14s
* feat(router.py): add retry headers to response makes it easy to add testing to ensure model-specific retries are respected * fix(add_retry_headers.py): clarify attempted retries vs. max retries * test(test_fallbacks.py): add test for checking if max retries set for model is respected * test(test_fallbacks.py): assert values for attempted retries and max retries are as expected * fix(utils.py): return timeout in litellm proxy response headers * test(test_fallbacks.py): add test to assert model specific timeout used on timeout error * test: add bad model with timeout to proxy * fix: fix linting error * fix(router.py): fix get model list from model alias * test: loosen test restriction - account for other events on proxy
40 lines
1 KiB
Python
40 lines
1 KiB
Python
from typing import Any, Optional, Union
|
|
|
|
from pydantic import BaseModel
|
|
|
|
from litellm.types.utils import HiddenParams
|
|
|
|
|
|
def add_retry_headers_to_response(
|
|
response: Any,
|
|
attempted_retries: int,
|
|
max_retries: Optional[int] = None,
|
|
) -> Any:
|
|
"""
|
|
Add retry headers to the request
|
|
"""
|
|
|
|
if response is None or not isinstance(response, BaseModel):
|
|
return response
|
|
|
|
retry_headers = {
|
|
"x-litellm-attempted-retries": attempted_retries,
|
|
}
|
|
if max_retries is not None:
|
|
retry_headers["x-litellm-max-retries"] = max_retries
|
|
|
|
hidden_params: Optional[Union[dict, HiddenParams]] = getattr(
|
|
response, "_hidden_params", {}
|
|
)
|
|
|
|
if hidden_params is None:
|
|
hidden_params = {}
|
|
elif isinstance(hidden_params, HiddenParams):
|
|
hidden_params = hidden_params.model_dump()
|
|
|
|
hidden_params.setdefault("additional_headers", {})
|
|
hidden_params["additional_headers"].update(retry_headers)
|
|
|
|
setattr(response, "_hidden_params", hidden_params)
|
|
|
|
return response
|