scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-10 03:28:53 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 84

test: keep response header verification out of cookbook

Browse source
This commit is contained in:
Yujong Lee 2026-09-02 13:01:12 -07:00
parent eda63531cf
commit d7290c3a4f
3 changed files with 82 additions and 144 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

55
cookbook/logging_observability/response_header_logger.md
View file

@ -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

86
cookbook/logging_observability/response_header_logger.py
View file

@ -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()

85
tests/local_testing/test_custom_logger.py
View file

@ -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()