mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-17 04:34:25 +00:00
fix(openai_like)!: surface provider response headers on non-streaming chat completions (#46928)
* fix(openai_like): surface provider response headers on non-streaming chat completions Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(openai_like): mark the header test locals Final * fix(proxy): leave out response header values HTTP cannot carry Non-streaming OpenAI-like chat now forwards every provider header, so on the httpx transport a provider header value outside latin-1 made the proxy answer 500 after the provider had served the call. get_custom_headers now leaves such a value out and keeps the other headers, which also ends the same 500 on streaming Adds integration cells for the provider request id on non-streaming chat and the Responses bridge, LiteLLM headers an upstream cannot overwrite, spend parity under an upstream cost header, the latin-1 case, and the logging callback payload * fix(cost): ignore an upstream response cost header that is negative or not a number * test(integration): cover invalid upstream cost headers and non-latin-1 headers on every endpoint * fix(prometheus): keep deployment metrics when a provider rate-limit header is not a number * test(prometheus): type the rate-limit header test and drop its clock reads * test(integration): cover deployment metrics under an unparseable provider rate-limit header --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
This commit is contained in:
jesus-berri
GitHub
parent
37b3a697c7
commit
7f603da9f1
10 changed files with 697 additions and 8 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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 "")
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
@ -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,)
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue