From d7290c3a4f92f42190b3ef8609d35b3ba9d10269 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Wed, 2 Sep 2026 13:01:12 -0700 Subject: [PATCH] test: keep response header verification out of cookbook --- .../response_header_logger.md | 55 ------------ .../response_header_logger.py | 86 ------------------- tests/local_testing/test_custom_logger.py | 85 +++++++++++++++++- 3 files changed, 82 insertions(+), 144 deletions(-) delete mode 100644 cookbook/logging_observability/response_header_logger.md delete mode 100644 cookbook/logging_observability/response_header_logger.py diff --git a/cookbook/logging_observability/response_header_logger.md b/cookbook/logging_observability/response_header_logger.md deleted file mode 100644 index 013fdaf147a..00000000000 --- a/cookbook/logging_observability/response_header_logger.md +++ /dev/null @@ -1,55 +0,0 @@ -# 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 diff --git a/cookbook/logging_observability/response_header_logger.py b/cookbook/logging_observability/response_header_logger.py deleted file mode 100644 index 749e337fbca..00000000000 --- a/cookbook/logging_observability/response_header_logger.py +++ /dev/null @@ -1,86 +0,0 @@ -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() diff --git a/tests/local_testing/test_custom_logger.py b/tests/local_testing/test_custom_logger.py index ca44a1cfa22..43ca7bddb3f 100644 --- a/tests/local_testing/test_custom_logger.py +++ b/tests/local_testing/test_custom_logger.py @@ -4,6 +4,10 @@ import inspect import os import time import traceback +from collections.abc import Mapping +from typing import Final, cast + +import httpx import pytest @@ -11,6 +15,7 @@ import pytest import litellm from litellm import completion, embedding from litellm.integrations.custom_logger import CustomLogger +from litellm.proxy._types import UserAPIKeyAuth class MyCustomHandler(CustomLogger): @@ -567,12 +572,88 @@ def test_redis_cache_completion_stream(): # test_redis_cache_completion_stream() +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()}, + ) + + + @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() @@ -602,7 +683,6 @@ async def test_response_header_logger_success_preserves_capture_and_metadata(met 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() @@ -639,7 +719,6 @@ async def test_response_header_logger_replaces_headers_after_fallback(fallback_h 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()