diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index 05dc2900349..d9ee6bdc270 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -2081,13 +2081,20 @@ def get_response_cost_from_hidden_params( additional_headers: Final = _hidden_params_dict.get("additional_headers", {}) if additional_headers and "llm_provider-x-litellm-response-cost" in additional_headers: - response_cost: Final = additional_headers["llm_provider-x-litellm-response-cost"] - if response_cost is None: - return None - return float(additional_headers["llm_provider-x-litellm-response-cost"]) + return _billable_header_cost(additional_headers["llm_provider-x-litellm-response-cost"]) return None +def _billable_header_cost(value: object) -> float | None: + if not isinstance(value, (str, int, float)): + return None + try: + cost: Final = float(value) + except ValueError: + return None + return cost if math.isfinite(cost) and cost >= 0 else None + + def response_cost_calculator( response_object: ModelResponse | EmbeddingResponse diff --git a/litellm/integrations/prometheus.py b/litellm/integrations/prometheus.py index 97a1073d7c0..517a28cf44a 100644 --- a/litellm/integrations/prometheus.py +++ b/litellm/integrations/prometheus.py @@ -3373,7 +3373,7 @@ class PrometheusLogger(CustomLogger): label_context=label_context, ) - if remaining_requests: + if isinstance(remaining_requests, int): """ "model_group", "api_provider", @@ -3387,7 +3387,7 @@ class PrometheusLogger(CustomLogger): ) self.litellm_remaining_requests_metric.labels(**_labels).set(remaining_requests) - if remaining_tokens: + if isinstance(remaining_tokens, int): _labels = prometheus_label_factory( supported_enum_labels=self.get_labels_for_metric(metric_name="litellm_remaining_tokens_metric"), enum_values=enum_values, diff --git a/litellm/llms/openai_like/chat/transformation.py b/litellm/llms/openai_like/chat/transformation.py index 83fcb43c9b8..e07b5627256 100644 --- a/litellm/llms/openai_like/chat/transformation.py +++ b/litellm/llms/openai_like/chat/transformation.py @@ -6,6 +6,7 @@ from typing import TYPE_CHECKING, Any, Final import httpx +from litellm.litellm_core_utils.core_helpers import set_provider_response_headers_in_hidden_params from litellm.secret_managers.main import get_secret_str from litellm.types.llms.openai import AllMessageValues, ChatCompletionAssistantMessage from litellm.types.utils import ModelResponse @@ -119,6 +120,7 @@ class OpenAILikeChatConfig(OpenAIGPTConfig): choice["message"] = message returned_response: Final = ModelResponse(**response_json) + set_provider_response_headers_in_hidden_params(returned_response, response.headers) if custom_llm_provider is not None: returned_response.model = custom_llm_provider + "/" + (returned_response.model or "") diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 2d387b43aa8..61c26cbd502 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -506,6 +506,10 @@ async def _cancel_pending_gather_tasks(tasks: Sequence["asyncio.Task[object]"]) pass +def _encodes_as_latin1(value: str) -> bool: + return all(ord(character) < 256 for character in value) + + @lru_cache(maxsize=512) def _litellm_model_supports_stream_options(litellm_model: str) -> bool: try: @@ -2027,10 +2031,11 @@ class ProxyBaseLLMRequestProcessing: headers.update(logging_caching_headers) try: - return {key: str(value) for key, value in headers.items() if value not in exclude_values} + present: Final = {key: str(value) for key, value in headers.items() if value not in exclude_values} except Exception as e: verbose_proxy_logger.error("Error setting custom headers: %s", e) return {} + return {key: value for key, value in present.items() if _encodes_as_latin1(value)} @staticmethod async def build_litellm_proxy_success_headers_from_llm_response( diff --git a/tests/integration/observability/test_callback_delivery.py b/tests/integration/observability/test_callback_delivery.py index 543e60da4a6..e786e655a13 100644 --- a/tests/integration/observability/test_callback_delivery.py +++ b/tests/integration/observability/test_callback_delivery.py @@ -11,7 +11,7 @@ import pytest import yaml from integration._support.client import Gateway, JsonValue, eventually, object_value, string_value from integration._support.database import read_rows -from integration._support.process import owned_proxy +from integration._support.process import graceful_stop_seconds, owned_proxy from integration._support.wire import Reply, Request, wire_server @@ -268,6 +268,89 @@ def test_streamed_responses_success_callback_carries_provider_apim_request_id(ga }, json.dumps(events[0]["hidden_params"]) +@pytest.mark.timeout(2 * graceful_stop_seconds() + 120) +def test_non_stream_openai_like_success_callback_carries_provider_headers(gateway: Gateway, tmp_path: Path) -> None: + marker: Final = uuid.uuid4().hex + request_id: Final = "req_" + marker + sink_secret: Final = "synthetic-sink-secret-" + marker + + def upstream(request: Request) -> Reply: + assert request.target.endswith("/chat/completions"), request.target + return Reply( + body=json.dumps( + { + "id": "chatcmpl-" + marker, + "object": "chat.completion", + "created": 1, + "model": "google.gemma-4-31b", + "choices": [ + {"index": 0, "message": {"role": "assistant", "content": marker}, "finish_reason": "stop"} + ], + "usage": {"prompt_tokens": 11, "completion_tokens": 4, "total_tokens": 15}, + } + ).encode(), + headers={"x-amzn-requestid": request_id, "x-ratelimit-remaining-requests": "997"}, + ) + + def sink(request: Request) -> Reply: + assert request.headers["authorization"] == f"Bearer {sink_secret}" + return Reply() + + with wire_server(upstream) as provider, wire_server(sink) as endpoint: + config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + config["litellm_settings"].update({"callbacks": ["generic_api"], "DEFAULT_FLUSH_INTERVAL_SECONDS": 1}) + path: Final = tmp_path / "callbacks.yaml" + path.write_text(yaml.safe_dump(config)) + with ( + owned_proxy( + gateway, + tmp_path, + { + "GENERIC_LOGGER_ENDPOINT": endpoint.url, + "GENERIC_LOGGER_HEADERS": f"Authorization=Bearer {sink_secret}", + }, + config=path, + ) as candidate, + candidate.scenario() as scenario, + ): + model: Final = scenario.model( + model="bedrock_mantle/google.gemma-4-31b", + api_base=provider.url, + api_key="synthetic-mantle-bearer", + aws_region_name="us-west-2", + ) + response: Final = candidate.request( + "POST", + "/v1/chat/completions", + { + "model": model, + "messages": [{"role": "user", "content": "header callback " + marker}], + "cache": {"no-cache": True}, + }, + ) + assert response.status_code == 200, response.text + assert response.json()["choices"][0]["message"]["content"] == marker, response.text + assert len(provider.drain()) == 1 + batches: Final[list[Request]] = [] # mutable-ok: drain() consumes the queue, later polls must keep earlier batches + + def delivered() -> tuple[dict, ...]: + batches.extend(endpoint.drain()) + return tuple( + event for batch in batches for event in json.loads(batch.body) if event.get("model_group") == model + ) + + events: Final = eventually(delivered, lambda values: len(values) == 1, seconds=10) + assert (events[0]["status"], events[0]["call_type"]) == ("success", "acompletion") + assert events[0]["litellm_call_id"] == response.headers["x-litellm-call-id"] + additional_headers: Final = events[0]["hidden_params"]["additional_headers"] or {} + assert { + name: additional_headers.get(name) + for name in ("llm_provider-x-amzn-requestid", "x_ratelimit_remaining_requests") + } == {"llm_provider-x-amzn-requestid": request_id, "x_ratelimit_remaining_requests": 997}, json.dumps( + events[0]["hidden_params"] + ) + + _RAISING_HOOK: Final = """ from litellm.integrations.custom_logger import CustomLogger diff --git a/tests/integration/providers/test_openai_like_response_headers_wire.py b/tests/integration/providers/test_openai_like_response_headers_wire.py new file mode 100644 index 00000000000..1ffa77cf326 --- /dev/null +++ b/tests/integration/providers/test_openai_like_response_headers_wire.py @@ -0,0 +1,435 @@ +import itertools +import json +import math +import uuid +from collections.abc import Callable, Mapping, Sequence +from hashlib import sha256 +from pathlib import Path +from types import MappingProxyType +from typing import Final, Literal, TypeAlias + +import httpx +import pytest +from integration._support.client import Gateway, Scenario, eventually, held +from integration._support.database import read_rows +from integration._support.process import graceful_stop_seconds, owned_proxy +from integration._support.prometheus_series import Sample, scrape, series_cap_config +from integration._support.vertex import service_account_json +from integration._support.wire import Reply, Request, Wire, wire_server +from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter + +Endpoint: TypeAlias = Literal["chat", "responses", "messages"] +Deploy: TypeAlias = Callable[[Gateway, Scenario, Wire], str] + +_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) +_PROJECT: Final = "scripted-project" +_REPLY_TEXT: Final = "header probe reply" +_USAGE: Final[JsonValue] = {"prompt_tokens": 13, "completion_tokens": 4, "total_tokens": 17} +_UPSTREAM_COST: Final = "0.0421" +_SPEND_SELECT: Final = 'SELECT spend FROM "LiteLLM_SpendLogs" WHERE api_key=%s' +_KEY_SPEND_SELECT: Final = 'SELECT spend FROM "LiteLLM_VerificationToken" WHERE token=%s' +_PATHS: Final[Mapping[Endpoint, str]] = MappingProxyType( + {"chat": "/v1/chat/completions", "responses": "/v1/responses", "messages": "/v1/messages"} +) +_RATE_LIMIT_HEADERS: Final[Mapping[str, str]] = MappingProxyType( + {"x-ratelimit-remaining-requests": "0", "x-ratelimit-remaining-tokens": "856, abc"} +) +_METERED_CALLS: Final[tuple[tuple[str, bool], ...]] = ( + ("groq", False), + ("groq", True), + ("openai", False), + ("openai", True), +) +_LATIN1_CALLS: Final[tuple[tuple[Endpoint, str, bool], ...]] = ( + ("chat", "bedrock_mantle", False), + ("chat", "bedrock_mantle", True), + ("responses", "groq", False), + ("responses", "groq", True), + ("messages", "groq", False), + ("messages", "groq", True), +) + + +def _keyed(model: str, **extra: JsonValue) -> Deploy: + def create(_gateway: Gateway, scenario: Scenario, wire: Wire) -> str: + return scenario.model(model=model, api_base=wire.url, api_key="synthetic-provider-key", **extra) + + return create + + +def _vertex(model: str, location: str) -> Deploy: + def create(gateway: Gateway, scenario: Scenario, wire: Wire) -> str: + return scenario.model( + model=f"vertex_ai/{model}", + api_base=wire.url, + api_key=None, + vertex_project=_PROJECT, + vertex_location=location, + vertex_credentials=service_account_json(_PROJECT, gateway.upstream_url.rstrip("/")), + ) + + return create + + +_MEMBERS: Final[Mapping[str, Deploy]] = MappingProxyType( + { + "bedrock_mantle": _keyed("bedrock_mantle/google.gemma-4-31b", aws_region_name="us-west-2"), + "groq": _keyed("groq/openai/gpt-oss-20b"), + "openai": _keyed("openai/gpt-5.4-mini"), + "watsonx": _keyed("watsonx/ibm/granite-3-8b-instruct", zen_api_key="synthetic-zen-key", project_id=_PROJECT), + "vertex_model_garden": _vertex("xai/grok-4.20-non-reasoning", "global"), + "maritalk": _keyed("maritalk/sabia-3"), + } +) + + +def _reply_json() -> bytes: + return json.dumps( + { + "id": f"chatcmpl-{uuid.uuid4().hex}", + "object": "chat.completion", + "created": 1700000000, + "model": "scripted", + "choices": [ + {"index": 0, "message": {"role": "assistant", "content": _REPLY_TEXT}, "finish_reason": "stop"} + ], + "usage": _USAGE, + "service_tier": "on_demand", + } + ).encode() + + +def _chunk(response_id: str, choices: list[JsonValue], **extra: JsonValue) -> bytes: + body: Final = { + "id": response_id, + "object": "chat.completion.chunk", + "created": 1700000000, + "model": "scripted", + "choices": choices, + **extra, + } + return f"data: {json.dumps(body)}\n\n".encode() + + +def _frames() -> tuple[bytes, ...]: + response_id: Final = f"chatcmpl-{uuid.uuid4().hex}" + return ( + _chunk( + response_id, [{"index": 0, "delta": {"role": "assistant", "content": _REPLY_TEXT}, "finish_reason": None}] + ), + _chunk(response_id, [{"index": 0, "delta": {}, "finish_reason": "stop"}]), + _chunk(response_id, [], usage=_USAGE), + b"data: [DONE]\n\n", + ) + + +def _peer(headers: Mapping[str, str]) -> Callable[[Request], Reply]: + def respond(request: Request) -> Reply: + if request.method != "POST": + return Reply(status=404) + if _JSON_OBJECT.validate_json(request.body).get("stream") is True: + return Reply(content_type="text/event-stream", chunks=_frames(), headers=headers) + return Reply(body=_reply_json(), headers=headers) + + return respond + + +def _request_id_headers(request_id: str) -> Mapping[str, str]: + return MappingProxyType({"x-request-id": request_id, "x-amzn-requestid": request_id}) + + +def _body(endpoint: Endpoint, model: str, stream: bool) -> dict[str, JsonValue]: + prompt: Final = f"header probe {uuid.uuid4().hex}" + match endpoint: + case "chat": + return { + "model": model, + "messages": [{"role": "user", "content": prompt}], + "stream": stream, + "cache": {"no-cache": True}, + **({"stream_options": {"include_usage": True}} if stream else {}), + } + case "responses": + return {"model": model, "input": prompt, "stream": stream} + case "messages": + return { + "model": model, + "max_tokens": 64, + "messages": [{"role": "user", "content": prompt}], + "stream": stream, + } + + +def _send(gateway: Gateway, endpoint: Endpoint, model: str, key: str, *, stream: bool = False) -> httpx.Response: + return gateway.request("POST", _PATHS[endpoint], _body(endpoint, model, stream), key=key) + + +class _Reply(BaseModel): + model_config = ConfigDict(frozen=True) + + +class _Text(_Reply): + text: str = "" + + +class _ChatMessage(_Reply): + content: str = "" + + +class _ChatChoice(_Reply): + message: _ChatMessage + + +class _ChatReply(_Reply): + choices: tuple[_ChatChoice, ...] + + +class _ResponsesItem(_Reply): + content: tuple[_Text, ...] = () + + +class _ResponsesReply(_Reply): + output: tuple[_ResponsesItem, ...] + + +class _MessagesReply(_Reply): + content: tuple[_Text, ...] + + +def _text(endpoint: Endpoint, response: httpx.Response) -> str: + match endpoint: + case "chat": + return _ChatReply.model_validate_json(response.content).choices[0].message.content + case "responses": + items: Final = _ResponsesReply.model_validate_json(response.content).output + return "".join(part.text for part in itertools.chain.from_iterable(item.content for item in items)) + case "messages": + return "".join(block.text for block in _MessagesReply.model_validate_json(response.content).content) + + +def _header_cost(response: httpx.Response) -> float: + values: Final = response.headers.get_list("x-litellm-response-cost") + return float(values[0]) if values else math.nan + + +def _llm_provider_request_ids(response: httpx.Response) -> dict[str, str | None]: + return {name: response.headers.get(name) for name in ("llm_provider-x-request-id", "llm_provider-x-amzn-requestid")} + + +def _spend(key: str) -> float: + digest: Final = sha256(key.encode()).hexdigest() + eventually(lambda: read_rows(_SPEND_SELECT, (digest,)), lambda rows: len(rows) >= 1, seconds=70) + rows: Final = held(lambda: read_rows(_SPEND_SELECT, (digest,)), lambda found: len(found) == 1, holding=3) + spend: Final = rows[0]["spend"] + assert isinstance(spend, (int, float)), rows + return float(spend) + + +def _model_values(samples: Sequence[Sample], name: str, model: str) -> tuple[float, ...]: + return tuple( + sample.value for sample in samples if sample.name == name and sample.labels.get("model_group") == model + ) + + +def _deployment_metrics( + samples: Sequence[Sample], models: Mapping[str, str] +) -> dict[str, tuple[float, tuple[float, ...], tuple[float, ...]]]: + return { + call: ( + sum(_model_values(samples, "litellm_deployment_success_responses_total", model)), + _model_values(samples, "litellm_remaining_requests_metric", model), + _model_values(samples, "litellm_remaining_tokens_metric", model), + ) + for call, model in models.items() + } + + +def _key_spend(key: str) -> float | None: + rows: Final = read_rows(_KEY_SPEND_SELECT, (sha256(key.encode()).hexdigest(),)) + spend: Final = rows[0]["spend"] if rows else None + return float(spend) if isinstance(spend, (int, float)) else None + + +@pytest.mark.parametrize( + ("endpoint", "member"), + [ + ("chat", "bedrock_mantle"), + ("chat", "groq"), + ("chat", "watsonx"), + ("chat", "vertex_model_garden"), + ("chat", "maritalk"), + ("responses", "groq"), + ], +) +def test_non_stream_returns_the_provider_request_id(gateway: Gateway, endpoint: Endpoint, member: str) -> None: + request_id: Final = f"req_{uuid.uuid4().hex}" + with wire_server(_peer(_request_id_headers(request_id))) as wire, gateway.scenario() as scenario: + model: Final = _MEMBERS[member](gateway, scenario, wire) + response: Final = _send(gateway, endpoint, model, scenario.key(models=[model])) + assert response.status_code == 200, response.text + assert _text(endpoint, response) == _REPLY_TEXT, response.text + assert _llm_provider_request_ids(response) == { + "llm_provider-x-request-id": request_id, + "llm_provider-x-amzn-requestid": request_id, + }, dict(response.headers) + assert len(wire.drain()) == 1 + + +def test_stream_chat_returns_the_provider_request_id(gateway: Gateway) -> None: + request_id: Final = f"req_{uuid.uuid4().hex}" + with wire_server(_peer(_request_id_headers(request_id))) as wire, gateway.scenario() as scenario: + model: Final = _MEMBERS["bedrock_mantle"](gateway, scenario, wire) + response: Final = _send(gateway, "chat", model, scenario.key(models=[model]), stream=True) + assert response.status_code == 200, response.text + assert _REPLY_TEXT in response.text, response.text + assert _llm_provider_request_ids(response) == { + "llm_provider-x-request-id": request_id, + "llm_provider-x-amzn-requestid": request_id, + }, dict(response.headers) + assert len(wire.drain()) == 1 + + +def test_upstream_cannot_overwrite_litellm_owned_headers(gateway: Gateway) -> None: + spoofed: Final = { + "x-litellm-call-id": "spoofed-call-id", + "x-litellm-model-id": "spoofed-model-id", + "x-litellm-attempted-fallbacks": "7", + } + with wire_server(_peer(spoofed)) as wire, gateway.scenario() as scenario: + model: Final = _MEMBERS["groq"](gateway, scenario, wire) + response: Final = _send(gateway, "chat", model, scenario.key(models=[model])) + assert response.status_code == 200, response.text + assert response.headers["x-litellm-call-id"] != spoofed["x-litellm-call-id"], dict(response.headers) + assert response.headers["x-litellm-model-id"] != spoofed["x-litellm-model-id"], dict(response.headers) + assert response.headers["x-litellm-attempted-fallbacks"] == "0", dict(response.headers) + assert len(wire.drain()) == 1 + + +@pytest.mark.parametrize(("endpoint", "member"), [("chat", "bedrock_mantle"), ("responses", "groq")]) +def test_non_stream_spend_matches_stream_spend_under_an_upstream_cost_header( + gateway: Gateway, endpoint: Endpoint, member: str +) -> None: + with wire_server(_peer({"x-litellm-response-cost": _UPSTREAM_COST})) as wire, gateway.scenario() as scenario: + model: Final = _MEMBERS[member](gateway, scenario, wire) + stream_key: Final = scenario.key(models=[model]) + non_stream_key: Final = scenario.key(models=[model]) + streamed: Final = _send(gateway, endpoint, model, stream_key, stream=True) + assert streamed.status_code == 200, streamed.text + assert _REPLY_TEXT in streamed.text, streamed.text + whole: Final = _send(gateway, endpoint, model, non_stream_key) + assert whole.status_code == 200, whole.text + assert _text(endpoint, whole) == _REPLY_TEXT, whole.text + stream_spend: Final = _spend(stream_key) + assert stream_spend > 0 + non_stream_spend: Final = _spend(non_stream_key) + assert math.isclose(non_stream_spend, stream_spend), (non_stream_spend, stream_spend) + assert len(wire.drain()) == 2 + + +@pytest.mark.parametrize("member", ["openai", "groq"]) +@pytest.mark.parametrize("upstream_cost", ["-5", "abc", "nan", "inf"], ids=["negative", "garbage", "nan", "infinite"]) +def test_an_invalid_upstream_cost_header_bills_the_price_map(gateway: Gateway, member: str, upstream_cost: str) -> None: + with ( + wire_server(_peer({"x-litellm-response-cost": upstream_cost})) as invalid_wire, + wire_server(_peer({})) as plain_wire, + gateway.scenario() as scenario, + ): + invalid_model: Final = _MEMBERS[member](gateway, scenario, invalid_wire) + plain_model: Final = _MEMBERS[member](gateway, scenario, plain_wire) + budget_key: Final = scenario.key(models=[invalid_model], max_budget=1) + reference_key: Final = scenario.key(models=[plain_model]) + billed: Final = _send(gateway, "chat", invalid_model, budget_key) + reference: Final = _send(gateway, "chat", plain_model, reference_key) + assert (billed.status_code, reference.status_code) == (200, 200), (billed.text, reference.text) + reference_spend: Final = _spend(reference_key) + assert reference_spend > 0 + key_spend: Final = eventually( + lambda: _key_spend(budget_key), + lambda spend: spend is not None and math.isclose(spend, reference_spend), + 70, + return_last_on_timeout=True, + ) + billed_costs: Final = { + "response_header": _header_cost(billed), + "spend_log": _spend(budget_key), + "key": key_spend if key_spend is not None else math.nan, + } + assert {name: math.isclose(cost, reference_spend) for name, cost in billed_costs.items()} == dict.fromkeys( + billed_costs, True + ), (billed_costs, reference_spend) + followup: Final = _send(gateway, "chat", invalid_model, budget_key) + assert followup.status_code == 200, followup.text + assert (len(invalid_wire.drain()), len(plain_wire.drain())) == (2, 1) + + +@pytest.mark.timeout(2 * graceful_stop_seconds() + 120) +def test_provider_header_outside_latin1_does_not_fail_the_call(gateway: Gateway, tmp_path: Path) -> None: + request_id: Final = f"req_{uuid.uuid4().hex}" + utf8_bytes_as_wire_text: Final = "東京リージョン".encode().decode("latin-1") + headers: Final = {**_request_id_headers(request_id), "x-provider-region": utf8_bytes_as_wire_text} + with ( + wire_server(_peer(headers)) as wire, + owned_proxy(gateway, tmp_path, {"DISABLE_AIOHTTP_TRANSPORT": "True"}) as httpx_transport, + httpx_transport.scenario() as scenario, + ): + models: Final = { + member: _MEMBERS[member](httpx_transport, scenario, wire) for member in ("bedrock_mantle", "groq") + } + key: Final = scenario.key(models=list(models.values())) + replies: Final = { + f"{endpoint}-{'stream' if stream else 'whole'}": _send( + httpx_transport, endpoint, models[member], key, stream=stream + ) + for endpoint, member, stream in _LATIN1_CALLS + } + assert {call: reply.status_code for call, reply in replies.items()} == dict.fromkeys(replies, 200), { + call: reply.text for call, reply in replies.items() + } + assert {call: reply.headers.get("llm_provider-x-provider-region") for call, reply in replies.items()} == ( + dict.fromkeys(replies, None) + ) + assert {call: _REPLY_TEXT in reply.text for call, reply in replies.items()} == dict.fromkeys(replies, True) + assert {endpoint: _text(endpoint, replies[f"{endpoint}-whole"]) for endpoint in _PATHS} == dict.fromkeys( + _PATHS, _REPLY_TEXT + ) + forwarding: Final = ("chat-whole", "chat-stream", "responses-whole") + assert {call: _llm_provider_request_ids(replies[call]) for call in forwarding} == dict.fromkeys( + forwarding, {"llm_provider-x-request-id": request_id, "llm_provider-x-amzn-requestid": request_id} + ) + assert len(wire.drain()) == len(_LATIN1_CALLS) + + +@pytest.mark.timeout(2 * graceful_stop_seconds() + 120) +def test_an_unparseable_rate_limit_header_keeps_the_deployment_metrics(gateway: Gateway, tmp_path: Path) -> None: + multiproc_dir: Final = tmp_path / "prometheus" + multiproc_dir.mkdir() + with ( + wire_server(_peer(_RATE_LIMIT_HEADERS)) as wire, + owned_proxy( + gateway, + tmp_path, + {"PROMETHEUS_MULTIPROC_DIR": str(multiproc_dir)}, + config=series_cap_config(tmp_path, {}), + workers=2, + ) as metered, + metered.scenario() as scenario, + ): + calls: Final = { + f"{member}-{'stream' if stream else 'whole'}": (_MEMBERS[member](metered, scenario, wire), stream) + for member, stream in _METERED_CALLS + } + models: Final = {call: model for call, (model, _) in calls.items()} + key: Final = scenario.key(models=list(models.values())) + statuses: Final = { + call: _send(metered, "chat", model, key, stream=stream).status_code + for call, (model, stream) in calls.items() + } + assert statuses == dict.fromkeys(calls, 200) + observed: Final = eventually( + lambda: _deployment_metrics(scrape(metered), models), + lambda metrics: all(successes >= 1 for successes, _, _ in metrics.values()), + 30, + return_last_on_timeout=True, + ) + assert observed == dict.fromkeys(calls, (1.0, (0.0,), ())), observed + assert len(wire.drain()) == len(calls) diff --git a/tests/unit/integrations/test_prometheus_remaining_tokens_router_fallback.py b/tests/unit/integrations/test_prometheus_remaining_tokens_router_fallback.py index 45b378d10fe..28d744cdb06 100644 --- a/tests/unit/integrations/test_prometheus_remaining_tokens_router_fallback.py +++ b/tests/unit/integrations/test_prometheus_remaining_tokens_router_fallback.py @@ -20,10 +20,14 @@ Tests cover: - llm_router unavailable / model_group missing / router raises → silent no-op. """ +from collections.abc import Iterator +from datetime import datetime +from typing import Final from unittest.mock import AsyncMock, MagicMock, patch import pytest from prometheus_client import REGISTRY +from prometheus_client.samples import Sample from litellm.integrations.prometheus import PrometheusLogger @@ -293,3 +297,42 @@ class TestRouterFallbackDefensivePaths: ) prometheus_logger.litellm_remaining_tokens_metric.labels.assert_not_called() + + +def _registry_samples() -> Iterator[Sample]: + for metric in REGISTRY.collect(): + yield from metric.samples + + +def _sample_values(sample_name: str) -> tuple[float, ...]: + return tuple(sample.value for sample in _registry_samples() if sample.name == sample_name) + + +def test_unparseable_provider_ratelimit_header_skips_only_its_own_gauge( + prometheus_logger: PrometheusLogger, +): + payload: Final = _build_payload( + model_group="groq-group", + custom_llm_provider="groq", + additional_headers={ + "x_ratelimit_remaining_requests": 0, + "x_ratelimit_remaining_tokens": "856, abc", + }, + ) + request_kwargs: Final = { + "litellm_params": {"custom_llm_provider": "groq", "metadata": {}}, + "model": "llama-3.1-8b-instant", + "standard_logging_object": payload, + } + + prometheus_logger.set_llm_deployment_success_metrics( + request_kwargs=request_kwargs, + start_time=datetime(2026, 1, 1), + end_time=datetime(2026, 1, 1, 0, 0, 1), + enum_values=_enum_values(model_group="groq-group"), + output_tokens=2.0, + ) + + assert _sample_values("litellm_remaining_requests_metric") == (0.0,) + assert _sample_values("litellm_remaining_tokens_metric") == () + assert _sample_values("litellm_deployment_success_responses_total") == (1.0,) diff --git a/tests/unit/llms/openai_like/chat/test_openai_like_chat_transformation.py b/tests/unit/llms/openai_like/chat/test_openai_like_chat_transformation.py index 88382b80d16..dbc6ec54a40 100644 --- a/tests/unit/llms/openai_like/chat/test_openai_like_chat_transformation.py +++ b/tests/unit/llms/openai_like/chat/test_openai_like_chat_transformation.py @@ -1,5 +1,11 @@ +from typing import Final +from unittest.mock import MagicMock + +import httpx import pytest + from litellm.llms.openai_like.chat.transformation import OpenAILikeChatConfig +from litellm.types.utils import ModelResponse def test_sanitize_usage_obj_handles_null_tokens(): @@ -47,3 +53,32 @@ def test_sanitize_usage_obj_valid_usage(): # Assert assert sanitized_json == original_json # The object should be unchanged + + +def test_transform_response_surfaces_provider_headers_as_llm_provider_headers(): + raw_response: Final = httpx.Response( + status_code=200, + headers={"x-amzn-requestid": "req_abc123", "content-type": "application/json"}, + json={ + "id": "chatcmpl-1", + "object": "chat.completion", + "created": 1, + "model": "google.gemma-4-31b", + "choices": [{"index": 0, "finish_reason": "stop", "message": {"role": "assistant", "content": "hi"}}], + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, + }, + ) + + response: Final = OpenAILikeChatConfig().transform_response( + model="google.gemma-4-31b", + raw_response=raw_response, + model_response=ModelResponse(), + logging_obj=MagicMock(), + request_data={}, + messages=[{"role": "user", "content": "hi"}], + optional_params={}, + litellm_params={}, + encoding=None, + ) + + assert response._hidden_params["additional_headers"]["llm_provider-x-amzn-requestid"] == "req_abc123" diff --git a/tests/unit/proxy/test_common_request_processing.py b/tests/unit/proxy/test_common_request_processing.py index 1d197602543..c3556bbaeb7 100644 --- a/tests/unit/proxy/test_common_request_processing.py +++ b/tests/unit/proxy/test_common_request_processing.py @@ -1724,6 +1724,29 @@ class TestProxyBaseLLMRequestProcessing: assert "x-litellm-classifier-cost" not in headers + def test_get_custom_headers_drops_a_provider_header_http_cannot_carry(self): + user_api_key_dict: Final = MagicMock(spec=UserAPIKeyAuth) + user_api_key_dict.tpm_limit = None + user_api_key_dict.rpm_limit = None + user_api_key_dict.max_budget = None + user_api_key_dict.spend = 0 + + headers: Final = ProxyBaseLLMRequestProcessing.get_custom_headers( + user_api_key_dict=user_api_key_dict, + call_id="call-1", + **{ + "llm_provider-x-amzn-requestid": "req-1", + "llm_provider-x-provider-region": "東京リージョン", + "llm_provider-x-provider-city": "Zürich", + }, + ) + response: Final = Response() + response.headers.update(headers) + + assert response.headers["llm_provider-x-amzn-requestid"] == "req-1" + assert response.headers["llm_provider-x-provider-city"] == "Zürich" + assert "llm_provider-x-provider-region" not in response.headers + def test_get_cost_breakdown_from_logging_obj_helper(self): """ Test the helper function that extracts cost breakdown information. diff --git a/tests/unit/test_cost_calculator.py b/tests/unit/test_cost_calculator.py index c61ebb3737c..2fbe6ec294b 100644 --- a/tests/unit/test_cost_calculator.py +++ b/tests/unit/test_cost_calculator.py @@ -229,6 +229,62 @@ def test_response_cost_calculator_custom_pricing_ignores_upstream_response_cost_ assert cost == pytest.approx(expected_cost) +def _header_cost_chat_response(additional_headers: dict[str, str]) -> ModelResponse: + response: Final = ModelResponse( + model="cc-header-cost-chat", + choices=[Choices(message=Message(content="ok"))], + usage=Usage(prompt_tokens=10, completion_tokens=20, total_tokens=30), + ) + response._hidden_params = {"additional_headers": additional_headers} + return response + + +def _cost_with_upstream_headers(monkeypatch: pytest.MonkeyPatch, additional_headers: dict[str, str]) -> float: + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url="")) + litellm.register_model( + model_cost={ + "cc-header-cost-chat": { + "litellm_provider": "openai", + "mode": "chat", + "input_cost_per_token": 0.001, + "output_cost_per_token": 0.002, + } + } + ) + return response_cost_calculator( + response_object=_header_cost_chat_response(additional_headers), + model="cc-header-cost-chat", + custom_llm_provider="openai", + call_type="completion", + optional_params={}, + ) + + +@pytest.mark.parametrize("header_cost", ("-5", "not-a-number", "nan", "inf")) +def test_response_cost_calculator_prices_from_the_map_when_upstream_cost_header_is_not_billable( + monkeypatch: pytest.MonkeyPatch, + header_cost: str, +) -> None: + price_map_cost: Final = _cost_with_upstream_headers(monkeypatch, {}) + + cost: Final = _cost_with_upstream_headers(monkeypatch, {"llm_provider-x-litellm-response-cost": header_cost}) + + assert price_map_cost > 0 + assert cost == pytest.approx(price_map_cost) + + +@pytest.mark.parametrize(("header_cost", "expected_cost"), (("0", 0.0), ("0.25", 0.25))) +def test_response_cost_calculator_bills_a_finite_non_negative_upstream_cost_header( + monkeypatch: pytest.MonkeyPatch, + header_cost: str, + expected_cost: float, +) -> None: + cost: Final = _cost_with_upstream_headers(monkeypatch, {"llm_provider-x-litellm-response-cost": header_cost}) + + assert cost == expected_cost + + def test_response_cost_calculator_keeps_optional_params_out_of_hidden_params(): class MockResponse(BaseModel): pass