mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-07 08:26:10 +00:00
fix(router): refresh streaming fallback headers and add spend logger proof
This commit is contained in:
parent
8d8266d01d
commit
eda63531cf
9 changed files with 438 additions and 39 deletions
|
|
@ -57,7 +57,7 @@
|
|||
"limit": 5601
|
||||
},
|
||||
"reportMissingTypeArgument": {
|
||||
"limit": 15306
|
||||
"limit": 15304
|
||||
},
|
||||
"reportMissingTypeStubs": {
|
||||
"limit": 40
|
||||
|
|
@ -99,7 +99,7 @@
|
|||
"limit": 0
|
||||
},
|
||||
"reportUnknownArgumentType": {
|
||||
"limit": 44364
|
||||
"limit": 44353
|
||||
},
|
||||
"reportUnknownLambdaType": {
|
||||
"limit": 109
|
||||
|
|
@ -108,10 +108,10 @@
|
|||
"limit": 38350
|
||||
},
|
||||
"reportUnknownParameterType": {
|
||||
"limit": 19626
|
||||
"limit": 19624
|
||||
},
|
||||
"reportUnknownVariableType": {
|
||||
"limit": 29890
|
||||
"limit": 29883
|
||||
},
|
||||
"reportUnnecessaryCast": {
|
||||
"limit": 111
|
||||
|
|
|
|||
55
cookbook/logging_observability/response_header_logger.md
Normal file
55
cookbook/logging_observability/response_header_logger.md
Normal file
|
|
@ -0,0 +1,55 @@
|
|||
# Store upstream response headers in spend logs
|
||||
|
||||
Copy `response_header_logger.py` next to your proxy configuration as `custom_callbacks.py`, then register the logger:
|
||||
|
||||
```yaml
|
||||
model_list:
|
||||
- model_name: openrouter-header-test
|
||||
litellm_params:
|
||||
model: openrouter/openai/gpt-5.6-luna
|
||||
api_key: os.environ/OPENROUTER_API_KEY
|
||||
litellm_settings:
|
||||
callbacks: [custom_callbacks.proxy_handler_instance]
|
||||
general_settings:
|
||||
master_key: os.environ/LITELLM_MASTER_KEY
|
||||
database_url: os.environ/DATABASE_URL
|
||||
proxy_batch_write_at: 1
|
||||
```
|
||||
|
||||
Start the proxy with this repository's code and your configuration:
|
||||
|
||||
```sh
|
||||
uv run --no-sync litellm --config config.yaml --port 4013
|
||||
```
|
||||
|
||||
The logger captures response headers before stream iteration, enriches successful spend logs through `async_logging_hook`, and recovers failure headers from the exception chain. A failure without recoverable headers retains the earlier capture. When a fallback response is captured, its headers replace the failed attempt’s headers. Request state stays in request metadata, so concurrent calls do not share headers on the logger instance
|
||||
|
||||
Send a request with existing custom metadata:
|
||||
|
||||
```sh
|
||||
curl --fail-with-body -sS -N http://localhost:4013/v1/chat/completions \
|
||||
-H "Authorization: Bearer $LITELLM_MASTER_KEY" \
|
||||
-H 'Content-Type: application/json' \
|
||||
-d '{"model":"openrouter-header-test","messages":[{"role":"user","content":"Reply exactly OK"}],"stream":true,"max_tokens":32,"metadata":{"spend_logs_metadata":{"verification_case":"stream","existing_field":"keep-me"}}}'
|
||||
```
|
||||
|
||||
After the background write, read the stored headers:
|
||||
|
||||
```sh
|
||||
curl --fail-with-body -sS http://localhost:4013/spend/logs \
|
||||
-H "Authorization: Bearer $LITELLM_MASTER_KEY" |
|
||||
jq '.[] | . as $row |
|
||||
(.metadata | if type == "string" then fromjson else . end).spend_logs_metadata as $m |
|
||||
{request_id: $row.request_id, case: $m.verification_case,
|
||||
existing_field: $m.existing_field, headers: $m.upstream_response_headers}'
|
||||
```
|
||||
|
||||
Headers appear in `metadata.spend_logs_metadata.upstream_response_headers`. Provider names receive the `llm_provider-` prefix; existing custom metadata and LiteLLM's internal-header protections are preserved
|
||||
|
||||
Set `stream` to false for the non-streaming control. To reproduce an upstream routing failure, add `"num_retries":0,"provider":{"only":["header-verification-nonexistent-provider"],"allow_fallbacks":false}` to the request. This returned HTTP 404 in verification
|
||||
|
||||
The logger was verified with OpenRouter chat completions, including simultaneous streams, a non-streaming control, and an upstream 404. The unit tests also cover a headerless error after header capture and both metadata keys used by the logging pipeline. Live commands, spend-log/database readbacks, and the exact tested revisions are in [PR #39376](https://github.com/BerriAI/litellm/pull/39376)
|
||||
|
||||
An earlier logger without early capture lost headers on a live streaming 429. That 429 did not recur after the logger update; retention for that error shape is covered by a unit test, not a second live 429 reproduction
|
||||
|
||||
Unmodified v1.98.0 still needs the shared HTTP-handler propagation fix as well as the router fix. The logger alone cannot recover headers dropped before the callback. These are headers returned by OpenRouter, which may differ from those of the provider behind it. Other providers, every retry/fallback attempt, and every failure type have not been live-tested
|
||||
86
cookbook/logging_observability/response_header_logger.py
Normal file
86
cookbook/logging_observability/response_header_logger.py
Normal file
|
|
@ -0,0 +1,86 @@
|
|||
from collections.abc import Mapping
|
||||
from typing import Final, cast
|
||||
|
||||
import httpx
|
||||
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
|
||||
|
||||
def _as_dict(value: object) -> dict[str, object]:
|
||||
return cast(dict[str, object], value) if isinstance(value, dict) else {}
|
||||
|
||||
|
||||
def _child(parent: dict[str, object], key: str) -> dict[str, object]:
|
||||
child: Final = _as_dict(parent.get(key))
|
||||
parent[key] = child # rebind-ok: callback metadata must be updated in place
|
||||
return child
|
||||
|
||||
|
||||
def _headers(value: object) -> dict[str, str]:
|
||||
if isinstance(value, httpx.Headers):
|
||||
return dict(value)
|
||||
return {name: header for name, header in _as_dict(value).items() if isinstance(header, str)}
|
||||
|
||||
|
||||
def get_error_headers(error: BaseException | None, seen: tuple[int, ...] = ()) -> dict[str, str]:
|
||||
if error is None or id(error) in seen:
|
||||
return {}
|
||||
response: Final[object] = getattr(error, "response", None)
|
||||
headers: Final = _headers(getattr(response, "headers", None)) or _headers(getattr(error, "headers", None))
|
||||
if headers:
|
||||
return headers
|
||||
visited: Final = (*seen, id(error))
|
||||
return get_error_headers(error.__cause__, visited) or get_error_headers(error.__context__, visited)
|
||||
|
||||
|
||||
def save_headers(metadata: dict[str, object], headers: Mapping[str, object]) -> None:
|
||||
existing: Final = _as_dict(metadata.get("spend_logs_metadata"))
|
||||
metadata["spend_logs_metadata"] = { # rebind-ok: callback metadata must be updated in place
|
||||
**existing,
|
||||
"upstream_response_headers": {
|
||||
**_as_dict(existing.get("upstream_response_headers")),
|
||||
**headers,
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
class ResponseHeaderLogger(CustomLogger):
|
||||
async def async_post_call_response_headers_hook(
|
||||
self,
|
||||
data: dict[str, object],
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
response: object,
|
||||
request_headers: dict[str, str] | None = None,
|
||||
litellm_call_info: dict[str, object] | None = None,
|
||||
) -> None:
|
||||
headers: Final = _headers(getattr(response, "_response_headers", None))
|
||||
spend_metadata: Final = _child(_child(data, "metadata"), "spend_logs_metadata")
|
||||
spend_metadata["upstream_response_headers"] = {f"llm_provider-{name}": value for name, value in headers.items()}
|
||||
|
||||
async def async_logging_hook(
|
||||
self, kwargs: dict[str, object], result: object, call_type: str
|
||||
) -> tuple[dict[str, object], object]:
|
||||
payload: Final = _as_dict(kwargs.get("standard_logging_object"))
|
||||
hidden: Final = _as_dict(payload.get("hidden_params"))
|
||||
headers: Final = _as_dict(hidden.get("additional_headers"))
|
||||
params: Final = _child(kwargs, "litellm_params")
|
||||
key: Final = "litellm_metadata" if params.get("litellm_metadata") else "metadata"
|
||||
save_headers(_child(params, key), headers)
|
||||
return kwargs, result
|
||||
|
||||
async def async_post_call_failure_hook(
|
||||
self,
|
||||
request_data: dict[str, object],
|
||||
original_exception: Exception,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
traceback_str: str | None = None,
|
||||
) -> None:
|
||||
headers: Final = get_error_headers(original_exception)
|
||||
save_headers(
|
||||
_child(request_data, "metadata"),
|
||||
{f"llm_provider-{name}": value for name, value in headers.items()},
|
||||
)
|
||||
|
||||
|
||||
proxy_handler_instance: Final = ResponseHeaderLogger()
|
||||
|
|
@ -719,7 +719,7 @@ class _UpstreamClosingStreamingResponse(StreamingResponse):
|
|||
content: AsyncGenerator[str, None],
|
||||
*,
|
||||
media_type: str | None = None,
|
||||
headers: dict | None = None,
|
||||
headers: Mapping[str, str] | None = None,
|
||||
status_code: int = status.HTTP_200_OK,
|
||||
upstream_generator: AsyncGenerator[str, None] | None = None,
|
||||
) -> None:
|
||||
|
|
@ -857,22 +857,23 @@ def _sse_error_frames(error_obj: Mapping[str, object]) -> tuple[str, str]:
|
|||
async def create_response(
|
||||
generator: AsyncGenerator[str, None],
|
||||
media_type: str,
|
||||
headers: dict,
|
||||
headers: Mapping[str, str],
|
||||
default_status_code: int = status.HTTP_200_OK,
|
||||
request: Request | None = None,
|
||||
refresh_headers: Callable[[], Awaitable[Mapping[str, str]]] | None = None,
|
||||
) -> StreamingResponse | JSONResponse:
|
||||
"""
|
||||
Create streaming response, checking if the first chunk is an error.
|
||||
If the first chunk is an error, return a standard JSON error response.
|
||||
Otherwise, return StreamingResponse and stream all content.
|
||||
"""
|
||||
|
||||
# Tell buffering reverse proxies (nginx, ingress-nginx, Envoy) to flush SSE
|
||||
# immediately instead of releasing the whole stream in one batch (issue #28384).
|
||||
streaming_headers: Final = {
|
||||
**headers,
|
||||
"Cache-Control": "no-cache",
|
||||
"X-Accel-Buffering": "no",
|
||||
}
|
||||
async def streaming_headers() -> Mapping[str, str]:
|
||||
current_headers: Final = await refresh_headers() if refresh_headers is not None else headers
|
||||
return MappingProxyType({**current_headers, "Cache-Control": "no-cache", "X-Accel-Buffering": "no"})
|
||||
|
||||
first_chunk_value: str | None = None
|
||||
final_status_code = default_status_code
|
||||
|
||||
|
|
@ -909,7 +910,7 @@ async def create_response(
|
|||
return JSONResponse(
|
||||
status_code=final_status_code,
|
||||
content={"error": error_dict},
|
||||
headers=headers,
|
||||
headers=await refresh_headers() if refresh_headers is not None else headers,
|
||||
)
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.debug("Error parsing first chunk value: %s", e)
|
||||
|
|
@ -938,7 +939,7 @@ async def create_response(
|
|||
return StreamingResponse(
|
||||
empty_gen(),
|
||||
media_type=media_type,
|
||||
headers=streaming_headers,
|
||||
headers=await streaming_headers(),
|
||||
status_code=default_status_code,
|
||||
)
|
||||
except Exception as e:
|
||||
|
|
@ -954,7 +955,7 @@ async def create_response(
|
|||
return StreamingResponse(
|
||||
error_gen_message(),
|
||||
media_type=media_type,
|
||||
headers=streaming_headers,
|
||||
headers=await streaming_headers(),
|
||||
status_code=error_status,
|
||||
)
|
||||
|
||||
|
|
@ -976,7 +977,7 @@ async def create_response(
|
|||
return _UpstreamClosingStreamingResponse(
|
||||
combined_generator(),
|
||||
media_type=media_type,
|
||||
headers=streaming_headers,
|
||||
headers=await streaming_headers(),
|
||||
status_code=final_status_code,
|
||||
upstream_generator=generator,
|
||||
)
|
||||
|
|
@ -2379,28 +2380,46 @@ class ProxyBaseLLMRequestProcessing:
|
|||
if self._is_streaming_request(
|
||||
data=self.data, is_streaming_request=is_streaming_request
|
||||
) or self._is_streaming_response(response): # use generate_responses to stream responses
|
||||
custom_headers: Final = ProxyBaseLLMRequestProcessing.get_custom_headers(
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
call_id=logging_obj.litellm_call_id,
|
||||
model_id=model_id,
|
||||
cache_key=cache_key,
|
||||
api_base=api_base,
|
||||
version=version,
|
||||
response_cost=response_cost,
|
||||
model_region=getattr(user_api_key_dict, "allowed_model_region", ""),
|
||||
fastest_response_batch_completion=fastest_response_batch_completion,
|
||||
request_data=self.data,
|
||||
hidden_params=hidden_params,
|
||||
litellm_logging_obj=logging_obj,
|
||||
**additional_headers,
|
||||
)
|
||||
|
||||
def get_stream_headers() -> dict[str, str]:
|
||||
current_hidden: Final = get_hidden_params_dict(response)
|
||||
return ProxyBaseLLMRequestProcessing.get_custom_headers(
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
call_id=logging_obj.litellm_call_id,
|
||||
model_id=self._get_model_id_from_response(current_hidden, self.data),
|
||||
cache_key=current_hidden.get("cache_key") or "",
|
||||
api_base=current_hidden.get("api_base") or "",
|
||||
version=version,
|
||||
response_cost=current_hidden.get("response_cost") or "",
|
||||
model_region=getattr(user_api_key_dict, "allowed_model_region", ""),
|
||||
fastest_response_batch_completion=current_hidden.get("fastest_response_batch_completion"),
|
||||
request_data=self.data,
|
||||
hidden_params=current_hidden,
|
||||
litellm_logging_obj=logging_obj,
|
||||
**(current_hidden.get("additional_headers") or MappingProxyType({})),
|
||||
)
|
||||
|
||||
custom_headers: Final = get_stream_headers()
|
||||
request_headers: Final = dict(request.headers)
|
||||
|
||||
async def refresh_stream_headers() -> Mapping[str, str]:
|
||||
if get_hidden_params_dict(response) is hidden_params:
|
||||
return custom_headers
|
||||
refreshed: Final = get_stream_headers()
|
||||
updated_callback_headers: Final = await proxy_logging_obj.post_call_response_headers_hook(
|
||||
data=self.data,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
response=response,
|
||||
request_headers=request_headers,
|
||||
)
|
||||
return MappingProxyType({**refreshed, **(updated_callback_headers or MappingProxyType({}))})
|
||||
|
||||
# Call response headers hook for streaming success
|
||||
callback_headers = await proxy_logging_obj.post_call_response_headers_hook(
|
||||
data=self.data,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
response=response,
|
||||
request_headers=dict(request.headers),
|
||||
request_headers=request_headers,
|
||||
)
|
||||
if callback_headers:
|
||||
custom_headers.update(callback_headers)
|
||||
|
|
@ -2508,6 +2527,8 @@ class ProxyBaseLLMRequestProcessing:
|
|||
)
|
||||
# Non-streaming response - fall through to normal response handling
|
||||
elif select_data_generator:
|
||||
from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper
|
||||
|
||||
selected_data_generator = select_data_generator(
|
||||
response=response,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
|
|
@ -2536,6 +2557,7 @@ class ProxyBaseLLMRequestProcessing:
|
|||
media_type="text/event-stream",
|
||||
headers=custom_headers,
|
||||
request=request,
|
||||
refresh_headers=refresh_stream_headers if isinstance(response, CustomStreamWrapper) else None,
|
||||
)
|
||||
|
||||
### CALL HOOKS ### - modify outgoing data
|
||||
|
|
|
|||
|
|
@ -2499,6 +2499,10 @@ class Router:
|
|||
if hasattr(model_response, "_hidden_params"):
|
||||
self._hidden_params = model_response._hidden_params.copy()
|
||||
|
||||
def update_headers(self, source: object) -> None:
|
||||
self._response_headers = getattr(source, "_response_headers", None)
|
||||
self._hidden_params = get_hidden_params_dict(source).copy()
|
||||
|
||||
def __aiter__(self):
|
||||
return self
|
||||
|
||||
|
|
@ -2552,8 +2556,10 @@ class Router:
|
|||
|
||||
# If fallback returns a streaming response, iterate over it
|
||||
if hasattr(fallback_response, "__aiter__"):
|
||||
prepared_fallback_hidden_params = Router._prepare_fallback_hidden_params(fallback_response)
|
||||
wrapped_response.update_headers(fallback_response)
|
||||
async for fallback_item in fallback_response:
|
||||
wrapped_response.update_headers(fallback_response)
|
||||
prepared_fallback_hidden_params = Router._prepare_fallback_hidden_params(fallback_response)
|
||||
Router._apply_fallback_hidden_params_to_item(fallback_item, prepared_fallback_hidden_params)
|
||||
if (
|
||||
fallback_item
|
||||
|
|
@ -2601,7 +2607,8 @@ class Router:
|
|||
e,
|
||||
)
|
||||
|
||||
return FallbackStreamWrapper(stream_with_fallbacks())
|
||||
wrapped_response: Final = FallbackStreamWrapper(stream_with_fallbacks())
|
||||
return wrapped_response
|
||||
|
||||
@staticmethod
|
||||
def _extract_partial_responses_usage(
|
||||
|
|
@ -3043,6 +3050,10 @@ class Router:
|
|||
if hasattr(model_response, "_hidden_params"):
|
||||
self._hidden_params = model_response._hidden_params.copy()
|
||||
|
||||
def update_headers(self, source: object) -> None:
|
||||
self._response_headers = getattr(source, "_response_headers", None)
|
||||
self._hidden_params = get_hidden_params_dict(source).copy()
|
||||
|
||||
def __iter__(self):
|
||||
return self
|
||||
|
||||
|
|
@ -3093,8 +3104,10 @@ class Router:
|
|||
)
|
||||
|
||||
if hasattr(fallback_response, "__iter__"):
|
||||
prepared_fallback_hidden_params = Router._prepare_fallback_hidden_params(fallback_response)
|
||||
wrapped_response.update_headers(fallback_response)
|
||||
for fallback_item in fallback_response:
|
||||
wrapped_response.update_headers(fallback_response)
|
||||
prepared_fallback_hidden_params = Router._prepare_fallback_hidden_params(fallback_response)
|
||||
Router._apply_fallback_hidden_params_to_item(fallback_item, prepared_fallback_hidden_params)
|
||||
if (
|
||||
fallback_item
|
||||
|
|
@ -3132,7 +3145,8 @@ class Router:
|
|||
close_err,
|
||||
)
|
||||
|
||||
return SyncFallbackStreamWrapper(stream_with_fallbacks())
|
||||
wrapped_response: Final = SyncFallbackStreamWrapper(stream_with_fallbacks())
|
||||
return wrapped_response
|
||||
|
||||
async def _silent_experiment_acompletion(self, silent_model: str, messages: Sequence[Mapping[str, str]], **kwargs):
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -565,3 +565,96 @@ def test_redis_cache_completion_stream():
|
|||
|
||||
|
||||
# test_redis_cache_completion_stream()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_response_header_logger_preserves_headers_on_headerless_stream_failure():
|
||||
from types import SimpleNamespace
|
||||
from typing import Final
|
||||
|
||||
from cookbook.logging_observability.response_header_logger import ResponseHeaderLogger
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
|
||||
logger: Final = ResponseHeaderLogger()
|
||||
auth: Final = UserAPIKeyAuth()
|
||||
data: Final = {"metadata": {"spend_logs_metadata": {"existing_field": "keep"}}}
|
||||
response: Final = SimpleNamespace(
|
||||
_response_headers={
|
||||
"x-generation-id": "stream-request",
|
||||
"x-litellm-attempted-fallbacks": "spoofed",
|
||||
}
|
||||
)
|
||||
await logger.async_post_call_response_headers_hook(data, auth, response)
|
||||
await logger.async_post_call_failure_hook(data, RuntimeError("stream failed"), auth)
|
||||
|
||||
assert data["metadata"]["spend_logs_metadata"] == {
|
||||
"existing_field": "keep",
|
||||
"upstream_response_headers": {
|
||||
"llm_provider-x-generation-id": "stream-request",
|
||||
"llm_provider-x-litellm-attempted-fallbacks": "spoofed",
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("metadata_key", ["metadata", "litellm_metadata"])
|
||||
async def test_response_header_logger_success_preserves_capture_and_metadata(metadata_key: str):
|
||||
from types import SimpleNamespace
|
||||
from typing import Final
|
||||
|
||||
from cookbook.logging_observability.response_header_logger import ResponseHeaderLogger
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
|
||||
logger: Final = ResponseHeaderLogger()
|
||||
metadata: Final = {"spend_logs_metadata": {"existing_field": "keep"}}
|
||||
await logger.async_post_call_response_headers_hook(
|
||||
{"metadata": metadata},
|
||||
UserAPIKeyAuth(),
|
||||
SimpleNamespace(_response_headers={"x-generation-id": "stream-request"}),
|
||||
)
|
||||
kwargs: Final = {
|
||||
"litellm_params": {metadata_key: metadata},
|
||||
"standard_logging_object": {
|
||||
"hidden_params": {
|
||||
"additional_headers": {
|
||||
"llm_provider-x-litellm-response-cost": 0.01,
|
||||
}
|
||||
}
|
||||
},
|
||||
}
|
||||
await logger.async_logging_hook(kwargs, None, "acompletion")
|
||||
|
||||
assert metadata["spend_logs_metadata"] == {
|
||||
"existing_field": "keep",
|
||||
"upstream_response_headers": {
|
||||
"llm_provider-x-generation-id": "stream-request",
|
||||
"llm_provider-x-litellm-response-cost": 0.01,
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("fallback_headers", [None, {"x-request-id": "winner"}])
|
||||
async def test_response_header_logger_replaces_headers_after_fallback(fallback_headers: dict[str, str] | None):
|
||||
from types import SimpleNamespace
|
||||
from typing import Final
|
||||
|
||||
from cookbook.logging_observability.response_header_logger import ResponseHeaderLogger
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
|
||||
logger: Final = ResponseHeaderLogger()
|
||||
auth: Final = UserAPIKeyAuth()
|
||||
data: Final = {"metadata": {"spend_logs_metadata": {"existing_field": "keep"}}}
|
||||
await logger.async_post_call_response_headers_hook(
|
||||
data,
|
||||
auth,
|
||||
SimpleNamespace(_response_headers={"x-request-id": "failed", "x-failed-only": "remove"}),
|
||||
)
|
||||
await logger.async_post_call_response_headers_hook(data, auth, SimpleNamespace(_response_headers=fallback_headers))
|
||||
await logger.async_post_call_failure_hook(data, RuntimeError("stream failed"), auth)
|
||||
assert data["metadata"]["spend_logs_metadata"] == {
|
||||
"existing_field": "keep",
|
||||
"upstream_response_headers": {
|
||||
f"llm_provider-{name}": value for name, value in (fallback_headers or {}).items()
|
||||
},
|
||||
}
|
||||
|
|
|
|||
|
|
@ -7665,3 +7665,35 @@ def test_log_llm_api_exception_traceback_only_for_unexpected_errors(exc, expect_
|
|||
records = [r for r in caplog.records if "_handle_llm_api_exception(): Exception occured" in r.getMessage()]
|
||||
assert len(records) == 1
|
||||
assert (records[0].exc_info is not None) is expect_traceback
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("stream_kind", ["success", "empty", "error"])
|
||||
async def test_create_response_refreshes_headers_before_committing(stream_kind: str):
|
||||
from typing import Final
|
||||
|
||||
headers: Final = {"llm_provider-x-request-id": "failed", "llm_provider-x-failed-only": "remove"}
|
||||
|
||||
async def stream():
|
||||
if stream_kind == "error":
|
||||
yield 'data: {"error":{"message":"failed","code":429}}\n\n'
|
||||
elif stream_kind == "success":
|
||||
yield 'data: {"id":"winner","choices":[]}\n\n'
|
||||
|
||||
async def refresh_headers():
|
||||
return {"llm_provider-x-request-id": "winner", "x-litellm-model-group": "fallback"}
|
||||
|
||||
response: Final = await create_response(
|
||||
stream(),
|
||||
"text/event-stream",
|
||||
headers,
|
||||
refresh_headers=refresh_headers,
|
||||
)
|
||||
assert response.headers["llm_provider-x-request-id"] == "winner"
|
||||
assert response.headers["x-litellm-model-group"] == "fallback"
|
||||
assert "llm_provider-x-failed-only" not in response.headers
|
||||
if stream_kind == "error":
|
||||
assert response.status_code == 429
|
||||
else:
|
||||
assert response.headers["cache-control"] == "no-cache"
|
||||
assert response.headers["x-accel-buffering"] == "no"
|
||||
|
|
|
|||
|
|
@ -2346,7 +2346,7 @@ async def test_completion_streaming_iterator_preserves_response_headers(
|
|||
_response_headers=response_headers,
|
||||
)
|
||||
router: Final = Router(model_list=[])
|
||||
upstream._hidden_params["additional_headers"]["x-litellm-model-group"] = "real-group"
|
||||
await router.set_response_headers(response=upstream, model_group="real-group")
|
||||
iterator_kwargs: Final = dict(
|
||||
model_response=upstream,
|
||||
messages=[{"role": "user", "content": "Hello"}],
|
||||
|
|
@ -2373,6 +2373,103 @@ async def test_completion_streaming_iterator_preserves_response_headers(
|
|||
assert StandardLoggingPayloadSetup.get_hidden_params(complete._hidden_params)["additional_headers"] == expected
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("is_async", [False, True])
|
||||
@pytest.mark.parametrize("fallback_headers", [None, {"x-request-id": "winner", "x-litellm-model-group": "spoofed"}])
|
||||
async def test_completion_streaming_iterator_replaces_failed_provider_headers(
|
||||
is_async: bool, fallback_headers: dict[str, str] | None
|
||||
) -> None:
|
||||
from typing import Final
|
||||
from unittest.mock import Mock
|
||||
|
||||
from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper
|
||||
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
|
||||
|
||||
client: Final = AsyncMock(spec=AsyncHTTPHandler) if is_async else Mock(spec=HTTPHandler)
|
||||
client.post.return_value = httpx.Response(
|
||||
200,
|
||||
content=(
|
||||
b'data: {"id":"winner","object":"chat.completion.chunk","created":1,'
|
||||
b'"model":"test-model","choices":[{"index":0,"delta":{"content":"Hello"},'
|
||||
b'"finish_reason":null}]}\n\n'
|
||||
b"data: [DONE]\n\n"
|
||||
),
|
||||
request=httpx.Request("POST", "https://provider.test/v1/chat/completions"),
|
||||
)
|
||||
request: Final = dict(
|
||||
model="hosted_vllm/test-model",
|
||||
messages=[{"role": "user", "content": "Hello"}],
|
||||
api_base="https://provider.test/v1",
|
||||
api_key="fake-key",
|
||||
stream=True,
|
||||
client=client,
|
||||
)
|
||||
original: Final = await litellm.acompletion(**request) if is_async else litellm.completion(**request)
|
||||
fallback: Final = CustomStreamWrapper(
|
||||
completion_stream=original.completion_stream,
|
||||
model=original.model,
|
||||
custom_llm_provider=original.custom_llm_provider,
|
||||
logging_obj=original.logging_obj,
|
||||
_response_headers=fallback_headers,
|
||||
)
|
||||
|
||||
class FallbackRouter(Router):
|
||||
async def async_function_with_fallbacks_common_utils(self, **kwargs):
|
||||
return fallback
|
||||
|
||||
def function_with_fallbacks(self, **kwargs):
|
||||
return fallback
|
||||
|
||||
class FailedStream:
|
||||
model = "failed-model"
|
||||
custom_llm_provider = "hosted_vllm"
|
||||
logging_obj = original.logging_obj
|
||||
chunks = []
|
||||
_response_headers = {"x-request-id": "failed", "x-failed-only": "remove"}
|
||||
_hidden_params = {
|
||||
"additional_headers": {"llm_provider-x-request-id": "failed", "llm_provider-x-failed-only": "remove"}
|
||||
}
|
||||
|
||||
def __iter__(self):
|
||||
return self
|
||||
|
||||
def __next__(self):
|
||||
raise MidStreamFallbackError(
|
||||
message="limited",
|
||||
model=self.model,
|
||||
llm_provider=self.custom_llm_provider,
|
||||
is_pre_first_chunk=True,
|
||||
)
|
||||
|
||||
def __aiter__(self):
|
||||
return self
|
||||
|
||||
async def __anext__(self):
|
||||
return self.__next__()
|
||||
|
||||
router: Final = FallbackRouter(model_list=[])
|
||||
await router.set_response_headers(response=fallback, model_group="fallback")
|
||||
iterator_kwargs: Final = dict(
|
||||
model_response=FailedStream(),
|
||||
messages=[{"role": "user", "content": "Hello"}],
|
||||
initial_kwargs={"model": "primary", "stream": True},
|
||||
)
|
||||
wrapped: Final = (
|
||||
await router._acompletion_streaming_iterator(**iterator_kwargs)
|
||||
if is_async
|
||||
else router._completion_streaming_iterator(**iterator_kwargs)
|
||||
)
|
||||
chunks: Final = [chunk async for chunk in wrapped] if is_async else list(wrapped)
|
||||
expected: Final = {
|
||||
**{f"llm_provider-{name}": value for name, value in (fallback_headers or {}).items()},
|
||||
"x-litellm-model-group": "fallback",
|
||||
}
|
||||
assert wrapped._response_headers == fallback_headers
|
||||
assert wrapped._hidden_params["additional_headers"] == expected
|
||||
assert all(chunk._hidden_params["additional_headers"] == expected for chunk in chunks)
|
||||
assert litellm.stream_chunk_builder(chunks).choices[0].message.content == "Hello"
|
||||
|
||||
|
||||
def test_completion_streaming_iterator_fallback_on_429():
|
||||
"""Sync streaming: MidStreamFallbackError (429 pre-first-chunk) triggers fallback.
|
||||
|
||||
|
|
|
|||
|
|
@ -1,9 +1,9 @@
|
|||
{
|
||||
"LIT001": {
|
||||
"limit": 22364
|
||||
"limit": 22363
|
||||
},
|
||||
"LIT002": {
|
||||
"limit": 26777
|
||||
"limit": 26776
|
||||
},
|
||||
"LIT003": {
|
||||
"limit": 269
|
||||
|
|
@ -27,7 +27,7 @@
|
|||
"limit": 0
|
||||
},
|
||||
"LIT010": {
|
||||
"limit": 16507
|
||||
"limit": 16505
|
||||
},
|
||||
"LIT011": {
|
||||
"limit": 5535
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue