litellm-mirror/litellm/router_utils/add_retry_fallback_headers.py
Ishaan Jaff 6b3bfa2b42
(Feat) - return x-litellm-attempted-fallbacks in responses from litellm proxy (#8558)
* add_fallback_headers_to_response

* test x-litellm-attempted-fallbacks

* unit test attempted fallbacks

* fix add_fallback_headers_to_response

* docs document response headers

* fix file name
2025-02-15 14:54:23 -08:00

68 lines
1.8 KiB
Python

from typing import Any, Optional, Union
from pydantic import BaseModel
from litellm.types.utils import HiddenParams
def _add_headers_to_response(response: Any, headers: dict) -> Any:
"""
Helper function to add headers to a response's hidden params
"""
if response is None or not isinstance(response, BaseModel):
return response
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(headers)
setattr(response, "_hidden_params", hidden_params)
return response
def add_retry_headers_to_response(
response: Any,
attempted_retries: int,
max_retries: Optional[int] = None,
) -> Any:
"""
Add retry headers to the request
"""
retry_headers = {
"x-litellm-attempted-retries": attempted_retries,
}
if max_retries is not None:
retry_headers["x-litellm-max-retries"] = max_retries
return _add_headers_to_response(response, retry_headers)
def add_fallback_headers_to_response(
response: Any,
attempted_fallbacks: int,
) -> Any:
"""
Add fallback headers to the response
Args:
response: The response to add the headers to
attempted_fallbacks: The number of fallbacks attempted
Returns:
The response with the headers added
Note: It's intentional that we don't add max_fallbacks in response headers
Want to avoid bloat in the response headers for performance.
"""
fallback_headers = {
"x-litellm-attempted-fallbacks": attempted_fallbacks,
}
return _add_headers_to_response(response, fallback_headers)