File size: 1,039 Bytes
e3278e4 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 |
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
|