From 184969df37fcab351fc828dd991ae8ccf068221a Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 23 Sep 2026 20:28:04 -0700 Subject: [PATCH 001/120] fix(models): add fireworks deprecation date for glm 5.2 fast serverless rows (#42866) Co-authored-by: kerry Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/model_prices_and_context_window_backup.json | 2 ++ model_prices_and_context_window.json | 2 ++ 2 files changed, 4 insertions(+) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 1a2501ab19b..60ac969781e 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -59469,6 +59469,7 @@ }, "fireworks_ai/glm-5p2-fast": { "cache_read_input_token_cost": 2.1e-07, + "deprecation_date": "2026-09-25", "input_cost_per_token": 2.1e-06, "litellm_provider": "fireworks_ai", "max_input_tokens": 1048576, @@ -59698,6 +59699,7 @@ }, "fireworks_ai/accounts/fireworks/routers/glm-5p2-fast": { "cache_read_input_token_cost": 2.1e-07, + "deprecation_date": "2026-09-25", "input_cost_per_token": 2.1e-06, "litellm_provider": "fireworks_ai", "max_input_tokens": 1048576, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 1a2501ab19b..60ac969781e 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -59469,6 +59469,7 @@ }, "fireworks_ai/glm-5p2-fast": { "cache_read_input_token_cost": 2.1e-07, + "deprecation_date": "2026-09-25", "input_cost_per_token": 2.1e-06, "litellm_provider": "fireworks_ai", "max_input_tokens": 1048576, @@ -59698,6 +59699,7 @@ }, "fireworks_ai/accounts/fireworks/routers/glm-5p2-fast": { "cache_read_input_token_cost": 2.1e-07, + "deprecation_date": "2026-09-25", "input_cost_per_token": 2.1e-06, "litellm_provider": "fireworks_ai", "max_input_tokens": 1048576, From 58e41e697bd2cdda98701932b5e85143c879e5df Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 23 Sep 2026 20:37:43 -0700 Subject: [PATCH 002/120] chore(vertex_ai): add deprecation dates for retired claude 3 and jamba 1.5 partner models (#42867) Co-authored-by: kerry Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/model_prices_and_context_window_backup.json | 13 +++++++++++++ model_prices_and_context_window.json | 13 +++++++++++++ 2 files changed, 26 insertions(+) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 60ac969781e..18c8578cec8 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -47710,6 +47710,7 @@ ] }, "vertex_ai/claude-3-5-haiku": { + "deprecation_date": "2026-07-05", "input_cost_per_token": 1e-06, "litellm_provider": "vertex_ai-anthropic_models", "max_input_tokens": 200000, @@ -47723,6 +47724,7 @@ "supports_tool_choice": true }, "vertex_ai/claude-3-5-haiku@20241022": { + "deprecation_date": "2026-07-05", "input_cost_per_token": 1e-06, "litellm_provider": "vertex_ai-anthropic_models", "max_input_tokens": 200000, @@ -47794,6 +47796,7 @@ "prompt_cache_min_tokens": 4096 }, "vertex_ai/claude-3-5-sonnet": { + "deprecation_date": "2026-02-19", "input_cost_per_token": 3e-06, "litellm_provider": "vertex_ai-anthropic_models", "max_input_tokens": 200000, @@ -47809,6 +47812,7 @@ "supports_vision": true }, "vertex_ai/claude-3-5-sonnet@20240620": { + "deprecation_date": "2026-02-19", "input_cost_per_token": 3e-06, "litellm_provider": "vertex_ai-anthropic_models", "max_input_tokens": 200000, @@ -47823,6 +47827,7 @@ "supports_vision": true }, "vertex_ai/claude-3-haiku": { + "deprecation_date": "2026-08-23", "input_cost_per_token": 2.5e-07, "litellm_provider": "vertex_ai-anthropic_models", "max_input_tokens": 200000, @@ -47836,6 +47841,7 @@ "supports_vision": true }, "vertex_ai/claude-3-haiku@20240307": { + "deprecation_date": "2026-08-23", "input_cost_per_token": 2.5e-07, "litellm_provider": "vertex_ai-anthropic_models", "max_input_tokens": 200000, @@ -47849,6 +47855,7 @@ "supports_vision": true }, "vertex_ai/claude-3-opus": { + "deprecation_date": "2025-08-01", "input_cost_per_token": 1.5e-05, "litellm_provider": "vertex_ai-anthropic_models", "max_input_tokens": 200000, @@ -47862,6 +47869,7 @@ "supports_vision": true }, "vertex_ai/claude-3-opus@20240229": { + "deprecation_date": "2025-08-01", "input_cost_per_token": 1.5e-05, "litellm_provider": "vertex_ai-anthropic_models", "max_input_tokens": 200000, @@ -49173,6 +49181,7 @@ "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing" }, "vertex_ai/jamba-1.5": { + "deprecation_date": "2026-02-27", "input_cost_per_token": 2e-07, "litellm_provider": "vertex_ai-ai21_models", "max_input_tokens": 256000, @@ -49183,6 +49192,7 @@ "supports_tool_choice": true }, "vertex_ai/jamba-1.5-large": { + "deprecation_date": "2026-02-27", "input_cost_per_token": 2e-06, "litellm_provider": "vertex_ai-ai21_models", "max_input_tokens": 256000, @@ -49193,6 +49203,7 @@ "supports_tool_choice": true }, "vertex_ai/jamba-1.5-large@001": { + "deprecation_date": "2026-02-27", "input_cost_per_token": 2e-06, "litellm_provider": "vertex_ai-ai21_models", "max_input_tokens": 256000, @@ -49203,6 +49214,7 @@ "supports_tool_choice": true }, "vertex_ai/jamba-1.5-mini": { + "deprecation_date": "2026-02-27", "input_cost_per_token": 2e-07, "litellm_provider": "vertex_ai-ai21_models", "max_input_tokens": 256000, @@ -49213,6 +49225,7 @@ "supports_tool_choice": true }, "vertex_ai/jamba-1.5-mini@001": { + "deprecation_date": "2026-02-27", "input_cost_per_token": 2e-07, "litellm_provider": "vertex_ai-ai21_models", "max_input_tokens": 256000, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 60ac969781e..18c8578cec8 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -47710,6 +47710,7 @@ ] }, "vertex_ai/claude-3-5-haiku": { + "deprecation_date": "2026-07-05", "input_cost_per_token": 1e-06, "litellm_provider": "vertex_ai-anthropic_models", "max_input_tokens": 200000, @@ -47723,6 +47724,7 @@ "supports_tool_choice": true }, "vertex_ai/claude-3-5-haiku@20241022": { + "deprecation_date": "2026-07-05", "input_cost_per_token": 1e-06, "litellm_provider": "vertex_ai-anthropic_models", "max_input_tokens": 200000, @@ -47794,6 +47796,7 @@ "prompt_cache_min_tokens": 4096 }, "vertex_ai/claude-3-5-sonnet": { + "deprecation_date": "2026-02-19", "input_cost_per_token": 3e-06, "litellm_provider": "vertex_ai-anthropic_models", "max_input_tokens": 200000, @@ -47809,6 +47812,7 @@ "supports_vision": true }, "vertex_ai/claude-3-5-sonnet@20240620": { + "deprecation_date": "2026-02-19", "input_cost_per_token": 3e-06, "litellm_provider": "vertex_ai-anthropic_models", "max_input_tokens": 200000, @@ -47823,6 +47827,7 @@ "supports_vision": true }, "vertex_ai/claude-3-haiku": { + "deprecation_date": "2026-08-23", "input_cost_per_token": 2.5e-07, "litellm_provider": "vertex_ai-anthropic_models", "max_input_tokens": 200000, @@ -47836,6 +47841,7 @@ "supports_vision": true }, "vertex_ai/claude-3-haiku@20240307": { + "deprecation_date": "2026-08-23", "input_cost_per_token": 2.5e-07, "litellm_provider": "vertex_ai-anthropic_models", "max_input_tokens": 200000, @@ -47849,6 +47855,7 @@ "supports_vision": true }, "vertex_ai/claude-3-opus": { + "deprecation_date": "2025-08-01", "input_cost_per_token": 1.5e-05, "litellm_provider": "vertex_ai-anthropic_models", "max_input_tokens": 200000, @@ -47862,6 +47869,7 @@ "supports_vision": true }, "vertex_ai/claude-3-opus@20240229": { + "deprecation_date": "2025-08-01", "input_cost_per_token": 1.5e-05, "litellm_provider": "vertex_ai-anthropic_models", "max_input_tokens": 200000, @@ -49173,6 +49181,7 @@ "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing" }, "vertex_ai/jamba-1.5": { + "deprecation_date": "2026-02-27", "input_cost_per_token": 2e-07, "litellm_provider": "vertex_ai-ai21_models", "max_input_tokens": 256000, @@ -49183,6 +49192,7 @@ "supports_tool_choice": true }, "vertex_ai/jamba-1.5-large": { + "deprecation_date": "2026-02-27", "input_cost_per_token": 2e-06, "litellm_provider": "vertex_ai-ai21_models", "max_input_tokens": 256000, @@ -49193,6 +49203,7 @@ "supports_tool_choice": true }, "vertex_ai/jamba-1.5-large@001": { + "deprecation_date": "2026-02-27", "input_cost_per_token": 2e-06, "litellm_provider": "vertex_ai-ai21_models", "max_input_tokens": 256000, @@ -49203,6 +49214,7 @@ "supports_tool_choice": true }, "vertex_ai/jamba-1.5-mini": { + "deprecation_date": "2026-02-27", "input_cost_per_token": 2e-07, "litellm_provider": "vertex_ai-ai21_models", "max_input_tokens": 256000, @@ -49213,6 +49225,7 @@ "supports_tool_choice": true }, "vertex_ai/jamba-1.5-mini@001": { + "deprecation_date": "2026-02-27", "input_cost_per_token": 2e-07, "litellm_provider": "vertex_ai-ai21_models", "max_input_tokens": 256000, From 3b715525d3717d636a9c663fcd755c1491f7c2de Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 23 Sep 2026 20:42:27 -0700 Subject: [PATCH 003/120] test(integration): assert /v1/models reports max_input_tokens and max_output_tokens (#42858) * test(integration): assert /v1/models reports max_input_tokens and max_output_tokens Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(integration): pin published gpt-4o-mini limits instead of reading the cost map in-test Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(integration): name the gpt-4o-mini deployment explicitly in the /v1/models test Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(integration): cite the source of the pinned gpt-4o-mini limits Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: kerry Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../test_model_listing_token_limits.py | 29 +++++++++++++++++++ 1 file changed, 29 insertions(+) create mode 100644 tests/integration/pricing/test_model_listing_token_limits.py diff --git a/tests/integration/pricing/test_model_listing_token_limits.py b/tests/integration/pricing/test_model_listing_token_limits.py new file mode 100644 index 00000000000..13702748ec5 --- /dev/null +++ b/tests/integration/pricing/test_model_listing_token_limits.py @@ -0,0 +1,29 @@ +import uuid +from typing import Final + +from integration._support.client import Gateway, object_value +from pydantic import JsonValue + + +def _listed_model(gateway: Gateway, model: str) -> dict[str, JsonValue]: + entries: Final = gateway.get("/v1/models")["data"] + assert isinstance(entries, list) + return next(object_value(entry) for entry in entries if object_value(entry)["id"] == model) + + +def test_v1_models_carries_cost_map_context_window_for_a_known_model(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + model: Final = scenario.model(model="openai/gpt-4o-mini") + listed: Final = _listed_model(gateway, model) + # OpenAI publishes these for gpt-4o-mini: https://platform.openai.com/docs/models/gpt-4o-mini (checked 2026-09-24) + assert listed["max_input_tokens"] == 128000, listed + assert listed["max_output_tokens"] == 16384, listed + + +def test_v1_models_carries_deployment_model_info_limits_for_an_unknown_model(gateway: Gateway) -> None: + unknown: Final = f"openai/custom-{uuid.uuid4().hex}" + with gateway.scenario() as scenario: + model: Final = scenario.model(model=unknown, model_info={"max_input_tokens": 4321, "max_output_tokens": 987}) + listed: Final = _listed_model(gateway, model) + assert listed["max_input_tokens"] == 4321, listed + assert listed["max_output_tokens"] == 987, listed From d248cc591430d6297ac29bcd9a560f20708a157f Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 23 Sep 2026 20:53:17 -0700 Subject: [PATCH 004/120] fix(fireworks_ai): route firerouter short names and bill pass-through legs at the routed model's rates (#42814) * fix(fireworks_ai): route firerouter short names and bill pass-through legs at the routed model's rates fireworks_ai/firerouter and fireworks_ai/firerouter/ resolve to accounts/fireworks/routers/... instead of a models/ path, and the cost calculator falls back to the routed model's own catalog entry before the Fireworks size buckets so a Claude leg is no longer priced at $0 Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(fireworks_ai): bill routed legs under the routed model's own provider Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(fireworks_ai): require the k suffix when parsing tiered input fields Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: kerry Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/llms/fireworks_ai/common_utils.py | 3 +- litellm/llms/fireworks_ai/cost_calculator.py | 9 +- .../test_fireworks_ai_router_slug_wire.py | 117 +++++++++++++++++- .../test_fireworks_ai_common_utils.py | 6 +- .../test_fireworks_ai_cost_calculator.py | 69 +++++++++++ 5 files changed, 199 insertions(+), 5 deletions(-) diff --git a/litellm/llms/fireworks_ai/common_utils.py b/litellm/llms/fireworks_ai/common_utils.py index 21a630a76d7..ae52b89aa58 100644 --- a/litellm/llms/fireworks_ai/common_utils.py +++ b/litellm/llms/fireworks_ai/common_utils.py @@ -59,6 +59,7 @@ def resolve_fireworks_api_key(api_key: str | None) -> str | None: AZURE_FOUNDRY_FIREWORKS_MODEL_ID_PREFIX: Final = "FW-" +FIREROUTER: Final = "firerouter" def resolve_fireworks_resource_name(model: str) -> str: @@ -67,7 +68,7 @@ def resolve_fireworks_resource_name(model: str) -> str: return stripped if stripped.startswith(("routers/", "models/")): return f"accounts/fireworks/{stripped}" - if stripped.endswith("-fast"): + if stripped.endswith("-fast") or stripped == FIREROUTER or stripped.startswith(f"{FIREROUTER}/"): return f"accounts/fireworks/routers/{stripped}" return f"accounts/fireworks/models/{stripped}" diff --git a/litellm/llms/fireworks_ai/cost_calculator.py b/litellm/llms/fireworks_ai/cost_calculator.py index 4b6ca7c9896..d42b4ddbbf7 100644 --- a/litellm/llms/fireworks_ai/cost_calculator.py +++ b/litellm/llms/fireworks_ai/cost_calculator.py @@ -59,6 +59,13 @@ def get_base_model_for_pricing(model_name: str) -> str: def _resolve_model_info(model: str) -> ModelInfo: try: return get_model_info(model=model, custom_llm_provider="fireworks_ai") + except Exception: + return _resolve_routed_model_info(model) + + +def _resolve_routed_model_info(model: str) -> ModelInfo: + try: + return get_model_info(model=model.removeprefix("fireworks_ai/")) except Exception: base_model: Final = get_base_model_for_pricing(model_name=model) return get_model_info(model=base_model, custom_llm_provider="fireworks_ai") @@ -81,7 +88,7 @@ def cost_per_token(model: str, usage: Usage, current_time: datetime | None = Non return generic_cost_per_token( model=model, usage=usage, - custom_llm_provider="fireworks_ai", + custom_llm_provider=model_info["litellm_provider"], model_info=model_info, current_time=current_time, ) diff --git a/tests/integration/providers/test_fireworks_ai_router_slug_wire.py b/tests/integration/providers/test_fireworks_ai_router_slug_wire.py index 4b4ad0b1243..55c83945ec7 100644 --- a/tests/integration/providers/test_fireworks_ai_router_slug_wire.py +++ b/tests/integration/providers/test_fireworks_ai_router_slug_wire.py @@ -1,16 +1,70 @@ import json +import uuid +from pathlib import Path from typing import Final import pytest -from integration._support.client import Gateway +from integration._support.client import Gateway, eventually +from integration._support.database import read_rows from integration._support.wire import Reply, Request, wire_server from pydantic import JsonValue, TypeAdapter _ROUTER_SLUG: Final = "routers/glm-latest" _ROUTER_RESOURCE: Final = "accounts/fireworks/routers/glm-latest" +_FIREROUTER_SLUGS: Final = ("firerouter", "firerouter/kimi-k3/deepseek-v4") _API_KEY: Final = "synthetic-fireworks-key" _PROMPT: Final = "route me through the router" +_COST_MAP_PATH: Final = Path(__file__).resolve().parents[3] / "model_prices_and_context_window.json" _JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) +_COST_MAP: Final = TypeAdapter(dict[str, dict[str, object]]) + + +def _positive_rate(entry: dict[str, object], field: str) -> bool: + value: Final = entry.get(field) + return isinstance(value, (int, float)) and value > 0 + + +def _pick_routed_model() -> str: + catalog: Final = _COST_MAP.validate_json(_COST_MAP_PATH.read_bytes()) + return next( + key + for key, entry in catalog.items() + if "/" not in key + and entry.get("litellm_provider") == "anthropic" + and _positive_rate(entry, "input_cost_per_token") + and _positive_rate(entry, "output_cost_per_token") + and f"fireworks_ai/{key}" not in catalog + ) + + +def _catalog_cost(model: str, field: str) -> float: + cost_value: Final = _COST_MAP.validate_json(_COST_MAP_PATH.read_bytes())[model][field] + assert isinstance(cost_value, (int, float)) + return float(cost_value) + + +_ROUTED_MODEL: Final = _pick_routed_model() + + +def _approx(value: float) -> object: + return pytest.approx(value, rel=1e-6) # pyright: ignore[reportUnknownMemberType] # pytest lacks typed approx stubs + + +def _chat_completion(identity: str, model: str, prompt_tokens: int, completion_tokens: int) -> bytes: + return json.dumps( + { + "id": identity, + "object": "chat.completion", + "created": 1, + "model": model, + "choices": [{"index": 0, "message": {"role": "assistant", "content": "routed"}, "finish_reason": "stop"}], + "usage": { + "prompt_tokens": prompt_tokens, + "completion_tokens": completion_tokens, + "total_tokens": prompt_tokens + completion_tokens, + }, + } + ).encode() def _provider_body(request: Request, target: str) -> dict[str, JsonValue]: @@ -82,3 +136,64 @@ def test_fireworks_router_slug_text_completion_sends_router_resource_not_models_ payload: Final = _JSON_OBJECT.validate_json(response.content) assert payload["choices"] == [{"index": 0, "text": "routed", "finish_reason": "stop", "logprobs": None}] assert [(request.method, request.target) for request in wire.drain()] == [("POST", "/completions")] + + +@pytest.mark.parametrize("slug", _FIREROUTER_SLUGS) +def test_fireworks_firerouter_short_name_sends_router_resource_not_models_path(gateway: Gateway, slug: str) -> None: + resource: Final = f"accounts/fireworks/routers/{slug}" + + def respond(request: Request) -> Reply: + body: Final = _provider_body(request, "/chat/completions") + assert body["model"] == resource, body + return Reply(body=_chat_completion(f"fw-{slug}", resource, 5, 1)) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=f"fireworks_ai/{slug}", api_base=wire.url, api_key=_API_KEY) + response: Final = gateway.request( + "POST", + "/v1/chat/completions", + {"model": model, "messages": [{"role": "user", "content": _PROMPT}]}, + ) + assert response.status_code == 200, response.text + payload: Final = _JSON_OBJECT.validate_json(response.content) + assert payload["choices"] == [ + {"finish_reason": "stop", "index": 0, "message": {"role": "assistant", "content": "routed"}} + ] + assert [(request.method, request.target) for request in wire.drain()] == [("POST", "/chat/completions")] + + +def test_fireworks_firerouter_claude_leg_is_charged_at_the_routed_models_own_rate(gateway: Gateway) -> None: + identity: Final = f"fw-firerouter-claude-{uuid.uuid4().hex}" + + def respond(request: Request) -> Reply: + body: Final = _provider_body(request, "/chat/completions") + assert body["model"] == "accounts/fireworks/routers/firerouter", body + assert request.headers["x-anthropic-api-key"] == "synthetic-anthropic-key" + return Reply(body=_chat_completion(identity, _ROUTED_MODEL, 23, 41)) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model="fireworks_ai/firerouter", + api_base=wire.url, + api_key=_API_KEY, + extra_headers={"x-anthropic-api-key": "synthetic-anthropic-key"}, + ) + response: Final = gateway.request( + "POST", + "/v1/chat/completions", + {"model": model, "messages": [{"role": "user", "content": _PROMPT}]}, + ) + assert response.status_code == 200, response.text + expected_cost: Final = 23 * _catalog_cost(_ROUTED_MODEL, "input_cost_per_token") + 41 * _catalog_cost( + _ROUTED_MODEL, "output_cost_per_token" + ) + assert expected_cost > 0 + assert float(response.headers["x-litellm-response-cost"]) == _approx(expected_cost) + rows: Final = eventually( + lambda: read_rows('SELECT spend FROM "LiteLLM_SpendLogs" WHERE request_id=%s', (identity,)), + lambda values: len(values) == 1, + seconds=70, + ) + spend: Final = rows[0]["spend"] + assert isinstance(spend, (int, float, str)) + assert float(spend) == _approx(expected_cost) diff --git a/tests/unit/llms/fireworks_ai/test_fireworks_ai_common_utils.py b/tests/unit/llms/fireworks_ai/test_fireworks_ai_common_utils.py index 16226a3ce74..e505f2ae8a6 100644 --- a/tests/unit/llms/fireworks_ai/test_fireworks_ai_common_utils.py +++ b/tests/unit/llms/fireworks_ai/test_fireworks_ai_common_utils.py @@ -1,7 +1,5 @@ - import pytest - from litellm.llms.fireworks_ai.common_utils import resolve_fireworks_resource_name @@ -16,6 +14,10 @@ from litellm.llms.fireworks_ai.common_utils import resolve_fireworks_resource_na ("glm-4p6", "accounts/fireworks/models/glm-4p6"), ("fireworks_ai/glm-4p6", "accounts/fireworks/models/glm-4p6"), ("kimi-k2p6-fast", "accounts/fireworks/routers/kimi-k2p6-fast"), + ("firerouter", "accounts/fireworks/routers/firerouter"), + ("fireworks_ai/firerouter", "accounts/fireworks/routers/firerouter"), + ("firerouter/kimi-k3/deepseek-v4", "accounts/fireworks/routers/firerouter/kimi-k3/deepseek-v4"), + ("firerouter-v2", "accounts/fireworks/models/firerouter-v2"), ( "accounts/fireworks/routers/glm-latest", "accounts/fireworks/routers/glm-latest", diff --git a/tests/unit/llms/fireworks_ai/test_fireworks_ai_cost_calculator.py b/tests/unit/llms/fireworks_ai/test_fireworks_ai_cost_calculator.py index c6096ba2745..c3aec979aa6 100644 --- a/tests/unit/llms/fireworks_ai/test_fireworks_ai_cost_calculator.py +++ b/tests/unit/llms/fireworks_ai/test_fireworks_ai_cost_calculator.py @@ -1,4 +1,5 @@ import math +import re from collections.abc import Generator from datetime import datetime, timezone from typing import Final @@ -326,3 +327,71 @@ def test_an_entry_without_an_input_rate_gets_no_cache_read_fallback(): assert prompt_cost == 0 assert completion_cost == 200 * 2e-06 + + +ROUTED_MODEL: Final = next( + key + for key, info in litellm.model_cost.items() + if "/" not in key + and info.get("litellm_provider") == "anthropic" + and (info.get("input_cost_per_token") or 0) > 0 + and (info.get("output_cost_per_token") or 0) > 0 + and f"fireworks_ai/{key}" not in litellm.model_cost +) + + +@pytest.mark.parametrize("model", [ROUTED_MODEL, f"fireworks_ai/{ROUTED_MODEL}"]) +def test_a_model_routed_to_another_provider_is_billed_at_that_models_own_rates(model: str): + own_rates: Final = litellm.get_model_info(model=ROUTED_MODEL, custom_llm_provider="anthropic") + usage: Final = _usage(prompt_tokens=23, cached_tokens=0, completion_tokens=41) + + prompt_cost, completion_cost = cost_per_token(model=model, usage=usage) + + assert prompt_cost == pytest.approx(23 * own_rates["input_cost_per_token"]) + assert completion_cost == pytest.approx(41 * own_rates["output_cost_per_token"]) + assert prompt_cost > 0 and completion_cost > 0 + + +def test_an_unknown_fireworks_model_still_falls_back_to_the_parameter_size_bucket(): + prompt_cost, completion_cost = cost_per_token( + model="accounts/fireworks/models/not-in-the-map-13b", + usage=_usage(prompt_tokens=100, cached_tokens=0, completion_tokens=10), + ) + bucket_prompt_cost, bucket_completion_cost = cost_per_token( + model="fireworks-ai-4.1b-to-16b", usage=_usage(prompt_tokens=100, cached_tokens=0, completion_tokens=10) + ) + + assert (prompt_cost, completion_cost) == (bucket_prompt_cost, bucket_completion_cost) + assert prompt_cost > 0 + + +_TIERED_INPUT_PATTERN: Final = re.compile(r"^input_cost_per_token_above_(\d+)k_tokens$") + + +def _threshold_tokens(field: str) -> int: + match: Final = _TIERED_INPUT_PATTERN.match(field) + assert match is not None, field + return int(match.group(1)) * 1000 + + +def test_a_routed_xai_model_keeps_xais_inclusive_token_threshold(): + candidate: Final = next( + ( + (key, field) + for key, info in litellm.model_cost.items() + if info.get("litellm_provider") == "xai" and f"fireworks_ai/{key}" not in litellm.model_cost + for field in info + if _TIERED_INPUT_PATTERN.match(field) + ), + None, + ) + if candidate is None: + pytest.skip("cost map has no xai entry with a tiered input rate") + key, field = candidate + usage: Final = _usage(prompt_tokens=_threshold_tokens(field), cached_tokens=0, completion_tokens=10) + + routed_prompt_cost, routed_completion_cost = cost_per_token(model=f"fireworks_ai/{key}", usage=usage) + + assert (routed_prompt_cost, routed_completion_cost) == generic_cost_per_token( + model=key, usage=usage, custom_llm_provider="xai" + ) From 1f1817ed399cc2129d7503cd504c1b2acba48dbc Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 23 Sep 2026 21:00:33 -0700 Subject: [PATCH 005/120] fix(models): add azure gpt-4o-mini-transcribe and gpt-4o-mini-tts deprecation dates (#42873) Co-authored-by: kerry Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/model_prices_and_context_window_backup.json | 2 ++ model_prices_and_context_window.json | 2 ++ 2 files changed, 4 insertions(+) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 18c8578cec8..15807127f13 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -6054,6 +6054,7 @@ "supports_tool_choice": true }, "azure/gpt-4o-mini-transcribe": { + "deprecation_date": "2027-06-15", "input_cost_per_audio_token": 1.25e-06, "input_cost_per_token": 1.25e-06, "litellm_provider": "azure", @@ -6066,6 +6067,7 @@ ] }, "azure/gpt-4o-mini-tts": { + "deprecation_date": "2027-06-15", "input_cost_per_token": 2.5e-06, "litellm_provider": "azure", "mode": "audio_speech", diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 18c8578cec8..15807127f13 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -6054,6 +6054,7 @@ "supports_tool_choice": true }, "azure/gpt-4o-mini-transcribe": { + "deprecation_date": "2027-06-15", "input_cost_per_audio_token": 1.25e-06, "input_cost_per_token": 1.25e-06, "litellm_provider": "azure", @@ -6066,6 +6067,7 @@ ] }, "azure/gpt-4o-mini-tts": { + "deprecation_date": "2027-06-15", "input_cost_per_token": 2.5e-06, "litellm_provider": "azure", "mode": "audio_speech", From 871f562f7288a92d91e9e98eab40be69b9f306bb Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 23 Sep 2026 21:01:10 -0700 Subject: [PATCH 006/120] fix(models): add openai deprecation date for gpt-5-chat-latest and gpt-5-chat (#42872) Co-authored-by: kerry Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/model_prices_and_context_window_backup.json | 2 ++ model_prices_and_context_window.json | 2 ++ 2 files changed, 4 insertions(+) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 15807127f13..e3fda66723f 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -34151,6 +34151,7 @@ }, "gpt-5-chat": { "cache_read_input_token_cost": 1.25e-07, + "deprecation_date": "2026-07-23", "input_cost_per_token": 1.25e-06, "litellm_provider": "openai", "max_input_tokens": 128000, @@ -34186,6 +34187,7 @@ }, "gpt-5-chat-latest": { "cache_read_input_token_cost": 1.25e-07, + "deprecation_date": "2026-07-23", "input_cost_per_token": 1.25e-06, "litellm_provider": "openai", "max_input_tokens": 128000, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 15807127f13..e3fda66723f 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -34151,6 +34151,7 @@ }, "gpt-5-chat": { "cache_read_input_token_cost": 1.25e-07, + "deprecation_date": "2026-07-23", "input_cost_per_token": 1.25e-06, "litellm_provider": "openai", "max_input_tokens": 128000, @@ -34186,6 +34187,7 @@ }, "gpt-5-chat-latest": { "cache_read_input_token_cost": 1.25e-07, + "deprecation_date": "2026-07-23", "input_cost_per_token": 1.25e-06, "litellm_provider": "openai", "max_input_tokens": 128000, From 15a2bd8b28fe88440e5f4f2ff79723e8363f61e9 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 23 Sep 2026 21:07:44 -0700 Subject: [PATCH 007/120] fix(models): add fireworks 2026-09-25 deprecation dates for glm 5.2, kimi k2.6, kimi k2.7 code, deepseek v4 and muse glimmer serverless rows (#42874) Co-authored-by: kerry Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../model_prices_and_context_window_backup.json | 14 ++++++++++++++ model_prices_and_context_window.json | 14 ++++++++++++++ 2 files changed, 28 insertions(+) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index e3fda66723f..9298900f1bc 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -24588,6 +24588,7 @@ "fireworks_ai/accounts/fireworks/models/deepseek-v4-pro-0813": { "cache_read_input_token_cost": 4.4e-08, "cache_read_input_token_cost_priority": 5.5e-08, + "deprecation_date": "2026-09-25", "input_cost_per_token": 1.32e-06, "input_cost_per_token_priority": 1.65e-06, "litellm_provider": "fireworks_ai", @@ -24607,6 +24608,7 @@ "fireworks_ai/deepseek-v4-pro-0813": { "cache_read_input_token_cost": 4.4e-08, "cache_read_input_token_cost_priority": 5.5e-08, + "deprecation_date": "2026-09-25", "input_cost_per_token": 1.32e-06, "input_cost_per_token_priority": 1.65e-06, "litellm_provider": "fireworks_ai", @@ -24712,6 +24714,7 @@ "fireworks_ai/accounts/fireworks/models/glm-5p2": { "cache_read_input_token_cost": 1.4e-07, "cache_read_input_token_cost_priority": 1.75e-07, + "deprecation_date": "2026-09-25", "input_cost_per_token": 1.4e-06, "input_cost_per_token_priority": 1.75e-06, "litellm_provider": "fireworks_ai", @@ -24820,6 +24823,7 @@ "fireworks_ai/accounts/fireworks/models/kimi-k2p6": { "cache_read_input_token_cost": 1.6e-07, "cache_read_input_token_cost_priority": 2.2e-07, + "deprecation_date": "2026-09-25", "input_cost_per_token": 9.5e-07, "input_cost_per_token_priority": 1.5e-06, "litellm_provider": "fireworks_ai", @@ -24839,6 +24843,7 @@ "fireworks_ai/accounts/fireworks/models/kimi-k2p7-code": { "cache_read_input_token_cost": 1.9e-07, "cache_read_input_token_cost_priority": 2.85e-07, + "deprecation_date": "2026-09-25", "input_cost_per_token": 9.5e-07, "input_cost_per_token_priority": 1.425e-06, "litellm_provider": "fireworks_ai", @@ -25109,6 +25114,7 @@ "fireworks_ai/glm-5p2": { "cache_read_input_token_cost": 1.4e-07, "cache_read_input_token_cost_priority": 1.75e-07, + "deprecation_date": "2026-09-25", "input_cost_per_token": 1.4e-06, "input_cost_per_token_priority": 1.75e-06, "litellm_provider": "fireworks_ai", @@ -25177,6 +25183,7 @@ "fireworks_ai/kimi-k2p6": { "cache_read_input_token_cost": 1.6e-07, "cache_read_input_token_cost_priority": 2.2e-07, + "deprecation_date": "2026-09-25", "input_cost_per_token": 9.5e-07, "input_cost_per_token_priority": 1.5e-06, "litellm_provider": "fireworks_ai", @@ -25213,6 +25220,7 @@ "fireworks_ai/kimi-k2p7-code": { "cache_read_input_token_cost": 1.9e-07, "cache_read_input_token_cost_priority": 2.85e-07, + "deprecation_date": "2026-09-25", "input_cost_per_token": 9.5e-07, "input_cost_per_token_priority": 1.425e-06, "litellm_provider": "fireworks_ai", @@ -59357,6 +59365,7 @@ "fireworks_ai/accounts/fireworks/models/deepseek-v4-flash-0731": { "cache_read_input_token_cost": 7e-09, "cache_read_input_token_cost_priority": 8.75e-09, + "deprecation_date": "2026-09-25", "input_cost_per_token": 2.2e-07, "input_cost_per_token_priority": 2.75e-07, "litellm_provider": "fireworks_ai", @@ -59395,6 +59404,7 @@ }, "fireworks_ai/accounts/fireworks/models/deepseek-v4-flash-vision-exp": { "cache_read_input_token_cost": 7e-09, + "deprecation_date": "2026-09-25", "input_cost_per_token": 2.2e-07, "litellm_provider": "fireworks_ai", "max_input_tokens": 1048576, @@ -59434,6 +59444,7 @@ "fireworks_ai/deepseek-v4-flash-0731": { "cache_read_input_token_cost": 7e-09, "cache_read_input_token_cost_priority": 8.75e-09, + "deprecation_date": "2026-09-25", "input_cost_per_token": 2.2e-07, "input_cost_per_token_priority": 2.75e-07, "litellm_provider": "fireworks_ai", @@ -59472,6 +59483,7 @@ }, "fireworks_ai/deepseek-v4-flash-vision-exp": { "cache_read_input_token_cost": 7e-09, + "deprecation_date": "2026-09-25", "input_cost_per_token": 2.2e-07, "litellm_provider": "fireworks_ai", "max_input_tokens": 1048576, @@ -59603,6 +59615,7 @@ }, "fireworks_ai/muse-glimmer-30b": { "cache_read_input_token_cost": 4e-08, + "deprecation_date": "2026-09-25", "input_cost_per_token": 3.5e-07, "litellm_provider": "fireworks_ai", "max_input_tokens": 131072, @@ -59651,6 +59664,7 @@ }, "fireworks_ai/accounts/fireworks/models/muse-glimmer-30b": { "cache_read_input_token_cost": 4e-08, + "deprecation_date": "2026-09-25", "input_cost_per_token": 3.5e-07, "litellm_provider": "fireworks_ai", "max_input_tokens": 131072, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index e3fda66723f..9298900f1bc 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -24588,6 +24588,7 @@ "fireworks_ai/accounts/fireworks/models/deepseek-v4-pro-0813": { "cache_read_input_token_cost": 4.4e-08, "cache_read_input_token_cost_priority": 5.5e-08, + "deprecation_date": "2026-09-25", "input_cost_per_token": 1.32e-06, "input_cost_per_token_priority": 1.65e-06, "litellm_provider": "fireworks_ai", @@ -24607,6 +24608,7 @@ "fireworks_ai/deepseek-v4-pro-0813": { "cache_read_input_token_cost": 4.4e-08, "cache_read_input_token_cost_priority": 5.5e-08, + "deprecation_date": "2026-09-25", "input_cost_per_token": 1.32e-06, "input_cost_per_token_priority": 1.65e-06, "litellm_provider": "fireworks_ai", @@ -24712,6 +24714,7 @@ "fireworks_ai/accounts/fireworks/models/glm-5p2": { "cache_read_input_token_cost": 1.4e-07, "cache_read_input_token_cost_priority": 1.75e-07, + "deprecation_date": "2026-09-25", "input_cost_per_token": 1.4e-06, "input_cost_per_token_priority": 1.75e-06, "litellm_provider": "fireworks_ai", @@ -24820,6 +24823,7 @@ "fireworks_ai/accounts/fireworks/models/kimi-k2p6": { "cache_read_input_token_cost": 1.6e-07, "cache_read_input_token_cost_priority": 2.2e-07, + "deprecation_date": "2026-09-25", "input_cost_per_token": 9.5e-07, "input_cost_per_token_priority": 1.5e-06, "litellm_provider": "fireworks_ai", @@ -24839,6 +24843,7 @@ "fireworks_ai/accounts/fireworks/models/kimi-k2p7-code": { "cache_read_input_token_cost": 1.9e-07, "cache_read_input_token_cost_priority": 2.85e-07, + "deprecation_date": "2026-09-25", "input_cost_per_token": 9.5e-07, "input_cost_per_token_priority": 1.425e-06, "litellm_provider": "fireworks_ai", @@ -25109,6 +25114,7 @@ "fireworks_ai/glm-5p2": { "cache_read_input_token_cost": 1.4e-07, "cache_read_input_token_cost_priority": 1.75e-07, + "deprecation_date": "2026-09-25", "input_cost_per_token": 1.4e-06, "input_cost_per_token_priority": 1.75e-06, "litellm_provider": "fireworks_ai", @@ -25177,6 +25183,7 @@ "fireworks_ai/kimi-k2p6": { "cache_read_input_token_cost": 1.6e-07, "cache_read_input_token_cost_priority": 2.2e-07, + "deprecation_date": "2026-09-25", "input_cost_per_token": 9.5e-07, "input_cost_per_token_priority": 1.5e-06, "litellm_provider": "fireworks_ai", @@ -25213,6 +25220,7 @@ "fireworks_ai/kimi-k2p7-code": { "cache_read_input_token_cost": 1.9e-07, "cache_read_input_token_cost_priority": 2.85e-07, + "deprecation_date": "2026-09-25", "input_cost_per_token": 9.5e-07, "input_cost_per_token_priority": 1.425e-06, "litellm_provider": "fireworks_ai", @@ -59357,6 +59365,7 @@ "fireworks_ai/accounts/fireworks/models/deepseek-v4-flash-0731": { "cache_read_input_token_cost": 7e-09, "cache_read_input_token_cost_priority": 8.75e-09, + "deprecation_date": "2026-09-25", "input_cost_per_token": 2.2e-07, "input_cost_per_token_priority": 2.75e-07, "litellm_provider": "fireworks_ai", @@ -59395,6 +59404,7 @@ }, "fireworks_ai/accounts/fireworks/models/deepseek-v4-flash-vision-exp": { "cache_read_input_token_cost": 7e-09, + "deprecation_date": "2026-09-25", "input_cost_per_token": 2.2e-07, "litellm_provider": "fireworks_ai", "max_input_tokens": 1048576, @@ -59434,6 +59444,7 @@ "fireworks_ai/deepseek-v4-flash-0731": { "cache_read_input_token_cost": 7e-09, "cache_read_input_token_cost_priority": 8.75e-09, + "deprecation_date": "2026-09-25", "input_cost_per_token": 2.2e-07, "input_cost_per_token_priority": 2.75e-07, "litellm_provider": "fireworks_ai", @@ -59472,6 +59483,7 @@ }, "fireworks_ai/deepseek-v4-flash-vision-exp": { "cache_read_input_token_cost": 7e-09, + "deprecation_date": "2026-09-25", "input_cost_per_token": 2.2e-07, "litellm_provider": "fireworks_ai", "max_input_tokens": 1048576, @@ -59603,6 +59615,7 @@ }, "fireworks_ai/muse-glimmer-30b": { "cache_read_input_token_cost": 4e-08, + "deprecation_date": "2026-09-25", "input_cost_per_token": 3.5e-07, "litellm_provider": "fireworks_ai", "max_input_tokens": 131072, @@ -59651,6 +59664,7 @@ }, "fireworks_ai/accounts/fireworks/models/muse-glimmer-30b": { "cache_read_input_token_cost": 4e-08, + "deprecation_date": "2026-09-25", "input_cost_per_token": 3.5e-07, "litellm_provider": "fireworks_ai", "max_input_tokens": 131072, From 7370650d91bea7ecdd04dd7350a6cf4229441048 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 23 Sep 2026 21:27:25 -0700 Subject: [PATCH 008/120] fix(model_prices): bedrock bare Claude ids priced at the Global SKU (aws-bedrock sync) (#42875) Co-authored-by: kerry Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- ...odel_prices_and_context_window_backup.json | 162 +++++++++--------- model_prices_and_context_window.json | 162 +++++++++--------- tests/test_litellm/test_utils.py | 10 +- 3 files changed, 167 insertions(+), 167 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 9298900f1bc..059c92f4b4b 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -812,18 +812,18 @@ "prompt_cache_min_tokens": 2048 }, "anthropic.claude-haiku-4-5-20251001-v1:0": { - "cache_creation_input_token_cost": 1.375e-06, - "cache_creation_input_token_cost_above_1hr": 2.2e-06, - "cache_read_input_token_cost": 1.1e-07, - "input_cost_per_token": 1.1e-06, + "cache_creation_input_token_cost": 1.25e-06, + "cache_creation_input_token_cost_above_1hr": 2e-06, + "cache_read_input_token_cost": 1e-07, + "input_cost_per_token": 1e-06, "litellm_provider": "bedrock_converse", "supports_tool_search": true, "max_input_tokens": 200000, "max_output_tokens": 64000, "max_tokens": 64000, "mode": "chat", - "output_cost_per_token": 5.5e-06, - "source": "https://aws.amazon.com/bedrock/pricing/", + "output_cost_per_token": 5e-06, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json", "supports_assistant_prefill": true, "supports_computer_use": true, "supports_function_calling": true, @@ -836,8 +836,8 @@ "supports_native_structured_output": true, "supports_parallel_tool_use_config": true, "prompt_cache_min_tokens": 4096, - "input_cost_per_token_batches": 5.5e-07, - "output_cost_per_token_batches": 2.75e-06 + "input_cost_per_token_batches": 5e-07, + "output_cost_per_token_batches": 2.5e-06 }, "anthropic.claude-haiku-4-5@20251001": { "cache_creation_input_token_cost": 1.25e-06, @@ -1028,17 +1028,17 @@ "prompt_cache_min_tokens": 1024 }, "anthropic.claude-opus-4-5-20251101-v1:0": { - "cache_creation_input_token_cost": 6.875e-06, - "cache_creation_input_token_cost_above_1hr": 1.1e-05, - "cache_read_input_token_cost": 5.5e-07, - "input_cost_per_token": 5.5e-06, + "cache_creation_input_token_cost": 6.25e-06, + "cache_creation_input_token_cost_above_1hr": 1e-05, + "cache_read_input_token_cost": 5e-07, + "input_cost_per_token": 5e-06, "litellm_provider": "bedrock_converse", "supports_tool_search": true, "max_input_tokens": 200000, "max_output_tokens": 64000, "max_tokens": 64000, "mode": "chat", - "output_cost_per_token": 2.75e-05, + "output_cost_per_token": 2.5e-05, "search_context_cost_per_query": { "search_context_size_high": 0.01, "search_context_size_low": 0.01, @@ -1058,22 +1058,22 @@ "bedrock_output_config_effort_ceiling": "high", "supports_parallel_tool_use_config": true, "prompt_cache_min_tokens": 4096, - "source": "https://aws.amazon.com/bedrock/pricing/" + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json" }, "anthropic.claude-opus-4-6-v1": { "supports_adaptive_thinking": true, "supports_legacy_thinking": true, - "cache_creation_input_token_cost": 6.875e-06, - "cache_creation_input_token_cost_above_1hr": 1.1e-05, - "cache_read_input_token_cost": 5.5e-07, - "input_cost_per_token": 5.5e-06, + "cache_creation_input_token_cost": 6.25e-06, + "cache_creation_input_token_cost_above_1hr": 1e-05, + "cache_read_input_token_cost": 5e-07, + "input_cost_per_token": 5e-06, "litellm_provider": "bedrock_converse", "supports_tool_search": true, "max_input_tokens": 1000000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", - "output_cost_per_token": 2.75e-05, + "output_cost_per_token": 2.5e-05, "search_context_cost_per_query": { "search_context_size_high": 0.01, "search_context_size_low": 0.01, @@ -1094,7 +1094,7 @@ "bedrock_output_config_effort_ceiling": "max", "supports_parallel_tool_use_config": true, "prompt_cache_min_tokens": 4096, - "source": "https://aws.amazon.com/bedrock/pricing/" + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json" }, "global.anthropic.claude-opus-4-6-v1": { "supports_adaptive_thinking": true, @@ -1243,17 +1243,17 @@ "anthropic.claude-opus-4-7": { "bedrock_converse_supports_strict_tools": false, "supports_adaptive_thinking": true, - "cache_creation_input_token_cost": 6.875e-06, - "cache_creation_input_token_cost_above_1hr": 1.1e-05, - "cache_read_input_token_cost": 5.5e-07, - "input_cost_per_token": 5.5e-06, + "cache_creation_input_token_cost": 6.25e-06, + "cache_creation_input_token_cost_above_1hr": 1e-05, + "cache_read_input_token_cost": 5e-07, + "input_cost_per_token": 5e-06, "litellm_provider": "bedrock_converse", "supports_tool_search": true, "max_input_tokens": 1000000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", - "output_cost_per_token": 2.75e-05, + "output_cost_per_token": 2.5e-05, "search_context_cost_per_query": { "search_context_size_high": 0.01, "search_context_size_low": 0.01, @@ -1276,7 +1276,7 @@ "bedrock_output_config_effort_ceiling": "xhigh", "supports_parallel_tool_use_config": true, "prompt_cache_min_tokens": 2048, - "source": "https://aws.amazon.com/bedrock/pricing/" + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json" }, "anthropic.claude-mythos-preview": { "input_cost_per_token": 0, @@ -1447,16 +1447,16 @@ "source": "https://aws.amazon.com/bedrock/pricing/" }, "anthropic.claude-fable-5": { - "cache_creation_input_token_cost": 1.375e-05, - "cache_creation_input_token_cost_above_1hr": 2.2e-05, - "cache_read_input_token_cost": 1.1e-06, - "input_cost_per_token": 1.1e-05, + "cache_creation_input_token_cost": 1.25e-05, + "cache_creation_input_token_cost_above_1hr": 2e-05, + "cache_read_input_token_cost": 1e-06, + "input_cost_per_token": 1e-05, "litellm_provider": "bedrock_converse", "max_input_tokens": 1000000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", - "output_cost_per_token": 5.5e-05, + "output_cost_per_token": 5e-05, "search_context_cost_per_query": { "search_context_size_high": 0.01, "search_context_size_low": 0.01, @@ -1482,19 +1482,19 @@ "bedrock_output_config_effort_ceiling": "xhigh", "supports_parallel_tool_use_config": true, "prompt_cache_min_tokens": 512, - "source": "https://aws.amazon.com/bedrock/pricing/" + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json" }, "anthropic.claude-fable-5-1": { - "cache_creation_input_token_cost": 1.375e-05, - "cache_creation_input_token_cost_above_1hr": 2.2e-05, - "cache_read_input_token_cost": 2.75e-07, - "input_cost_per_token": 1.1e-05, + "cache_creation_input_token_cost": 1.25e-05, + "cache_creation_input_token_cost_above_1hr": 2e-05, + "cache_read_input_token_cost": 2.5e-07, + "input_cost_per_token": 1e-05, "litellm_provider": "bedrock_converse", "max_input_tokens": 1000000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", - "output_cost_per_token": 5.5e-05, + "output_cost_per_token": 5e-05, "search_context_cost_per_query": { "search_context_size_high": 0.01, "search_context_size_low": 0.01, @@ -1521,7 +1521,7 @@ "bedrock_output_config_effort_ceiling": "xhigh", "supports_parallel_tool_use_config": true, "prompt_cache_min_tokens": 512, - "source": "https://aws.amazon.com/bedrock/pricing/" + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json" }, "global.anthropic.claude-fable-5": { "cache_creation_input_token_cost": 1.25e-05, @@ -1757,17 +1757,17 @@ "bedrock_converse_supports_strict_tools": false, "supports_adaptive_thinking": true, "supports_mid_conversation_system": true, - "cache_creation_input_token_cost": 6.875e-06, - "cache_creation_input_token_cost_above_1hr": 1.1e-05, - "cache_read_input_token_cost": 5.5e-07, - "input_cost_per_token": 5.5e-06, + "cache_creation_input_token_cost": 6.25e-06, + "cache_creation_input_token_cost_above_1hr": 1e-05, + "cache_read_input_token_cost": 5e-07, + "input_cost_per_token": 5e-06, "litellm_provider": "bedrock_converse", "supports_tool_search": true, "max_input_tokens": 1000000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", - "output_cost_per_token": 2.75e-05, + "output_cost_per_token": 2.5e-05, "search_context_cost_per_query": { "search_context_size_high": 0.01, "search_context_size_low": 0.01, @@ -1789,23 +1789,23 @@ "supports_output_config": true, "supports_parallel_tool_use_config": true, "prompt_cache_min_tokens": 512, - "source": "https://aws.amazon.com/bedrock/pricing/" + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json" }, "anthropic.claude-opus-5-5": { "bedrock_converse_supports_strict_tools": false, "supports_adaptive_thinking": true, "supports_mid_conversation_system": true, - "cache_creation_input_token_cost": 5.5e-06, - "cache_creation_input_token_cost_above_1hr": 8.8e-06, - "cache_read_input_token_cost": 2.2e-07, - "input_cost_per_token": 4.4e-06, + "cache_creation_input_token_cost": 5e-06, + "cache_creation_input_token_cost_above_1hr": 8e-06, + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 4e-06, "litellm_provider": "bedrock_converse", "supports_tool_search": true, "max_input_tokens": 1000000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", - "output_cost_per_token": 2.2e-05, + "output_cost_per_token": 2e-05, "search_context_cost_per_query": { "search_context_size_high": 0.01, "search_context_size_low": 0.01, @@ -1827,7 +1827,7 @@ "supports_output_config": true, "supports_parallel_tool_use_config": true, "prompt_cache_min_tokens": 512, - "source": "https://aws.amazon.com/bedrock/pricing/", + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json", "thinking_always_on": true, "supports_forced_tool_use": false }, @@ -2225,17 +2225,17 @@ "bedrock_converse_supports_strict_tools": false, "supports_adaptive_thinking": true, "supports_mid_conversation_system": true, - "cache_creation_input_token_cost": 6.875e-06, - "cache_creation_input_token_cost_above_1hr": 1.1e-05, - "cache_read_input_token_cost": 5.5e-07, - "input_cost_per_token": 5.5e-06, + "cache_creation_input_token_cost": 6.25e-06, + "cache_creation_input_token_cost_above_1hr": 1e-05, + "cache_read_input_token_cost": 5e-07, + "input_cost_per_token": 5e-06, "litellm_provider": "bedrock_converse", "supports_tool_search": true, "max_input_tokens": 1000000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", - "output_cost_per_token": 2.75e-05, + "output_cost_per_token": 2.5e-05, "search_context_cost_per_query": { "search_context_size_high": 0.01, "search_context_size_low": 0.01, @@ -2258,7 +2258,7 @@ "bedrock_output_config_effort_ceiling": "xhigh", "supports_parallel_tool_use_config": true, "prompt_cache_min_tokens": 1024, - "source": "https://aws.amazon.com/bedrock/pricing/" + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json" }, "global.anthropic.claude-opus-4-8": { "bedrock_converse_supports_strict_tools": false, @@ -2494,17 +2494,17 @@ }, "anthropic.claude-sonnet-5": { "bedrock_converse_supports_strict_tools": false, - "cache_creation_input_token_cost": 2.75e-06, - "cache_creation_input_token_cost_above_1hr": 4.4e-06, - "cache_read_input_token_cost": 2.2e-07, - "input_cost_per_token": 2.2e-06, + "cache_creation_input_token_cost": 2.5e-06, + "cache_creation_input_token_cost_above_1hr": 4e-06, + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 2e-06, "litellm_provider": "bedrock_converse", "supports_tool_search": true, "max_input_tokens": 1000000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", - "output_cost_per_token": 1.1e-05, + "output_cost_per_token": 1e-05, "search_context_cost_per_query": { "search_context_size_high": 0.01, "search_context_size_low": 0.01, @@ -2529,7 +2529,7 @@ "bedrock_output_config_effort_ceiling": "xhigh", "supports_parallel_tool_use_config": true, "prompt_cache_min_tokens": 1024, - "source": "https://aws.amazon.com/bedrock/pricing/" + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json" }, "global.anthropic.claude-sonnet-5": { "bedrock_converse_supports_strict_tools": false, @@ -2729,17 +2729,17 @@ "anthropic.claude-sonnet-4-6": { "supports_adaptive_thinking": true, "supports_legacy_thinking": true, - "cache_creation_input_token_cost": 4.125e-06, - "cache_creation_input_token_cost_above_1hr": 6.6e-06, - "cache_read_input_token_cost": 3.3e-07, - "input_cost_per_token": 3.3e-06, + "cache_creation_input_token_cost": 3.75e-06, + "cache_creation_input_token_cost_above_1hr": 6e-06, + "cache_read_input_token_cost": 3e-07, + "input_cost_per_token": 3e-06, "litellm_provider": "bedrock_converse", "supports_tool_search": true, "max_input_tokens": 1000000, "max_output_tokens": 64000, "max_tokens": 64000, "mode": "chat", - "output_cost_per_token": 1.65e-05, + "output_cost_per_token": 1.5e-05, "search_context_cost_per_query": { "search_context_size_high": 0.01, "search_context_size_low": 0.01, @@ -2759,7 +2759,7 @@ "supports_output_config": true, "supports_parallel_tool_use_config": true, "prompt_cache_min_tokens": 1024, - "source": "https://aws.amazon.com/bedrock/pricing/" + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json" }, "global.anthropic.claude-sonnet-4-6": { "supports_adaptive_thinking": true, @@ -2970,22 +2970,22 @@ "source": "https://aws.amazon.com/bedrock/pricing/" }, "anthropic.claude-sonnet-4-5-20250929-v1:0": { - "cache_creation_input_token_cost": 4.125e-06, - "cache_creation_input_token_cost_above_1hr": 6.6e-06, - "cache_read_input_token_cost": 3.3e-07, - "input_cost_per_token": 3.3e-06, - "input_cost_per_token_above_200k_tokens": 6.6e-06, - "output_cost_per_token_above_200k_tokens": 2.475e-05, - "cache_creation_input_token_cost_above_200k_tokens": 8.25e-06, - "cache_creation_input_token_cost_above_1hr_above_200k_tokens": 1.32e-05, - "cache_read_input_token_cost_above_200k_tokens": 6.6e-07, + "cache_creation_input_token_cost": 3.75e-06, + "cache_creation_input_token_cost_above_1hr": 6e-06, + "cache_read_input_token_cost": 3e-07, + "input_cost_per_token": 3e-06, + "input_cost_per_token_above_200k_tokens": 6e-06, + "output_cost_per_token_above_200k_tokens": 2.25e-05, + "cache_creation_input_token_cost_above_200k_tokens": 7.5e-06, + "cache_creation_input_token_cost_above_1hr_above_200k_tokens": 1.2e-05, + "cache_read_input_token_cost_above_200k_tokens": 6e-07, "litellm_provider": "bedrock_converse", "supports_tool_search": true, "max_input_tokens": 200000, "max_output_tokens": 64000, "max_tokens": 64000, "mode": "chat", - "output_cost_per_token": 1.65e-05, + "output_cost_per_token": 1.5e-05, "search_context_cost_per_query": { "search_context_size_high": 0.01, "search_context_size_low": 0.01, @@ -3003,9 +3003,9 @@ "supports_native_structured_output": true, "supports_parallel_tool_use_config": true, "prompt_cache_min_tokens": 1024, - "input_cost_per_token_batches": 1.65e-06, - "output_cost_per_token_batches": 8.25e-06, - "source": "https://aws.amazon.com/bedrock/pricing/" + "input_cost_per_token_batches": 1.5e-06, + "output_cost_per_token_batches": 7.5e-06, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json" }, "anthropic.claude-v1": { "input_cost_per_token": 8e-06, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 9298900f1bc..059c92f4b4b 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -812,18 +812,18 @@ "prompt_cache_min_tokens": 2048 }, "anthropic.claude-haiku-4-5-20251001-v1:0": { - "cache_creation_input_token_cost": 1.375e-06, - "cache_creation_input_token_cost_above_1hr": 2.2e-06, - "cache_read_input_token_cost": 1.1e-07, - "input_cost_per_token": 1.1e-06, + "cache_creation_input_token_cost": 1.25e-06, + "cache_creation_input_token_cost_above_1hr": 2e-06, + "cache_read_input_token_cost": 1e-07, + "input_cost_per_token": 1e-06, "litellm_provider": "bedrock_converse", "supports_tool_search": true, "max_input_tokens": 200000, "max_output_tokens": 64000, "max_tokens": 64000, "mode": "chat", - "output_cost_per_token": 5.5e-06, - "source": "https://aws.amazon.com/bedrock/pricing/", + "output_cost_per_token": 5e-06, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json", "supports_assistant_prefill": true, "supports_computer_use": true, "supports_function_calling": true, @@ -836,8 +836,8 @@ "supports_native_structured_output": true, "supports_parallel_tool_use_config": true, "prompt_cache_min_tokens": 4096, - "input_cost_per_token_batches": 5.5e-07, - "output_cost_per_token_batches": 2.75e-06 + "input_cost_per_token_batches": 5e-07, + "output_cost_per_token_batches": 2.5e-06 }, "anthropic.claude-haiku-4-5@20251001": { "cache_creation_input_token_cost": 1.25e-06, @@ -1028,17 +1028,17 @@ "prompt_cache_min_tokens": 1024 }, "anthropic.claude-opus-4-5-20251101-v1:0": { - "cache_creation_input_token_cost": 6.875e-06, - "cache_creation_input_token_cost_above_1hr": 1.1e-05, - "cache_read_input_token_cost": 5.5e-07, - "input_cost_per_token": 5.5e-06, + "cache_creation_input_token_cost": 6.25e-06, + "cache_creation_input_token_cost_above_1hr": 1e-05, + "cache_read_input_token_cost": 5e-07, + "input_cost_per_token": 5e-06, "litellm_provider": "bedrock_converse", "supports_tool_search": true, "max_input_tokens": 200000, "max_output_tokens": 64000, "max_tokens": 64000, "mode": "chat", - "output_cost_per_token": 2.75e-05, + "output_cost_per_token": 2.5e-05, "search_context_cost_per_query": { "search_context_size_high": 0.01, "search_context_size_low": 0.01, @@ -1058,22 +1058,22 @@ "bedrock_output_config_effort_ceiling": "high", "supports_parallel_tool_use_config": true, "prompt_cache_min_tokens": 4096, - "source": "https://aws.amazon.com/bedrock/pricing/" + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json" }, "anthropic.claude-opus-4-6-v1": { "supports_adaptive_thinking": true, "supports_legacy_thinking": true, - "cache_creation_input_token_cost": 6.875e-06, - "cache_creation_input_token_cost_above_1hr": 1.1e-05, - "cache_read_input_token_cost": 5.5e-07, - "input_cost_per_token": 5.5e-06, + "cache_creation_input_token_cost": 6.25e-06, + "cache_creation_input_token_cost_above_1hr": 1e-05, + "cache_read_input_token_cost": 5e-07, + "input_cost_per_token": 5e-06, "litellm_provider": "bedrock_converse", "supports_tool_search": true, "max_input_tokens": 1000000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", - "output_cost_per_token": 2.75e-05, + "output_cost_per_token": 2.5e-05, "search_context_cost_per_query": { "search_context_size_high": 0.01, "search_context_size_low": 0.01, @@ -1094,7 +1094,7 @@ "bedrock_output_config_effort_ceiling": "max", "supports_parallel_tool_use_config": true, "prompt_cache_min_tokens": 4096, - "source": "https://aws.amazon.com/bedrock/pricing/" + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json" }, "global.anthropic.claude-opus-4-6-v1": { "supports_adaptive_thinking": true, @@ -1243,17 +1243,17 @@ "anthropic.claude-opus-4-7": { "bedrock_converse_supports_strict_tools": false, "supports_adaptive_thinking": true, - "cache_creation_input_token_cost": 6.875e-06, - "cache_creation_input_token_cost_above_1hr": 1.1e-05, - "cache_read_input_token_cost": 5.5e-07, - "input_cost_per_token": 5.5e-06, + "cache_creation_input_token_cost": 6.25e-06, + "cache_creation_input_token_cost_above_1hr": 1e-05, + "cache_read_input_token_cost": 5e-07, + "input_cost_per_token": 5e-06, "litellm_provider": "bedrock_converse", "supports_tool_search": true, "max_input_tokens": 1000000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", - "output_cost_per_token": 2.75e-05, + "output_cost_per_token": 2.5e-05, "search_context_cost_per_query": { "search_context_size_high": 0.01, "search_context_size_low": 0.01, @@ -1276,7 +1276,7 @@ "bedrock_output_config_effort_ceiling": "xhigh", "supports_parallel_tool_use_config": true, "prompt_cache_min_tokens": 2048, - "source": "https://aws.amazon.com/bedrock/pricing/" + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json" }, "anthropic.claude-mythos-preview": { "input_cost_per_token": 0, @@ -1447,16 +1447,16 @@ "source": "https://aws.amazon.com/bedrock/pricing/" }, "anthropic.claude-fable-5": { - "cache_creation_input_token_cost": 1.375e-05, - "cache_creation_input_token_cost_above_1hr": 2.2e-05, - "cache_read_input_token_cost": 1.1e-06, - "input_cost_per_token": 1.1e-05, + "cache_creation_input_token_cost": 1.25e-05, + "cache_creation_input_token_cost_above_1hr": 2e-05, + "cache_read_input_token_cost": 1e-06, + "input_cost_per_token": 1e-05, "litellm_provider": "bedrock_converse", "max_input_tokens": 1000000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", - "output_cost_per_token": 5.5e-05, + "output_cost_per_token": 5e-05, "search_context_cost_per_query": { "search_context_size_high": 0.01, "search_context_size_low": 0.01, @@ -1482,19 +1482,19 @@ "bedrock_output_config_effort_ceiling": "xhigh", "supports_parallel_tool_use_config": true, "prompt_cache_min_tokens": 512, - "source": "https://aws.amazon.com/bedrock/pricing/" + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json" }, "anthropic.claude-fable-5-1": { - "cache_creation_input_token_cost": 1.375e-05, - "cache_creation_input_token_cost_above_1hr": 2.2e-05, - "cache_read_input_token_cost": 2.75e-07, - "input_cost_per_token": 1.1e-05, + "cache_creation_input_token_cost": 1.25e-05, + "cache_creation_input_token_cost_above_1hr": 2e-05, + "cache_read_input_token_cost": 2.5e-07, + "input_cost_per_token": 1e-05, "litellm_provider": "bedrock_converse", "max_input_tokens": 1000000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", - "output_cost_per_token": 5.5e-05, + "output_cost_per_token": 5e-05, "search_context_cost_per_query": { "search_context_size_high": 0.01, "search_context_size_low": 0.01, @@ -1521,7 +1521,7 @@ "bedrock_output_config_effort_ceiling": "xhigh", "supports_parallel_tool_use_config": true, "prompt_cache_min_tokens": 512, - "source": "https://aws.amazon.com/bedrock/pricing/" + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json" }, "global.anthropic.claude-fable-5": { "cache_creation_input_token_cost": 1.25e-05, @@ -1757,17 +1757,17 @@ "bedrock_converse_supports_strict_tools": false, "supports_adaptive_thinking": true, "supports_mid_conversation_system": true, - "cache_creation_input_token_cost": 6.875e-06, - "cache_creation_input_token_cost_above_1hr": 1.1e-05, - "cache_read_input_token_cost": 5.5e-07, - "input_cost_per_token": 5.5e-06, + "cache_creation_input_token_cost": 6.25e-06, + "cache_creation_input_token_cost_above_1hr": 1e-05, + "cache_read_input_token_cost": 5e-07, + "input_cost_per_token": 5e-06, "litellm_provider": "bedrock_converse", "supports_tool_search": true, "max_input_tokens": 1000000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", - "output_cost_per_token": 2.75e-05, + "output_cost_per_token": 2.5e-05, "search_context_cost_per_query": { "search_context_size_high": 0.01, "search_context_size_low": 0.01, @@ -1789,23 +1789,23 @@ "supports_output_config": true, "supports_parallel_tool_use_config": true, "prompt_cache_min_tokens": 512, - "source": "https://aws.amazon.com/bedrock/pricing/" + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json" }, "anthropic.claude-opus-5-5": { "bedrock_converse_supports_strict_tools": false, "supports_adaptive_thinking": true, "supports_mid_conversation_system": true, - "cache_creation_input_token_cost": 5.5e-06, - "cache_creation_input_token_cost_above_1hr": 8.8e-06, - "cache_read_input_token_cost": 2.2e-07, - "input_cost_per_token": 4.4e-06, + "cache_creation_input_token_cost": 5e-06, + "cache_creation_input_token_cost_above_1hr": 8e-06, + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 4e-06, "litellm_provider": "bedrock_converse", "supports_tool_search": true, "max_input_tokens": 1000000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", - "output_cost_per_token": 2.2e-05, + "output_cost_per_token": 2e-05, "search_context_cost_per_query": { "search_context_size_high": 0.01, "search_context_size_low": 0.01, @@ -1827,7 +1827,7 @@ "supports_output_config": true, "supports_parallel_tool_use_config": true, "prompt_cache_min_tokens": 512, - "source": "https://aws.amazon.com/bedrock/pricing/", + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json", "thinking_always_on": true, "supports_forced_tool_use": false }, @@ -2225,17 +2225,17 @@ "bedrock_converse_supports_strict_tools": false, "supports_adaptive_thinking": true, "supports_mid_conversation_system": true, - "cache_creation_input_token_cost": 6.875e-06, - "cache_creation_input_token_cost_above_1hr": 1.1e-05, - "cache_read_input_token_cost": 5.5e-07, - "input_cost_per_token": 5.5e-06, + "cache_creation_input_token_cost": 6.25e-06, + "cache_creation_input_token_cost_above_1hr": 1e-05, + "cache_read_input_token_cost": 5e-07, + "input_cost_per_token": 5e-06, "litellm_provider": "bedrock_converse", "supports_tool_search": true, "max_input_tokens": 1000000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", - "output_cost_per_token": 2.75e-05, + "output_cost_per_token": 2.5e-05, "search_context_cost_per_query": { "search_context_size_high": 0.01, "search_context_size_low": 0.01, @@ -2258,7 +2258,7 @@ "bedrock_output_config_effort_ceiling": "xhigh", "supports_parallel_tool_use_config": true, "prompt_cache_min_tokens": 1024, - "source": "https://aws.amazon.com/bedrock/pricing/" + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json" }, "global.anthropic.claude-opus-4-8": { "bedrock_converse_supports_strict_tools": false, @@ -2494,17 +2494,17 @@ }, "anthropic.claude-sonnet-5": { "bedrock_converse_supports_strict_tools": false, - "cache_creation_input_token_cost": 2.75e-06, - "cache_creation_input_token_cost_above_1hr": 4.4e-06, - "cache_read_input_token_cost": 2.2e-07, - "input_cost_per_token": 2.2e-06, + "cache_creation_input_token_cost": 2.5e-06, + "cache_creation_input_token_cost_above_1hr": 4e-06, + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 2e-06, "litellm_provider": "bedrock_converse", "supports_tool_search": true, "max_input_tokens": 1000000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", - "output_cost_per_token": 1.1e-05, + "output_cost_per_token": 1e-05, "search_context_cost_per_query": { "search_context_size_high": 0.01, "search_context_size_low": 0.01, @@ -2529,7 +2529,7 @@ "bedrock_output_config_effort_ceiling": "xhigh", "supports_parallel_tool_use_config": true, "prompt_cache_min_tokens": 1024, - "source": "https://aws.amazon.com/bedrock/pricing/" + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json" }, "global.anthropic.claude-sonnet-5": { "bedrock_converse_supports_strict_tools": false, @@ -2729,17 +2729,17 @@ "anthropic.claude-sonnet-4-6": { "supports_adaptive_thinking": true, "supports_legacy_thinking": true, - "cache_creation_input_token_cost": 4.125e-06, - "cache_creation_input_token_cost_above_1hr": 6.6e-06, - "cache_read_input_token_cost": 3.3e-07, - "input_cost_per_token": 3.3e-06, + "cache_creation_input_token_cost": 3.75e-06, + "cache_creation_input_token_cost_above_1hr": 6e-06, + "cache_read_input_token_cost": 3e-07, + "input_cost_per_token": 3e-06, "litellm_provider": "bedrock_converse", "supports_tool_search": true, "max_input_tokens": 1000000, "max_output_tokens": 64000, "max_tokens": 64000, "mode": "chat", - "output_cost_per_token": 1.65e-05, + "output_cost_per_token": 1.5e-05, "search_context_cost_per_query": { "search_context_size_high": 0.01, "search_context_size_low": 0.01, @@ -2759,7 +2759,7 @@ "supports_output_config": true, "supports_parallel_tool_use_config": true, "prompt_cache_min_tokens": 1024, - "source": "https://aws.amazon.com/bedrock/pricing/" + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json" }, "global.anthropic.claude-sonnet-4-6": { "supports_adaptive_thinking": true, @@ -2970,22 +2970,22 @@ "source": "https://aws.amazon.com/bedrock/pricing/" }, "anthropic.claude-sonnet-4-5-20250929-v1:0": { - "cache_creation_input_token_cost": 4.125e-06, - "cache_creation_input_token_cost_above_1hr": 6.6e-06, - "cache_read_input_token_cost": 3.3e-07, - "input_cost_per_token": 3.3e-06, - "input_cost_per_token_above_200k_tokens": 6.6e-06, - "output_cost_per_token_above_200k_tokens": 2.475e-05, - "cache_creation_input_token_cost_above_200k_tokens": 8.25e-06, - "cache_creation_input_token_cost_above_1hr_above_200k_tokens": 1.32e-05, - "cache_read_input_token_cost_above_200k_tokens": 6.6e-07, + "cache_creation_input_token_cost": 3.75e-06, + "cache_creation_input_token_cost_above_1hr": 6e-06, + "cache_read_input_token_cost": 3e-07, + "input_cost_per_token": 3e-06, + "input_cost_per_token_above_200k_tokens": 6e-06, + "output_cost_per_token_above_200k_tokens": 2.25e-05, + "cache_creation_input_token_cost_above_200k_tokens": 7.5e-06, + "cache_creation_input_token_cost_above_1hr_above_200k_tokens": 1.2e-05, + "cache_read_input_token_cost_above_200k_tokens": 6e-07, "litellm_provider": "bedrock_converse", "supports_tool_search": true, "max_input_tokens": 200000, "max_output_tokens": 64000, "max_tokens": 64000, "mode": "chat", - "output_cost_per_token": 1.65e-05, + "output_cost_per_token": 1.5e-05, "search_context_cost_per_query": { "search_context_size_high": 0.01, "search_context_size_low": 0.01, @@ -3003,9 +3003,9 @@ "supports_native_structured_output": true, "supports_parallel_tool_use_config": true, "prompt_cache_min_tokens": 1024, - "input_cost_per_token_batches": 1.65e-06, - "output_cost_per_token_batches": 8.25e-06, - "source": "https://aws.amazon.com/bedrock/pricing/" + "input_cost_per_token_batches": 1.5e-06, + "output_cost_per_token_batches": 7.5e-06, + "source": "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json" }, "anthropic.claude-v1": { "input_cost_per_token": 8e-06, diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index 3fca6935d57..5dc09db4535 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -1227,17 +1227,17 @@ def test_get_model_info_bedrock_regional_inference_profile_pricing(local_model_c "anthropic.claude-sonnet-5", ], ) -def test_bedrock_bare_claude_id_is_priced_in_region(local_model_cost_map, bare_key): - """A bare Bedrock Claude id is an in-region invocation, so it carries the same - in-region rate as its us. inference profile, not the cheaper global. rate.""" +def test_bedrock_bare_claude_id_is_priced_global(local_model_cost_map, bare_key): + """A bare Bedrock Claude id is billed at the Global SKU, so it carries the same + rate as its global. inference profile and sits below the regional us. rate.""" bare = litellm.model_cost[bare_key] us = litellm.model_cost[f"us.{bare_key}"] global_ = litellm.model_cost[f"global.{bare_key}"] cost_fields = [f for f in bare if "cost" in f] assert cost_fields for field in cost_fields: - assert bare[field] == us[field], field - assert bare["input_cost_per_token"] > global_["input_cost_per_token"] + assert bare[field] == global_[field], field + assert bare["input_cost_per_token"] < us["input_cost_per_token"] def test_get_model_info_bedrock_mantle_region_prefix_falls_back_to_the_mantle_row(local_model_cost_map): From e99c5d30b6238f417d0370b450bded6f7e403692 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 23 Sep 2026 21:32:27 -0700 Subject: [PATCH 009/120] feat(cost-map): add retired azure gpt-5 chat and o1-preview data zone rows (#42878) Co-authored-by: kerry Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- ...odel_prices_and_context_window_backup.json | 249 ++++++++++++++++++ model_prices_and_context_window.json | 249 ++++++++++++++++++ 2 files changed, 498 insertions(+) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 059c92f4b4b..2c4de480f8b 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -6627,6 +6627,255 @@ "output_cost_per_token_priority": 2e-05, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" }, + "azure/gpt-5-chat": { + "cache_read_input_token_cost": 1.25e-07, + "deprecation_date": "2026-05-13", + "input_cost_per_token": 1.25e-06, + "litellm_provider": "azure", + "max_input_tokens": 128000, + "max_output_tokens": 16384, + "max_tokens": 16384, + "mode": "chat", + "output_cost_per_token": 1e-05, + "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_prompt_caching": true, + "supports_vision": true + }, + "azure/gpt-5.1-chat": { + "cache_read_input_token_cost": 1.25e-07, + "deprecation_date": "2026-06-29", + "input_cost_per_token": 1.25e-06, + "litellm_provider": "azure", + "max_input_tokens": 111616, + "max_output_tokens": 16384, + "max_tokens": 16384, + "mode": "chat", + "output_cost_per_token": 1e-05, + "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses" + ], + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, + "azure/gpt-5.2-chat": { + "cache_read_input_token_cost": 1.75e-07, + "deprecation_date": "2026-06-29", + "input_cost_per_token": 1.75e-06, + "litellm_provider": "azure", + "max_input_tokens": 111616, + "max_output_tokens": 16384, + "max_tokens": 16384, + "mode": "chat", + "output_cost_per_token": 1.4e-05, + "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses" + ], + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_prompt_caching": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, + "azure/gpt-5.3-chat": { + "cache_read_input_token_cost": 1.75e-07, + "deprecation_date": "2026-06-29", + "input_cost_per_token": 1.75e-06, + "litellm_provider": "azure", + "max_input_tokens": 111616, + "max_output_tokens": 16384, + "max_tokens": 16384, + "mode": "chat", + "output_cost_per_token": 1.4e-05, + "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses" + ], + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_prompt_caching": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, + "azure/us/gpt-5.1-chat": { + "cache_read_input_token_cost": 1.375e-07, + "deprecation_date": "2026-06-29", + "input_cost_per_token": 1.375e-06, + "litellm_provider": "azure", + "max_input_tokens": 111616, + "max_output_tokens": 16384, + "max_tokens": 16384, + "mode": "chat", + "output_cost_per_token": 1.1e-05, + "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses" + ], + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, + "azure/us/gpt-5.2-chat": { + "cache_read_input_token_cost": 1.925e-07, + "deprecation_date": "2026-06-29", + "input_cost_per_token": 1.925e-06, + "litellm_provider": "azure", + "max_input_tokens": 111616, + "max_output_tokens": 16384, + "max_tokens": 16384, + "mode": "chat", + "output_cost_per_token": 1.54e-05, + "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses" + ], + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_prompt_caching": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, + "azure/us/gpt-5.3-chat": { + "cache_read_input_token_cost": 1.925e-07, + "deprecation_date": "2026-06-29", + "input_cost_per_token": 1.925e-06, + "litellm_provider": "azure", + "max_input_tokens": 111616, + "max_output_tokens": 16384, + "max_tokens": 16384, + "mode": "chat", + "output_cost_per_token": 1.54e-05, + "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses" + ], + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_prompt_caching": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, + "azure/eu/gpt-5.1-chat": { + "cache_read_input_token_cost": 1.375e-07, + "deprecation_date": "2026-06-29", + "input_cost_per_token": 1.375e-06, + "litellm_provider": "azure", + "max_input_tokens": 111616, + "max_output_tokens": 16384, + "max_tokens": 16384, + "mode": "chat", + "output_cost_per_token": 1.1e-05, + "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses" + ], + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, + "azure/eu/gpt-5.2-chat": { + "cache_read_input_token_cost": 1.925e-07, + "deprecation_date": "2026-06-29", + "input_cost_per_token": 1.925e-06, + "litellm_provider": "azure", + "max_input_tokens": 111616, + "max_output_tokens": 16384, + "max_tokens": 16384, + "mode": "chat", + "output_cost_per_token": 1.54e-05, + "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses" + ], + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_prompt_caching": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, + "azure/eu/gpt-5.3-chat": { + "cache_read_input_token_cost": 1.925e-07, + "deprecation_date": "2026-06-29", + "input_cost_per_token": 1.925e-06, + "litellm_provider": "azure", + "max_input_tokens": 111616, + "max_output_tokens": 16384, + "max_tokens": 16384, + "mode": "chat", + "output_cost_per_token": 1.54e-05, + "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses" + ], + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_prompt_caching": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, + "azure/us/o1-preview": { + "cache_read_input_token_cost": 8.25e-06, + "deprecation_date": "2025-07-28", + "input_cost_per_token": 1.65e-05, + "litellm_provider": "azure", + "max_input_tokens": 128000, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "chat", + "output_cost_per_token": 6.6e-05, + "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", + "supported_endpoints": [ + "/v1/chat/completions" + ], + "supports_prompt_caching": true + }, + "azure/eu/o1-preview": { + "cache_read_input_token_cost": 8.25e-06, + "deprecation_date": "2025-07-28", + "input_cost_per_token": 1.65e-05, + "litellm_provider": "azure", + "max_input_tokens": 128000, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "chat", + "output_cost_per_token": 6.6e-05, + "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", + "supported_endpoints": [ + "/v1/chat/completions" + ], + "supports_prompt_caching": true + }, "azure/gpt-5.1-codex": { "deprecation_date": "2027-05-15", "cache_read_input_token_cost": 1.25e-07, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 059c92f4b4b..2c4de480f8b 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -6627,6 +6627,255 @@ "output_cost_per_token_priority": 2e-05, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" }, + "azure/gpt-5-chat": { + "cache_read_input_token_cost": 1.25e-07, + "deprecation_date": "2026-05-13", + "input_cost_per_token": 1.25e-06, + "litellm_provider": "azure", + "max_input_tokens": 128000, + "max_output_tokens": 16384, + "max_tokens": 16384, + "mode": "chat", + "output_cost_per_token": 1e-05, + "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_prompt_caching": true, + "supports_vision": true + }, + "azure/gpt-5.1-chat": { + "cache_read_input_token_cost": 1.25e-07, + "deprecation_date": "2026-06-29", + "input_cost_per_token": 1.25e-06, + "litellm_provider": "azure", + "max_input_tokens": 111616, + "max_output_tokens": 16384, + "max_tokens": 16384, + "mode": "chat", + "output_cost_per_token": 1e-05, + "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses" + ], + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, + "azure/gpt-5.2-chat": { + "cache_read_input_token_cost": 1.75e-07, + "deprecation_date": "2026-06-29", + "input_cost_per_token": 1.75e-06, + "litellm_provider": "azure", + "max_input_tokens": 111616, + "max_output_tokens": 16384, + "max_tokens": 16384, + "mode": "chat", + "output_cost_per_token": 1.4e-05, + "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses" + ], + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_prompt_caching": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, + "azure/gpt-5.3-chat": { + "cache_read_input_token_cost": 1.75e-07, + "deprecation_date": "2026-06-29", + "input_cost_per_token": 1.75e-06, + "litellm_provider": "azure", + "max_input_tokens": 111616, + "max_output_tokens": 16384, + "max_tokens": 16384, + "mode": "chat", + "output_cost_per_token": 1.4e-05, + "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses" + ], + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_prompt_caching": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, + "azure/us/gpt-5.1-chat": { + "cache_read_input_token_cost": 1.375e-07, + "deprecation_date": "2026-06-29", + "input_cost_per_token": 1.375e-06, + "litellm_provider": "azure", + "max_input_tokens": 111616, + "max_output_tokens": 16384, + "max_tokens": 16384, + "mode": "chat", + "output_cost_per_token": 1.1e-05, + "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses" + ], + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, + "azure/us/gpt-5.2-chat": { + "cache_read_input_token_cost": 1.925e-07, + "deprecation_date": "2026-06-29", + "input_cost_per_token": 1.925e-06, + "litellm_provider": "azure", + "max_input_tokens": 111616, + "max_output_tokens": 16384, + "max_tokens": 16384, + "mode": "chat", + "output_cost_per_token": 1.54e-05, + "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses" + ], + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_prompt_caching": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, + "azure/us/gpt-5.3-chat": { + "cache_read_input_token_cost": 1.925e-07, + "deprecation_date": "2026-06-29", + "input_cost_per_token": 1.925e-06, + "litellm_provider": "azure", + "max_input_tokens": 111616, + "max_output_tokens": 16384, + "max_tokens": 16384, + "mode": "chat", + "output_cost_per_token": 1.54e-05, + "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses" + ], + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_prompt_caching": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, + "azure/eu/gpt-5.1-chat": { + "cache_read_input_token_cost": 1.375e-07, + "deprecation_date": "2026-06-29", + "input_cost_per_token": 1.375e-06, + "litellm_provider": "azure", + "max_input_tokens": 111616, + "max_output_tokens": 16384, + "max_tokens": 16384, + "mode": "chat", + "output_cost_per_token": 1.1e-05, + "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses" + ], + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, + "azure/eu/gpt-5.2-chat": { + "cache_read_input_token_cost": 1.925e-07, + "deprecation_date": "2026-06-29", + "input_cost_per_token": 1.925e-06, + "litellm_provider": "azure", + "max_input_tokens": 111616, + "max_output_tokens": 16384, + "max_tokens": 16384, + "mode": "chat", + "output_cost_per_token": 1.54e-05, + "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses" + ], + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_prompt_caching": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, + "azure/eu/gpt-5.3-chat": { + "cache_read_input_token_cost": 1.925e-07, + "deprecation_date": "2026-06-29", + "input_cost_per_token": 1.925e-06, + "litellm_provider": "azure", + "max_input_tokens": 111616, + "max_output_tokens": 16384, + "max_tokens": 16384, + "mode": "chat", + "output_cost_per_token": 1.54e-05, + "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses" + ], + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_prompt_caching": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, + "azure/us/o1-preview": { + "cache_read_input_token_cost": 8.25e-06, + "deprecation_date": "2025-07-28", + "input_cost_per_token": 1.65e-05, + "litellm_provider": "azure", + "max_input_tokens": 128000, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "chat", + "output_cost_per_token": 6.6e-05, + "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", + "supported_endpoints": [ + "/v1/chat/completions" + ], + "supports_prompt_caching": true + }, + "azure/eu/o1-preview": { + "cache_read_input_token_cost": 8.25e-06, + "deprecation_date": "2025-07-28", + "input_cost_per_token": 1.65e-05, + "litellm_provider": "azure", + "max_input_tokens": 128000, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "chat", + "output_cost_per_token": 6.6e-05, + "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", + "supported_endpoints": [ + "/v1/chat/completions" + ], + "supports_prompt_caching": true + }, "azure/gpt-5.1-codex": { "deprecation_date": "2027-05-15", "cache_read_input_token_cost": 1.25e-07, From 9d12c217afec1e16b003b9015e88c010d2b2584f Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 23 Sep 2026 21:42:23 -0700 Subject: [PATCH 010/120] fix(models): correct gemini robotics er 2 preview audio input price (#42877) Co-authored-by: kerry Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/model_prices_and_context_window_backup.json | 4 ++-- model_prices_and_context_window.json | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 2c4de480f8b..bcf77bc16d4 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -27530,7 +27530,7 @@ "gemini/gemini-robotics-er-2-preview": { "cache_read_input_token_cost": 1e-07, "cache_read_input_token_cost_batches": 5e-08, - "input_cost_per_audio_token": 2e-06, + "input_cost_per_audio_token": 1e-06, "input_cost_per_token": 1e-06, "input_cost_per_token_batches": 5e-07, "litellm_provider": "gemini", @@ -59160,7 +59160,7 @@ } }, "gemini/gemini-robotics-er-2-streaming-preview": { - "input_cost_per_audio_token": 2e-06, + "input_cost_per_audio_token": 1e-06, "input_cost_per_token": 1e-06, "litellm_provider": "gemini", "mode": "chat", diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 2c4de480f8b..bcf77bc16d4 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -27530,7 +27530,7 @@ "gemini/gemini-robotics-er-2-preview": { "cache_read_input_token_cost": 1e-07, "cache_read_input_token_cost_batches": 5e-08, - "input_cost_per_audio_token": 2e-06, + "input_cost_per_audio_token": 1e-06, "input_cost_per_token": 1e-06, "input_cost_per_token_batches": 5e-07, "litellm_provider": "gemini", @@ -59160,7 +59160,7 @@ } }, "gemini/gemini-robotics-er-2-streaming-preview": { - "input_cost_per_audio_token": 2e-06, + "input_cost_per_audio_token": 1e-06, "input_cost_per_token": 1e-06, "litellm_provider": "gemini", "mode": "chat", From b21b20ed13b0cb531bbf4ef56db263983f1511ce Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 23 Sep 2026 21:55:55 -0700 Subject: [PATCH 011/120] feat(vertex_ai): add gemini-3.8-flash-cyber pricing (#42879) * feat(vertex_ai): add gemini-3.8-flash-cyber pricing Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(vertex_ai): mark gemini-3.8-flash-cyber minimal reasoning unsupported Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: kerry Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- ...odel_prices_and_context_window_backup.json | 92 +++++++++++++++++++ model_prices_and_context_window.json | 92 +++++++++++++++++++ 2 files changed, 184 insertions(+) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index bcf77bc16d4..66b72090631 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -27336,6 +27336,52 @@ "web_search_billing_unit": "per_query", "google_maps_grounding_cost_per_query": 0.014 }, + "vertex_ai/gemini-3.8-flash-cyber": { + "cache_read_input_token_cost": 1.5e-07, + "cache_read_input_token_cost_flex": 7.5e-08, + "cache_read_input_token_cost_priority": 2.7e-07, + "input_cost_per_token": 1.5e-06, + "input_cost_per_token_flex": 7.5e-07, + "input_cost_per_token_priority": 2.7e-06, + "litellm_provider": "vertex_ai", + "max_input_tokens": 1048576, + "max_output_tokens": 65536, + "max_tokens": 65536, + "mode": "chat", + "output_cost_per_reasoning_token": 7.5e-06, + "output_cost_per_token": 7.5e-06, + "output_cost_per_token_flex": 3.75e-06, + "output_cost_per_token_priority": 1.35e-05, + "regional_endpoint_uplift_multiplier": 1.1, + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/completions" + ], + "supported_modalities": [ + "text", + "image", + "audio", + "video" + ], + "supported_output_modalities": [ + "text" + ], + "supports_audio_input": true, + "supports_function_calling": false, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_minimal_reasoning_effort": false, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": false, + "supports_url_context": false, + "supports_video_input": true, + "supports_vision": true, + "supports_web_search": false, + "supports_native_streaming": true + }, "vertex_ai/gemini-3.1-pro-preview": { "prompt_cache_min_tokens": 4096, "cache_read_input_token_cost": 2e-07, @@ -29508,6 +29554,52 @@ "web_search_billing_unit": "per_query", "google_maps_grounding_cost_per_query": 0.014 }, + "gemini-3.8-flash-cyber": { + "cache_read_input_token_cost": 1.5e-07, + "cache_read_input_token_cost_flex": 7.5e-08, + "cache_read_input_token_cost_priority": 2.7e-07, + "input_cost_per_token": 1.5e-06, + "input_cost_per_token_flex": 7.5e-07, + "input_cost_per_token_priority": 2.7e-06, + "litellm_provider": "vertex_ai-language-models", + "max_input_tokens": 1048576, + "max_output_tokens": 65536, + "max_tokens": 65536, + "mode": "chat", + "output_cost_per_reasoning_token": 7.5e-06, + "output_cost_per_token": 7.5e-06, + "output_cost_per_token_flex": 3.75e-06, + "output_cost_per_token_priority": 1.35e-05, + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/completions" + ], + "supported_modalities": [ + "text", + "image", + "audio", + "video" + ], + "supported_output_modalities": [ + "text" + ], + "supports_audio_output": false, + "supports_audio_input": true, + "supports_function_calling": false, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_minimal_reasoning_effort": false, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": false, + "supports_url_context": false, + "supports_video_input": true, + "supports_vision": true, + "supports_web_search": false, + "supports_native_streaming": true + }, "gemini/gemini-2.5-pro-preview-tts": { "cache_read_input_token_cost": 1.25e-07, "input_cost_per_audio_token": 7e-07, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index bcf77bc16d4..66b72090631 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -27336,6 +27336,52 @@ "web_search_billing_unit": "per_query", "google_maps_grounding_cost_per_query": 0.014 }, + "vertex_ai/gemini-3.8-flash-cyber": { + "cache_read_input_token_cost": 1.5e-07, + "cache_read_input_token_cost_flex": 7.5e-08, + "cache_read_input_token_cost_priority": 2.7e-07, + "input_cost_per_token": 1.5e-06, + "input_cost_per_token_flex": 7.5e-07, + "input_cost_per_token_priority": 2.7e-06, + "litellm_provider": "vertex_ai", + "max_input_tokens": 1048576, + "max_output_tokens": 65536, + "max_tokens": 65536, + "mode": "chat", + "output_cost_per_reasoning_token": 7.5e-06, + "output_cost_per_token": 7.5e-06, + "output_cost_per_token_flex": 3.75e-06, + "output_cost_per_token_priority": 1.35e-05, + "regional_endpoint_uplift_multiplier": 1.1, + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/completions" + ], + "supported_modalities": [ + "text", + "image", + "audio", + "video" + ], + "supported_output_modalities": [ + "text" + ], + "supports_audio_input": true, + "supports_function_calling": false, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_minimal_reasoning_effort": false, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": false, + "supports_url_context": false, + "supports_video_input": true, + "supports_vision": true, + "supports_web_search": false, + "supports_native_streaming": true + }, "vertex_ai/gemini-3.1-pro-preview": { "prompt_cache_min_tokens": 4096, "cache_read_input_token_cost": 2e-07, @@ -29508,6 +29554,52 @@ "web_search_billing_unit": "per_query", "google_maps_grounding_cost_per_query": 0.014 }, + "gemini-3.8-flash-cyber": { + "cache_read_input_token_cost": 1.5e-07, + "cache_read_input_token_cost_flex": 7.5e-08, + "cache_read_input_token_cost_priority": 2.7e-07, + "input_cost_per_token": 1.5e-06, + "input_cost_per_token_flex": 7.5e-07, + "input_cost_per_token_priority": 2.7e-06, + "litellm_provider": "vertex_ai-language-models", + "max_input_tokens": 1048576, + "max_output_tokens": 65536, + "max_tokens": 65536, + "mode": "chat", + "output_cost_per_reasoning_token": 7.5e-06, + "output_cost_per_token": 7.5e-06, + "output_cost_per_token_flex": 3.75e-06, + "output_cost_per_token_priority": 1.35e-05, + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/completions" + ], + "supported_modalities": [ + "text", + "image", + "audio", + "video" + ], + "supported_output_modalities": [ + "text" + ], + "supports_audio_output": false, + "supports_audio_input": true, + "supports_function_calling": false, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_minimal_reasoning_effort": false, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": false, + "supports_url_context": false, + "supports_video_input": true, + "supports_vision": true, + "supports_web_search": false, + "supports_native_streaming": true + }, "gemini/gemini-2.5-pro-preview-tts": { "cache_read_input_token_cost": 1.25e-07, "input_cost_per_audio_token": 7e-07, From b8154bcbc0907a99b916f07a4b05c0622c4965ee Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 23 Sep 2026 22:08:23 -0700 Subject: [PATCH 012/120] fix(presidio): stream non-Anthropic raw SSE through the post_call hook unbuffered (#42777) * fix(presidio): stream non-Anthropic raw SSE through the post_call hook unbuffered Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(presidio): keep the pytest.raises block to a single await Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * refactor(presidio): move raw SSE format check into a helper to keep hook complexity flat Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(presidio): fold the raw SSE format check into the existing bytes branch to stay within the complexity budget Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(presidio): decide raw SSE stream shape on a complete first frame, not a transport fragment Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(integration): cover presidio post_call streaming for native gemini passthrough and anthropic messages Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(presidio): cap first SSE frame coalescing at 64 KiB so an unterminated first event cannot buffer unbounded Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(presidio): name raw SSE passthrough in the skipped output masking warning Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: yucheng Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../guardrails/guardrail_hooks/presidio.py | 54 +- .../test_presidio_streaming_output.py | 538 ++++++++++++++++++ .../guardrail_hooks/test_presidio.py | 143 ++++- 3 files changed, 730 insertions(+), 5 deletions(-) create mode 100644 tests/integration/observability/test_presidio_streaming_output.py diff --git a/litellm/proxy/guardrails/guardrail_hooks/presidio.py b/litellm/proxy/guardrails/guardrail_hooks/presidio.py index eb07c19a580..794bf08729e 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/presidio.py +++ b/litellm/proxy/guardrails/guardrail_hooks/presidio.py @@ -42,6 +42,7 @@ from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.guardrails.anthropic_sse import ( anthropic_sse_chunks_from_response, assemble_anthropic_sse_stream, + is_anthropic_sse_stream, model_response_text, ) from litellm.types.guardrails import ( @@ -93,6 +94,42 @@ def _json_escaped_len(text: str) -> int: return len(json.dumps(text).encode("utf-8")) - 2 # strip the surrounding quotes +_MAX_FIRST_SSE_FRAME_BYTES: Final = 64 * 1024 + + +def _holds_complete_sse_frame(raw: bytes) -> bool: + """Whether ``raw`` holds one blank-line terminated SSE event, or is too large to keep joining.""" + return b"\n\n" in raw or b"\r\n\r\n" in raw or len(raw) >= _MAX_FIRST_SSE_FRAME_BYTES + + +async def _coalesce_first_sse_frame(stream: AsyncIterator[object]) -> AsyncGenerator[object, None]: + """ + Join leading raw ``bytes`` chunks until they hold one complete SSE event, so + the stream shape is decided on a whole frame rather than a transport fragment. + Everything after that first frame is forwarded untouched. + """ + pending = b"" + try: + async for chunk in stream: + if not isinstance(chunk, bytes): + yield chunk + continue + pending += chunk + if _holds_complete_sse_frame(pending): + break + else: + if pending: + yield pending + return + except Exception: + if pending: + yield pending + raise + yield pending + async for chunk in stream: + yield chunk + + class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): user_api_key_cache = None ad_hoc_recognizers: list[str] | None = None @@ -1356,7 +1393,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): all_chunks: list[ModelResponseStream] = [] passthrough_due_to_unknown_stream_shape = False try: - stream: Final = response.__aiter__() + stream: Final = _coalesce_first_sse_frame(response.__aiter__()) async for chunk in stream: if isinstance(chunk, ModelResponseStream): if passthrough_due_to_unknown_stream_shape: @@ -1364,7 +1401,15 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): else: all_chunks.append(chunk) elif isinstance(chunk, bytes): - if passthrough_due_to_unknown_stream_shape or all_chunks: + first_frame_is_anthropic = ( + not passthrough_due_to_unknown_stream_shape + and not all_chunks + and is_anthropic_sse_stream((chunk,)) + ) + if not first_frame_is_anthropic: + passthrough_due_to_unknown_stream_shape = ( + passthrough_due_to_unknown_stream_shape or not all_chunks + ) yield chunk continue for masked_chunk in await self._mask_anthropic_sse_stream(chunk, stream, request_data): @@ -1387,8 +1432,9 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): yield chunk if passthrough_due_to_unknown_stream_shape: verbose_proxy_logger.warning( - "Presidio apply_to_output: streaming response contained unknown event objects " - "(e.g. /v1/responses events). Output PII masking was skipped for this response." + "Presidio apply_to_output: streaming response was not a parsed chat completion stream " + "(raw non-Anthropic SSE passthrough or /v1/responses events). " + "Output PII masking was skipped for this response." ) return if not all_chunks: diff --git a/tests/integration/observability/test_presidio_streaming_output.py b/tests/integration/observability/test_presidio_streaming_output.py new file mode 100644 index 00000000000..5bc0a46427b --- /dev/null +++ b/tests/integration/observability/test_presidio_streaming_output.py @@ -0,0 +1,538 @@ +import json +import re +import signal +import threading +import uuid +from collections.abc import Callable, Iterator, Mapping +from concurrent.futures import ThreadPoolExecutor +from contextlib import ExitStack, contextmanager +from dataclasses import dataclass +from pathlib import Path +from typing import Final + +import psutil +import yaml +from integration._support.client import Gateway, eventually +from integration._support.process import OwnedProxy, group_members, owned_proxy_process +from integration._support.wire import Reply, Request, Wire, wire_server +from openai import OpenAI +from pydantic import BaseModel + +PERSON: Final = "John Smith" +MASK: Final = "" +GEMINI_MODEL: Final = "gemini-2.5-flash" + + +def gemini_frame(text: str) -> bytes: + payload: Final = { + "candidates": [{"content": {"parts": [{"text": text}], "role": "model"}, "index": 0}], + "usageMetadata": {"promptTokenCount": 10, "candidatesTokenCount": 5, "totalTokenCount": 15}, + "modelVersion": GEMINI_MODEL, + } + return b"data: " + json.dumps(payload).encode() + b"\r\n\r\n" + + +class GeminiPart(BaseModel): + text: str + + +class GeminiContent(BaseModel): + parts: list[GeminiPart] + + +class GeminiCandidate(BaseModel): + content: GeminiContent + + +class GeminiFrame(BaseModel): + candidates: list[GeminiCandidate] + + +def data_payloads(raw: bytes) -> tuple[dict[str, object], ...]: + """JSON payload of each ``data:`` frame, whatever line ending the sender used.""" + return tuple(json.loads(line[len("data: ") :]) for line in raw.decode().splitlines() if line.startswith("data: ")) + + +def gemini_text(payload: Mapping[str, object]) -> str: + return GeminiFrame.model_validate(payload).candidates[0].content.parts[0].text + + +def gemini_texts(raw: bytes) -> tuple[str, ...]: + return tuple(gemini_text(payload) for payload in data_payloads(raw)) + + +def anthropic_frame(event_type: str, payload: dict[str, object]) -> bytes: + return f"event: {event_type}\ndata: {json.dumps(payload)}\n\n".encode() + + +def anthropic_stream(identity: str, text: str) -> tuple[bytes, ...]: + return ( + anthropic_frame( + "message_start", + { + "type": "message_start", + "message": { + "id": identity, + "type": "message", + "role": "assistant", + "model": "claude-sonnet-4-5-20250929", + "content": [], + "stop_reason": None, + "usage": {"input_tokens": 11, "output_tokens": 0}, + }, + }, + ), + anthropic_frame( + "content_block_start", + {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}, + ), + anthropic_frame( + "content_block_delta", + {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": text}}, + ), + anthropic_frame("content_block_stop", {"type": "content_block_stop", "index": 0}), + anthropic_frame( + "message_delta", + { + "type": "message_delta", + "delta": {"stop_reason": "end_turn", "stop_sequence": None}, + "usage": {"output_tokens": 4}, + }, + ), + anthropic_frame("message_stop", {"type": "message_stop"}), + ) + + +def openai_frame(identity: str, delta: dict[str, str], finish: str | None = None) -> bytes: + payload: Final = { + "id": identity, + "object": "chat.completion.chunk", + "created": 1, + "model": "gpt-4o-mini", + "choices": [{"index": 0, "delta": delta, "finish_reason": finish}], + } + return b"data: " + json.dumps(payload).encode() + b"\n\n" + + +def analyzer(request: Request) -> Reply: + assert request.target == "/analyze", request.target + text: Final = json.loads(request.body)["text"] + findings: Final = [ + {"entity_type": "PERSON", "start": match.start(), "end": match.end(), "score": 0.85} + for match in re.finditer(re.escape(PERSON), text) + ] + return Reply(body=json.dumps(findings).encode()) + + +def anonymizer(request: Request) -> Reply: + assert request.target == "/anonymize", request.target + body: Final = json.loads(request.body) + text: Final = body["text"] + items: Final = [ + {"entity_type": "PERSON", "start": item["start"], "end": item["end"], "operator": "replace", "text": MASK} + for item in body["analyzer_results"] + ] + return Reply(body=json.dumps({"text": text.replace(PERSON, MASK), "items": items}).encode()) + + +def broken(request: Request) -> Reply: + return Reply(status=500, body=b'{"error": "scripted outage"}') + + +@dataclass(frozen=True, slots=True) +class Received: + status: int + frames: tuple[bytes, ...] + + @property + def text(self) -> str: + return b"".join(self.frames).decode() + + +@dataclass(frozen=True, slots=True) +class Rig: + proxy: OwnedProxy + upstream: Wire + analyzer: Wire + anonymizer: Wire + guardrail: str + gemini: str + anthropic: str + openai: str + + @property + def gateway(self) -> Gateway: + return self.proxy.gateway + + def stream(self, path: str, body: dict[str, object] | None = None, *, key: str | None = None) -> Received: + with self.gateway.client.stream( + "POST", path, json=body, headers={"Authorization": f"Bearer {key or self.gateway.key}"} + ) as response: + return Received(response.status_code, tuple(response.iter_raw())) + + def gemini_path(self) -> str: + return f"/v1beta/models/{self.gemini}:streamGenerateContent?alt=sse" + + def gemini_body(self) -> dict[str, object]: + return {"contents": [{"role": "user", "parts": [{"text": "who designed it"}]}]} + + def messages_body(self, *, guardrails: tuple[str, ...] | None = None) -> dict[str, object]: + return { + "model": self.anthropic, + "max_tokens": 64, + "stream": True, + "messages": [{"role": "user", "content": "who designed it"}], + **({"guardrails": list(guardrails)} if guardrails is not None else {}), + } + + +def anthropic_text(received: Received) -> str: + events: Final = tuple( + json.loads(line.removeprefix("data: ")) for line in received.text.split("\n") if line.startswith("data: ") + ) + return "".join(event["delta"]["text"] for event in events if event.get("type") == "content_block_delta") + + +@contextmanager +def presidio_rig( + gateway: Gateway, + tmp_path: Path, + provider: Callable[[Request], Reply], + *, + analyze: Callable[[Request], Reply] = analyzer, + anonymize: Callable[[Request], Reply] = anonymizer, + default_on: bool = True, +) -> Iterator[Rig]: + guardrail: Final = "presidio" + uuid.uuid4().hex + with ExitStack() as stack: + upstream: Final = stack.enter_context(wire_server(provider)) + analyze_sink: Final = stack.enter_context(wire_server(analyze)) + anonymize_sink: Final = stack.enter_context(wire_server(anonymize)) + config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + config["guardrails"] = [ + { + "guardrail_name": guardrail, + "litellm_params": { + "guardrail": "presidio", + "mode": "post_call", + "default_on": default_on, + "presidio_analyzer_api_base": analyze_sink.url, + "presidio_anonymizer_api_base": anonymize_sink.url, + "presidio_filter_scope": "output", + }, + } + ] + path: Final = tmp_path / f"{guardrail}.yaml" + path.write_text(yaml.safe_dump(config)) + proxy: Final = stack.enter_context(owned_proxy_process(gateway, tmp_path, {}, config=path, workers=2)) + scenario: Final = stack.enter_context(proxy.gateway.scenario()) + yield Rig( + proxy=proxy, + upstream=upstream, + analyzer=analyze_sink, + anonymizer=anonymize_sink, + guardrail=guardrail, + gemini=scenario.model( + model=f"gemini/{GEMINI_MODEL}", api_base=upstream.url, api_key="synthetic-gemini-key" + ), + anthropic=scenario.model( + model="anthropic/claude-sonnet-4-5-20250929", api_base=upstream.url, api_key="synthetic-anthropic-key" + ), + openai=scenario.model(model="openai/gpt-4o-mini", api_base=upstream.url + "/v1", api_key="synthetic-key"), + ) + + +def gemini_provider(reply: Reply) -> Callable[[Request], Reply]: + def provider(request: Request) -> Reply: + assert "streamGenerateContent" in request.target, request.target + return reply + + return provider + + +def test_native_gemini_first_frame_reaches_caller_before_upstream_sends_the_second( + gateway: Gateway, tmp_path: Path +) -> None: + gate: Final = threading.Event() + first: Final = gemini_frame("first ") + second: Final = gemini_frame("second ") + provider: Final = gemini_provider( + Reply(content_type="text/event-stream", chunks=(first, second), gate_after_first=gate) + ) + with presidio_rig(gateway, tmp_path, provider) as rig: + with rig.gateway.client.stream( + "POST", rig.gemini_path(), json=rig.gemini_body(), headers={"Authorization": f"Bearer {rig.gateway.key}"} + ) as response: + assert response.status_code == 200, response.read().decode() + chunks: Final = response.iter_raw() + arrived: Final = next(chunks) + assert gemini_texts(arrived) == ("first ",), f"first chunk while upstream is gated: {arrived!r}" + gate.set() + rest: Final = b"".join(chunks) + assert gemini_texts(rest) == ("second ",), rest + assert len(rig.upstream.drain()) == 1 + assert rig.analyzer.drain() == () and rig.anonymizer.drain() == () + + +def test_native_gemini_frames_received_before_upstream_abort_reach_caller(gateway: Gateway, tmp_path: Path) -> None: + frames: Final = (gemini_frame(f"chunk {index} from {PERSON}. ") for index in range(3)) + provider: Final = gemini_provider( + Reply(content_type="text/event-stream", chunks=tuple(frames), abort_after=2, pause_between_chunks=0.2) + ) + with presidio_rig(gateway, tmp_path, provider) as rig: + received: Final = rig.stream(rig.gemini_path(), rig.gemini_body()) + assert received.status == 200, received.text + *frames_before_abort, trailer = data_payloads(b"".join(received.frames)) + assert [gemini_text(frame) for frame in frames_before_abort] == [ + f"chunk 0 from {PERSON}. ", + f"chunk 1 from {PERSON}. ", + ], received.text + assert "candidates" not in trailer and json.dumps(trailer).count('"code": "500"') == 1, received.text + assert len(rig.upstream.drain()) == 1 + + +def test_native_gemini_first_frame_split_into_transport_fragments_streams_every_byte( + gateway: Gateway, tmp_path: Path +) -> None: + first: Final = gemini_frame(f"fragmented {PERSON}") + second: Final = gemini_frame("whole") + chunks: Final = (first[:7], first[7:19], first[19:], second) + provider: Final = gemini_provider(Reply(content_type="text/event-stream", chunks=chunks)) + with presidio_rig(gateway, tmp_path, provider) as rig: + received: Final = rig.stream(rig.gemini_path(), rig.gemini_body()) + assert received.status == 200, received.text + assert gemini_texts(b"".join(received.frames)) == (f"fragmented {PERSON}", "whole") + + +def test_native_gemini_non_json_frame_passes_through_unchanged(gateway: Gateway, tmp_path: Path) -> None: + frames: Final = (b"data: not json at all\r\n\r\n", gemini_frame("after")) + provider: Final = gemini_provider(Reply(content_type="text/event-stream", chunks=frames)) + with presidio_rig(gateway, tmp_path, provider) as rig: + received: Final = rig.stream(rig.gemini_path(), rig.gemini_body()) + assert received.status == 200, received.text + assert received.text.replace("\r\n", "\n") == b"".join(frames).decode().replace("\r\n", "\n") + + +def test_native_gemini_empty_stream_returns_200_with_no_body(gateway: Gateway, tmp_path: Path) -> None: + provider: Final = gemini_provider(Reply(content_type="text/event-stream", chunks=())) + with presidio_rig(gateway, tmp_path, provider) as rig: + received: Final = rig.stream(rig.gemini_path(), rig.gemini_body()) + assert received.status == 200, received.text + assert received.text == "" + + +def test_native_gemini_streams_while_presidio_analyzer_is_down(gateway: Gateway, tmp_path: Path) -> None: + frames: Final = (gemini_frame(f"{PERSON} one. "), gemini_frame("two.")) + provider: Final = gemini_provider(Reply(content_type="text/event-stream", chunks=frames)) + with presidio_rig(gateway, tmp_path, provider, analyze=broken) as rig: + received: Final = rig.stream(rig.gemini_path(), rig.gemini_body()) + assert received.status == 200, received.text + assert gemini_texts(b"".join(received.frames)) == (f"{PERSON} one. ", "two.") + assert rig.analyzer.drain() == () + + +def test_native_gemini_unauthenticated_request_is_rejected_before_upstream(gateway: Gateway, tmp_path: Path) -> None: + provider: Final = gemini_provider(Reply(content_type="text/event-stream", chunks=(gemini_frame("never"),))) + with presidio_rig(gateway, tmp_path, provider) as rig: + received: Final = rig.stream(rig.gemini_path(), rig.gemini_body(), key="sk-not-a-key") + assert received.status == 401, received.text + assert rig.upstream.drain() == () + + +def anthropic_provider(chunks: tuple[bytes, ...]) -> Callable[[Request], Reply]: + def provider(request: Request) -> Reply: + assert request.target == "/v1/messages", request.target + return Reply(content_type="text/event-stream", chunks=chunks) + + return provider + + +def test_anthropic_messages_stream_masks_person_in_text_delta(gateway: Gateway, tmp_path: Path) -> None: + identity: Final = "msg_" + uuid.uuid4().hex + provider: Final = anthropic_provider(anthropic_stream(identity, f"{PERSON} designed it.")) + with presidio_rig(gateway, tmp_path, provider) as rig: + received: Final = rig.stream("/v1/messages", rig.messages_body()) + assert received.status == 200, received.text + assert anthropic_text(received) == f"{MASK} designed it." + assert PERSON not in received.text + assert identity in received.text + analyzed: Final = rig.analyzer.drain() + anonymized: Final = rig.anonymizer.drain() + assert len(analyzed) == len(anonymized) == 1 + assert json.loads(analyzed[0].body)["text"] == f"{PERSON} designed it." + + +def test_anthropic_messages_first_frame_split_across_transport_chunks_is_still_masked( + gateway: Gateway, tmp_path: Path +) -> None: + identity: Final = "msg_" + uuid.uuid4().hex + whole: Final = anthropic_stream(identity, f"{PERSON} designed it.") + split_at: Final = whole[0].index(b'"message_') + len(b'"message_') + chunks: Final = (whole[0][:split_at], whole[0][split_at:], *whole[1:]) + with presidio_rig(gateway, tmp_path, anthropic_provider(chunks)) as rig: + received: Final = rig.stream("/v1/messages", rig.messages_body()) + assert received.status == 200, received.text + assert anthropic_text(received) == f"{MASK} designed it." + assert received.text.count("event: message_start") == 1 + + +def test_anthropic_messages_stream_fails_closed_when_analyzer_is_down(gateway: Gateway, tmp_path: Path) -> None: + identity: Final = "msg_" + uuid.uuid4().hex + provider: Final = anthropic_provider(anthropic_stream(identity, f"{PERSON} designed it.")) + with presidio_rig(gateway, tmp_path, provider, analyze=broken) as rig: + received: Final = rig.stream("/v1/messages", rig.messages_body()) + assert PERSON not in received.text, received.text + assert "Presidio analyzer" in received.text, received.text + assert rig.anonymizer.drain() == () + + +def test_anthropic_messages_per_request_guardrails_selects_masking(gateway: Gateway, tmp_path: Path) -> None: + identity: Final = "msg_" + uuid.uuid4().hex + provider: Final = anthropic_provider(anthropic_stream(identity, f"{PERSON} designed it.")) + with presidio_rig(gateway, tmp_path, provider, default_on=False) as rig: + unguarded: Final = rig.stream("/v1/messages", rig.messages_body()) + assert unguarded.status == 200, unguarded.text + assert anthropic_text(unguarded) == f"{PERSON} designed it." + assert rig.analyzer.drain() == () + guarded: Final = rig.stream("/v1/messages", rig.messages_body(guardrails=(rig.guardrail,))) + assert guarded.status == 200, guarded.text + assert anthropic_text(guarded) == f"{MASK} designed it." + assert len(rig.analyzer.drain()) == 1 + + +def openai_provider(identity: str) -> Callable[[Request], Reply]: + def provider(request: Request) -> Reply: + assert request.target == "/v1/chat/completions", request.target + if json.loads(request.body).get("stream"): + return Reply( + content_type="text/event-stream", + chunks=( + openai_frame(identity, {"role": "assistant", "content": ""}), + openai_frame(identity, {"content": f"{PERSON} designed"}), + openai_frame(identity, {"content": " it."}, "stop"), + b"data: [DONE]\n\n", + ), + ) + return Reply( + body=json.dumps( + { + "id": identity, + "object": "chat.completion", + "created": 1, + "model": "gpt-4o-mini", + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": f"{PERSON} designed it."}, + "finish_reason": "stop", + } + ], + "usage": {"prompt_tokens": 11, "completion_tokens": 4, "total_tokens": 15}, + } + ).encode() + ) + + return provider + + +def test_chat_completions_openai_sdk_stream_and_non_stream_are_masked(gateway: Gateway, tmp_path: Path) -> None: + identity: Final = "chatcmpl-" + uuid.uuid4().hex + with presidio_rig(gateway, tmp_path, openai_provider(identity)) as rig: + client: Final = OpenAI(api_key=rig.gateway.key, base_url=f"{rig.gateway.client.base_url}/v1", max_retries=0) + streamed: Final = client.chat.completions.create( + model=rig.openai, messages=[{"role": "user", "content": "who designed it"}], stream=True + ) + pieces: Final = tuple( + chunk.choices[0].delta.content for chunk in streamed if chunk.choices and chunk.choices[0].delta.content + ) + assert "".join(pieces) == f"{MASK} designed it.", pieces + whole: Final = client.chat.completions.create( + model=rig.openai, messages=[{"role": "user", "content": "who designed it"}] + ) + assert whole.id == identity + assert whole.choices[0].message.content == f"{MASK} designed it." + assert len(rig.upstream.drain()) == 2 + assert len(rig.analyzer.drain()) == len(rig.anonymizer.drain()) == 2 + + +def test_mixed_burst_survives_anonymizer_outage_and_recovers(gateway: Gateway, tmp_path: Path) -> None: + outage: Final = threading.Event() + + def flaky_anonymizer(request: Request) -> Reply: + return Reply(status=503, body=b'{"error": "scripted outage"}') if outage.is_set() else anonymizer(request) + + def provider(request: Request) -> Reply: + if request.target == "/v1/messages": + identity: Final = "msg_" + json.loads(request.body)["messages"][0]["content"] + return Reply(content_type="text/event-stream", chunks=anthropic_stream(identity, f"{PERSON} designed it.")) + return Reply( + content_type="text/event-stream", + chunks=(gemini_frame(f"{PERSON} "), gemini_frame("designed it.")), + pause_between_chunks=0.05, + ) + + with presidio_rig(gateway, tmp_path, provider, anonymize=flaky_anonymizer) as rig: + + def gemini_call(index: int) -> tuple[str, str, int]: + received: Final = rig.stream(rig.gemini_path(), rig.gemini_body()) + return ( + "gemini", + f"g{index}", + received.status if gemini_texts(b"".join(received.frames)) == (f"{PERSON} ", "designed it.") else -1, + ) + + def anthropic_call(index: int) -> tuple[str, str, int]: + body: Final = {**rig.messages_body(), "messages": [{"role": "user", "content": f"a{index}"}]} + received: Final = rig.stream("/v1/messages", body) + leaked: Final = PERSON in received.text + return ("anthropic", f"a{index}", -1 if leaked else (1 if MASK in received.text else 0)) + + def phase(offset: int) -> tuple[tuple[str, str, int], ...]: + with ThreadPoolExecutor(max_workers=12) as pool: + futures: Final = tuple( + pool.submit(gemini_call if index % 2 == 0 else anthropic_call, offset + index) + for index in range(12) + ) + return tuple(future.result() for future in futures) + + healthy_before: Final = phase(0) + outage.set() + during: Final = phase(100) + outage.clear() + healthy_after: Final = phase(200) + + for name, results in (("before", healthy_before), ("during", during), ("after", healthy_after)): + assert all(status == 200 for kind, _, status in results if kind == "gemini"), (name, results) + assert all(status == 1 for kind, _, status in healthy_before + healthy_after if kind == "anthropic"), ( + healthy_before, + healthy_after, + ) + assert all(status == 0 for kind, _, status in during if kind == "anthropic"), during + identities: Final = tuple(identity for _, identity, _ in healthy_before + during + healthy_after) + assert len(identities) == len(set(identities)) == 36 + + +def test_native_gemini_keeps_streaming_after_one_worker_is_killed(gateway: Gateway, tmp_path: Path) -> None: + frames: Final = (gemini_frame("alive "), gemini_frame("still.")) + provider: Final = gemini_provider(Reply(content_type="text/event-stream", chunks=frames, pause_between_chunks=0.05)) + with presidio_rig(gateway, tmp_path, provider) as rig: + workers: Final = eventually( + lambda: tuple( + member for member in group_members(rig.proxy.process.pid) if member.pid != rig.proxy.process.pid + ), + lambda members: len(members) >= 2, + seconds=30, + ) + victim: Final = workers[0] + with ThreadPoolExecutor(max_workers=8) as pool: + futures: Final = tuple(pool.submit(rig.stream, rig.gemini_path(), rig.gemini_body()) for _ in range(8)) + victim.send_signal(signal.SIGKILL) + psutil.wait_procs((victim,), timeout=10) + first_wave: Final = tuple(future.result() for future in futures) + survivors: Final = tuple(received for received in first_wave if received.status == 200) + assert survivors, [received.text[:200] for received in first_wave] + assert all(gemini_texts(b"".join(received.frames)) == ("alive ", "still.") for received in survivors) + second_wave: Final = tuple(rig.stream(rig.gemini_path(), rig.gemini_body()) for _ in range(6)) + assert all(received.status == 200 for received in second_wave), [r.text[:200] for r in second_wave] + assert all(gemini_texts(b"".join(received.frames)) == ("alive ", "still.") for received in second_wave) + assert rig.proxy.process.poll() is None diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_presidio.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_presidio.py index 89e72debbf5..1b696669724 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_presidio.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_presidio.py @@ -2261,7 +2261,7 @@ async def test_apply_to_output_streaming_mixed_chunks_flushes_and_warns(): assert mock_logger.warning.call_count == 2 warning_messages = [call.args[0] for call in mock_logger.warning.call_args_list] assert any("mixed stream detected" in msg for msg in warning_messages) - assert any("unknown event objects" in msg for msg in warning_messages) + assert any("Output PII masking was skipped" in msg for msg in warning_messages) # --------------------------------------------------------------------------- @@ -2519,6 +2519,147 @@ async def test_apply_to_output_streaming_anthropic_sse_bytes_without_pii_are_for assert collected == byte_chunks +def _gemini_sse(text: str) -> bytes: + payload = {"candidates": [{"content": {"parts": [{"text": text}], "role": "model"}, "index": 0}]} + return f"data: {json.dumps(payload)}\n\n".encode() + + +@pytest.mark.asyncio +async def test_apply_to_output_streaming_gemini_sse_bytes_are_forwarded_incrementally_until_upstream_aborts(): + guardrail = _OPTIONAL_PresidioPIIMasking( + mock_testing=True, + apply_to_output=True, + mock_redacted_text={"text": ""}, + ) + frames = [_gemini_sse("Partial one from John Smith. "), _gemini_sse("Partial two. ")] + collected: list[object] = [] + + async def mock_stream(): + for frame in frames: + yield frame + raise ConnectionError("upstream closed mid-stream") + + async def collect() -> None: + async for chunk in guardrail.async_post_call_streaming_iterator_hook( + user_api_key_dict=UserAPIKeyAuth(api_key="test-key"), + response=mock_stream(), + request_data={}, + ): + collected.append(chunk) + + with pytest.raises(ConnectionError): + await collect() + + assert collected == frames + + +@pytest.mark.asyncio +async def test_apply_to_output_streaming_anthropic_first_frame_split_across_transport_chunks_is_still_masked(): + guardrail = _OPTIONAL_PresidioPIIMasking( + mock_testing=True, + apply_to_output=True, + mock_redacted_text={"text": ""}, + ) + message_start = _anthropic_sse( + "message_start", + {"type": "message_start", "message": {"id": "msg_1", "model": "claude", "content": [], "usage": {}}}, + ) + split_at = message_start.index(b'"message_') + len(b'"message_') + byte_chunks = [ + message_start[:split_at], + message_start[split_at:], + _anthropic_sse( + "content_block_start", + {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}, + ), + _anthropic_sse( + "content_block_delta", + {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "John Smith"}}, + ), + _anthropic_sse("content_block_stop", {"type": "content_block_stop", "index": 0}), + _anthropic_sse("message_delta", {"type": "message_delta", "delta": {"stop_reason": "end_turn"}, "usage": {}}), + _anthropic_sse("message_stop", {"type": "message_stop"}), + ] + + async def mock_stream(): + for b in byte_chunks: + yield b + + collected = [] + async for chunk in guardrail.async_post_call_streaming_iterator_hook( + user_api_key_dict=UserAPIKeyAuth(api_key="test-key"), + response=mock_stream(), + request_data={}, + ): + collected.append(chunk) + + joined = b"".join(collected).decode() + assert "John Smith" not in joined, joined + assert "".join(text for _, text in _anthropic_text_deltas(collected)) == "" + assert joined.count("event: message_start") == 1 + + +@pytest.mark.asyncio +async def test_apply_to_output_streaming_gemini_first_frame_split_across_transport_chunks_streams_incrementally(): + guardrail = _OPTIONAL_PresidioPIIMasking( + mock_testing=True, + apply_to_output=True, + mock_redacted_text={"text": ""}, + ) + first = _gemini_sse("Partial one from John Smith. ") + second = _gemini_sse("Partial two. ") + collected: list[object] = [] + + async def mock_stream(): + yield first[:20] + yield first[20:] + yield second + raise ConnectionError("upstream closed mid-stream") + + async def collect() -> None: + async for chunk in guardrail.async_post_call_streaming_iterator_hook( + user_api_key_dict=UserAPIKeyAuth(api_key="test-key"), + response=mock_stream(), + request_data={}, + ): + collected.append(chunk) + + with pytest.raises(ConnectionError): + await collect() + + assert collected == [first, second] + + +@pytest.mark.asyncio +async def test_apply_to_output_streaming_unterminated_first_frame_is_released_once_it_exceeds_the_cap(): + guardrail = _OPTIONAL_PresidioPIIMasking( + mock_testing=True, + apply_to_output=True, + mock_redacted_text={"text": ""}, + ) + piece = b"data: " + b"x" * 1023 + b"\n" + pieces_to_cap = -(-(64 * 1024) // len(piece)) + released_at: list[int] = [] + + async def mock_stream(): + for index in range(pieces_to_cap * 4): + if collected: + released_at.append(index) + yield piece + + collected: list[object] = [] + async for chunk in guardrail.async_post_call_streaming_iterator_hook( + user_api_key_dict=UserAPIKeyAuth(api_key="test-key"), + response=mock_stream(), + request_data={}, + ): + collected.append(chunk) + + assert released_at, "nothing reached the caller before the upstream finished" + assert released_at[0] == pieces_to_cap, released_at[:3] + assert b"".join(collected) == piece * (pieces_to_cap * 4) + + @pytest.mark.asyncio async def test_apply_to_output_streaming_anthropic_sse_bytes_fail_closed_when_presidio_is_unreachable(): """ From 0fb999b613b126d077b7eead27bb709ea449414a Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 23 Sep 2026 22:18:28 -0700 Subject: [PATCH 013/120] fix(cost-map): halve openrouter deepseek-v4-flash-0731 output price (#42881) Co-authored-by: kerry Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/model_prices_and_context_window_backup.json | 2 +- model_prices_and_context_window.json | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 66b72090631..6d577fcf41b 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -65505,7 +65505,7 @@ }, "openrouter/deepseek/deepseek-v4-flash-0731": { "input_cost_per_token": 4e-08, - "output_cost_per_token": 6.4e-07, + "output_cost_per_token": 3.2e-07, "cache_read_input_token_cost": 1.6e-08, "litellm_provider": "openrouter", "max_input_tokens": 1310720, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 66b72090631..6d577fcf41b 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -65505,7 +65505,7 @@ }, "openrouter/deepseek/deepseek-v4-flash-0731": { "input_cost_per_token": 4e-08, - "output_cost_per_token": 6.4e-07, + "output_cost_per_token": 3.2e-07, "cache_read_input_token_cost": 1.6e-08, "litellm_provider": "openrouter", "max_input_tokens": 1310720, From 001179a6368b416810b9c47f3bee40861e85df7b Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 23 Sep 2026 22:21:05 -0700 Subject: [PATCH 014/120] fix(cost-map): sync vertex-ai deprecation dates from Vertex model lifecycle pages (#42882) Co-authored-by: kerry Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- ...model_prices_and_context_window_backup.json | 18 ++++++++++++++++-- model_prices_and_context_window.json | 18 ++++++++++++++++-- 2 files changed, 32 insertions(+), 4 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 6d577fcf41b..e9046f345a1 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -26057,7 +26057,7 @@ "supports_image_size": false }, "gemini-2.5-flash-image": { - "deprecation_date": "2026-10-02", + "deprecation_date": "2027-03-15", "cache_read_input_token_cost": 3e-08, "input_cost_per_audio_token": 1e-06, "input_cost_per_token": 3e-07, @@ -49110,6 +49110,7 @@ }, "vertex_ai/deepseek-ai/deepseek-v3.1-maas": { "cache_read_input_token_cost": 6e-08, + "deprecation_date": "2026-10-21", "input_cost_per_token": 6e-07, "input_cost_per_token_batches": 3e-07, "litellm_provider": "vertex_ai-deepseek_models", @@ -49131,6 +49132,7 @@ }, "vertex_ai/deepseek-ai/deepseek-v3.2-maas": { "cache_read_input_token_cost": 5.6e-08, + "deprecation_date": "2026-10-21", "input_cost_per_token": 5.6e-07, "input_cost_per_token_batches": 2.8e-07, "litellm_provider": "vertex_ai-deepseek_models", @@ -49151,6 +49153,7 @@ "supports_tool_choice": true }, "vertex_ai/deepseek-ai/deepseek-r1-0528-maas": { + "deprecation_date": "2026-10-21", "input_cost_per_token": 1.35e-06, "input_cost_per_token_batches": 6.75e-07, "litellm_provider": "vertex_ai-deepseek_models", @@ -49171,7 +49174,7 @@ "supports_tool_choice": true }, "vertex_ai/gemini-2.5-flash-image": { - "deprecation_date": "2026-10-02", + "deprecation_date": "2027-03-15", "cache_read_input_token_cost": 3e-08, "input_cost_per_audio_token": 1e-06, "input_cost_per_token": 3e-07, @@ -49858,6 +49861,7 @@ }, "vertex_ai/minimaxai/minimax-m2-maas": { "cache_read_input_token_cost": 3e-08, + "deprecation_date": "2026-10-21", "input_cost_per_token": 3e-07, "litellm_provider": "vertex_ai-minimax_models", "max_input_tokens": 196608, @@ -49871,6 +49875,7 @@ }, "vertex_ai/moonshotai/kimi-k2-thinking-maas": { "cache_read_input_token_cost": 6e-08, + "deprecation_date": "2026-10-21", "input_cost_per_token": 6e-07, "litellm_provider": "vertex_ai-moonshot_models", "max_input_tokens": 256000, @@ -49885,6 +49890,7 @@ }, "vertex_ai/zai-org/glm-4.7-maas": { "cache_read_input_token_cost": 6e-08, + "deprecation_date": "2026-10-21", "input_cost_per_token": 6e-07, "litellm_provider": "vertex_ai-zai_models", "max_input_tokens": 200000, @@ -49902,6 +49908,7 @@ }, "vertex_ai/zai-org/glm-5-maas": { "cache_read_input_token_cost": 1e-07, + "deprecation_date": "2026-10-21", "input_cost_per_token": 1e-06, "litellm_provider": "vertex_ai-zai_models", "max_input_tokens": 200000, @@ -50067,6 +50074,7 @@ "source": "https://cloud.google.com/generative-ai-app-builder/pricing" }, "vertex_ai/deepseek-ai/deepseek-ocr-maas": { + "deprecation_date": "2026-10-21", "litellm_provider": "vertex_ai", "mode": "ocr", "input_cost_per_token": 3e-07, @@ -50107,6 +50115,7 @@ "supports_reasoning": true }, "vertex_ai/openai/gpt-oss-20b-maas": { + "deprecation_date": "2026-10-21", "input_cost_per_token": 7e-08, "litellm_provider": "vertex_ai-openai_models", "max_input_tokens": 131072, @@ -50237,6 +50246,7 @@ "supports_vision": true }, "vertex_ai/qwen/qwen3-235b-a22b-instruct-2507-maas": { + "deprecation_date": "2026-10-21", "input_cost_per_token": 2.2e-07, "input_cost_per_token_batches": 1.1e-07, "litellm_provider": "vertex_ai-qwen_models", @@ -50256,6 +50266,7 @@ }, "vertex_ai/qwen/qwen3-coder-480b-a35b-instruct-maas": { "cache_read_input_token_cost": 2.2e-08, + "deprecation_date": "2026-10-21", "input_cost_per_token": 2.2e-07, "input_cost_per_token_batches": 1.1e-07, "litellm_provider": "vertex_ai-qwen_models", @@ -50273,6 +50284,7 @@ "supports_tool_choice": true }, "vertex_ai/qwen/qwen3-next-80b-a3b-instruct-maas": { + "deprecation_date": "2026-10-21", "input_cost_per_token": 1.5e-07, "litellm_provider": "vertex_ai-qwen_models", "max_input_tokens": 262144, @@ -50288,6 +50300,7 @@ "supports_tool_choice": true }, "vertex_ai/qwen/qwen3-next-80b-a3b-thinking-maas": { + "deprecation_date": "2026-10-21", "input_cost_per_token": 1.5e-07, "litellm_provider": "vertex_ai-qwen_models", "max_input_tokens": 262144, @@ -75522,6 +75535,7 @@ "supports_tool_choice": true }, "vertex_ai/meta/llama-3.3-70b-instruct-maas": { + "deprecation_date": "2026-10-21", "input_cost_per_token": 7.2e-07, "input_cost_per_token_batches": 3.6e-07, "litellm_provider": "vertex_ai-llama_models", diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 6d577fcf41b..e9046f345a1 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -26057,7 +26057,7 @@ "supports_image_size": false }, "gemini-2.5-flash-image": { - "deprecation_date": "2026-10-02", + "deprecation_date": "2027-03-15", "cache_read_input_token_cost": 3e-08, "input_cost_per_audio_token": 1e-06, "input_cost_per_token": 3e-07, @@ -49110,6 +49110,7 @@ }, "vertex_ai/deepseek-ai/deepseek-v3.1-maas": { "cache_read_input_token_cost": 6e-08, + "deprecation_date": "2026-10-21", "input_cost_per_token": 6e-07, "input_cost_per_token_batches": 3e-07, "litellm_provider": "vertex_ai-deepseek_models", @@ -49131,6 +49132,7 @@ }, "vertex_ai/deepseek-ai/deepseek-v3.2-maas": { "cache_read_input_token_cost": 5.6e-08, + "deprecation_date": "2026-10-21", "input_cost_per_token": 5.6e-07, "input_cost_per_token_batches": 2.8e-07, "litellm_provider": "vertex_ai-deepseek_models", @@ -49151,6 +49153,7 @@ "supports_tool_choice": true }, "vertex_ai/deepseek-ai/deepseek-r1-0528-maas": { + "deprecation_date": "2026-10-21", "input_cost_per_token": 1.35e-06, "input_cost_per_token_batches": 6.75e-07, "litellm_provider": "vertex_ai-deepseek_models", @@ -49171,7 +49174,7 @@ "supports_tool_choice": true }, "vertex_ai/gemini-2.5-flash-image": { - "deprecation_date": "2026-10-02", + "deprecation_date": "2027-03-15", "cache_read_input_token_cost": 3e-08, "input_cost_per_audio_token": 1e-06, "input_cost_per_token": 3e-07, @@ -49858,6 +49861,7 @@ }, "vertex_ai/minimaxai/minimax-m2-maas": { "cache_read_input_token_cost": 3e-08, + "deprecation_date": "2026-10-21", "input_cost_per_token": 3e-07, "litellm_provider": "vertex_ai-minimax_models", "max_input_tokens": 196608, @@ -49871,6 +49875,7 @@ }, "vertex_ai/moonshotai/kimi-k2-thinking-maas": { "cache_read_input_token_cost": 6e-08, + "deprecation_date": "2026-10-21", "input_cost_per_token": 6e-07, "litellm_provider": "vertex_ai-moonshot_models", "max_input_tokens": 256000, @@ -49885,6 +49890,7 @@ }, "vertex_ai/zai-org/glm-4.7-maas": { "cache_read_input_token_cost": 6e-08, + "deprecation_date": "2026-10-21", "input_cost_per_token": 6e-07, "litellm_provider": "vertex_ai-zai_models", "max_input_tokens": 200000, @@ -49902,6 +49908,7 @@ }, "vertex_ai/zai-org/glm-5-maas": { "cache_read_input_token_cost": 1e-07, + "deprecation_date": "2026-10-21", "input_cost_per_token": 1e-06, "litellm_provider": "vertex_ai-zai_models", "max_input_tokens": 200000, @@ -50067,6 +50074,7 @@ "source": "https://cloud.google.com/generative-ai-app-builder/pricing" }, "vertex_ai/deepseek-ai/deepseek-ocr-maas": { + "deprecation_date": "2026-10-21", "litellm_provider": "vertex_ai", "mode": "ocr", "input_cost_per_token": 3e-07, @@ -50107,6 +50115,7 @@ "supports_reasoning": true }, "vertex_ai/openai/gpt-oss-20b-maas": { + "deprecation_date": "2026-10-21", "input_cost_per_token": 7e-08, "litellm_provider": "vertex_ai-openai_models", "max_input_tokens": 131072, @@ -50237,6 +50246,7 @@ "supports_vision": true }, "vertex_ai/qwen/qwen3-235b-a22b-instruct-2507-maas": { + "deprecation_date": "2026-10-21", "input_cost_per_token": 2.2e-07, "input_cost_per_token_batches": 1.1e-07, "litellm_provider": "vertex_ai-qwen_models", @@ -50256,6 +50266,7 @@ }, "vertex_ai/qwen/qwen3-coder-480b-a35b-instruct-maas": { "cache_read_input_token_cost": 2.2e-08, + "deprecation_date": "2026-10-21", "input_cost_per_token": 2.2e-07, "input_cost_per_token_batches": 1.1e-07, "litellm_provider": "vertex_ai-qwen_models", @@ -50273,6 +50284,7 @@ "supports_tool_choice": true }, "vertex_ai/qwen/qwen3-next-80b-a3b-instruct-maas": { + "deprecation_date": "2026-10-21", "input_cost_per_token": 1.5e-07, "litellm_provider": "vertex_ai-qwen_models", "max_input_tokens": 262144, @@ -50288,6 +50300,7 @@ "supports_tool_choice": true }, "vertex_ai/qwen/qwen3-next-80b-a3b-thinking-maas": { + "deprecation_date": "2026-10-21", "input_cost_per_token": 1.5e-07, "litellm_provider": "vertex_ai-qwen_models", "max_input_tokens": 262144, @@ -75522,6 +75535,7 @@ "supports_tool_choice": true }, "vertex_ai/meta/llama-3.3-70b-instruct-maas": { + "deprecation_date": "2026-10-21", "input_cost_per_token": 7.2e-07, "input_cost_per_token_batches": 3.6e-07, "litellm_provider": "vertex_ai-llama_models", From dfb61e9b48bdadf80eab828c0e2b64d2a1d17445 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 23 Sep 2026 22:24:37 -0700 Subject: [PATCH 015/120] fix(cost-map): add azure gpt-realtime-mini deprecation date (#42883) Co-authored-by: kerry Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/model_prices_and_context_window_backup.json | 1 + model_prices_and_context_window.json | 1 + 2 files changed, 2 insertions(+) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index e9046f345a1..1d15bea40ac 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -5991,6 +5991,7 @@ "cache_creation_input_audio_token_cost": 3e-07, "cache_read_input_audio_token_cost": 3e-07, "cache_read_input_token_cost": 6e-08, + "deprecation_date": "2027-06-15", "input_cost_per_audio_token": 1e-05, "input_cost_per_image_token": 8e-07, "input_cost_per_token": 6e-07, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index e9046f345a1..1d15bea40ac 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -5991,6 +5991,7 @@ "cache_creation_input_audio_token_cost": 3e-07, "cache_read_input_audio_token_cost": 3e-07, "cache_read_input_token_cost": 6e-08, + "deprecation_date": "2027-06-15", "input_cost_per_audio_token": 1e-05, "input_cost_per_image_token": 8e-07, "input_cost_per_token": 6e-07, From 6b1ec4cd3abeae890b52e506cf96aa592b8935af Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 23 Sep 2026 22:26:59 -0700 Subject: [PATCH 016/120] fix(models): correct fireworks kimi k3 us pricing to the published rate (#42884) Co-authored-by: kerry Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/model_prices_and_context_window_backup.json | 12 ++++++------ model_prices_and_context_window.json | 12 ++++++------ 2 files changed, 12 insertions(+), 12 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 1d15bea40ac..b75e7ec7430 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -59931,14 +59931,14 @@ "supports_vision": true }, "fireworks_ai/kimi-k3-us": { - "cache_read_input_token_cost": 3.3e-07, - "input_cost_per_token": 3.3e-06, + "cache_read_input_token_cost": 4.5e-07, + "input_cost_per_token": 4.5e-06, "litellm_provider": "fireworks_ai", "max_input_tokens": 1048576, "max_output_tokens": 131072, "max_tokens": 131072, "mode": "chat", - "output_cost_per_token": 1.65e-05, + "output_cost_per_token": 2.25e-05, "reasoning_effort_levels": [ "low", "high", @@ -60139,14 +60139,14 @@ "supports_vision": true }, "fireworks_ai/accounts/fireworks/routers/kimi-k3-us": { - "cache_read_input_token_cost": 3.3e-07, - "input_cost_per_token": 3.3e-06, + "cache_read_input_token_cost": 4.5e-07, + "input_cost_per_token": 4.5e-06, "litellm_provider": "fireworks_ai", "max_input_tokens": 1048576, "max_output_tokens": 131072, "max_tokens": 131072, "mode": "chat", - "output_cost_per_token": 1.65e-05, + "output_cost_per_token": 2.25e-05, "reasoning_effort_levels": [ "low", "high", diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 1d15bea40ac..b75e7ec7430 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -59931,14 +59931,14 @@ "supports_vision": true }, "fireworks_ai/kimi-k3-us": { - "cache_read_input_token_cost": 3.3e-07, - "input_cost_per_token": 3.3e-06, + "cache_read_input_token_cost": 4.5e-07, + "input_cost_per_token": 4.5e-06, "litellm_provider": "fireworks_ai", "max_input_tokens": 1048576, "max_output_tokens": 131072, "max_tokens": 131072, "mode": "chat", - "output_cost_per_token": 1.65e-05, + "output_cost_per_token": 2.25e-05, "reasoning_effort_levels": [ "low", "high", @@ -60139,14 +60139,14 @@ "supports_vision": true }, "fireworks_ai/accounts/fireworks/routers/kimi-k3-us": { - "cache_read_input_token_cost": 3.3e-07, - "input_cost_per_token": 3.3e-06, + "cache_read_input_token_cost": 4.5e-07, + "input_cost_per_token": 4.5e-06, "litellm_provider": "fireworks_ai", "max_input_tokens": 1048576, "max_output_tokens": 131072, "max_tokens": 131072, "mode": "chat", - "output_cost_per_token": 1.65e-05, + "output_cost_per_token": 2.25e-05, "reasoning_effort_levels": [ "low", "high", From 0094cff47a9890a434a7d47b774f37e5f779d3fc Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 23 Sep 2026 22:38:31 -0700 Subject: [PATCH 017/120] fix(cost-map): add azure gpt-realtime-mini-2025-10-06 deprecation date (#42885) Co-authored-by: kerry Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/model_prices_and_context_window_backup.json | 1 + model_prices_and_context_window.json | 1 + 2 files changed, 2 insertions(+) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index b75e7ec7430..0b606f5db4e 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -6025,6 +6025,7 @@ "cache_creation_input_audio_token_cost": 3e-07, "cache_read_input_audio_token_cost": 3e-07, "cache_read_input_token_cost": 6e-08, + "deprecation_date": "2027-04-06", "input_cost_per_audio_token": 1e-05, "input_cost_per_image_token": 8e-07, "input_cost_per_token": 6e-07, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index b75e7ec7430..0b606f5db4e 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -6025,6 +6025,7 @@ "cache_creation_input_audio_token_cost": 3e-07, "cache_read_input_audio_token_cost": 3e-07, "cache_read_input_token_cost": 6e-08, + "deprecation_date": "2027-04-06", "input_cost_per_audio_token": 1e-05, "input_cost_per_image_token": 8e-07, "input_cost_per_token": 6e-07, From 251fdf03089be0bb3d27c5306300b0e6df572a4f Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 23 Sep 2026 22:44:23 -0700 Subject: [PATCH 018/120] fix(logging): scan the exceeded budget wording linearly so a crafted error message cannot stall the proxy (#42778) * fix(logging): bound the exceeded budget regex so a crafted error message cannot stall the proxy Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(integration): audit normalized_error clustering on long messages with a real two worker proxy Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(integration): poll for the budget denial and correlate upstream 503 bursts by request marker Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(logging): replace the bounded exceeded budget regex with a linear scan that keeps the original semantics Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(integration): compare upstream error wording against the decoded message Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(integration): tolerate a reaped worker while listing proxy children Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: yucheng Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../litellm_core_utils/error_normalization.py | 14 +- .../test_normalized_error_long_message.py | 463 ++++++++++++++++++ .../test_error_normalization.py | 35 ++ 3 files changed, 510 insertions(+), 2 deletions(-) create mode 100644 tests/integration/spend/test_normalized_error_long_message.py diff --git a/litellm/litellm_core_utils/error_normalization.py b/litellm/litellm_core_utils/error_normalization.py index 2a9ff748899..61eb47a9675 100644 --- a/litellm/litellm_core_utils/error_normalization.py +++ b/litellm/litellm_core_utils/error_normalization.py @@ -61,7 +61,7 @@ class _HasProxyErrorType(Protocol): _MESSAGE_PATTERNS: Final[tuple[tuple[re.Pattern[str], str], ...]] = ( ( - re.compile(r"budget has been exceeded|max budget|exceeded.*budget|crossed budget", re.IGNORECASE), + re.compile(r"budget has been exceeded|max budget|crossed budget", re.IGNORECASE), BUDGET_EXCEEDED, ), (re.compile(r"no healthy deployments?|no deployments available", re.IGNORECASE), NO_HEALTHY_DEPLOYMENTS), @@ -155,6 +155,14 @@ _CLASS_CODE_TABLE: Final[tuple[tuple[tuple[type[BaseException], ...], str], ...] ) +def _exceeded_before_budget(message: str) -> bool: + """Linear-time equivalent of ``re.search(r"exceeded.*budget", message, re.IGNORECASE)``.""" + return any( + (start := line.find("exceeded")) != -1 and line.find("budget", start + len("exceeded")) != -1 + for line in message.lower().split("\n") + ) + + def _classify_by_message(message: str, patterns: tuple[tuple[re.Pattern[str], str], ...]) -> str | None: return next((code for pattern, code in patterns if pattern.search(message)), None) @@ -183,7 +191,9 @@ def normalize_error(exc: Exception | None, status_code: str, message: str) -> st by_proxy_type: Final = _PROXY_ERROR_TYPE_MAP.get(proxy_type) if isinstance(proxy_type, str) else None if by_proxy_type is not None: return by_proxy_type - by_message: Final = _classify_by_message(message, _MESSAGE_PATTERNS) + by_message: Final = ( + BUDGET_EXCEEDED if _exceeded_before_budget(message) else _classify_by_message(message, _MESSAGE_PATTERNS) + ) if by_message is not None: return by_message by_class: Final = _classify_by_class(exc) diff --git a/tests/integration/spend/test_normalized_error_long_message.py b/tests/integration/spend/test_normalized_error_long_message.py new file mode 100644 index 00000000000..accf160b03f --- /dev/null +++ b/tests/integration/spend/test_normalized_error_long_message.py @@ -0,0 +1,463 @@ +import asyncio +import json +import os +import signal +import threading +import time +import uuid +from collections.abc import Callable, Iterator, Mapping +from concurrent.futures import FIRST_COMPLETED, ThreadPoolExecutor, wait +from contextlib import contextmanager +from dataclasses import dataclass +from hashlib import sha256 +from pathlib import Path +from typing import Final + +import anthropic +import httpx +import openai +import psutil +import pytest +from integration._support.client import Gateway, eventually, object_value, string_value +from integration._support.database import read_rows +from integration._support.process import OwnedProxy, owned_proxy_process +from integration._support.wire import Reply, Request, wire_server +from pydantic import JsonValue + +CRAFTED_MODEL: Final = ("exceeded " * 32_000)[:288_000] +HOSTILE_5KB_MODEL: Final = ("exceeded budget " * 400)[:5_000] +FAST_SECONDS: Final = 10.0 +LIVELINESS_MAX_SECONDS: Final = 5.0 +ROW_SECONDS: Final = 70 +CHAT: Final = "/v1/chat/completions" +MESSAGES: Final = "/v1/messages" +RESPONSES: Final = "/v1/responses" + + +def _body(path: str, model: str, marker: str, stream: bool = False) -> dict[str, JsonValue]: + content: Final = f"normalized error audit {marker}" + match path: + case "/v1/messages": + return { + "model": model, + "max_tokens": 8, + "messages": [{"role": "user", "content": content}], + "stream": stream, + } + case "/v1/responses": + return {"model": model, "input": content, "stream": stream} + case _: + return {"model": model, "messages": [{"role": "user", "content": content}], "stream": stream} + + +@dataclass(frozen=True, slots=True) +class _Timed: + response: httpx.Response + seconds: float + + +def _timed_post(client: httpx.Client, path: str, body: Mapping[str, JsonValue], key: str) -> _Timed: + started: Final = time.perf_counter() + response: Final = client.post(path, json=body, headers={"Authorization": f"Bearer {key}"}) + return _Timed(response, time.perf_counter() - started) + + +@contextmanager +def _patient_client(gateway: Gateway) -> Iterator[httpx.Client]: + with httpx.Client(base_url=str(gateway.client.base_url), timeout=120, trust_env=False) as client: + yield client + + +def _error_information(call_id: str) -> dict[str, JsonValue]: + rows: Final = eventually( + lambda: read_rows( + "SELECT status, metadata->'error_information' AS info FROM \"LiteLLM_SpendLogs\" WHERE request_id=%s", + (call_id,), + ), + lambda values: len(values) == 1, + seconds=ROW_SECONDS, + ) + assert rows[0]["status"] == "failure", rows + return object_value(rows[0]["info"]) + + +def _assert_crafted_failure(timed: _Timed, expected_status: int = 400) -> None: + response: Final = timed.response + assert response.status_code == expected_status, response.text[:300] + assert "Invalid model name passed in" in response.text, response.text[:300] + assert timed.seconds < FAST_SECONDS, f"crafted 288 KB model took {timed.seconds:.2f}s" + info: Final = _error_information(response.headers["x-litellm-call-id"]) + assert info["normalized_error"] == "400_INVALID_REQUEST" and info["error_code"] == "400", info + + +@pytest.mark.parametrize("path", [CHAT, MESSAGES, RESPONSES]) +def test_crafted_288kb_model_fails_fast_and_logs_invalid_request(gateway: Gateway, path: str) -> None: + with gateway.scenario() as scenario, _patient_client(gateway) as client: + key: Final = scenario.key() + _assert_crafted_failure(_timed_post(client, path, _body(path, CRAFTED_MODEL, uuid.uuid4().hex), key)) + + +@pytest.mark.parametrize("path", [CHAT, MESSAGES, RESPONSES]) +def test_crafted_288kb_model_with_stream_true_fails_fast(gateway: Gateway, path: str) -> None: + with gateway.scenario() as scenario, _patient_client(gateway) as client: + key: Final = scenario.key() + body: Final = _body(path, CRAFTED_MODEL, uuid.uuid4().hex, stream=True) + _assert_crafted_failure(_timed_post(client, path, body, key)) + + +def test_crafted_288kb_model_through_async_openai_sdk_fails_fast(gateway: Gateway) -> None: + async def call(key: str) -> tuple[openai.BadRequestError, float]: + client: Final = openai.AsyncOpenAI(base_url=f"{gateway.client.base_url}/v1", api_key=key, timeout=120) + started: Final = time.perf_counter() + try: + with pytest.raises(openai.BadRequestError) as raised: + await client.chat.completions.create( + model=CRAFTED_MODEL, messages=[{"role": "user", "content": f"audit {uuid.uuid4().hex}"}] + ) + return raised.value, time.perf_counter() - started + finally: + await client.close() + + with gateway.scenario() as scenario: + error, seconds = asyncio.run(call(scenario.key())) + assert seconds < FAST_SECONDS, f"crafted 288 KB model took {seconds:.2f}s" + assert "Invalid model name passed in" in str(error), str(error)[:300] + info: Final = _error_information(error.response.headers["x-litellm-call-id"]) + assert info["normalized_error"] == "400_INVALID_REQUEST", info + + +def test_crafted_288kb_model_through_anthropic_sdk_fails_fast(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + client: Final = anthropic.Anthropic(base_url=str(gateway.client.base_url), api_key=scenario.key(), timeout=120) + started: Final = time.perf_counter() + with pytest.raises(anthropic.BadRequestError) as raised: + client.messages.create( + model=CRAFTED_MODEL, max_tokens=8, messages=[{"role": "user", "content": f"audit {uuid.uuid4().hex}"}] + ) + seconds: Final = time.perf_counter() - started + assert seconds < FAST_SECONDS, f"crafted 288 KB model took {seconds:.2f}s" + assert "Invalid model name passed in" in str(raised.value), str(raised.value)[:300] + info: Final = _error_information(raised.value.response.headers["x-litellm-call-id"]) + assert info["normalized_error"] == "400_INVALID_REQUEST", info + + +def _poll_liveliness(client: httpx.Client, stop: threading.Event) -> list[float]: + latencies: Final[list[float]] = [] # mutable-ok: thread-local sample buffer drained once by the caller + while not stop.is_set(): + started = time.perf_counter() + assert client.get("/health/liveliness").status_code == 200 + latencies.append(time.perf_counter() - started) + stop.wait(0.1) + return latencies + + +def _cmdline(process: psutil.Process) -> str: + try: + return " ".join(process.cmdline()) + except psutil.Error: + return "" + + +def _worker_pids(owned: OwnedProxy) -> tuple[int, ...]: + return tuple(child.pid for child in psutil.Process(owned.process.pid).children() if "spawn_main" in _cmdline(child)) + + +def test_two_concurrent_crafted_requests_do_not_stall_liveliness_on_a_two_worker_proxy( + gateway: Gateway, tmp_path: Path +) -> None: + with owned_proxy_process(gateway, tmp_path, {}, workers=2) as owned, _patient_client(owned.gateway) as client: + assert len(_worker_pids(owned)) == 2, _worker_pids(owned) + stop: Final = threading.Event() + with ThreadPoolExecutor(max_workers=3) as pool: + liveliness: Final = pool.submit(_poll_liveliness, client, stop) + crafted: Final = tuple( + pool.submit(_timed_post, client, CHAT, _body(CHAT, CRAFTED_MODEL, uuid.uuid4().hex), gateway.key) + for _ in range(2) + ) + results: Final = tuple(future.result() for future in crafted) + stop.set() + latencies: Final = liveliness.result() + for timed in results: + _assert_crafted_failure(timed) + assert latencies and max(latencies) < LIVELINESS_MAX_SECONDS, f"liveliness max {max(latencies):.2f}s" + + +def _completion(request: Request) -> Reply: + body: Final = object_value(json.loads(request.body or b"{}")) + return Reply( + body=json.dumps( + { + "id": "chatcmpl-" + uuid.uuid4().hex, + "object": "chat.completion", + "created": 1, + "model": body.get("model", "unknown"), + "choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 20, "completion_tokens": 20, "total_tokens": 40}, + } + ).encode() + ) + + +def _rate_limited(message: str) -> Callable[[Request], Reply]: + def respond(_request: Request) -> Reply: + return Reply( + status=429, + body=json.dumps({"error": {"message": message, "type": "rate_limit_error", "code": "429"}}).encode(), + ) + + return respond + + +def _budget_denied_row(key: str) -> dict[str, JsonValue]: + rows: Final = eventually( + lambda: read_rows( + "SELECT request_id, metadata->'error_information' AS info FROM \"LiteLLM_SpendLogs\" " + "WHERE api_key=%s AND status='failure'", + (sha256(key.encode()).hexdigest(),), + ), + lambda values: len(values) == 1, + seconds=ROW_SECONDS, + ) + return object_value(rows[0]["info"]) + + +def _exhaust( + gateway: Gateway, client: httpx.Client, model: str, key: str, table: str, column: str, identity: str +) -> None: + first: Final = gateway.chat(model, key=key, text=f"spend {uuid.uuid4().hex}") + assert object_value(first["usage"])["total_tokens"] == 40, first + eventually( + lambda: read_rows(f'SELECT spend FROM "{table}" WHERE {column}=%s', (identity,)), + lambda values: len(values) == 1 and float(string_value(str(values[0]["spend"]))) >= 0.06, + seconds=ROW_SECONDS, + ) + denied: Final = eventually( + lambda: client.post( + CHAT, json=_body(CHAT, model, uuid.uuid4().hex), headers={"Authorization": f"Bearer {key}"} + ), + lambda response: response.status_code in {400, 422}, + seconds=ROW_SECONDS, + ) + assert denied.json()["error"]["type"] == "budget_exceeded", denied.text + info: Final = _budget_denied_row(key) + assert info["normalized_error"] == "429_BUDGET_EXCEEDED", info + assert "budget" in string_value(info["error_message"]).lower(), info + + +def test_exhausted_key_budget_denial_clusters_as_budget_exceeded(gateway: Gateway) -> None: + with gateway.scenario() as scenario, _patient_client(gateway) as client: + model: Final = scenario.model(input_cost_per_token=0.001, output_cost_per_token=0.002) + key: Final = scenario.key(models=[model], max_budget=0.06) + _exhaust(gateway, client, model, key, "LiteLLM_VerificationToken", "token", sha256(key.encode()).hexdigest()) + + +def test_exhausted_team_budget_denial_clusters_as_budget_exceeded(gateway: Gateway) -> None: + with gateway.scenario() as scenario, _patient_client(gateway) as client: + model: Final = scenario.model(input_cost_per_token=0.001, output_cost_per_token=0.002) + team: Final = scenario.team(models=[model], max_budget=0.06) + key: Final = scenario.key(team_id=team, models=[model]) + _exhaust(gateway, client, model, key, "LiteLLM_TeamTable", "team_id", team) + + +def _upstream_failure_row(gateway: Gateway, message: str) -> dict[str, JsonValue]: + with wire_server(_rate_limited(message)) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(api_base=wire.url + "/v1", num_retries=0) + key: Final = scenario.key(models=[model]) + failed: Final = gateway.request("POST", CHAT, _body(CHAT, model, uuid.uuid4().hex), key=key) + assert failed.status_code == 429 and message in failed.json()["error"]["message"], failed.text[:300] + assert len(wire.drain()) == 1 + return _error_information(failed.headers["x-litellm-call-id"]) + + +@pytest.mark.parametrize( + "message", + [ + "Budget has been exceeded! Current cost: 11.0, Max budget: 10.0", + "ExceededBudget: User=audit over budget. Spend=12.5, Budget=10.0", + "Exceeded budget for provider openai: 105.2 >= 100.0", + "exceeded" + "x" * 64 + "budget", + ], +) +def test_upstream_budget_wording_clusters_as_budget_exceeded(gateway: Gateway, message: str) -> None: + info: Final = _upstream_failure_row(gateway, message) + assert info["normalized_error"] == "429_BUDGET_EXCEEDED", info + + +def test_upstream_exceeded_and_budget_65_chars_apart_still_clusters_as_budget_exceeded(gateway: Gateway) -> None: + info: Final = _upstream_failure_row(gateway, "exceeded" + "x" * 65 + "budget") + assert info["normalized_error"] == "429_BUDGET_EXCEEDED", info + + +def test_upstream_exceeded_and_budget_on_different_lines_cluster_by_exception_class(gateway: Gateway) -> None: + info: Final = _upstream_failure_row(gateway, "exceeded the limit\nbudget unaffected") + assert info["normalized_error"] == "429_RATE_LIMIT_EXCEEDED", info + + +def test_hostile_model_values_are_rejected_without_taking_the_proxy_down(gateway: Gateway) -> None: + with gateway.scenario() as scenario, _patient_client(gateway) as client: + key: Final = scenario.key() + hostile_values: Final[tuple[JsonValue, ...]] = (5, ["gpt-4o-mini"]) + for hostile in hostile_values: + rejected = _timed_post(client, CHAT, {"model": hostile, "messages": []}, key) + assert rejected.response.status_code == 400 and "must be a string" in rejected.response.text + assert _error_information(rejected.response.headers["x-litellm-call-id"])["normalized_error"] == ( + "400_INVALID_REQUEST" + ) + empty: Final = _timed_post(client, CHAT, _body(CHAT, "", uuid.uuid4().hex), key) + assert empty.response.status_code == 400, empty.response.text + assert _error_information(empty.response.headers["x-litellm-call-id"])["normalized_error"] == ( + "400_INVALID_REQUEST" + ) + repeated: Final = tuple( + _timed_post(client, CHAT, _body(CHAT, HOSTILE_5KB_MODEL, uuid.uuid4().hex), key) for _ in range(2) + ) + call_ids: Final = tuple(timed.response.headers["x-litellm-call-id"] for timed in repeated) + assert len(set(call_ids)) == 2 and all(timed.response.status_code == 400 for timed in repeated) + assert all(timed.seconds < FAST_SECONDS for timed in repeated), [timed.seconds for timed in repeated] + codes: Final = tuple(_error_information(call_id)["normalized_error"] for call_id in call_ids) + assert len(set(codes)) == 1 and codes[0] in {"400_INVALID_REQUEST", "429_BUDGET_EXCEEDED"}, codes + unauthenticated: Final = client.post(CHAT, json=_body(CHAT, "gpt-4o-mini", "x")) + assert unauthenticated.status_code == 401, unauthenticated.text + assert client.get("/health/liveliness").status_code == 200 + + +@dataclass(frozen=True, slots=True) +class _BurstResult: + label: str + status: int | None + call_id: str | None + response_id: str | None + seconds: float + + +def _burst_call(client: httpx.Client, label: str, path: str, body: Mapping[str, JsonValue], key: str) -> _BurstResult: + started: Final = time.perf_counter() + try: + response: Final = client.post(path, json=body, headers={"Authorization": f"Bearer {key}"}) + except httpx.TransportError: + return _BurstResult(label, None, None, None, time.perf_counter() - started) + identity: Final = object_value(response.json()).get("id") if response.status_code == 200 else None + return _BurstResult( + label, + response.status_code, + response.headers.get("x-litellm-call-id"), + identity if isinstance(identity, str) else None, + time.perf_counter() - started, + ) + + +def _burst( + client: httpx.Client, happy_model: str, happy_key: str, open_key: str, during: Callable[[], None] +) -> tuple[_BurstResult, ...]: + crafted: Final = tuple( + (f"crafted-{path}-{index}", path, _body(path, CRAFTED_MODEL, uuid.uuid4().hex, stream=index % 2 == 1), open_key) + for path in (CHAT, MESSAGES, RESPONSES) + for index in range(4) + ) + happy: Final = tuple( + (f"happy-{index}", CHAT, _body(CHAT, happy_model, f"happy-{index}"), happy_key) for index in range(8) + ) + late: Final = tuple( + (f"late-{index}", CHAT, _body(CHAT, happy_model, f"late-{index}"), happy_key) for index in range(8) + ) + with ( + ThreadPoolExecutor(max_workers=28) as pool, + httpx.Client(base_url=client.base_url, timeout=client.timeout, trust_env=False) as fresh, + ): + first: Final = tuple( + pool.submit(_burst_call, client, label, path, body, key) for label, path, body, key in crafted + happy + ) + wait(first, return_when=FIRST_COMPLETED) + during() + second: Final = tuple( + pool.submit(_burst_call, fresh, label, path, body, key) for label, path, body, key in late + ) + return tuple(future.result() for future in first + second) + + +def _assert_rows_land_exactly_once(results: tuple[_BurstResult, ...], prefix: str) -> None: + landed: Final = tuple(result for result in results if result.label.startswith(prefix) and result.status == 200) + assert landed, results + response_ids: Final = tuple(string_value(result.response_id) for result in landed) + assert len(set(response_ids)) == len(response_ids), response_ids + rows: Final = eventually( + lambda: read_rows( + 'SELECT request_id, status FROM "LiteLLM_SpendLogs" WHERE request_id = ANY(%s::text[])', + ("{" + ",".join(response_ids) + "}",), + ), + lambda values: len(values) == len(response_ids), + seconds=ROW_SECONDS, + ) + assert sorted(string_value(row["request_id"]) for row in rows) == sorted(response_ids), rows + assert all(row["status"] == "success" for row in rows), rows + + +def test_killing_one_worker_mid_burst_leaves_the_other_serving_crafted_and_happy_traffic( + gateway: Gateway, tmp_path: Path +) -> None: + with ( + wire_server(_completion) as wire, + owned_proxy_process(gateway, tmp_path, {}, workers=2) as owned, + owned.gateway.scenario() as scenario, + _patient_client(owned.gateway) as client, + ): + model: Final = scenario.model(api_base=wire.url + "/v1", num_retries=0) + key: Final = scenario.key(models=[model]) + workers: Final = _worker_pids(owned) + assert len(workers) == 2, workers + + def kill_one_worker() -> None: + os.kill(workers[0], signal.SIGKILL) + + results: Final = _burst(client, model, key, scenario.key(), kill_one_worker) + dropped: Final = tuple(result for result in results if result.status is None) + assert len(dropped) < len(results), results + crafted: Final = tuple(result for result in results if result.label.startswith("crafted") and result.status) + assert crafted and all(result.status == 400 and result.seconds < FAST_SECONDS for result in crafted), crafted + late: Final = tuple(result for result in results if result.label.startswith("late")) + assert all(result.status == 200 for result in late), late + _assert_rows_land_exactly_once(results, "late") + survivor: Final = tuple(pid for pid in _worker_pids(owned) if pid != workers[0]) + assert survivor, "no worker left serving" + after: Final = owned.gateway.chat(model, key=key, text=f"after kill {uuid.uuid4().hex}") + assert isinstance(after["id"], str) and after["id"].startswith("chatcmpl-"), after + assert client.get("/health/liveliness").status_code == 200 + + +def test_upstream_returning_503_mid_burst_logs_every_failure_with_its_own_cluster_key( + gateway: Gateway, tmp_path: Path +) -> None: + def overloaded(_request: Request) -> Reply: + return Reply( + status=503, + body=b'{"error":{"message":"Controlled provider outage","type":"server_error","code":"503"}}', + ) + + with ( + wire_server(overloaded) as wire, + owned_proxy_process(gateway, tmp_path, {}, workers=2) as owned, + owned.gateway.scenario() as scenario, + _patient_client(owned.gateway) as client, + ): + model: Final = scenario.model(api_base=wire.url + "/v1", num_retries=0) + key: Final = scenario.key(models=[model]) + results: Final = _burst(client, model, key, scenario.key(), lambda: None) + assert all(result.status is not None for result in results), results + happy: Final = tuple(result for result in results if result.label.startswith("happy")) + assert all(result.status == 503 for result in happy), happy + late: Final = tuple(result for result in results if result.label.startswith("late")) + assert all(result.status == 503 for result in late), late + seen: Final = tuple(request.body.decode() for request in wire.drain()) + assert all(any(f"normalized error audit {result.label}" in body for body in seen) for result in happy + late), ( + seen + ) + crafted: Final = tuple(result for result in results if result.label.startswith("crafted")) + assert all(result.status == 400 and result.seconds < FAST_SECONDS for result in crafted), crafted + codes: Final = { + result.label: _error_information(string_value(result.call_id))["normalized_error"] for result in results + } + assert all(code == "503_PROVIDER_OVERLOADED" for label, code in codes.items() if label.startswith("happy")), ( + codes + ) + assert all(code == "400_INVALID_REQUEST" for label, code in codes.items() if label.startswith("crafted")), codes + assert client.get("/health/liveliness").status_code == 200 diff --git a/tests/test_litellm/litellm_core_utils/test_error_normalization.py b/tests/test_litellm/litellm_core_utils/test_error_normalization.py index 9d5469ddb4c..d65b6d316ac 100644 --- a/tests/test_litellm/litellm_core_utils/test_error_normalization.py +++ b/tests/test_litellm/litellm_core_utils/test_error_normalization.py @@ -1,3 +1,5 @@ +import time + import httpx import pytest @@ -229,3 +231,36 @@ def test_normalized_error_never_embeds_dynamic_parts() -> None: info = StandardLoggingPayloadSetup.get_error_information(exc) assert info["error_message"] == "No team has access to anthropic.claude-sonnet-4-5" assert "claude" not in (info["normalized_error"] or "") + + +def test_repeated_exceeded_in_a_288kb_message_classifies_in_linear_time() -> None: + model = ("exceeded " * 32_000)[:288_000] + message = ( + f"/chat/completions: Invalid model name passed in model={model}. Call `/v1/models` to view available models" + ) + exc = litellm.BadRequestError(message=message, model="unknown-model", llm_provider="openai") + started = time.perf_counter() + code = normalize_error(exc, "400", message) + elapsed = time.perf_counter() - started + assert code == "400_INVALID_REQUEST", code + assert elapsed < 1.0, f"normalize_error took {elapsed:.2f}s on a 288 KB message" + + +@pytest.mark.parametrize( + "message", + [ + "ExceededBudget: User=abc over budget. Spend=12.5, Budget=10.0", + "Exceeded budget for provider openai: 105.2 >= 100.0", + "LiteLLM Team: team-1, exceeded budget for model=gpt-4o-mini", + "ExceededBudget: Key over 1d budget. Spend=3.0, Budget=2.0", + "Budget has been exceeded! Key=sk-... Current cost: 11.0, Max budget: 10.0", + "EXCEEDED " + "x" * 65 + " BuDgEt", + ], +) +def test_real_budget_wordings_still_cluster_as_budget_exceeded(message: str) -> None: + assert normalize_error(Exception(message), "400", message) == "429_BUDGET_EXCEEDED" + + +@pytest.mark.parametrize("message", ["budget then exceeded", "exceeded the limit\nbudget unaffected", "exceededbudge"]) +def test_exceeded_without_a_following_budget_on_the_same_line_is_not_budget(message: str) -> None: + assert normalize_error(Exception(message), "400", message) == "400_INVALID_REQUEST" From 1519032d9042d9e9b3541612def76231082ec50b Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 23 Sep 2026 22:48:04 -0700 Subject: [PATCH 019/120] fix(proxy): keep deployment labels on cache-hit post_call guardrail rejections (#42780) * fix(proxy): keep deployment labels on cache-hit post_call rejections A post-call failure on a response served from the litellm cache set no first_api_call_start_time, so the failure hook flagged it as rejected before routing and dropped the model_id and provider labels Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): read the cache hit from caching_details in the failure hook model_call_details[cache_hit] is stamped inside the enqueued success handler, so a post-call failure can observe it too early; logging_obj.caching_details is set synchronously before the cached response returns Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(integration): cover cache-hit guardrail reject deployment labels across endpoints, modes and chaos Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(integration): bound the worker-kill reject count by in-flight losses Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(proxy): assert provider and model labels on the cache-hit regression test Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: yucheng Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/utils.py | 10 +- .../test_cache_hit_guardrail_metrics.py | 652 ++++++++++++++++++ .../test_cache_hit_guardrail_metrics_chaos.py | 348 ++++++++++ .../test_post_call_failure_hook.py | 60 ++ 4 files changed, 1069 insertions(+), 1 deletion(-) create mode 100644 tests/integration/observability/test_cache_hit_guardrail_metrics.py create mode 100644 tests/integration/observability/test_cache_hit_guardrail_metrics_chaos.py diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 78d6c25a336..bc64293c9b3 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -974,6 +974,14 @@ def _failure_usage_to_lift( _EMPTY_LIFT: Final = MappingProxyType({}) +def _reached_deployment(litellm_logging_obj: Logging) -> bool: + """A provider handoff or a cached response both mean the router selected a deployment.""" + caching_details: Final = litellm_logging_obj.caching_details + return litellm_logging_obj.model_call_details.get("first_api_call_start_time") is not None or ( + caching_details is not None and caching_details.get("cache_hit") is True + ) + + def _stamp_deployment_attribution( litellm_params: dict[str, object], model_group: str | None, team_id: str | None, dispatched: bool ) -> Mapping[str, object]: @@ -3325,7 +3333,7 @@ class ProxyLogging: _litellm_params, request_data.get("model"), user_api_key_dict.team_id, - dispatched=litellm_logging_obj.model_call_details.get("first_api_call_start_time") is not None, + dispatched=_reached_deployment(litellm_logging_obj), ) litellm_logging_obj.update_environment_variables( diff --git a/tests/integration/observability/test_cache_hit_guardrail_metrics.py b/tests/integration/observability/test_cache_hit_guardrail_metrics.py new file mode 100644 index 00000000000..888cfd7ab9d --- /dev/null +++ b/tests/integration/observability/test_cache_hit_guardrail_metrics.py @@ -0,0 +1,652 @@ +import asyncio +import json +import subprocess +import uuid +from collections.abc import Callable, Generator, Mapping +from contextlib import contextmanager +from dataclasses import dataclass +from pathlib import Path +from typing import Final + +import anthropic +import httpx +import openai +import pytest +import yaml +from integration._support.client import Gateway, Scenario, eventually, object_value, string_value +from integration._support.database import read_rows +from integration._support.process import owned_proxy_process +from integration._support.wire import Reply, Request, Wire, wire_server +from prometheus_client.parser import text_string_to_metric_families + +GUARDRAIL_PATH: Final = "/beta/litellm_basic_guardrail_api" +DEPLOYMENT_FAILURE: Final = "litellm_deployment_failure_responses_total" +DEPLOYMENT_REQUESTS: Final = "litellm_deployment_total_requests_total" +DEPLOYMENT_STATE: Final = "litellm_deployment_state" +PROXY_FAILED: Final = "litellm_proxy_failed_requests_metric_total" + + +def _chat_sse(marker: str) -> tuple[bytes, ...]: + chunk: Final = { + "id": "chatcmpl_" + marker, + "object": "chat.completion.chunk", + "created": 1, + "model": "gpt-4o-mini", + } + frames: Final = ( + {**chunk, "choices": [{"index": 0, "delta": {"role": "assistant", "content": ""}, "finish_reason": None}]}, + {**chunk, "choices": [{"index": 0, "delta": {"content": "provider control"}, "finish_reason": None}]}, + {**chunk, "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}]}, + ) + return tuple(f"data: {json.dumps(frame)}".encode() for frame in frames) + (b"data: [DONE]",) + + +def _provider_body(target: str, marker: str, streamed: bool) -> Reply: + match target: + case "/v1/chat/completions": + if streamed: + return Reply(content_type="text/event-stream", chunks=_chat_sse(marker)) + body: dict = { + "id": "chatcmpl_" + marker, + "object": "chat.completion", + "created": 1, + "model": "gpt-4o-mini", + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": "provider control " + marker}, + "finish_reason": "stop", + } + ], + "usage": {"prompt_tokens": 11, "completion_tokens": 4, "total_tokens": 15}, + } + case "/v1/messages": + body = { + "id": "msg_" + marker, + "type": "message", + "role": "assistant", + "model": "claude-sonnet-4-5-20250929", + "content": [{"type": "text", "text": "provider control " + marker}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 11, "output_tokens": 4}, + } + case "/v1/responses": + body = { + "id": "resp_" + marker, + "object": "response", + "created_at": 1, + "status": "completed", + "model": "gpt-4o-mini", + "output": [ + { + "type": "message", + "id": "msg_" + marker, + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": "provider control " + marker, "annotations": []}], + } + ], + "usage": {"input_tokens": 11, "output_tokens": 4, "total_tokens": 15}, + } + case "/v1/embeddings": + body = { + "object": "list", + "data": [{"object": "embedding", "index": 0, "embedding": [0.1, 0.2, 0.3]}], + "model": "text-embedding-3-small", + "usage": {"prompt_tokens": 3, "total_tokens": 3}, + } + case _: + return Reply(status=404, body=json.dumps({"error": "unexpected provider target " + target}).encode()) + return Reply(body=json.dumps(body).encode()) + + +def _provider(marker: str) -> Callable[[Request], Reply]: + def respond(request: Request) -> Reply: + streamed: Final = b'"stream":true' in request.body.replace(b" ", b"") + return _provider_body(request.target.split("?", 1)[0], marker, streamed) + + return respond + + +def _blocking_sink(request: Request) -> Reply: + assert request.target == GUARDRAIL_PATH, request.target + return Reply(body=json.dumps({"action": "BLOCKED", "blocked_reason": "synthetic block"}).encode()) + + +def _failing_sink(status: int) -> Callable[[Request], Reply]: + def respond(request: Request) -> Reply: + assert request.target == GUARDRAIL_PATH, request.target + return Reply(status=status, body=json.dumps({"error": "synthetic guardrail outage"}).encode()) + + return respond + + +def _first_call_pass_sink() -> Callable[[Request], Reply]: + calls: list[int] = [] # mutable-ok: the wire handler must remember call order across requests + + def respond(request: Request) -> Reply: + assert request.target == GUARDRAIL_PATH, request.target + calls.append(1) + action: dict = ( + {"action": "NONE"} if len(calls) == 1 else {"action": "BLOCKED", "blocked_reason": "synthetic block"} + ) + return Reply(body=json.dumps(action).encode()) + + return respond + + +def _guardrail_config( + tmp_path: Path, + name: str, + sink_url: str, + *, + mode: str = "post_call", + default_on: bool = False, + local_cache: bool = False, + ttl: int | None = None, +) -> Path: + config: dict = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + config["litellm_settings"]["callbacks"] = ["prometheus"] + if local_cache: + config["litellm_settings"]["cache_params"] = {"type": "local"} + if ttl is not None: + config["litellm_settings"]["cache_params"]["ttl"] = ttl + config["guardrails"] = [ + { + "guardrail_name": name, + "litellm_params": { + "guardrail": "generic_guardrail_api", + "mode": mode, + "default_on": default_on, + "api_base": sink_url, + "api_key": "synthetic-guardrail-key", + }, + } + ] + path: Final = tmp_path / "guardrail.yaml" + path.write_text(yaml.safe_dump(config)) + return path + + +@dataclass(frozen=True, slots=True) +class Rig: + candidate: Gateway + scenario: Scenario + model_name: str + deployment_id: str + guardrail_name: str + policy: Wire + provider: Wire + process: subprocess.Popen[bytes] + + +@contextmanager +def _rig( + gateway: Gateway, + tmp_path: Path, + marker: str, + *, + sink: Callable[[Request], Reply] = _blocking_sink, + mode: str = "post_call", + default_on: bool = False, + local_cache: bool = False, + ttl: int | None = None, + workers: int = 1, + upstream_model: str = "openai/gpt-4o-mini", + api_base_suffix: str = "/v1", + env: Mapping[str, str] | None = None, +) -> Generator[Rig, None, None]: + identity: Final = "guardrail-" + marker + with wire_server(sink) as policy, wire_server(_provider(marker)) as provider: + config: Final = _guardrail_config( + tmp_path, identity, policy.url, mode=mode, default_on=default_on, local_cache=local_cache, ttl=ttl + ) + prom_dir: Final = tmp_path / "prom" + prom_dir.mkdir() + with ( + owned_proxy_process( + gateway, + tmp_path, + {"PROMETHEUS_MULTIPROC_DIR": str(prom_dir), **(env or {})}, + config=config, + workers=workers, + ) as owned, + owned.gateway.scenario() as scenario, + ): + model: Final = scenario.model( + model=upstream_model, api_base=provider.url + api_base_suffix, api_key="synthetic-provider-key" + ) + entries: Final = owned.gateway.get("/model/info")["data"] + assert isinstance(entries, list) + entry: Final = next(item for item in entries if object_value(item)["model_name"] == model) + yield Rig( + owned.gateway, + scenario, + model, + string_value(object_value(object_value(entry)["model_info"])["id"]), + identity, + policy, + provider, + owned.process, + ) + + +def _metric_samples(candidate: Gateway, model_name: str) -> tuple: + response: Final = candidate.client.request( + "GET", "/metrics", headers={"Authorization": f"Bearer {candidate.key}"}, follow_redirects=True + ) + assert response.status_code == 200, f"GET /metrics: {response.status_code} {response.text[:300]}" + return tuple( + sample + for family in text_string_to_metric_families(response.text) + for sample in family.samples + if sample.labels.get("requested_model") == model_name + or (sample.name == DEPLOYMENT_STATE and sample.labels.get("model_id") != "") + ) + + +def _count(samples: tuple, name: str, model_id: str) -> float: + return float( + sum(sample.value for sample in samples if sample.name == name and sample.labels.get("model_id") == model_id) + ) + + +def _populated_failures(samples: tuple, rig: Rig, api_provider: str) -> float: + return float( + sum( + sample.value + for sample in samples + if sample.name == DEPLOYMENT_FAILURE + and sample.labels.get("model_id") == rig.deployment_id + and sample.labels.get("api_provider") == api_provider + and sample.labels.get("litellm_model_name") != "" + ) + ) + + +def _expect_metrics( + rig: Rig, + populated: float, + blank: float, + *, + api_provider: str = "openai", + pf_id: str | None = None, + pf_populated: float | None = None, + pf_blank: float | None = None, +) -> tuple: + expected_id: Final = rig.deployment_id if pf_id is None else pf_id + expected_pf_populated: Final = populated if pf_populated is None else pf_populated + expected_pf_blank: Final = blank if pf_blank is None else pf_blank + + def read() -> tuple: + samples: Final = _metric_samples(rig.candidate, rig.model_name) + satisfied: Final = ( + _populated_failures(samples, rig, api_provider) == populated + and _count(samples, DEPLOYMENT_FAILURE, "") == blank + and _count(samples, PROXY_FAILED, expected_id) == expected_pf_populated + and _count(samples, PROXY_FAILED, "") == expected_pf_blank + ) + return samples if satisfied else () + + return eventually(read, bool, seconds=70) + + +def _spend_rows(call_id: str) -> tuple[dict, ...]: + rows: Final = eventually( + lambda: read_rows( + 'SELECT request_id, custom_llm_provider, model_id, status FROM "LiteLLM_SpendLogs" ' + "WHERE request_id = %s OR request_id LIKE %s", + (call_id, call_id + "\\_%"), + ), + lambda values: len(values) >= 1, + seconds=70, + ) + return tuple(dict(row) for row in rows) + + +def _assert_spend(call_id: str, rig: Rig, api_provider: str = "openai") -> None: + rows: Final = _spend_rows(call_id) + failures: Final = tuple(row for row in rows if row["status"] == "failure") + assert len(failures) == 1, rows + assert (failures[0]["custom_llm_provider"], failures[0]["model_id"]) == (api_provider, rig.deployment_id), rows + + +def _call_id(reject: httpx.Response) -> str: + return reject.headers["x-litellm-call-id"] + + +def _chat_body(model: str, text: str, guardrail: str | None, stream: bool = False) -> dict: + body: dict = {"model": model, "messages": [{"role": "user", "content": text}]} + if stream: + body["stream"] = True + if guardrail is not None: + body["guardrails"] = [guardrail] + return body + + +def test_cache_hit_post_call_reject_keeps_deployment_labels(gateway: Gateway, tmp_path: Path) -> None: + """H1: warm then identical post_call-rejected cache hit keeps populated deployment labels.""" + marker: Final = uuid.uuid4().hex + with _rig(gateway, tmp_path, marker) as rig: + text: Final = "cache hit control h1 " + marker + warm: Final = rig.candidate.request("POST", "/v1/chat/completions", _chat_body(rig.model_name, text, None)) + assert warm.status_code == 200, warm.text + reject: Final = rig.candidate.request( + "POST", "/v1/chat/completions", _chat_body(rig.model_name, text, rig.guardrail_name) + ) + assert reject.status_code == 400, reject.text + assert rig.provider.received.qsize() == 1, rig.provider.drain() + _expect_metrics(rig, 1, 0) + _assert_spend(_call_id(reject), rig) + + +def test_cache_hit_post_call_reject_keeps_deployment_labels_openai_sdk(gateway: Gateway, tmp_path: Path) -> None: + """H2: same as H1 through the openai AsyncOpenAI client.""" + marker: Final = uuid.uuid4().hex + with _rig(gateway, tmp_path, marker) as rig: + text: Final = "cache hit control h2 " + marker + sdk: Final = openai.AsyncOpenAI( + base_url=str(rig.candidate.client.base_url) + "/v1", + api_key=rig.candidate.key, + http_client=httpx.AsyncClient(trust_env=False, timeout=15), + ) + + async def run() -> int: + await sdk.chat.completions.create(model=rig.model_name, messages=[{"role": "user", "content": text}]) + try: + await sdk.chat.completions.create( + model=rig.model_name, + messages=[{"role": "user", "content": text}], + extra_body={"guardrails": [rig.guardrail_name]}, + ) + return 200 + except openai.BadRequestError: + return 400 + + assert asyncio.run(run()) == 400 + assert rig.provider.received.qsize() == 1, rig.provider.drain() + _expect_metrics(rig, 1, 0) + + +def test_cache_hit_post_call_reject_streaming(gateway: Gateway, tmp_path: Path) -> None: + """H3: streamed responses are not cached; the reject call hits upstream again and no failure hook fires.""" + marker: Final = uuid.uuid4().hex + with _rig(gateway, tmp_path, marker) as rig: + text: Final = "cache hit control h3 " + marker + warm: Final = rig.candidate.request( + "POST", "/v1/chat/completions", _chat_body(rig.model_name, text, None, stream=True) + ) + assert warm.status_code == 200, warm.text + reject: Final = rig.candidate.request( + "POST", "/v1/chat/completions", _chat_body(rig.model_name, text, rig.guardrail_name, stream=True) + ) + assert reject.status_code == 200, reject.text + assert rig.provider.received.qsize() == 2, rig.provider.drain() + samples: Final = _metric_samples(rig.candidate, rig.model_name) + assert _populated_failures(samples, rig, "openai") == 0, samples + assert _count(samples, DEPLOYMENT_FAILURE, "") == 0, samples + + +def test_cache_hit_post_call_reject_keeps_deployment_labels_anthropic(gateway: Gateway, tmp_path: Path) -> None: + """H4: /v1/messages cache hit reject through the anthropic SDK.""" + marker: Final = uuid.uuid4().hex + with _rig( + gateway, tmp_path, marker, upstream_model="anthropic/claude-sonnet-4-5-20250929", api_base_suffix="" + ) as rig: + text: Final = "cache hit control h4 " + marker + sdk: Final = anthropic.Anthropic( + base_url=str(rig.candidate.client.base_url), + api_key=rig.candidate.key, + http_client=httpx.Client(trust_env=False, timeout=15), + ) + sdk.messages.create(model=rig.model_name, max_tokens=16, messages=[{"role": "user", "content": text}]) + raised: bool = False # mutable-ok: a flag set inside the except block cannot be Final + try: + sdk.messages.create( + model=rig.model_name, + max_tokens=16, + messages=[{"role": "user", "content": text}], + extra_body={"guardrails": [rig.guardrail_name]}, + ) + except anthropic.BadRequestError: + raised = True + assert raised, "cache-hit post_call guardrail did not reject /v1/messages" + assert rig.provider.received.qsize() == 1, rig.provider.drain() + _expect_metrics(rig, 1, 0, api_provider="anthropic", pf_id="None") + + +def test_cache_hit_post_call_reject_keeps_deployment_labels_responses(gateway: Gateway, tmp_path: Path) -> None: + """H5: /v1/responses cache hit reject.""" + marker: Final = uuid.uuid4().hex + with _rig(gateway, tmp_path, marker) as rig: + text: Final = "cache hit control h5 " + marker + warm: Final = rig.candidate.request("POST", "/v1/responses", {"model": rig.model_name, "input": text}) + assert warm.status_code == 200, warm.text + reject: Final = rig.candidate.request( + "POST", + "/v1/responses", + {"model": rig.model_name, "input": text, "guardrails": [rig.guardrail_name]}, + ) + assert reject.status_code == 400, reject.text + assert rig.provider.received.qsize() == 1, rig.provider.drain() + _expect_metrics(rig, 1, 0, pf_id="None") + _assert_spend(_call_id(reject), rig) + + +def test_cache_hit_post_call_reject_embeddings(gateway: Gateway, tmp_path: Path) -> None: + """H6: post_call guardrails do not run on embeddings; the cached response returns 200 unguarded.""" + marker: Final = uuid.uuid4().hex + with _rig(gateway, tmp_path, marker) as rig: + text: Final = "cache hit control h6 " + marker + warm: Final = rig.candidate.request("POST", "/v1/embeddings", {"model": rig.model_name, "input": text}) + assert warm.status_code == 200, warm.text + reject: Final = rig.candidate.request( + "POST", + "/v1/embeddings", + {"model": rig.model_name, "input": text, "guardrails": [rig.guardrail_name]}, + ) + assert reject.status_code == 200, reject.text + assert rig.provider.received.qsize() == 1, rig.provider.drain() + samples: Final = _metric_samples(rig.candidate, rig.model_name) + assert _populated_failures(samples, rig, "openai") == 0, samples + assert _count(samples, DEPLOYMENT_FAILURE, "") == 0, samples + + +def test_cache_hit_during_call_reject_keeps_deployment_labels(gateway: Gateway, tmp_path: Path) -> None: + """H7: during_call guardrail reject on a cache hit.""" + marker: Final = uuid.uuid4().hex + with _rig(gateway, tmp_path, marker, mode="during_call") as rig: + text: Final = "cache hit control h7 " + marker + warm: Final = rig.candidate.request("POST", "/v1/chat/completions", _chat_body(rig.model_name, text, None)) + assert warm.status_code == 200, warm.text + reject: Final = rig.candidate.request( + "POST", "/v1/chat/completions", _chat_body(rig.model_name, text, rig.guardrail_name) + ) + assert reject.status_code == 400, reject.text + _expect_metrics(rig, 1, 0) + + +def test_pre_call_reject_on_cache_hit_stays_blank(gateway: Gateway, tmp_path: Path) -> None: + """C1: pre_call reject never reaches the deployment; labels stay blank on both legs.""" + marker: Final = uuid.uuid4().hex + with _rig(gateway, tmp_path, marker, mode="pre_call") as rig: + text: Final = "cache hit control c1 " + marker + warm: Final = rig.candidate.request("POST", "/v1/chat/completions", _chat_body(rig.model_name, text, None)) + assert warm.status_code == 200, warm.text + reject: Final = rig.candidate.request( + "POST", "/v1/chat/completions", _chat_body(rig.model_name, text, rig.guardrail_name) + ) + assert reject.status_code == 400, reject.text + _expect_metrics(rig, 0, 1, pf_populated=1, pf_blank=0) + + +def test_post_call_reject_without_cache_keeps_deployment_labels(gateway: Gateway, tmp_path: Path) -> None: + """C2: a real provider call rejected post_call keeps populated labels on both legs.""" + marker: Final = uuid.uuid4().hex + with _rig(gateway, tmp_path, marker) as rig: + text: Final = "non cache control c2 " + marker + reject: Final = rig.candidate.request( + "POST", "/v1/chat/completions", _chat_body(rig.model_name, text, rig.guardrail_name) + ) + assert reject.status_code == 400, reject.text + assert rig.provider.received.qsize() == 1, rig.provider.drain() + _expect_metrics(rig, 1, 0) + _assert_spend(_call_id(reject), rig) + + +def test_cache_hit_post_call_reject_default_on(gateway: Gateway, tmp_path: Path) -> None: + """C3: default_on post_call guardrail rejects the cached response (sink passes the warm call).""" + marker: Final = uuid.uuid4().hex + with _rig(gateway, tmp_path, marker, sink=_first_call_pass_sink(), default_on=True) as rig: + text: Final = "cache hit control c3 " + marker + warm: Final = rig.candidate.request("POST", "/v1/chat/completions", _chat_body(rig.model_name, text, None)) + assert warm.status_code == 200, warm.text + reject: Final = rig.candidate.request("POST", "/v1/chat/completions", _chat_body(rig.model_name, text, None)) + assert reject.status_code == 400, reject.text + assert rig.provider.received.qsize() == 1, rig.provider.drain() + _expect_metrics(rig, 1, 0) + + +def test_cache_hit_post_call_reject_key_metadata_guardrails(gateway: Gateway, tmp_path: Path) -> None: + """C4: guardrail attached via key metadata guardrails on a cache hit.""" + marker: Final = uuid.uuid4().hex + with _rig(gateway, tmp_path, marker, sink=_first_call_pass_sink()) as rig: + key: Final = rig.candidate.post("/key/generate", {"metadata": {"guardrails": [rig.guardrail_name]}})["key"] + text: Final = "cache hit control c4 " + marker + warm: Final = rig.candidate.request( + "POST", "/v1/chat/completions", _chat_body(rig.model_name, text, None), key=key + ) + assert warm.status_code == 200, warm.text + reject: Final = rig.candidate.request( + "POST", "/v1/chat/completions", _chat_body(rig.model_name, text, None), key=key + ) + assert reject.status_code == 400, reject.text + assert rig.provider.received.qsize() == 1, rig.provider.drain() + assert rig.policy.received.qsize() == 2 + _expect_metrics(rig, 1, 0) + + +def test_cache_hit_post_call_reject_local_cache(gateway: Gateway, tmp_path: Path) -> None: + """C5: same cache-hit reject with cache_params type local.""" + marker: Final = uuid.uuid4().hex + with _rig(gateway, tmp_path, marker, local_cache=True) as rig: + text: Final = "cache hit control c5 " + marker + warm: Final = rig.candidate.request("POST", "/v1/chat/completions", _chat_body(rig.model_name, text, None)) + assert warm.status_code == 200, warm.text + reject: Final = rig.candidate.request( + "POST", "/v1/chat/completions", _chat_body(rig.model_name, text, rig.guardrail_name) + ) + assert reject.status_code == 400, reject.text + assert rig.provider.received.qsize() == 1, rig.provider.drain() + _expect_metrics(rig, 1, 0) + + +@pytest.mark.parametrize("status", (500, 403)) +def test_cache_hit_post_call_guardrail_outage_keeps_deployment_labels( + gateway: Gateway, tmp_path: Path, status: int +) -> None: + """S1/S2: guardrail sink answers 500/403 on the cache-hit call; failure hook still counts as dispatched.""" + marker: Final = uuid.uuid4().hex + with _rig(gateway, tmp_path, marker, sink=_failing_sink(status)) as rig: + text: Final = "cache hit control s " + marker + warm: Final = rig.candidate.request("POST", "/v1/chat/completions", _chat_body(rig.model_name, text, None)) + assert warm.status_code == 200, warm.text + reject: Final = rig.candidate.request( + "POST", "/v1/chat/completions", _chat_body(rig.model_name, text, rig.guardrail_name) + ) + assert reject.status_code >= 400, reject.text + assert rig.provider.received.qsize() == 1, rig.provider.drain() + _expect_metrics(rig, 0, 0, pf_populated=1, pf_blank=0) + + +def test_two_identical_cache_hit_rejects_increment_populated_series(gateway: Gateway, tmp_path: Path) -> None: + """E1: two identical cache-hit rejects count +2 on the populated series, two spend rows.""" + marker: Final = uuid.uuid4().hex + with _rig(gateway, tmp_path, marker) as rig: + text: Final = "cache hit control e1 " + marker + warm: Final = rig.candidate.request("POST", "/v1/chat/completions", _chat_body(rig.model_name, text, None)) + assert warm.status_code == 200, warm.text + rejects: Final = tuple( + rig.candidate.request("POST", "/v1/chat/completions", _chat_body(rig.model_name, text, rig.guardrail_name)) + for _ in range(2) + ) + assert all(response.status_code == 400 for response in rejects), [r.text for r in rejects] + assert rig.provider.received.qsize() == 1, rig.provider.drain() + _expect_metrics(rig, 2, 0) + + +def test_two_identical_cache_hit_rejects_write_matching_spend_rows(gateway: Gateway, tmp_path: Path) -> None: + """E1b: both cache-hit rejects land a failure spend row.""" + pytest.skip("BUG: roughly one in four back-to-back cache-hit rejects never lands its LiteLLM_SpendLogs row") + marker: Final = uuid.uuid4().hex + with _rig(gateway, tmp_path, marker) as rig: + text: Final = "cache hit control e1b " + marker + warm: Final = rig.candidate.request("POST", "/v1/chat/completions", _chat_body(rig.model_name, text, None)) + assert warm.status_code == 200, warm.text + rejects: Final = tuple( + rig.candidate.request("POST", "/v1/chat/completions", _chat_body(rig.model_name, text, rig.guardrail_name)) + for _ in range(2) + ) + assert all(response.status_code == 400 for response in rejects), [r.text for r in rejects] + for response in rejects: + _assert_spend(_call_id(response), rig) + + +def test_cache_hit_reject_after_ttl_expiry_is_a_miss(gateway: Gateway, tmp_path: Path) -> None: + """E2: cache_params ttl=1; post-expiry the same body misses, hits upstream again, labels populated.""" + marker: Final = uuid.uuid4().hex + with _rig(gateway, tmp_path, marker, ttl=1) as rig: + text: Final = "cache hit control e2 " + marker + warm: Final = rig.candidate.request("POST", "/v1/chat/completions", _chat_body(rig.model_name, text, None)) + assert warm.status_code == 200, warm.text + assert rig.provider.received.qsize() == 1 + + rejects: list[int] = [] # mutable-ok: the poll helper must remember how many rejects it issued + + def expired_miss() -> int: + rig.candidate.request("POST", "/v1/chat/completions", _chat_body(rig.model_name, text, rig.guardrail_name)) + rejects.append(1) + return rig.provider.received.qsize() + + eventually(lambda: expired_miss() == 2, bool, seconds=70) + _expect_metrics(rig, len(rejects), 0) + + +def test_cache_hit_reject_metrics_aggregate_across_workers(gateway: Gateway, tmp_path: Path) -> None: + """E3: workers=2, 8 cache-hit rejects, aggregated /metrics shows +8 on the populated series.""" + marker: Final = uuid.uuid4().hex + with _rig(gateway, tmp_path, marker, workers=2) as rig: + text: Final = "cache hit control e3 " + marker + warm: Final = rig.candidate.request("POST", "/v1/chat/completions", _chat_body(rig.model_name, text, None)) + assert warm.status_code == 200, warm.text + rejects: Final = tuple( + rig.candidate.request("POST", "/v1/chat/completions", _chat_body(rig.model_name, text, rig.guardrail_name)) + for _ in range(8) + ) + assert all(response.status_code == 400 for response in rejects), [r.text for r in rejects] + _expect_metrics(rig, 8, 0) + + +def test_cache_hit_reject_deployment_metric_set_diff(gateway: Gateway, tmp_path: Path) -> None: + """E4: exact expected label sets on litellm_deployment_* and litellm_proxy_failed_requests_metric.""" + marker: Final = uuid.uuid4().hex + with _rig(gateway, tmp_path, marker) as rig: + text: Final = "cache hit control e4 " + marker + warm: Final = rig.candidate.request("POST", "/v1/chat/completions", _chat_body(rig.model_name, text, None)) + assert warm.status_code == 200, warm.text + reject: Final = rig.candidate.request( + "POST", "/v1/chat/completions", _chat_body(rig.model_name, text, rig.guardrail_name) + ) + assert reject.status_code == 400, reject.text + samples: Final = _expect_metrics(rig, 1, 0) + blank: Final = tuple(sample for sample in samples if sample.labels.get("model_id") == "") + assert blank == (), blank + states: Final = tuple( + sample.value + for sample in samples + if sample.name == DEPLOYMENT_STATE + and sample.labels.get("model_id") == rig.deployment_id + and sample.labels.get("api_base") == "" + ) + assert states == (1.0,), states diff --git a/tests/integration/observability/test_cache_hit_guardrail_metrics_chaos.py b/tests/integration/observability/test_cache_hit_guardrail_metrics_chaos.py new file mode 100644 index 00000000000..faaa1fc7325 --- /dev/null +++ b/tests/integration/observability/test_cache_hit_guardrail_metrics_chaos.py @@ -0,0 +1,348 @@ +import json +import signal +import socket +import subprocess +import threading +import uuid +from collections.abc import Callable, Generator +from concurrent.futures import ThreadPoolExecutor +from contextlib import contextmanager +from pathlib import Path +from typing import Final + +import httpx +import psutil +from integration._support.client import Gateway, eventually, object_value, string_value +from integration._support.process import owned_proxy_process +from integration._support.wire import Reply, Request, wire_server +from prometheus_client.parser import text_string_to_metric_families +from test_cache_hit_guardrail_metrics import ( + DEPLOYMENT_FAILURE, + GUARDRAIL_PATH, + PROXY_FAILED, + Rig, + _blocking_sink, + _chat_body, + _guardrail_config, + _provider, + _rig, +) + +BURST: Final = 10 + + +@contextmanager +def _redis(port: int) -> Generator[subprocess.Popen[bytes], None, None]: + process: Final = subprocess.Popen(["redis-server", "--port", str(port), "--save", ""], stdout=subprocess.DEVNULL) + try: + yield process + finally: + process.kill() + process.wait(timeout=10) + + +def _free_port() -> int: + with socket.socket() as reserve: + reserve.bind(("127.0.0.1", 0)) + return reserve.getsockname()[1] + + +def _deployment_id(candidate: Gateway, model_name: str) -> str: + entries: Final = candidate.get("/model/info")["data"] + assert isinstance(entries, list) + entry: Final = next(item for item in entries if object_value(item)["model_name"] == model_name) + return string_value(object_value(object_value(entry)["model_info"])["id"]) + + +def _stall_sink(stall: threading.Event, release: threading.Event) -> Callable[[Request], Reply]: + def respond(request: Request) -> Reply: + assert request.target == GUARDRAIL_PATH, request.target + if stall.is_set(): + release.wait(timeout=60) + return Reply(body=json.dumps({"action": "BLOCKED", "blocked_reason": "synthetic block"}).encode()) + + return respond + + +def _samples(candidate: Gateway, model_names: tuple[str, ...]) -> tuple: + response: Final = candidate.client.request( + "GET", "/metrics", headers={"Authorization": f"Bearer {candidate.key}"}, follow_redirects=True + ) + assert response.status_code == 200, f"GET /metrics: {response.status_code}" + return tuple( + sample + for family in text_string_to_metric_families(response.text) + for sample in family.samples + if sample.labels.get("requested_model") in model_names + ) + + +def _populated(samples: tuple, deployment_id: str) -> float: + return float( + sum( + sample.value + for sample in samples + if sample.name == DEPLOYMENT_FAILURE and sample.labels.get("model_id") == deployment_id + ) + ) + + +def _blank(samples: tuple) -> float: + return float( + sum( + sample.value + for sample in samples + if sample.name == DEPLOYMENT_FAILURE and sample.labels.get("model_id") == "" + ) + ) + + +def _proxy_failed(samples: tuple) -> float: + return float(sum(sample.value for sample in samples if sample.name == PROXY_FAILED)) + + +def _burst_bodies(rig: Rig, marker: str, anthropic_name: str | None) -> tuple[tuple[str, dict], ...]: + chat: Final = tuple( + ("/v1/chat/completions", _chat_body(rig.model_name, f"burst {marker} {index}", rig.guardrail_name)) + for index in range(BURST) + ) + responses: Final = tuple( + ( + "/v1/responses", + {"model": rig.model_name, "input": f"burst {marker} r{index}", "guardrails": [rig.guardrail_name]}, + ) + for index in range(BURST) + ) + messages: Final = ( + tuple( + ( + "/v1/messages", + { + "model": anthropic_name, + "max_tokens": 16, + "messages": [{"role": "user", "content": f"burst {marker} m{index}"}], + "guardrails": [rig.guardrail_name], + }, + ) + for index in range(BURST) + ) + if anthropic_name is not None + else () + ) + return chat + responses + messages + + +def _warm(rig: Rig, bodies: tuple[tuple[str, dict], ...]) -> None: + for path, body in bodies: + warmed: Final = dict(body) + warmed.pop("guardrails", None) + response: Final = rig.candidate.request("POST", path, warmed) + assert response.status_code == 200, f"warm {path}: {response.status_code} {response.text}" + + +def _fire(rig: Rig, bodies: tuple[tuple[str, dict], ...]) -> tuple[tuple[int, str | None], ...]: + def call(item: tuple[str, dict]) -> tuple[int, str | None]: + path, body = item + try: + response: Final = rig.candidate.request("POST", path, body) + return response.status_code, response.headers.get("x-litellm-call-id") + except httpx.HTTPError: + return -1, None + + with ThreadPoolExecutor(max_workers=8) as pool: + return tuple(pool.map(call, bodies)) + + +def _expect_counted_within( + rig: Rig, model_names: tuple[str, ...], deployment_ids: tuple[str, ...], low: int, high: int +) -> None: + def converged() -> tuple: + samples: Final = _samples(rig.candidate, model_names) + populated: Final = sum(_populated(samples, deployment) for deployment in deployment_ids) + if low <= populated <= high and _blank(samples) == 0: + return samples + return () + + eventually(converged, bool, seconds=70) + + +def _expect_exactly_once(rig: Rig, model_names: tuple[str, ...], deployment_ids: tuple[str, ...], four_xx: int) -> None: + _expect_counted_within(rig, model_names, deployment_ids, four_xx, four_xx) + + +def test_burst_cache_hit_rejects_count_exactly_once(gateway: Gateway, tmp_path: Path) -> None: + """X0: 30 mixed-endpoint cache-hit rejects across two deployments, each counted once.""" + marker: Final = uuid.uuid4().hex + with _rig(gateway, tmp_path, marker) as rig: + anthropic_name: Final = rig.scenario.model( + model="anthropic/claude-sonnet-4-5-20250929", api_base=rig.provider.url, api_key="synthetic-provider-key" + ) + anthropic_id: Final = _deployment_id(rig.candidate, anthropic_name) + bodies: Final = _burst_bodies(rig, marker, anthropic_name) + _warm(rig, bodies) + outcomes: Final = _fire(rig, bodies) + rejected: Final = sum(1 for status, _ in outcomes if status >= 400) + assert all(status == 400 for status, _ in outcomes), outcomes + _expect_exactly_once(rig, (rig.model_name, anthropic_name), (rig.deployment_id, anthropic_id), rejected) + + +def test_stalled_guardrail_sink_recovers_and_counts(gateway: Gateway, tmp_path: Path) -> None: + """X1: guardrail sink stalls mid-burst; requests fail exactly once, then recovery counts again.""" + marker: Final = uuid.uuid4().hex + stall: Final = threading.Event() + release: Final = threading.Event() + with _rig(gateway, tmp_path, marker, sink=_stall_sink(stall, release)) as rig: + bodies: Final = _burst_bodies(rig, marker, None) + _warm(rig, bodies) + stall.set() + with ThreadPoolExecutor(max_workers=8) as pool: + futures: Final = tuple( + pool.submit(lambda b: rig.candidate.request("POST", b[0], b[1]), body) for body in bodies + ) + eventually(lambda: rig.policy.received.qsize() >= 5, bool, seconds=30) + release.set() + outcomes: Final = tuple( + (future.result().status_code, future.result().headers.get("x-litellm-call-id")) for future in futures + ) + assert all(status >= 400 for status, _ in outcomes), outcomes + blocked: Final = sum(1 for status, _ in outcomes if status == 400) + outages: Final = sum(1 for status, _ in outcomes if status >= 500) + assert blocked + outages == len(bodies), outcomes + samples: Final = eventually( + lambda: _samples(rig.candidate, (rig.model_name,)), + lambda observed: _proxy_failed(observed) == blocked + outages, + seconds=70, + ) + assert _proxy_failed(samples) == blocked + outages, (samples, outcomes) + follow_up: Final = rig.candidate.request( + "POST", + "/v1/chat/completions", + _chat_body(rig.model_name, "post stall unrelated " + marker, None), + ) + assert follow_up.status_code == 200, follow_up.text + _expect_exactly_once(rig, (rig.model_name,), (rig.deployment_id,), blocked) + + +def test_redis_outage_keeps_serving_in_memory_hits(gateway: Gateway, tmp_path: Path) -> None: + """X2: the redis cache keeps an in-memory shadow, so a redis kill does not stop cache-hit rejects.""" + marker: Final = uuid.uuid4().hex + port: Final = _free_port() + with _redis(port) as redis_one: + with _rig(gateway, tmp_path, marker, env={"REDIS_HOST": "127.0.0.1", "REDIS_PORT": str(port)}) as rig: + bodies: Final = _burst_bodies(rig, marker, None)[:BURST] + _warm(rig, bodies) + reject: Final = rig.candidate.request("POST", *bodies[0]) + assert reject.status_code == 400, reject.text + warmed_hits: Final = rig.provider.received.qsize() + redis_one.kill() + redis_one.wait(timeout=10) + outcomes: Final = _fire(rig, bodies[1:]) + assert all(status == 400 for status, _ in outcomes), outcomes + assert rig.provider.received.qsize() == warmed_hits, ( + "redis outage reached the provider", + warmed_hits, + rig.provider.received.qsize(), + ) + with _redis(port): + recovered: Final = rig.candidate.request( + "POST", + "/v1/chat/completions", + _chat_body(rig.model_name, "x2 rehit " + marker, rig.guardrail_name), + ) + assert recovered.status_code == 400, recovered.text + _expect_exactly_once(rig, (rig.model_name,), (rig.deployment_id,), 1 + len(bodies)) + + +def test_worker_kill_mid_burst_keeps_counting(gateway: Gateway, tmp_path: Path) -> None: + """X3: workers=2, SIGKILL one uvicorn child mid-burst; survivors keep rejecting; the count is answered plus at most the in-flight requests the killed worker had already counted.""" + marker: Final = uuid.uuid4().hex + with _rig(gateway, tmp_path, marker, workers=2) as rig: + bodies: Final = _burst_bodies(rig, marker, None) + _warm(rig, bodies) + children: Final = psutil.Process(rig.process.pid).children(recursive=True) + assert children, "no uvicorn worker children found" + with ThreadPoolExecutor(max_workers=8) as pool: + futures: Final = tuple( + pool.submit(lambda b: rig.candidate.request("POST", b[0], b[1]), body) for body in bodies + ) + eventually(lambda: rig.policy.received.qsize() >= 3, bool, seconds=30) + children[0].send_signal(signal.SIGKILL) + statuses: list[int] = [] # mutable-ok: collect per-request outcomes from concurrent futures + for future in futures: + try: + statuses.append(future.result().status_code) + except httpx.HTTPError: + statuses.append(-1) + answered: Final = sum(1 for status in statuses if status >= 0) + transport_lost: Final = sum(1 for status in statuses if status == -1) + assert all(status == 400 for status in statuses if status >= 0), ( + statuses, + transport_lost, + ) + _expect_counted_within(rig, (rig.model_name,), (rig.deployment_id,), answered, answered + transport_lost) + + +def test_proxy_restart_mid_burst_keeps_counting(gateway: Gateway, tmp_path: Path) -> None: + """X4: restart the owned proxy between the two halves; pre-restart count asserted, then recounted.""" + marker: Final = uuid.uuid4().hex + prom_dir: Final = tmp_path / "prom" + prom_dir.mkdir() + with wire_server(_blocking_sink) as policy, wire_server(_provider(marker)) as provider: + config: Final = _guardrail_config(tmp_path, "guardrail-" + marker, policy.url) + bodies: Final = tuple( + ( + "/v1/chat/completions", + _chat_body("pending-model", f"burst {marker} {index}", "guardrail-" + marker), + ) + for index in range(BURST) + ) + with owned_proxy_process( + gateway, tmp_path, {"PROMETHEUS_MULTIPROC_DIR": str(prom_dir)}, config=config + ) as owned_one: + model: Final = "restart-" + marker + owned_one.gateway.post( + "/model/new", + { + "model_name": model, + "litellm_params": { + "model": "openai/gpt-4o-mini", + "api_base": provider.url + "/v1", + "api_key": "synthetic-provider-key", + }, + }, + ) + deployment: Final = _deployment_id(owned_one.gateway, model) + named: Final = tuple((path, {**body, "model": model}) for path, body in bodies) + first_half, second_half = named[: BURST // 2], named[BURST // 2 :] + for path, body in named: + warmed: Final = dict(body) + warmed.pop("guardrails", None) + assert owned_one.gateway.request("POST", path, warmed).status_code == 200 + outcomes_one: Final = tuple(owned_one.gateway.request("POST", path, body) for path, body in first_half) + assert all(response.status_code == 400 for response in outcomes_one), [r.text for r in outcomes_one] + pre: Final = eventually( + lambda: ( + _populated(_samples(owned_one.gateway, (model,)), deployment), + _blank(_samples(owned_one.gateway, (model,))), + ), + lambda observed: observed[0] == len(first_half) and observed[1] == 0, + seconds=70, + ) + with owned_proxy_process( + gateway, tmp_path, {"PROMETHEUS_MULTIPROC_DIR": str(prom_dir)}, config=config + ) as owned_two: + outcomes_two: Final = tuple(owned_two.gateway.request("POST", path, body) for path, body in second_half) + assert all(response.status_code == 400 for response in outcomes_two), ( + pre, + [(r.status_code, r.text[:200]) for r in outcomes_two], + ) + post: Final = eventually( + lambda: ( + _populated(_samples(owned_two.gateway, (model,)), deployment), + _blank(_samples(owned_two.gateway, (model,))), + ), + lambda observed: observed[0] == len(named) and observed[1] == 0, + seconds=70, + ) + assert post[0] == len(named), (pre, post, outcomes_two) + owned_two.gateway.post("/model/delete", {"id": deployment}) diff --git a/tests/test_litellm/proxy/utils/proxy_logging/test_post_call_failure_hook.py b/tests/test_litellm/proxy/utils/proxy_logging/test_post_call_failure_hook.py index 0046721fd9c..51145ca687b 100644 --- a/tests/test_litellm/proxy/utils/proxy_logging/test_post_call_failure_hook.py +++ b/tests/test_litellm/proxy/utils/proxy_logging/test_post_call_failure_hook.py @@ -18,6 +18,7 @@ from litellm.exceptions import GuardrailRaisedException from litellm.integrations.custom_logger import CustomLogger from litellm.proxy._types import ProxyErrorTypes, UserAPIKeyAuth from litellm.proxy.utils import ProxyLogging +from litellm.types.utils import CachingDetails @pytest.fixture(autouse=True) @@ -249,6 +250,65 @@ async def test_post_call_failure_hook_keeps_router_stamped_metadata_for_post_cal assert kwargs["standard_logging_object"]["model_id"] == "routed-deployment" +@pytest.mark.asyncio +async def test_post_call_failure_hook_keeps_deployment_attribution_for_cache_hit_post_call_failures( + proxy_logging, make_user_api_key_auth, monkeypatch +): + """A post-call guardrail blocks a response served from the litellm cache. No provider call was made, + so ``first_api_call_start_time`` is unset, but the router did pick the deployment: the pre-routing + flag must stay off so ``litellm_deployment_failure_responses`` keeps its model_id and provider labels.""" + from litellm.proxy import proxy_server + + recorded: list[dict] = [] + + class _RecordingLogger(CustomLogger): + async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): + recorded.append(kwargs) + + monkeypatch.setattr( + proxy_server, + "llm_router", + litellm.Router( + model_list=[ + { + "model_name": "internal-model", + "litellm_params": {"model": "openai/gpt-4.1", "api_key": "sk-test"}, + "model_info": {"id": "routed-deployment"}, + } + ] + ), + ) + monkeypatch.setattr(litellm, "callbacks", [_RecordingLogger()]) + proxy_logging.alert_types = [] + + request_data = { + "litellm_call_id": "cache-hit-post-call-guardrail", + "model": "internal-model", + "messages": [{"role": "user", "content": "hi"}], + "metadata": {"model_info": {"id": "routed-deployment"}}, + } + logging_obj, request_data = litellm.utils.function_setup( + original_function="acompletion", rules_obj=litellm.utils.Rules(), start_time=datetime.now(), **request_data + ) + logging_obj.caching_details = CachingDetails(cache_hit=True, cache_duration_ms=1.0) + request_data["litellm_logging_obj"] = logging_obj + + await proxy_logging.post_call_failure_hook( + request_data=request_data, + original_exception=GuardrailRaisedException(guardrail_name="g", message="response blocked"), + user_api_key_dict=make_user_api_key_auth(request_route="/chat/completions"), + route="/chat/completions", + ) + + assert len(recorded) == 1 + kwargs = recorded[0] + assert PROXY_REJECTED_BEFORE_ROUTING_KEY not in kwargs["litellm_params"], kwargs["litellm_params"] + assert kwargs["standard_logging_object"]["model_id"] == "routed-deployment" + assert kwargs["standard_logging_object"]["custom_llm_provider"] == "openai" + assert kwargs["model"] == "internal-model" + assert kwargs["litellm_params"]["custom_llm_provider"] == "openai" + + @pytest.mark.asyncio async def test_post_call_failure_hook_flags_pre_routing_reject_despite_caller_model_info( proxy_logging, make_user_api_key_auth, monkeypatch From 093ceb576d7d9743bbe6ca061791884084143916 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 23 Sep 2026 22:54:19 -0700 Subject: [PATCH 020/120] fix(proxy): do not requeue a daily spend batch whose commit already left for postgres (#42786) * fix(proxy): do not requeue a daily spend batch whose commit already left for postgres Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): settle an interrupted daily spend commit from the shutdown flush instead of blocking the cancelled tick Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(integration): burst two workers and SIGTERM during daily spend COMMIT, expect exactly once Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: yucheng Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/db/db_spend_update_writer.py | 114 ++++++++- .../daily_spend_update_queue.py | 14 ++ .../integration/spend/test_shutdown_flush.py | 132 +++++++++- .../proxy/db/test_db_spend_update_writer.py | 238 ++++++++++++++++++ 4 files changed, 486 insertions(+), 12 deletions(-) diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index e38214c98a6..d3ac37a3e0f 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -12,7 +12,8 @@ import os import random import time import traceback -from collections.abc import Callable, Mapping, Sequence +from collections.abc import Callable, Coroutine, Mapping, Sequence +from contextvars import ContextVar from datetime import datetime, timedelta, timezone from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, TypeAlias, TypeVar, cast, overload @@ -269,6 +270,80 @@ def _spend_update_tx(prisma_client: PrismaClient) -> _SpendTransactionManager: return tx +_daily_spend_commit_started: Final[ContextVar[asyncio.Event | None]] = ContextVar( + "_daily_spend_commit_started", default=None +) + + +def _mark_daily_spend_commit_started() -> None: + started: Final = _daily_spend_commit_started.get() + if started is not None: + started.set() + + +def _mark_daily_spend_commit_finished() -> None: + started: Final = _daily_spend_commit_started.get() + if started is not None: + started.clear() + + +def _start_daily_spend_commit( + commit_started: asyncio.Event, commit: Callable[[], Coroutine[object, object, None]] +) -> "asyncio.Task[None]": + token: Final = _daily_spend_commit_started.set(commit_started) + try: + return asyncio.ensure_future(commit()) + finally: + _daily_spend_commit_started.reset(token) + + +def _track_interrupted_commit(commits: set[asyncio.Task[None]], settle: Coroutine[object, object, None]) -> None: + task: Final = asyncio.ensure_future(settle) + commits.add(task) + task.add_done_callback(commits.discard) + + +async def _settle_interrupted_commits(commits: set[asyncio.Task[None]]) -> None: + while commits: + await asyncio.wait(tuple(commits)) + + +async def _restore_tag_spend_the_commit_left_behind( + commit_task: "asyncio.Task[None]", + redis_update_buffer: RedisUpdateBuffer, + transactions: dict[str, DailyTagSpendTransaction], +) -> None: + await asyncio.wait({commit_task}) + if commit_task.cancelled() or commit_task.exception() is None: + return + await redis_update_buffer.restore_transactions_to_redis( + daily_tag_spend_update_transactions=transactions, + ) + + +async def _requeue_daily_spend_the_commit_left_behind( + commit_task: "asyncio.Task[None]", + queue: DailySpendUpdateQueue, + entity_type: str, + transactions: dict[str, BaseDailySpendTransaction], +) -> None: + await asyncio.wait({commit_task}) + if commit_task.cancelled() or not transactions: + return + failure: Final = commit_task.exception() + if failure is None: + return + spend_log_error( + "Spend tracking - daily %s spend commit interrupted by shutdown failed. Re-queued %d rows for the " + "shutdown flush. Error: %s", + entity_type, + len(transactions), + str(failure), + exc=failure, + ) + await queue.add_update(transactions) + + # The per-team advisory lock the team endpoints hold while changing a roster (TEAM_ADVISORY_LOCK_SQL), # so the roster check below cannot interleave with their writes. A row lock would deadlock with the # access-group endpoints, which lock a team row after an access-group lock. @@ -391,6 +466,9 @@ class DBSpendUpdateWriter: self.daily_org_spend_update_queue = DailySpendUpdateQueue() self.daily_tag_spend_update_queue = DailySpendUpdateQueue() self.window_spend_update_queue = WindowSpendUpdateQueue() + self.interrupted_tag_commits: set[asyncio.Task[None]] = ( + set() + ) # mutable-ok: same registry as DailySpendUpdateQueue.interrupted_commits async def update_database( # LiteLLM management object fields @@ -1606,17 +1684,24 @@ class DBSpendUpdateWriter: proxy_logging_obj: ProxyLogging, ) -> None: transactions: Final = await queue.flush_and_get_aggregated_daily_spend_update_transactions() - commit_task: Final = asyncio.ensure_future( - commit( + commit_started: Final = asyncio.Event() + commit_task: Final = _start_daily_spend_commit( + commit_started, + lambda: commit( n_retry_times=n_retry_times, prisma_client=prisma_client, proxy_logging_obj=proxy_logging_obj, daily_spend_transactions=cast(dict[str, _DailySpendTransactionT], transactions), - ) + ), ) try: await asyncio.shield(commit_task) except asyncio.CancelledError: + if commit_started.is_set(): + queue.track_interrupted_commit( + _requeue_daily_spend_the_commit_left_behind(commit_task, queue, entity_type, transactions) + ) + raise commit_task.cancel() if transactions: await queue.add_update(transactions) @@ -1841,23 +1926,36 @@ class DBSpendUpdateWriter: The drain is destructive, so a failed commit must push the transactions back for the next tick or their spend is lost permanently. """ + await _settle_interrupted_commits(self.interrupted_tag_commits) daily_tag_spend_update_transactions: Final = ( await self.redis_update_buffer.get_all_daily_tag_spend_update_transactions_from_redis_buffer() ) if not daily_tag_spend_update_transactions: return - commit_task: Final = asyncio.ensure_future( - DBSpendUpdateWriter.update_daily_tag_spend( + commit_started: Final = asyncio.Event() + commit_task: Final = _start_daily_spend_commit( + commit_started, + lambda: DBSpendUpdateWriter.update_daily_tag_spend( n_retry_times=n_retry_times, prisma_client=prisma_client, proxy_logging_obj=proxy_logging_obj, daily_spend_transactions=daily_tag_spend_update_transactions, - ) + ), ) try: await asyncio.shield(commit_task) except BaseException: # noqa: BLE001 # a cancel must restore the drained rows before its rollback returns + if commit_started.is_set(): + _track_interrupted_commit( + self.interrupted_tag_commits, + _restore_tag_spend_the_commit_left_behind( + commit_task, + self.redis_update_buffer, + daily_tag_spend_update_transactions, + ), + ) + raise commit_task.cancel() await self.redis_update_buffer.restore_transactions_to_redis( daily_tag_spend_update_transactions=daily_tag_spend_update_transactions, @@ -2382,6 +2480,8 @@ class DBSpendUpdateWriter: sql, params = build_bulk_upsert(table=table, batch=merged_batch) async with _spend_update_tx(prisma_client) as transaction: await transaction.execute_raw(sql, *params) + _mark_daily_spend_commit_started() + _mark_daily_spend_commit_finished() except Exception as batch_error: if _spend_commit_failure_is_requeue_safe(batch_error): spend_log_error( diff --git a/litellm/proxy/db/db_transaction_queue/daily_spend_update_queue.py b/litellm/proxy/db/db_transaction_queue/daily_spend_update_queue.py index c6381cd070b..f911d5a6767 100644 --- a/litellm/proxy/db/db_transaction_queue/daily_spend_update_queue.py +++ b/litellm/proxy/db/db_transaction_queue/daily_spend_update_queue.py @@ -1,4 +1,5 @@ import asyncio +from collections.abc import Coroutine from copy import deepcopy from typing import Final @@ -57,6 +58,18 @@ class DailySpendUpdateQueue(BaseUpdateQueue): self.update_queue: asyncio.Queue[dict[str, BaseDailySpendTransaction]] = asyncio.Queue( maxsize=LITELLM_ASYNCIO_QUEUE_MAXSIZE ) + self.interrupted_commits: set[asyncio.Task[None]] = ( + set() + ) # mutable-ok: registry of in-flight commit outcomes, entries leave via their done callback + + def track_interrupted_commit(self, settle: Coroutine[object, object, None]) -> None: + task: Final = asyncio.ensure_future(settle) + self.interrupted_commits.add(task) + task.add_done_callback(self.interrupted_commits.discard) + + async def settle_interrupted_commits(self) -> None: + while self.interrupted_commits: + await asyncio.wait(tuple(self.interrupted_commits)) async def add_update(self, update: dict[str, BaseDailySpendTransaction]): """Enqueue an update.""" @@ -81,6 +94,7 @@ class DailySpendUpdateQueue(BaseUpdateQueue): self, ) -> dict[str, BaseDailySpendTransaction]: """Get all updates from the queue and return all updates aggregated by daily_transaction_key. Works for both user and team spend updates.""" + await self.settle_interrupted_commits() updates: Final = await self.flush_all_updates_from_in_memory_queue() if len(updates) > 0: verbose_proxy_logger.info( diff --git a/tests/integration/spend/test_shutdown_flush.py b/tests/integration/spend/test_shutdown_flush.py index d27275f5213..b744c1f4b5f 100644 --- a/tests/integration/spend/test_shutdown_flush.py +++ b/tests/integration/spend/test_shutdown_flush.py @@ -4,6 +4,7 @@ import signal import threading import uuid from collections.abc import Callable, Iterator +from concurrent.futures import ThreadPoolExecutor from contextlib import contextmanager from dataclasses import dataclass from pathlib import Path @@ -13,16 +14,18 @@ import httpx import psycopg import pytest import yaml - from integration._support.client import Gateway, delete_key_if_present, eventually, string_value from integration._support.database import read_rows from integration._support.process import OwnedProxy, owned_proxy_process from integration._support.wire import Reply, Request, wire_server +from psycopg import sql REQUESTS_WHILE_BLOCKED: Final = 6 CANCEL_LOG_LINE: Final = "in-flight scheduled job(s) for shutdown" BATCH_DRAINED_LOG_LINE: Final = f"flushed {REQUESTS_WHILE_BLOCKED} daily spend update items from in-memory queue" MODEL_INSERT_ARRIVED_LOG_LINE: Final = "path=/model/new" +COMMIT_DELAY_SECONDS: Final = 15 +BURST_REQUESTS: Final = 30 def _api_requests(table: str, column: str, identity: str) -> int: @@ -44,6 +47,50 @@ def _waiting_on(table: str) -> int: return waiting +def _committing_daily_user_spend() -> int: + rows: Final = read_rows( + "SELECT count(*)::int AS committing FROM pg_stat_activity " + "WHERE query='COMMIT' AND state='active' AND wait_event='PgSleep' AND pid IN " + "(SELECT pid FROM pg_locks WHERE relation = %s::regclass AND mode='RowExclusiveLock')", + ('"LiteLLM_DailyUserSpend"',), + ) + committing: Final = rows[0]["committing"] + assert isinstance(committing, int) + return committing + + +def _install_slow_commit(user_id: str, fails_once: bool) -> str: + suffix: Final = f"slow_commit_{uuid.uuid4().hex}" + with psycopg.connect(os.environ["DATABASE_URL"], autocommit=True) as connection: + connection.execute(sql.SQL("CREATE SEQUENCE {}").format(sql.Identifier(suffix))) + connection.execute( + sql.SQL( + "CREATE FUNCTION {}() RETURNS trigger LANGUAGE plpgsql AS $slow$ " + "BEGIN PERFORM pg_sleep({}); " + "IF {} AND nextval({}) = 1 THEN RAISE EXCEPTION 'integration: first COMMIT fails'; END IF; " + "RETURN NULL; END $slow$" + ).format( + sql.Identifier(suffix), sql.Literal(COMMIT_DELAY_SECONDS), sql.Literal(fails_once), sql.Literal(suffix) + ) + ) + connection.execute( + sql.SQL( + 'CREATE CONSTRAINT TRIGGER {} AFTER INSERT OR UPDATE ON "LiteLLM_DailyUserSpend" ' + "DEFERRABLE INITIALLY DEFERRED FOR EACH ROW WHEN (NEW.user_id = {}) EXECUTE FUNCTION {}()" + ).format(sql.Identifier(suffix), sql.Literal(user_id), sql.Identifier(suffix)) + ) + return suffix + + +def _drop_slow_commit(suffix: str) -> None: + with psycopg.connect(os.environ["DATABASE_URL"], autocommit=True) as connection: + connection.execute( + sql.SQL('DROP TRIGGER IF EXISTS {} ON "LiteLLM_DailyUserSpend"').format(sql.Identifier(suffix)) + ) + connection.execute(sql.SQL("DROP FUNCTION IF EXISTS {}()").format(sql.Identifier(suffix))) + connection.execute(sql.SQL("DROP SEQUENCE IF EXISTS {}").format(sql.Identifier(suffix))) + + def _provider(request: Request) -> Reply: if request.method != "POST": return Reply(status=404, body=b'{"error":"not scripted"}') @@ -77,6 +124,17 @@ class _Shutdown: def daily_user_requests(self) -> int: return _api_requests("LiteLLM_DailyUserSpend", "user_id", self.owner) + def spend_logs(self) -> int: + rows: Final = read_rows('SELECT count(*)::int AS total FROM "LiteLLM_SpendLogs" WHERE "user"=%s', (self.owner,)) + total: Final = rows[0]["total"] + assert isinstance(total, int) + return total + + def burst(self, requests: int) -> None: + with ThreadPoolExecutor(max_workers=8) as pool: + for outcome in pool.map(lambda _: self.chat(), range(requests)): + assert outcome is None + def logged(self, line: str, times: int = 1) -> bool: return self.owned.log.read_text(errors="replace").count(line) >= times @@ -121,7 +179,15 @@ def _config_with_pool_limit(tmp_path: Path, pool_limit: int) -> Path: @contextmanager -def _proxy_with_one_seeded_row(gateway: Gateway, tmp_path: Path, pool_limit: int) -> Iterator[_Shutdown]: +def _proxy_with_one_seeded_row( + gateway: Gateway, + tmp_path: Path, + pool_limit: int, + cancel_timeout_seconds: int = 5, + settle_seconds: int = 0, + requests: int = REQUESTS_WHILE_BLOCKED, + workers: int = 1, +) -> Iterator[_Shutdown]: owner: Final = f"integration-owner-{uuid.uuid4().hex}" with gateway.scenario() as scenario, wire_server(_provider) as wire: model: Final = scenario.model(api_base=wire.url + "/v1", num_retries=0) @@ -133,9 +199,10 @@ def _proxy_with_one_seeded_row(gateway: Gateway, tmp_path: Path, pool_limit: int "LITELLM_LOG": "DEBUG", "GRACEFUL_SHUTDOWN_TIMEOUT": "1", "SCHEDULED_JOB_SHUTDOWN_FINISH_TIMEOUT_SECONDS": "1", - "SCHEDULED_JOB_SHUTDOWN_CANCEL_TIMEOUT_SECONDS": "5", + "SCHEDULED_JOB_SHUTDOWN_CANCEL_TIMEOUT_SECONDS": str(cancel_timeout_seconds), }, config=_config_with_pool_limit(tmp_path, pool_limit), + workers=workers, ) as owned: key: Final = string_value( owned.gateway.post("/key/generate", {"user_id": owner, "team_id": team, "models": [model]})["key"] @@ -145,8 +212,18 @@ def _proxy_with_one_seeded_row(gateway: Gateway, tmp_path: Path, pool_limit: int shutdown.chat() eventually(shutdown.daily_user_requests, lambda total: total == 1, seconds=60) yield shutdown - assert _api_requests("LiteLLM_DailyUserSpend", "user_id", owner) == 1 + REQUESTS_WHILE_BLOCKED - assert _api_requests("LiteLLM_DailyTeamSpend", "team_id", team) == 1 + REQUESTS_WHILE_BLOCKED + written: Final = 1 + requests + if settle_seconds: + eventually( + lambda: ( + _api_requests("LiteLLM_DailyUserSpend", "user_id", owner), + _api_requests("LiteLLM_DailyTeamSpend", "team_id", team), + ), + lambda totals: totals == (written, written), + seconds=settle_seconds, + ) + assert _api_requests("LiteLLM_DailyUserSpend", "user_id", owner) == written + assert _api_requests("LiteLLM_DailyTeamSpend", "team_id", team) == written @pytest.mark.covers("quota_management.spend_tracking.shutdown_cancel_keeps_in_flight_daily_batch") @@ -187,3 +264,48 @@ def test_daily_spend_batch_cancelled_while_waiting_for_a_row_lock_is_written_exa lambda: shutdown.logged(BATCH_DRAINED_LOG_LINE) and _waiting_on("LiteLLM_DailyUserSpend") == 1, holder.rollback, ) + + +@pytest.mark.parametrize( + ("cancel_timeout_seconds", "commit_fails_once"), + [ + pytest.param(60, False, id="cancel_budget_outlives_commit"), + pytest.param(5, False, id="commit_outlives_cancel_budget"), + pytest.param(60, True, id="commit_fails_within_cancel_budget"), + pytest.param(5, True, id="commit_fails_after_cancel_budget"), + ], +) +def test_daily_spend_batch_cancelled_while_postgres_is_committing_it_is_written_exactly_once( + gateway: Gateway, tmp_path: Path, cancel_timeout_seconds: int, commit_fails_once: bool +) -> None: + with ( + _proxy_with_one_seeded_row( + gateway, tmp_path, pool_limit=10, cancel_timeout_seconds=cancel_timeout_seconds, settle_seconds=90 + ) as shutdown, + psycopg.connect(os.environ["DATABASE_URL"]) as memberships, + ): + suffix: Final = _install_slow_commit(shutdown.owner, fails_once=commit_fails_once) + try: + shutdown.chat_while_spend_update_is_blocked(memberships, "LiteLLM_TeamMembership") + memberships.rollback() + shutdown.terminate_once( + lambda: shutdown.logged(BATCH_DRAINED_LOG_LINE) and _committing_daily_user_spend() == 1, + lambda: None, + ) + finally: + _drop_slow_commit(suffix) + + +def test_daily_spend_burst_across_two_workers_survives_shutdown_during_commit_exactly_once( + gateway: Gateway, tmp_path: Path +) -> None: + with _proxy_with_one_seeded_row( + gateway, tmp_path, pool_limit=10, settle_seconds=120, requests=BURST_REQUESTS, workers=2 + ) as shutdown: + suffix: Final = _install_slow_commit(shutdown.owner, fails_once=False) + try: + shutdown.burst(BURST_REQUESTS) + eventually(shutdown.spend_logs, lambda total: total == 1 + BURST_REQUESTS, seconds=60) + shutdown.terminate_once(lambda: _committing_daily_user_spend() >= 1, lambda: None) + finally: + _drop_slow_commit(suffix) diff --git a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py index daa07c8224a..bf3b9aed234 100644 --- a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py +++ b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py @@ -4501,3 +4501,241 @@ async def test_tag_batch_drained_from_redis_and_cancelled_mid_flight_is_restored await asyncio.wait_for(db.rolled_back.wait(), timeout=5) assert db.transaction_outcomes == ["rollback"] assert _daily_upserts(db, "LiteLLM_DailyTagSpend") == [] + + +class _CommittingDailySpendFakeDB(_DailySpendFakeDB): + """Runs the daily upsert at once but holds the COMMIT until released. A COMMIT that has left + the client lands on the server whether or not the client keeps waiting for the reply.""" + + def __init__(self) -> None: + super().__init__(failing_table=None) + self.committing = asyncio.Event() + self.commit_release = asyncio.Event() + self.transaction_outcomes: list[str] = [] + + @asynccontextmanager + async def _tx(self) -> AsyncIterator["_CommittingDailySpendFakeDB"]: + try: + yield self + except BaseException: + self.transaction_outcomes.append("rollback") + raise + self.committing.set() + try: + await self.commit_release.wait() + finally: + self.transaction_outcomes.append("commit") + + +@pytest.mark.parametrize(("queue_name", "entity_type", "entity_id_field", "table"), _DAILY_SPEND_ENTITIES) +@pytest.mark.asyncio +async def test_cancel_that_lands_while_the_daily_batch_is_committing_waits_for_the_commit_and_does_not_requeue_it( + queue_name: str, entity_type: str, entity_id_field: str, table: str +): + """Shutdown cancels the tick after the COMMIT has left for Postgres. The server finishes that + commit whatever the client does, so putting the batch back on the queue makes the final flush + write the same spend a second time. The tick has to wait for the commit's outcome instead.""" + db_writer = DBSpendUpdateWriter() + queue = _DAILY_SPEND_QUEUES[queue_name](db_writer) + await queue.add_update({"key-a": _daily_entity_txn(entity_id_field)}) + await queue.add_update({"key-a": _daily_entity_txn(entity_id_field)}) + db = _CommittingDailySpendFakeDB() + + def flush(prisma_db: _DailySpendFakeDB): + return db_writer._flush_daily_spend_queue( + queue=queue, + entity_type=entity_type, + commit=_DAILY_SPEND_COMMITS[entity_type], + n_retry_times=0, + prisma_client=_WindowSpendFakePrisma(prisma_db), + proxy_logging_obj=MagicMock(), + ) + + tick = asyncio.ensure_future(flush(db)) + await asyncio.wait_for(db.committing.wait(), timeout=5) + tick.cancel() + finished, _ = await asyncio.wait({tick}, timeout=0.2) + assert finished == {tick}, "the cancelled tick must hand the in-flight commit's outcome to the next flush" + with pytest.raises(asyncio.CancelledError): + tick.result() + assert len(queue.interrupted_commits) == 1 + + db.commit_release.set() + await queue.settle_interrupted_commits() + + assert db.transaction_outcomes == ["commit"] + (upsert,) = _daily_upserts(db, table) + assert _row_values(upsert, "api_requests") == [2] + assert queue.update_queue.empty(), "a batch whose COMMIT already left for the server must not be requeued" + + final_db = _DailySpendFakeDB(failing_table=None) + await flush(final_db) + assert _daily_upserts(final_db, table) == [], "the final flush must not write the committed batch again" + + +@pytest.mark.asyncio +async def test_tag_batch_drained_from_redis_and_cancelled_while_committing_is_not_restored(): + """Same in-flight COMMIT as the in-memory path, but the drained rows live in Redis. Restoring + them after the server committed writes the tag spend twice on the next tick.""" + db_writer = DBSpendUpdateWriter() + drained = {"key-a": cast(DailyTagSpendTransaction, _daily_entity_txn("tag"))} + redis_buffer = _DrainedTagRedisBuffer(drained) + db_writer.redis_update_buffer = cast(RedisUpdateBuffer, redis_buffer) + db = _CommittingDailySpendFakeDB() + + tick = asyncio.ensure_future( + db_writer._drain_and_commit_daily_tag_spend_from_redis( + prisma_client=_WindowSpendFakePrisma(db), + n_retry_times=0, + proxy_logging_obj=MagicMock(), + ) + ) + await asyncio.wait_for(db.committing.wait(), timeout=5) + tick.cancel() + finished, _ = await asyncio.wait({tick}, timeout=0.2) + assert finished == {tick}, "the cancelled drain must hand the in-flight commit's outcome to the next drain" + with pytest.raises(asyncio.CancelledError): + tick.result() + assert len(db_writer.interrupted_tag_commits) == 1 + + db.commit_release.set() + (settle,) = tuple(db_writer.interrupted_tag_commits) + await settle + + assert db.transaction_outcomes == ["commit"] + assert redis_buffer.restored == [], ( + "a tag batch whose COMMIT already left for the server must not be restored to Redis" + ) + + +class _CommitFailingDailySpendFakeDB(_CommittingDailySpendFakeDB): + """COMMIT leaves for the server but the reply comes back as a failure.""" + + @asynccontextmanager + async def _tx(self) -> AsyncIterator["_CommitFailingDailySpendFakeDB"]: + yield self + self.committing.set() + await self.commit_release.wait() + self.transaction_outcomes.append("commit_failed") + raise Exception("connection reset") + + +@pytest.mark.asyncio +async def test_cancel_while_committing_requeues_the_batch_when_the_commit_itself_fails(): + """Waiting for the in-flight commit's outcome must not swallow a real commit failure: + the batch still goes back on the queue and the next flush writes it once.""" + db_writer = DBSpendUpdateWriter() + queue = db_writer.daily_spend_update_queue + await queue.add_update({"key-a": _daily_txn()}) + await queue.add_update({"key-a": _daily_txn()}) + db = _CommitFailingDailySpendFakeDB() + + def flush(prisma_db: _DailySpendFakeDB): + return db_writer._flush_daily_spend_queue( + queue=queue, + entity_type="user", + commit=DBSpendUpdateWriter.update_daily_user_spend, + n_retry_times=0, + prisma_client=_WindowSpendFakePrisma(prisma_db), + proxy_logging_obj=MagicMock(), + ) + + tick = asyncio.ensure_future(flush(db)) + await asyncio.wait_for(db.committing.wait(), timeout=5) + tick.cancel() + finished, _ = await asyncio.wait({tick}, timeout=0.2) + assert finished == {tick}, "the cancelled tick must not eat the shutdown budget waiting on the commit" + with pytest.raises(asyncio.CancelledError): + tick.result() + + db.commit_release.set() + await queue.settle_interrupted_commits() + + assert db.transaction_outcomes == ["commit_failed"] + assert not queue.update_queue.empty(), "a batch whose COMMIT came back failed must be requeued" + + final_db = _DailySpendFakeDB(failing_table=None) + await flush(final_db) + (upsert,) = _daily_upserts(final_db, "LiteLLM_DailyUserSpend") + assert _row_values(upsert, "api_requests") == [2] + assert queue.update_queue.empty() + + +@pytest.mark.asyncio +async def test_shutdown_flush_that_lands_before_the_interrupted_commit_resolves_still_writes_a_failed_batch_once(): + """The cancelled tick returns right away, so a COMMIT can still be in flight when the + shutdown flush runs. If that commit later fails, the flush must first settle it, pick the + requeued rows back up, and write them exactly once instead of losing them.""" + db_writer = DBSpendUpdateWriter() + queue = db_writer.daily_spend_update_queue + await queue.add_update({"key-a": _daily_txn()}) + await queue.add_update({"key-a": _daily_txn()}) + db = _CommitFailingDailySpendFakeDB() + + def flush(prisma_db: _DailySpendFakeDB): + return db_writer._flush_daily_spend_queue( + queue=queue, + entity_type="user", + commit=DBSpendUpdateWriter.update_daily_user_spend, + n_retry_times=0, + prisma_client=_WindowSpendFakePrisma(prisma_db), + proxy_logging_obj=MagicMock(), + ) + + tick = asyncio.ensure_future(flush(db)) + await asyncio.wait_for(db.committing.wait(), timeout=5) + tick.cancel() + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(tick, timeout=5) + assert db.transaction_outcomes == [], "the COMMIT is still on the wire when the shutdown flush starts" + + final_db = _DailySpendFakeDB(failing_table=None) + shutdown_flush = asyncio.ensure_future(flush(final_db)) + finished, _ = await asyncio.wait({shutdown_flush}, timeout=0.2) + assert finished == set(), "the shutdown flush must wait for the interrupted commit's outcome" + assert _daily_upserts(final_db, "LiteLLM_DailyUserSpend") == [] + + db.commit_release.set() + await asyncio.wait_for(shutdown_flush, timeout=5) + + (upsert,) = _daily_upserts(final_db, "LiteLLM_DailyUserSpend") + assert _row_values(upsert, "api_requests") == [2] + assert queue.update_queue.empty() + + +@pytest.mark.asyncio +async def test_shutdown_drain_that_lands_before_the_interrupted_tag_commit_resolves_restores_a_failed_batch(): + """Same ordering for the Redis tag path: the shutdown drain must settle the interrupted + commit before the destructive drain, or a commit that fails late is never restored.""" + db_writer = DBSpendUpdateWriter() + drained = {"key-a": cast(DailyTagSpendTransaction, _daily_entity_txn("tag"))} + redis_buffer = _DrainedTagRedisBuffer(drained) + db_writer.redis_update_buffer = cast(RedisUpdateBuffer, redis_buffer) + db = _CommitFailingDailySpendFakeDB() + + def drain(prisma_db: _DailySpendFakeDB): + return db_writer._drain_and_commit_daily_tag_spend_from_redis( + prisma_client=_WindowSpendFakePrisma(prisma_db), + n_retry_times=0, + proxy_logging_obj=MagicMock(), + ) + + tick = asyncio.ensure_future(drain(db)) + await asyncio.wait_for(db.committing.wait(), timeout=5) + tick.cancel() + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(tick, timeout=5) + assert db.transaction_outcomes == [] + + final_db = _DailySpendFakeDB(failing_table=None) + shutdown_drain = asyncio.ensure_future(drain(final_db)) + finished, _ = await asyncio.wait({shutdown_drain}, timeout=0.2) + assert finished == set(), "the shutdown drain must wait for the interrupted commit's outcome" + assert _daily_upserts(final_db, "LiteLLM_DailyTagSpend") == [] + + db.commit_release.set() + await asyncio.wait_for(shutdown_drain, timeout=5) + + assert redis_buffer.restored == [drained], "a tag batch whose COMMIT came back failed must be restored to Redis" + (upsert,) = _daily_upserts(final_db, "LiteLLM_DailyTagSpend") + assert _row_values(upsert, "api_requests") == [1] From 8f7cea5fcd87229f4768e3cc4628fa68d034c17a Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 23 Sep 2026 23:21:59 -0700 Subject: [PATCH 021/120] fix(cost-map): sync openrouter prices and add fireworks ember-1 (#42889) Co-authored-by: kerry Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- ...odel_prices_and_context_window_backup.json | 27 ++++++++++++++++--- model_prices_and_context_window.json | 27 ++++++++++++++++--- 2 files changed, 48 insertions(+), 6 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 0b606f5db4e..e19ddaf72a7 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -41253,6 +41253,27 @@ "supports_vision": false, "supports_web_search": false }, + "openrouter/fireworks/ember-1": { + "input_cost_per_token": 3e-06, + "output_cost_per_token": 1.5e-05, + "cache_read_input_token_cost": 3e-07, + "litellm_provider": "openrouter", + "max_input_tokens": 1048576, + "max_output_tokens": 943718, + "max_tokens": 943718, + "mode": "chat", + "source": "https://openrouter.ai/api/v1/models", + "supports_audio_input": false, + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_parallel_function_calling": false, + "supports_pdf_input": false, + "supports_vision": true, + "supports_prompt_caching": true, + "supports_web_search": false + }, "openrouter/google/gemini-2.5-flash": { "cache_creation_input_token_cost": 8.33333333333333e-08, "cache_read_input_audio_token_cost": 1e-07, @@ -65519,7 +65540,7 @@ "supports_prompt_caching": true }, "openrouter/deepseek/deepseek-v4-flash-0731": { - "input_cost_per_token": 4e-08, + "input_cost_per_token": 3e-08, "output_cost_per_token": 3.2e-07, "cache_read_input_token_cost": 1.6e-08, "litellm_provider": "openrouter", @@ -67140,8 +67161,8 @@ "supports_web_search": false }, "openrouter/qwen/qwen3-30b-a3b-instruct-2507": { - "input_cost_per_token": 4.815e-08, - "output_cost_per_token": 1.9305e-07, + "input_cost_per_token": 1e-07, + "output_cost_per_token": 3e-07, "litellm_provider": "openrouter", "max_input_tokens": 262144, "max_output_tokens": 32000, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 0b606f5db4e..e19ddaf72a7 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -41253,6 +41253,27 @@ "supports_vision": false, "supports_web_search": false }, + "openrouter/fireworks/ember-1": { + "input_cost_per_token": 3e-06, + "output_cost_per_token": 1.5e-05, + "cache_read_input_token_cost": 3e-07, + "litellm_provider": "openrouter", + "max_input_tokens": 1048576, + "max_output_tokens": 943718, + "max_tokens": 943718, + "mode": "chat", + "source": "https://openrouter.ai/api/v1/models", + "supports_audio_input": false, + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_parallel_function_calling": false, + "supports_pdf_input": false, + "supports_vision": true, + "supports_prompt_caching": true, + "supports_web_search": false + }, "openrouter/google/gemini-2.5-flash": { "cache_creation_input_token_cost": 8.33333333333333e-08, "cache_read_input_audio_token_cost": 1e-07, @@ -65519,7 +65540,7 @@ "supports_prompt_caching": true }, "openrouter/deepseek/deepseek-v4-flash-0731": { - "input_cost_per_token": 4e-08, + "input_cost_per_token": 3e-08, "output_cost_per_token": 3.2e-07, "cache_read_input_token_cost": 1.6e-08, "litellm_provider": "openrouter", @@ -67140,8 +67161,8 @@ "supports_web_search": false }, "openrouter/qwen/qwen3-30b-a3b-instruct-2507": { - "input_cost_per_token": 4.815e-08, - "output_cost_per_token": 1.9305e-07, + "input_cost_per_token": 1e-07, + "output_cost_per_token": 3e-07, "litellm_provider": "openrouter", "max_input_tokens": 262144, "max_output_tokens": 32000, From 42d8d08815236862963fac4d2f3302bb3a51bf49 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 23 Sep 2026 23:25:24 -0700 Subject: [PATCH 022/120] fix(cost-map): source and chat completions endpoint for bedrock mantle gpt-5.4 and gpt-5.5 (#42890) Co-authored-by: kerry Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/model_prices_and_context_window_backup.json | 8 ++++++-- model_prices_and_context_window.json | 8 ++++++-- 2 files changed, 12 insertions(+), 4 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index e19ddaf72a7..c01f7c2f396 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -57528,6 +57528,7 @@ "mode": "responses", "use_openai_responses_path": true, "supported_endpoints": [ + "/v1/chat/completions", "/v1/responses" ], "supported_modalities": [ @@ -57545,7 +57546,8 @@ "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, - "supports_web_search": true + "supports_web_search": true, + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-55.html" }, "bedrock_mantle/openai.gpt-5.4": { "input_cost_per_token": 2.75e-06, @@ -57566,6 +57568,7 @@ "mode": "responses", "use_openai_responses_path": true, "supported_endpoints": [ + "/v1/chat/completions", "/v1/responses" ], "supported_modalities": [ @@ -57583,7 +57586,8 @@ "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, - "supports_web_search": true + "supports_web_search": true, + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-54.html" }, "bedrock_mantle/google.gemma-4-31b": { "input_cost_per_token": 1.4e-07, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index e19ddaf72a7..c01f7c2f396 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -57528,6 +57528,7 @@ "mode": "responses", "use_openai_responses_path": true, "supported_endpoints": [ + "/v1/chat/completions", "/v1/responses" ], "supported_modalities": [ @@ -57545,7 +57546,8 @@ "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, - "supports_web_search": true + "supports_web_search": true, + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-55.html" }, "bedrock_mantle/openai.gpt-5.4": { "input_cost_per_token": 2.75e-06, @@ -57566,6 +57568,7 @@ "mode": "responses", "use_openai_responses_path": true, "supported_endpoints": [ + "/v1/chat/completions", "/v1/responses" ], "supported_modalities": [ @@ -57583,7 +57586,8 @@ "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, - "supports_web_search": true + "supports_web_search": true, + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-54.html" }, "bedrock_mantle/google.gemma-4-31b": { "input_cost_per_token": 1.4e-07, From a5b9b6da4dc7bd8d95dc0af4a920a77a4a1f85b4 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 23 Sep 2026 23:41:10 -0700 Subject: [PATCH 023/120] feat(cost-map): add vertex_ai/gemini-3.8-live (#42891) * feat(cost-map): add vertex_ai/gemini-3.8-live Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(cost-map): price vertex_ai/gemini-3.8-live video tokens Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: kerry Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- ...odel_prices_and_context_window_backup.json | 36 +++++++++++++++++++ model_prices_and_context_window.json | 36 +++++++++++++++++++ 2 files changed, 72 insertions(+) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index c01f7c2f396..b73f12bb673 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -27384,6 +27384,42 @@ "supports_web_search": false, "supports_native_streaming": true }, + "vertex_ai/gemini-3.8-live": { + "input_cost_per_audio_token": 3e-06, + "input_cost_per_image_token": 1e-06, + "input_cost_per_token": 7.5e-07, + "input_cost_per_video_per_second": 3.3333333333333335e-05, + "input_cost_per_video_token": 1e-06, + "litellm_provider": "vertex_ai", + "max_input_tokens": 131072, + "max_output_tokens": 65536, + "max_tokens": 65536, + "mode": "realtime", + "output_cost_per_audio_token": 1.2e-05, + "output_cost_per_token": 4.5e-06, + "source": "https://ai.google.dev/gemini-api/docs/pricing", + "supported_endpoints": [ + "/vertex_ai/live", + "/v1/realtime" + ], + "supported_modalities": [ + "text", + "image", + "audio", + "video" + ], + "supported_output_modalities": [ + "text", + "audio" + ], + "supports_audio_input": true, + "supports_audio_output": true, + "supports_function_calling": true, + "supports_response_schema": false, + "supports_vision": true, + "supports_web_search": true, + "gemini_audio_only_live": true + }, "vertex_ai/gemini-3.1-pro-preview": { "prompt_cache_min_tokens": 4096, "cache_read_input_token_cost": 2e-07, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index c01f7c2f396..b73f12bb673 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -27384,6 +27384,42 @@ "supports_web_search": false, "supports_native_streaming": true }, + "vertex_ai/gemini-3.8-live": { + "input_cost_per_audio_token": 3e-06, + "input_cost_per_image_token": 1e-06, + "input_cost_per_token": 7.5e-07, + "input_cost_per_video_per_second": 3.3333333333333335e-05, + "input_cost_per_video_token": 1e-06, + "litellm_provider": "vertex_ai", + "max_input_tokens": 131072, + "max_output_tokens": 65536, + "max_tokens": 65536, + "mode": "realtime", + "output_cost_per_audio_token": 1.2e-05, + "output_cost_per_token": 4.5e-06, + "source": "https://ai.google.dev/gemini-api/docs/pricing", + "supported_endpoints": [ + "/vertex_ai/live", + "/v1/realtime" + ], + "supported_modalities": [ + "text", + "image", + "audio", + "video" + ], + "supported_output_modalities": [ + "text", + "audio" + ], + "supports_audio_input": true, + "supports_audio_output": true, + "supports_function_calling": true, + "supports_response_schema": false, + "supports_vision": true, + "supports_web_search": true, + "gemini_audio_only_live": true + }, "vertex_ai/gemini-3.1-pro-preview": { "prompt_cache_min_tokens": 4096, "cache_read_input_token_cost": 2e-07, From 6dbd65b23098644c75d948ed2c26bd74b5e0e73c Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 23 Sep 2026 23:42:24 -0700 Subject: [PATCH 024/120] fix(passthrough): log upstream 4xx/5xx error bodies and carry them into the failure hook (#42695) * test(integration): reproduce passthrough upstream error body missing from logs and spend row Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(passthrough): log upstream 4xx/5xx error bodies and carry them into the failure hook Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(error_normalization): let the passthrough prefix win over upstream body text Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(passthrough): honor message redaction for upstream error bodies Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(passthrough): bound the upstream error body read and sanitize it before logging Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * refactor(passthrough): use the Sequence import directly in the allowed-routes cast Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(passthrough): rechunk the upstream error stream so the preview read stays bounded Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(integration): audit matrix for passthrough upstream error visibility Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(passthrough): drop the restating docstring on the upstream failure logger Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(integration): drop the retired covers markers from the passthrough error tests Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(passthrough): keep the upstream status when the error body peek fails Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(passthrough): cover the relay aclose in the mid-read failure test Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(passthrough): relay decoded partial body on mid-read failure Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: yucheng Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/constants.py | 1 + .../litellm_core_utils/error_normalization.py | 2 +- .../pass_through_endpoints.py | 144 +++- tests/integration/_support/wire.py | 6 +- .../test_passthrough_upstream_error_chaos.py | 150 ++++ ...t_passthrough_upstream_error_visibility.py | 600 +++++++++++++++ .../test_error_normalization.py | 12 + .../test_pass_through_endpoints.py | 716 +++++++++++++++++- 8 files changed, 1595 insertions(+), 36 deletions(-) create mode 100644 tests/integration/observability/test_passthrough_upstream_error_chaos.py create mode 100644 tests/integration/observability/test_passthrough_upstream_error_visibility.py diff --git a/litellm/constants.py b/litellm/constants.py index e5b662bd515..dcc12ef2ba9 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1518,6 +1518,7 @@ PROMETHEUS_BUDGET_METRICS_REFRESH_INTERVAL_MINUTES: Final = int( CLOUDZERO_EXPORT_INTERVAL_MINUTES: Final = int(os.getenv("CLOUDZERO_EXPORT_INTERVAL_MINUTES", 60)) MCP_TOOL_NAME_PREFIX: Final = "mcp_tool" MAXIMUM_TRACEBACK_LINES_TO_LOG: Final = int(os.getenv("MAXIMUM_TRACEBACK_LINES_TO_LOG", 100)) +PASSTHROUGH_UPSTREAM_ERROR_BODY_MAX_LOG_CHARS: Final = 4096 # Headers to control callbacks X_LITELLM_DISABLE_CALLBACKS: Final = "x-litellm-disable-callbacks" diff --git a/litellm/litellm_core_utils/error_normalization.py b/litellm/litellm_core_utils/error_normalization.py index 61eb47a9675..be0098ec34b 100644 --- a/litellm/litellm_core_utils/error_normalization.py +++ b/litellm/litellm_core_utils/error_normalization.py @@ -60,13 +60,13 @@ class _HasProxyErrorType(Protocol): _MESSAGE_PATTERNS: Final[tuple[tuple[re.Pattern[str], str], ...]] = ( + (re.compile(r"upstream passthrough request failed", re.IGNORECASE), UPSTREAM_PASSTHROUGH), ( re.compile(r"budget has been exceeded|max budget|crossed budget", re.IGNORECASE), BUDGET_EXCEEDED, ), (re.compile(r"no healthy deployments?|no deployments available", re.IGNORECASE), NO_HEALTHY_DEPLOYMENTS), (re.compile(r"not allowed to access model due to tags configuration", re.IGNORECASE), MODEL_ACCESS_DENIED), - (re.compile(r"upstream passthrough request failed", re.IGNORECASE), UPSTREAM_PASSTHROUGH), (re.compile(r"is not supported for provider|not implemented", re.IGNORECASE), UNSUPPORTED_OPERATION), ( re.compile(r"context window|context length|(prompt|input) is too long|tokens? ?> ?\d+ ?maximum", re.IGNORECASE), diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index f985c1d49d1..a119335ba46 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -5,7 +5,7 @@ import json import posixpath import traceback from base64 import b64encode -from collections.abc import AsyncGenerator, Callable, Iterable, Mapping, Sequence +from collections.abc import AsyncGenerator, AsyncIterator, Callable, Iterable, Mapping, Sequence from dataclasses import dataclass from datetime import datetime from itertools import count, groupby @@ -41,6 +41,8 @@ from litellm._logging import verbose_proxy_logger from litellm._uuid import uuid from litellm.constants import ( MAXIMUM_TRACEBACK_LINES_TO_LOG, + PASSTHROUGH_UPSTREAM_ERROR_BODY_MAX_LOG_CHARS, + REDACTED_BY_LITELLM, SESSION_ID_OMITTED_METADATA_KEY, WEBSOCKET_CLOSE_REASON_MAX_BYTES, ) @@ -56,6 +58,7 @@ from litellm.litellm_core_utils.internal_call_metadata import MODEL_ACCESS_GROUP from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.litellm_logging import _get_masked_values from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER +from litellm.litellm_core_utils.redact_messages import should_redact_message_logging from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.llms.base_llm.managed_resources.utils import ( resolve_passthrough_managed_id_provider, @@ -849,23 +852,106 @@ def _resolve_team_callback_wiring( ) +def _truncate_upstream_error_body(body: str) -> str: + if len(body) <= PASSTHROUGH_UPSTREAM_ERROR_BODY_MAX_LOG_CHARS: + return body + return ( + f"{body[:PASSTHROUGH_UPSTREAM_ERROR_BODY_MAX_LOG_CHARS]}... " + f"(truncated at {PASSTHROUGH_UPSTREAM_ERROR_BODY_MAX_LOG_CHARS} chars)" + ) + + +def _sanitize_upstream_error_body(body: str) -> str: + return " ".join("".join(char if char.isprintable() else " " for char in body).split()) + + +class _PrefixReplayStream(httpx.AsyncByteStream): + def __init__(self, prefix: bytes, rest: AsyncIterator[bytes], upstream: httpx.Response) -> None: + self._prefix: Final = prefix + self._rest: Final = rest + self._upstream: Final = upstream + + async def __aiter__(self) -> AsyncIterator[bytes]: + if self._prefix: + yield self._prefix + async for chunk in self._rest: + yield chunk + + async def aclose(self) -> None: + await self._upstream.aclose() + + +async def _no_more_chunks() -> AsyncIterator[bytes]: + return + yield b"" + + +async def _read_error_body_preview( + stream: AsyncIterator[bytes], +) -> tuple[bytes, AsyncIterator[bytes]]: + collected: Final[list[bytes]] = [] # mutable-ok: accumulated until the preview byte budget, then joined once + total = 0 # rebind-ok: running byte count against the preview budget + try: + async for chunk in stream: + collected.append(chunk) + total += len(chunk) + if total > PASSTHROUGH_UPSTREAM_ERROR_BODY_MAX_LOG_CHARS: + break + except httpx.HTTPError as err: + partial: Final = b"".join(collected) + verbose_proxy_logger.warning( + "pass_through_endpoint: upstream error body read failed after %d bytes: %s", + len(partial), + type(err).__name__, + ) + return partial, _no_more_chunks() + return b"".join(collected), stream + + +def _headers_without_body_framing(headers: httpx.Headers) -> httpx.Headers: + return httpx.Headers( + [(name, value) for name, value in headers.raw if name.lower() not in (b"content-encoding", b"content-length")] + ) + + +async def _error_body_preview_and_relay(response: httpx.Response) -> tuple[str, httpx.Response]: + if response.is_stream_consumed: + return response.text, response + body_iter: Final = response.aiter_bytes() + prefix, rest = await _read_error_body_preview(body_iter) + preview_text: Final = prefix.decode(response.encoding or "utf-8", errors="replace") + return preview_text, httpx.Response( + status_code=response.status_code, + headers=_headers_without_body_framing(response.headers), + stream=_PrefixReplayStream(prefix=prefix, rest=rest, upstream=response), + request=response.request, + extensions=response.extensions, + ) + + async def _log_passthrough_upstream_failure( response: httpx.Response, user_api_key_dict: UserAPIKeyAuth, request_payload: dict, -) -> None: - """Fire LiteLLM-side failure hooks (spend tracking, alerting callbacks) for - an upstream 4xx/5xx passthrough response. - - Passthrough must return the upstream status/body/headers to the client - unchanged, so this never raises or transforms the response - it only - mirrors the monitoring side effect that ``post_call_failure_hook`` would - have received had the error originated inside LiteLLM. - """ + logging_obj: LiteLLMLoggingObj, +) -> httpx.Response: if response.status_code < 400: - return + return response from litellm.proxy.proxy_server import proxy_logging_obj + preview_text, relay_response = await _error_body_preview_and_relay(response) + upstream_error_body: Final = ( + REDACTED_BY_LITELLM + if should_redact_message_logging(logging_obj.model_call_details) + else _truncate_upstream_error_body(_sanitize_upstream_error_body(preview_text)) + ) + verbose_proxy_logger.warning( + "pass_through_endpoint: upstream %s %s returned %s: %s", + response.request.method, + response.url.copy_with(query=None, fragment=None), + response.status_code, + upstream_error_body, + ) try: response.raise_for_status() except httpx.HTTPStatusError: @@ -878,7 +964,7 @@ async def _log_passthrough_upstream_failure( # rate-limit errors already are. synthetic_exception: Final = HTTPException( status_code=response.status_code, - detail=f"Upstream passthrough request failed with status {response.status_code}", + detail=f"Upstream passthrough request failed with status {response.status_code}: {upstream_error_body}", ) try: await proxy_logging_obj.post_call_failure_hook( @@ -892,6 +978,7 @@ async def _log_passthrough_upstream_failure( "pass_through_endpoint: post_call_failure_hook raised for upstream error", exc_info=True, ) + return relay_response async def _relay_reporting_failures( @@ -1321,7 +1408,7 @@ async def pass_through_request( headers=response.headers, ) - await _log_passthrough_upstream_failure( + relay_response: Final = await _log_passthrough_upstream_failure( response=response, user_api_key_dict=user_api_key_dict, request_payload=_build_passthrough_failure_request_payload( @@ -1331,17 +1418,18 @@ async def pass_through_request( custom_llm_provider=custom_llm_provider, upstream_usage=upstream_usage, ), + logging_obj=logging_obj, ) # Call response headers hook for streaming pass-through _response_headers = HttpPassThroughEndpointHelpers.get_response_headers( - headers=response.headers, + headers=relay_response.headers, litellm_call_id=litellm_call_id, ) callback_headers = await proxy_logging_obj.post_call_response_headers_hook( data=_parsed_body or {}, user_api_key_dict=user_api_key_dict, - response=response, + response=relay_response, request_headers=dict(request.headers), ) if callback_headers: @@ -1352,7 +1440,7 @@ async def pass_through_request( stream=_own_streamed_managed_ids( stream=_relay_reporting_failures( stream=PassThroughStreamingHandler.chunk_processor( - response=response, + response=relay_response, request_body=_parsed_body, litellm_logging_obj=logging_obj, endpoint_type=endpoint_type, @@ -1360,7 +1448,7 @@ async def pass_through_request( passthrough_success_handler_obj=pass_through_endpoint_logging, url_route=str(url), ), - upstream_status=response.status_code, + upstream_status=relay_response.status_code, user_api_key_dict=user_api_key_dict, request_payload=_build_passthrough_failure_request_payload( parsed_body=_parsed_body, @@ -1374,10 +1462,10 @@ async def pass_through_request( user_api_key_dict=user_api_key_dict, ), ping_interval_seconds=litellm.sse_keepalive_ping_interval_seconds, - upstream_headers=response.headers, + upstream_headers=relay_response.headers, ), headers=_response_headers, - status_code=response.status_code, + status_code=relay_response.status_code, ) if state_raw_body is not None: @@ -1412,7 +1500,7 @@ async def pass_through_request( logging_obj.stream = True logging_obj.model_call_details["stream"] = True - await _log_passthrough_upstream_failure( + detected_relay_response: Final = await _log_passthrough_upstream_failure( response=response, user_api_key_dict=user_api_key_dict, request_payload=_build_passthrough_failure_request_payload( @@ -1422,17 +1510,18 @@ async def pass_through_request( custom_llm_provider=custom_llm_provider, upstream_usage=upstream_usage, ), + logging_obj=logging_obj, ) # Call response headers hook for detected streaming pass-through _response_headers = HttpPassThroughEndpointHelpers.get_response_headers( - headers=response.headers, + headers=detected_relay_response.headers, litellm_call_id=litellm_call_id, ) callback_headers = await proxy_logging_obj.post_call_response_headers_hook( data=_parsed_body or {}, user_api_key_dict=user_api_key_dict, - response=response, + response=detected_relay_response, request_headers=dict(request.headers), ) if callback_headers: @@ -1443,7 +1532,7 @@ async def pass_through_request( stream=_own_streamed_managed_ids( stream=_relay_reporting_failures( stream=PassThroughStreamingHandler.chunk_processor( - response=response, + response=detected_relay_response, request_body=_parsed_body, litellm_logging_obj=logging_obj, endpoint_type=endpoint_type, @@ -1451,7 +1540,7 @@ async def pass_through_request( passthrough_success_handler_obj=pass_through_endpoint_logging, url_route=str(url), ), - upstream_status=response.status_code, + upstream_status=detected_relay_response.status_code, user_api_key_dict=user_api_key_dict, request_payload=_build_passthrough_failure_request_payload( parsed_body=_parsed_body, @@ -1465,10 +1554,10 @@ async def pass_through_request( user_api_key_dict=user_api_key_dict, ), ping_interval_seconds=litellm.sse_keepalive_ping_interval_seconds, - upstream_headers=response.headers, + upstream_headers=detected_relay_response.headers, ), headers=_response_headers, - status_code=response.status_code, + status_code=detected_relay_response.status_code, ) if not _should_buffer_passthrough_response(response): @@ -1526,6 +1615,7 @@ async def pass_through_request( response=response, user_api_key_dict=user_api_key_dict, request_payload=failure_request_payload, + logging_obj=logging_obj, ) if response.status_code < 400 and response_body is not None and guardrails_to_run: @@ -3435,7 +3525,7 @@ async def _filter_endpoints_by_team_allowed_routes( for endpoint in pass_through_endpoints if endpoint.path in cast( # cast-ok: guarded above; team metadata stores this key as a list of route paths - "Sequence[str]", team_metadata.get("allowed_passthrough_routes") + Sequence[str], team_metadata.get("allowed_passthrough_routes") ) ] diff --git a/tests/integration/_support/wire.py b/tests/integration/_support/wire.py index 5052c34e021..ed96d4e4e83 100644 --- a/tests/integration/_support/wire.py +++ b/tests/integration/_support/wire.py @@ -43,7 +43,9 @@ class Wire: @contextmanager -def wire_server(respond: Callable[[Request], Reply], tls: ssl.SSLContext | None = None) -> Generator[Wire, None, None]: +def wire_server( + respond: Callable[[Request], Reply], tls: ssl.SSLContext | None = None, port: int = 0 +) -> Generator[Wire, None, None]: """Owned TCP peer; requests traverse the real HTTP client and serialization.""" received: Final[SimpleQueue[Request]] = SimpleQueue() errors: Final[SimpleQueue[Exception]] = SimpleQueue() @@ -114,7 +116,7 @@ def wire_server(respond: Callable[[Request], Reply], tls: ssl.SSLContext | None if tls is not None: self.socket = tls.wrap_socket(self.socket, server_side=True) - with OwnedHTTPServer(("127.0.0.1", 0), Handler) as server: + with OwnedHTTPServer(("127.0.0.1", port), Handler) as server: thread: Final = threading.Thread(target=server.serve_forever, kwargs={"poll_interval": 0.05}) thread.start() try: diff --git a/tests/integration/observability/test_passthrough_upstream_error_chaos.py b/tests/integration/observability/test_passthrough_upstream_error_chaos.py new file mode 100644 index 00000000000..d94b3b24954 --- /dev/null +++ b/tests/integration/observability/test_passthrough_upstream_error_chaos.py @@ -0,0 +1,150 @@ +import asyncio +import json +import re +import signal +from pathlib import Path +from typing import Final + +import httpx +import psutil +import pytest +import yaml +from integration._support.client import Gateway, eventually, object_value +from integration._support.database import read_rows +from integration._support.process import owned_proxy_process +from integration._support.wire import Reply, Request, wire_server +from pydantic import JsonValue + +_GENERATE_CONTENT: Final[dict[str, JsonValue]] = {"contents": [{"role": "user", "parts": [{"text": "hi"}]}]} +_NOT_FOUND_BODY: Final = json.dumps( + { + "error": { + "code": 404, + "message": "models/nope-9 is not found for this scripted upstream", + "status": "NOT_FOUND", + } + } +).encode() +_INTERNAL_BODY: Final = ( + '{"error":{"code":500,"message":"' + "chunked upstream failure body " * 200 + '","status":"INTERNAL"}}' +).encode() +_OK_CHUNKS: Final = tuple(f"data: ok-{index}\n\n".encode() for index in range(3)) +_STARTED_WORKER: Final = re.compile(r"Started server process \[(\d+)\]") + + +def _chaos_reply(request: Request) -> Reply: + if "streamGenerateContent" in request.target: + return Reply(status=500, chunks=tuple(_INTERNAL_BODY[i : i + 512] for i in range(0, len(_INTERNAL_BODY), 512))) + if "healthy-model" in request.target: + return Reply(status=200, chunks=_OK_CHUNKS, content_type="text/event-stream") + return Reply(status=404, body=_NOT_FOUND_BODY) + + +def _error_information(call_id: str) -> dict[str, JsonValue]: + rows: Final = eventually( + lambda: read_rows('SELECT metadata FROM "LiteLLM_SpendLogs" WHERE request_id=%s', (call_id,)), + lambda values: len(values) == 1, + seconds=70, + ) + metadata: Final = rows[0]["metadata"] + parsed: Final = json.loads(metadata) if isinstance(metadata, str) else object_value(metadata) + return object_value(parsed["error_information"]) + + +def _single_spend_row(call_id: str) -> None: + rows: Final = eventually( + lambda: read_rows('SELECT request_id FROM "LiteLLM_SpendLogs" WHERE request_id=%s', (call_id,)), + lambda values: len(values) == 1, + seconds=70, + ) + assert len(rows) == 1, call_id + + +async def _fire_burst( + base_url: str, key: str, count: int, *, tolerate_transport_errors: bool = False +) -> tuple[httpx.Response, ...]: + async def one(client: httpx.AsyncClient, index: int) -> httpx.Response: + if index % 3 == 0: + path: Final = "/gemini/v1beta/models/nope-9:generateContent" + elif index % 3 == 1: + path = "/gemini/v1beta/models/nope-9:streamGenerateContent?alt=sse" + else: + path = "/gemini/v1beta/models/healthy-model:streamGenerateContent?alt=sse" + return await client.post( + path, + json=_GENERATE_CONTENT, + headers={"Authorization": f"Bearer {key}", "x-goog-api-key": key}, + ) + + async with httpx.AsyncClient(base_url=base_url, timeout=30, trust_env=False) as client: + results: Final = await asyncio.gather( + *(one(client, index) for index in range(count)), return_exceptions=tolerate_transport_errors + ) + for result in results: + assert not isinstance(result, BaseException) or isinstance(result, httpx.TransportError), repr(result) + return tuple(result for result in results if isinstance(result, httpx.Response)) + + +async def test_passthrough_upstream_outage_mid_burst_still_logs_errors_once(gateway: Gateway, tmp_path: Path) -> None: + config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + path: Final = tmp_path / "chaos-outage.yaml" + with wire_server(_chaos_reply) as wire: + port: Final = int(wire.url.rsplit(":", 1)[1]) + config["environment_variables"] = {"GEMINI_API_BASE": wire.url, "GEMINI_API_KEY": "scripted"} + path.write_text(yaml.safe_dump(config)) + with owned_proxy_process(gateway, tmp_path, {}, config=path, workers=2) as owned: + candidate: Final = owned.gateway + burst: Final = asyncio.create_task(_fire_burst(str(candidate.client.base_url), candidate.key, 30)) + await asyncio.to_thread(eventually, lambda: wire.received.qsize(), lambda size: size >= 10, 30) + with wire_server(_chaos_reply, port=port): + responses: Final = await burst + assert len(responses) == 30 + for response in responses: + assert response.status_code in (200, 404, 500, 502), response.status_code + assert "x-litellm-call-id" in response.headers, response.status_code + assert len(_STARTED_WORKER.findall(owned.log.read_text())) >= 2 + for response in responses: + _single_spend_row(response.headers["x-litellm-call-id"]) + if response.status_code == 404: + error_information: Final = _error_information(response.headers["x-litellm-call-id"]) + assert "not found for this scripted upstream" in str(error_information["error_message"]), response.text + elif response.status_code == 500: + assert "chunked upstream failure body" in str( + _error_information(response.headers["x-litellm-call-id"])["error_message"] + ), response.text + + +async def test_passthrough_worker_sigkill_leaves_sibling_serving_and_logging(gateway: Gateway, tmp_path: Path) -> None: + config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + path: Final = tmp_path / "chaos-kill.yaml" + with wire_server(_chaos_reply) as wire: + config["environment_variables"] = {"GEMINI_API_BASE": wire.url, "GEMINI_API_KEY": "scripted"} + path.write_text(yaml.safe_dump(config)) + with owned_proxy_process(gateway, tmp_path, {}, config=path, workers=2) as owned: + candidate: Final = owned.gateway + workers: Final = eventually( + lambda: tuple(int(pid) for pid in _STARTED_WORKER.findall(owned.log.read_text())), + lambda pids: len(pids) == 2, + seconds=30, + ) + burst: Final = asyncio.create_task( + _fire_burst(str(candidate.client.base_url), candidate.key, 20, tolerate_transport_errors=True) + ) + await asyncio.to_thread(eventually, lambda: wire.received.qsize(), lambda size: size >= 5, 30) + psutil.Process(workers[0]).send_signal(signal.SIGKILL) + responses: Final = await burst + for response in responses: + assert response.status_code in (200, 404, 500, 502), response.status_code + follow_up: Final = candidate.request( + "POST", + "/gemini/v1beta/models/nope-9:generateContent", + _GENERATE_CONTENT, + headers={"x-goog-api-key": candidate.key}, + ) + assert follow_up.status_code == 404, follow_up.text + assert follow_up.json() == json.loads(_NOT_FOUND_BODY), follow_up.text + for response in responses: + if "x-litellm-call-id" in response.headers: + _single_spend_row(response.headers["x-litellm-call-id"]) + error_information: Final = _error_information(follow_up.headers["x-litellm-call-id"]) + assert "not found for this scripted upstream" in str(error_information["error_message"]), follow_up.text diff --git a/tests/integration/observability/test_passthrough_upstream_error_visibility.py b/tests/integration/observability/test_passthrough_upstream_error_visibility.py new file mode 100644 index 00000000000..bb18add2f2f --- /dev/null +++ b/tests/integration/observability/test_passthrough_upstream_error_visibility.py @@ -0,0 +1,600 @@ +import gzip +import json +from hashlib import sha256 +from pathlib import Path +from typing import Final + +import httpx +import pytest +import yaml +from integration._support.client import Gateway, eventually, object_value +from integration._support.database import read_rows +from integration._support.process import owned_proxy_process +from integration._support.wire import Reply, Request, wire_server +from openai import AsyncOpenAI, NotFoundError, OpenAI +from pydantic import JsonValue + +_UPSTREAM_ERROR: Final[dict[str, JsonValue]] = { + "error": { + "code": 404, + "message": "Publisher Model `publishers/anthropic/models/claude-nope-9` was not found or your project does not have access to it. Please ensure you are using a valid model version.", + "status": "NOT_FOUND", + } +} + + +def test_gemini_passthrough_upstream_error_body_reaches_proxy_log_and_spend_row( + gateway: Gateway, tmp_path: Path +) -> None: + def respond(request: Request) -> Reply: + return Reply(status=404, body=json.dumps(_UPSTREAM_ERROR).encode()) + + config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + path: Final = tmp_path / "gemini-passthrough.yaml" + with wire_server(respond) as wire: + config["environment_variables"] = {"GEMINI_API_BASE": wire.url, "GEMINI_API_KEY": "scripted"} + path.write_text(yaml.safe_dump(config)) + with owned_proxy_process(gateway, tmp_path, {}, config=path) as owned: + candidate: Final = owned.gateway + response: Final = candidate.request( + "POST", + "/gemini/v1beta/models/claude-nope-9:generateContent", + {"contents": [{"role": "user", "parts": [{"text": "hi"}]}]}, + headers={"x-goog-api-key": candidate.key}, + ) + assert response.status_code == 404, response.text + assert response.json() == _UPSTREAM_ERROR, response.text + try: + eventually( + lambda: owned.log.read_text(), + lambda text: "was not found or your project" in text, + seconds=30, + ) + except AssertionError: + pytest.fail( + f"upstream 404 body never reached the proxy log after {response.status_code} passthrough; " + f"log tail: {owned.log.read_text()[-2000:]}" + ) + rows: Final = eventually( + lambda: read_rows( + 'SELECT metadata FROM "LiteLLM_SpendLogs" WHERE request_id=%s', + (response.headers["x-litellm-call-id"],), + ), + lambda values: len(values) == 1, + seconds=70, + ) + metadata: Final = rows[0]["metadata"] + parsed: Final = json.loads(metadata) if isinstance(metadata, str) else object_value(metadata) + error_information: Final = object_value(parsed["error_information"]) + assert "was not found or your project" in str(error_information["error_message"]), response.text + assert error_information["error_code"] == "404", response.text + + +_GEMINI_MODEL_PATH: Final = "/gemini/v1beta/models/claude-nope-9:generateContent" +_GEMINI_STREAM_PATH: Final = "/gemini/v1beta/models/claude-nope-9:streamGenerateContent" +_GENERATE_CONTENT: Final[dict[str, JsonValue]] = {"contents": [{"role": "user", "parts": [{"text": "hi"}]}]} +_UPSTREAM_500_BODY: Final = ( + '{"error":{"code":500,"message":"' + "chunked upstream failure body " * 200 + '","status":"INTERNAL"}}' +).encode() + + +def _gemini_config(path: Path, wire_url: str) -> None: + config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + config["environment_variables"] = {"GEMINI_API_BASE": wire_url, "GEMINI_API_KEY": "scripted"} + path.write_text(yaml.safe_dump(config)) + + +def _gemini_headers(candidate: Gateway) -> dict[str, str]: + return {"Authorization": f"Bearer {candidate.key}", "x-goog-api-key": candidate.key} + + +def _spend_error_information(call_id: str) -> dict[str, JsonValue]: + rows: Final = eventually( + lambda: read_rows('SELECT metadata FROM "LiteLLM_SpendLogs" WHERE request_id=%s', (call_id,)), + lambda values: len(values) == 1, + seconds=70, + ) + metadata: Final = rows[0]["metadata"] + parsed: Final = json.loads(metadata) if isinstance(metadata, str) else object_value(metadata) + return object_value(parsed["error_information"]) + + +def _spend_status(call_id: str) -> str: + rows: Final = eventually( + lambda: read_rows('SELECT status FROM "LiteLLM_SpendLogs" WHERE request_id=%s', (call_id,)), + lambda values: len(values) == 1, + seconds=70, + ) + return str(rows[0]["status"]) + + +def _upstream_warning(log: Path, needle: str = "pass_through_endpoint: upstream") -> str: + text: Final = eventually(lambda: log.read_text(), lambda content: needle in content, seconds=30) + return next(line for line in text.splitlines() if needle in line) + + +def _upstream_warnings(log: Path, needle: str = "pass_through_endpoint: upstream") -> tuple[str, ...]: + return tuple(line for line in log.read_text().splitlines() if needle in line) + + +async def test_gemini_passthrough_async_client_404_body_reaches_proxy_log_and_spend_row( + gateway: Gateway, tmp_path: Path +) -> None: + def respond(request: Request) -> Reply: + return Reply(status=404, body=json.dumps(_UPSTREAM_ERROR).encode()) + + path: Final = tmp_path / "gemini-async.yaml" + with wire_server(respond) as wire: + _gemini_config(path, wire.url) + with owned_proxy_process(gateway, tmp_path, {}, config=path, workers=2) as owned: + candidate: Final = owned.gateway + async with httpx.AsyncClient( + base_url=str(candidate.client.base_url), timeout=15, trust_env=False + ) as async_client: + response: Final = await async_client.post( + _GEMINI_MODEL_PATH, json=_GENERATE_CONTENT, headers=_gemini_headers(candidate) + ) + assert response.status_code == 404, response.text + assert response.json() == _UPSTREAM_ERROR, response.text + warning: Final = _upstream_warning(owned.log) + assert "was not found or your project" in warning, warning + error_information: Final = _spend_error_information(response.headers["x-litellm-call-id"]) + assert "was not found or your project" in str(error_information["error_message"]), response.text + assert error_information["error_code"] == "404", response.text + + +def test_gemini_passthrough_streaming_500_relays_full_body_and_logs_bounded_preview( + gateway: Gateway, tmp_path: Path +) -> None: + body: Final = _UPSTREAM_500_BODY + assert len(body) == 6055 + chunks: Final = tuple(body[index * 512 : (index + 1) * 512] for index in range(11)) + (body[5632:],) + + def respond(request: Request) -> Reply: + return Reply(status=500, chunks=chunks) + + path: Final = tmp_path / "gemini-stream-500.yaml" + with wire_server(respond) as wire: + _gemini_config(path, wire.url) + with owned_proxy_process(gateway, tmp_path, {}, config=path, workers=2) as owned: + candidate: Final = owned.gateway + with candidate.client.stream( + "POST", + _GEMINI_STREAM_PATH, + params={"alt": "sse"}, + json=_GENERATE_CONTENT, + headers=_gemini_headers(candidate), + ) as response: + assert response.status_code == 500, response.text + streamed: Final = response.read() + assert streamed == body + warning: Final = _upstream_warning(owned.log) + assert warning.endswith("... (truncated at 4096 chars)"), warning + error_information: Final = _spend_error_information(response.headers["x-litellm-call-id"]) + error_message: Final = str(error_information["error_message"]) + assert error_message.endswith("... (truncated at 4096 chars)"), error_message + assert error_information["error_code"] == "500", error_message + + +def test_gemini_passthrough_success_logs_nothing_and_spend_row_is_success(gateway: Gateway, tmp_path: Path) -> None: + upstream_ok: Final = { + "candidates": [{"content": {"parts": [{"text": "hello"}], "role": "model"}}], + "usageMetadata": {"promptTokenCount": 3, "candidatesTokenCount": 2, "totalTokenCount": 5}, + } + + def respond(request: Request) -> Reply: + return Reply(status=200, body=json.dumps(upstream_ok).encode()) + + path: Final = tmp_path / "gemini-200.yaml" + with wire_server(respond) as wire: + _gemini_config(path, wire.url) + with owned_proxy_process(gateway, tmp_path, {}, config=path, workers=2) as owned: + candidate: Final = owned.gateway + response: Final = candidate.request( + "POST", _GEMINI_MODEL_PATH, _GENERATE_CONTENT, headers={"x-goog-api-key": candidate.key} + ) + assert response.status_code == 200, response.text + assert response.json() == upstream_ok, response.text + assert _spend_status(response.headers["x-litellm-call-id"]) == "success" + assert not _upstream_warnings(owned.log), owned.log.read_text()[-2000:] + + +def test_gemini_passthrough_streaming_200_relays_every_chunk(gateway: Gateway, tmp_path: Path) -> None: + chunks: Final = tuple(f"data: chunk-{index}\n\n".encode() for index in range(5)) + + def respond(request: Request) -> Reply: + return Reply(status=200, chunks=chunks, content_type="text/event-stream") + + path: Final = tmp_path / "gemini-stream-200.yaml" + with wire_server(respond) as wire: + _gemini_config(path, wire.url) + with owned_proxy_process(gateway, tmp_path, {}, config=path, workers=2) as owned: + candidate: Final = owned.gateway + with candidate.client.stream( + "POST", + _GEMINI_STREAM_PATH, + params={"alt": "sse"}, + json=_GENERATE_CONTENT, + headers=_gemini_headers(candidate), + ) as response: + assert response.status_code == 200 + streamed: Final = response.read() + assert streamed == b"".join(chunks) + assert not _upstream_warnings(owned.log), owned.log.read_text()[-2000:] + + +def test_config_pass_through_route_logs_body_and_strips_query(gateway: Gateway, tmp_path: Path) -> None: + upstream_error: Final = {"error": {"message": "max budget reached for this deployment"}} + + def respond(request: Request) -> Reply: + return Reply(status=403, body=json.dumps(upstream_error).encode()) + + config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + path: Final = tmp_path / "config-route.yaml" + with wire_server(respond) as wire: + config["general_settings"]["pass_through_endpoints"] = [ + { + "path": "/audit-pt", + "target": f"{wire.url}/upstream?trace=secret-q", + "include_subpath": True, + "headers": {"Authorization": "Bearer scripted"}, + } + ] + path.write_text(yaml.safe_dump(config)) + with owned_proxy_process(gateway, tmp_path, {}, config=path, workers=2) as owned: + candidate: Final = owned.gateway + response: Final = candidate.request("POST", "/audit-pt", _GENERATE_CONTENT) + assert response.status_code == 403, response.text + assert response.json() == upstream_error, response.text + warning: Final = _upstream_warning(owned.log) + assert "max budget reached for this deployment" in warning, warning + assert "?" not in warning and "secret-q" not in warning, warning + error_information: Final = _spend_error_information(response.headers["x-litellm-call-id"]) + assert error_information["normalized_error"] == "500_UPSTREAM_PASSTHROUGH", response.text + assert "max budget reached for this deployment" in str(error_information["error_message"]), response.text + + +_OPENAI_UPSTREAM_404: Final[dict[str, JsonValue]] = { + "error": {"message": "The model `nope-9` does not exist", "type": "invalid_request_error"} +} + + +def _openai_config(path: Path, wire_url: str) -> None: + config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + config["environment_variables"] = {"OPENAI_API_BASE": wire_url, "OPENAI_API_KEY": "scripted"} + path.write_text(yaml.safe_dump(config)) + + +def test_openai_passthrough_sdk_error_body_reaches_proxy_log_and_spend_row(gateway: Gateway, tmp_path: Path) -> None: + def respond(request: Request) -> Reply: + return Reply(status=404, body=json.dumps(_OPENAI_UPSTREAM_404).encode()) + + path: Final = tmp_path / "openai-404.yaml" + with wire_server(respond) as wire: + _openai_config(path, wire.url) + with owned_proxy_process(gateway, tmp_path, {}, config=path, workers=2) as owned: + candidate: Final = owned.gateway + with OpenAI( + api_key=candidate.key, + base_url=f"{str(candidate.client.base_url).rstrip('/')}/openai", + max_retries=0, + http_client=httpx.Client(timeout=15, trust_env=False), + ) as sdk: + with pytest.raises(NotFoundError) as raised: + sdk.chat.completions.create(model="nope-9", messages=[{"role": "user", "content": "hi"}]) + assert "does not exist" in str(raised.value), raised.value + warning: Final = _upstream_warning(owned.log) + assert "does not exist" in warning, warning + error_information: Final = _spend_error_information(raised.value.response.headers["x-litellm-call-id"]) + assert "does not exist" in str(error_information["error_message"]) + + +async def test_openai_passthrough_async_sdk_error_body_reaches_proxy_log_and_spend_row( + gateway: Gateway, tmp_path: Path +) -> None: + def respond(request: Request) -> Reply: + return Reply(status=404, body=json.dumps(_OPENAI_UPSTREAM_404).encode()) + + path: Final = tmp_path / "openai-async-404.yaml" + with wire_server(respond) as wire: + _openai_config(path, wire.url) + with owned_proxy_process(gateway, tmp_path, {}, config=path, workers=2) as owned: + candidate: Final = owned.gateway + async with AsyncOpenAI( + api_key=candidate.key, + base_url=f"{str(candidate.client.base_url).rstrip('/')}/openai", + max_retries=0, + http_client=httpx.AsyncClient(timeout=15, trust_env=False), + ) as sdk: + with pytest.raises(NotFoundError) as raised: + await sdk.chat.completions.create(model="nope-9", messages=[{"role": "user", "content": "hi"}]) + assert "does not exist" in str(raised.value), raised.value + warning: Final = _upstream_warning(owned.log) + assert "does not exist" in warning, warning + error_information: Final = _spend_error_information(raised.value.response.headers["x-litellm-call-id"]) + assert "does not exist" in str(error_information["error_message"]) + + +def test_gemini_passthrough_control_characters_cannot_forge_log_lines(gateway: Gateway, tmp_path: Path) -> None: + forged: Final = b'{"error": "line one"}\n2026-01-01 FAKE LOG LINE\x1b[31m\r' + b"x" * 4943 + b"\x00tail" + assert len(forged) == 5000 + + def respond(request: Request) -> Reply: + return Reply(status=502, body=forged, content_type="text/html") + + path: Final = tmp_path / "gemini-forged.yaml" + with wire_server(respond) as wire: + _gemini_config(path, wire.url) + with owned_proxy_process(gateway, tmp_path, {}, config=path, workers=2) as owned: + candidate: Final = owned.gateway + response: Final = candidate.request( + "POST", _GEMINI_MODEL_PATH, _GENERATE_CONTENT, headers={"x-goog-api-key": candidate.key} + ) + assert response.status_code == 502, response.text + assert response.content == forged, response.text + warning: Final = _upstream_warning(owned.log) + assert "\n" not in warning and "\x1b" not in warning, warning + assert "line one" in warning and "FAKE LOG LINE" in warning, warning + assert warning.endswith("... (truncated at 4096 chars)"), warning + error_information: Final = _spend_error_information(response.headers["x-litellm-call-id"]) + assert error_information["error_code"] == "502", response.text + + +def test_gemini_passthrough_empty_error_body_still_logged_and_proxy_serves(gateway: Gateway, tmp_path: Path) -> None: + def respond(request: Request) -> Reply: + if "claude-nope-9" in request.target: + return Reply(status=404, body=b"") + return Reply(status=200, body=b'{"candidates": [{"content": {"parts": [{"text": "ok"}]}}]}') + + path: Final = tmp_path / "gemini-empty.yaml" + with wire_server(respond) as wire: + _gemini_config(path, wire.url) + with owned_proxy_process(gateway, tmp_path, {}, config=path, workers=2) as owned: + candidate: Final = owned.gateway + response: Final = candidate.request( + "POST", _GEMINI_MODEL_PATH, _GENERATE_CONTENT, headers={"x-goog-api-key": candidate.key} + ) + assert response.status_code == 404, response.text + assert response.content == b"", response.text + warning: Final = _upstream_warning(owned.log) + assert "returned 404" in warning, warning + error_information: Final = _spend_error_information(response.headers["x-litellm-call-id"]) + assert error_information["error_code"] == "404", response.text + follow_up: Final = candidate.request( + "POST", + "/gemini/v1beta/models/healthy-model:generateContent", + _GENERATE_CONTENT, + headers={"x-goog-api-key": candidate.key}, + ) + assert follow_up.status_code == 200, follow_up.text + + +def test_gemini_passthrough_gzip_error_body_decoded_for_log_and_client(gateway: Gateway, tmp_path: Path) -> None: + upstream_error: Final = {"error": {"message": "gzipped upstream says the model is gone"}} + + def respond(request: Request) -> Reply: + return Reply( + status=400, + body=gzip.compress(json.dumps(upstream_error).encode()), + headers={"content-encoding": "gzip"}, + ) + + path: Final = tmp_path / "gemini-gzip.yaml" + with wire_server(respond) as wire: + _gemini_config(path, wire.url) + with owned_proxy_process(gateway, tmp_path, {}, config=path, workers=2) as owned: + candidate: Final = owned.gateway + response: Final = candidate.request( + "POST", _GEMINI_MODEL_PATH, _GENERATE_CONTENT, headers={"x-goog-api-key": candidate.key} + ) + assert response.status_code == 400, response.text + assert response.json() == upstream_error, response.text + warning: Final = _upstream_warning(owned.log) + assert "gzipped upstream says the model is gone" in warning, warning + + +def test_gemini_passthrough_streaming_gzip_error_body_decoded_for_log_and_client( + gateway: Gateway, tmp_path: Path +) -> None: + upstream_error: Final = {"error": {"message": "streamed gzip upstream denies the deployment"}} + compressed: Final = gzip.compress(json.dumps(upstream_error).encode()) + third: Final = len(compressed) // 3 + + def respond(request: Request) -> Reply: + return Reply( + status=403, + chunks=(compressed[:third], compressed[third : 2 * third], compressed[2 * third :]), + headers={"content-encoding": "gzip"}, + ) + + path: Final = tmp_path / "gemini-stream-gzip.yaml" + with wire_server(respond) as wire: + _gemini_config(path, wire.url) + with owned_proxy_process(gateway, tmp_path, {}, config=path, workers=2) as owned: + candidate: Final = owned.gateway + with candidate.client.stream( + "POST", + _GEMINI_STREAM_PATH, + params={"alt": "sse"}, + json=_GENERATE_CONTENT, + headers=_gemini_headers(candidate), + ) as response: + assert response.status_code == 403 + streamed: Final = response.read() + assert json.loads(streamed) == upstream_error, streamed + warning: Final = _upstream_warning(owned.log) + assert "streamed gzip upstream denies the deployment" in warning, warning + + +def test_gemini_passthrough_error_body_redacted_when_message_logging_off(gateway: Gateway, tmp_path: Path) -> None: + upstream_error: Final = {"error": {"message": "sensitive upstream explanation"}} + + def respond(request: Request) -> Reply: + return Reply(status=404, body=json.dumps(upstream_error).encode()) + + config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + path: Final = tmp_path / "gemini-redacted.yaml" + with wire_server(respond) as wire: + config["environment_variables"] = {"GEMINI_API_BASE": wire.url, "GEMINI_API_KEY": "scripted"} + config["litellm_settings"]["turn_off_message_logging"] = True + path.write_text(yaml.safe_dump(config)) + with owned_proxy_process(gateway, tmp_path, {}, config=path, workers=2) as owned: + candidate: Final = owned.gateway + response: Final = candidate.request( + "POST", _GEMINI_MODEL_PATH, _GENERATE_CONTENT, headers={"x-goog-api-key": candidate.key} + ) + assert response.status_code == 404, response.text + assert response.json() == upstream_error, response.text + warning: Final = _upstream_warning(owned.log) + assert "redacted-by-litellm" in warning, warning + assert "sensitive upstream explanation" not in warning, warning + error_information: Final = _spend_error_information(response.headers["x-litellm-call-id"]) + error_message: Final = str(error_information["error_message"]) + assert "redacted-by-litellm" in error_message, error_message + assert "sensitive upstream explanation" not in error_message, error_message + + +def test_gemini_passthrough_exact_4096_byte_body_logged_without_marker(gateway: Gateway, tmp_path: Path) -> None: + body: Final = b'{"error": "' + b"y" * 4083 + b'"}' + assert len(body) == 4096 + + def respond(request: Request) -> Reply: + return Reply(status=404, body=body) + + path: Final = tmp_path / "gemini-exact.yaml" + with wire_server(respond) as wire: + _gemini_config(path, wire.url) + with owned_proxy_process(gateway, tmp_path, {}, config=path, workers=2) as owned: + candidate: Final = owned.gateway + response: Final = candidate.request( + "POST", _GEMINI_MODEL_PATH, _GENERATE_CONTENT, headers={"x-goog-api-key": candidate.key} + ) + assert response.status_code == 404, response.text + warning: Final = _upstream_warning(owned.log) + assert body[:512].decode() in warning, warning + assert "(truncated at 4096 chars)" not in warning, warning + + +def test_gemini_passthrough_4097_byte_body_truncated_with_marker(gateway: Gateway, tmp_path: Path) -> None: + body: Final = b'{"error": "' + b"y" * 4084 + b'"}' + assert len(body) == 4097 + + def respond(request: Request) -> Reply: + return Reply(status=404, body=body) + + path: Final = tmp_path / "gemini-over.yaml" + with wire_server(respond) as wire: + _gemini_config(path, wire.url) + with owned_proxy_process(gateway, tmp_path, {}, config=path, workers=2) as owned: + candidate: Final = owned.gateway + response: Final = candidate.request( + "POST", _GEMINI_MODEL_PATH, _GENERATE_CONTENT, headers={"x-goog-api-key": candidate.key} + ) + assert response.status_code == 404, response.text + warning: Final = _upstream_warning(owned.log) + assert body[:512].decode() in warning, warning + assert warning.endswith("... (truncated at 4096 chars)"), warning + + +def test_gemini_passthrough_one_byte_stream_chunks_reassembled_and_logged(gateway: Gateway, tmp_path: Path) -> None: + body: Final = json.dumps(_UPSTREAM_ERROR).encode() + + def respond(request: Request) -> Reply: + return Reply(status=404, chunks=tuple(bytes([byte]) for byte in body)) + + path: Final = tmp_path / "gemini-one-byte.yaml" + with wire_server(respond) as wire: + _gemini_config(path, wire.url) + with owned_proxy_process(gateway, tmp_path, {}, config=path, workers=2) as owned: + candidate: Final = owned.gateway + with candidate.client.stream( + "POST", + _GEMINI_STREAM_PATH, + params={"alt": "sse"}, + json=_GENERATE_CONTENT, + headers=_gemini_headers(candidate), + ) as response: + assert response.status_code == 404 + streamed: Final = response.read() + assert streamed == body + warning: Final = _upstream_warning(owned.log) + assert "was not found or your project" in warning, warning + + +def test_gemini_passthrough_repeated_errors_each_get_row_and_log_line(gateway: Gateway, tmp_path: Path) -> None: + def respond(request: Request) -> Reply: + return Reply(status=404, body=json.dumps(_UPSTREAM_ERROR).encode()) + + path: Final = tmp_path / "gemini-twice.yaml" + with wire_server(respond) as wire: + _gemini_config(path, wire.url) + with owned_proxy_process(gateway, tmp_path, {}, config=path, workers=2) as owned: + candidate: Final = owned.gateway + responses: Final = tuple( + candidate.request( + "POST", _GEMINI_MODEL_PATH, _GENERATE_CONTENT, headers={"x-goog-api-key": candidate.key} + ) + for _ in range(2) + ) + call_ids: Final = tuple(response.headers["x-litellm-call-id"] for response in responses) + assert len(set(call_ids)) == 2 + for response in responses: + assert response.status_code == 404, response.text + error_information: Final = _spend_error_information(response.headers["x-litellm-call-id"]) + assert "was not found or your project" in str(error_information["error_message"]), response.text + eventually( + lambda: _upstream_warnings(owned.log, "returned 404"), + lambda lines: len(lines) == 2, + seconds=30, + ) + + +def test_budget_rejected_call_keeps_budget_normalized_error(gateway: Gateway, tmp_path: Path) -> None: + path: Final = tmp_path / "budget.yaml" + path.write_text(Path("tests/integration/proxy_config.yaml").read_text()) + with owned_proxy_process(gateway, tmp_path, {}, config=path, workers=2) as owned: + candidate: Final = owned.gateway + with candidate.scenario() as scenario: + model: Final = scenario.model() + key: Final = scenario.key(max_budget=0.000001) + first: Final = candidate.chat(model, key=key) + assert "id" in first, first + rejected: Final = candidate.request( + "POST", + "/v1/chat/completions", + {"model": model, "messages": [{"role": "user", "content": "over budget"}]}, + key=key, + ) + assert rejected.status_code == 422 and "budget_exceeded" in rejected.text, rejected.text + digest: Final = sha256(key.encode()).hexdigest() + rows: Final = eventually( + lambda: read_rows( + 'SELECT metadata FROM "LiteLLM_SpendLogs" WHERE api_key=%s', + (digest,), + ), + lambda values: any( + "BUDGET_EXCEEDED" + in str( + object_value( + json.loads(row["metadata"]) + if isinstance(row["metadata"], str) + else object_value(row["metadata"]) + )["error_information"] + ) + for row in values + ), + seconds=70, + ) + budget_rows: Final = tuple( + row + for row in rows + if "BUDGET_EXCEEDED" + in str( + object_value( + json.loads(row["metadata"]) + if isinstance(row["metadata"], str) + else object_value(row["metadata"]) + )["error_information"] + ) + ) + assert len(budget_rows) == 1, budget_rows diff --git a/tests/test_litellm/litellm_core_utils/test_error_normalization.py b/tests/test_litellm/litellm_core_utils/test_error_normalization.py index d65b6d316ac..87b9463cb01 100644 --- a/tests/test_litellm/litellm_core_utils/test_error_normalization.py +++ b/tests/test_litellm/litellm_core_utils/test_error_normalization.py @@ -165,6 +165,18 @@ def test_variants_of_one_failure_share_a_normalized_error(messages: tuple[Except assert normalized == {expected} +def test_normalize_error_passthrough_prefix_wins_over_upstream_body_text() -> None: + from fastapi import HTTPException + + for detail in ( + 'Upstream passthrough request failed with status 400: {"error": {"message": "no deployments available for this model"}}', + 'Upstream passthrough request failed with status 400: {"error": {"message": "max budget reached"}}', + ): + exc = HTTPException(status_code=400, detail=detail) + message = f"400: {detail}" + assert normalize_error(exc, "400", message) == "500_UPSTREAM_PASSTHROUGH", message + + def test_router_no_healthy_deployment_wording_clusters_as_no_healthy_deployments() -> None: for message in (RouterErrors.no_healthy_deployments.value, "No healthy deployments found."): exc = litellm.BadRequestError(message, llm_provider="openai", model="gpt-4o") diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py b/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py index 6ad850866b7..7a64d5f2218 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py @@ -1,8 +1,10 @@ import asyncio +import gzip import json import logging import os import sys +import zlib from collections.abc import Callable from contextlib import ExitStack, contextmanager from io import BytesIO @@ -12,12 +14,14 @@ from unittest.mock import AsyncMock, MagicMock, patch import httpx import pytest -from fastapi import Request, Response, UploadFile +from fastapi import HTTPException, Request, Response, UploadFile +from fastapi.responses import StreamingResponse from pydantic import ValidationError from starlette.datastructures import FormData, Headers, QueryParams from starlette.datastructures import UploadFile as StarletteUploadFile import litellm +from litellm._logging import verbose_proxy_logger from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.proxy._types import ProxyException, UserAPIKeyAuth @@ -27,6 +31,8 @@ from litellm.proxy.pass_through_endpoints.pass_through_endpoints import ( HttpPassThroughEndpointHelpers, InitPassThroughEndpointHelpers, _registered_pass_through_routes, + _truncate_upstream_error_body, + _with_trace_context, chat_completion_pass_through_endpoint, create_pass_through_route, initialize_pass_through_endpoints, @@ -34,7 +40,6 @@ from litellm.proxy.pass_through_endpoints.pass_through_endpoints import ( resolve_llm_passthrough_timeout, resolve_pass_through_request_timeout, websocket_passthrough_request, - _with_trace_context, ) from litellm.proxy.pass_through_endpoints.success_handler import ( PassThroughEndpointLogging, @@ -4126,6 +4131,705 @@ async def test_pass_through_request_streaming_upstream_error_returned_unchanged( assert failure_call_kwargs["original_exception"].status_code == 403 +class _UpstreamErrorBodyStream(httpx.AsyncByteStream): + def __init__(self, body: bytes) -> None: + self._body: Final = body + + async def __aiter__(self): + yield self._body + + +def _upstream_error_request() -> MagicMock: + mock_request: Final = MagicMock(spec=Request) + mock_request.method = "POST" + mock_request.url = "http://test-proxy.com/mock-upstream/v1beta/models/claude-nope-9:generateContent" + mock_request.body = AsyncMock(return_value=b'{"contents": []}') + mock_request.headers = Headers({"content-type": "application/json"}) + mock_request.query_params = QueryParams({}) + return mock_request + + +@pytest.mark.asyncio +async def test_pass_through_request_non_streaming_upstream_error_body_logged_and_in_failure_detail( + caplog: pytest.LogCaptureFixture, +): + upstream_body: Final = { + "error": { + "code": 404, + "message": "Publisher Model `publishers/anthropic/models/claude-nope-9` was not found or your project does not have access", + "status": "NOT_FOUND", + } + } + upstream_content: Final = json.dumps(upstream_body).encode("utf-8") + upstream_response: Final = httpx.Response( + status_code=404, + headers={"content-type": "application/json"}, + content=upstream_content, + request=httpx.Request("POST", "http://target-api.com/v1beta/models/claude-nope-9:generateContent"), + ) + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + with patch("litellm.proxy.proxy_server.proxy_logging_obj") as mock_proxy_logging: + with patch( + "litellm.proxy.pass_through_endpoints.pass_through_endpoints.get_async_httpx_client" + ) as mock_get_client: + with patch( + "litellm.proxy.pass_through_endpoints.pass_through_endpoints.ProxyBaseLLMRequestProcessing" + ) as mock_processing: + mock_proxy_logging.pre_call_hook = AsyncMock(return_value={}) + mock_proxy_logging.post_call_failure_hook = AsyncMock() + mock_proxy_logging.post_call_response_headers_hook = AsyncMock(return_value=None) + mock_processing.get_custom_headers.return_value = {} + + async_client: Final = MagicMock() + async_client.build_request = MagicMock(return_value=MagicMock()) + async_client.send = AsyncMock(return_value=upstream_response) + mock_get_client.return_value = MagicMock(client=async_client) + + response: Final = await pass_through_request( + request=_upstream_error_request(), + target="http://target-api.com/v1beta/models/claude-nope-9:generateContent", + custom_headers={}, + user_api_key_dict=MagicMock(), + ) + + warning_messages: Final = [record.getMessage() for record in caplog.records if record.levelno == logging.WARNING] + upstream_warnings: Final = [ + message for message in warning_messages if "upstream" in message and "returned 404" in message + ] + assert len(upstream_warnings) == 1, warning_messages + assert "was not found or your project" in upstream_warnings[0] + assert "/v1beta/models/claude-nope-9:generateContent" in upstream_warnings[0] + + assert response.status_code == 404 + assert response.body == upstream_content + + mock_proxy_logging.post_call_failure_hook.assert_called_once() + failure_call_kwargs: Final = mock_proxy_logging.post_call_failure_hook.call_args.kwargs + original_exception: Final = failure_call_kwargs["original_exception"] + assert isinstance(original_exception, HTTPException) + assert original_exception.status_code == 404 + assert "was not found or your project" in original_exception.detail + + +@pytest.mark.asyncio +async def test_pass_through_request_streaming_upstream_error_body_reaches_client_and_failure_detail(): + upstream_content: Final = ( + b'data: {"error": {"code": 403, "message": "stream access was not found or your project lacks"}}\n\n' + ) + upstream_response: Final = httpx.Response( + status_code=403, + headers={"content-type": "text/event-stream"}, + stream=_UpstreamErrorBodyStream(upstream_content), + request=httpx.Request("POST", "http://target-api.com/v1beta/models/claude-nope-9:streamGenerateContent"), + ) + + with patch("litellm.proxy.proxy_server.proxy_logging_obj") as mock_proxy_logging: + with patch( + "litellm.proxy.pass_through_endpoints.pass_through_endpoints.get_async_httpx_client" + ) as mock_get_client: + with patch( + "litellm.proxy.pass_through_endpoints.pass_through_endpoints.pass_through_endpoint_logging.pass_through_async_success_handler" + ) as mock_success_handler: + mock_proxy_logging.pre_call_hook = AsyncMock(return_value={}) + mock_proxy_logging.post_call_failure_hook = AsyncMock() + mock_proxy_logging.post_call_response_headers_hook = AsyncMock(return_value=None) + mock_success_handler.return_value = None + + async_client: Final = MagicMock() + async_client.build_request = MagicMock(return_value=MagicMock()) + async_client.send = AsyncMock(return_value=upstream_response) + mock_get_client.return_value = MagicMock(client=async_client) + + response: Final = await pass_through_request( + request=_upstream_error_request(), + target="http://target-api.com/v1beta/models/claude-nope-9:streamGenerateContent", + custom_headers={}, + user_api_key_dict=MagicMock(), + stream=True, + ) + + assert isinstance(response, StreamingResponse) + assert response.status_code == 403 + streamed_chunks: Final = [chunk async for chunk in response.body_iterator] + streamed_bytes: Final = b"".join( + chunk if isinstance(chunk, bytes) else chunk.encode("utf-8") for chunk in streamed_chunks + ) + assert streamed_bytes == upstream_content + + mock_proxy_logging.post_call_failure_hook.assert_called_once() + original_exception: Final = mock_proxy_logging.post_call_failure_hook.call_args.kwargs["original_exception"] + assert "was not found or your project" in original_exception.detail + + +@pytest.mark.asyncio +async def test_truncate_upstream_error_body_caps_at_log_limit(): + short_body: Final = "x" * 4096 + assert _truncate_upstream_error_body(short_body) == short_body + + long_body: Final = "a" * 5000 + truncated: Final = _truncate_upstream_error_body(long_body) + assert truncated == f"{'a' * 4096}... (truncated at 4096 chars)" + + upstream_response: Final = httpx.Response( + status_code=500, + headers={"content-type": "text/plain"}, + content=long_body.encode("utf-8"), + request=httpx.Request("POST", "http://target-api.com/api/big-error"), + ) + with patch("litellm.proxy.proxy_server.proxy_logging_obj") as mock_proxy_logging: + with patch( + "litellm.proxy.pass_through_endpoints.pass_through_endpoints.get_async_httpx_client" + ) as mock_get_client: + with patch( + "litellm.proxy.pass_through_endpoints.pass_through_endpoints.ProxyBaseLLMRequestProcessing" + ) as mock_processing: + mock_proxy_logging.pre_call_hook = AsyncMock(return_value={}) + mock_proxy_logging.post_call_failure_hook = AsyncMock() + mock_proxy_logging.post_call_response_headers_hook = AsyncMock(return_value=None) + mock_processing.get_custom_headers.return_value = {} + + async_client: Final = MagicMock() + async_client.build_request = MagicMock(return_value=MagicMock()) + async_client.send = AsyncMock(return_value=upstream_response) + mock_get_client.return_value = MagicMock(client=async_client) + + await pass_through_request( + request=_upstream_error_request(), + target="http://target-api.com/api/big-error", + custom_headers={}, + user_api_key_dict=MagicMock(), + ) + + detail: Final = mock_proxy_logging.post_call_failure_hook.call_args.kwargs["original_exception"].detail + assert detail == f"Upstream passthrough request failed with status 500: {'a' * 4096}... (truncated at 4096 chars)" + + +@pytest.mark.asyncio +async def test_pass_through_request_upstream_error_log_strips_provider_key_from_url(): + upstream_content: Final = b'{"error": "denied"}' + upstream_response: Final = httpx.Response( + status_code=404, + headers={"content-type": "application/json"}, + content=upstream_content, + request=httpx.Request( + "POST", + "http://target-api.com/v1beta/models/claude-nope-9:generateContent?key=AIzaSySecretProviderKey123", + ), + ) + + with patch.object(verbose_proxy_logger, "warning") as mock_warning: + with patch("litellm.proxy.proxy_server.proxy_logging_obj") as mock_proxy_logging: + with patch( + "litellm.proxy.pass_through_endpoints.pass_through_endpoints.get_async_httpx_client" + ) as mock_get_client: + with patch( + "litellm.proxy.pass_through_endpoints.pass_through_endpoints.ProxyBaseLLMRequestProcessing" + ) as mock_processing: + mock_proxy_logging.pre_call_hook = AsyncMock(return_value={}) + mock_proxy_logging.post_call_failure_hook = AsyncMock() + mock_proxy_logging.post_call_response_headers_hook = AsyncMock(return_value=None) + mock_processing.get_custom_headers.return_value = {} + + async_client: Final = MagicMock() + async_client.build_request = MagicMock(return_value=MagicMock()) + async_client.send = AsyncMock(return_value=upstream_response) + mock_get_client.return_value = MagicMock(client=async_client) + + await pass_through_request( + request=_upstream_error_request(), + target="http://target-api.com/v1beta/models/claude-nope-9:generateContent", + custom_headers={}, + user_api_key_dict=MagicMock(), + ) + + upstream_warnings: Final = [ + call + for call in mock_warning.call_args_list + if call.args[0] == "pass_through_endpoint: upstream %s %s returned %s: %s" + ] + assert len(upstream_warnings) == 1, mock_warning.call_args_list + logged_url: Final = str(upstream_warnings[0].args[2]) + assert "/v1beta/models/claude-nope-9:generateContent" in logged_url + assert "AIzaSySecretProviderKey123" not in logged_url + assert "key=" not in logged_url + + +@pytest.mark.asyncio +@pytest.mark.parametrize("turn_off_message_logging", [True, False]) +async def test_passthrough_upstream_error_body_redacted_when_message_logging_off( + turn_off_message_logging: bool, +): + upstream_content: Final = b'{"error": {"message": "upstream body says the project was not found"}}' + upstream_response: Final = httpx.Response( + status_code=404, + headers={"content-type": "application/json"}, + content=upstream_content, + request=httpx.Request("POST", "http://target-api.com/v1beta/models/claude-nope-9:generateContent"), + ) + user_api_key_dict: Final = MagicMock() + user_api_key_dict.metadata = { + "logging": [ + { + "callback_name": "prometheus", + "callback_type": "success_and_failure", + "callback_vars": {"turn_off_message_logging": turn_off_message_logging}, + } + ] + } + user_api_key_dict.team_metadata = None + user_api_key_dict.team_id = None + + with patch.object(verbose_proxy_logger, "warning") as mock_warning: + with patch("litellm.proxy.proxy_server.proxy_logging_obj") as mock_proxy_logging: + with patch( + "litellm.proxy.pass_through_endpoints.pass_through_endpoints.get_async_httpx_client" + ) as mock_get_client: + with patch( + "litellm.proxy.pass_through_endpoints.pass_through_endpoints.ProxyBaseLLMRequestProcessing" + ) as mock_processing: + mock_proxy_logging.pre_call_hook = AsyncMock(return_value={}) + mock_proxy_logging.post_call_failure_hook = AsyncMock() + mock_proxy_logging.post_call_response_headers_hook = AsyncMock(return_value=None) + mock_processing.get_custom_headers.return_value = {} + + async_client: Final = MagicMock() + async_client.build_request = MagicMock(return_value=MagicMock()) + async_client.send = AsyncMock(return_value=upstream_response) + mock_get_client.return_value = MagicMock(client=async_client) + + response: Final = await pass_through_request( + request=_upstream_error_request(), + target="http://target-api.com/v1beta/models/claude-nope-9:generateContent", + custom_headers={}, + user_api_key_dict=user_api_key_dict, + ) + + assert response.status_code == 404 + assert response.body == upstream_content + + upstream_warnings: Final = [ + call + for call in mock_warning.call_args_list + if call.args[0] == "pass_through_endpoint: upstream %s %s returned %s: %s" + ] + assert len(upstream_warnings) == 1, mock_warning.call_args_list + logged_body: Final = str(upstream_warnings[0].args[4]) + + mock_proxy_logging.post_call_failure_hook.assert_called_once() + detail: Final = mock_proxy_logging.post_call_failure_hook.call_args.kwargs["original_exception"].detail + + if turn_off_message_logging: + assert logged_body == "redacted-by-litellm" + assert "upstream body says the project was not found" not in logged_body + assert detail == "Upstream passthrough request failed with status 404: redacted-by-litellm" + else: + assert "upstream body says the project was not found" in logged_body + assert detail == f"Upstream passthrough request failed with status 404: {upstream_content.decode()}" + + +class _ChunkedUpstreamErrorBodyStream(httpx.AsyncByteStream): + def __init__(self, chunks: tuple[bytes, ...]) -> None: + self._chunks: Final = chunks + self.served: int = 0 + + async def __aiter__(self): + for chunk in self._chunks: + self.served += 1 + yield chunk + + +@pytest.mark.asyncio +async def test_pass_through_request_streaming_upstream_error_reads_only_preview_and_relays_full_body(): + chunk_size: Final = 1024 + chunks: Final = tuple(b"x" * chunk_size for _ in range(10)) + upstream_content: Final = b"".join(chunks) + body_stream: Final = _ChunkedUpstreamErrorBodyStream(chunks) + upstream_response: Final = httpx.Response( + status_code=500, + headers={"content-type": "text/plain"}, + stream=body_stream, + request=httpx.Request("POST", "http://target-api.com/v1beta/models/claude-nope-9:streamGenerateContent"), + ) + + served_at_warning: list[int] = [] + real_warning: Final = verbose_proxy_logger.warning + + def _recording_warning(*args, **kwargs): + if args and args[0] == "pass_through_endpoint: upstream %s %s returned %s: %s": + served_at_warning.append(body_stream.served) + return real_warning(*args, **kwargs) + + with patch.object(verbose_proxy_logger, "warning", side_effect=_recording_warning): + with patch("litellm.proxy.proxy_server.proxy_logging_obj") as mock_proxy_logging: + with patch( + "litellm.proxy.pass_through_endpoints.pass_through_endpoints.get_async_httpx_client" + ) as mock_get_client: + with patch( + "litellm.proxy.pass_through_endpoints.pass_through_endpoints.pass_through_endpoint_logging.pass_through_async_success_handler" + ) as mock_success_handler: + mock_proxy_logging.pre_call_hook = AsyncMock(return_value={}) + mock_proxy_logging.post_call_failure_hook = AsyncMock() + mock_proxy_logging.post_call_response_headers_hook = AsyncMock(return_value=None) + mock_success_handler.return_value = None + + async_client: Final = MagicMock() + async_client.build_request = MagicMock(return_value=MagicMock()) + async_client.send = AsyncMock(return_value=upstream_response) + mock_get_client.return_value = MagicMock(client=async_client) + + response: Final = await pass_through_request( + request=_upstream_error_request(), + target="http://target-api.com/v1beta/models/claude-nope-9:streamGenerateContent", + custom_headers={}, + user_api_key_dict=MagicMock(), + stream=True, + ) + + assert isinstance(response, StreamingResponse) + assert response.status_code == 500 + streamed_chunks: Final = [chunk async for chunk in response.body_iterator] + streamed_bytes: Final = b"".join( + chunk if isinstance(chunk, bytes) else chunk.encode("utf-8") for chunk in streamed_chunks + ) + assert streamed_bytes == upstream_content + + assert served_at_warning == [5], ( + "each raw chunk is yielded as-is; five 1024-byte chunks are the first point the preview budget is exceeded" + ) + expected_body: Final = f"{'x' * 4096}... (truncated at 4096 chars)" + assert ( + mock_proxy_logging.post_call_failure_hook.call_args.kwargs["original_exception"].detail + == f"Upstream passthrough request failed with status 500: {expected_body}" + ) + + +@pytest.mark.asyncio +async def test_pass_through_request_streaming_upstream_error_single_large_chunk_stays_bounded(): + first_chunk: Final = b"x" * 65536 + second_chunk: Final = b'{"error": "tail"}' + upstream_content: Final = first_chunk + second_chunk + body_stream: Final = _ChunkedUpstreamErrorBodyStream((first_chunk, second_chunk)) + upstream_response: Final = httpx.Response( + status_code=500, + headers={"content-type": "text/plain"}, + stream=body_stream, + request=httpx.Request("POST", "http://target-api.com/v1beta/models/claude-nope-9:streamGenerateContent"), + ) + + served_at_warning: list[int] = [] + real_warning: Final = verbose_proxy_logger.warning + + def _recording_warning(*args, **kwargs): + if args and args[0] == "pass_through_endpoint: upstream %s %s returned %s: %s": + served_at_warning.append(body_stream.served) + return real_warning(*args, **kwargs) + + with patch.object(verbose_proxy_logger, "warning", side_effect=_recording_warning): + with patch("litellm.proxy.proxy_server.proxy_logging_obj") as mock_proxy_logging: + with patch( + "litellm.proxy.pass_through_endpoints.pass_through_endpoints.get_async_httpx_client" + ) as mock_get_client: + with patch( + "litellm.proxy.pass_through_endpoints.pass_through_endpoints.pass_through_endpoint_logging.pass_through_async_success_handler" + ) as mock_success_handler: + mock_proxy_logging.pre_call_hook = AsyncMock(return_value={}) + mock_proxy_logging.post_call_failure_hook = AsyncMock() + mock_proxy_logging.post_call_response_headers_hook = AsyncMock(return_value=None) + mock_success_handler.return_value = None + + async_client: Final = MagicMock() + async_client.build_request = MagicMock(return_value=MagicMock()) + async_client.send = AsyncMock(return_value=upstream_response) + mock_get_client.return_value = MagicMock(client=async_client) + + response: Final = await pass_through_request( + request=_upstream_error_request(), + target="http://target-api.com/v1beta/models/claude-nope-9:streamGenerateContent", + custom_headers={}, + user_api_key_dict=MagicMock(), + stream=True, + ) + + assert isinstance(response, StreamingResponse) + assert response.status_code == 500 + streamed_chunks: Final = [chunk async for chunk in response.body_iterator] + streamed_bytes: Final = b"".join( + chunk if isinstance(chunk, bytes) else chunk.encode("utf-8") for chunk in streamed_chunks + ) + assert streamed_bytes == upstream_content + + assert served_at_warning == [1], ( + "the rechunked preview is served from the first raw chunk; the second must not be pulled before the warning" + ) + expected_body: Final = f"{'x' * 4096}... (truncated at 4096 chars)" + assert ( + mock_proxy_logging.post_call_failure_hook.call_args.kwargs["original_exception"].detail + == f"Upstream passthrough request failed with status 500: {expected_body}" + ) + + +class _UpstreamErrorBodyStreamDropping(httpx.AsyncByteStream): + async def __aiter__(self): + yield b'{"error": "half' + raise httpx.ReadError("peer reset") + + +@pytest.mark.asyncio +async def test_pass_through_request_streaming_upstream_error_body_read_failure_keeps_status_and_partial_body(): + """ + Regression: a 502 whose upstream dies while the error preview is being read + must still reach the client with status 502 and the bytes already received; + the read failure must not escape as a ProxyException 500. + """ + upstream_response: Final = httpx.Response( + status_code=502, + headers={"content-type": "application/json"}, + stream=_UpstreamErrorBodyStreamDropping(), + request=httpx.Request("POST", "http://target-api.com/v1beta/models/claude-nope-9:streamGenerateContent"), + ) + + recorded_warnings: list[tuple] = [] + real_warning: Final = verbose_proxy_logger.warning + + def _recording_warning(*args, **kwargs): + if args and str(args[0]).startswith("pass_through_endpoint: upstream"): + recorded_warnings.append(args) + return real_warning(*args, **kwargs) + + with patch.object(verbose_proxy_logger, "warning", side_effect=_recording_warning): + with patch("litellm.proxy.proxy_server.proxy_logging_obj") as mock_proxy_logging: + with patch( + "litellm.proxy.pass_through_endpoints.pass_through_endpoints.get_async_httpx_client" + ) as mock_get_client: + with patch( + "litellm.proxy.pass_through_endpoints.pass_through_endpoints.pass_through_endpoint_logging.pass_through_async_success_handler" + ) as mock_success_handler: + mock_proxy_logging.pre_call_hook = AsyncMock(return_value={}) + mock_proxy_logging.post_call_failure_hook = AsyncMock() + mock_proxy_logging.post_call_response_headers_hook = AsyncMock(return_value=None) + mock_success_handler.return_value = None + + async_client: Final = MagicMock() + async_client.build_request = MagicMock(return_value=MagicMock()) + async_client.send = AsyncMock(return_value=upstream_response) + mock_get_client.return_value = MagicMock(client=async_client) + + response: Final = await pass_through_request( + request=_upstream_error_request(), + target="http://target-api.com/v1beta/models/claude-nope-9:streamGenerateContent", + custom_headers={}, + user_api_key_dict=MagicMock(), + stream=True, + ) + + assert isinstance(response, StreamingResponse) + assert response.status_code == 502 + streamed_chunks: Final = [chunk async for chunk in response.body_iterator] + streamed_bytes: Final = b"".join( + chunk if isinstance(chunk, bytes) else chunk.encode("utf-8") for chunk in streamed_chunks + ) + assert streamed_bytes == b'{"error": "half' + await upstream_response.aclose() + + rendered: Final = [str(args[0]) for args in recorded_warnings] + formats: Final = [args[0] for args in recorded_warnings] + assert any( + fmt == "pass_through_endpoint: upstream %s %s returned %s: %s" and '{"error": "half' in str(args[4]) + for args, fmt in zip(recorded_warnings, formats) + ), rendered + assert any( + fmt == "pass_through_endpoint: upstream error body read failed after %d bytes: %s" + and args[1] == 15 + and args[2] == "ReadError" + for args, fmt in zip(recorded_warnings, formats) + ), rendered + + +class _UpstreamErrorGzipStreamDropping(httpx.AsyncByteStream): + def __init__(self, flushed_prefix: bytes) -> None: + self._flushed_prefix: Final = flushed_prefix + + async def __aiter__(self): + yield self._flushed_prefix + raise httpx.ReadError("peer reset") + + +@pytest.mark.asyncio +async def test_pass_through_request_streaming_upstream_error_gzip_read_failure_relays_decoded_partial(): + """ + Regression: a mid-read failure on a gzip upstream must relay the decoded + plaintext, not the compressed bytes; the relay strips content-encoding so + raw compressed bytes would reach the client as garbage. + """ + plaintext: Final = b'{"error": "half' + compressor: Final = zlib.compressobj(level=6, wbits=31) + flushed_prefix: Final = compressor.compress(plaintext) + compressor.flush(zlib.Z_SYNC_FLUSH) + upstream_response: Final = httpx.Response( + status_code=502, + headers={"content-type": "application/json", "content-encoding": "gzip"}, + stream=_UpstreamErrorGzipStreamDropping(flushed_prefix), + request=httpx.Request("POST", "http://target-api.com/v1beta/models/claude-nope-9:streamGenerateContent"), + ) + + recorded_warnings: list[tuple] = [] + real_warning: Final = verbose_proxy_logger.warning + + def _recording_warning(*args, **kwargs): + if args and str(args[0]).startswith("pass_through_endpoint: upstream"): + recorded_warnings.append(args) + return real_warning(*args, **kwargs) + + with patch.object(verbose_proxy_logger, "warning", side_effect=_recording_warning): + with patch("litellm.proxy.proxy_server.proxy_logging_obj") as mock_proxy_logging: + with patch( + "litellm.proxy.pass_through_endpoints.pass_through_endpoints.get_async_httpx_client" + ) as mock_get_client: + with patch( + "litellm.proxy.pass_through_endpoints.pass_through_endpoints.pass_through_endpoint_logging.pass_through_async_success_handler" + ) as mock_success_handler: + mock_proxy_logging.pre_call_hook = AsyncMock(return_value={}) + mock_proxy_logging.post_call_failure_hook = AsyncMock() + mock_proxy_logging.post_call_response_headers_hook = AsyncMock(return_value=None) + mock_success_handler.return_value = None + + async_client: Final = MagicMock() + async_client.build_request = MagicMock(return_value=MagicMock()) + async_client.send = AsyncMock(return_value=upstream_response) + mock_get_client.return_value = MagicMock(client=async_client) + + response: Final = await pass_through_request( + request=_upstream_error_request(), + target="http://target-api.com/v1beta/models/claude-nope-9:streamGenerateContent", + custom_headers={}, + user_api_key_dict=MagicMock(), + stream=True, + ) + + assert isinstance(response, StreamingResponse) + assert response.status_code == 502 + assert "content-encoding" not in response.headers + streamed_chunks: Final = [chunk async for chunk in response.body_iterator] + streamed_bytes: Final = b"".join( + chunk if isinstance(chunk, bytes) else chunk.encode("utf-8") for chunk in streamed_chunks + ) + assert streamed_bytes == plaintext + await upstream_response.aclose() + + rendered: Final = [str(args[0]) for args in recorded_warnings] + assert any( + args[0] == "pass_through_endpoint: upstream %s %s returned %s: %s" and plaintext.decode() in str(args[4]) + for args in recorded_warnings + ), rendered + + +@pytest.mark.asyncio +async def test_pass_through_request_streaming_upstream_error_gzip_body_decoded_for_log_and_client(): + upstream_content: Final = b'{"error": {"message": "gzipped upstream says the project was not found"}}' + compressed: Final = gzip.compress(upstream_content) + upstream_response: Final = httpx.Response( + status_code=502, + headers={"content-type": "text/event-stream", "content-encoding": "gzip"}, + stream=_ChunkedUpstreamErrorBodyStream((compressed[:10], compressed[10:])), + request=httpx.Request("POST", "http://target-api.com/v1beta/models/claude-nope-9:streamGenerateContent"), + ) + + with patch.object(verbose_proxy_logger, "warning") as mock_warning: + with patch("litellm.proxy.proxy_server.proxy_logging_obj") as mock_proxy_logging: + with patch( + "litellm.proxy.pass_through_endpoints.pass_through_endpoints.get_async_httpx_client" + ) as mock_get_client: + with patch( + "litellm.proxy.pass_through_endpoints.pass_through_endpoints.pass_through_endpoint_logging.pass_through_async_success_handler" + ) as mock_success_handler: + mock_proxy_logging.pre_call_hook = AsyncMock(return_value={}) + mock_proxy_logging.post_call_failure_hook = AsyncMock() + mock_proxy_logging.post_call_response_headers_hook = AsyncMock(return_value=None) + mock_success_handler.return_value = None + + async_client: Final = MagicMock() + async_client.build_request = MagicMock(return_value=MagicMock()) + async_client.send = AsyncMock(return_value=upstream_response) + mock_get_client.return_value = MagicMock(client=async_client) + + response: Final = await pass_through_request( + request=_upstream_error_request(), + target="http://target-api.com/v1beta/models/claude-nope-9:streamGenerateContent", + custom_headers={}, + user_api_key_dict=MagicMock(), + stream=True, + ) + + assert isinstance(response, StreamingResponse) + streamed_chunks: Final = [chunk async for chunk in response.body_iterator] + streamed_bytes: Final = b"".join( + chunk if isinstance(chunk, bytes) else chunk.encode("utf-8") for chunk in streamed_chunks + ) + assert streamed_bytes == upstream_content + + upstream_warnings: Final = [ + call + for call in mock_warning.call_args_list + if call.args[0] == "pass_through_endpoint: upstream %s %s returned %s: %s" + ] + assert len(upstream_warnings) == 1, mock_warning.call_args_list + logged_body: Final = str(upstream_warnings[0].args[4]) + assert "gzipped upstream says the project was not found" in logged_body + + +@pytest.mark.asyncio +async def test_pass_through_request_upstream_error_body_sanitized_against_log_forging(): + upstream_content: Final = b'{"error": "line one"}\n2026-01-01 FAKE LOG LINE\x1b[31m' + upstream_response: Final = httpx.Response( + status_code=404, + headers={"content-type": "application/json"}, + content=upstream_content, + request=httpx.Request("POST", "http://target-api.com/v1beta/models/claude-nope-9:generateContent"), + ) + + with patch.object(verbose_proxy_logger, "warning") as mock_warning: + with patch("litellm.proxy.proxy_server.proxy_logging_obj") as mock_proxy_logging: + with patch( + "litellm.proxy.pass_through_endpoints.pass_through_endpoints.get_async_httpx_client" + ) as mock_get_client: + with patch( + "litellm.proxy.pass_through_endpoints.pass_through_endpoints.ProxyBaseLLMRequestProcessing" + ) as mock_processing: + mock_proxy_logging.pre_call_hook = AsyncMock(return_value={}) + mock_proxy_logging.post_call_failure_hook = AsyncMock() + mock_proxy_logging.post_call_response_headers_hook = AsyncMock(return_value=None) + mock_processing.get_custom_headers.return_value = {} + + async_client: Final = MagicMock() + async_client.build_request = MagicMock(return_value=MagicMock()) + async_client.send = AsyncMock(return_value=upstream_response) + mock_get_client.return_value = MagicMock(client=async_client) + + await pass_through_request( + request=_upstream_error_request(), + target="http://target-api.com/v1beta/models/claude-nope-9:generateContent", + custom_headers={}, + user_api_key_dict=MagicMock(), + ) + + upstream_warnings: Final = [ + call + for call in mock_warning.call_args_list + if call.args[0] == "pass_through_endpoint: upstream %s %s returned %s: %s" + ] + assert len(upstream_warnings) == 1, mock_warning.call_args_list + logged_body: Final = str(upstream_warnings[0].args[4]) + assert logged_body == '{"error": "line one"} 2026-01-01 FAKE LOG LINE [31m' + assert "\n" not in logged_body + assert "\x1b" not in logged_body + + detail: Final = mock_proxy_logging.post_call_failure_hook.call_args.kwargs["original_exception"].detail + assert ( + detail + == 'Upstream passthrough request failed with status 404: {"error": "line one"} 2026-01-01 FAKE LOG LINE [31m' + ) + + class _UpstreamDroppingMidStream(httpx.AsyncByteStream): async def __aiter__(self): yield b'data: {"id": "chatcmpl-1", "choices": [{"delta": {"content": "hi"}}]}\n\n' @@ -4287,7 +4991,9 @@ async def test_pass_through_request_claims_the_budget_reservation_only_when_its_ mock_proxy_logging.post_call_failure_hook = AsyncMock() mock_proxy_logging.post_call_response_headers_hook = AsyncMock(return_value=None) mock_processing.get_custom_headers.return_value = {} - mock_worker.ensure_initialized_and_enqueue = MagicMock(side_effect=lambda async_coroutine: async_coroutine.close()) + mock_worker.ensure_initialized_and_enqueue = MagicMock( + side_effect=lambda async_coroutine: async_coroutine.close() + ) async_client = MagicMock() async_client.build_request = MagicMock(return_value=MagicMock()) async_client.send = AsyncMock(return_value=upstream_response) @@ -5304,9 +6010,7 @@ async def test_websocket_passthrough_propagates_active_trace_context( mock_proxy_logging.post_call_success_hook = AsyncMock() mock_proxy_logging.post_call_failure_hook = AsyncMock() mock_worker = MagicMock() - mock_worker.ensure_initialized_and_enqueue = MagicMock( - side_effect=lambda async_coroutine: async_coroutine.close() - ) + mock_worker.ensure_initialized_and_enqueue = MagicMock(side_effect=lambda async_coroutine: async_coroutine.close()) monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", mock_proxy_logging) monkeypatch.setattr( "litellm.proxy.pass_through_endpoints.pass_through_endpoints.connect", From 3fb6f8740b109d638c55c19436cc9440595642f4 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 24 Sep 2026 00:25:03 -0700 Subject: [PATCH 025/120] test(integration): add read-replica routing harness to the CircleCI integration suite (#42692) * test(integration): add read-replica routing harness * refactor(integration): hoist the maintenance url imports * fix(integration): keep per-test databases and the witness sequence readable under replica roles * fix(integration): opt bespoke database and pool tests out of the injected read replica * test(integration): commit recorded replica routing expectations * fix(integration): judge routing by role containment so shrinking role sets do not fail * fix(integration): run the pool-limit shutdown choreography on the superuser database url * ci(integration): add the mcp group to the replica matrix * fix(integration): judge routing by exact role sets with a named either-role allowlist * test(integration): drop containment-era routing expectations for re-recording * chore(integration): drop docstrings from the replica harness scripts * docs(integration): describe exact routing matching and the either-role list * test(integration): record exact replica routing expectations * test(integration): allow the SELECT 1 health probe on either role * test(integration): replace committed routing expectations with an on-demand base-vs-head parity run * test(integration): fix parity env scope, readme wording, and seed-deterministic serialization test * test(integration): make the sorted-role serialization test deterministic in-process * test(integration): swap all product code in parity runs and pin role gains * ci(integration): force tracked-file removal before parity checkout --------- Co-authored-by: yuneng --- .circleci/config.yml | 124 ++++- .circleci/scripts/prepare_replica_roles.py | 52 +++ .circleci/scripts/run_integration.sh | 31 +- tests/integration/README.md | 2 + tests/integration/_support/process.py | 18 +- tests/integration/_support/routing.py | 333 +++++++++++++ tests/integration/conftest.py | 3 + .../database/test_transaction_atomicity.py | 1 + ..._user_updates_wedged_coordination_redis.py | 1 + .../test_vector_store_config_ownership.py | 9 +- tests/integration/routing/either_role.json | 3 + .../routing/test_redis_recovery.py | 2 +- .../spend/test_daily_rollup_retry.py | 1 + .../integration/spend/test_shutdown_flush.py | 2 + tests/unit/integration_support/__init__.py | 0 .../unit/integration_support/test_routing.py | 438 ++++++++++++++++++ 16 files changed, 1008 insertions(+), 12 deletions(-) create mode 100644 .circleci/scripts/prepare_replica_roles.py create mode 100644 tests/integration/_support/routing.py create mode 100644 tests/integration/routing/either_role.json create mode 100644 tests/unit/integration_support/__init__.py create mode 100644 tests/unit/integration_support/test_routing.py diff --git a/.circleci/config.yml b/.circleci/config.yml index eb76244c1ab..370424dca86 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -12,6 +12,9 @@ parameters: migration_source_sha: type: string default: "" + routing_parity_base: + type: string + default: "" orbs: codecov: codecov/codecov@4.0.1 node: circleci/node@5.1.0 # Add this line to declare the node orb @@ -176,6 +179,9 @@ commands: image: type: string default: postgres:14@sha256:6a70deda415ec296f977890e11aba04a0db9f632a362e3fce45e845e3db74f26 + server_args: + type: string + default: "" steps: - run: name: Start PostgreSQL @@ -186,7 +192,7 @@ commands: -e POSTGRES_PASSWORD=postgres \ -e POSTGRES_DB=<< parameters.db_name >> \ -p 5432:5432 \ - << parameters.image >> + << parameters.image >> << parameters.server_args >> - wait_for_service: url: tcp://localhost:5432 timeout: "60" @@ -3108,6 +3114,10 @@ jobs: parameters: suite: type: string + mode: + type: enum + enum: [standard, replica] + default: standard machine: image: ubuntu-2204:2024.04.1 resource_class: large @@ -3142,18 +3152,19 @@ jobs: command: cd ui/litellm-dashboard && NEXT_TELEMETRY_DISABLED=1 npm run build - start_postgres: image: postgres:16@sha256:e17e86066e5ef83e0952a9347f5c792b7ece00972e2aa787a6986f471b3dd3d5 + server_args: "-c shared_preload_libraries=pg_stat_statements -c pg_stat_statements.track=all -c pg_stat_statements.max=20000" - start_redis - run: name: Run owned integration contracts - command: bash .circleci/scripts/run_integration.sh << parameters.suite >> + command: bash .circleci/scripts/run_integration.sh << parameters.suite >> << parameters.mode >> no_output_timeout: 15m - run: name: Stop owned database and Redis when: always command: | - mkdir -p test-results/integration-<< parameters.suite >> - docker logs postgres-db > test-results/integration-<< parameters.suite >>/postgres.log 2>&1 || true - docker logs redis-cache > test-results/integration-<< parameters.suite >>/redis.log 2>&1 || true + mkdir -p test-results/services-<< parameters.suite >>-<< parameters.mode >> + docker logs postgres-db > test-results/services-<< parameters.suite >>-<< parameters.mode >>/postgres.log 2>&1 || true + docker logs redis-cache > test-results/services-<< parameters.suite >>-<< parameters.mode >>/redis.log 2>&1 || true docker rm -f postgres-db redis-cache test -z "$(docker ps -aq --filter name=postgres-db --filter name=redis-cache)" - store_test_results: @@ -3161,6 +3172,76 @@ jobs: - store_artifacts: path: test-results + routing_parity: + parameters: + suite: + type: string + machine: + image: ubuntu-2204:2024.04.1 + resource_class: large + working_directory: ~/project + steps: + - setup_litellm_test_deps + - run: + name: Check out base product code + environment: + ROUTING_PARITY_BASE: << pipeline.parameters.routing_parity_base >> + command: | + [[ "$ROUTING_PARITY_BASE" =~ ^[0-9a-f]{40}$ ]] || exit 1 + git fetch --depth 1 origin "$ROUTING_PARITY_BASE" + git rm -r -f --quiet litellm enterprise litellm-proxy-extras + git checkout "$ROUTING_PARITY_BASE" -- litellm enterprise litellm-proxy-extras + git reset --quiet + test -f litellm/rust_bridge/_native.abi3.so + - start_postgres: + image: postgres:16@sha256:e17e86066e5ef83e0952a9347f5c792b7ece00972e2aa787a6986f471b3dd3d5 + server_args: "-c shared_preload_libraries=pg_stat_statements -c pg_stat_statements.track=all -c pg_stat_statements.max=20000" + - start_redis + - run: + name: Run base side + command: bash .circleci/scripts/run_integration.sh << parameters.suite >> parity base + no_output_timeout: 15m + - run: + name: Stop base database and Redis + when: always + command: | + mkdir -p test-results/services-<< parameters.suite >>-parity-base + docker logs postgres-db > test-results/services-<< parameters.suite >>-parity-base/postgres.log 2>&1 || true + docker logs redis-cache > test-results/services-<< parameters.suite >>-parity-base/redis.log 2>&1 || true + docker rm -f postgres-db redis-cache + test -z "$(docker ps -aq --filter name=postgres-db --filter name=redis-cache)" + - run: + name: Check out head product code + command: | + git rm -r -f --quiet litellm enterprise litellm-proxy-extras + git checkout "$CIRCLE_SHA1" -- litellm enterprise litellm-proxy-extras + git reset --quiet + test -f litellm/rust_bridge/_native.abi3.so + - start_postgres: + image: postgres:16@sha256:e17e86066e5ef83e0952a9347f5c792b7ece00972e2aa787a6986f471b3dd3d5 + server_args: "-c shared_preload_libraries=pg_stat_statements -c pg_stat_statements.track=all -c pg_stat_statements.max=20000" + - start_redis + - run: + name: Run head side + command: bash .circleci/scripts/run_integration.sh << parameters.suite >> parity head + no_output_timeout: 15m + - run: + name: Stop head database and Redis + when: always + command: | + mkdir -p test-results/services-<< parameters.suite >>-parity-head + docker logs postgres-db > test-results/services-<< parameters.suite >>-parity-head/postgres.log 2>&1 || true + docker logs redis-cache > test-results/services-<< parameters.suite >>-parity-head/redis.log 2>&1 || true + docker rm -f postgres-db redis-cache + test -z "$(docker ps -aq --filter name=postgres-db --filter name=redis-cache)" + - run: + name: Compare routing parity + command: PYTHONPATH="$PWD/tests" .venv/bin/python -m integration._support.routing check test-results/parity-<< parameters.suite >>/base test-results/parity-<< parameters.suite >>/head + - store_test_results: + path: test-results + - store_artifacts: + path: test-results + unit: machine: image: ubuntu-2204:2024.04.1 @@ -3224,8 +3305,22 @@ workflows: branches: only: main jobs: *migration_jobs + routing_parity: + when: + not: + equal: ["", << pipeline.parameters.routing_parity_base >>] + jobs: + - routing_parity: + name: routing-parity-<< matrix.suite >> + matrix: + parameters: + suite: [management, accounting, database, providers, extensions, cost, mcp] integration: - unless: << pipeline.parameters.run_migration_tests >> + unless: + or: + - << pipeline.parameters.run_migration_tests >> + - not: + equal: ["", << pipeline.parameters.routing_parity_base >>] jobs: - integration_contracts: name: integration-<< matrix.suite >> @@ -3237,8 +3332,23 @@ workflows: only: - main - /litellm_.*/ + - integration_contracts: + name: integration-<< matrix.suite >>-replica + matrix: + parameters: + suite: [management, database] + mode: [replica] + filters: + branches: + only: + - main + - /litellm_.*/ build_and_test: - unless: << pipeline.parameters.run_migration_tests >> + unless: + or: + - << pipeline.parameters.run_migration_tests >> + - not: + equal: ["", << pipeline.parameters.routing_parity_base >>] jobs: - using_litellm_on_windows: filters: &main_branches diff --git a/.circleci/scripts/prepare_replica_roles.py b/.circleci/scripts/prepare_replica_roles.py new file mode 100644 index 00000000000..fb4b7fcae97 --- /dev/null +++ b/.circleci/scripts/prepare_replica_roles.py @@ -0,0 +1,52 @@ +from __future__ import annotations + +import os +from typing import Final +from urllib.parse import urlsplit, urlunsplit + +import psycopg + +DATABASE_URL: Final = os.environ["DATABASE_URL"] + + +def postgres_url() -> str: + parsed: Final = urlsplit(DATABASE_URL) + return urlunsplit(parsed._replace(path="/postgres")) + + +def main() -> None: + with psycopg.connect(postgres_url(), autocommit=True) as admin: + admin.execute("CREATE EXTENSION IF NOT EXISTS pg_stat_statements") + admin.execute("CREATE ROLE litellm_writer LOGIN PASSWORD 'litellm-writer' NOSUPERUSER") + admin.execute("CREATE ROLE litellm_reader LOGIN PASSWORD 'litellm-reader' NOSUPERUSER NOINHERIT") + admin.execute("ALTER ROLE litellm_reader SET default_transaction_read_only = on") + admin.execute("ALTER DATABASE circle_test OWNER TO litellm_writer") + admin.execute("GRANT CONNECT ON DATABASE circle_test TO litellm_reader") + with psycopg.connect(DATABASE_URL, autocommit=True) as admin: + admin.execute("GRANT USAGE ON SCHEMA public TO litellm_reader") + admin.execute( + "ALTER DEFAULT PRIVILEGES FOR ROLE litellm_writer IN SCHEMA public GRANT SELECT ON TABLES TO litellm_reader" + ) + admin.execute("GRANT SELECT ON ALL TABLES IN SCHEMA public TO litellm_reader") + + parsed: Final = urlsplit(DATABASE_URL) + reader_url: Final = urlunsplit( + parsed._replace(netloc=f"litellm_reader:litellm-reader@{parsed.hostname}:{parsed.port}") + ) + writer_url: Final = urlunsplit( + parsed._replace(netloc=f"litellm_writer:litellm-writer@{parsed.hostname}:{parsed.port}") + ) + with psycopg.connect(reader_url, autocommit=True) as reader: + assert reader.execute("SHOW transaction_read_only").fetchone() == ("on",) + try: + reader.execute("CREATE TABLE integration_readonly_probe (id int)") + except psycopg.errors.ReadOnlySqlTransaction: + pass + else: + raise AssertionError("litellm_reader executed a write statement") + with psycopg.connect(writer_url, autocommit=True) as writer: + assert writer.execute("SELECT current_user").fetchone() == ("litellm_writer",) + + +if __name__ == "__main__": + main() diff --git a/.circleci/scripts/run_integration.sh b/.circleci/scripts/run_integration.sh index d16ac9cd124..b617a79946c 100644 --- a/.circleci/scripts/run_integration.sh +++ b/.circleci/scripts/run_integration.sh @@ -7,7 +7,15 @@ if [ "${GITHUB_ACTIONS:-}" = true ]; then fi suite="${1:?integration suite required}" -results="test-results/integration-${suite}" +mode="${2:-standard}" +side="${3:-}" +if [ "$mode" = replica ]; then + results="test-results/integration-${suite}-replica" +elif [ "$mode" = parity ]; then + results="test-results/parity-${suite}/${side:?parity side required}" +else + results="test-results/integration-${suite}" +fi mkdir -p "$results" integration_identity="$(.venv/bin/python -c 'import uuid; print(uuid.uuid4().hex)')" upstream_pid="" @@ -80,6 +88,18 @@ export INTEGRATION_ORDER_SEED="$INTEGRATION_SEED" uv run --no-sync prisma generate --schema litellm/proxy/schema.prisma > "$results/prisma-generate.log" 2>&1 +export INTEGRATION_PROXY_DATABASE_URL="" +export INTEGRATION_PROXY_READ_REPLICA_URL="" +export INTEGRATION_ROUTING="" +if [ "$mode" = replica ] || [ "$mode" = parity ]; then + .venv/bin/python .circleci/scripts/prepare_replica_roles.py > "$results/prepare-replica-roles.log" 2>&1 + export INTEGRATION_PROXY_DATABASE_URL="postgresql://litellm_writer:litellm-writer@127.0.0.1:5432/circle_test" + export INTEGRATION_PROXY_READ_REPLICA_URL="postgresql://litellm_reader:litellm-reader@127.0.0.1:5432/circle_test" +fi +if [ "$mode" = parity ]; then + export INTEGRATION_ROUTING=capture +fi + sudo iptables -N integration_only guard_created=true sudo iptables -A integration_only -o lo -j ACCEPT @@ -137,8 +157,12 @@ start_proxy() { else cost_map_env=("LITELLM_LOCAL_MODEL_COST_MAP=True") fi + local -a database_env=("DATABASE_URL=${INTEGRATION_PROXY_DATABASE_URL:-$DATABASE_URL}") + if [ -n "$INTEGRATION_PROXY_READ_REPLICA_URL" ]; then + database_env+=("DATABASE_URL_READ_REPLICA=$INTEGRATION_PROXY_READ_REPLICA_URL") + fi setsid env -i PATH="$PATH" HOME="$HOME" PYTHONPATH="$PYTHONPATH" INTEGRATION_RUN_ID="$integration_identity" \ - DATABASE_URL="$DATABASE_URL" REDIS_HOST="$REDIS_HOST" REDIS_PORT="$REDIS_PORT" \ + "${database_env[@]}" REDIS_HOST="$REDIS_HOST" REDIS_PORT="$REDIS_PORT" \ INTEGRATION_UPSTREAM_URL="$INTEGRATION_UPSTREAM_URL" \ LITELLM_MASTER_KEY="$LITELLM_MASTER_KEY" LITELLM_SALT_KEY="$LITELLM_SALT_KEY" LITELLM_UI_PATH="$LITELLM_UI_PATH" PROXY_BASE_URL="http://127.0.0.1:$port" \ LITELLM_MODE=PRODUCTION STORE_MODEL_IN_DB=True "${cost_map_env[@]}" \ @@ -195,6 +219,9 @@ env -i PATH="$PATH" HOME="$HOME" PYTHONPATH="$PYTHONPATH" \ INTEGRATION_SEED="$INTEGRATION_SEED" \ INTEGRATION_ORDER_SEED="$INTEGRATION_ORDER_SEED" \ LITELLM_LOCAL_MODEL_COST_MAP=True AWS_EC2_METADATA_DISABLED=true DO_NOT_TRACK=1 \ + INTEGRATION_PROXY_DATABASE_URL="$INTEGRATION_PROXY_DATABASE_URL" \ + INTEGRATION_PROXY_READ_REPLICA_URL="$INTEGRATION_PROXY_READ_REPLICA_URL" \ + INTEGRATION_ROUTING="$INTEGRATION_ROUTING" \ .venv/bin/python tests/integration/run.py "$suite" --results "$results" if [ "${INTEGRATION_COVERAGE:-0}" = 1 ]; then diff --git a/tests/integration/README.md b/tests/integration/README.md index f7c1305ad2e..ac9b01786b9 100644 --- a/tests/integration/README.md +++ b/tests/integration/README.md @@ -35,3 +35,5 @@ The extensions shard uses the built-in generic callback and guardrail transports The mcp shard runs the MCP gateway against SDK peers owned by each test (`_support/mcp.py`): streamable HTTP, SSE and stdio peers, an OpenAPI-spec app, and an OAuth 2.1 authorization-server double. Every peer records the requests it receives so a test can assert what reached the peer, not only what the proxy answered. The shard runs with `INTEGRATION_WORKERS` set and with `INTEGRATION_COVERAGE=1`, which starts the proxy under `coverage run --parallel-mode` limited to the MCP modules and stores `coverage.txt` plus an HTML report with the job artifacts. A test that fails because the product is wrong is skipped with `pytest.skip("BUG: ")` so the skip list in `execution.json` is the open MCP bug list Browser contracts live in `tests/e2e/ui/tests/integrationCritical` and run only through `tests/e2e/ui/integration.config.ts`. The expected browser results are listed in `expected.json` in that directory and checked by `.circleci/scripts/verify_integration_browser.py`. The CircleCI browser shard builds the checked-out dashboard, starts the owned proxy with that build, and verifies one exact browser result without retries or skips. The default Playwright selection excludes this directory. The focused project flow asserts the submitted create and clear values, fresh SQL state and actual blocked/restored serving while preserving model restrictions + +Two always-on `-replica` CircleCI jobs (management, database) run their groups in replica mode, where every proxy connects through a real `litellm_writer` role and a real read-only `litellm_reader` role against the same PostgreSQL. Nothing is captured there: the job passes when the tests pass, and a write routed to the read-only reader fails the test that issued it. A deeper check runs on demand as the `routing_parity` workflow, triggered through the CircleCI API v2 pipeline endpoint on the PR branch with `{"parameters": {"routing_parity_base": "<40-hex merge-base sha>"}}`. The workflow fans out over the seven groups, and each `routing-parity-` job runs its own group twice against the same test harness, once with `litellm/`, `enterprise/`, and `litellm-proxy-extras/` checked out from the base revision and once from the head, with a pytest plugin snapshotting `pg_stat_statements` into `routing-observed.json` per side. The `check` step then compares the two observations and writes `routing-diff.txt`: a statement seen on both sides fails when its role set changed, globally or for the same test (per-test capture is skipped under xdist), unless it is listed in `tests/integration/routing/either_role.json`, where each entry names the statement and a one-line reason it legitimately runs on whichever role asks for it, printed under `== either role ==`. Queries seen on only one side are listed, never failed, `pg_stat_statements` evictions and a role that never ran a statement are failures diff --git a/tests/integration/_support/process.py b/tests/integration/_support/process.py index 5c44beaa570..fcbaf7c8d8c 100644 --- a/tests/integration/_support/process.py +++ b/tests/integration/_support/process.py @@ -9,6 +9,7 @@ from collections.abc import Iterator, Mapping from contextlib import contextmanager from dataclasses import dataclass from pathlib import Path +from types import MappingProxyType from typing import Final import httpx @@ -16,6 +17,17 @@ import psutil from integration._support.client import Gateway +def proxy_database_environment() -> Mapping[str, str]: + writer: Final = os.environ.get("INTEGRATION_PROXY_DATABASE_URL", "") + reader: Final = os.environ.get("INTEGRATION_PROXY_READ_REPLICA_URL", "") + return MappingProxyType( + { + **({"DATABASE_URL": writer} if writer else {}), + **({"DATABASE_URL_READ_REPLICA": reader} if reader else {}), + } + ) + + def in_group(process: psutil.Process, group: int) -> bool: try: return os.getpgid(process.pid) == group @@ -83,7 +95,11 @@ def owned_proxy_process( port: Final = reserve.getsockname()[1] root: Final = Path(os.environ.get("INTEGRATION_PROXY_ROOT") or Path(__file__).resolve().parents[3]) environment: Final = { - **{name: value for name, value in os.environ.items() if name not in remove_environment}, + **{ + name: value + for name, value in {**os.environ, **proxy_database_environment()}.items() + if name not in remove_environment + }, "LITELLM_MASTER_KEY": gateway.key, "LITELLM_SALT_KEY": os.environ.get("LITELLM_SALT_KEY", "sk-integration-salt"), "STORE_MODEL_IN_DB": "True", diff --git a/tests/integration/_support/routing.py b/tests/integration/_support/routing.py new file mode 100644 index 00000000000..14f4741367a --- /dev/null +++ b/tests/integration/_support/routing.py @@ -0,0 +1,333 @@ +from __future__ import annotations + +import argparse +import itertools +import json +import os +import re +import sys +from collections.abc import Iterator, Mapping +from dataclasses import dataclass +from pathlib import Path +from types import MappingProxyType +from typing import Final +from urllib.parse import urlsplit, urlunsplit + +import psycopg +import pytest +from pydantic import TypeAdapter + +WRITER_ROLE: Final = "litellm_writer" +READER_ROLE: Final = "litellm_reader" +ROLES: Final = (READER_ROLE, WRITER_ROLE) +DATABASE_NAME: Final = "circle_test" +OBSERVED_FILE: Final = "routing-observed.json" +DIFF_FILE: Final = "routing-diff.txt" +EITHER_ROLE_FILE: Final = Path(__file__).resolve().parents[1] / "routing" / "either_role.json" + +RoleSet = frozenset[str] +RoutingMap = Mapping[str, frozenset[str]] +Snapshot = Mapping[tuple[str, str], int] + +_PLACEHOLDERS: Final = re.compile(r"\$\d+(?:\s*,\s*\$\d+)*") +_QUERIES: Final = TypeAdapter(dict[str, tuple[str, ...]]) +_OBSERVED: Final = TypeAdapter(dict[str, object]) + + +def normalize(query: str) -> str: + return _PLACEHOLDERS.sub("$n", " ".join(query.split())) + + +@dataclass(frozen=True, slots=True) +class Observation: + queries: RoutingMap + tests: Mapping[str, RoutingMap] + calls: Mapping[str, int] + dealloc: int + + +@dataclass(frozen=True, slots=True) +class Mismatch: + test: str | None + query: str + base: tuple[str, ...] + head: tuple[str, ...] + + +@dataclass(frozen=True, slots=True) +class Report: + mismatches: tuple[Mismatch, ...] + only_base: tuple[str, ...] + only_head: tuple[str, ...] + calls: Mapping[str, Mapping[str, int]] + dealloc: Mapping[str, int] + either_role: tuple[str, ...] = () + + def failures(self) -> tuple[str, ...]: + mismatch_failures: Final = tuple( + f"{mismatch.test if mismatch.test is not None else 'global'}: {mismatch.query}: " + f"base [{', '.join(mismatch.base)}] head [{', '.join(mismatch.head)}]" + for mismatch in self.mismatches + ) + side_failures: Final = tuple( + failure + for side in ("base", "head") + for failure in ( + *( + (f"{side}: pg_stat_statements evicted {self.dealloc[side]} entries (dealloc > 0)",) + if self.dealloc[side] > 0 + else () + ), + *(f"{side}: no {role} calls observed" for role in ROLES if self.calls[side].get(role, 0) == 0), + ) + ) + return (*mismatch_failures, *side_failures) + + +def _sorted_map(value: RoutingMap) -> RoutingMap: + return MappingProxyType(dict(sorted(value.items()))) + + +def compare(base: Observation, head: Observation, either_role: frozenset[str] = frozenset()) -> Report: + mismatches: Final = ( + *( + Mismatch( + None, + query, + tuple(sorted(base_roles)), + tuple(sorted(head.queries[query])), + ) + for query, base_roles in base.queries.items() + if query in head.queries and head.queries[query] != base_roles and query not in either_role + ), + *( + Mismatch( + test, + query, + tuple(sorted(base_roles)), + tuple(sorted(head.tests[test][query])), + ) + for test, queries in base.tests.items() + if test in head.tests + for query, base_roles in queries.items() + if query in head.tests[test] and head.tests[test][query] != base_roles and query not in either_role + ), + ) + varying: Final = frozenset( + query + for query in either_role + if (query in base.queries and query in head.queries and head.queries[query] != base.queries[query]) + or any( + query in base.tests[test] + and query in head.tests[test] + and head.tests[test][query] != base.tests[test][query] + for test in frozenset(base.tests) & frozenset(head.tests) + ) + ) + return Report( + mismatches, + tuple(sorted(query for query in base.queries if query not in head.queries)), + tuple(sorted(query for query in head.queries if query not in base.queries)), + MappingProxyType({"base": base.calls, "head": head.calls}), + MappingProxyType({"base": base.dealloc, "head": head.dealloc}), + tuple(sorted(varying)), + ) + + +def render(report: Report) -> str: + failures: Final = report.failures() + lines: Final = ( + "== failures ==", + *(failures or ("none",)), + "", + "== either role ==", + *(report.either_role or ("none",)), + "", + "== queries only in base ==", + *(report.only_base or ("none",)), + "", + "== queries only in head ==", + *(report.only_head or ("none",)), + "", + "== calls ==", + *( + line + for side in ("base", "head") + for line in ( + *(f"{side} {role}: {report.calls[side].get(role, 0)}" for role in ROLES), + f"{side} dealloc: {report.dealloc[side]}", + ) + ), + ) + return "\n".join(lines) + "\n" + + +def _roles(document: Mapping[str, tuple[str, ...]]) -> RoutingMap: + return _sorted_map({query: frozenset(roles) for query, roles in document.items()}) + + +def _tests(document: Mapping[str, Mapping[str, tuple[str, ...]]]) -> Mapping[str, RoutingMap]: + return MappingProxyType({node: _roles(queries) for node, queries in document.items()}) + + +def load_observation(path: Path) -> Observation: + document: Final = _OBSERVED.validate_python(json.loads(path.read_text())) + queries: Final = _QUERIES.validate_python(document.get("queries", {})) + tests: Final = TypeAdapter(dict[str, dict[str, tuple[str, ...]]]).validate_python(document.get("tests", {})) + calls: Final = TypeAdapter(dict[str, int]).validate_python(document.get("calls", {})) + dealloc: Final = TypeAdapter(int).validate_python(document.get("dealloc", 0)) + return Observation(_roles(queries), _tests(tests), MappingProxyType(calls), dealloc) + + +def load_either_role(path: Path) -> frozenset[str]: + if not path.exists(): + return frozenset() + document: Final = TypeAdapter(dict[str, str]).validate_python(json.loads(path.read_text())) + return frozenset(document) + + +def _serializable(queries: RoutingMap, tests: Mapping[str, RoutingMap]) -> dict[str, object]: + return { + "queries": {query: sorted(roles) for query, roles in queries.items()}, + "tests": {node: {query: sorted(roles) for query, roles in mapping.items()} for node, mapping in tests.items()}, + } + + +def dump_observation(observation: Observation) -> str: + document: Final = _serializable(observation.queries, observation.tests) + return ( + json.dumps( + {**document, "calls": dict(observation.calls), "dealloc": observation.dealloc}, + indent=2, + sort_keys=True, + ) + + "\n" + ) + + +def _maintenance_url() -> str: + parsed: Final = urlsplit(os.environ["DATABASE_URL"]) + return urlunsplit(parsed._replace(path="/postgres")) + + +def snapshot(connection: psycopg.Connection[object]) -> Mapping[tuple[str, str], int]: + rows: Final = connection.execute( + """ + SELECT r.rolname, s.query, s.calls + FROM pg_stat_statements s + JOIN pg_roles r ON r.oid = s.userid + WHERE s.dbid = (SELECT oid FROM pg_database WHERE datname = %s) + AND r.rolname = ANY(%s) + """, + (DATABASE_NAME, list(ROLES)), + ).fetchall() + return MappingProxyType( + { + key: sum(calls for _, _, calls in grouped) + for key, grouped in itertools.groupby( + sorted((str(role), normalize(str(query)), int(calls)) for role, query, calls in rows), + key=lambda row: (row[0], row[1]), + ) + } + ) + + +def delta(before: Snapshot, after: Snapshot) -> RoutingMap: + pairs: Final = {key: after.get(key, 0) - before.get(key, 0) for key in frozenset(before) | frozenset(after)} + queries: Final = frozenset(query for (_, query), change in pairs.items() if change > 0) + return MappingProxyType( + {query: frozenset(role for role in ROLES if pairs.get((role, query), 0) > 0) for query in sorted(queries)} + ) + + +def role_calls(before: Snapshot, after: Snapshot) -> Mapping[str, int]: + return MappingProxyType( + { + role: sum( + max(after.get((role, query), 0) - before.get((role, query), 0), 0) + for query in frozenset(q for _, q in before) | frozenset(q for _, q in after) + ) + for role in ROLES + } + ) + + +def dealloc(connection: psycopg.Connection[object]) -> int: + return int(connection.execute("SELECT dealloc FROM pg_stat_statements_info").fetchone()[0]) + + +class RoutingPlugin: + def __init__(self, config: pytest.Config) -> None: + self.config = config + self._session_start: Snapshot | None = None + self._tests: tuple[tuple[str, RoutingMap], ...] = () + + def _snapshot(self) -> Snapshot: + with psycopg.connect(_maintenance_url(), autocommit=True) as connection: + return snapshot(connection) + + def pytest_sessionstart(self, session: pytest.Session) -> None: + if hasattr(self.config, "workerinput"): + return + self._session_start = self._snapshot() + + @pytest.hookimpl(hookwrapper=True) + def pytest_runtest_protocol(self, item: pytest.Item, nextitem: pytest.Item | None) -> Iterator[None]: + if self.config.getoption("numprocesses", default=None) or hasattr(self.config, "workerinput"): + yield + return + before: Final = self._snapshot() + yield + after: Final = self._snapshot() + self._tests = (*self._tests, (item.nodeid, delta(before, after))) + + def pytest_sessionfinish(self, session: pytest.Session, exitstatus: int) -> None: + if hasattr(self.config, "workerinput"): + return + end: Final = self._snapshot() + with psycopg.connect(_maintenance_url(), autocommit=True) as connection: + evictions: Final = dealloc(connection) + start: Final = self._session_start or {} + tests: Final = MappingProxyType({node: mapping for node, mapping in self._tests}) + destination: Final = Path(os.environ["INTEGRATION_RESULTS_DIR"]) + destination.mkdir(parents=True, exist_ok=True) + (destination / OBSERVED_FILE).write_text( + dump_observation( + Observation( + delta(start, end), + tests, + role_calls(start, end), + evictions, + ) + ) + ) + + +def main(argv: tuple[str, ...] | list[str]) -> int: + parser: Final = argparse.ArgumentParser() + parser.add_argument("command", choices=("check",)) + parser.add_argument("base_dir", type=Path) + parser.add_argument("head_dir", type=Path) + parser.add_argument("--either-role", type=Path, default=EITHER_ROLE_FILE) + parser.add_argument("--diff", type=Path, default=None) + options: Final = parser.parse_args(argv) + base_path: Final = options.base_dir / OBSERVED_FILE + head_path: Final = options.head_dir / OBSERVED_FILE + for path in (base_path, head_path): + if not path.exists(): + sys.stderr.write(f"observed routing file missing: {path}\n") + if not base_path.exists() or not head_path.exists(): + return 1 + report: Final = compare( + load_observation(base_path), + load_observation(head_path), + load_either_role(options.either_role), + ) + diff: Final = render(report) + (options.diff or options.head_dir.parent / DIFF_FILE).write_text(diff) + sys.stdout.write(diff) + return 1 if report.failures() else 0 + + +if __name__ == "__main__": + raise SystemExit(main(sys.argv[1:])) diff --git a/tests/integration/conftest.py b/tests/integration/conftest.py index 4986f5ddcc0..368b3ebee75 100644 --- a/tests/integration/conftest.py +++ b/tests/integration/conftest.py @@ -15,6 +15,7 @@ from redis import Redis from tests.integration._support.client import Gateway, eventually, gateway_from_environment from tests.integration._support.generation import LIFECYCLE_SETTINGS from tests.integration._support.manifest import OWNED_DIRECTORIES +from tests.integration._support.routing import RoutingPlugin COLLECTED: Final = pytest.StashKey[tuple[str, ...]]() REPORTS: Final = pytest.StashKey[list[pytest.TestReport]]() @@ -29,6 +30,8 @@ def pytest_configure(config: pytest.Config) -> None: config.addinivalue_line("markers", "covers(*ids): legacy contract IDs kept for existing tests, not enforced") config.stash[REPORTS] = [] config.pluginmanager.register(IntegrationReportPlugin(config)) + if os.environ.get("INTEGRATION_ROUTING"): + config.pluginmanager.register(RoutingPlugin(config)) class IntegrationReportPlugin: diff --git a/tests/integration/database/test_transaction_atomicity.py b/tests/integration/database/test_transaction_atomicity.py index c150354d9a6..030f0f3e445 100644 --- a/tests/integration/database/test_transaction_atomicity.py +++ b/tests/integration/database/test_transaction_atomicity.py @@ -43,6 +43,7 @@ def test_access_group_second_key_constraint_failure_rolls_back_all_writes(gatewa ) with psycopg.connect(os.environ["DATABASE_URL"], autocommit=True) as connection, ExitStack() as cleanup: connection.execute(sql.SQL("CREATE SEQUENCE {}").format(sql.Identifier(witness))) + connection.execute(sql.SQL("GRANT USAGE ON SEQUENCE {} TO PUBLIC").format(sql.Identifier(witness))) cleanup.callback(connection.execute, sql.SQL("DROP SEQUENCE {}").format(sql.Identifier(witness))) connection.execute( sql.SQL( diff --git a/tests/integration/management/test_user_updates_wedged_coordination_redis.py b/tests/integration/management/test_user_updates_wedged_coordination_redis.py index d8c84a70778..0715059744e 100644 --- a/tests/integration/management/test_user_updates_wedged_coordination_redis.py +++ b/tests/integration/management/test_user_updates_wedged_coordination_redis.py @@ -121,6 +121,7 @@ def test_user_budget_updates_return_promptly_while_coordination_redis_is_wedged( "REDIS_PORT": str(coordination.port), }, config=Path("tests/integration/coordination_redis_proxy_config.yaml"), + remove_environment=("DATABASE_URL_READ_REPLICA",), workers=2, ) as candidate, Redis(host=coordination.host, port=coordination.port, socket_timeout=1) as subscriber_client, diff --git a/tests/integration/management/test_vector_store_config_ownership.py b/tests/integration/management/test_vector_store_config_ownership.py index e1e7ac42472..e4ea4e324ff 100644 --- a/tests/integration/management/test_vector_store_config_ownership.py +++ b/tests/integration/management/test_vector_store_config_ownership.py @@ -318,7 +318,14 @@ def test_redis_outage_keeps_config_store_served_and_recovers( "REDIS_PORT": str(cache.port), "REDIS_CIRCUIT_BREAKER_RECOVERY_TIMEOUT": "1", } - with owned_proxy(gateway, tmp_path, overrides, config=PROXY_CONFIG, workers=2) as candidate: + with owned_proxy( + gateway, + tmp_path, + overrides, + config=PROXY_CONFIG, + workers=2, + remove_environment=("DATABASE_URL_READ_REPLICA",), + ) as candidate: db_store_id: Final = f"vs_db_{uuid.uuid4().hex}" for phase in ("before", "during", "after"): if phase == "during": diff --git a/tests/integration/routing/either_role.json b/tests/integration/routing/either_role.json new file mode 100644 index 00000000000..68824b50e11 --- /dev/null +++ b/tests/integration/routing/either_role.json @@ -0,0 +1,3 @@ +{ + "SELECT $n": "SELECT 1 health probe: the database watchdog and probe target query whichever pool reader_unavailable selects and the reconnect smoke test always uses the writer, all timer driven, so it lands under whichever test is in flight" +} diff --git a/tests/integration/routing/test_redis_recovery.py b/tests/integration/routing/test_redis_recovery.py index 81d27a190b0..9b41d44926f 100644 --- a/tests/integration/routing/test_redis_recovery.py +++ b/tests/integration/routing/test_redis_recovery.py @@ -26,7 +26,7 @@ def test_owned_redis_outage_recovers_requests_and_real_response_cache(gateway: G try: with owned_redis(tmp_path) as cache, monkeypatch.context() as environment: environment.setenv("DATABASE_URL", database_url) - with owned_proxy(gateway, tmp_path, {"DATABASE_URL": database_url, "REDIS_HOST": cache.host, "REDIS_PORT": str(cache.port), "REDIS_CIRCUIT_BREAKER_RECOVERY_TIMEOUT": "1"}) as candidate, candidate.scenario() as scenario, httpx.Client(base_url=gateway.upstream_url, timeout=5, trust_env=False) as upstream: + with owned_proxy(gateway, tmp_path, {"DATABASE_URL": database_url, "REDIS_HOST": cache.host, "REDIS_PORT": str(cache.port), "REDIS_CIRCUIT_BREAKER_RECOVERY_TIMEOUT": "1"}, remove_environment=("DATABASE_URL_READ_REPLICA",)) as candidate, candidate.scenario() as scenario, httpx.Client(base_url=gateway.upstream_url, timeout=5, trust_env=False) as upstream: model: Final = scenario.model() key: Final = scenario.key(models=[model]) for generation in ("before", "after"): diff --git a/tests/integration/spend/test_daily_rollup_retry.py b/tests/integration/spend/test_daily_rollup_retry.py index cf1b989639a..5b48d841f75 100644 --- a/tests/integration/spend/test_daily_rollup_retry.py +++ b/tests/integration/spend/test_daily_rollup_retry.py @@ -26,6 +26,7 @@ def _install_daily_user_rollup_fault(user_id: str) -> str: _execute( ( sql.SQL("CREATE SEQUENCE {}").format(sequence), + sql.SQL("GRANT USAGE ON SEQUENCE {} TO PUBLIC").format(sequence), sql.SQL( "CREATE FUNCTION {}() RETURNS trigger LANGUAGE plpgsql AS $fault$ " "BEGIN PERFORM nextval({}); " diff --git a/tests/integration/spend/test_shutdown_flush.py b/tests/integration/spend/test_shutdown_flush.py index b744c1f4b5f..5441f70d59e 100644 --- a/tests/integration/spend/test_shutdown_flush.py +++ b/tests/integration/spend/test_shutdown_flush.py @@ -196,12 +196,14 @@ def _proxy_with_one_seeded_row( gateway, tmp_path, { + "DATABASE_URL": os.environ["DATABASE_URL"], "LITELLM_LOG": "DEBUG", "GRACEFUL_SHUTDOWN_TIMEOUT": "1", "SCHEDULED_JOB_SHUTDOWN_FINISH_TIMEOUT_SECONDS": "1", "SCHEDULED_JOB_SHUTDOWN_CANCEL_TIMEOUT_SECONDS": str(cancel_timeout_seconds), }, config=_config_with_pool_limit(tmp_path, pool_limit), + remove_environment=("DATABASE_URL_READ_REPLICA",), workers=workers, ) as owned: key: Final = string_value( diff --git a/tests/unit/integration_support/__init__.py b/tests/unit/integration_support/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/integration_support/test_routing.py b/tests/unit/integration_support/test_routing.py new file mode 100644 index 00000000000..22a74423eb7 --- /dev/null +++ b/tests/unit/integration_support/test_routing.py @@ -0,0 +1,438 @@ +from __future__ import annotations + +import json +from collections.abc import Iterator +from pathlib import Path +from types import MappingProxyType +from typing import Final + +import pytest + +from tests.integration._support.routing import ( + DIFF_FILE, + OBSERVED_FILE, + READER_ROLE, + WRITER_ROLE, + Mismatch, + Observation, + compare, + delta, + dump_observation, + load_either_role, + load_observation, + main, + normalize, + render, + role_calls, +) + +TOKEN_QUERY: Final = 'UPDATE "LiteLLM_VerificationToken" SET token = $n WHERE token = $n' +NODE_ID: Final = "tests/integration/management/test_keys.py::test_generate" + + +def _routing(entries: dict[str, tuple[str, ...]]) -> MappingProxyType[str, frozenset[str]]: + return MappingProxyType({query: frozenset(roles) for query, roles in entries.items()}) + + +def _observation( + queries: dict[str, tuple[str, ...]], + tests: dict[str, dict[str, tuple[str, ...]]] | None = None, + calls: dict[str, int] | None = None, + dealloc: int = 0, +) -> Observation: + return Observation( + _routing(queries), + MappingProxyType({node: _routing(mapping) for node, mapping in (tests or {}).items()}), + MappingProxyType(calls if calls is not None else {"litellm_reader": 3, "litellm_writer": 7}), + dealloc, + ) + + +@pytest.mark.parametrize( + ("raw", "expected"), + [ + ("SELECT a\n FROM t", "SELECT a FROM t"), + ("SELECT * FROM t WHERE id IN ($1, $2, $3)", "SELECT * FROM t WHERE id IN ($n)"), + ("SELECT * FROM t WHERE id IN ($1,$2)", "SELECT * FROM t WHERE id IN ($n)"), + ("SELECT * FROM t WHERE id IN ($4)", "SELECT * FROM t WHERE id IN ($n)"), + ( + "INSERT INTO t VALUES ($1, $2) ON CONFLICT ($3, $4, $5) DO NOTHING", + "INSERT INTO t VALUES ($n) ON CONFLICT ($n) DO NOTHING", + ), + ], +) +def test_normalize_collapses_whitespace_and_placeholders(raw: str, expected: str) -> None: + assert normalize(raw) == expected + + +def test_compare_reports_global_role_mismatch() -> None: + base: Final = _observation({TOKEN_QUERY: ("litellm_reader",), "SELECT 1": ("litellm_writer",)}) + head: Final = _observation({TOKEN_QUERY: ("litellm_writer",), "SELECT 1": ("litellm_writer",)}) + report: Final = compare(base, head) + assert report.mismatches == (Mismatch(None, TOKEN_QUERY, ("litellm_reader",), ("litellm_writer",)),) + assert report.failures() == (f"global: {TOKEN_QUERY}: base [litellm_reader] head [litellm_writer]",) + + +def test_compare_reports_global_shrink_mismatch() -> None: + base: Final = _observation({TOKEN_QUERY: ("litellm_reader", "litellm_writer")}) + head: Final = _observation({TOKEN_QUERY: ("litellm_writer",)}) + report: Final = compare(base, head) + assert report.mismatches == ( + Mismatch(None, TOKEN_QUERY, ("litellm_reader", "litellm_writer"), ("litellm_writer",)), + ) + assert report.failures() == (f"global: {TOKEN_QUERY}: base [litellm_reader, litellm_writer] head [litellm_writer]",) + + +def test_compare_reports_per_test_mismatch_with_nodeid() -> None: + base: Final = _observation( + {TOKEN_QUERY: ("litellm_reader",)}, + {NODE_ID: {TOKEN_QUERY: ("litellm_reader",)}}, + ) + head: Final = _observation( + {TOKEN_QUERY: ("litellm_reader",)}, + {NODE_ID: {TOKEN_QUERY: ("litellm_writer",)}}, + ) + report: Final = compare(base, head) + assert report.mismatches == (Mismatch(NODE_ID, TOKEN_QUERY, ("litellm_reader",), ("litellm_writer",)),) + assert report.failures() == (f"{NODE_ID}: {TOKEN_QUERY}: base [litellm_reader] head [litellm_writer]",) + + +def test_compare_per_test_mismatch_ignores_global_observation() -> None: + base: Final = _observation( + {TOKEN_QUERY: ("litellm_reader", "litellm_writer")}, + {NODE_ID: {TOKEN_QUERY: ("litellm_reader",)}}, + ) + head: Final = _observation( + {TOKEN_QUERY: ("litellm_reader", "litellm_writer")}, + {NODE_ID: {TOKEN_QUERY: ("litellm_writer",)}}, + ) + report: Final = compare(base, head) + assert report.mismatches == (Mismatch(NODE_ID, TOKEN_QUERY, ("litellm_reader",), ("litellm_writer",)),) + assert report.failures() == (f"{NODE_ID}: {TOKEN_QUERY}: base [litellm_reader] head [litellm_writer]",) + + +def test_compare_reports_per_test_shrink_mismatch() -> None: + base: Final = _observation( + {TOKEN_QUERY: ("litellm_reader", "litellm_writer")}, + {NODE_ID: {TOKEN_QUERY: ("litellm_reader", "litellm_writer")}}, + ) + head: Final = _observation( + {TOKEN_QUERY: ("litellm_reader", "litellm_writer")}, + {NODE_ID: {TOKEN_QUERY: ("litellm_reader",)}}, + ) + report: Final = compare(base, head) + assert report.mismatches == ( + Mismatch(NODE_ID, TOKEN_QUERY, ("litellm_reader", "litellm_writer"), ("litellm_reader",)), + ) + assert report.failures() == ( + f"{NODE_ID}: {TOKEN_QUERY}: base [litellm_reader, litellm_writer] head [litellm_reader]", + ) + + +def test_compare_reports_global_gain_mismatch() -> None: + base: Final = _observation({TOKEN_QUERY: ("litellm_reader",)}) + head: Final = _observation({TOKEN_QUERY: ("litellm_reader", "litellm_writer")}) + report: Final = compare(base, head) + assert report.mismatches == ( + Mismatch(None, TOKEN_QUERY, ("litellm_reader",), ("litellm_reader", "litellm_writer")), + ) + assert report.failures() == (f"global: {TOKEN_QUERY}: base [litellm_reader] head [litellm_reader, litellm_writer]",) + + +def test_compare_reports_per_test_gain_mismatch() -> None: + base: Final = _observation( + {TOKEN_QUERY: ("litellm_reader",)}, + {NODE_ID: {TOKEN_QUERY: ("litellm_reader",)}}, + ) + head: Final = _observation( + {TOKEN_QUERY: ("litellm_reader",)}, + {NODE_ID: {TOKEN_QUERY: ("litellm_reader", "litellm_writer")}}, + ) + report: Final = compare(base, head) + assert report.mismatches == ( + Mismatch(NODE_ID, TOKEN_QUERY, ("litellm_reader",), ("litellm_reader", "litellm_writer")), + ) + assert report.failures() == ( + f"{NODE_ID}: {TOKEN_QUERY}: base [litellm_reader] head [litellm_reader, litellm_writer]", + ) + + +def test_compare_either_role_suppresses_and_reports_variance() -> None: + base: Final = _observation( + {TOKEN_QUERY: ("litellm_reader", "litellm_writer"), "SELECT quiet": ("litellm_reader",)}, + {NODE_ID: {TOKEN_QUERY: ("litellm_reader",)}}, + ) + head: Final = _observation( + {TOKEN_QUERY: ("litellm_writer",), "SELECT quiet": ("litellm_reader",)}, + {NODE_ID: {TOKEN_QUERY: ("litellm_writer",)}}, + ) + report: Final = compare(base, head, either_role=frozenset({TOKEN_QUERY, "SELECT quiet"})) + assert report.mismatches == () + assert report.failures() == () + assert report.either_role == (TOKEN_QUERY,) + assert "== either role ==\n" + TOKEN_QUERY + "\n" in render(report) + + +def test_compare_either_role_matches_exact_keys_only() -> None: + base: Final = _observation( + { + "SELECT $n": ("litellm_reader",), + "SELECT $n FROM x": ("litellm_reader",), + "SELECT $n FROM x WHERE y = $n": ("litellm_reader",), + } + ) + head: Final = _observation( + { + "SELECT $n": ("litellm_writer",), + "SELECT $n FROM x": ("litellm_writer",), + "SELECT $n FROM x WHERE y = $n": ("litellm_writer",), + } + ) + report: Final = compare(base, head, either_role=frozenset({"SELECT $n FROM x"})) + assert frozenset(mismatch.query for mismatch in report.mismatches) == frozenset( + {"SELECT $n", "SELECT $n FROM x WHERE y = $n"} + ) + other: Final = compare(base, head, either_role=frozenset({"SELECT $n"})) + assert frozenset(mismatch.query for mismatch in other.mismatches) == frozenset( + {"SELECT $n FROM x", "SELECT $n FROM x WHERE y = $n"} + ) + + +def test_compare_one_sided_queries_are_listed_not_failed() -> None: + base: Final = _observation({"SELECT a": ("litellm_reader",), "SELECT gone": ("litellm_writer",)}) + head: Final = _observation({"SELECT a": ("litellm_reader",), "SELECT new": ("litellm_writer",)}) + report: Final = compare(base, head) + assert report.only_base == ("SELECT gone",) + assert report.only_head == ("SELECT new",) + assert report.mismatches == () + assert report.failures() == () + + +def test_failures_flags_dealloc_evictions_on_base() -> None: + report: Final = compare(_observation({}, dealloc=1), _observation({})) + assert report.failures() == ("base: pg_stat_statements evicted 1 entries (dealloc > 0)",) + + +def test_failures_flags_dealloc_evictions_on_head() -> None: + report: Final = compare(_observation({}), _observation({}, dealloc=1)) + assert report.failures() == ("head: pg_stat_statements evicted 1 entries (dealloc > 0)",) + + +def test_failures_flags_silent_reader_on_base() -> None: + report: Final = compare( + _observation({}, calls={"litellm_reader": 0, "litellm_writer": 5}), + _observation({}), + ) + assert report.failures() == ("base: no litellm_reader calls observed",) + + +def test_failures_flags_silent_reader_on_head() -> None: + report: Final = compare( + _observation({}), + _observation({}, calls={"litellm_reader": 0, "litellm_writer": 5}), + ) + assert report.failures() == ("head: no litellm_reader calls observed",) + + +def test_failures_flags_silent_writer() -> None: + report: Final = compare( + _observation({}, calls={"litellm_reader": 5, "litellm_writer": 0}), + _observation({}), + ) + assert report.failures() == ("base: no litellm_writer calls observed",) + + +def test_failures_counts_missing_role_as_silent() -> None: + report: Final = compare(_observation({}), _observation({}, calls={"litellm_writer": 5})) + assert report.failures() == ("head: no litellm_reader calls observed",) + + +def test_compare_skips_per_test_mismatches_for_xdist_shape() -> None: + base: Final = _observation( + {TOKEN_QUERY: ("litellm_reader",)}, + {NODE_ID: {TOKEN_QUERY: ("litellm_reader",)}}, + ) + head: Final = _observation({TOKEN_QUERY: ("litellm_reader",)}) + assert head.tests == {} + report: Final = compare(base, head) + assert report.mismatches == () + assert report.failures() == () + + +class _WriterFirst(frozenset[str]): + def __iter__(self) -> Iterator[str]: + return iter((WRITER_ROLE, READER_ROLE)) + + +def test_dump_observation_sorts_role_lists_and_round_trips(tmp_path: Path) -> None: + queries: Final = [f"SELECT {index}" for index in range(4)] + observation: Final = Observation( + MappingProxyType({query: _WriterFirst({WRITER_ROLE, READER_ROLE}) for query in queries}), + MappingProxyType( + {NODE_ID: MappingProxyType({query: _WriterFirst({WRITER_ROLE, READER_ROLE}) for query in queries})} + ), + MappingProxyType({READER_ROLE: 1, WRITER_ROLE: 2}), + 0, + ) + expected: Final = ( + json.dumps( + { + "queries": {query: ["litellm_reader", "litellm_writer"] for query in queries}, + "tests": {NODE_ID: {query: ["litellm_reader", "litellm_writer"] for query in queries}}, + "calls": {"litellm_reader": 1, "litellm_writer": 2}, + "dealloc": 0, + }, + sort_keys=True, + indent=2, + ) + + "\n" + ) + dumped: Final = dump_observation(observation) + assert dumped == expected + path: Final = tmp_path / OBSERVED_FILE + path.write_text(dumped) + loaded: Final = load_observation(path) + assert loaded.queries == _routing({query: (WRITER_ROLE, READER_ROLE) for query in queries}) + assert loaded.tests == {NODE_ID: loaded.queries} + + +def _write_observed(results: Path, observation: Observation) -> None: + results.mkdir(parents=True, exist_ok=True) + (results / OBSERVED_FILE).write_text(dump_observation(observation)) + + +def test_main_check_returns_zero_for_matching_routes(tmp_path: Path) -> None: + base_dir: Final = tmp_path / "base" + head_dir: Final = tmp_path / "head" + observation: Final = _observation({TOKEN_QUERY: ("litellm_reader",)}) + _write_observed(base_dir, observation) + _write_observed(head_dir, observation) + assert main(["check", str(base_dir), str(head_dir)]) == 0 + diff: Final = (tmp_path / DIFF_FILE).read_text() + assert "== failures ==\nnone\n" in diff + + +def test_main_check_returns_one_and_writes_exact_diff(tmp_path: Path) -> None: + base_dir: Final = tmp_path / "parity" / "base" + head_dir: Final = tmp_path / "parity" / "head" + _write_observed( + base_dir, + _observation({TOKEN_QUERY: ("litellm_reader",), "SELECT absent": ("litellm_writer",)}), + ) + _write_observed( + head_dir, + _observation( + {TOKEN_QUERY: ("litellm_writer",)}, + calls={"litellm_reader": 0, "litellm_writer": 5}, + dealloc=2, + ), + ) + assert main(["check", str(base_dir), str(head_dir)]) == 1 + assert (head_dir.parent / DIFF_FILE).read_text() == ( + "== failures ==\n" + f"global: {TOKEN_QUERY}: base [litellm_reader] head [litellm_writer]\n" + "head: pg_stat_statements evicted 2 entries (dealloc > 0)\n" + "head: no litellm_reader calls observed\n" + "\n" + "== either role ==\n" + "none\n" + "\n" + "== queries only in base ==\n" + "SELECT absent\n" + "\n" + "== queries only in head ==\n" + "none\n" + "\n" + "== calls ==\n" + "base litellm_reader: 3\n" + "base litellm_writer: 7\n" + "base dealloc: 0\n" + "head litellm_reader: 0\n" + "head litellm_writer: 5\n" + "head dealloc: 2\n" + ) + + +def test_main_check_missing_observed_returns_one(tmp_path: Path, capsys: pytest.CaptureFixture[str]) -> None: + base_dir: Final = tmp_path / "base" + head_dir: Final = tmp_path / "head" + _write_observed(base_dir, _observation({})) + head_dir.mkdir() + assert main(["check", str(base_dir), str(head_dir)]) == 1 + assert "observed routing file missing" in capsys.readouterr().err + + +def test_main_check_either_role_suppresses_shrink(tmp_path: Path) -> None: + base_dir: Final = tmp_path / "base" + head_dir: Final = tmp_path / "head" + _write_observed(base_dir, _observation({TOKEN_QUERY: ("litellm_reader", "litellm_writer")})) + _write_observed(head_dir, _observation({TOKEN_QUERY: ("litellm_writer",)})) + argv: Final = ["check", str(base_dir), str(head_dir)] + allowlist: Final = tmp_path / "either.json" + allowlist.write_text(json.dumps({TOKEN_QUERY: "timer probe may use either pool"})) + assert main([*argv, "--either-role", str(allowlist)]) == 0 + assert "== either role ==\n" + TOKEN_QUERY + "\n" in (tmp_path / DIFF_FILE).read_text() + assert main(argv) == 1 + + +def test_delta_maps_positive_increases_per_role() -> None: + before: Final = MappingProxyType( + { + ("litellm_reader", "SELECT both"): 1, + ("litellm_writer", "SELECT both"): 2, + ("litellm_reader", "SELECT reader"): 3, + ("litellm_writer", "SELECT gone"): 4, + ("litellm_reader", "SELECT same"): 5, + } + ) + after: Final = MappingProxyType( + { + ("litellm_reader", "SELECT both"): 2, + ("litellm_writer", "SELECT both"): 5, + ("litellm_reader", "SELECT reader"): 6, + ("litellm_reader", "SELECT same"): 5, + ("litellm_writer", "SELECT writer"): 7, + } + ) + assert delta(before, after) == { + "SELECT both": frozenset({"litellm_reader", "litellm_writer"}), + "SELECT reader": frozenset({"litellm_reader"}), + "SELECT writer": frozenset({"litellm_writer"}), + } + + +def test_role_calls_sums_positive_increases_per_role() -> None: + before: Final = MappingProxyType( + { + ("litellm_reader", "SELECT a"): 10, + ("litellm_reader", "SELECT b"): 4, + ("litellm_writer", "SELECT a"): 1, + } + ) + after: Final = MappingProxyType( + { + ("litellm_reader", "SELECT a"): 11, + ("litellm_reader", "SELECT b"): 2, + ("litellm_writer", "SELECT a"): 1, + ("litellm_writer", "SELECT c"): 6, + } + ) + assert role_calls(before, after) == {"litellm_reader": 1, "litellm_writer": 6} + + +def test_load_either_role_missing_path_returns_empty(tmp_path: Path) -> None: + assert load_either_role(tmp_path / "absent.json") == frozenset() + + +def test_load_either_role_reads_query_keys(tmp_path: Path) -> None: + path: Final = tmp_path / "either.json" + path.write_text(json.dumps({"SELECT $n": "probe", "SELECT now()": "clock"})) + assert load_either_role(path) == frozenset({"SELECT $n", "SELECT now()"}) + + +def test_load_observation_reads_calls_and_dealloc(tmp_path: Path) -> None: + observation: Final = _observation({TOKEN_QUERY: ("litellm_reader",)}, dealloc=0) + path: Final = tmp_path / OBSERVED_FILE + path.write_text(dump_observation(observation)) + loaded: Final = load_observation(path) + assert loaded == observation From 09ebb28473e6e9e09c80ce2f822b88d8bc24f2e4 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 24 Sep 2026 07:32:00 +0000 Subject: [PATCH 026/120] fix(s3_v2): bound concurrent S3 uploads per flush and add opt-in JSONL batch files (#41258) * fix(s3_v2): bound concurrent S3 uploads per flush and add opt-in JSONL batch files Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(s3_v2): keep failed uploads queued, parse env-backed flags, add integration coverage Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * refactor(s3_v2): type test helpers and honor constructor bound when config value is null Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * refactor(s3): annotate required casts for the type-discipline gate Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(s3_v2): keep tenant prefixes, stable retries and cold storage safety in batch file mode Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(s3_v2): cover root-level batch file keys for codecov patch target Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(s3_v2): audit matrix across chat, messages and responses surfaces with sink faults Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(s3_v2): read sink objects under the lock in the audit cells Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(s3_v2): rebind the retry queue instead of slicing in place Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(s3_v2): ignore stray non-POST requests in the surface upstream Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(s3_v2): keep fake upload state on the fake client instead of nonlocal counters Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Co-authored-by: yucheng --- litellm/constants.py | 1 + litellm/integrations/s3.py | 35 +- litellm/integrations/s3_v2.py | 126 +++- litellm/types/integrations/s3_v2.py | 2 + .../observability/_s3_v2_support.py | 344 ++++++++++ .../test_s3_v2_flush_surfaces.py | 97 +++ .../observability/test_s3_v2_upload_fanout.py | 630 ++++++++++++++++++ tests/test_litellm/integrations/test_s3_v2.py | 450 +++++++++++++ 8 files changed, 1665 insertions(+), 20 deletions(-) create mode 100644 tests/integration/observability/_s3_v2_support.py create mode 100644 tests/integration/observability/test_s3_v2_flush_surfaces.py create mode 100644 tests/integration/observability/test_s3_v2_upload_fanout.py diff --git a/litellm/constants.py b/litellm/constants.py index dcc12ef2ba9..3ff80d8b7dd 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -48,6 +48,7 @@ DEFAULT_BATCH_SIZE: Final = int(os.getenv("DEFAULT_BATCH_SIZE", 512)) DEFAULT_FLUSH_INTERVAL_SECONDS: Final = int(os.getenv("DEFAULT_FLUSH_INTERVAL_SECONDS", 5)) DEFAULT_S3_FLUSH_INTERVAL_SECONDS: Final = int(os.getenv("DEFAULT_S3_FLUSH_INTERVAL_SECONDS", 10)) DEFAULT_S3_BATCH_SIZE: Final = int(os.getenv("DEFAULT_S3_BATCH_SIZE", 512)) +DEFAULT_S3_MAX_CONCURRENT_UPLOADS: Final = int(os.getenv("DEFAULT_S3_MAX_CONCURRENT_UPLOADS", "16")) # https://docs.aws.amazon.com/AmazonS3/latest/userguide/object-keys.html MAX_S3_OBJECT_KEY_BYTES: Final = 1024 S3_BOUNDED_OBJECT_KEY_HEAD_BYTES: Final = 64 diff --git a/litellm/integrations/s3.py b/litellm/integrations/s3.py index 54ec876fd0d..3c8619e82b2 100644 --- a/litellm/integrations/s3.py +++ b/litellm/integrations/s3.py @@ -20,7 +20,8 @@ from litellm.constants import ( ) from litellm.types.utils import StandardLoggingPayload -_S3_LOG_PROMPTS_ONLY: Final = TypeAdapter(bool) +_S3_BOOL: Final = TypeAdapter(bool) +_UPLOAD_BOUND: Final = TypeAdapter(int) def resolve_s3_log_prompts_only(configured: object, environ: Mapping[str, str] | None = None) -> bool: @@ -29,12 +30,42 @@ def resolve_s3_log_prompts_only(configured: object, environ: Mapping[str, str] | if raw is None or raw == "": return False try: - return _S3_LOG_PROMPTS_ONLY.validate_python(raw.strip() if isinstance(raw, str) else raw) + return _S3_BOOL.validate_python(raw.strip() if isinstance(raw, str) else raw) except ValidationError: verbose_logger.warning("s3 logging: s3_log_prompts_only=%r is not a boolean, logging prompts only", raw) return True +def resolve_s3_max_concurrent_uploads(configured: object, fallback: int) -> int: + if configured is None or configured == "": + return fallback + try: + bound: Final = _UPLOAD_BOUND.validate_python(configured.strip() if isinstance(configured, str) else configured) + except ValidationError: + verbose_logger.warning( + "s3 logging: s3_max_concurrent_uploads=%r is not an integer, using %s", configured, fallback + ) + return fallback + if bound < 1: + verbose_logger.warning( + "s3 logging: s3_max_concurrent_uploads=%r must be at least 1, using %s", configured, fallback + ) + return fallback + return bound + + +def resolve_s3_batch_file_upload(configured: object) -> bool: + if configured is None or configured == "": + return False + try: + return _S3_BOOL.validate_python(configured.strip() if isinstance(configured, str) else configured) + except ValidationError: + verbose_logger.warning( + "s3 logging: s3_batch_file_upload=%r is not a boolean, keeping per-request objects", configured + ) + return False + + def prompts_only_payload(payload: StandardLoggingPayload) -> StandardLoggingPayload: return {**payload, "response": None} diff --git a/litellm/integrations/s3_v2.py b/litellm/integrations/s3_v2.py index 826f55cc798..dc33fe6c2bd 100644 --- a/litellm/integrations/s3_v2.py +++ b/litellm/integrations/s3_v2.py @@ -3,26 +3,33 @@ s3 Bucket Logging Integration async_log_success_event: Processes the event, stores it in memory for DEFAULT_S3_FLUSH_INTERVAL_SECONDS seconds or until DEFAULT_S3_BATCH_SIZE and then flushes to s3 async_log_failure_event: Processes the event, stores it in memory for DEFAULT_S3_FLUSH_INTERVAL_SECONDS seconds or until DEFAULT_S3_BATCH_SIZE and then flushes to s3 -NOTE 1: S3 does not provide a BATCH PUT API endpoint, so we create tasks to upload each element individually +NOTE 1: S3 does not provide a BATCH PUT API endpoint; by default each element is uploaded concurrently (bounded by s3_max_concurrent_uploads), or with s3_batch_file_upload the whole flush is written as one .jsonl file """ import asyncio import time from collections.abc import Mapping -from datetime import datetime +from datetime import datetime, timezone from typing import TYPE_CHECKING, Final, cast from urllib.parse import quote +from uuid import uuid4 import httpx import litellm from litellm._logging import print_verbose, verbose_logger -from litellm.constants import DEFAULT_S3_BATCH_SIZE, DEFAULT_S3_FLUSH_INTERVAL_SECONDS +from litellm.constants import ( + DEFAULT_S3_BATCH_SIZE, + DEFAULT_S3_FLUSH_INTERVAL_SECONDS, + DEFAULT_S3_MAX_CONCURRENT_UPLOADS, +) from litellm.integrations.s3 import ( get_s3_object_download_filename, get_s3_object_key, prompts_only_payload, + resolve_s3_batch_file_upload, resolve_s3_log_prompts_only, + resolve_s3_max_concurrent_uploads, resolve_sse_params, ) from litellm.litellm_core_utils.aws_partition import get_aws_dns_suffix @@ -43,7 +50,20 @@ if TYPE_CHECKING: from botocore.credentials import Credentials +def _s3_key_parent(s3_object_key: str) -> str: + return s3_object_key.rsplit("/", 1)[0] if "/" in s3_object_key else "" + + +class S3BatchUploadError(Exception): + def __init__(self, failed: int, total: int) -> None: + self.failed = failed + self.total = total + super().__init__(f"{failed} of {total} S3 uploads failed; events kept in queue for the next flush") + + class S3Logger(CustomBatchLogger, BaseAWSLLM): + preserve_events_added_during_flush = True + def __init__( self, s3_bucket_name: str | None = None, @@ -71,6 +91,8 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM): s3_server_side_encryption: str | None = None, s3_sse_kms_key_id: str | None = None, s3_log_prompts_only: bool | None = None, + s3_max_concurrent_uploads: int = DEFAULT_S3_MAX_CONCURRENT_UPLOADS, + s3_batch_file_upload: bool = False, s3_callback_params_override: dict | None = None, **kwargs, ): @@ -112,7 +134,10 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM): s3_server_side_encryption=s3_server_side_encryption, s3_sse_kms_key_id=s3_sse_kms_key_id, s3_log_prompts_only=s3_log_prompts_only, + s3_max_concurrent_uploads=s3_max_concurrent_uploads, + s3_batch_file_upload=s3_batch_file_upload, ) + self._upload_semaphore = asyncio.Semaphore(self.s3_max_concurrent_uploads) verbose_logger.debug("s3 logger using endpoint url %s", s3_endpoint_url) # IMPORTANT @@ -168,6 +193,8 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM): s3_server_side_encryption: str | None = None, s3_sse_kms_key_id: str | None = None, s3_log_prompts_only: bool | None = None, + s3_max_concurrent_uploads: int = DEFAULT_S3_MAX_CONCURRENT_UPLOADS, + s3_batch_file_upload: bool = False, params_source: dict | None = None, ): """ @@ -226,6 +253,16 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM): params.get("s3_sse_kms_key_id") or s3_sse_kms_key_id, ) + configured_bound: Final = params.get("s3_max_concurrent_uploads") + self.s3_max_concurrent_uploads = resolve_s3_max_concurrent_uploads( + s3_max_concurrent_uploads if configured_bound is None or configured_bound == "" else configured_bound, + DEFAULT_S3_MAX_CONCURRENT_UPLOADS, + ) + + self.s3_batch_file_upload = s3_batch_file_upload or resolve_s3_batch_file_upload( + params.get("s3_batch_file_upload") + ) + def _build_object_url(self, s3_object_key: str) -> str: """ Build the exact URL that is both signed and sent, with the key percent-encoded once. @@ -347,7 +384,7 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM): verbose_logger.exception("s3 Layer Error - %s", e) self.handle_callback_failure(callback_name="S3Logger") - async def async_upload_data_to_s3(self, batch_logging_element: s3BatchLoggingElement): + async def async_upload_data_to_s3(self, batch_logging_element: s3BatchLoggingElement) -> bool: try: import base64 import hashlib @@ -364,7 +401,11 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM): url: Final = self._build_object_url(batch_logging_element.s3_object_key) # Convert JSON to string - json_string: Final = safe_dumps(batch_logging_element.payload) + json_string: Final = ( + batch_logging_element.body + if batch_logging_element.body is not None + else safe_dumps(batch_logging_element.payload) + ) # Calculate SHA256 hash of the content content_hash: Final = hashlib.sha256(json_string.encode("utf-8")).hexdigest() @@ -374,7 +415,7 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM): # Prepare the request headers: Final = { - "Content-Type": "application/json", + "Content-Type": batch_logging_element.content_type, "Content-MD5": content_md5, "x-amz-content-sha256": content_hash, "Content-Language": "en", @@ -421,27 +462,72 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM): except Exception as e: verbose_logger.exception("Error uploading to s3: %s", e) self.handle_callback_failure(callback_name="S3Logger") + return False + return True - async def async_send_batch(self): + async def async_send_batch(self) -> None: """ + Sends runs from self.log_queue. - Sends runs from self.log_queue - - Returns: None - - Raises: Does not raise an exception, will only verbose_logger.exception() + Raises S3BatchUploadError when any upload failed; CustomBatchLogger.flush_queue + keeps the surviving queue entries for the next flush. """ - verbose_logger.debug("s3_v2 logger - sending batch of %s", len(self.log_queue)) - if not self.log_queue: + batch: Final = tuple(self.log_queue) + if not batch: return + verbose_logger.debug("s3_v2 logger - sending batch of %s", len(batch)) ######################################################### # Flush the log queue to s3 # the log queue can be bounded by DEFAULT_S3_BATCH_SIZE # see custom_batch_logger.py which triggers the flush ######################################################### - for payload in self.log_queue: - asyncio.create_task(self.async_upload_data_to_s3(payload)) + uploads: Final = self._batch_file_elements(batch) if self._batch_file_mode_active() else batch + results: Final = await asyncio.gather(*(self._upload_bounded(element) for element in uploads)) + failed: Final = tuple(element for element, ok in zip(uploads, results, strict=True) if not ok) + if not failed: + return + self.log_queue = [*failed, *self.log_queue[len(batch) :]] + raise S3BatchUploadError(failed=len(failed), total=len(uploads)) + + def _batch_file_mode_active(self) -> bool: + if not self.s3_batch_file_upload: + return False + if litellm.cold_storage_custom_logger == "s3_v2": + verbose_logger.warning( + "s3 logging: s3_batch_file_upload is ignored because s3_v2 is the cold storage logger; " + "per-request objects are required for spend log lookups" + ) + return False + return True + + async def _upload_bounded(self, element: s3BatchLoggingElement) -> bool: + async with self._upload_semaphore: + return await self.async_upload_data_to_s3(element) + + def _batch_file_elements(self, batch: tuple[s3BatchLoggingElement, ...]) -> tuple[s3BatchLoggingElement, ...]: + now: Final = datetime.now(timezone.utc) + groups: Final = { + parent: tuple( + element for element in batch if element.body is None and _s3_key_parent(element.s3_object_key) == parent + ) + for parent in sorted({_s3_key_parent(element.s3_object_key) for element in batch if element.body is None}) + } + return tuple(element for element in batch if element.body is not None) + tuple( + self._build_batch_file_element(elements, parent, now) for parent, elements in groups.items() + ) + + def _build_batch_file_element( + self, elements: tuple[s3BatchLoggingElement, ...], parent: str, now: datetime + ) -> s3BatchLoggingElement: + batch_name: Final = f"batch_{now.strftime('%H-%M-%S')}_{uuid4().hex}" + return s3BatchLoggingElement( + payload={}, + body="\n".join(safe_dumps(element.payload) for element in elements), + content_type="application/x-ndjson", + s3_object_key=f"{parent}/{batch_name}.jsonl" if parent else f"{batch_name}.jsonl", + s3_object_download_filename=f"{batch_name}.jsonl", + ) def create_s3_batch_logging_element( self, @@ -521,7 +607,11 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM): url: Final = self._build_object_url(batch_logging_element.s3_object_key) # Convert JSON to string - json_string: Final = safe_dumps(batch_logging_element.payload) + json_string: Final = ( + batch_logging_element.body + if batch_logging_element.body is not None + else safe_dumps(batch_logging_element.payload) + ) # Calculate SHA256 hash of the content content_hash: Final = hashlib.sha256(json_string.encode("utf-8")).hexdigest() @@ -531,7 +621,7 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM): # Prepare the request headers: Final = { - "Content-Type": "application/json", + "Content-Type": batch_logging_element.content_type, "Content-MD5": content_md5, "x-amz-content-sha256": content_hash, "Content-Language": "en", diff --git a/litellm/types/integrations/s3_v2.py b/litellm/types/integrations/s3_v2.py index 32864bf5b8c..555b16dc141 100644 --- a/litellm/types/integrations/s3_v2.py +++ b/litellm/types/integrations/s3_v2.py @@ -9,3 +9,5 @@ class s3BatchLoggingElement(BaseModel): payload: dict s3_object_key: str s3_object_download_filename: str + body: str | None = None + content_type: str = "application/json" diff --git a/tests/integration/observability/_s3_v2_support.py b/tests/integration/observability/_s3_v2_support.py new file mode 100644 index 00000000000..104c0eda863 --- /dev/null +++ b/tests/integration/observability/_s3_v2_support.py @@ -0,0 +1,344 @@ +import asyncio +import json +import threading +import time +from collections.abc import Mapping +from concurrent.futures import ThreadPoolExecutor +from dataclasses import dataclass, field +from pathlib import Path +from types import MappingProxyType +from typing import Final + +import anthropic +import openai +import yaml +from integration._support.client import Gateway, JsonValue, eventually, object_value +from integration._support.wire import Reply, Request + +BUCKET: Final = "integration-bucket" +PREFIX: Final = "integration-logs" + + +@dataclass(slots=True) +class RecordingS3Sink: + """Records every accepted PUT body by target, tracks peak concurrency, and can reject a leading + run of PUT attempts with a chosen status before accepting. Serves stored bodies back on GET.""" + + fail_attempts: int = 0 + fail_until: float = 0.0 + fail_status: int = 503 + delay_seconds: float = 0.5 + lock: threading.Lock = field(default_factory=threading.Lock) + in_flight: int = 0 + peak: int = 0 + attempts: int = 0 + store: dict[str, bytes] = field(default_factory=dict) # mutable-ok: GET reads must see writes from earlier PUTs + + def respond(self, request: Request) -> Reply: + if request.method == "GET": + body: Final = self.store.get(request.target) + if body is None: + return Reply(status=404) + return Reply(body=body) + assert request.method == "PUT", request.method + assert request.target.startswith(f"/{BUCKET}/{PREFIX}/"), request.target + with self.lock: + self.attempts += 1 + if self.attempts <= self.fail_attempts or time.time() < self.fail_until: + return Reply( + status=self.fail_status, + body=b"SinkFailure", + content_type="application/xml", + ) + self.in_flight += 1 + self.peak = max(self.peak, self.in_flight) + self.store[request.target] = request.body + time.sleep(self.delay_seconds) + with self.lock: + self.in_flight -= 1 + return Reply() + + def objects(self) -> Mapping[str, bytes]: + with self.lock: + return MappingProxyType(dict(self.store)) + + def payloads(self) -> tuple[dict[str, JsonValue], ...]: + return tuple(object_value(json.loads(line)) for body in self.objects().values() for line in body.splitlines()) + + +def s3_config( + path: Path, sink_url: str, extra: Mapping[str, JsonValue], settings: Mapping[str, JsonValue] | None = None +) -> Path: + config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + config["litellm_settings"].update( + { + "callbacks": ["s3_v2"], + "s3_callback_params": { + "s3_bucket_name": BUCKET, + "s3_region_name": "us-east-1", + "s3_endpoint_url": sink_url, + "s3_path": PREFIX, + "s3_aws_access_key_id": "AKIAIOSFODNN7EXAMPLE", + "s3_aws_secret_access_key": "wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY", + **extra, + }, + **(settings or {}), + } + ) + target: Final = path / "s3_v2.yaml" + target.write_text(yaml.safe_dump(config)) + return target + + +def _chat_completion(identity: str) -> dict[str, JsonValue]: + return { + "id": identity, + "object": "chat.completion", + "created": 1, + "model": "gpt-4o-mini", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 11, "completion_tokens": 4, "total_tokens": 15}, + } + + +def _chat_stream_frames(identity: str) -> tuple[bytes, ...]: + chunks: Final = ( + { + "id": identity, + "object": "chat.completion.chunk", + "created": 1, + "model": "gpt-4o-mini", + "choices": [{"index": 0, "delta": {"role": "assistant", "content": "ok"}, "finish_reason": None}], + }, + { + "id": identity, + "object": "chat.completion.chunk", + "created": 1, + "model": "gpt-4o-mini", + "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 11, "completion_tokens": 4, "total_tokens": 15}, + }, + ) + return tuple(f"data: {json.dumps(chunk)}\n\n".encode() for chunk in chunks) + (b"data: [DONE]\n\n",) + + +def _messages_completion(identity: str) -> dict[str, JsonValue]: + return { + "id": identity, + "type": "message", + "role": "assistant", + "model": "claude-sonnet-4-5-20250929", + "content": [{"type": "text", "text": "ok"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 11, "output_tokens": 4}, + } + + +def _messages_stream_frames(identity: str) -> tuple[bytes, ...]: + events: Final = ( + ( + "message_start", + { + "type": "message_start", + "message": { + "id": identity, + "type": "message", + "role": "assistant", + "model": "claude-sonnet-4-5-20250929", + "content": [], + "stop_reason": None, + "usage": {"input_tokens": 11, "output_tokens": 1}, + }, + }, + ), + ( + "content_block_start", + {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}, + ), + ( + "content_block_delta", + {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "ok"}}, + ), + ("content_block_stop", {"type": "content_block_stop", "index": 0}), + ( + "message_delta", + {"type": "message_delta", "delta": {"stop_reason": "end_turn"}, "usage": {"output_tokens": 4}}, + ), + ("message_stop", {"type": "message_stop"}), + ) + return tuple(f"event: {name}\ndata: {json.dumps(payload)}\n\n".encode() for name, payload in events) + + +def _responses_completion(identity: str) -> dict[str, JsonValue]: + return { + "id": identity, + "object": "response", + "created_at": 1, + "status": "completed", + "model": "gpt-4o-mini", + "output": [ + { + "type": "message", + "id": f"msg_{identity}", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": "ok", "annotations": []}], + } + ], + "usage": {"input_tokens": 11, "output_tokens": 4, "total_tokens": 15}, + } + + +def _responses_stream_frames(identity: str) -> tuple[bytes, ...]: + events: Final = ( + ( + "response.created", + { + "type": "response.created", + "response": {**_responses_completion(identity), "status": "in_progress", "output": []}, + }, + ), + ( + "response.output_text.delta", + { + "type": "response.output_text.delta", + "item_id": f"msg_{identity}", + "output_index": 0, + "content_index": 0, + "delta": "ok", + }, + ), + ("response.completed", {"type": "response.completed", "response": _responses_completion(identity)}), + ) + return tuple(f"event: {name}\ndata: {json.dumps(payload)}\n\n".encode() for name, payload in events) + + +def surface_reply(request: Request) -> Reply: + """Scripted upstream that echoes the caller's marker string back as the response id.""" + if request.method != "POST" or not request.body: + return Reply(status=404) + body: Final = json.loads(request.body) + if request.target.endswith("/chat/completions"): + identity: Final = body["messages"][0]["content"] + if body.get("stream"): + return Reply(content_type="text/event-stream", chunks=_chat_stream_frames(identity)) + return Reply(body=json.dumps(_chat_completion(identity)).encode()) + if request.target.endswith("/messages"): + identity_messages: Final = body["messages"][0]["content"] + if body.get("stream"): + return Reply(content_type="text/event-stream", chunks=_messages_stream_frames(identity_messages)) + return Reply(body=json.dumps(_messages_completion(identity_messages)).encode()) + assert request.target.endswith("/responses"), request.target + identity_responses: Final = body["input"] + if body.get("stream"): + return Reply(content_type="text/event-stream", chunks=_responses_stream_frames(identity_responses)) + return Reply(body=json.dumps(_responses_completion(identity_responses)).encode()) + + +SURFACES: Final = ("chat", "chat_stream", "messages", "messages_stream", "responses", "responses_stream") + + +def call_surface( + candidate: Gateway, surface: str, openai_model: str, anthropic_model: str, key: str, marker: str +) -> tuple[str, str | None]: + """Drive one request through the given surface; return (client-visible response id, x-litellm-call-id).""" + base: Final = str(candidate.client.base_url).rstrip("/") + headers: Final = {"Authorization": f"Bearer {key}"} + if surface == "chat": + reply: Final = openai.OpenAI(base_url=f"{base}/v1", api_key=key).chat.completions.create( + model=openai_model, + messages=[{"role": "user", "content": marker}], + extra_body={"cache": {"no-cache": True}}, + ) + return reply.id, None + + async def chat_stream() -> str: + stream = await openai.AsyncOpenAI(base_url=f"{base}/v1", api_key=key).chat.completions.create( + model=openai_model, + messages=[{"role": "user", "content": marker}], + stream=True, + extra_body={"cache": {"no-cache": True}}, + ) + seen = "" + async for chunk in stream: + seen = chunk.id # rebind-ok: the stream yields one chunk at a time + return seen + + if surface == "chat_stream": + return asyncio.run(chat_stream()), None + if surface in ("messages", "messages_stream"): + client: Final = anthropic.Anthropic(base_url=base, api_key="anthropic-placeholder", default_headers=headers) + if surface == "messages": + reply_messages: Final = client.messages.create( + model=anthropic_model, max_tokens=16, messages=[{"role": "user", "content": marker}] + ) + return reply_messages.id, None + with client.messages.stream( + model=anthropic_model, max_tokens=16, messages=[{"role": "user", "content": marker}] + ) as stream: + final: Final = stream.get_final_message() + return final.id, None + if surface == "responses": + response: Final = candidate.request( + "POST", + "/v1/responses", + {"model": openai_model, "input": marker, "cache": {"no-cache": True}}, + key=key, + ) + assert response.status_code == 200, response.text + return str(response.json()["id"]), response.headers.get("x-litellm-call-id") + assert surface == "responses_stream", surface + with candidate.client.stream( + "POST", + "/v1/responses", + json={"model": openai_model, "input": marker, "stream": True}, + headers=headers, + ) as response: + text: Final = response.read().decode() + assert response.status_code == 200, text + call_id: Final = response.headers.get("x-litellm-call-id") + assert marker in text, text + return marker, call_id + + +def collect_payloads(sink: RecordingS3Sink, count: int, seconds: float = 60) -> tuple[dict[str, JsonValue], ...]: + """Wait until `count` stored payload lines exist, then return every stored payload object.""" + + def delivered() -> int: + return sum(len(body.splitlines()) for body in sink.objects().values()) + + eventually(delivered, lambda total: total >= count, seconds=seconds) + return sink.payloads() + + +def mixed_burst( + candidate: Gateway, openai_model: str, anthropic_model: str, key: str, marker: str, per_surface: int = 8 +) -> tuple[tuple[str, str | None], ...]: + """Fire `per_surface` requests on every surface; returns (response id, x-litellm-call-id) per request.""" + jobs: Final = tuple( + (surface, f"{marker}-{surface}-{index}") for surface in SURFACES for index in range(per_surface) + ) + + def call(job: tuple[str, str]) -> tuple[str, str | None]: + surface, identity = job + return call_surface(candidate, surface, openai_model, anthropic_model, key, identity) + + with ThreadPoolExecutor(max_workers=48) as pool: + return tuple(pool.map(call, jobs)) + + +def matched_ids( + payloads: tuple[dict[str, JsonValue], ...], answered: tuple[tuple[str, str | None], ...] +) -> frozenset[str]: + """Every payload must be accountable to an answered request by response id or litellm_call_id.""" + response_ids: Final = frozenset(observed for observed, _ in answered) + call_ids: Final = frozenset(call_id for _, call_id in answered if call_id is not None) + landed: Final = [] + for payload in payloads: + if payload["id"] in response_ids: + landed.append(payload["id"]) + continue + assert payload["litellm_call_id"] in call_ids, f"unmatched payload {payload['id']!r}" + landed.append(str(payload["id"])) + return frozenset(landed) diff --git a/tests/integration/observability/test_s3_v2_flush_surfaces.py b/tests/integration/observability/test_s3_v2_flush_surfaces.py new file mode 100644 index 00000000000..2e0b7260a13 --- /dev/null +++ b/tests/integration/observability/test_s3_v2_flush_surfaces.py @@ -0,0 +1,97 @@ +import re +import uuid +from pathlib import Path +from typing import Final + +import pytest +from _s3_v2_support import ( + BUCKET, + PREFIX, + RecordingS3Sink, + collect_payloads, + matched_ids, + mixed_burst, + s3_config, + surface_reply, +) +from integration._support.client import Gateway +from integration._support.process import owned_proxy +from integration._support.wire import wire_server + +PER_REQUEST_KEY: Final = re.compile(rf"^/{BUCKET}/{PREFIX}/\d{{4}}-\d{{2}}-\d{{2}}/.+\.json$") +BATCH_KEY: Final = re.compile( + rf"^/{BUCKET}/{PREFIX}/\d{{4}}-\d{{2}}-\d{{2}}/batch_\d{{2}}-\d{{2}}-\d{{2}}_[0-9a-f]{{32}}\.jsonl$" +) + + +@pytest.mark.covers("other.observability.s3_v2.mixed_surface_burst_bounds_puts_one_object_per_response_id") +def test_s3_v2_mixed_surface_burst_bounds_puts_one_object_per_response_id(gateway: Gateway, tmp_path: Path) -> None: + marker: Final = "s3mix" + uuid.uuid4().hex[:8] + sink: Final = RecordingS3Sink() + with wire_server(surface_reply) as provider, wire_server(sink.respond) as bucket: + config: Final = s3_config(tmp_path, bucket.url, {}) + with ( + owned_proxy(gateway, tmp_path, {"DEFAULT_S3_FLUSH_INTERVAL_SECONDS": "3"}, config=config) as candidate, + candidate.scenario() as scenario, + ): + openai_model: Final = scenario.model(api_base=provider.url + "/v1", api_key="synthetic-provider-key") + anthropic_model: Final = scenario.model( + model="anthropic/claude-sonnet-4-5-20250929", api_base=provider.url, api_key="synthetic-provider-key" + ) + key: Final = scenario.key(models=[openai_model, anthropic_model]) + answered: Final = mixed_burst(candidate, openai_model, anthropic_model, key, marker) + payloads: Final = collect_payloads(sink, len(answered)) + targets: Final = tuple(sink.objects()) + assert sum(1 for r in provider.drain() if r.method == "POST") == 48 + assert sink.peak <= 16, f"peak concurrent PUTs {sink.peak} exceeded the default bound" + assert all(PER_REQUEST_KEY.match(target) for target in targets), list(targets) + assert len(targets) == 48 + assert matched_ids(payloads, answered) + + +@pytest.mark.covers("other.observability.s3_v2.mixed_surface_batch_writes_ndjson_lines_per_response_id") +def test_s3_v2_mixed_surface_batch_writes_ndjson_lines_per_response_id(gateway: Gateway, tmp_path: Path) -> None: + marker: Final = "s3mixb" + uuid.uuid4().hex[:8] + sink: Final = RecordingS3Sink() + with wire_server(surface_reply) as provider, wire_server(sink.respond) as bucket: + config: Final = s3_config(tmp_path, bucket.url, {"s3_batch_file_upload": True}) + with ( + owned_proxy(gateway, tmp_path, {"DEFAULT_S3_FLUSH_INTERVAL_SECONDS": "3"}, config=config) as candidate, + candidate.scenario() as scenario, + ): + openai_model: Final = scenario.model(api_base=provider.url + "/v1", api_key="synthetic-provider-key") + anthropic_model: Final = scenario.model( + model="anthropic/claude-sonnet-4-5-20250929", api_base=provider.url, api_key="synthetic-provider-key" + ) + key: Final = scenario.key(models=[openai_model, anthropic_model]) + answered: Final = mixed_burst(candidate, openai_model, anthropic_model, key, marker) + payloads: Final = collect_payloads(sink, len(answered)) + targets: Final = tuple(sink.objects()) + puts: Final = bucket.drain() + assert sum(1 for r in provider.drain() if r.method == "POST") == 48 + assert all(BATCH_KEY.match(target) for target in targets), list(targets) + assert all(put.headers["content-type"] == "application/x-ndjson" for put in puts), [put.headers for put in puts] + assert matched_ids(payloads, answered) + assert len(payloads) == 48 + + +@pytest.mark.covers("other.observability.s3_v2.sink_outage_mid_mixed_burst_recovers_every_response_id") +def test_s3_v2_sink_outage_mid_mixed_burst_recovers_every_response_id(gateway: Gateway, tmp_path: Path) -> None: + marker: Final = "s3mixo" + uuid.uuid4().hex[:8] + sink: Final = RecordingS3Sink(fail_attempts=30, fail_status=503, delay_seconds=0.2) + with wire_server(surface_reply) as provider, wire_server(sink.respond) as bucket: + config: Final = s3_config(tmp_path, bucket.url, {}) + with ( + owned_proxy(gateway, tmp_path, {"DEFAULT_S3_FLUSH_INTERVAL_SECONDS": "3"}, config=config) as candidate, + candidate.scenario() as scenario, + ): + openai_model: Final = scenario.model(api_base=provider.url + "/v1", api_key="synthetic-provider-key") + anthropic_model: Final = scenario.model( + model="anthropic/claude-sonnet-4-5-20250929", api_base=provider.url, api_key="synthetic-provider-key" + ) + key: Final = scenario.key(models=[openai_model, anthropic_model]) + answered: Final = mixed_burst(candidate, openai_model, anthropic_model, key, marker) + payloads: Final = collect_payloads(sink, len(answered), seconds=90) + assert sum(1 for r in provider.drain() if r.method == "POST") == 48 + assert matched_ids(payloads, answered) + assert len(payloads) == 48, "a stored id was overwritten or duplicated" diff --git a/tests/integration/observability/test_s3_v2_upload_fanout.py b/tests/integration/observability/test_s3_v2_upload_fanout.py new file mode 100644 index 00000000000..3ebca152327 --- /dev/null +++ b/tests/integration/observability/test_s3_v2_upload_fanout.py @@ -0,0 +1,630 @@ +import json +import re +import threading +import time +import uuid +from collections.abc import Mapping +from concurrent.futures import ThreadPoolExecutor +from dataclasses import dataclass, field +from pathlib import Path +from typing import Final + +import httpx +import pytest +import yaml +from _s3_v2_support import RecordingS3Sink, collect_payloads +from _s3_v2_support import s3_config as _recording_s3_config +from integration._support.client import Gateway, JsonValue, eventually +from integration._support.process import group_members, owned_proxy, owned_proxy_process +from integration._support.wire import Reply, Request, Wire, wire_server + +BUCKET: Final = "integration-bucket" +PREFIX: Final = "integration-logs" +REQUESTS: Final = 64 +PUT_DELAY_SECONDS: Final = 0.5 + + +@dataclass(slots=True) +class S3Sink: + """Accepts every PUT after a fixed delay and records the peak number of PUTs in flight.""" + + lock: threading.Lock = field(default_factory=threading.Lock) + in_flight: int = 0 + peak: int = 0 + + def respond(self, request: Request) -> Reply: + assert request.method == "PUT", request.method + assert request.target.startswith(f"/{BUCKET}/{PREFIX}/"), request.target + with self.lock: + self.in_flight += 1 + self.peak = max(self.peak, self.in_flight) + time.sleep(PUT_DELAY_SECONDS) + with self.lock: + self.in_flight -= 1 + return Reply() + + +def _chat_reply(request: Request) -> Reply: + if request.method != "POST" or not request.body: + return Reply(status=404) + text: Final = json.loads(request.body)["messages"][0]["content"] + return Reply( + body=json.dumps( + { + "id": text, + "object": "chat.completion", + "created": 1, + "model": "gpt-4o-mini", + "choices": [{"index": 0, "message": {"role": "assistant", "content": text}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 11, "completion_tokens": 4, "total_tokens": 15}, + } + ).encode() + ) + + +def _s3_config(path: Path, sink_url: str, extra: Mapping[str, JsonValue]) -> Path: + config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + config["litellm_settings"].update( + { + "callbacks": ["s3_v2"], + "s3_callback_params": { + "s3_bucket_name": BUCKET, + "s3_region_name": "us-east-1", + "s3_endpoint_url": sink_url, + "s3_path": PREFIX, + "s3_aws_access_key_id": "AKIAIOSFODNN7EXAMPLE", + "s3_aws_secret_access_key": "wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY", + **extra, + }, + } + ) + target: Final = path / "s3_v2.yaml" + target.write_text(yaml.safe_dump(config)) + return target + + +def _burst(candidate: Gateway, model: str, key: str, marker: str) -> frozenset[str]: + ids: Final = tuple(f"{marker}-{index}" for index in range(REQUESTS)) + + def request(identity: str) -> str: + response: Final = candidate.request( + "POST", + "/v1/chat/completions", + {"model": model, "messages": [{"role": "user", "content": identity}], "cache": {"no-cache": True}}, + key=key, + ) + assert response.status_code == 200, response.text + return response.json()["id"] + + with ThreadPoolExecutor(max_workers=32) as pool: + returned: Final = frozenset(pool.map(request, ids)) + assert returned == frozenset(ids) + return returned + + +def _collect(bucket: Wire, count_lines: bool, expected: int) -> tuple[Request, ...]: + puts: Final[list[Request]] = [] # mutable-ok: drain() consumes the queue, later polls must keep earlier PUTs + + def delivered() -> int: + puts.extend(bucket.drain()) + return sum(len(put.body.splitlines()) if count_lines else 1 for put in puts) + + eventually(delivered, lambda total: total >= expected, seconds=30) + return tuple(puts) + + +PER_REQUEST_KEY: Final = re.compile(rf"^/{BUCKET}/{PREFIX}/\d{{4}}-\d{{2}}-\d{{2}}/.+\.json$") +BATCH_KEY: Final = re.compile( + rf"^/{BUCKET}/{PREFIX}/\d{{4}}-\d{{2}}-\d{{2}}/batch_\d{{2}}-\d{{2}}-\d{{2}}_[0-9a-f]{{32}}\.jsonl$" +) + + +@pytest.mark.covers("other.observability.s3_v2.flush_bounds_concurrent_puts_to_default_and_keeps_every_log") +def test_s3_v2_flush_bounds_concurrent_puts_to_the_default_of_sixteen(gateway: Gateway, tmp_path: Path) -> None: + marker: Final = "s3fan" + uuid.uuid4().hex[:8] + sink: Final = S3Sink() + with wire_server(_chat_reply) as provider, wire_server(sink.respond) as bucket: + config: Final = _s3_config(tmp_path, bucket.url, {}) + with ( + owned_proxy(gateway, tmp_path, {"DEFAULT_S3_FLUSH_INTERVAL_SECONDS": "3"}, config=config) as candidate, + candidate.scenario() as scenario, + ): + model: Final = scenario.model(api_base=provider.url + "/v1", api_key="synthetic-provider-key") + key: Final = scenario.key(models=[model]) + ids: Final = _burst(candidate, model, key, marker) + puts: Final = _collect(bucket, count_lines=False, expected=REQUESTS) + assert sum(1 for r in provider.drain() if r.method == "POST") == REQUESTS + assert sink.peak <= 16, f"peak concurrent PUTs {sink.peak} exceeded the default bound for {REQUESTS} queued logs" + assert all(PER_REQUEST_KEY.match(put.target) for put in puts), [put.target for put in puts] + assert frozenset(json.loads(put.body)["id"] for put in puts) == ids + assert len({put.target for put in puts}) == REQUESTS + + +@pytest.mark.covers("other.observability.s3_v2.configured_bound_and_env_backed_false_keeps_per_request_objects") +def test_s3_v2_honors_configured_bound_and_env_backed_false_batch_flag(gateway: Gateway, tmp_path: Path) -> None: + marker: Final = "s3cap" + uuid.uuid4().hex[:8] + sink: Final = S3Sink() + with wire_server(_chat_reply) as provider, wire_server(sink.respond) as bucket: + config: Final = _s3_config( + tmp_path, + bucket.url, + {"s3_max_concurrent_uploads": 4, "s3_batch_file_upload": "os.environ/INTEGRATION_S3_BATCH_FILE_UPLOAD"}, + ) + with ( + owned_proxy( + gateway, + tmp_path, + {"DEFAULT_S3_FLUSH_INTERVAL_SECONDS": "3", "INTEGRATION_S3_BATCH_FILE_UPLOAD": "false"}, + config=config, + ) as candidate, + candidate.scenario() as scenario, + ): + model: Final = scenario.model(api_base=provider.url + "/v1", api_key="synthetic-provider-key") + key: Final = scenario.key(models=[model]) + ids: Final = _burst(candidate, model, key, marker) + puts: Final = _collect(bucket, count_lines=False, expected=REQUESTS) + assert sum(1 for r in provider.drain() if r.method == "POST") == REQUESTS + assert sink.peak <= 4, f"peak concurrent PUTs {sink.peak} exceeded s3_max_concurrent_uploads=4" + assert all(PER_REQUEST_KEY.match(put.target) for put in puts), [put.target for put in puts] + assert frozenset(json.loads(put.body)["id"] for put in puts) == ids + + +@pytest.mark.covers("other.observability.s3_v2.batch_file_upload_writes_one_ndjson_object_per_flush") +def test_s3_v2_batch_file_upload_writes_one_jsonl_object_per_flush(gateway: Gateway, tmp_path: Path) -> None: + marker: Final = "s3jsonl" + uuid.uuid4().hex[:8] + sink: Final = S3Sink() + with wire_server(_chat_reply) as provider, wire_server(sink.respond) as bucket: + config: Final = _s3_config(tmp_path, bucket.url, {"s3_batch_file_upload": True}) + with ( + owned_proxy(gateway, tmp_path, {"DEFAULT_S3_FLUSH_INTERVAL_SECONDS": "3"}, config=config) as candidate, + candidate.scenario() as scenario, + ): + model: Final = scenario.model(api_base=provider.url + "/v1", api_key="synthetic-provider-key") + key: Final = scenario.key(models=[model]) + ids: Final = _burst(candidate, model, key, marker) + puts: Final = _collect(bucket, count_lines=True, expected=REQUESTS) + assert sum(1 for r in provider.drain() if r.method == "POST") == REQUESTS + assert len(puts) <= 2, f"{len(puts)} PUTs for {REQUESTS} logs; batch mode must write one object per flush" + assert all(BATCH_KEY.match(put.target) for put in puts), [put.target for put in puts] + assert all(put.headers["content-type"] == "application/x-ndjson" for put in puts), [put.headers for put in puts] + lines: Final = tuple(line for put in puts for line in put.body.decode().splitlines()) + assert frozenset(json.loads(line)["id"] for line in lines) == ids + assert len(lines) == REQUESTS + + +@pytest.mark.covers("other.observability.s3_v2.batch_file_upload_keeps_team_prefix_in_object_key") +def test_s3_v2_batch_file_upload_keeps_team_alias_prefix(gateway: Gateway, tmp_path: Path) -> None: + marker: Final = "s3team" + uuid.uuid4().hex[:8] + team_alias: Final = f"alpha-{uuid.uuid4().hex[:8]}" + team_batch_key: Final = re.compile( + rf"^/{BUCKET}/{PREFIX}/{team_alias}/\d{{4}}-\d{{2}}-\d{{2}}/batch_\d{{2}}-\d{{2}}-\d{{2}}_[0-9a-f]{{32}}\.jsonl$" + ) + sink: Final = S3Sink() + with wire_server(_chat_reply) as provider, wire_server(sink.respond) as bucket: + config: Final = _s3_config(tmp_path, bucket.url, {"s3_batch_file_upload": True, "s3_use_team_prefix": True}) + with ( + owned_proxy(gateway, tmp_path, {"DEFAULT_S3_FLUSH_INTERVAL_SECONDS": "3"}, config=config) as candidate, + candidate.scenario() as scenario, + ): + model: Final = scenario.model(api_base=provider.url + "/v1", api_key="synthetic-provider-key") + team: Final = scenario.team(team_alias=team_alias, models=[model]) + key: Final = scenario.key(team_id=team, models=[model]) + ids: Final = _burst(candidate, model, key, marker) + puts: Final = _collect(bucket, count_lines=True, expected=REQUESTS) + assert sum(1 for r in provider.drain() if r.method == "POST") == REQUESTS + assert len(puts) >= 1 + assert all(team_batch_key.match(put.target) for put in puts), [put.target for put in puts] + lines: Final = tuple(line for put in puts for line in put.body.decode().splitlines()) + assert frozenset(json.loads(line)["id"] for line in lines) == ids + assert len(lines) == REQUESTS + + +@pytest.mark.covers("other.observability.s3_v2.upstream_failure_events_land_alongside_successes") +def test_s3_v2_upstream_failure_events_land_alongside_successes(gateway: Gateway, tmp_path: Path) -> None: + marker: Final = "s3fail" + uuid.uuid4().hex[:8] + sink: Final = RecordingS3Sink() + + def provider(request: Request) -> Reply: + text: Final = json.loads(request.body)["messages"][0]["content"] + if text.endswith("-fail"): + return Reply( + status=401, + body=b'{"error": {"message": "synthetic upstream rejection", "code": "synthetic_401"}}', + ) + return _chat_reply(request) + + with wire_server(provider) as upstream, wire_server(sink.respond) as bucket: + config: Final = _s3_config(tmp_path, bucket.url, {}) + with ( + owned_proxy(gateway, tmp_path, {"DEFAULT_S3_FLUSH_INTERVAL_SECONDS": "3"}, config=config) as candidate, + candidate.scenario() as scenario, + ): + model: Final = scenario.model(api_base=upstream.url + "/v1", api_key="synthetic-provider-key") + key: Final = scenario.key(models=[model]) + success_ids: Final = tuple(f"{marker}-{index}" for index in range(8)) + failure_ids: Final = tuple(f"{marker}-{index}-fail" for index in range(4)) + + def send(identity: str) -> httpx.Response: + return candidate.request( + "POST", + "/v1/chat/completions", + {"model": model, "messages": [{"role": "user", "content": identity}], "cache": {"no-cache": True}}, + key=key, + ) + + with ThreadPoolExecutor(max_workers=12) as pool: + responses: Final = tuple(pool.map(send, (*success_ids, *failure_ids))) + ok: Final = responses[:8] + rejected: Final = responses[8:] + assert all(response.status_code == 200 for response in ok), [r.text for r in ok] + assert tuple(response.json()["id"] for response in ok) == success_ids + for response in rejected: + assert response.status_code in (400, 401), response.status_code + assert "synthetic upstream rejection" in response.text, response.text + failure_call_ids: Final = frozenset(response.headers["x-litellm-call-id"] for response in rejected) + payloads: Final = collect_payloads(sink, len(success_ids) + len(failure_ids)) + assert len(upstream.drain()) == len(success_ids) + len(failure_ids) + delivered: Final = frozenset(payload["id"] for payload in payloads if payload["status"] == "success") + assert delivered == frozenset(success_ids) + failures: Final = tuple(payload for payload in payloads if payload["status"] == "failure") + assert len(failures) == len(failure_ids) + assert frozenset(payload["litellm_call_id"] for payload in failures) == failure_call_ids + assert all("synthetic upstream rejection" in json.dumps(payload["error_information"]) for payload in failures) + + +@pytest.mark.covers("other.observability.s3_v2.invalid_or_empty_bound_falls_back_to_sixteen") +@pytest.mark.parametrize( + ("bad", "warns"), + [ + pytest.param("abc", True, id="non_integer"), + pytest.param(0, True, id="below_one"), + pytest.param("", False, id="empty"), + ], +) +def test_s3_v2_invalid_or_empty_bound_falls_back_to_sixteen( + gateway: Gateway, tmp_path: Path, bad: JsonValue, warns: bool +) -> None: + marker: Final = "s3bound" + uuid.uuid4().hex[:8] + sink: Final = RecordingS3Sink() + with wire_server(_chat_reply) as provider, wire_server(sink.respond) as bucket: + config: Final = _s3_config(tmp_path, bucket.url, {"s3_max_concurrent_uploads": bad}) + with ( + owned_proxy_process(gateway, tmp_path, {"DEFAULT_S3_FLUSH_INTERVAL_SECONDS": "3"}, config=config) as owned, + owned.gateway.scenario() as scenario, + ): + model: Final = scenario.model(api_base=provider.url + "/v1", api_key="synthetic-provider-key") + key: Final = scenario.key(models=[model]) + ids: Final = _burst(owned.gateway, model, key, marker) + payloads: Final = collect_payloads(sink, REQUESTS) + if warns: + eventually( + lambda: owned.log.read_text(), + lambda text: "s3_max_concurrent_uploads" in text, + seconds=15, + ) + else: + assert "s3_max_concurrent_uploads" not in owned.log.read_text() + assert sum(1 for r in provider.drain() if r.method == "POST") == REQUESTS + assert sink.peak <= 16, f"peak concurrent PUTs {sink.peak} exceeded the fallback bound" + assert frozenset(payload["id"] for payload in payloads) == ids + + +@pytest.mark.covers("other.observability.s3_v2.sink_rejection_requeues_and_delivers_every_id_once") +def test_s3_v2_sink_rejection_requeues_and_delivers_every_id_once(gateway: Gateway, tmp_path: Path) -> None: + marker: Final = "s3deny" + uuid.uuid4().hex[:8] + sink: Final = RecordingS3Sink(fail_status=403, delay_seconds=0.2) + with wire_server(_chat_reply) as provider, wire_server(sink.respond) as bucket: + config: Final = _s3_config(tmp_path, bucket.url, {}) + with ( + owned_proxy_process(gateway, tmp_path, {"DEFAULT_S3_FLUSH_INTERVAL_SECONDS": "3"}, config=config) as owned, + owned.gateway.scenario() as scenario, + ): + model: Final = scenario.model(api_base=provider.url + "/v1", api_key="synthetic-provider-key") + key: Final = scenario.key(models=[model]) + sink.fail_until = time.time() + 10 + ids: Final = _burst(owned.gateway, model, key, marker) + payloads: Final = collect_payloads(sink, REQUESTS, seconds=90) + eventually( + lambda: owned.log.read_text(), + lambda text: "S3BatchUploadError" in text, + seconds=15, + ) + readiness: Final = owned.gateway.client.get("/health/readiness") + assert readiness.status_code == 200, readiness.text + assert sum(1 for r in provider.drain() if r.method == "POST") == REQUESTS + assert len(sink.objects()) == REQUESTS + assert frozenset(payload["id"] for payload in payloads) == ids + + +@pytest.mark.covers("other.observability.s3_v2.batch_retry_resends_identical_key_and_body") +def test_s3_v2_batch_retry_resends_identical_key_and_body(gateway: Gateway, tmp_path: Path) -> None: + marker: Final = "s3retry" + uuid.uuid4().hex[:8] + sink: Final = RecordingS3Sink(fail_status=500, delay_seconds=0.2) + with wire_server(_chat_reply) as provider, wire_server(sink.respond) as bucket: + config: Final = _s3_config(tmp_path, bucket.url, {"s3_batch_file_upload": True}) + with ( + owned_proxy(gateway, tmp_path, {"DEFAULT_S3_FLUSH_INTERVAL_SECONDS": "3"}, config=config) as candidate, + candidate.scenario() as scenario, + ): + model: Final = scenario.model(api_base=provider.url + "/v1", api_key="synthetic-provider-key") + key: Final = scenario.key(models=[model]) + sink.fail_until = time.time() + 8 + ids: Final = _burst(candidate, model, key, marker) + payloads: Final = collect_payloads(sink, REQUESTS, seconds=90) + puts: Final = bucket.drain() + assert sum(1 for r in provider.drain() if r.method == "POST") == REQUESTS + by_target: Final = {} + for put in puts: + by_target.setdefault(put.target, set()).add(put.body) # mutable-ok: grouping attempts seen so far per target + assert all(len(bodies) == 1 for bodies in by_target.values()), "a retried batch PUT changed key or body" + assert max(sum(1 for put in puts if put.target == target) for target in by_target) >= 2, "no retried PUT observed" + assert frozenset(payload["id"] for payload in payloads) == ids + assert len(payloads) == REQUESTS + + +@pytest.mark.covers("other.observability.s3_v2.unknown_model_rejection_keeps_other_requests_logging") +def test_s3_v2_unknown_model_rejection_keeps_other_requests_logging(gateway: Gateway, tmp_path: Path) -> None: + marker: Final = "s3ghost" + uuid.uuid4().hex[:8] + sink: Final = RecordingS3Sink() + with wire_server(_chat_reply) as provider, wire_server(sink.respond) as bucket: + config: Final = _s3_config(tmp_path, bucket.url, {}) + with ( + owned_proxy(gateway, tmp_path, {"DEFAULT_S3_FLUSH_INTERVAL_SECONDS": "3"}, config=config) as candidate, + candidate.scenario() as scenario, + ): + model: Final = scenario.model(api_base=provider.url + "/v1", api_key="synthetic-provider-key") + key: Final = scenario.key(models=[model]) + ghost: Final = candidate.request( + "POST", + "/v1/chat/completions", + {"model": f"ghost-{uuid.uuid4().hex}", "messages": [{"role": "user", "content": "hi"}]}, + key=key, + ) + assert ghost.status_code in (400, 403, 404), ghost.text + ids: Final = _burst(candidate, model, key, marker) + eventually( + lambda: frozenset(payload["id"] for payload in sink.payloads()), + lambda landed: ids <= landed, + seconds=90, + ) + payloads: Final = sink.payloads() + assert sum(1 for r in provider.drain() if r.method == "POST") == REQUESTS + assert ids <= frozenset(payload["id"] for payload in payloads) + extras: Final = tuple(payload for payload in payloads if payload["id"] not in ids) + assert all(payload["status"] == "failure" for payload in extras), extras + + +@pytest.mark.covers("other.observability.s3_v2.batch_flag_ignored_when_s3_v2_is_cold_storage_logger") +def test_s3_v2_batch_flag_ignored_when_s3_v2_is_cold_storage_logger(gateway: Gateway, tmp_path: Path) -> None: + marker: Final = "s3cold" + uuid.uuid4().hex[:8] + sink: Final = RecordingS3Sink() + with wire_server(_chat_reply) as provider, wire_server(sink.respond) as bucket: + config: Final = _recording_s3_config( + tmp_path, + bucket.url, + {"s3_batch_file_upload": True}, + {"cold_storage_custom_logger": "s3_v2"}, + ) + with ( + owned_proxy_process(gateway, tmp_path, {"DEFAULT_S3_FLUSH_INTERVAL_SECONDS": "3"}, config=config) as owned, + owned.gateway.scenario() as scenario, + ): + model: Final = scenario.model(api_base=provider.url + "/v1", api_key="synthetic-provider-key") + key: Final = scenario.key(models=[model]) + response: Final = owned.gateway.request( + "POST", + "/v1/chat/completions", + {"model": model, "messages": [{"role": "user", "content": marker}], "cache": {"no-cache": True}}, + key=key, + ) + assert response.status_code == 200, response.text + request_id: Final = str(response.json()["id"]) + payloads: Final = collect_payloads(sink, 1) + assert all(PER_REQUEST_KEY.match(target) for target in sink.objects()), list(sink.objects()) + eventually( + lambda: owned.log.read_text(), + lambda text: "s3_batch_file_upload is ignored because s3_v2 is the cold storage logger" in text, + seconds=15, + ) + spend: Final = eventually( + lambda: owned.gateway.request("GET", f"/spend/logs/ui/{request_id}"), + lambda reply: reply.status_code == 200 and bool((reply.json() or {}).get("messages")), + seconds=60, + ) + assert spend.status_code == 200, spend.text + body: Final = spend.json() + assert body["messages"], spend.text + assert body["response"], spend.text + assert payloads[0]["id"] == request_id + + +@pytest.mark.covers("other.observability.s3_v2.identical_requests_land_distinct_objects") +def test_s3_v2_identical_requests_land_distinct_objects(gateway: Gateway, tmp_path: Path) -> None: + marker: Final = "s3same" + uuid.uuid4().hex[:8] + sink: Final = RecordingS3Sink() + with wire_server(_chat_reply) as provider, wire_server(sink.respond) as bucket: + config: Final = _s3_config(tmp_path, bucket.url, {}) + with ( + owned_proxy(gateway, tmp_path, {"DEFAULT_S3_FLUSH_INTERVAL_SECONDS": "3"}, config=config) as candidate, + candidate.scenario() as scenario, + ): + model: Final = scenario.model(api_base=provider.url + "/v1", api_key="synthetic-provider-key") + key: Final = scenario.key(models=[model]) + + def send(_: int) -> str: + response: Final = candidate.request( + "POST", + "/v1/chat/completions", + {"model": model, "messages": [{"role": "user", "content": marker}], "cache": {"no-cache": True}}, + key=key, + ) + assert response.status_code == 200, response.text + return str(response.json()["id"]) + + with ThreadPoolExecutor(max_workers=16) as pool: + returned: Final = frozenset(pool.map(send, range(16))) + payloads: Final = collect_payloads(sink, 16) + assert sum(1 for r in provider.drain() if r.method == "POST") == 16 + assert returned == {marker}, "the upstream echo keeps the same id for identical requests" + assert len(sink.objects()) == 16, "identical requests must still land as distinct objects" + assert all(payload["id"] == marker for payload in payloads) + + +@pytest.mark.covers("other.observability.s3_v2.two_workers_bound_and_deliver_every_id") +def test_s3_v2_two_workers_bound_and_deliver_every_id(gateway: Gateway, tmp_path: Path) -> None: + marker: Final = "s3work" + uuid.uuid4().hex[:8] + sink: Final = RecordingS3Sink() + with wire_server(_chat_reply) as provider, wire_server(sink.respond) as bucket: + config: Final = _s3_config(tmp_path, bucket.url, {}) + with ( + owned_proxy( + gateway, tmp_path, {"DEFAULT_S3_FLUSH_INTERVAL_SECONDS": "3"}, config=config, workers=2 + ) as candidate, + candidate.scenario() as scenario, + ): + model: Final = scenario.model(api_base=provider.url + "/v1", api_key="synthetic-provider-key") + key: Final = scenario.key(models=[model]) + ids: Final = _burst(candidate, model, key, marker) + payloads: Final = collect_payloads(sink, REQUESTS) + assert sum(1 for r in provider.drain() if r.method == "POST") == REQUESTS + assert sink.peak <= 32, f"peak concurrent PUTs {sink.peak} exceeded two workers at the default bound" + assert len(sink.objects()) == REQUESTS + assert frozenset(payload["id"] for payload in payloads) == ids + + +@pytest.mark.covers("other.observability.s3_v2.slow_sink_never_duplicates_or_stalls_readiness") +def test_s3_v2_slow_sink_never_duplicates_or_stalls_readiness(gateway: Gateway, tmp_path: Path) -> None: + marker: Final = "s3slow" + uuid.uuid4().hex[:8] + sink: Final = RecordingS3Sink(delay_seconds=1.5) + with wire_server(_chat_reply) as provider, wire_server(sink.respond) as bucket: + config: Final = _s3_config(tmp_path, bucket.url, {"s3_batch_file_upload": True}) + with ( + owned_proxy(gateway, tmp_path, {"DEFAULT_S3_FLUSH_INTERVAL_SECONDS": "1"}, config=config) as candidate, + candidate.scenario() as scenario, + ): + model: Final = scenario.model(api_base=provider.url + "/v1", api_key="synthetic-provider-key") + key: Final = scenario.key(models=[model]) + ids: Final = _burst(candidate, model, key, marker) + + def delivered() -> int: + readiness: Final = candidate.client.get("/health/readiness") + assert readiness.status_code == 200, readiness.text + return sum(len(body.splitlines()) for body in sink.objects().values()) + + eventually(delivered, lambda total: total >= REQUESTS, seconds=90) + payloads: Final = sink.payloads() + puts: Final = bucket.drain() + targets: Final = tuple(put.target for put in puts) + assert sum(1 for r in provider.drain() if r.method == "POST") == REQUESTS + assert len(set(targets)) == len(targets), "the same object was PUT more than once" + assert frozenset(payload["id"] for payload in payloads) == ids + assert len(payloads) == REQUESTS + + +@pytest.mark.covers("other.observability.s3_v2.worker_kill_mid_burst_keeps_surviving_deliveries") +def test_s3_v2_worker_kill_mid_burst_keeps_surviving_deliveries(gateway: Gateway, tmp_path: Path) -> None: + marker: Final = "s3kill" + uuid.uuid4().hex[:8] + sink: Final = RecordingS3Sink() + with wire_server(_chat_reply) as provider, wire_server(sink.respond) as bucket: + config: Final = _s3_config(tmp_path, bucket.url, {}) + with ( + owned_proxy_process( + gateway, tmp_path, {"DEFAULT_S3_FLUSH_INTERVAL_SECONDS": "3"}, config=config, workers=2 + ) as owned, + owned.gateway.scenario() as scenario, + ): + model: Final = scenario.model(api_base=provider.url + "/v1", api_key="synthetic-provider-key") + key: Final = scenario.key(models=[model]) + sent: Final = tuple(f"{marker}-{index}" for index in range(REQUESTS)) + + def send(identity: str) -> tuple[str, bool]: + try: + response: Final = owned.gateway.request( + "POST", + "/v1/chat/completions", + { + "model": model, + "messages": [{"role": "user", "content": identity}], + "cache": {"no-cache": True}, + }, + key=key, + ) + except Exception: + return identity, False + return identity, response.status_code == 200 + + with ThreadPoolExecutor(max_workers=32) as pool: + futures: Final = tuple(pool.submit(send, identity) for identity in sent) + time.sleep(0.5) + children: Final = tuple( + process for process in group_members(owned.process.pid) if process.pid != owned.process.pid + ) + assert children, "no worker children found to kill" + children[0].kill() + results: Final = tuple(future.result() for future in futures) + survivors: Final = frozenset(identity for identity, ok in results if ok) + assert survivors, "no request survived the worker kill" + readiness: Final = owned.gateway.client.get("/health/readiness") + assert readiness.status_code == 200, readiness.text + payloads: Final = collect_payloads(sink, len(survivors), seconds=90) + landed: Final = frozenset(payload["id"] for payload in payloads) + assert survivors <= landed, "an id whose response succeeded never landed" + assert landed <= frozenset(sent), "an id that was never sent landed" + + +@pytest.mark.covers("other.observability.s3_v2.sigterm_mid_burst_loses_only_inflight_without_duplicates") +def test_s3_v2_sigterm_mid_burst_loses_only_inflight_without_duplicates(gateway: Gateway, tmp_path: Path) -> None: + marker: Final = "s3term" + uuid.uuid4().hex[:8] + sink: Final = RecordingS3Sink() + with wire_server(_chat_reply) as provider, wire_server(sink.respond) as bucket: + config: Final = _s3_config(tmp_path, bucket.url, {}) + owned: Final = owned_proxy_process(gateway, tmp_path, {"DEFAULT_S3_FLUSH_INTERVAL_SECONDS": "3"}, config=config) + candidate_owned: Final = owned.__enter__() + try: + created: Final = candidate_owned.gateway.post( + "/model/new", + { + "model_name": f"integration-{marker}", + "litellm_params": { + "model": "openai/gpt-4o-mini", + "api_key": "synthetic-provider-key", + "api_base": provider.url + "/v1", + }, + "model_info": {}, + }, + ) + model: Final = str(created["model_name"]) + key: Final = str(candidate_owned.gateway.post("/key/generate", {"models": [model]})["key"]) + sent: Final = tuple(f"{marker}-{index}" for index in range(REQUESTS)) + + def send(identity: str) -> tuple[str, bool]: + try: + response: Final = candidate_owned.gateway.request( + "POST", + "/v1/chat/completions", + { + "model": model, + "messages": [{"role": "user", "content": identity}], + "cache": {"no-cache": True}, + }, + key=key, + ) + except Exception: + return identity, False + return identity, response.status_code == 200 + + with ThreadPoolExecutor(max_workers=32) as pool: + futures: Final = tuple(pool.submit(send, identity) for identity in sent) + time.sleep(0.5) + candidate_owned.process.terminate() + results: Final = tuple(future.result() for future in futures) + candidate_owned.process.wait(timeout=30) + finally: + owned.__exit__(None, None, None) + answered: Final = frozenset(identity for identity, ok in results if ok) + landed: Final = frozenset(payload["id"] for payload in sink.payloads()) + assert landed <= answered, ( + "a delivered object has no matching answered request; lost in-flight ids are expected, extras are not" + ) + targets: Final = tuple(sink.objects()) + assert len(set(targets)) == len(targets), "the same object was PUT more than once" diff --git a/tests/test_litellm/integrations/test_s3_v2.py b/tests/test_litellm/integrations/test_s3_v2.py index a9d13038180..c67eaa45112 100644 --- a/tests/test_litellm/integrations/test_s3_v2.py +++ b/tests/test_litellm/integrations/test_s3_v2.py @@ -2468,3 +2468,453 @@ def test_prompts_only_toggle_is_exposed_to_admin_ui_for_both_s3_callbacks(callba from litellm.integrations.custom_logger import CustomLogger assert "S3_LOG_PROMPTS_ONLY" in CustomLogger.get_callback_env_vars(callback_name) + + +def _element(payload: dict[str, object], key_suffix: str) -> s3BatchLoggingElement: + return s3BatchLoggingElement( + s3_object_key=f"2025-09-14/test-{key_suffix}.json", + payload=payload, + s3_object_download_filename=f"test-{key_suffix}.json", + ) + + +def _ok_response() -> MagicMock: + response = MagicMock() + response.status_code = 200 + response.raise_for_status = MagicMock() + return response + + +class _CountingPut: + def __init__(self) -> None: + self.in_flight = 0 + self.peak = 0 + self.calls = 0 + + async def __call__(self, url: str, data: str | None = None, headers: dict[str, str] | None = None) -> MagicMock: + self.in_flight += 1 + self.peak = max(self.peak, self.in_flight) + self.calls += 1 + await asyncio.sleep(0.01) + self.in_flight -= 1 + return _ok_response() + + +class _RecordingPut: + def __init__(self) -> None: + self.calls: tuple[tuple[str, str | None, dict[str, str] | None], ...] = () + + async def __call__(self, url: str, data: str | None = None, headers: dict[str, str] | None = None) -> MagicMock: + self.calls = (*self.calls, (url, data, headers)) + return _ok_response() + + +class _LateAppendingPut: + def __init__(self, logger: S3Logger, element: s3BatchLoggingElement, fail_first: bool = False) -> None: + self.logger = logger + self.element = element + self.fail_first = fail_first + self.appended = False + + async def __call__(self, url: str, data: str | None = None, headers: dict[str, str] | None = None) -> MagicMock: + if not self.appended: + self.appended = True + self.logger.log_queue.append(self.element) + if self.fail_first: + return _failure_response() + return _ok_response() + + +class _FailOnSuffixPut: + def __init__(self, suffixes: tuple[str, ...]) -> None: + self.failing = True + self.suffixes = suffixes + + async def __call__(self, url: str, data: str | None = None, headers: dict[str, str] | None = None) -> MagicMock: + if self.failing and url.endswith(self.suffixes): + return _failure_response() + return _ok_response() + + +class _FailUntilClearedPut: + def __init__(self) -> None: + self.failing = True + self.calls: tuple[tuple[str, str | None], ...] = () + + async def __call__(self, url: str, data: str | None = None, headers: dict[str, str] | None = None) -> MagicMock: + self.calls = (*self.calls, (url, data)) + if self.failing: + return _failure_response() + return _ok_response() + + +@pytest.mark.asyncio +async def test_async_send_batch_bounds_concurrent_uploads() -> None: + logger = S3Logger( + s3_bucket_name="test-bucket", + s3_aws_access_key_id="test-key", + s3_aws_secret_access_key="test-secret", + s3_region_name="us-east-1", + s3_max_concurrent_uploads=4, + ) + + put = _CountingPut() + logger.async_httpx_client = AsyncMock() + logger.async_httpx_client.put = put + + logger.log_queue = [_element({"i": i}, f"{i}") for i in range(40)] + + await logger.async_send_batch() + + assert put.peak == 4 + assert put.calls == 40 + + +@pytest.mark.asyncio +async def test_async_send_batch_uploads_single_jsonl_file() -> None: + import json + + logger = S3Logger( + s3_bucket_name="test-bucket", + s3_aws_access_key_id="test-key", + s3_aws_secret_access_key="test-secret", + s3_region_name="us-east-1", + s3_batch_file_upload=True, + ) + + put = _RecordingPut() + logger.async_httpx_client = AsyncMock() + logger.async_httpx_client.put = put + + payloads = [{"id": "req-1"}, {"id": "req-2"}, {"id": "req-3"}] + logger.log_queue = [_element(payload, f"{i}") for i, payload in enumerate(payloads)] + + await logger.async_send_batch() + + assert len(put.calls) == 1 + url, data, headers = put.calls[0] + assert url.endswith(".jsonl") + assert data is not None + assert headers is not None + assert [json.loads(line) for line in data.splitlines()] == payloads + assert headers["Content-Type"] == "application/x-ndjson" + + +@pytest.mark.asyncio +async def test_flush_queue_preserves_events_added_during_upload() -> None: + logger = S3Logger( + s3_bucket_name="test-bucket", + s3_aws_access_key_id="test-key", + s3_aws_secret_access_key="test-secret", + s3_region_name="us-east-1", + ) + + late_element = _element({"id": "late"}, "late") + + logger.async_httpx_client = AsyncMock() + logger.async_httpx_client.put = _LateAppendingPut(logger, late_element) + + logger.log_queue = [_element({"id": "first"}, "first")] + + await logger.flush_queue() + + assert logger.log_queue == [late_element] + + +def _override_logger(**overrides: object) -> S3Logger: + return S3Logger( + s3_bucket_name="test-bucket", + s3_aws_access_key_id="test-key", + s3_aws_secret_access_key="test-secret", + s3_region_name="us-east-1", + s3_callback_params_override=overrides, + ) + + +def test_env_backed_false_string_keeps_per_request_uploads() -> None: + assert _override_logger(s3_batch_file_upload="false").s3_batch_file_upload is False + assert _override_logger(s3_batch_file_upload="true").s3_batch_file_upload is True + logger = S3Logger( + s3_bucket_name="test-bucket", + s3_aws_access_key_id="test-key", + s3_aws_secret_access_key="test-secret", + s3_region_name="us-east-1", + s3_batch_file_upload=True, + s3_callback_params_override={"s3_batch_file_upload": "false"}, + ) + assert logger.s3_batch_file_upload is True + + +@pytest.mark.parametrize("bad", [0, -3, "0", "abc", ""]) +def test_invalid_concurrency_falls_back_to_default(bad: object) -> None: + from litellm.constants import DEFAULT_S3_MAX_CONCURRENT_UPLOADS + + logger = _override_logger(s3_max_concurrent_uploads=bad) + + assert logger.s3_max_concurrent_uploads == DEFAULT_S3_MAX_CONCURRENT_UPLOADS + assert logger._upload_semaphore._value == DEFAULT_S3_MAX_CONCURRENT_UPLOADS + + +def test_env_backed_concurrency_string_is_parsed() -> None: + logger = _override_logger(s3_max_concurrent_uploads="4") + + assert logger.s3_max_concurrent_uploads == 4 + assert logger._upload_semaphore._value == 4 + + +@pytest.mark.parametrize("empty", [None, ""]) +def test_empty_config_concurrency_falls_back_to_constructor_value(empty: object) -> None: + logger = S3Logger( + s3_bucket_name="test-bucket", + s3_aws_access_key_id="test-key", + s3_aws_secret_access_key="test-secret", + s3_region_name="us-east-1", + s3_max_concurrent_uploads=4, + s3_callback_params_override={"s3_max_concurrent_uploads": empty}, + ) + + assert logger.s3_max_concurrent_uploads == 4 + assert logger._upload_semaphore._value == 4 + + +def _failure_response() -> MagicMock: + response = MagicMock() + response.status_code = 400 + response.raise_for_status = MagicMock(side_effect=Exception("s3 rejected the object")) + return response + + +@pytest.mark.asyncio +async def test_failed_uploads_stay_queued_for_next_flush() -> None: + logger = S3Logger( + s3_bucket_name="test-bucket", + s3_aws_access_key_id="test-key", + s3_aws_secret_access_key="test-secret", + s3_region_name="us-east-1", + ) + + elements = [_element({"i": i}, f"{i}") for i in range(5)] + put = _FailOnSuffixPut(("test-2.json", "test-4.json")) + + logger.async_httpx_client = AsyncMock() + logger.async_httpx_client.put = put + logger.log_queue = list(elements) + + await logger.flush_queue() + + assert logger.log_queue == [elements[2], elements[4]] + + put.failing = False + await logger.flush_queue() + + assert logger.log_queue == [] + + +@pytest.mark.asyncio +async def test_batch_file_upload_failure_keeps_whole_batch() -> None: + logger = S3Logger( + s3_bucket_name="test-bucket", + s3_aws_access_key_id="test-key", + s3_aws_secret_access_key="test-secret", + s3_region_name="us-east-1", + s3_batch_file_upload=True, + ) + + put = _FailUntilClearedPut() + + logger.async_httpx_client = AsyncMock() + logger.async_httpx_client.put = put + + elements = [_element({"i": i}, f"{i}") for i in range(3)] + logger.log_queue = list(elements) + + await logger.flush_queue() + + assert len(put.calls) == 1 + assert len(logger.log_queue) == 1 + assert logger.log_queue[0].body == "\n".join(json.dumps(element.payload) for element in elements) + + +@pytest.mark.asyncio +async def test_events_appended_during_failed_flush_survive() -> None: + logger = S3Logger( + s3_bucket_name="test-bucket", + s3_aws_access_key_id="test-key", + s3_aws_secret_access_key="test-secret", + s3_region_name="us-east-1", + ) + + late = _element({"id": "late"}, "late") + + logger.async_httpx_client = AsyncMock() + logger.async_httpx_client.put = _LateAppendingPut(logger, late, fail_first=True) + + first = _element({"id": "first"}, "first") + logger.log_queue = [first] + + await logger.flush_queue() + + assert logger.log_queue == [first, late] + + +@pytest.mark.asyncio +async def test_batch_file_key_shape() -> None: + logger = S3Logger( + s3_bucket_name="test-bucket", + s3_aws_access_key_id="test-key", + s3_aws_secret_access_key="test-secret", + s3_region_name="us-east-1", + s3_path="logs", + s3_batch_file_upload=True, + ) + + put = _RecordingPut() + + logger.async_httpx_client = AsyncMock() + logger.async_httpx_client.put = put + logger.log_queue = [_element({"id": "req-1"}, "0")] + + await logger.async_send_batch() + + ((url, _data, headers),) = put.calls + assert headers is not None + assert re.search(r".*/2025-09-14/batch_\d{2}-\d{2}-\d{2}_[0-9a-f]{32}\.jsonl$", url) + assert headers["Content-Disposition"].endswith('.jsonl"') + + +@pytest.mark.asyncio +async def test_batch_file_groups_raw_elements_by_key_parent() -> None: + logger = S3Logger( + s3_bucket_name="test-bucket", + s3_aws_access_key_id="test-key", + s3_aws_secret_access_key="test-secret", + s3_region_name="us-east-1", + s3_batch_file_upload=True, + ) + + put = _RecordingPut() + + logger.async_httpx_client = AsyncMock() + logger.async_httpx_client.put = put + + alpha = s3BatchLoggingElement( + s3_object_key="logs/alpha/2026-01-01/a.json", payload={"id": "a"}, s3_object_download_filename="a.json" + ) + beta = s3BatchLoggingElement( + s3_object_key="logs/beta/2026-01-01/b.json", payload={"id": "b"}, s3_object_download_filename="b.json" + ) + plain = s3BatchLoggingElement( + s3_object_key="logs/2026-01-01/c.json", payload={"id": "c"}, s3_object_download_filename="c.json" + ) + root = s3BatchLoggingElement( + s3_object_key="solo.json", payload={"id": "d"}, s3_object_download_filename="solo.json" + ) + logger.log_queue = [alpha, beta, plain, root] + + await logger.async_send_batch() + + assert len(put.calls) == 4 + by_parent = { + re.sub(r"(^|/)batch_\d{2}-\d{2}-\d{2}_[0-9a-f]{32}\.jsonl$", "", url.split(".com/", 1)[-1]): (url, data) + for url, data, _headers in put.calls + } + assert sorted(by_parent) == ["", "logs/2026-01-01", "logs/alpha/2026-01-01", "logs/beta/2026-01-01"] + assert [line for line in by_parent[""][1].splitlines()] == [json.dumps({"id": "d"})] + assert [line for line in by_parent["logs/alpha/2026-01-01"][1].splitlines()] == [json.dumps({"id": "a"})] + assert [line for line in by_parent["logs/beta/2026-01-01"][1].splitlines()] == [json.dumps({"id": "b"})] + assert [line for line in by_parent["logs/2026-01-01"][1].splitlines()] == [json.dumps({"id": "c"})] + + +@pytest.mark.asyncio +async def test_failed_batch_file_is_requeued_and_resent_unchanged() -> None: + logger = S3Logger( + s3_bucket_name="test-bucket", + s3_aws_access_key_id="test-key", + s3_aws_secret_access_key="test-secret", + s3_region_name="us-east-1", + s3_batch_file_upload=True, + ) + + put = _FailUntilClearedPut() + + logger.async_httpx_client = AsyncMock() + logger.async_httpx_client.put = put + logger.log_queue = [_element({"i": i}, f"{i}") for i in range(3)] + + await logger.flush_queue() + + assert len(logger.log_queue) == 1 + assert logger.log_queue[0].body is not None + assert logger.log_queue[0].s3_object_key.endswith(".jsonl") + + put.failing = False + await logger.flush_queue() + + assert logger.log_queue == [] + assert len(put.calls) == 2 + assert put.calls[0] == put.calls[1] + + +@pytest.mark.asyncio +async def test_elements_appended_after_failed_batch_file_get_their_own_file() -> None: + logger = S3Logger( + s3_bucket_name="test-bucket", + s3_aws_access_key_id="test-key", + s3_aws_secret_access_key="test-secret", + s3_region_name="us-east-1", + s3_batch_file_upload=True, + ) + + put = _FailUntilClearedPut() + + logger.async_httpx_client = AsyncMock() + logger.async_httpx_client.put = put + logger.log_queue = [_element({"id": "first"}, "first")] + + await logger.flush_queue() + + late = _element({"id": "late"}, "late") + logger.log_queue.append(late) + + put.failing = False + await logger.flush_queue() + + assert logger.log_queue == [] + assert len(put.calls) == 3 + assert put.calls[0] == put.calls[1] + assert put.calls[2][0] != put.calls[0][0] + assert put.calls[2][1] == json.dumps({"id": "late"}) + + +@pytest.mark.asyncio +async def test_batch_file_mode_disabled_when_s3_v2_is_cold_storage_logger(monkeypatch: pytest.MonkeyPatch) -> None: + logger = S3Logger( + s3_bucket_name="test-bucket", + s3_aws_access_key_id="test-key", + s3_aws_secret_access_key="test-secret", + s3_region_name="us-east-1", + s3_batch_file_upload=True, + ) + + put = _RecordingPut() + + logger.async_httpx_client = AsyncMock() + logger.async_httpx_client.put = put + + import litellm + + monkeypatch.setattr(litellm, "cold_storage_custom_logger", "s3_v2") + logger.log_queue = [_element({"id": "req-1"}, "0")] + + await logger.async_send_batch() + + assert len(put.calls) == 1 + assert put.calls[0][0].endswith("test-0.json") + + monkeypatch.setattr(litellm, "cold_storage_custom_logger", None) + logger.log_queue = [_element({"id": "req-2"}, "1")] + + await logger.async_send_batch() + + assert len(put.calls) == 2 + assert put.calls[1][0].endswith(".jsonl") From d11705a24d4aabacefc40b2f8c85a2a21c23bf5c Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 24 Sep 2026 01:04:31 -0700 Subject: [PATCH 027/120] fix(cost-map): source for bedrock mantle gpt-5.6 luna, sol, terra and grok-4.6 (#42898) Co-authored-by: kerry Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/model_prices_and_context_window_backup.json | 12 ++++++++---- model_prices_and_context_window.json | 12 ++++++++---- 2 files changed, 16 insertions(+), 8 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index b73f12bb673..36866bca3ff 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -56802,7 +56802,8 @@ "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, - "supports_web_search": true + "supports_web_search": true, + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-56-sol.html" }, "bedrock_mantle/openai.gpt-5.6-terra": { "input_cost_per_token": 2.2e-06, @@ -56844,7 +56845,8 @@ "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, - "supports_web_search": true + "supports_web_search": true, + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-56-terra.html" }, "bedrock_mantle/openai.gpt-5.6-cyber": { "input_cost_per_token": 1.375e-05, @@ -56951,7 +56953,8 @@ "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, - "supports_web_search": true + "supports_web_search": true, + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-56-luna.html" }, "us.openai.gpt-5.6-sol": { "input_cost_per_token": 4.4e-06, @@ -57722,7 +57725,8 @@ "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-xai-grok-4-6.html" }, "bedrock_mantle/anthropic.claude-haiku-4-5": { "cache_creation_input_token_cost": 1.25e-06, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index b73f12bb673..36866bca3ff 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -56802,7 +56802,8 @@ "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, - "supports_web_search": true + "supports_web_search": true, + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-56-sol.html" }, "bedrock_mantle/openai.gpt-5.6-terra": { "input_cost_per_token": 2.2e-06, @@ -56844,7 +56845,8 @@ "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, - "supports_web_search": true + "supports_web_search": true, + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-56-terra.html" }, "bedrock_mantle/openai.gpt-5.6-cyber": { "input_cost_per_token": 1.375e-05, @@ -56951,7 +56953,8 @@ "supports_tool_choice": true, "supports_vision": true, "supports_xhigh_reasoning_effort": true, - "supports_web_search": true + "supports_web_search": true, + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-56-luna.html" }, "us.openai.gpt-5.6-sol": { "input_cost_per_token": 4.4e-06, @@ -57722,7 +57725,8 @@ "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-xai-grok-4-6.html" }, "bedrock_mantle/anthropic.claude-haiku-4-5": { "cache_creation_input_token_cost": 1.25e-06, From 265f874eaf3f1de13a77f819546555e11e378212 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Thu, 24 Sep 2026 01:05:55 -0700 Subject: [PATCH 028/120] test(integration): edge-case matrices for malformed token limits and callback_settings shapes (#42895) * test(integration): edge-case matrices for malformed token limits and callback_settings shapes Extends the integration suite so two classes of issues found by gauntlet reviews are caught end to end against a real proxy: - non-numeric or odd model_info token limits (from /model/new and from config YAML) must be listed as absent on /v1/models, /models, /v1/models/{id} and /model/info, keep sibling models listed, and still serve chat - every callback_settings shape (top level and per consumer) must let the proxy boot, register the configured callbacks and serve chat Four product bugs on main surfaced by the matrices are recorded as BUG skips per the suite convention: chat 500 and /model_group/info 500 on non-numeric token limits, a startup crash on a non-object callback_settings, and otel silently dropped on a non-object callback_settings.otel * test(integration): pin the exact coerced value for numeric-edge token limits Addresses review feedback: the numeric-edge matrix only asserted 'int or absent'. It now asserts the listed value for each case on /v1/models, /models and /v1/models/{id}, which also lets the listing helper drop its optional-expectation branch. --- .../test_callback_settings_boot.py | 154 +++++++++++++ .../test_model_listing_token_limits.py | 215 ++++++++++++++++++ 2 files changed, 369 insertions(+) create mode 100644 tests/integration/configuration/test_callback_settings_boot.py diff --git a/tests/integration/configuration/test_callback_settings_boot.py b/tests/integration/configuration/test_callback_settings_boot.py new file mode 100644 index 00000000000..53e81b74572 --- /dev/null +++ b/tests/integration/configuration/test_callback_settings_boot.py @@ -0,0 +1,154 @@ +import json +import uuid +from collections.abc import Mapping +from pathlib import Path +from typing import Final + +import pytest +from pydantic import JsonValue + +from tests.integration._support.client import Gateway +from tests.integration._support.process import owned_proxy + +SERVING_CONSUMERS: Final = { + "compression_interception": "CompressionInterceptionLogger", + "code_interpreter_interception": "CodeInterpreterInterceptionLogger", + "websearch_interception": "WebSearchInterceptionLogger", +} +OTEL_CONSUMER: Final = {"otel": "OpenTelemetry"} +GUARDRAIL_CONSUMERS: Final = { + "presidio": "_OPTIONAL_PresidioPIIMasking", + "lakera_prompt_injection": "lakeraAI_Moderation", +} + +TOP_LEVEL_SHAPES: Final = ( + pytest.param({}, id="empty-object"), + pytest.param(None, id="null"), + pytest.param("otel", id="string"), + pytest.param(["otel"], id="list"), + pytest.param(True, id="bool"), + pytest.param(0, id="zero"), +) + +CONSUMER_SHAPES: Final = ( + pytest.param({}, id="empty-object"), + pytest.param(None, id="null"), + pytest.param("on", id="string"), + pytest.param(True, id="bool"), + pytest.param([], id="empty-list"), + pytest.param(["on"], id="list"), + pytest.param(7, id="int"), +) + +TOP_LEVEL_BOOT_CRASH: Final = ( + "BUG: a non-object callback_settings is stored verbatim and proxy startup crashes calling .get on it" +) +TOP_LEVEL_BOOT_CRASH_IDS: Final = frozenset({"string", "list", "bool"}) +OTEL_DROPPED: Final = ( + "BUG: a non-object callback_settings.otel fails dict() and the otel callback is silently not registered" +) +OTEL_DROPPED_IDS: Final = frozenset({"null", "string", "bool", "int"}) + + +def _write_config( + directory: Path, upstream_url: str, model: str, callbacks: tuple[str, ...], callback_settings: JsonValue +) -> Path: + config: Final = directory / f"callback_settings_{uuid.uuid4().hex}.yaml" + config.write_text( + json.dumps( + { + "model_list": [ + { + "model_name": model, + "litellm_params": { + "model": f"openai/{model}", + "api_base": f"{upstream_url}/v1", + "api_key": "integration-provider-key", + }, + } + ], + "litellm_settings": {"callbacks": list(callbacks)}, + "callback_settings": callback_settings, + "general_settings": { + "master_key": "os.environ/LITELLM_MASTER_KEY", + "database_url": "os.environ/DATABASE_URL", + }, + } + ) + ) + return config + + +def _assert_registered(candidate: Gateway, consumers: Mapping[str, str]) -> None: + response: Final = candidate.request("GET", "/active/callbacks") + assert response.status_code == 200, response.text + missing: Final = sorted(name for name, class_name in consumers.items() if class_name not in response.text) + assert missing == [], response.text + + +@pytest.mark.parametrize("callback_settings", TOP_LEVEL_SHAPES) +def test_top_level_callback_settings_shape_boots_registers_and_serves_chat( + gateway: Gateway, tmp_path: Path, callback_settings: JsonValue, request: pytest.FixtureRequest +) -> None: + if request.node.callspec.id in TOP_LEVEL_BOOT_CRASH_IDS: + pytest.skip(TOP_LEVEL_BOOT_CRASH) + consumers: Final = {**SERVING_CONSUMERS, **OTEL_CONSUMER} + model: Final = f"integration-callback-settings-{uuid.uuid4().hex}" + config: Final = _write_config(tmp_path, gateway.upstream_url, model, tuple(consumers), callback_settings) + with owned_proxy(gateway, tmp_path, {"STORE_MODEL_IN_DB": "False"}, config=config) as candidate: + _assert_registered(candidate, consumers) + reply: Final = candidate.chat(model, text=f"callback settings {uuid.uuid4().hex}") + assert reply["model"] == model, reply + + +@pytest.mark.parametrize("value", CONSUMER_SHAPES) +def test_serving_consumer_settings_shape_boots_registers_and_serves_chat( + gateway: Gateway, tmp_path: Path, value: JsonValue +) -> None: + model: Final = f"integration-callback-settings-{uuid.uuid4().hex}" + config: Final = _write_config( + tmp_path, + gateway.upstream_url, + model, + tuple(SERVING_CONSUMERS), + {consumer: value for consumer in SERVING_CONSUMERS}, + ) + with owned_proxy(gateway, tmp_path, {"STORE_MODEL_IN_DB": "False"}, config=config) as candidate: + _assert_registered(candidate, SERVING_CONSUMERS) + reply: Final = candidate.chat(model, text=f"callback settings {uuid.uuid4().hex}") + assert reply["model"] == model, reply + + +@pytest.mark.parametrize("value", CONSUMER_SHAPES) +def test_otel_settings_shape_boots_registers_and_serves_chat( + gateway: Gateway, tmp_path: Path, value: JsonValue, request: pytest.FixtureRequest +) -> None: + if request.node.callspec.id in OTEL_DROPPED_IDS: + pytest.skip(OTEL_DROPPED) + model: Final = f"integration-callback-settings-{uuid.uuid4().hex}" + config: Final = _write_config(tmp_path, gateway.upstream_url, model, tuple(OTEL_CONSUMER), {"otel": value}) + with owned_proxy(gateway, tmp_path, {"STORE_MODEL_IN_DB": "False"}, config=config) as candidate: + _assert_registered(candidate, OTEL_CONSUMER) + reply: Final = candidate.chat(model, text=f"callback settings {uuid.uuid4().hex}") + assert reply["model"] == model, reply + + +@pytest.mark.parametrize("value", CONSUMER_SHAPES) +def test_guardrail_consumer_settings_shape_boots_and_registers( + gateway: Gateway, tmp_path: Path, value: JsonValue +) -> None: + model: Final = f"integration-callback-settings-{uuid.uuid4().hex}" + config: Final = _write_config( + tmp_path, + gateway.upstream_url, + model, + tuple(GUARDRAIL_CONSUMERS), + {consumer: value for consumer in GUARDRAIL_CONSUMERS}, + ) + environment: Final = { + "STORE_MODEL_IN_DB": "False", + "PRESIDIO_ANALYZER_API_BASE": gateway.upstream_url, + "PRESIDIO_ANONYMIZER_API_BASE": gateway.upstream_url, + } + with owned_proxy(gateway, tmp_path, environment, config=config) as candidate: + _assert_registered(candidate, GUARDRAIL_CONSUMERS) diff --git a/tests/integration/pricing/test_model_listing_token_limits.py b/tests/integration/pricing/test_model_listing_token_limits.py index 13702748ec5..5f75b5bb38a 100644 --- a/tests/integration/pricing/test_model_listing_token_limits.py +++ b/tests/integration/pricing/test_model_listing_token_limits.py @@ -1,9 +1,93 @@ +import json import uuid +from collections.abc import Mapping +from pathlib import Path from typing import Final +import pytest from integration._support.client import Gateway, object_value +from integration._support.process import owned_proxy from pydantic import JsonValue +SIBLING_LIMITS: Final = {"max_input_tokens": 4321, "max_output_tokens": 987} + +NON_NUMERIC_LIMITS: Final = ( + pytest.param("", id="empty-string"), + pytest.param(" ", id="blank-string"), + pytest.param("128,000", id="thousands-separator"), + pytest.param("unlimited", id="word"), + pytest.param("NaN", id="nan-string"), + pytest.param("inf", id="inf-string"), + pytest.param([], id="empty-list"), + pytest.param([4096], id="list"), + pytest.param({}, id="empty-object"), + pytest.param({"tokens": 4096}, id="object"), + pytest.param(True, id="bool"), + pytest.param(None, id="null"), +) + +NUMERIC_EDGE_LIMITS: Final = ( + pytest.param(0, id="zero"), + pytest.param(-1, id="negative"), + pytest.param(1.5, id="float"), + pytest.param("1.5", id="float-string"), + pytest.param("1e9", id="exponent-string"), + pytest.param(10**12, id="huge"), +) +NUMERIC_EDGE_EXPECTED: Final = { + "zero": 0, + "negative": -1, + "float": 1, + "float-string": 1, + "exponent-string": 1_000_000_000, + "huge": 10**12, +} + +MODEL_GROUP_INFO_500: Final = ( + "BUG: /model_group/info returns 500 for every caller when one deployment's token limit is non-numeric" +) +CHAT_500: Final = ( + "BUG: chat completions return 500 from ModelGroupInfo validation when the deployment's token limit is non-numeric" +) +MODEL_GROUP_INFO_500_IDS: Final = frozenset( + {"empty-string", "blank-string", "thousands-separator", "word", "nan-string", "inf-string"} + | {"empty-list", "list", "empty-object", "object"} +) +CHAT_500_IDS: Final = frozenset( + {"empty-string", "blank-string", "thousands-separator", "word", "empty-list", "list", "empty-object", "object"} +) + + +def _listed(gateway: Gateway, path: str) -> dict[str, dict[str, JsonValue]]: + entries: Final = gateway.get(path)["data"] + assert isinstance(entries, list) + return {str(object_value(entry)["id"]): object_value(entry) for entry in entries} + + +def _limits(entry: Mapping[str, JsonValue]) -> tuple[JsonValue, JsonValue]: + return entry.get("max_input_tokens"), entry.get("max_output_tokens") + + +def _assert_listing_spares_the_sibling( + gateway: Gateway, broken: str, sibling: str, broken_limits: tuple[JsonValue, JsonValue] +) -> None: + for path in ("/v1/models", "/models"): + listed: Final = _listed(gateway, path) + assert _limits(listed[sibling]) == (4321, 987), (path, listed[sibling]) + assert _limits(listed[broken]) == broken_limits, (path, listed[broken]) + single: Final = gateway.get(f"/v1/models/{broken}") + assert single["id"] == broken, single + assert _limits(single) == broken_limits, single + registered: Final = gateway.get("/model/info")["data"] + assert isinstance(registered, list) + assert {broken, sibling} <= {str(object_value(entry)["model_name"]) for entry in registered} + + +def _assert_serves_chat(gateway: Gateway, *models: str) -> None: + for model in models: + reply: Final = gateway.chat(model, text=f"token limit edge {uuid.uuid4().hex}") + assert reply["model"] == model, reply + def _listed_model(gateway: Gateway, model: str) -> dict[str, JsonValue]: entries: Final = gateway.get("/v1/models")["data"] @@ -27,3 +111,134 @@ def test_v1_models_carries_deployment_model_info_limits_for_an_unknown_model(gat listed: Final = _listed_model(gateway, model) assert listed["max_input_tokens"] == 4321, listed assert listed["max_output_tokens"] == 987, listed + + +def test_numeric_string_token_limit_is_coerced_to_an_int(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + model: Final = scenario.model( + model=f"openai/custom-{uuid.uuid4().hex}", + model_info={"max_input_tokens": "4096", "max_output_tokens": "512"}, + ) + assert _limits(_listed_model(gateway, model)) == (4096, 512) + + +@pytest.mark.parametrize("value", NON_NUMERIC_LIMITS) +def test_non_numeric_token_limit_is_listed_as_absent_without_breaking_the_listing( + gateway: Gateway, value: JsonValue +) -> None: + with gateway.scenario() as scenario: + sibling: Final = scenario.model(model=f"openai/custom-{uuid.uuid4().hex}", model_info=SIBLING_LIMITS) + broken: Final = scenario.model( + model=f"openai/custom-{uuid.uuid4().hex}", + model_info={"max_input_tokens": value, "max_output_tokens": value}, + ) + _assert_listing_spares_the_sibling(gateway, broken, sibling, (None, None)) + + +@pytest.mark.parametrize("value", NON_NUMERIC_LIMITS) +def test_non_numeric_token_limit_still_serves_chat( + gateway: Gateway, value: JsonValue, request: pytest.FixtureRequest +) -> None: + if request.node.callspec.id in CHAT_500_IDS: + pytest.skip(CHAT_500) + with gateway.scenario() as scenario: + sibling: Final = scenario.model(model=f"openai/custom-{uuid.uuid4().hex}", model_info=SIBLING_LIMITS) + broken: Final = scenario.model( + model=f"openai/custom-{uuid.uuid4().hex}", + model_info={"max_input_tokens": value, "max_output_tokens": value}, + ) + _assert_serves_chat(gateway, broken, sibling) + + +@pytest.mark.parametrize("value", NUMERIC_EDGE_LIMITS) +def test_numeric_edge_token_limit_is_listed_as_its_integer_without_breaking_the_listing( + gateway: Gateway, value: JsonValue, request: pytest.FixtureRequest +) -> None: + expected: Final = NUMERIC_EDGE_EXPECTED[request.node.callspec.id] + with gateway.scenario() as scenario: + sibling: Final = scenario.model(model=f"openai/custom-{uuid.uuid4().hex}", model_info=SIBLING_LIMITS) + broken: Final = scenario.model( + model=f"openai/custom-{uuid.uuid4().hex}", + model_info={"max_input_tokens": value, "max_output_tokens": value}, + ) + _assert_listing_spares_the_sibling(gateway, broken, sibling, (expected, expected)) + _assert_serves_chat(gateway, broken, sibling) + + +@pytest.mark.parametrize("field", ("max_input_tokens", "max_output_tokens")) +def test_one_malformed_limit_does_not_disturb_the_other(gateway: Gateway, field: str) -> None: + other: Final = "max_output_tokens" if field == "max_input_tokens" else "max_input_tokens" + with gateway.scenario() as scenario: + model: Final = scenario.model( + model=f"openai/custom-{uuid.uuid4().hex}", model_info={field: "128,000", other: 2048} + ) + listed: Final = _listed_model(gateway, model) + assert listed.get(field) is None, listed + assert listed[other] == 2048, listed + + +@pytest.mark.parametrize("value", NON_NUMERIC_LIMITS + NUMERIC_EDGE_LIMITS) +def test_malformed_token_limit_keeps_model_group_info_serving( + gateway: Gateway, value: JsonValue, request: pytest.FixtureRequest +) -> None: + if request.node.callspec.id in MODEL_GROUP_INFO_500_IDS: + pytest.skip(MODEL_GROUP_INFO_500) + with gateway.scenario() as scenario: + sibling: Final = scenario.model(model=f"openai/custom-{uuid.uuid4().hex}", model_info=SIBLING_LIMITS) + broken: Final = scenario.model( + model=f"openai/custom-{uuid.uuid4().hex}", + model_info={"max_input_tokens": value, "max_output_tokens": value}, + ) + groups: Final = gateway.get("/model_group/info")["data"] + assert isinstance(groups, list) + assert {broken, sibling} <= {str(object_value(group)["model_group"]) for group in groups} + single: Final = gateway.get("/model_group/info", {"model_group": broken})["data"] + assert isinstance(single, list) + assert [object_value(group)["model_group"] for group in single] == [broken] + + +def _yaml_deployment(name: str, upstream_url: str, model_info: Mapping[str, JsonValue]) -> dict[str, JsonValue]: + return { + "model_name": name, + "litellm_params": { + "model": f"openai/custom-{uuid.uuid4().hex}", + "api_base": f"{upstream_url}/v1", + "api_key": "integration-provider-key", + }, + "model_info": dict(model_info), + } + + +def test_non_numeric_token_limits_in_config_yaml_are_listed_as_absent(gateway: Gateway, tmp_path: Path) -> None: + run: Final = uuid.uuid4().hex + sibling: Final = f"integration-yaml-sibling-{run}" + broken: Final = {f"integration-yaml-{parameter.id}-{run}": parameter.values[0] for parameter in NON_NUMERIC_LIMITS} + serving: Final = tuple( + f"integration-yaml-{parameter.id}-{run}" for parameter in NON_NUMERIC_LIMITS if parameter.id not in CHAT_500_IDS + ) + config: Final = tmp_path / "malformed_token_limits.yaml" + config.write_text( + json.dumps( + { + "model_list": [ + _yaml_deployment(sibling, gateway.upstream_url, SIBLING_LIMITS), + *( + _yaml_deployment( + name, gateway.upstream_url, {"max_input_tokens": value, "max_output_tokens": value} + ) + for name, value in broken.items() + ), + ], + "general_settings": { + "master_key": "os.environ/LITELLM_MASTER_KEY", + "database_url": "os.environ/DATABASE_URL", + "store_model_in_db": True, + }, + "router_settings": {"disable_cooldowns": True}, + } + ) + ) + with owned_proxy(gateway, tmp_path, {}, config=config) as candidate: + for name in broken: + _assert_listing_spares_the_sibling(candidate, name, sibling, (None, None)) + _assert_serves_chat(candidate, sibling, *serving) From 1c8a0ff6023b3eba6b01e955c49edcfa97bf0d22 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 24 Sep 2026 01:06:54 -0700 Subject: [PATCH 029/120] fix(cost-map): add azure deprecation dates for gpt-6 and gpt-realtime-whisper (#42897) * fix(cost-map): add azure deprecation dates for gpt-6 and gpt-realtime-whisper Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(cost-map): add azure deprecation dates to dated gpt-6 keys Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: kerry Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/model_prices_and_context_window_backup.json | 7 +++++++ model_prices_and_context_window.json | 7 +++++++ 2 files changed, 14 insertions(+) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 36866bca3ff..d2aec118d31 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -6159,6 +6159,7 @@ ] }, "azure/gpt-realtime-whisper": { + "deprecation_date": "2027-05-06", "input_cost_per_second": 0.0002833333333333333, "litellm_provider": "azure", "mode": "audio_transcription", @@ -8060,6 +8061,7 @@ "cache_creation_input_token_cost_above_272k_tokens": 2.5e-05, "cache_read_input_token_cost": 1e-06, "cache_read_input_token_cost_above_272k_tokens": 2e-06, + "deprecation_date": "2028-01-11", "input_cost_per_token": 1e-05, "input_cost_per_token_above_272k_tokens": 2e-05, "litellm_provider": "azure", @@ -8108,6 +8110,7 @@ "cache_creation_input_token_cost_above_272k_tokens": 2.5e-05, "cache_read_input_token_cost": 1e-06, "cache_read_input_token_cost_above_272k_tokens": 2e-06, + "deprecation_date": "2028-01-11", "input_cost_per_token": 1e-05, "input_cost_per_token_above_272k_tokens": 2e-05, "litellm_provider": "azure", @@ -8156,6 +8159,7 @@ "cache_creation_input_token_cost_above_272k_tokens": 2.5e-07, "cache_read_input_token_cost": 1e-08, "cache_read_input_token_cost_above_272k_tokens": 2e-08, + "deprecation_date": "2028-03-11", "input_cost_per_token": 1e-07, "input_cost_per_token_above_272k_tokens": 2e-07, "litellm_provider": "azure", @@ -8204,6 +8208,7 @@ "cache_creation_input_token_cost_above_272k_tokens": 2.5e-07, "cache_read_input_token_cost": 1e-08, "cache_read_input_token_cost_above_272k_tokens": 2e-08, + "deprecation_date": "2028-03-11", "input_cost_per_token": 1e-07, "input_cost_per_token_above_272k_tokens": 2e-07, "litellm_provider": "azure", @@ -8252,6 +8257,7 @@ "cache_creation_input_token_cost_above_272k_tokens": 5e-06, "cache_read_input_token_cost": 2e-07, "cache_read_input_token_cost_above_272k_tokens": 4e-07, + "deprecation_date": "2028-03-11", "input_cost_per_token": 2e-06, "input_cost_per_token_above_272k_tokens": 4e-06, "litellm_provider": "azure", @@ -8300,6 +8306,7 @@ "cache_creation_input_token_cost_above_272k_tokens": 5e-06, "cache_read_input_token_cost": 2e-07, "cache_read_input_token_cost_above_272k_tokens": 4e-07, + "deprecation_date": "2028-03-11", "input_cost_per_token": 2e-06, "input_cost_per_token_above_272k_tokens": 4e-06, "litellm_provider": "azure", diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 36866bca3ff..d2aec118d31 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -6159,6 +6159,7 @@ ] }, "azure/gpt-realtime-whisper": { + "deprecation_date": "2027-05-06", "input_cost_per_second": 0.0002833333333333333, "litellm_provider": "azure", "mode": "audio_transcription", @@ -8060,6 +8061,7 @@ "cache_creation_input_token_cost_above_272k_tokens": 2.5e-05, "cache_read_input_token_cost": 1e-06, "cache_read_input_token_cost_above_272k_tokens": 2e-06, + "deprecation_date": "2028-01-11", "input_cost_per_token": 1e-05, "input_cost_per_token_above_272k_tokens": 2e-05, "litellm_provider": "azure", @@ -8108,6 +8110,7 @@ "cache_creation_input_token_cost_above_272k_tokens": 2.5e-05, "cache_read_input_token_cost": 1e-06, "cache_read_input_token_cost_above_272k_tokens": 2e-06, + "deprecation_date": "2028-01-11", "input_cost_per_token": 1e-05, "input_cost_per_token_above_272k_tokens": 2e-05, "litellm_provider": "azure", @@ -8156,6 +8159,7 @@ "cache_creation_input_token_cost_above_272k_tokens": 2.5e-07, "cache_read_input_token_cost": 1e-08, "cache_read_input_token_cost_above_272k_tokens": 2e-08, + "deprecation_date": "2028-03-11", "input_cost_per_token": 1e-07, "input_cost_per_token_above_272k_tokens": 2e-07, "litellm_provider": "azure", @@ -8204,6 +8208,7 @@ "cache_creation_input_token_cost_above_272k_tokens": 2.5e-07, "cache_read_input_token_cost": 1e-08, "cache_read_input_token_cost_above_272k_tokens": 2e-08, + "deprecation_date": "2028-03-11", "input_cost_per_token": 1e-07, "input_cost_per_token_above_272k_tokens": 2e-07, "litellm_provider": "azure", @@ -8252,6 +8257,7 @@ "cache_creation_input_token_cost_above_272k_tokens": 5e-06, "cache_read_input_token_cost": 2e-07, "cache_read_input_token_cost_above_272k_tokens": 4e-07, + "deprecation_date": "2028-03-11", "input_cost_per_token": 2e-06, "input_cost_per_token_above_272k_tokens": 4e-06, "litellm_provider": "azure", @@ -8300,6 +8306,7 @@ "cache_creation_input_token_cost_above_272k_tokens": 5e-06, "cache_read_input_token_cost": 2e-07, "cache_read_input_token_cost_above_272k_tokens": 4e-07, + "deprecation_date": "2028-03-11", "input_cost_per_token": 2e-06, "input_cost_per_token_above_272k_tokens": 4e-06, "litellm_provider": "azure", From 04f3ade124476fa1eba61e99a7c8e34ae0b6cc0e Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 24 Sep 2026 02:07:12 -0700 Subject: [PATCH 030/120] feat(cost-map): add wandb DeepSeek-V4.1-Flash and gemma-4-26B-A4B-it (#42924) * feat(cost-map): add wandb DeepSeek-V4.1-Flash and gemma-4-26B-A4B-it Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(cost-map): mark wandb gemma-4-26B-A4B-it as reasoning capable Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: kerry Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- ...odel_prices_and_context_window_backup.json | 28 +++++++++++++++++++ model_prices_and_context_window.json | 28 +++++++++++++++++++ 2 files changed, 56 insertions(+) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index d2aec118d31..5ebc0416aff 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -70438,6 +70438,34 @@ "output_cost_per_token": 0.0, "source": "https://docs.typesafe.ai/models" }, + "wandb/deepseek-ai/DeepSeek-V4.1-Flash": { + "cache_read_input_token_cost": 3e-08, + "input_cost_per_token": 2e-07, + "litellm_provider": "wandb", + "max_input_tokens": 1049000, + "max_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 6.5e-07, + "source": "https://wandb.ai/site/pricing/tokens/", + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_vision": true + }, + "wandb/google/gemma-4-26B-A4B-it": { + "cache_read_input_token_cost": 5e-08, + "input_cost_per_token": 1e-07, + "litellm_provider": "wandb", + "max_input_tokens": 262000, + "max_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 3e-07, + "source": "https://wandb.ai/site/pricing/tokens/", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true + }, "wandb/zai-org/GLM-5.3-Flash": { "cache_read_input_token_cost": 5e-08, "input_cost_per_token": 1.5e-07, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index d2aec118d31..5ebc0416aff 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -70438,6 +70438,34 @@ "output_cost_per_token": 0.0, "source": "https://docs.typesafe.ai/models" }, + "wandb/deepseek-ai/DeepSeek-V4.1-Flash": { + "cache_read_input_token_cost": 3e-08, + "input_cost_per_token": 2e-07, + "litellm_provider": "wandb", + "max_input_tokens": 1049000, + "max_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 6.5e-07, + "source": "https://wandb.ai/site/pricing/tokens/", + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_vision": true + }, + "wandb/google/gemma-4-26B-A4B-it": { + "cache_read_input_token_cost": 5e-08, + "input_cost_per_token": 1e-07, + "litellm_provider": "wandb", + "max_input_tokens": 262000, + "max_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 3e-07, + "source": "https://wandb.ai/site/pricing/tokens/", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true + }, "wandb/zai-org/GLM-5.3-Flash": { "cache_read_input_token_cost": 5e-08, "input_cost_per_token": 1.5e-07, From b550db1b4fcfa954326da18c1126e37d0c3a7bd7 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 24 Sep 2026 02:58:44 -0700 Subject: [PATCH 031/120] fix(cost-map): update openrouter kimi-k2.7-code input price (#42932) Co-authored-by: kerry Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/model_prices_and_context_window_backup.json | 2 +- model_prices_and_context_window.json | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 5ebc0416aff..af66504c341 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -65843,7 +65843,7 @@ "supports_web_search": false }, "openrouter/moonshotai/kimi-k2.7-code": { - "input_cost_per_token": 7.062e-07, + "input_cost_per_token": 6.562e-07, "output_cost_per_token": 3.3e-06, "cache_read_input_token_cost": 1.8e-07, "litellm_provider": "openrouter", diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 5ebc0416aff..af66504c341 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -65843,7 +65843,7 @@ "supports_web_search": false }, "openrouter/moonshotai/kimi-k2.7-code": { - "input_cost_per_token": 7.062e-07, + "input_cost_per_token": 6.562e-07, "output_cost_per_token": 3.3e-06, "cache_read_input_token_cost": 1.8e-07, "litellm_provider": "openrouter", From 8bbe7edb711ee894deb5ad7e3cc10391a9837de9 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 24 Sep 2026 03:03:53 -0700 Subject: [PATCH 032/120] fix(cost-map): add azure deprecation dates for regional gpt-6 rows (#42933) Co-authored-by: kerry Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/model_prices_and_context_window_backup.json | 6 ++++++ model_prices_and_context_window.json | 6 ++++++ 2 files changed, 12 insertions(+) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index af66504c341..54475f38364 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -8647,6 +8647,7 @@ "supports_minimal_reasoning_effort": false }, "azure/us/gpt-6-astra": { + "deprecation_date": "2028-01-11", "cache_creation_input_token_cost": 1.375e-05, "cache_creation_input_token_cost_above_272k_tokens": 2.75e-05, "cache_read_input_token_cost": 1.1e-06, @@ -8695,6 +8696,7 @@ "supports_xhigh_reasoning_effort": true }, "azure/us/gpt-6-luna": { + "deprecation_date": "2028-03-11", "cache_creation_input_token_cost": 1.375e-07, "cache_creation_input_token_cost_above_272k_tokens": 2.75e-07, "cache_read_input_token_cost": 1.1e-08, @@ -8743,6 +8745,7 @@ "supports_xhigh_reasoning_effort": true }, "azure/us/gpt-6-sol": { + "deprecation_date": "2028-03-11", "cache_creation_input_token_cost": 2.75e-06, "cache_creation_input_token_cost_above_272k_tokens": 5.5e-06, "cache_read_input_token_cost": 2.2e-07, @@ -68811,6 +68814,7 @@ "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" }, "azure/eu/gpt-6-astra": { + "deprecation_date": "2028-01-11", "cache_creation_input_token_cost": 1.375e-05, "cache_creation_input_token_cost_above_272k_tokens": 2.75e-05, "cache_read_input_token_cost": 1.1e-06, @@ -68824,6 +68828,7 @@ "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" }, "azure/eu/gpt-6-luna": { + "deprecation_date": "2028-03-11", "cache_creation_input_token_cost": 1.5e-07, "cache_creation_input_token_cost_above_272k_tokens": 3e-07, "cache_read_input_token_cost": 1.2e-08, @@ -68872,6 +68877,7 @@ "supports_xhigh_reasoning_effort": true }, "azure/eu/gpt-6-sol": { + "deprecation_date": "2028-03-11", "cache_creation_input_token_cost": 3e-06, "cache_creation_input_token_cost_above_272k_tokens": 6e-06, "cache_read_input_token_cost": 2.4e-07, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index af66504c341..54475f38364 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -8647,6 +8647,7 @@ "supports_minimal_reasoning_effort": false }, "azure/us/gpt-6-astra": { + "deprecation_date": "2028-01-11", "cache_creation_input_token_cost": 1.375e-05, "cache_creation_input_token_cost_above_272k_tokens": 2.75e-05, "cache_read_input_token_cost": 1.1e-06, @@ -8695,6 +8696,7 @@ "supports_xhigh_reasoning_effort": true }, "azure/us/gpt-6-luna": { + "deprecation_date": "2028-03-11", "cache_creation_input_token_cost": 1.375e-07, "cache_creation_input_token_cost_above_272k_tokens": 2.75e-07, "cache_read_input_token_cost": 1.1e-08, @@ -8743,6 +8745,7 @@ "supports_xhigh_reasoning_effort": true }, "azure/us/gpt-6-sol": { + "deprecation_date": "2028-03-11", "cache_creation_input_token_cost": 2.75e-06, "cache_creation_input_token_cost_above_272k_tokens": 5.5e-06, "cache_read_input_token_cost": 2.2e-07, @@ -68811,6 +68814,7 @@ "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" }, "azure/eu/gpt-6-astra": { + "deprecation_date": "2028-01-11", "cache_creation_input_token_cost": 1.375e-05, "cache_creation_input_token_cost_above_272k_tokens": 2.75e-05, "cache_read_input_token_cost": 1.1e-06, @@ -68824,6 +68828,7 @@ "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" }, "azure/eu/gpt-6-luna": { + "deprecation_date": "2028-03-11", "cache_creation_input_token_cost": 1.5e-07, "cache_creation_input_token_cost_above_272k_tokens": 3e-07, "cache_read_input_token_cost": 1.2e-08, @@ -68872,6 +68877,7 @@ "supports_xhigh_reasoning_effort": true }, "azure/eu/gpt-6-sol": { + "deprecation_date": "2028-03-11", "cache_creation_input_token_cost": 3e-06, "cache_creation_input_token_cost_above_272k_tokens": 6e-06, "cache_read_input_token_cost": 2.4e-07, From d17c0d77240e24f89de1003c5cf68ea331f9e258 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 24 Sep 2026 04:17:23 -0700 Subject: [PATCH 033/120] refactor: daily fresh tech debt cleanup, rolling PR (#42710) * refactor: clear fresh tech debt from the last 24 hours (2026-09-05, 2026-09-06) Drop the TID251 cast import and both cast-ok casts from the refusal message_delta rebuild by narrowing the TypedDict union on its type literal, drop the redundant Mapping cast after the isinstance check in _mapping_field, and type the Lyria predict read-only helpers as Mapping[str, object] instead of a bare dict with mutable-ok. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * refactor: clear fresh tech debt from the last 24 hours (2026-09-09) Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * refactor: clear fresh tech debt from the last 24 hours (2026-09-10) Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * refactor: keep the pre-existing cost-estimate comment and usage cost read out of the cleanup Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * refactor: clear fresh tech debt from the last 24 hours (2026-09-13) Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * refactor: keep the model info pricing helper out of the cleanup Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * refactor: drop suppressions that no longer suppress anything (2026-09-16) Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * refactor: keep the rebind-ok reason inside the line limit Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(anthropic): rebuild the refusal message_delta by spreading the chunk Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * refactor: clear fresh tech debt from the last 24 hours (2026-09-17) Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * refactor: clear fresh tech debt from the last 24 hours (2026-09-18) Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * refactor: clear fresh tech debt from the last 24 hours (2026-09-19) Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * refactor: keep the pre-existing protected-resource return type out of the cleanup Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * refactor: clear fresh tech debt from the last 24 hours (2026-09-20) Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * refactor: type fresh getattr, Any, and bare dict debt from 2026-09-22 Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(mcp): keep the string guard on tools/list next_cursor Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * refactor(vercel_ai_gateway): type the embedding error headers dict Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * chore(techdebt): fix inert suppressions and missing Final in 2026-09-22 changes Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * chore(techdebt): shorten suppression reason to fit line length Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * chore(techdebt): format provider spread so its suppression sits on the literal Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * chore(techdebt): drop the logger extras suppression that LIT013 now flags as inert Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * chore(techdebt): clear fresh suppressions, Any aliases and slop from the 24h window Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * refactor(vercel): take a read-only headers mapping in get_error_class Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: mateo Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/cost_calculator.py | 8 ++++--- litellm/experimental_mcp_client/tools.py | 2 +- litellm/integrations/arize/arize.py | 2 +- .../mavvrik_focus/mavvrik_focus_logger.py | 24 ++++--------------- .../websearch_interception/handler.py | 2 +- litellm/litellm_core_utils/core_helpers.py | 2 +- litellm/litellm_core_utils/litellm_logging.py | 4 ++-- .../llm_cost_calc/zero_cost_diagnostic.py | 11 ++++++--- .../prompt_templates/common_utils.py | 4 ++-- .../litellm_core_utils/streaming_handler.py | 2 +- litellm/llms/a2a/chat/transformation.py | 2 +- .../chat/guardrail_translation/handler.py | 2 +- .../adapters/streaming_iterator.py | 24 ++++++++----------- .../messages/utils.py | 2 +- .../responses_adapters/transformation.py | 4 ++-- .../guardrail_translation/handler.py | 2 +- .../embedding/transformation.py | 7 ++++-- .../vertex_and_google_ai_studio_gemini.py | 2 +- litellm/main.py | 4 ++-- litellm/passthrough/main.py | 4 +++- .../mcp_server/bridge_token_flow.py | 4 ++-- .../mcp_server/discoverable_endpoints.py | 4 +++- .../mcp_server/mcp_server_manager.py | 2 +- .../_experimental/mcp_server/tool_search.py | 5 ++-- litellm/proxy/auth/auth_checks.py | 2 +- litellm/proxy/common_request_processing.py | 2 +- litellm/proxy/db/baseline_accounting.py | 4 ++-- .../guardrails/guardrail_initializers.py | 2 +- .../proxy/hooks/proxy_track_cost_callback.py | 2 +- .../management_endpoints/common_utils.py | 8 +++---- .../management_helpers/bulk_user_creation.py | 4 ++-- .../management_helpers/bulk_user_deletion.py | 2 +- .../llm_passthrough_endpoints.py | 2 +- .../vertex_passthrough_logging_handler.py | 4 ++-- .../proxy/response_api_endpoints/endpoints.py | 4 ++-- .../daily_global_spend_rollup.py | 4 ---- litellm/proxy/utils.py | 1 - .../transformation.py | 2 +- litellm/responses/streaming_iterator.py | 2 +- .../complexity_router/jev_classifier.py | 4 ++-- .../encrypted_content_affinity_check.py | 6 ++--- litellm/rust_bridge/dispatch.py | 10 ++++---- litellm/rust_bridge/response_metadata.py | 4 ++-- litellm/types/llms/bedrock.py | 1 - litellm/types/utils.py | 4 ++-- 45 files changed, 96 insertions(+), 107 deletions(-) diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index 7990832dc48..6cc0d9444cd 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -2017,8 +2017,8 @@ def _deployment_model_info( return cast(ModelInfo, registered_deployment_info) # cast-ok: router registers deployment prices under its id if litellm_logging_obj is None: return None - litellm_params: Final = getattr(litellm_logging_obj, "litellm_params", None) - if litellm_params is None: + litellm_params: Final = litellm_logging_obj.litellm_params + if not litellm_params: return None return next( ( @@ -2036,7 +2036,9 @@ def _ocr_model_info( router_model_id: str | None, ) -> OCRPricing | None: deployment_info: Final = _deployment_model_info(litellm_logging_obj, custom_pricing, router_model_id) - litellm_params: Final = getattr(litellm_logging_obj, "litellm_params", None) if custom_pricing else None + litellm_params: Final = ( + litellm_logging_obj.litellm_params if custom_pricing and litellm_logging_obj is not None else None + ) if litellm_params is None: return deployment_info return _layered_ocr_pricing(litellm_params, deployment_info) diff --git a/litellm/experimental_mcp_client/tools.py b/litellm/experimental_mcp_client/tools.py index a9ee851d529..df644fd7f4a 100644 --- a/litellm/experimental_mcp_client/tools.py +++ b/litellm/experimental_mcp_client/tools.py @@ -129,7 +129,7 @@ async def list_tools_with_pagination( ) tools.extend(result.tools) - next_cursor = getattr(result, "next_cursor", None) + next_cursor = result.next_cursor if not isinstance(next_cursor, str) or not next_cursor: return tools if next_cursor in seen_cursors: diff --git a/litellm/integrations/arize/arize.py b/litellm/integrations/arize/arize.py index 4ab8d9796b2..57e60fea759 100644 --- a/litellm/integrations/arize/arize.py +++ b/litellm/integrations/arize/arize.py @@ -112,7 +112,7 @@ class ArizeLogger(OpenTelemetry): if value is None or value in ("", "None"): return None try: - rate = float(value) + rate: Final = float(value) except (TypeError, ValueError): verbose_logger.warning( "ArizeLogger: %s value %r is not a number; exporting the request", diff --git a/litellm/integrations/mavvrik_focus/mavvrik_focus_logger.py b/litellm/integrations/mavvrik_focus/mavvrik_focus_logger.py index 7e3c4cc3ce8..4c2f75bb4d7 100644 --- a/litellm/integrations/mavvrik_focus/mavvrik_focus_logger.py +++ b/litellm/integrations/mavvrik_focus/mavvrik_focus_logger.py @@ -21,7 +21,7 @@ from __future__ import annotations import os from datetime import datetime, timedelta, timezone -from typing import TYPE_CHECKING, Any, Final, Protocol +from typing import TYPE_CHECKING, Any, Final import litellm from litellm._logging import verbose_proxy_logger @@ -35,17 +35,6 @@ else: AsyncIOScheduler = Any -class _PodLockManager(Protocol): - """The subset of PodLockManager this logger drives to serialize the export across pods.""" - - @property - def redis_cache(self) -> object: ... - - async def acquire_lock(self, cronjob_id: str) -> bool | None: ... - - async def release_lock(self, cronjob_id: str) -> None: ... - - def _parse_metrics_marker( marker: object | None, ) -> datetime | None: @@ -237,13 +226,10 @@ class MavvrikFocusLogger(FocusLogger): """Scheduler entry point — uses Mavvrik-specific pod-lock key.""" from litellm.proxy.proxy_server import proxy_logging_obj # noqa: PLC0415 - pod_lock_manager: _PodLockManager | None = None - if proxy_logging_obj is not None: - writer: Final[object] = getattr(proxy_logging_obj, "db_spend_update_writer", None) - if writer is not None: - pod_lock_manager = getattr(writer, "pod_lock_manager", None) - - if pod_lock_manager and pod_lock_manager.redis_cache: + pod_lock_manager: Final = ( + proxy_logging_obj.db_spend_update_writer.pod_lock_manager if proxy_logging_obj is not None else None + ) + if pod_lock_manager is not None and pod_lock_manager.redis_cache: acquired: Final = await pod_lock_manager.acquire_lock(cronjob_id=MAVVRIK_FOCUS_EXPORT_JOB_NAME) if not acquired: verbose_proxy_logger.debug("Mavvrik FOCUS export: unable to acquire pod lock") diff --git a/litellm/integrations/websearch_interception/handler.py b/litellm/integrations/websearch_interception/handler.py index 54578323fa4..90bbd5a00d8 100644 --- a/litellm/integrations/websearch_interception/handler.py +++ b/litellm/integrations/websearch_interception/handler.py @@ -1849,7 +1849,7 @@ class WebSearchInterceptionLogger(CustomLogger): for tool_call in tool_calls: # Handle both Anthropic-style input and OpenAI-style function.arguments query = None - tool_args: dict | None = None # mutable-ok: the tool call's own arguments dict + tool_args: dict[str, object] | None = None # mutable-ok: the tool call's own arguments dict if "input" in tool_call and isinstance(tool_call["input"], dict): tool_args = tool_call["input"] query = tool_args.get("query") diff --git a/litellm/litellm_core_utils/core_helpers.py b/litellm/litellm_core_utils/core_helpers.py index a2d40279c49..3afa6a913b5 100644 --- a/litellm/litellm_core_utils/core_helpers.py +++ b/litellm/litellm_core_utils/core_helpers.py @@ -365,7 +365,7 @@ def _budget_reservation_on_auth_object(user_api_key_auth: object) -> object: return getattr(user_api_key_auth, "budget_reservation", None) -def budget_reservation_from_metadata(metadata: Mapping[str, object]) -> dict | None: +def budget_reservation_from_metadata(metadata: Mapping[str, object]) -> dict[str, object] | None: stamped: Final = metadata.get("user_api_key_budget_reservation") if isinstance(stamped, dict): return stamped diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 0603414cabd..2cecec729c2 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -5191,7 +5191,7 @@ def _maybe_construct_otel_v2(callback_name: str, _in_memory_loggers: list[Custom for callback in _in_memory_loggers: if ( isinstance(callback, OpenTelemetryV2) - and getattr(callback, "callback_name", None) == callback_name + and callback.callback_name == callback_name and (serves_a_destination or not _exports_nowhere(callback.config)) ): return callback @@ -6663,7 +6663,7 @@ def get_standard_logging_object_payload( cost_breakdown=request_cost_breakdown, autorouter_savings=autorouter_savings, autorouter_savings_estimate=( - { + { # mutable-ok: spend-log JSON serialization requires plain mappings "version": 3, "status": "unknown", "reason": "pending_projection", diff --git a/litellm/litellm_core_utils/llm_cost_calc/zero_cost_diagnostic.py b/litellm/litellm_core_utils/llm_cost_calc/zero_cost_diagnostic.py index 6331d815bdc..9688b511ea8 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/zero_cost_diagnostic.py +++ b/litellm/litellm_core_utils/llm_cost_calc/zero_cost_diagnostic.py @@ -5,7 +5,12 @@ from typing import Final from pydantic import TypeAdapter, ValidationError from typing_extensions import assert_never -from litellm.types.utils import StandardLoggingZeroCostDiagnostic, Usage +from litellm.types.utils import ( + CompletionTokensDetailsWrapper, + PromptTokensDetailsWrapper, + StandardLoggingZeroCostDiagnostic, + Usage, +) ZERO_COST_COUNTER_NAME: Final = "litellm_zero_cost_requests_total" @@ -18,8 +23,8 @@ _NESTED_PRICING: Final = TypeAdapter(Mapping[str, object] | tuple[object, ...]) _MAX_PRICING_DEPTH: Final = 4 -def _audio_tokens(details: object) -> int: - audio_tokens: Final = getattr(details, "audio_tokens", None) +def _audio_tokens(details: PromptTokensDetailsWrapper | CompletionTokensDetailsWrapper | None) -> int: + audio_tokens: Final = details.audio_tokens if details is not None else None return audio_tokens if isinstance(audio_tokens, int) and audio_tokens > 0 else 0 diff --git a/litellm/litellm_core_utils/prompt_templates/common_utils.py b/litellm/litellm_core_utils/prompt_templates/common_utils.py index 378295e1b7a..8e5d2cd0a17 100644 --- a/litellm/litellm_core_utils/prompt_templates/common_utils.py +++ b/litellm/litellm_core_utils/prompt_templates/common_utils.py @@ -2003,11 +2003,11 @@ def strip_encrypted_reasoning_from_messages(messages: object) -> None: """ if not isinstance(messages, list): return - for content in _anthropic_content_lists(cast(list[object], messages)): # cast-ok: untyped client json + for content in anthropic_content_lists(cast(list[object], messages)): # cast-ok: untyped client json _strip_encrypted_reasoning_from_blocks(content) -def _anthropic_content_lists(messages: Sequence[object]) -> Iterator[object]: +def anthropic_content_lists(messages: Sequence[object]) -> Iterator[object]: return ( cast(list[object], content) # cast-ok: narrowed by isinstance for message in messages diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index f97a274708f..fa687b585f5 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -1329,7 +1329,7 @@ class CustomStreamWrapper: "is_finished": chunk_finish_reason is not None, "finish_reason": chunk_finish_reason, "original_chunk": cached_chunk, - "tool_calls": (getattr(cached_choice.delta, "tool_calls", None) if cached_choice is not None else None), + "tool_calls": cached_choice.delta.tool_calls if cached_choice is not None else None, } completion_obj["content"] = response_obj["text"] diff --git a/litellm/llms/a2a/chat/transformation.py b/litellm/llms/a2a/chat/transformation.py index c5a71daaba5..0813e0827d2 100644 --- a/litellm/llms/a2a/chat/transformation.py +++ b/litellm/llms/a2a/chat/transformation.py @@ -48,7 +48,7 @@ def _registry_api_key(agent_litellm_params: Mapping[str, object]) -> str | None: return configured_api_key if isinstance(configured_api_key, str) else None -def _registry_headers(agent_litellm_params: Mapping[str, object]) -> dict[str, Any] | None: +def _registry_headers(agent_litellm_params: Mapping[str, object]) -> dict[str, object] | None: stored_headers: Final = agent_litellm_params.get("headers") if not isinstance(stored_headers, Mapping): return None diff --git a/litellm/llms/anthropic/chat/guardrail_translation/handler.py b/litellm/llms/anthropic/chat/guardrail_translation/handler.py index a78f633f5d7..e1c727ad235 100644 --- a/litellm/llms/anthropic/chat/guardrail_translation/handler.py +++ b/litellm/llms/anthropic/chat/guardrail_translation/handler.py @@ -685,7 +685,7 @@ class AnthropicMessagesHandler(BaseTranslation): return data - def _hoisted_top_level_system_message(self, data: dict) -> AllMessageValues | None: + def _hoisted_top_level_system_message(self, data: Mapping[str, object]) -> AllMessageValues | None: """Return the system message produced by translating the top-level prompt.""" system: Final = data.get("system") if not system: diff --git a/litellm/llms/anthropic/experimental_pass_through/adapters/streaming_iterator.py b/litellm/llms/anthropic/experimental_pass_through/adapters/streaming_iterator.py index 4486eb0985a..20753afee5c 100644 --- a/litellm/llms/anthropic/experimental_pass_through/adapters/streaming_iterator.py +++ b/litellm/llms/anthropic/experimental_pass_through/adapters/streaming_iterator.py @@ -11,7 +11,6 @@ from typing import ( Final, Literal, Protocol, - cast, # noqa: TID251 # rebuilt message_delta dict spans the ContentBlockDelta/MessageBlockDelta union get_args, ) @@ -27,6 +26,7 @@ from litellm.types.llms.anthropic import ( ContentBlockDelta, ContextManagementResponse, MessageBlockDelta, + MessageDelta, StreamingContentBlockDeltaType, UsageDelta, UsageIteration, @@ -1028,26 +1028,22 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper): self, processed_chunk: ContentBlockDelta | MessageBlockDelta, ) -> ContentBlockDelta | MessageBlockDelta: - if processed_chunk.get("type") != "message_delta" or not self._refusal_text: + if processed_chunk["type"] != "message_delta" or not self._refusal_text: return processed_chunk - delta: Final = cast(Mapping[str, object], processed_chunk["delta"]) # cast-ok: keys checked before use + delta: Final = processed_chunk["delta"] if delta.get("stop_reason") == "max_tokens": return processed_chunk from litellm.llms.anthropic.experimental_pass_through.messages.utils import ( refusal_stop_details, ) - return cast( # cast-ok: rebuilt dict matches the message_delta TypedDict shape for this branch - ContentBlockDelta | MessageBlockDelta, - { # mutable-ok: fresh translation payload; never mutated after construction - **processed_chunk, - "delta": { # mutable-ok: fresh message_delta payload; never mutated after construction - **delta, - "stop_reason": "refusal", - "stop_details": refusal_stop_details(self._refusal_text), - }, - }, - ) + refusal_delta: Final[MessageDelta] = { + **delta, + "stop_reason": "refusal", + "stop_details": refusal_stop_details(self._refusal_text), + } + refusal_chunk: Final[MessageBlockDelta] = {**processed_chunk, "delta": refusal_delta} + return refusal_chunk @staticmethod def _delta_has_content(processed_chunk: Mapping[str, object]) -> bool: diff --git a/litellm/llms/anthropic/experimental_pass_through/messages/utils.py b/litellm/llms/anthropic/experimental_pass_through/messages/utils.py index 7545dff1408..89105c00428 100644 --- a/litellm/llms/anthropic/experimental_pass_through/messages/utils.py +++ b/litellm/llms/anthropic/experimental_pass_through/messages/utils.py @@ -37,7 +37,7 @@ def _mapping_field(container: object, key: str) -> object | None: """One key of a raw provider payload, or None when the payload is not a mapping.""" if not isinstance(container, Mapping): return None - return cast(Mapping[str, object], container).get(key) # cast-ok: raw payload, callers re-check every value + return container.get(key) def _mapping_str_field(container: object, key: str) -> str | None: diff --git a/litellm/llms/anthropic/experimental_pass_through/responses_adapters/transformation.py b/litellm/llms/anthropic/experimental_pass_through/responses_adapters/transformation.py index e3d3425f8a6..6a31173a9c6 100644 --- a/litellm/llms/anthropic/experimental_pass_through/responses_adapters/transformation.py +++ b/litellm/llms/anthropic/experimental_pass_through/responses_adapters/transformation.py @@ -169,7 +169,7 @@ class LiteLLMAnthropicToResponsesAPIAdapter: cls, summary: Iterable[object], encrypted_content: object, - ) -> dict[str, Any] | None: # mutable-ok: API message payload + ) -> dict[str, object] | None: # mutable-ok: API message payload """The one Anthropic block for a Responses reasoning item. The item's encrypted reasoning rides the block's opaque field (`signature`, or @@ -198,7 +198,7 @@ class LiteLLMAnthropicToResponsesAPIAdapter: @classmethod def _assistant_group_to_input_items( cls, group: tuple[Mapping[str, object], ...] - ) -> tuple[dict[str, Any], ...]: # mutable-ok: API message payload + ) -> tuple[dict[str, object], ...]: # mutable-ok: API message payload first: Final = group[0] btype: Final = first.get("type") if btype in ("thinking", "redacted_thinking"): diff --git a/litellm/llms/openai/responses/guardrail_translation/handler.py b/litellm/llms/openai/responses/guardrail_translation/handler.py index ad8e29ad4ac..66cebe0175d 100644 --- a/litellm/llms/openai/responses/guardrail_translation/handler.py +++ b/litellm/llms/openai/responses/guardrail_translation/handler.py @@ -994,7 +994,7 @@ class OpenAIResponsesHandler(BaseTranslation): def _spread_text_rewrite_over_stream_events( self, - stream_events: Sequence[Any], + stream_events: Sequence[object], rewritten_text: str, guardrail_name: str, ) -> None: diff --git a/litellm/llms/vercel_ai_gateway/embedding/transformation.py b/litellm/llms/vercel_ai_gateway/embedding/transformation.py index fc9c6bcc19f..9cfaaab89f0 100644 --- a/litellm/llms/vercel_ai_gateway/embedding/transformation.py +++ b/litellm/llms/vercel_ai_gateway/embedding/transformation.py @@ -7,6 +7,7 @@ Vercel AI Gateway is OpenAI-compatible and supports embeddings via the /v1/embed Docs: https://vercel.com/docs/ai-gateway/openai-compat/embeddings """ +from collections.abc import Mapping from typing import TYPE_CHECKING, Any, Final import httpx @@ -161,12 +162,14 @@ class VercelAIGatewayEmbeddingConfig(BaseEmbeddingConfig): optional_params[param] = value return optional_params - def get_error_class(self, error_message: str, status_code: int, headers: Any) -> BaseLLMException: + def get_error_class( + self, error_message: str, status_code: int, headers: Mapping[str, str] | httpx.Headers + ) -> BaseLLMException: """ Get the error class for Vercel AI Gateway errors. """ return VercelAIGatewayException( message=error_message, status_code=status_code, - headers=headers, + headers=headers if isinstance(headers, httpx.Headers) else httpx.Headers(headers), ) diff --git a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py index a50665c71df..941ec4ad419 100644 --- a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py +++ b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py @@ -286,7 +286,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): """ Check if the model is Gemini 3 or newer. """ - model_name = model.split("/")[-1].lower() + model_name: Final = model.split("/")[-1].lower() is_vertex_fine_tuned_model: Final = model_name.isdigit() or ( model.startswith("gemini/") and not model_name.startswith("gemini-") ) diff --git a/litellm/main.py b/litellm/main.py index 4c40d864169..ceac729d3f0 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -1182,7 +1182,7 @@ def _is_claude_tool_target(custom_llm_provider: str | None, model: str) -> bool: return False -def _without_anthropic_only_tool_keys(tool: dict) -> dict: +def _without_anthropic_only_tool_keys(tool: dict[str, object]) -> dict[str, object]: kept: Final = {key: value for key, value in tool.items() if key not in _ANTHROPIC_ONLY_TOOL_KEYS} function: Final = tool.get("function") if not isinstance(function, dict): @@ -1193,7 +1193,7 @@ def _without_anthropic_only_tool_keys(tool: dict) -> dict: } -def _drop_anthropic_only_tool_keys(tools: list[dict] | None) -> list[dict] | None: +def _drop_anthropic_only_tool_keys(tools: list[dict[str, object]] | None) -> list[dict[str, object]] | None: if tools is None: return None return [_without_anthropic_only_tool_keys(tool) if isinstance(tool, dict) else tool for tool in tools] diff --git a/litellm/passthrough/main.py b/litellm/passthrough/main.py index ef931827d85..45ab52690bf 100644 --- a/litellm/passthrough/main.py +++ b/litellm/passthrough/main.py @@ -540,7 +540,9 @@ def llm_passthrough_route( ) ## IS STREAMING REQUEST - _streaming_request_data: dict = data if isinstance(data, dict) else (json if isinstance(json, dict) else {}) + _streaming_request_data: Final[dict[str, object]] = ( + data if isinstance(data, dict) else (json if isinstance(json, dict) else {}) + ) is_streaming_request: Final = provider_config.is_streaming_request( endpoint=endpoint, request_data=_streaming_request_data, diff --git a/litellm/proxy/_experimental/mcp_server/bridge_token_flow.py b/litellm/proxy/_experimental/mcp_server/bridge_token_flow.py index e5d11271c67..77bdbd26b35 100644 --- a/litellm/proxy/_experimental/mcp_server/bridge_token_flow.py +++ b/litellm/proxy/_experimental/mcp_server/bridge_token_flow.py @@ -86,8 +86,8 @@ async def oauth_authorization_uses_gateway_credential(request: Request) -> bool: async def _opaque_bearer_is_gateway_credential(token: str) -> bool: - from litellm.proxy._experimental.mcp_server.outbound_credentials.envelope import ( - is_envelope, # noqa: PLC0415 # envelope imports bridge types + from litellm.proxy._experimental.mcp_server.outbound_credentials.envelope import ( # noqa: PLC0415 # envelope imports bridge types + is_envelope, is_refresh_envelope, ) from litellm.proxy._types import hash_token # noqa: PLC0415 # proxy import cycle diff --git a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py index ade829a1b67..d42c1c6b879 100644 --- a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py @@ -2634,7 +2634,9 @@ def _build_aggregate_protected_resource_response(request: Request) -> dict: } -def _build_aggregate_authorization_server_response(request: Request, token_exchange_available: bool) -> dict: +def _build_aggregate_authorization_server_response( + request: Request, token_exchange_available: bool +) -> dict[str, object]: """RFC 8414 metadata for the gateway as the aggregate authorization server. The issuer is ``{base}/mcp`` and must stay equal to the value the diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 0c520142fb3..312dcb27d89 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -3490,7 +3490,7 @@ class MCPServerManager: passthrough_server_ids: Final = [ server.server_id for server in self.get_registry().values() - if getattr(server, "auth_type", None) == MCPAuth.true_passthrough + if server.auth_type == MCPAuth.true_passthrough ] combined_servers.update(passthrough_server_ids) diff --git a/litellm/proxy/_experimental/mcp_server/tool_search.py b/litellm/proxy/_experimental/mcp_server/tool_search.py index 3c060752934..71e46f8df25 100644 --- a/litellm/proxy/_experimental/mcp_server/tool_search.py +++ b/litellm/proxy/_experimental/mcp_server/tool_search.py @@ -128,14 +128,15 @@ def with_mcp_proxy_identity(tool: Tool, server_id: str) -> Tool: def _mcp_proxy_identity(tool: Tool) -> MCPProxyToolIdentity: - identity: Final = (tool.meta or {}).get(_MCP_PROXY_IDENTITY_META_KEY) # mutable-ok: absent metadata default + identity: Final = None if tool.meta is None else tool.meta.get(_MCP_PROXY_IDENTITY_META_KEY) if not isinstance(identity, Mapping): raise TypeError("MCP proxy tool identity is missing") server_id: Final = identity.get("server_id") tool_name: Final = identity.get("tool_name") if not isinstance(server_id, str) or not isinstance(tool_name, str): raise TypeError("MCP proxy tool identity is invalid") - return {"server_id": server_id, "tool_name": tool_name} # mutable-ok: TypedDict identity payload + resolved: Final[MCPProxyToolIdentity] = {"server_id": server_id, "tool_name": tool_name} + return resolved def mcp_proxy_tool_id(tool: Tool) -> str: diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index d2fed1eb421..f5e98b40ca7 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -4073,7 +4073,7 @@ async def get_org_object_for_request( ) except OrganizationNotFoundError: return None - except Exception as e: # noqa: BLE001 # only a DB outage may fail auth here, anything else degrades to no org limits + except Exception as e: if not PrismaDBExceptionHandler.is_database_service_unavailable_error_in_chain(e): verbose_proxy_logger.debug("org lookup failed, continuing without org limits", exc_info=True) return None diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 2e43c6b0d22..904070cfadd 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -2323,7 +2323,7 @@ class ProxyBaseLLMRequestProcessing: return fallbacks if isinstance(fallbacks, list) and fallbacks else None @staticmethod - def _resolve_fallback_models(model: str, fallbacks: list) -> list | None: + def _resolve_fallback_models(model: str, fallbacks: list) -> list[str] | None: from litellm.router_utils.fallback_event_handlers import get_fallback_model_group fallback_model_group, generic_fallback_idx = get_fallback_model_group( diff --git a/litellm/proxy/db/baseline_accounting.py b/litellm/proxy/db/baseline_accounting.py index 78036237993..05a4a989152 100644 --- a/litellm/proxy/db/baseline_accounting.py +++ b/litellm/proxy/db/baseline_accounting.py @@ -486,7 +486,7 @@ class BaselineAccountingStore: async def _pages( self, db: SupportsRawQueries, scope: str, after_revision: int, withdraw_from: float | None = None ) -> AsyncIterator[tuple[_StoredRecord, ...]]: - cursor: float | None = None + cursor: float | None = None # rebind-ok: keyset pagination advances after each complete timestamp group while page := _RECORDS.validate_python( tuple(await db.query_raw(_READ_PAGE, scope, after_revision, cursor, _PAGE_TIMESTAMPS, withdraw_from)) ): @@ -627,7 +627,7 @@ async def flush_baseline_accounting(client: PrismaClient) -> None: more_queued: Final = bool(client.baseline_accounting_transactions) try: remaining: Final = await asyncio.wait_for(_flush_records(store, batch), timeout=5) - except (Exception, asyncio.CancelledError) as error: # noqa: BLE001 # unknown acknowledgements can be replayed safely + except (Exception, asyncio.CancelledError) as error: async with client.baseline_accounting_lock: client.baseline_accounting_transactions.extend(batch) if isinstance(error, asyncio.CancelledError): diff --git a/litellm/proxy/guardrails/guardrail_initializers.py b/litellm/proxy/guardrails/guardrail_initializers.py index 7bafad26569..c422902d30d 100644 --- a/litellm/proxy/guardrails/guardrail_initializers.py +++ b/litellm/proxy/guardrails/guardrail_initializers.py @@ -120,7 +120,7 @@ def initialize_presidio(litellm_params: LitellmParams, guardrail: Guardrail) -> _OPTIONAL_PresidioPIIMasking, ) - explicit_filter_scope: Final = getattr(litellm_params, "presidio_filter_scope", None) + explicit_filter_scope: Final = litellm_params.presidio_filter_scope filter_scope: Final = explicit_filter_scope or ("input" if _is_mcp_only_mode(litellm_params.mode) else "both") run_input: Final = filter_scope in ("input", "both") run_output: Final = filter_scope in ("output", "both") diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index 81894a5ff12..5dae9e8bb10 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -770,7 +770,7 @@ async def _reconcile_budget_reservation_before_db_update( "Failed to invalidate budget reservation counters after pre-persist reconcile failed" ) finally: - budget_reservation["finalized"] = True # rebind-ok: the counter update reads the stamp off the shared dict + budget_reservation["finalized"] = True # rebind-ok: stamps the caller's shared dict for the counter update async def _release_budget_reservation(budget_reservation: dict | None) -> None: diff --git a/litellm/proxy/management_endpoints/common_utils.py b/litellm/proxy/management_endpoints/common_utils.py index 0bd4eb5a5d8..59c06a3f888 100644 --- a/litellm/proxy/management_endpoints/common_utils.py +++ b/litellm/proxy/management_endpoints/common_utils.py @@ -528,7 +528,7 @@ def _prisma_value(value: object) -> object: return list(value) if isinstance(value, tuple) else value -def member_budget_patch(source: BaseModel) -> dict[str, Any]: +def member_budget_patch(source: BaseModel) -> Mapping[str, object]: """Map the per-member limit fields a request actually set to their budget-table columns (merge-patch: a sent value updates, an explicit null clears, an absent field is left untouched).""" @@ -561,7 +561,7 @@ async def _upsert_budget_and_membership( user_id: str, existing_budget_id: str | None, user_api_key_dict: UserAPIKeyAuth, - budget_patch: dict[str, Any], + budget_patch: Mapping[str, object], team_default_budget_id: str | None = None, shared_budget_ids: frozenset[str] | None = None, ): @@ -624,9 +624,9 @@ async def _upsert_budget_and_membership( if is_shared_default and not temp_only else None ) - source: Final[Mapping[str, Any]] = source_row.model_dump() if source_row is not None else MappingProxyType({}) + source: Final[Mapping[str, object]] = source_row.model_dump() if source_row is not None else MappingProxyType({}) - create_data: Final[dict[str, Any]] = { # mutable-ok: Prisma create payloads are dict-shaped + create_data: Final[dict[str, object]] = { # mutable-ok: Prisma create payloads are dict-shaped "created_by": user_api_key_dict.user_id or "", "updated_by": user_api_key_dict.user_id or "", **MappingProxyType( diff --git a/litellm/proxy/management_helpers/bulk_user_creation.py b/litellm/proxy/management_helpers/bulk_user_creation.py index dd3f4ff1b12..ec8fd312766 100644 --- a/litellm/proxy/management_helpers/bulk_user_creation.py +++ b/litellm/proxy/management_helpers/bulk_user_creation.py @@ -348,7 +348,7 @@ async def _prepare_user(user: _PendingUser, prisma_client: PrismaClient) -> _Pre data: Final = {**dumped, "user_id": user.user_id} # mutable-ok: /user/new defaults helper mutates in place data_json: Final = _JSON_OBJECT.validate_python(_update_internal_new_user_params(data, user.request)) with_permission: Final = _JSON_OBJECT.validate_python( - await _set_object_permission(data_json=data_json, prisma_client=prisma_client) # pyright: ignore[reportUnknownArgumentType] # validated by the adapter + await _set_object_permission(data_json=data_json, prisma_client=prisma_client) ) return _PreparedUser(user, _USER_ROW.validate_python(with_permission)) except Exception as exc: # noqa: BLE001 # any preparation failure is reported on this row only @@ -509,7 +509,7 @@ class _TeamsData(TypedDict): def _default_member_budget_id(team: LiteLLM_TeamTable) -> str | None: metadata: Final = ( _JSON_OBJECT.validate_python( - team.metadata # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # LiteLLM_TeamTable.metadata is a bare dict; validated by the adapter + team.metadata # pyright: ignore[reportUnknownMemberType] # LiteLLM_TeamTable.metadata is a bare dict; validated by the adapter ) if team.metadata # pyright: ignore[reportUnknownMemberType] # same bare dict else None diff --git a/litellm/proxy/management_helpers/bulk_user_deletion.py b/litellm/proxy/management_helpers/bulk_user_deletion.py index 8b5b601fe8a..c7b89a6dd6c 100644 --- a/litellm/proxy/management_helpers/bulk_user_deletion.py +++ b/litellm/proxy/management_helpers/bulk_user_deletion.py @@ -204,7 +204,7 @@ def _error_message(exc: BaseException) -> str: if isinstance(exc, HTTPException) and isinstance(exc.detail, dict): return str(exc.detail.get("error", exc.detail)) # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # HTTPException.detail is untyped if isinstance(exc, HTTPException): - return str(exc.detail) # pyright: ignore[reportUnknownArgumentType] # HTTPException.detail is untyped + return str(exc.detail) return str(exc) or type(exc).__name__ diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index 4e60c318f03..eaa03b67b40 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -3902,7 +3902,7 @@ async def handle_gigachat_passthrough_router_model( is_streaming: Final = request_body.get("stream", False) # pyright: ignore[reportUnknownVariableType] # request_body is dict[Unknown, Unknown] - data: dict[str, Any] = await _read_request_body(request=request) # Any needed for proxy pipeline + data: Final[dict[str, object]] = await _read_request_body(request=request) if user_api_key_dict is not None: auth_metadata: Final = { metadata_key: value diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py index 2cdeddbea30..040250637ea 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py @@ -458,7 +458,7 @@ class VertexPassthroughLoggingHandler: @staticmethod def _is_audio_predict_response( model: str, - json_response: dict, # mutable-ok: predicate inspects the decoded provider response dictionary without mutation + json_response: Mapping[str, object], ) -> bool: return ( VertexPassthroughLoggingHandler._get_audio_prediction_count(json_response=json_response) > 0 @@ -467,7 +467,7 @@ class VertexPassthroughLoggingHandler: @staticmethod def _get_audio_prediction_count( - json_response: dict, # mutable-ok: counter inspects the decoded provider response dictionary without mutation + json_response: Mapping[str, object], ) -> int: predictions: Final = json_response.get("predictions") if not isinstance(predictions, list): diff --git a/litellm/proxy/response_api_endpoints/endpoints.py b/litellm/proxy/response_api_endpoints/endpoints.py index ca4008dba3d..75eefb2e73b 100644 --- a/litellm/proxy/response_api_endpoints/endpoints.py +++ b/litellm/proxy/response_api_endpoints/endpoints.py @@ -6,7 +6,7 @@ from collections.abc import AsyncIterator, Awaitable, Mapping, Sequence from enum import Enum from functools import partial from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final, NamedTuple, Protocol, cast, get_args +from typing import TYPE_CHECKING, Any, Final, NamedTuple, Protocol, TypeAlias, cast, get_args from uuid import uuid4 import fastapi @@ -49,7 +49,7 @@ if TYPE_CHECKING: router: Final = APIRouter() -_ResponseDocSchemas = dict[int | str, dict[str, Any]] # pyright: ignore[reportExplicitAny] # fastapi's responses kwarg +_ResponseDocSchemas: TypeAlias = dict[int | str, dict[str, object]] # fastapi's responses kwarg RESPONSES_API_RESPONSE_SCHEMAS: Final[_ResponseDocSchemas] = {200: {"model": ResponsesAPIResponse}} RESPONSES_API_CREATE_RESPONSE_SCHEMAS: Final[_ResponseDocSchemas] = { diff --git a/litellm/proxy/spend_tracking/daily_global_spend_rollup.py b/litellm/proxy/spend_tracking/daily_global_spend_rollup.py index 376b113ed02..b81b6c1943e 100644 --- a/litellm/proxy/spend_tracking/daily_global_spend_rollup.py +++ b/litellm/proxy/spend_tracking/daily_global_spend_rollup.py @@ -198,10 +198,6 @@ async def _scan_pending(prisma_client: "PrismaClient") -> _PendingScan: return _PendingScan(marker, db_now.now, tuple(_DateRow.model_validate(row).date for row in rows)) -async def pending_days(prisma_client: "PrismaClient") -> tuple[str, ...]: - return (await _scan_pending(prisma_client)).days - - async def reconcile_day(prisma_client: "PrismaClient", day: str) -> None: """Rewrite one day of the global table from the per-key sums. Idempotent: a rerun overwrites every group with the same totals.""" diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index bc64293c9b3..fe161f5d50b 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -2389,7 +2389,6 @@ class ProxyLogging: ) try: - # Execute guardrail pipelines before the normal callback loop if not skip_guardrails: data, _ = await self._maybe_execute_pipelines( # rebind-ok: pipeline edits feed the callback loop below data=data, diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index e421cae0724..bd239922fd3 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -2254,7 +2254,7 @@ class LiteLLMCompletionResponsesConfig: ) -> Mapping[str, ResponseFunctionWebSearch]: calls: Final[dict[str, ResponseFunctionWebSearch]] = {} # mutable-ok: indexes provider-built calls for choice in chat_completion_response.choices: - provider_fields = getattr(choice.message, "provider_specific_fields", None) + provider_fields = choice.message.provider_specific_fields if not isinstance(provider_fields, Mapping): continue web_search_calls = provider_fields.get("web_search_calls") diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index 64989c4cf1c..70f2a7db6da 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -1360,7 +1360,7 @@ def _billed_terminal_response( return None usage: Final[object] = response_obj.get("usage") # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # a model_constructed terminal event leaves response as an untyped dict return ResponsesAPIResponse.model_construct( - **{**response_obj, "usage": usage if usage is not None or estimate is None else estimate()} # pyright: ignore[reportUnknownArgumentType, reportArgumentType] # same untyped dict spread + **{**response_obj, "usage": usage if usage is not None or estimate is None else estimate()} # pyright: ignore[reportArgumentType] # same untyped dict spread ) diff --git a/litellm/router_strategy/complexity_router/jev_classifier.py b/litellm/router_strategy/complexity_router/jev_classifier.py index c0d0d1de8e3..02e57975626 100644 --- a/litellm/router_strategy/complexity_router/jev_classifier.py +++ b/litellm/router_strategy/complexity_router/jev_classifier.py @@ -1,7 +1,7 @@ from collections.abc import Mapping from datetime import datetime, timezone from types import MappingProxyType -from typing import Annotated, Final, Literal, NamedTuple, Protocol +from typing import Annotated, Final, Literal, NamedTuple, Protocol, TypeAlias from uuid import uuid4 import httpx @@ -24,7 +24,7 @@ from litellm.proxy.pass_through_endpoints.llm_provider_handlers.typesafe_passthr from litellm.router_strategy.complexity_router.config import DEFAULT_JEV_INSTRUCTIONS as _DEFAULT_JEV_INSTRUCTIONS from litellm.types.utils import AUTOROUTER_CLASSIFIER_CALL_ORIGIN -JevProbability = Annotated[float, Field(ge=0.0, le=1.0)] +JevProbability: TypeAlias = Annotated[float, Field(ge=0.0, le=1.0)] DEFAULT_JEV_INSTRUCTIONS: Final = _DEFAULT_JEV_INSTRUCTIONS diff --git a/litellm/router_utils/pre_call_checks/encrypted_content_affinity_check.py b/litellm/router_utils/pre_call_checks/encrypted_content_affinity_check.py index cb5c3089685..e5e40d7d6f5 100644 --- a/litellm/router_utils/pre_call_checks/encrypted_content_affinity_check.py +++ b/litellm/router_utils/pre_call_checks/encrypted_content_affinity_check.py @@ -50,6 +50,7 @@ from litellm.exceptions import ( from litellm.integrations.custom_logger import CustomLogger, Span from litellm.litellm_core_utils.credential_accessor import CredentialAccessor from litellm.litellm_core_utils.prompt_templates.common_utils import ( + anthropic_content_lists, encrypted_content_of_block, strip_encrypted_reasoning_from_messages, ) @@ -155,10 +156,7 @@ class EncryptedContentAffinityCheck(CustomLogger): return iter(()) return ( cast(Mapping[str, object], block) # cast-ok: narrowed by isinstance - for message in cast(list[object], messages) # cast-ok: narrowed by isinstance - if isinstance(message, Mapping) - for content in (cast(Mapping[str, object], message).get("content"),) # cast-ok: narrowed by isinstance - if isinstance(content, list) + for content in anthropic_content_lists(cast(list[object], messages)) # cast-ok: narrowed by isinstance for block in cast(list[object], content) # cast-ok: narrowed by isinstance if isinstance(block, Mapping) ) diff --git a/litellm/rust_bridge/dispatch.py b/litellm/rust_bridge/dispatch.py index 076b7759c6d..94ccddc92b7 100644 --- a/litellm/rust_bridge/dispatch.py +++ b/litellm/rust_bridge/dispatch.py @@ -2,7 +2,7 @@ from __future__ import annotations from collections.abc import Awaitable, Callable, Mapping from dataclasses import dataclass -from typing import Final, Generic, TypeVar +from typing import Final, Generic, TypeAlias, TypeVar from litellm.rust_bridge import catalog, runtime from litellm.rust_bridge.bindings import NativeBinding @@ -10,11 +10,11 @@ from litellm.rust_bridge.catalog import Route, RouteContext, RouteRule, Rules from litellm.rust_bridge.configuration import Decision from litellm.rust_bridge.configuration import decision as rollout_decision -RequestT = TypeVar("RequestT") -NativeT = TypeVar("NativeT") -ResultT = TypeVar("ResultT") +RequestT: Final = TypeVar("RequestT") +NativeT: Final = TypeVar("NativeT") +ResultT: Final = TypeVar("ResultT") -NativeHook = Callable[[RequestT, tuple[object, ...], Mapping[str, object]], ResultT] +NativeHook: TypeAlias = Callable[[RequestT, tuple[object, ...], Mapping[str, object]], ResultT] def call_hook( diff --git a/litellm/rust_bridge/response_metadata.py b/litellm/rust_bridge/response_metadata.py index 1c03515720e..ae459710b34 100644 --- a/litellm/rust_bridge/response_metadata.py +++ b/litellm/rust_bridge/response_metadata.py @@ -1,10 +1,10 @@ -from typing import TypeVar +from typing import Final, TypeVar from litellm.router_utils.add_retry_fallback_headers import ( _add_headers_to_response, # pyright: ignore[reportPrivateUsage] # reuse the proxy's identity-preserving response metadata writer ) -ResultT = TypeVar("ResultT") +ResultT: Final = TypeVar("ResultT") def mark_rust_response(response: ResultT) -> ResultT: diff --git a/litellm/types/llms/bedrock.py b/litellm/types/llms/bedrock.py index 3674bb670d5..e4c41c3ee5b 100644 --- a/litellm/types/llms/bedrock.py +++ b/litellm/types/llms/bedrock.py @@ -558,7 +558,6 @@ class AmazonTitanMultimodalEmbeddingResponse(TypedDict): message: str # Specifies any errors that occur during generation. -# TwelveLabs Marengo Embed types TWELVELABS_EMBEDDING_INPUT_TYPES = Literal["text", "image", "video", "audio"] TWELVELABS_EMBEDDING_OPTIONS = Literal["visual-text", "visual-image", "audio"] diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 064e3040054..caf88e5d517 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -3860,7 +3860,7 @@ def without_server_derived_pricing(model_info: Mapping[str, Any]) -> Mapping[str ) -def echoed_cost_map_pricing_fields(model_info: Mapping[str, Any]) -> tuple[str, ...]: +def echoed_cost_map_pricing_fields(model_info: Mapping[str, object]) -> tuple[str, ...]: """Pricing fields a stored ``model_info`` blob copied from a ``/model/info`` response. Only ``litellm.get_model_info`` emits ``key`` (the resolved cost-map entry), so a stored @@ -3891,7 +3891,7 @@ def echoed_cost_map_fields( ) -def pricing_override_fields(*sources: Mapping[str, Any]) -> tuple[str, ...]: +def pricing_override_fields(*sources: Mapping[str, object]) -> tuple[str, ...]: return tuple( sorted( frozenset( From 2eeb16266b906494d3a961d0d03da20810b49924 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 24 Sep 2026 04:39:09 -0700 Subject: [PATCH 034/120] fix(cost-map): sync vertex-ai rows (gemma 4 maas cache price, chirp_2) (#42942) * fix(cost-map): sync vertex-ai rows (gemma 4 maas cache price, chirp_2) Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(vertex-ai): expect chirp_2 as speech-to-text model now that catalog row exists Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: kerry Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../model_prices_and_context_window_backup.json | 15 +++++++++++++++ model_prices_and_context_window.json | 15 +++++++++++++++ .../test_vertex_ai_realtime_transformation.py | 2 +- 3 files changed, 31 insertions(+), 1 deletion(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 54475f38364..298ec16fc55 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -48131,6 +48131,20 @@ "/v1/realtime" ] }, + "vertex_ai/chirp_2": { + "input_cost_per_second": 0.00026667, + "litellm_provider": "vertex_ai", + "metadata": { + "calculation": "$0.016/60 seconds = $0.00026667 per second", + "original_pricing_per_minute": 0.016 + }, + "mode": "audio_transcription", + "source": "https://cloud.google.com/speech-to-text/pricing", + "supported_endpoints": [ + "/v1/audio/transcriptions", + "/v1/realtime" + ] + }, "vertex_ai/claude-3-5-haiku": { "deprecation_date": "2026-07-05", "input_cost_per_token": 1e-06, @@ -50155,6 +50169,7 @@ ] }, "vertex_ai/google/gemma-4-26b-a4b-it-maas": { + "cache_read_input_token_cost": 1.5e-08, "input_cost_per_token": 1.5e-07, "litellm_provider": "vertex_ai-openai_models", "max_input_tokens": 262144, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 54475f38364..298ec16fc55 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -48131,6 +48131,20 @@ "/v1/realtime" ] }, + "vertex_ai/chirp_2": { + "input_cost_per_second": 0.00026667, + "litellm_provider": "vertex_ai", + "metadata": { + "calculation": "$0.016/60 seconds = $0.00026667 per second", + "original_pricing_per_minute": 0.016 + }, + "mode": "audio_transcription", + "source": "https://cloud.google.com/speech-to-text/pricing", + "supported_endpoints": [ + "/v1/audio/transcriptions", + "/v1/realtime" + ] + }, "vertex_ai/claude-3-5-haiku": { "deprecation_date": "2026-07-05", "input_cost_per_token": 1e-06, @@ -50155,6 +50169,7 @@ ] }, "vertex_ai/google/gemma-4-26b-a4b-it-maas": { + "cache_read_input_token_cost": 1.5e-08, "input_cost_per_token": 1.5e-07, "litellm_provider": "vertex_ai-openai_models", "max_input_tokens": 262144, diff --git a/tests/test_litellm/llms/vertex_ai/audio_transcription/test_vertex_ai_realtime_transformation.py b/tests/test_litellm/llms/vertex_ai/audio_transcription/test_vertex_ai_realtime_transformation.py index 84c3a4e244a..b50d6e19458 100644 --- a/tests/test_litellm/llms/vertex_ai/audio_transcription/test_vertex_ai_realtime_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/audio_transcription/test_vertex_ai_realtime_transformation.py @@ -115,7 +115,7 @@ def _commands(config: VertexChirpRealtimeConfig, payload: str) -> list[object]: [ ("vertex_ai/chirp_3", True), ("chirp_3", True), - ("chirp_2", False), + ("chirp_2", True), ("gemini-live-2.5-flash", False), ("vertex_ai/gemini-2.0-flash-live-preview-04-09", False), ("vertex_ai/gemini-3.5-transcribe-live-preview", False), From ddc7ee6838726f2c329202bc4bef9abca58e1b35 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 24 Sep 2026 04:54:10 -0700 Subject: [PATCH 035/120] feat(bedrock): add gpt-5.4 and gpt-5.5 us and global inference profile pricing (#42941) * feat(bedrock): add gpt-5.4 and gpt-5.5 us and global inference profile pricing Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(bedrock): drop supports_max_reasoning_effort from gpt-5.4 and gpt-5.5 rows Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: kerry Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- ...odel_prices_and_context_window_backup.json | 124 ++++++++++++++++++ model_prices_and_context_window.json | 124 ++++++++++++++++++ 2 files changed, 248 insertions(+) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 298ec16fc55..0f2782efffa 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -57151,6 +57151,130 @@ "/v1/responses" ] }, + "us.openai.gpt-5.4": { + "input_cost_per_token": 2.75e-06, + "input_cost_per_token_above_272k_tokens": 5.5e-06, + "cache_read_input_token_cost": 2.75e-07, + "cache_read_input_token_cost_above_272k_tokens": 5.5e-07, + "output_cost_per_token": 1.65e-05, + "output_cost_per_token_above_272k_tokens": 2.475e-05, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-54.html", + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_minimal_reasoning_effort": false, + "supports_tool_choice": true, + "supports_reasoning": true, + "supports_xhigh_reasoning_effort": true, + "supports_vision": true, + "supports_sampling_params": false, + "supported_endpoints": [ + "/v1/responses" + ] + }, + "global.openai.gpt-5.4": { + "input_cost_per_token": 2.75e-06, + "input_cost_per_token_above_272k_tokens": 5.5e-06, + "cache_read_input_token_cost": 2.75e-07, + "cache_read_input_token_cost_above_272k_tokens": 5.5e-07, + "output_cost_per_token": 1.65e-05, + "output_cost_per_token_above_272k_tokens": 2.475e-05, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-54.html", + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_minimal_reasoning_effort": false, + "supports_tool_choice": true, + "supports_reasoning": true, + "supports_xhigh_reasoning_effort": true, + "supports_vision": true, + "supports_sampling_params": false, + "supported_endpoints": [ + "/v1/responses" + ] + }, + "us.openai.gpt-5.5": { + "input_cost_per_token": 5.5e-06, + "input_cost_per_token_above_272k_tokens": 1.1e-05, + "cache_read_input_token_cost": 5.5e-07, + "cache_read_input_token_cost_above_272k_tokens": 1.1e-06, + "output_cost_per_token": 3.3e-05, + "output_cost_per_token_above_272k_tokens": 4.95e-05, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-55.html", + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_minimal_reasoning_effort": false, + "supports_tool_choice": true, + "supports_reasoning": true, + "supports_xhigh_reasoning_effort": true, + "supports_vision": true, + "supports_sampling_params": false, + "supported_endpoints": [ + "/v1/responses" + ] + }, + "global.openai.gpt-5.5": { + "input_cost_per_token": 5.5e-06, + "input_cost_per_token_above_272k_tokens": 1.1e-05, + "cache_read_input_token_cost": 5.5e-07, + "cache_read_input_token_cost_above_272k_tokens": 1.1e-06, + "output_cost_per_token": 3.3e-05, + "output_cost_per_token_above_272k_tokens": 4.95e-05, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-55.html", + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_minimal_reasoning_effort": false, + "supports_tool_choice": true, + "supports_reasoning": true, + "supports_xhigh_reasoning_effort": true, + "supports_vision": true, + "supports_sampling_params": false, + "supported_endpoints": [ + "/v1/responses" + ] + }, "global.openai.gpt-5.6-luna": { "input_cost_per_token": 2e-07, "input_cost_per_token_above_272k_tokens": 4e-07, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 298ec16fc55..0f2782efffa 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -57151,6 +57151,130 @@ "/v1/responses" ] }, + "us.openai.gpt-5.4": { + "input_cost_per_token": 2.75e-06, + "input_cost_per_token_above_272k_tokens": 5.5e-06, + "cache_read_input_token_cost": 2.75e-07, + "cache_read_input_token_cost_above_272k_tokens": 5.5e-07, + "output_cost_per_token": 1.65e-05, + "output_cost_per_token_above_272k_tokens": 2.475e-05, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-54.html", + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_minimal_reasoning_effort": false, + "supports_tool_choice": true, + "supports_reasoning": true, + "supports_xhigh_reasoning_effort": true, + "supports_vision": true, + "supports_sampling_params": false, + "supported_endpoints": [ + "/v1/responses" + ] + }, + "global.openai.gpt-5.4": { + "input_cost_per_token": 2.75e-06, + "input_cost_per_token_above_272k_tokens": 5.5e-06, + "cache_read_input_token_cost": 2.75e-07, + "cache_read_input_token_cost_above_272k_tokens": 5.5e-07, + "output_cost_per_token": 1.65e-05, + "output_cost_per_token_above_272k_tokens": 2.475e-05, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-54.html", + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_minimal_reasoning_effort": false, + "supports_tool_choice": true, + "supports_reasoning": true, + "supports_xhigh_reasoning_effort": true, + "supports_vision": true, + "supports_sampling_params": false, + "supported_endpoints": [ + "/v1/responses" + ] + }, + "us.openai.gpt-5.5": { + "input_cost_per_token": 5.5e-06, + "input_cost_per_token_above_272k_tokens": 1.1e-05, + "cache_read_input_token_cost": 5.5e-07, + "cache_read_input_token_cost_above_272k_tokens": 1.1e-06, + "output_cost_per_token": 3.3e-05, + "output_cost_per_token_above_272k_tokens": 4.95e-05, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-55.html", + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_minimal_reasoning_effort": false, + "supports_tool_choice": true, + "supports_reasoning": true, + "supports_xhigh_reasoning_effort": true, + "supports_vision": true, + "supports_sampling_params": false, + "supported_endpoints": [ + "/v1/responses" + ] + }, + "global.openai.gpt-5.5": { + "input_cost_per_token": 5.5e-06, + "input_cost_per_token_above_272k_tokens": 1.1e-05, + "cache_read_input_token_cost": 5.5e-07, + "cache_read_input_token_cost_above_272k_tokens": 1.1e-06, + "output_cost_per_token": 3.3e-05, + "output_cost_per_token_above_272k_tokens": 4.95e-05, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-openai-gpt-55.html", + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_minimal_reasoning_effort": false, + "supports_tool_choice": true, + "supports_reasoning": true, + "supports_xhigh_reasoning_effort": true, + "supports_vision": true, + "supports_sampling_params": false, + "supported_endpoints": [ + "/v1/responses" + ] + }, "global.openai.gpt-5.6-luna": { "input_cost_per_token": 2e-07, "input_cost_per_token_above_272k_tokens": 4e-07, From c4b56b6adad3c2fc8f9af691b5c0e325c5702a9e Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 24 Sep 2026 06:15:55 -0700 Subject: [PATCH 036/120] fix(cost-map): update azure gpt-4.1-nano and gpt-4o-2024-05-13 retirement dates (#42947) Co-authored-by: kerry Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../model_prices_and_context_window_backup.json | 16 ++++++++-------- model_prices_and_context_window.json | 16 ++++++++-------- 2 files changed, 16 insertions(+), 16 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 0f2782efffa..bba80677a80 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -5442,7 +5442,7 @@ "supports_web_search": false }, "azure/gpt-4.1-nano": { - "deprecation_date": "2027-04-14", + "deprecation_date": "2026-10-14", "cache_read_input_token_cost": 2.5e-08, "input_cost_per_token": 1e-07, "input_cost_per_token_batches": 5e-08, @@ -5476,7 +5476,7 @@ "supports_vision": true }, "azure/gpt-4.1-nano-2025-04-14": { - "deprecation_date": "2027-04-14", + "deprecation_date": "2026-10-14", "cache_read_input_token_cost": 2.5e-08, "input_cost_per_token": 1e-07, "input_cost_per_token_batches": 5e-08, @@ -5546,7 +5546,7 @@ "supports_vision": true }, "azure/gpt-4o-2024-05-13": { - "deprecation_date": "2026-10-01", + "deprecation_date": "2026-12-09", "input_cost_per_token": 5e-06, "input_cost_per_token_batches": 2.5e-06, "litellm_provider": "azure", @@ -10660,7 +10660,7 @@ "supports_web_search": false }, "azure/us/gpt-4.1-nano-2025-04-14": { - "deprecation_date": "2027-04-14", + "deprecation_date": "2026-10-14", "cache_read_input_token_cost": 2.8e-08, "input_cost_per_token": 1.1e-07, "input_cost_per_token_batches": 5.5e-08, @@ -68776,7 +68776,7 @@ "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" }, "azure/eu/gpt-4.1-nano": { - "deprecation_date": "2027-04-14", + "deprecation_date": "2026-10-14", "cache_read_input_token_cost": 2.8e-08, "input_cost_per_token": 1.1e-07, "input_cost_per_token_batches": 5.5e-08, @@ -68787,7 +68787,7 @@ "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" }, "azure/eu/gpt-4o-2024-05-13": { - "deprecation_date": "2026-10-01", + "deprecation_date": "2026-12-09", "input_cost_per_token": 5.5e-06, "input_cost_per_token_batches": 2.75e-06, "litellm_provider": "azure", @@ -69211,7 +69211,7 @@ "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" }, "azure/us/gpt-4.1-nano": { - "deprecation_date": "2027-04-14", + "deprecation_date": "2026-10-14", "cache_read_input_token_cost": 2.8e-08, "input_cost_per_token": 1.1e-07, "input_cost_per_token_batches": 5.5e-08, @@ -69222,7 +69222,7 @@ "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" }, "azure/us/gpt-4o-2024-05-13": { - "deprecation_date": "2026-10-01", + "deprecation_date": "2026-12-09", "input_cost_per_token": 5.5e-06, "input_cost_per_token_batches": 2.75e-06, "litellm_provider": "azure", diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 0f2782efffa..bba80677a80 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -5442,7 +5442,7 @@ "supports_web_search": false }, "azure/gpt-4.1-nano": { - "deprecation_date": "2027-04-14", + "deprecation_date": "2026-10-14", "cache_read_input_token_cost": 2.5e-08, "input_cost_per_token": 1e-07, "input_cost_per_token_batches": 5e-08, @@ -5476,7 +5476,7 @@ "supports_vision": true }, "azure/gpt-4.1-nano-2025-04-14": { - "deprecation_date": "2027-04-14", + "deprecation_date": "2026-10-14", "cache_read_input_token_cost": 2.5e-08, "input_cost_per_token": 1e-07, "input_cost_per_token_batches": 5e-08, @@ -5546,7 +5546,7 @@ "supports_vision": true }, "azure/gpt-4o-2024-05-13": { - "deprecation_date": "2026-10-01", + "deprecation_date": "2026-12-09", "input_cost_per_token": 5e-06, "input_cost_per_token_batches": 2.5e-06, "litellm_provider": "azure", @@ -10660,7 +10660,7 @@ "supports_web_search": false }, "azure/us/gpt-4.1-nano-2025-04-14": { - "deprecation_date": "2027-04-14", + "deprecation_date": "2026-10-14", "cache_read_input_token_cost": 2.8e-08, "input_cost_per_token": 1.1e-07, "input_cost_per_token_batches": 5.5e-08, @@ -68776,7 +68776,7 @@ "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" }, "azure/eu/gpt-4.1-nano": { - "deprecation_date": "2027-04-14", + "deprecation_date": "2026-10-14", "cache_read_input_token_cost": 2.8e-08, "input_cost_per_token": 1.1e-07, "input_cost_per_token_batches": 5.5e-08, @@ -68787,7 +68787,7 @@ "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" }, "azure/eu/gpt-4o-2024-05-13": { - "deprecation_date": "2026-10-01", + "deprecation_date": "2026-12-09", "input_cost_per_token": 5.5e-06, "input_cost_per_token_batches": 2.75e-06, "litellm_provider": "azure", @@ -69211,7 +69211,7 @@ "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" }, "azure/us/gpt-4.1-nano": { - "deprecation_date": "2027-04-14", + "deprecation_date": "2026-10-14", "cache_read_input_token_cost": 2.8e-08, "input_cost_per_token": 1.1e-07, "input_cost_per_token_batches": 5.5e-08, @@ -69222,7 +69222,7 @@ "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" }, "azure/us/gpt-4o-2024-05-13": { - "deprecation_date": "2026-10-01", + "deprecation_date": "2026-12-09", "input_cost_per_token": 5.5e-06, "input_cost_per_token_batches": 2.75e-06, "litellm_provider": "azure", From 8477fe4108b742acbf7e40aabb79e3147fb9d329 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 24 Sep 2026 06:41:13 -0700 Subject: [PATCH 037/120] fix(proxy): list key and team model aliases in GET /v1/models (#42908) * fix(proxy): list key and team model aliases in GET /v1/models Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): keep alias listing helpers within the type discipline budget Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(proxy): cover alias rows on GET /v1/models and /v1/models/{id} Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): apply team then key aliases like chat completions and keep the alias as the retrieved id Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): apply key aliases twice like chat completions and skip only malformed alias entries Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): apply the global model_alias_map between the key alias passes like chat completions Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): list only the caller's own aliases and never rewrite a listed model id on retrieval Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * style(proxy): ruff format model_info alias lookup Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): hide undiscoverable names from model retrieval so an alias named like one resolves to its target Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): keep undiscoverable models retrievable by id while excluding them from the alias guard Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * refactor(proxy): pass an immutable name sequence into the model_info alias guard Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(proxy): type the model list alias test helpers Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(proxy): annotate the new alias listing test fixtures and helpers Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../proxy/common_utils/model_listing_utils.py | 80 +++++++- litellm/proxy/proxy_server.py | 39 +++- tests/e2e/coverage_registry/other.yaml | 1 + tests/e2e/models.py | 1 + tests/e2e/other/other_client.py | 16 +- tests/e2e/other/test_jwt_auth_e2e.py | 47 ++++- .../common_utils/test_model_listing_utils.py | 177 +++++++++++++++--- .../proxy/test_model_list_aliases.py | 151 +++++++++++++++ 8 files changed, 477 insertions(+), 35 deletions(-) create mode 100644 tests/test_litellm/proxy/test_model_list_aliases.py diff --git a/litellm/proxy/common_utils/model_listing_utils.py b/litellm/proxy/common_utils/model_listing_utils.py index 4b3ff2711b3..3c6555662e6 100644 --- a/litellm/proxy/common_utils/model_listing_utils.py +++ b/litellm/proxy/common_utils/model_listing_utils.py @@ -13,9 +13,12 @@ from __future__ import annotations import re from collections.abc import Container, Mapping, Sequence from dataclasses import dataclass +from functools import reduce from types import MappingProxyType from typing import TYPE_CHECKING, Final, cast +from pydantic import TypeAdapter, ValidationError + import litellm if TYPE_CHECKING: @@ -28,6 +31,8 @@ CLAUDE_CODE_CLIENT: Final = "claude-code" _CLAUDE_CODE_ALIAS_PREFIX: Final = "claude-router-" _ONE_MILLION_SUFFIX: Final = "[1m]" _ONE_MILLION_TOKENS: Final = 1_000_000 +_ALIAS_ENTRIES: Final = TypeAdapter(Mapping[object, object]) +_NO_ALIASES: Final[Mapping[str, str]] = MappingProxyType({}) def configured_display_names( @@ -152,6 +157,77 @@ class ClaudeCodeRoutingNames: ) +@dataclass(frozen=True, slots=True) +class CallerAliases: + """`own` are the caller's key and team alias maps, the names `/v1/models` lists for it. + `rewrite` are the maps `/chat/completions` rewrites its model through, in the order it + applies them: the team's, the key's in `add_litellm_data_to_request`, then the global + `model_alias_map` and the key's again in `common_processing_pre_call_logic`.""" + + own: tuple[object, ...] + rewrite: tuple[object, ...] + + +def caller_alias_maps( + key_aliases: object, + team_aliases: object, + key_team_id: str | None, + listed_team_id: str | None, +) -> CallerAliases: + """Team aliases count only when listing the team the key authenticated as.""" + if listed_team_id is not None and listed_team_id != key_team_id: + return CallerAliases((key_aliases,), (key_aliases, litellm.model_alias_map, key_aliases)) + return CallerAliases((team_aliases, key_aliases), (team_aliases, key_aliases, litellm.model_alias_map, key_aliases)) + + +def _alias_map(aliases: object) -> Mapping[str, str]: + try: + entries: Final = _ALIAS_ENTRIES.validate_python(aliases, strict=True) + except ValidationError: + return _NO_ALIASES + return MappingProxyType( + {alias: target for alias, target in entries.items() if isinstance(alias, str) and isinstance(target, str)} + ) + + +def _alias_names(alias_maps: Sequence[Mapping[str, str]]) -> tuple[str, ...]: + return tuple(dict.fromkeys(alias for aliases in alias_maps for alias in aliases)) + + +def _rewrite(model_id: str, alias_maps: Sequence[Mapping[str, str]]) -> str | None: + target: Final = reduce(lambda name, aliases: aliases.get(name, name), alias_maps, model_id) + return None if target == model_id else target + + +def alias_target(model_id: str, aliases: CallerAliases, listed: Container[str] = frozenset()) -> str | None: + """The model group `/chat/completions` rewrites `model_id` to, else None. A `model_id` + already `listed` keeps its own row, so it is never rewritten.""" + if model_id in listed: + return None + return _rewrite(model_id, tuple(_alias_map(alias_map) for alias_map in aliases.rewrite)) + + +def alias_listing_entries( + entries: Sequence[tuple[str, str]], + aliases: CallerAliases, +) -> tuple[tuple[str, str], ...]: + """`entries` plus one `(alias, lookup_id)` row per key or team alias whose target is + listed. An alias colliding with a listed id keeps the listed entry.""" + maps: Final = tuple(_alias_map(alias_map) for alias_map in aliases.rewrite) + own: Final = tuple(_alias_map(alias_map) for alias_map in aliases.own) + lookup_by_response: Final = MappingProxyType(dict(entries)) + lookup_ids: Final = frozenset(lookup_by_response.values()) + targets: Final = MappingProxyType( + {alias: _rewrite(alias, maps) for alias in _alias_names(own) if alias not in lookup_by_response} + ) + added: Final = tuple( + (alias, lookup_by_response.get(target, target)) + for alias, target in targets.items() + if target is not None and (target in lookup_by_response or target in lookup_ids) + ) + return (*entries, *added) + + def claude_code_requested_group( requested: str, llm_router: Router, @@ -218,7 +294,7 @@ class TeamModelNameTranslator: @staticmethod def _response_to_lookup_map( - model_names: list[str], + model_names: Sequence[str], internal_to_public: dict[str, str], ) -> dict[str, str]: """Map each public response id to the first internal lookup id seen in @@ -235,7 +311,7 @@ class TeamModelNameTranslator: @staticmethod def listing_entries( - model_names: list[str], + model_names: Sequence[str], llm_router: Router | None, general_settings: Mapping[str, object], ) -> list[tuple[str, str]]: diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 7e18db742cd..9fb72d2d1b5 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -422,6 +422,9 @@ from litellm.proxy.common_utils.model_deprecation import collect_model_deprecati from litellm.proxy.common_utils.model_listing_utils import ( ClaudeCodeRoutingNames, TeamModelNameTranslator, + alias_listing_entries, + alias_target, + caller_alias_maps, claude_code_view_ids, configured_display_names, is_claude_code_client, @@ -11172,14 +11175,13 @@ async def model_list( view_aliases: Final = ( view_router_settings.get("model_group_alias") if isinstance(view_router_settings, Mapping) else None ) + caller_aliases: Final = caller_alias_maps( + user_api_key_dict.aliases, user_api_key_dict.team_model_aliases, user_api_key_dict.team_id, team_id + ) routing_names: Final = ClaudeCodeRoutingNames( llm_router, team_id or user_api_key_dict.team_id, - ( - user_api_key_dict.aliases, - user_api_key_dict.team_model_aliases, - view_aliases, - ), + (*caller_aliases.rewrite, view_aliases), ) # Validate scope parameter if provided @@ -11307,7 +11309,9 @@ async def model_list( # The internal routing key drives the metadata/fallback lookup, while the # public name is what the client sees as the model id. model_data = [] - entries: Final = TeamModelNameTranslator.listing_entries(all_models, llm_router, settings) + entries: Final = alias_listing_entries( + TeamModelNameTranslator.listing_entries(all_models, llm_router, settings), caller_aliases + ) for response_id, lookup_id in entries: model_info = create_model_info_response( model_id=lookup_id, @@ -11391,7 +11395,8 @@ async def model_info( ) # Mirror /v1/models' visibility filter so first-occurrence resolution - # cannot land on a deployment the listing had hidden. + # cannot land on a deployment the listing had hidden. Undiscoverable + # models stay retrievable by id, they only drop out of the alias guard. blocked_names: Final = llm_router.get_fully_blocked_model_names() if llm_router is not None else set() unhealthy_names: Final = await get_hidden_unhealthy_model_names( healthy_only=healthy_only, @@ -11401,10 +11406,25 @@ async def model_info( hidden_names: Final = blocked_names | unhealthy_names if hidden_names: all_models = [m for m in all_models if m not in hidden_names] + undiscoverable_names: Final = undiscoverable_model_names( + all_models, llm_router, user_api_key_dict, team_id or user_api_key_dict.team_id + ) internal_to_public: Final = TeamModelNameTranslator.build_internal_to_public_map(llm_router, settings) + aliased_model_id: Final = alias_target( + model_id, + caller_alias_maps( + user_api_key_dict.aliases, user_api_key_dict.team_model_aliases, user_api_key_dict.team_id, team_id + ), + frozenset( + response_id + for response_id, _ in TeamModelNameTranslator.listing_entries( + tuple(m for m in all_models if m not in undiscoverable_names), llm_router, settings + ) + ), + ) resolved_model_id: Final = TeamModelNameTranslator.resolve_public_name( - model_id=model_id, + model_id=aliased_model_id or model_id, available_models=all_models, llm_router=llm_router, general_settings=settings, @@ -11434,7 +11454,8 @@ async def model_info( fallback_type=None, llm_router=llm_router, ) - return {**response, "id": internal_to_public.get(resolved_model_id, model_id)} # mutable-ok: response id differs + response_id: Final = model_id if aliased_model_id else internal_to_public.get(resolved_model_id, model_id) + return {**response, "id": response_id} # mutable-ok: response id differs def _blocked_response_usage(original_response: object | None) -> "litellm.Usage": diff --git a/tests/e2e/coverage_registry/other.yaml b/tests/e2e/coverage_registry/other.yaml index f695af3cb11..3bd98ff5b0b 100644 --- a/tests/e2e/coverage_registry/other.yaml +++ b/tests/e2e/coverage_registry/other.yaml @@ -17,6 +17,7 @@ - {id: other.auth.jwt.virtual_key_unaffected, module: other, tier: P0, area: auth, assertions: [virtual_key_unaffected], source: "handle_jwt.py:213 is_jwt / user_api_key_auth.py:1332-1333", rationale: "enable_jwt_auth only routes three-segment bearer tokens into the JWT branch, so sk- virtual keys keep working on the same proxy"} - {id: other.auth.jwt.team_header_alias_binds_team, module: other, tier: P0, area: auth, assertions: [team_header_alias_binds_team], source: "handle_jwt.py JWTAuthManager.resolve_team_from_header / LIT-7181", fail_before_fix: proven, rationale: "x-litellm-team-id carrying the team alias binds and attributes the same team as the team id, so a managed client can pin a stable alias instead of a uuid"} - {id: other.auth.jwt.team_header_non_member_alias_denied, module: other, tier: P0, area: auth, assertions: [team_header_non_member_alias_denied], source: "handle_jwt.py JWTAuthManager.resolve_team_from_header / LIT-7181", rationale: "x-litellm-team-id naming the alias of a team the JWT does not grant is denied 403 with the same body as an unknown value, so the response does not reveal whether that team exists"} +- {id: other.auth.jwt.team_model_alias_listed_and_routes, module: other, tier: P1, area: auth, assertions: [team_model_alias_listed_and_routes], source: "proxy_server.py model_list / common_utils/model_listing_utils.py alias_listing_entries / LIT-8515", fail_before_fix: proven, rationale: "A team model_aliases name the JWT caller can complete on is also listed by GET /v1/models for that caller, in the OpenAI and the Anthropic (Claude Code) shapes, next to its target, so a managed client can discover the alias it is meant to send"} - {id: other.auth.model_access_group.wildcard_bare_name_allowed, module: other, tier: P0, area: auth, assertions: [wildcard_bare_name_allowed], source: "auth_checks.py:3232 / LIT-5813", fail_before_fix: proven, rationale: "A grant of a group holding a wildcard deployment covers the bare model names callers actually send, not only the provider-prefixed spelling"} - {id: other.auth.model_access_group.member_allowed, module: other, tier: P0, area: auth, assertions: [member_allowed], source: "auth_checks.py:3232", rationale: "A key whose allow-list is a model access group can call the deployments in that group"} - {id: other.auth.model_access_group.non_member_denied, module: other, tier: P0, area: auth, assertions: [non_member_denied], source: "auth_checks.py:3232", rationale: "That same grant reaches nothing outside the group, including provider models the group's wildcard does not cover"} diff --git a/tests/e2e/models.py b/tests/e2e/models.py index bd7f5171172..c96f4b0bef1 100644 --- a/tests/e2e/models.py +++ b/tests/e2e/models.py @@ -1426,6 +1426,7 @@ class TeamNewBody(BaseModel): team_id: str | None = None organization_id: str | None = None metadata: TeamMetadata | None = None + model_aliases: dict[str, str] | None = None class TeamNewResponse(BaseModel): diff --git a/tests/e2e/other/other_client.py b/tests/e2e/other/other_client.py index 057a735e5e5..93c198586f6 100644 --- a/tests/e2e/other/other_client.py +++ b/tests/e2e/other/other_client.py @@ -14,12 +14,15 @@ own endpoints, so no test ever holds a signing key. from __future__ import annotations from dataclasses import dataclass +from typing import Final -from e2e_http import AuthHeaders, NoBody, ProbeResult, Result +from e2e_http import AnthropicHeaders, AuthHeaders, NoBody, ProbeResult, Result from idp import Keycloak, keycloak_from_env from models import ( ChatBody, ChatResponse, + ModelsListParams, + ModelsListResponse, ReadinessDetailsResponse, ReadinessResponse, UserListParams, @@ -88,6 +91,17 @@ class OtherClient: response_type=ChatResponse, ) + def list_models_as(self, token: str, *, anthropic: bool = False) -> Result[ModelsListResponse]: + """GET /v1/models under `token`, in the OpenAI shape or, with `anthropic`, the + Anthropic Models API shape Claude Code reads. Both carry `data[].id`.""" + bearer: Final = self.proxy.transport.bearer(token) + return self.proxy.transport.get( + "/v1/models", + headers=AnthropicHeaders(authorization=bearer.authorization) if anthropic else bearer, + params=ModelsListParams(return_wildcard_routes=False), + response_type=ModelsListResponse, + ) + def list_users_as(self, key: str) -> Result[UserListResponse]: """GET /user/list under `key`. Admin-only, so it doubles as the master key's authorization proof: the master key (proxy admin) reads it, a diff --git a/tests/e2e/other/test_jwt_auth_e2e.py b/tests/e2e/other/test_jwt_auth_e2e.py index de06c376f76..a8aa474f04f 100644 --- a/tests/e2e/other/test_jwt_auth_e2e.py +++ b/tests/e2e/other/test_jwt_auth_e2e.py @@ -78,9 +78,35 @@ def bound_team(client: OtherClient, resources: ResourceManager) -> BoundTeam: return BoundTeam(identity=provisioned, team_id=provisioned.group, team_alias=team_alias) -def _ping() -> ChatBody: +@dataclass(frozen=True, slots=True) +class AliasedTeam: + identity: Identity + alias: str + target: str + + +@pytest.fixture +def aliased_team(client: OtherClient, resources: ResourceManager) -> AliasedTeam: + """An identity whose team carries a model_aliases entry, the name a managed + client such as Claude Code sends and the team rewrites to a real model group.""" + marker: Final = unique_marker() + provisioned: Final = _provision(client, resources, marker=marker) + alias: Final = f"e2e-jwt-model-alias-{marker}" + team_id: Final = client.proxy.create_team( + TeamNewBody( + team_alias=f"e2e-jwt-aliased-{marker}", + team_id=provisioned.group, + models=[CHEAP_OPENAI_MODEL], + model_aliases={alias: CHEAP_OPENAI_MODEL}, + ) + ) + resources.defer(lambda: client.proxy.delete_team(team_id)) + return AliasedTeam(identity=provisioned, alias=alias, target=CHEAP_OPENAI_MODEL) + + +def _ping(model: str = CHEAP_OPENAI_MODEL) -> ChatBody: return ChatBody( - model=CHEAP_OPENAI_MODEL, + model=model, messages=[ChatMessage(role="user", content=f"Reply with the single word pong. {unique_marker()}")], max_tokens=16, ) @@ -221,6 +247,23 @@ class TestJwtTeamHeader: f"{bound_team.team_id!r}, got {by_alias!r}" ) + @pytest.mark.covers("other.auth.jwt.team_model_alias_listed_and_routes") + @pytest.mark.parametrize("anthropic", [False, True], ids=["openai_shape", "anthropic_shape"]) + def test_team_model_alias_is_listed_by_v1_models_under_the_same_token_that_routes_it( + self, client: OtherClient, aliased_team: AliasedTeam, anthropic: bool + ) -> None: + token: Final = client.idp.access_token(aliased_team.identity) + + routed: Final = unwrap(client.proxy.chat(token, _ping(model=aliased_team.alias))) + assert routed.choices, f"precondition: /chat/completions must route the team alias, got {routed}" + + listed: Final = tuple(entry.id for entry in unwrap(client.list_models_as(token, anthropic=anthropic)).data) + assert aliased_team.alias in listed, ( + f"/v1/models must list team alias {aliased_team.alias!r} that the same token routes on " + f"/chat/completions, got {listed}" + ) + assert aliased_team.target in listed, f"the alias target {aliased_team.target!r} must stay listed, got {listed}" + @pytest.mark.covers("other.auth.jwt.team_header_non_member_alias_denied") def test_team_header_with_the_alias_of_a_team_the_caller_is_not_in_is_rejected_like_an_unknown_value( self, client: OtherClient, resources: ResourceManager, bound_team: BoundTeam diff --git a/tests/test_litellm/proxy/common_utils/test_model_listing_utils.py b/tests/test_litellm/proxy/common_utils/test_model_listing_utils.py index 7ef03140093..97e38d5e17a 100644 --- a/tests/test_litellm/proxy/common_utils/test_model_listing_utils.py +++ b/tests/test_litellm/proxy/common_utils/test_model_listing_utils.py @@ -4,9 +4,14 @@ from itertools import combinations import pytest +import litellm from litellm import Router from litellm.proxy.common_utils.model_listing_utils import ( + CallerAliases, ClaudeCodeRoutingNames, + alias_listing_entries, + alias_target, + caller_alias_maps, claude_code_group_name, claude_code_model_id, claude_code_requested_group, @@ -22,6 +27,10 @@ def _marked(name): return f"{_encoded(name)}[1m]" +def _caller(*maps: object) -> CallerAliases: + return CallerAliases(maps, maps) + + def _row(name, limit=1000000): return {"id": name, "object": "model", "created": 0, "owned_by": "openai", "max_input_tokens": limit} @@ -29,15 +38,16 @@ def _row(name, limit=1000000): def _router(*names, aliases=None): return Router( model_list=[ - {"model_name": name, "litellm_params": {"model": "openai/gpt-4o", "api_key": "sk-fake"}} - for name in names + {"model_name": name, "litellm_params": {"model": "openai/gpt-4o", "api_key": "sk-fake"}} for name in names ], model_group_alias=aliases, ) @pytest.mark.parametrize("limit", [None, 999999, 1000000]) -@pytest.mark.parametrize("name", ["foo", "foo[1m]", "foo[1M]", "a/b: 世界", "claude-router-foo", "claude-opus-5", "claude-opus-5[1m]"]) +@pytest.mark.parametrize( + "name", ["foo", "foo[1m]", "foo[1M]", "a/b: 世界", "claude-router-foo", "claude-opus-5", "claude-opus-5[1m]"] +) def test_listing_round_trips_entire_source_name(name, limit): names = frozenset({name}) view = claude_code_model_id(name, limit, names) @@ -48,7 +58,15 @@ def test_listing_round_trips_entire_source_name(name, limit): def test_collision_matrix_round_trips_without_duplicate_ids(): - universe = ("foo", "foo[1m]", "claude-router-foo", _encoded("foo"), _encoded("foo") + "[1m]", "claude-opus-5", "claude-opus-5[1m]") + universe = ( + "foo", + "foo[1m]", + "claude-router-foo", + _encoded("foo"), + _encoded("foo") + "[1m]", + "claude-opus-5", + "claude-opus-5[1m]", + ) for pair in combinations(universe, 2): for visible in (pair, pair[:1], pair[1:]): names = frozenset(pair) @@ -57,18 +75,31 @@ def test_collision_matrix_round_trips_without_duplicate_ids(): assert all((claude_code_group_name(shown, names) or shown) == source for source, shown in view.items()) -@pytest.mark.parametrize("spelling", ["claude-router-foo", "claude-router-ff", "claude-router-66 6f6f", "claude-router-666F6F", "claude-router-", _encoded("missing")]) +@pytest.mark.parametrize( + "spelling", + [ + "claude-router-foo", + "claude-router-ff", + "claude-router-66 6f6f", + "claude-router-666F6F", + "claude-router-", + _encoded("missing"), + ], +) def test_unknown_or_noncanonical_ids_are_never_guessed(spelling): assert claude_code_group_name(spelling, frozenset({"foo"})) is None -@pytest.mark.parametrize("headers,enabled", [ - ({"user-agent": "claude-code/2.1.267"}, True), - ({"user-agent": "claude-cli/2.1.267 (external, sdk-cli)"}, True), - ({"x-gateway-client": "Claude-Code"}, True), - ({"user-agent": "anthropic-sdk-python/0.40"}, False), - ({}, False), -]) +@pytest.mark.parametrize( + "headers,enabled", + [ + ({"user-agent": "claude-code/2.1.267"}, True), + ({"user-agent": "claude-cli/2.1.267 (external, sdk-cli)"}, True), + ({"x-gateway-client": "Claude-Code"}, True), + ({"user-agent": "anthropic-sdk-python/0.40"}, False), + ({}, False), + ], +) def test_only_claude_code_gets_the_view(headers, enabled): rows = (_row("foo"), _row("claude-opus-5")) view = claude_code_view_ids(rows, headers, frozenset(row["id"] for row in rows)) @@ -77,12 +108,15 @@ def test_only_claude_code_gets_the_view(headers, enabled): @pytest.mark.parametrize("layer", ["literal", "global", "router", "key", "team", "wildcard"]) def test_configured_names_outrank_generated_ids_even_when_hidden_from_listing(monkeypatch, layer): - import litellm - encoded = _encoded("foo") alias = {encoded: "other"} monkeypatch.setattr(litellm, "model_alias_map", alias if layer == "global" else {}) - router = _router("foo", "other", *( (encoded,) if layer == "literal" else ("*",) if layer == "wildcard" else ()), aliases=alias if layer == "router" else None) + router = _router( + "foo", + "other", + *((encoded,) if layer == "literal" else ("*",) if layer == "wildcard" else ()), + aliases=alias if layer == "router" else None, + ) maps = (alias,) if layer in ("key", "team") else () names = ClaudeCodeRoutingNames(router, None, maps) assert claude_code_requested_group(encoded, router, None, maps) is None @@ -98,12 +132,113 @@ def test_mutation_breaking_the_hex_name_cannot_route_to_the_source(source): assert claude_code_requested_group(_marked(source), router, None) == source +def test_team_alias_is_listed_under_its_target_metadata_and_only_when_the_target_is_accessible() -> None: + entries = [("gpt-4.1-mini", "gpt-4.1-mini"), ("team-public", "model_name_team_1_abc")] + aliases = ( + {"gpt-4.1-mini": "team-public"}, + None, + {"claude-sonnet-4-5": "gpt-4.1-mini", "via-public": "team-public", "not-granted": "gpt-4.1"}, + ) + assert alias_listing_entries(entries, _caller(*aliases)) == ( + *entries, + ("claude-sonnet-4-5", "gpt-4.1-mini"), + ("via-public", "model_name_team_1_abc"), + ) + assert alias_listing_entries(entries, _caller(None, {})) == tuple(entries) + + +def test_alias_target_resolves_the_requested_alias_across_key_and_team_maps() -> None: + maps = _caller({"o": "gpt-4.1"}, {"claude-sonnet-4-5": "gpt-4.1-mini"}) + assert alias_target("claude-sonnet-4-5", maps) == "gpt-4.1-mini" + assert alias_target("gpt-4.1-mini", _caller(None, {"claude-sonnet-4-5": "gpt-4.1-mini"})) is None + + +def test_alias_colliding_with_a_listed_id_keeps_the_listed_model_at_list_and_retrieval() -> None: + entries = [("fast", "fast"), ("gpt-4.1-mini", "gpt-4.1-mini")] + maps = _caller({"fast": "gpt-4.1-mini"}) + listed = frozenset(response_id for response_id, _ in entries) + + assert alias_listing_entries(entries, maps) == tuple(entries) + assert alias_target("fast", maps, listed) is None + assert alias_target("fast", maps) == "gpt-4.1-mini" + + +def test_alias_maps_apply_in_the_order_chat_completions_applies_them() -> None: + team_then_key = _caller({"fast": "gpt-4.1-mini", "hop": "mid"}, {"fast": "gpt-4.1", "mid": "gpt-4.1"}) + entries = [("gpt-4.1-mini", "gpt-4.1-mini"), ("gpt-4.1", "gpt-4.1")] + + assert alias_target("fast", team_then_key) == "gpt-4.1-mini" + assert alias_target("hop", team_then_key) == "gpt-4.1" + assert alias_listing_entries(entries, team_then_key) == ( + *entries, + ("fast", "gpt-4.1-mini"), + ("hop", "gpt-4.1"), + ("mid", "gpt-4.1"), + ) + + +def test_one_bad_alias_entry_hides_only_itself() -> None: + aliases = {"fast": "gpt-4.1-mini", "broken": 5, 7: "gpt-4.1-mini"} + entries = [("gpt-4.1-mini", "gpt-4.1-mini")] + + assert alias_listing_entries(entries, _caller(aliases)) == (*entries, ("fast", "gpt-4.1-mini")) + assert alias_target("fast", _caller(aliases)) == "gpt-4.1-mini" + + +def test_chained_key_alias_is_listed_only_when_its_final_target_is_listable() -> None: + key_aliases = {"a": "b", "b": "hidden"} + entries = [("b", "b")] + + assert alias_listing_entries(entries, caller_alias_maps(key_aliases, None, "team-a", None)) == (*entries,) + assert alias_target("a", caller_alias_maps(key_aliases, None, "team-a", None)) == "hidden" + + +def test_team_aliases_only_apply_when_listing_the_team_the_key_authenticated_as( + monkeypatch: pytest.MonkeyPatch, +) -> None: + key_aliases, team_aliases, global_aliases = {"k": "gpt-4.1"}, {"t": "gpt-4.1-mini"}, {"g": "gpt-4.1"} + monkeypatch.setattr(litellm, "model_alias_map", global_aliases) + own_team = CallerAliases((team_aliases, key_aliases), (team_aliases, key_aliases, global_aliases, key_aliases)) + assert caller_alias_maps(key_aliases, team_aliases, "team-a", None) == own_team + assert caller_alias_maps(key_aliases, team_aliases, "team-a", "team-a") == own_team + assert caller_alias_maps(key_aliases, team_aliases, "team-a", "team-b") == CallerAliases( + (key_aliases,), (key_aliases, global_aliases, key_aliases) + ) + + +def test_global_alias_rewrites_between_the_two_key_passes_like_chat_completions( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "model_alias_map", {"b": "d"}) + key_aliases = {"a": "b", "b": "c"} + entries = [("c", "c"), ("d", "d")] + maps = caller_alias_maps(key_aliases, None, "team-a", None) + + assert alias_target("a", maps) == "d" + assert alias_listing_entries(entries, maps) == (*entries, ("a", "d"), ("b", "c")) + + +def test_global_aliases_rewrite_but_are_not_listed_as_caller_rows(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm, "model_alias_map", {"g": "gpt-4.1-mini"}) + entries = [("gpt-4.1-mini", "gpt-4.1-mini")] + maps = caller_alias_maps({"k": "g"}, None, "team-a", None) + + assert alias_listing_entries(entries, maps) == (*entries, ("k", "gpt-4.1-mini")) + assert alias_target("g", maps) == "gpt-4.1-mini" + + def test_team_public_name_uses_the_same_scope_at_list_and_request(): - router = Router(model_list=[{ - "model_name": "model_name_team-a_id", - "litellm_params": {"model": "openai/gpt-4o", "api_key": "sk-fake"}, - "model_info": {"team_id": "team-a", "team_public_model_name": "shared"}, - }]) - shown = claude_code_view_ids((_row("shared"),), {"user-agent": "claude-code/2.1.267"}, ClaudeCodeRoutingNames(router, "team-a"))["shared"] + router = Router( + model_list=[ + { + "model_name": "model_name_team-a_id", + "litellm_params": {"model": "openai/gpt-4o", "api_key": "sk-fake"}, + "model_info": {"team_id": "team-a", "team_public_model_name": "shared"}, + } + ] + ) + shown = claude_code_view_ids( + (_row("shared"),), {"user-agent": "claude-code/2.1.267"}, ClaudeCodeRoutingNames(router, "team-a") + )["shared"] assert claude_code_requested_group(shown, router, "team-a") == "shared" assert claude_code_requested_group(shown, router, "team-b") is None diff --git a/tests/test_litellm/proxy/test_model_list_aliases.py b/tests/test_litellm/proxy/test_model_list_aliases.py new file mode 100644 index 00000000000..25a941bf2ae --- /dev/null +++ b/tests/test_litellm/proxy/test_model_list_aliases.py @@ -0,0 +1,151 @@ +""" +Tests for key and team `model_aliases` on the model listing endpoints: GET /v1/models +(`model_list`, OpenAI and Anthropic shapes) and GET /v1/models/{id} (`model_info`). +An alias the caller can complete on is listed next to its target and resolves by name. +""" + +import pytest +from starlette.requests import Request + +from litellm import Router +from litellm.proxy import proxy_server +from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth + + +def _deployment(model_name: str, model: str = "openai/gpt-4.1-mini", **model_info: str | bool) -> dict[str, object]: + return { + "model_name": model_name, + "litellm_params": {"model": model, "api_key": "sk-fake"}, + "model_info": {"id": f"{model_name}-id", **model_info}, + } + + +@pytest.fixture +def router(monkeypatch: pytest.MonkeyPatch) -> Router: + router = Router( + model_list=[ + _deployment("gpt-4.1-mini"), + _deployment("gpt-4.1", model="openai/gpt-4.1"), + _deployment("model_name_team1_abc", team_id="team1", team_public_model_name="team-chat"), + _deployment("hidden", model="anthropic/claude-sonnet-4-5", discoverable=False), + ] + ) + monkeypatch.setattr(proxy_server, "llm_router", router) + monkeypatch.setattr(proxy_server, "llm_model_list", router.model_list) + monkeypatch.setattr(proxy_server, "prisma_client", None) + monkeypatch.setattr(proxy_server, "general_settings", {}) + monkeypatch.setattr(proxy_server, "user_model", None) + return router + + +def _team_member( + team_id: str = "team1", models: list[str] | None = None, **aliases: dict[str, str] | None +) -> UserAPIKeyAuth: + return UserAPIKeyAuth( + api_key="sk-test", + user_id="u", + user_role=LitellmUserRoles.INTERNAL_USER, + team_id=team_id, + team_models=["gpt-4.1-mini", "model_name_team1_abc"], + models=models or ["gpt-4.1-mini", "model_name_team1_abc"], + **aliases, + ) + + +def _anthropic_request(*extra_headers: tuple[bytes, bytes]) -> Request: + return Request( + scope={ + "type": "http", + "method": "GET", + "path": "/v1/models", + "query_string": b"", + "headers": [(b"anthropic-version", b"2023-06-01"), *extra_headers], + } + ) + + +def _claude_code_request() -> Request: + return _anthropic_request((b"user-agent", b"claude-cli/2.1.267 (external, cli)")) + + +async def _v1_models(user_api_key_dict: UserAPIKeyAuth, request: Request | None = None) -> list[str]: + response = await proxy_server.model_list(user_api_key_dict=user_api_key_dict, request=request) + return [m["id"] for m in response["data"]] + + +@pytest.mark.asyncio +async def test_v1_models_lists_team_alias_next_to_its_target_in_both_shapes(router: Router) -> None: + caller = _team_member(team_model_aliases={"claude-sonnet-4-5": "gpt-4.1-mini"}) + + assert await _v1_models(caller) == ["gpt-4.1-mini", "team-chat", "claude-sonnet-4-5"] + assert await _v1_models(caller, request=_anthropic_request()) == ["gpt-4.1-mini", "team-chat", "claude-sonnet-4-5"] + + +@pytest.mark.asyncio +async def test_claude_code_picker_lists_the_alias_under_its_own_name(router: Router) -> None: + caller = _team_member(team_model_aliases={"claude-sonnet-4-5": "gpt-4.1-mini"}) + + picker_ids = await _v1_models(caller, request=_claude_code_request()) + assert any(picker_id.startswith("claude-sonnet-4-5") for picker_id in picker_ids), picker_ids + + +@pytest.mark.asyncio +async def test_v1_models_lists_key_alias_and_hides_alias_to_a_model_the_caller_cannot_list(router: Router) -> None: + caller = _team_member(aliases={"mini": "gpt-4.1-mini", "big": "gpt-4.1"}) + + assert await _v1_models(caller) == ["gpt-4.1-mini", "team-chat", "mini"] + + +@pytest.mark.asyncio +async def test_v1_models_resolves_a_team_alias_through_the_key_alias_like_chat_completions_does(router: Router) -> None: + caller = _team_member(team_model_aliases={"fast": "mid"}, aliases={"fast": "gpt-4.1", "mid": "gpt-4.1-mini"}) + + assert await _v1_models(caller) == ["gpt-4.1-mini", "team-chat", "fast", "mid"] + response = await proxy_server.model_info(model_id="fast", user_api_key_dict=caller) + assert response["id"] == "fast" + + +@pytest.mark.asyncio +async def test_v1_models_skips_only_the_malformed_alias_entries(router: Router) -> None: + caller = _team_member(team_model_aliases={"claude-sonnet-4-5": 5, "fast": "gpt-4.1-mini"}) + + assert await _v1_models(caller) == ["gpt-4.1-mini", "team-chat", "fast"] + + +@pytest.mark.asyncio +async def test_v1_models_by_id_resolves_a_team_alias_to_its_target_metadata(router: Router) -> None: + caller = _team_member(team_model_aliases={"claude-sonnet-4-5": "gpt-4.1-mini"}) + + response = await proxy_server.model_info(model_id="claude-sonnet-4-5", user_api_key_dict=caller) + assert response["id"] == "claude-sonnet-4-5" + assert response["owned_by"] == "openai" + + +@pytest.mark.asyncio +async def test_v1_models_by_id_retrieves_the_listed_model_when_an_alias_collides_with_its_id(router: Router) -> None: + caller = _team_member(aliases={"team-chat": "gpt-4.1"}) + + assert await _v1_models(caller) == ["gpt-4.1-mini", "team-chat"] + response = await proxy_server.model_info(model_id="team-chat", user_api_key_dict=caller) + assert response["id"] == "team-chat" + + +@pytest.mark.asyncio +async def test_v1_models_by_id_resolves_an_alias_named_like_an_undiscoverable_model_to_the_alias_target( + router: Router, +) -> None: + caller = _team_member(aliases={"hidden": "gpt-4.1-mini"}, models=["gpt-4.1-mini", "model_name_team1_abc", "hidden"]) + + assert await _v1_models(caller) == ["gpt-4.1-mini", "team-chat", "hidden"] + target = await proxy_server.model_info(model_id="gpt-4.1-mini", user_api_key_dict=caller) + response = await proxy_server.model_info(model_id="hidden", user_api_key_dict=caller) + assert response == {**target, "id": "hidden"} + + +@pytest.mark.asyncio +async def test_v1_models_by_id_keeps_the_alias_as_id_when_it_targets_a_team_scoped_model(router: Router) -> None: + caller = _team_member(team_model_aliases={"chat": "team-chat"}) + + assert "chat" in await _v1_models(caller) + response = await proxy_server.model_info(model_id="chat", user_api_key_dict=caller) + assert response["id"] == "chat" From 785d4974ba177f90f4368530751b50193c85cdfd Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 24 Sep 2026 06:49:17 -0700 Subject: [PATCH 038/120] fix(cost-map): sync openrouter prices for deepseek v4 and glm-5.3 (#42952) Co-authored-by: kerry Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- ...odel_prices_and_context_window_backup.json | 22 +++++++++---------- model_prices_and_context_window.json | 22 +++++++++---------- 2 files changed, 22 insertions(+), 22 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index bba80677a80..a1878c00678 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -41236,21 +41236,21 @@ "supports_web_search": false }, "openrouter/deepseek/deepseek-v4-pro": { - "input_cost_per_token": 9.396e-07, + "input_cost_per_token": 9.31074e-07, "input_cost_per_token_cache_hit": 4.4e-08, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 384000, "max_tokens": 384000, "mode": "chat", - "output_cost_per_token": 1.8792e-06, + "output_cost_per_token": 1.862148e-06, "source": "https://openrouter.ai/api/v1/models", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "cache_read_input_token_cost": 7.83e-08, + "cache_read_input_token_cost": 7.75895e-08, "supports_audio_input": false, "supports_pdf_input": false, "supports_vision": false, @@ -65596,13 +65596,13 @@ "supports_web_search": false }, "openrouter/z-ai/glm-5.3": { - "input_cost_per_token": 8.4e-07, - "output_cost_per_token": 2.64e-06, - "cache_read_input_token_cost": 1.56e-07, + "input_cost_per_token": 1.4e-06, + "output_cost_per_token": 4.4e-06, + "cache_read_input_token_cost": 2.6e-07, "litellm_provider": "openrouter", "max_input_tokens": 1310720, - "max_output_tokens": 131072, - "max_tokens": 131072, + "max_output_tokens": 943717, + "max_tokens": 943717, "mode": "chat", "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, @@ -66307,9 +66307,9 @@ "supports_web_search": true }, "openrouter/deepseek/deepseek-v4-flash": { - "input_cost_per_token": 8.8606e-08, - "output_cost_per_token": 1.77212e-07, - "cache_read_input_token_cost": 1.77212e-08, + "input_cost_per_token": 8.554e-08, + "output_cost_per_token": 1.7108e-07, + "cache_read_input_token_cost": 1.7108e-08, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 384000, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index bba80677a80..a1878c00678 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -41236,21 +41236,21 @@ "supports_web_search": false }, "openrouter/deepseek/deepseek-v4-pro": { - "input_cost_per_token": 9.396e-07, + "input_cost_per_token": 9.31074e-07, "input_cost_per_token_cache_hit": 4.4e-08, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 384000, "max_tokens": 384000, "mode": "chat", - "output_cost_per_token": 1.8792e-06, + "output_cost_per_token": 1.862148e-06, "source": "https://openrouter.ai/api/v1/models", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "cache_read_input_token_cost": 7.83e-08, + "cache_read_input_token_cost": 7.75895e-08, "supports_audio_input": false, "supports_pdf_input": false, "supports_vision": false, @@ -65596,13 +65596,13 @@ "supports_web_search": false }, "openrouter/z-ai/glm-5.3": { - "input_cost_per_token": 8.4e-07, - "output_cost_per_token": 2.64e-06, - "cache_read_input_token_cost": 1.56e-07, + "input_cost_per_token": 1.4e-06, + "output_cost_per_token": 4.4e-06, + "cache_read_input_token_cost": 2.6e-07, "litellm_provider": "openrouter", "max_input_tokens": 1310720, - "max_output_tokens": 131072, - "max_tokens": 131072, + "max_output_tokens": 943717, + "max_tokens": 943717, "mode": "chat", "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, @@ -66307,9 +66307,9 @@ "supports_web_search": true }, "openrouter/deepseek/deepseek-v4-flash": { - "input_cost_per_token": 8.8606e-08, - "output_cost_per_token": 1.77212e-07, - "cache_read_input_token_cost": 1.77212e-08, + "input_cost_per_token": 8.554e-08, + "output_cost_per_token": 1.7108e-07, + "cache_read_input_token_cost": 1.7108e-08, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 384000, From 3acdfda19c37b88d2011982d3fb6d5850ce43d49 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 24 Sep 2026 06:52:06 -0700 Subject: [PATCH 039/120] fix(cost-map): add batch prices for vertex gemini-3.8-flash-cyber (#42953) Co-authored-by: kerry Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/model_prices_and_context_window_backup.json | 6 ++++++ model_prices_and_context_window.json | 6 ++++++ 2 files changed, 12 insertions(+) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index a1878c00678..b5282f12539 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -27350,9 +27350,11 @@ }, "vertex_ai/gemini-3.8-flash-cyber": { "cache_read_input_token_cost": 1.5e-07, + "cache_read_input_token_cost_batches": 7.5e-08, "cache_read_input_token_cost_flex": 7.5e-08, "cache_read_input_token_cost_priority": 2.7e-07, "input_cost_per_token": 1.5e-06, + "input_cost_per_token_batches": 7.5e-07, "input_cost_per_token_flex": 7.5e-07, "input_cost_per_token_priority": 2.7e-06, "litellm_provider": "vertex_ai", @@ -27362,6 +27364,7 @@ "mode": "chat", "output_cost_per_reasoning_token": 7.5e-06, "output_cost_per_token": 7.5e-06, + "output_cost_per_token_batches": 3.75e-06, "output_cost_per_token_flex": 3.75e-06, "output_cost_per_token_priority": 1.35e-05, "regional_endpoint_uplift_multiplier": 1.1, @@ -29604,9 +29607,11 @@ }, "gemini-3.8-flash-cyber": { "cache_read_input_token_cost": 1.5e-07, + "cache_read_input_token_cost_batches": 7.5e-08, "cache_read_input_token_cost_flex": 7.5e-08, "cache_read_input_token_cost_priority": 2.7e-07, "input_cost_per_token": 1.5e-06, + "input_cost_per_token_batches": 7.5e-07, "input_cost_per_token_flex": 7.5e-07, "input_cost_per_token_priority": 2.7e-06, "litellm_provider": "vertex_ai-language-models", @@ -29616,6 +29621,7 @@ "mode": "chat", "output_cost_per_reasoning_token": 7.5e-06, "output_cost_per_token": 7.5e-06, + "output_cost_per_token_batches": 3.75e-06, "output_cost_per_token_flex": 3.75e-06, "output_cost_per_token_priority": 1.35e-05, "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index a1878c00678..b5282f12539 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -27350,9 +27350,11 @@ }, "vertex_ai/gemini-3.8-flash-cyber": { "cache_read_input_token_cost": 1.5e-07, + "cache_read_input_token_cost_batches": 7.5e-08, "cache_read_input_token_cost_flex": 7.5e-08, "cache_read_input_token_cost_priority": 2.7e-07, "input_cost_per_token": 1.5e-06, + "input_cost_per_token_batches": 7.5e-07, "input_cost_per_token_flex": 7.5e-07, "input_cost_per_token_priority": 2.7e-06, "litellm_provider": "vertex_ai", @@ -27362,6 +27364,7 @@ "mode": "chat", "output_cost_per_reasoning_token": 7.5e-06, "output_cost_per_token": 7.5e-06, + "output_cost_per_token_batches": 3.75e-06, "output_cost_per_token_flex": 3.75e-06, "output_cost_per_token_priority": 1.35e-05, "regional_endpoint_uplift_multiplier": 1.1, @@ -29604,9 +29607,11 @@ }, "gemini-3.8-flash-cyber": { "cache_read_input_token_cost": 1.5e-07, + "cache_read_input_token_cost_batches": 7.5e-08, "cache_read_input_token_cost_flex": 7.5e-08, "cache_read_input_token_cost_priority": 2.7e-07, "input_cost_per_token": 1.5e-06, + "input_cost_per_token_batches": 7.5e-07, "input_cost_per_token_flex": 7.5e-07, "input_cost_per_token_priority": 2.7e-06, "litellm_provider": "vertex_ai-language-models", @@ -29616,6 +29621,7 @@ "mode": "chat", "output_cost_per_reasoning_token": 7.5e-06, "output_cost_per_token": 7.5e-06, + "output_cost_per_token_batches": 3.75e-06, "output_cost_per_token_flex": 3.75e-06, "output_cost_per_token_priority": 1.35e-05, "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", From 77c7a870d542131ac6dd0fdea603d99bb88cdc84 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 24 Sep 2026 06:55:14 -0700 Subject: [PATCH 040/120] fix(cost-map): add azure realtime, audio and partner model rows (#42954) Co-authored-by: kerry Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- ...odel_prices_and_context_window_backup.json | 217 ++++++++++++++++++ model_prices_and_context_window.json | 217 ++++++++++++++++++ 2 files changed, 434 insertions(+) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index b5282f12539..d87a21e6d1f 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -63384,6 +63384,115 @@ "supports_tool_choice": true, "supports_vision": true }, + "azure_ai/FW-DeepSeek-V4.1-Flash": { + "cache_read_input_token_cost": 8e-09, + "input_cost_per_token": 3.75e-07, + "litellm_provider": "azure_ai", + "max_input_tokens": 1000000, + "max_output_tokens": 384000, + "max_tokens": 384000, + "mode": "chat", + "output_cost_per_token": 1.5e-06, + "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true + }, + "azure_ai/FW-DeepSeek-V4-Flash": { + "cache_read_input_token_cost": 3e-08, + "input_cost_per_token": 1.5e-07, + "litellm_provider": "azure_ai", + "max_input_tokens": 1000000, + "max_output_tokens": 384000, + "max_tokens": 384000, + "mode": "chat", + "output_cost_per_token": 3.1e-07, + "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true + }, + "azure_ai/FW-GLM-5.3": { + "cache_read_input_token_cost": 3.25e-07, + "input_cost_per_token": 1.75e-06, + "litellm_provider": "azure_ai", + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 5.5e-06, + "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true + }, + "azure_ai/FW-GLM-5.3-Flash": { + "cache_read_input_token_cost": 3.8e-08, + "input_cost_per_token": 1.88e-07, + "litellm_provider": "azure_ai", + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 6.25e-07, + "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true + }, + "azure_ai/FW-GPT-OSS-120B": { + "cache_read_input_token_cost": 8.2e-08, + "input_cost_per_token": 1.65e-07, + "litellm_provider": "azure_ai", + "max_input_tokens": 131072, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 6.6e-07, + "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, + "azure_ai/Cohere-command-a-plus-05-2026": { + "input_cost_per_token": 8e-07, + "litellm_provider": "azure_ai", + "max_input_tokens": 128000, + "max_output_tokens": 64000, + "max_tokens": 64000, + "mode": "chat", + "output_cost_per_token": 3.2e-06, + "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true + }, + "azure_ai/mistral-medium-3-5": { + "input_cost_per_token": 1.5e-06, + "litellm_provider": "azure_ai", + "max_input_tokens": 128000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 7.5e-06, + "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_response_schema": true, + "supports_vision": true + }, "bedrock/us-gov-west-1/nvidia.nemotron-nano-3-30b": { "input_cost_per_token": 7.2e-08, "litellm_provider": "bedrock", @@ -69433,6 +69542,114 @@ "mode": "embedding", "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" }, + "azure/gpt-realtime-2": { + "cache_creation_input_audio_token_cost": 4e-07, + "cache_read_input_audio_token_cost": 4e-07, + "cache_read_input_token_cost": 4e-07, + "input_cost_per_audio_token": 3.2e-05, + "input_cost_per_image_token": 5e-06, + "input_cost_per_token": 4e-06, + "litellm_provider": "azure", + "max_input_tokens": 32000, + "max_output_tokens": 4096, + "max_tokens": 4096, + "mode": "realtime", + "output_cost_per_audio_token": 6.4e-05, + "output_cost_per_token": 2.4e-05, + "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", + "supported_endpoints": [ + "/v1/realtime" + ], + "supported_modalities": [ + "text", + "image", + "audio" + ], + "supported_output_modalities": [ + "text", + "audio" + ], + "supports_audio_input": true, + "supports_audio_output": true, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_system_messages": true, + "supports_tool_choice": true + }, + "azure/gpt-live-1": { + "input_cost_per_second": 0.000833333333333, + "litellm_provider": "azure", + "mode": "realtime", + "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", + "supported_modalities": [ + "text", + "audio" + ], + "supported_output_modalities": [ + "text", + "audio" + ], + "supports_audio_input": true, + "supports_audio_output": true, + "supports_function_calling": true + }, + "azure/gpt-live-transcribe": { + "input_cost_per_second": 0.000283333333333, + "litellm_provider": "azure", + "max_input_tokens": 32000, + "max_output_tokens": 4096, + "max_tokens": 4096, + "mode": "audio_transcription", + "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", + "supported_endpoints": [ + "/v1/realtime", + "/v1/realtime/transcription_sessions" + ], + "supported_modalities": [ + "text", + "audio" + ], + "supported_output_modalities": [ + "text" + ], + "supports_audio_input": true + }, + "azure/gpt-transcribe": { + "input_cost_per_second": 7.5e-05, + "litellm_provider": "azure", + "mode": "audio_transcription", + "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", + "supported_endpoints": [ + "/v1/audio/transcriptions", + "/v1/realtime/transcription_sessions" + ], + "supported_modalities": [ + "text", + "audio" + ], + "supported_output_modalities": [ + "text" + ], + "supports_audio_input": true + }, + "azure/gpt-realtime-translate": { + "input_cost_per_second": 0.000566666666667, + "litellm_provider": "azure", + "max_input_tokens": 32000, + "max_output_tokens": 4096, + "max_tokens": 4096, + "mode": "realtime", + "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", + "supported_modalities": [ + "audio" + ], + "supported_output_modalities": [ + "text", + "audio" + ], + "supports_audio_input": true, + "supports_audio_output": true + }, "aihubmix/agnes-2.5-flash": { "input_cost_per_token": 3e-08, "litellm_provider": "aihubmix", diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index b5282f12539..d87a21e6d1f 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -63384,6 +63384,115 @@ "supports_tool_choice": true, "supports_vision": true }, + "azure_ai/FW-DeepSeek-V4.1-Flash": { + "cache_read_input_token_cost": 8e-09, + "input_cost_per_token": 3.75e-07, + "litellm_provider": "azure_ai", + "max_input_tokens": 1000000, + "max_output_tokens": 384000, + "max_tokens": 384000, + "mode": "chat", + "output_cost_per_token": 1.5e-06, + "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true + }, + "azure_ai/FW-DeepSeek-V4-Flash": { + "cache_read_input_token_cost": 3e-08, + "input_cost_per_token": 1.5e-07, + "litellm_provider": "azure_ai", + "max_input_tokens": 1000000, + "max_output_tokens": 384000, + "max_tokens": 384000, + "mode": "chat", + "output_cost_per_token": 3.1e-07, + "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true + }, + "azure_ai/FW-GLM-5.3": { + "cache_read_input_token_cost": 3.25e-07, + "input_cost_per_token": 1.75e-06, + "litellm_provider": "azure_ai", + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 5.5e-06, + "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true + }, + "azure_ai/FW-GLM-5.3-Flash": { + "cache_read_input_token_cost": 3.8e-08, + "input_cost_per_token": 1.88e-07, + "litellm_provider": "azure_ai", + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 6.25e-07, + "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true + }, + "azure_ai/FW-GPT-OSS-120B": { + "cache_read_input_token_cost": 8.2e-08, + "input_cost_per_token": 1.65e-07, + "litellm_provider": "azure_ai", + "max_input_tokens": 131072, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 6.6e-07, + "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, + "azure_ai/Cohere-command-a-plus-05-2026": { + "input_cost_per_token": 8e-07, + "litellm_provider": "azure_ai", + "max_input_tokens": 128000, + "max_output_tokens": 64000, + "max_tokens": 64000, + "mode": "chat", + "output_cost_per_token": 3.2e-06, + "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true + }, + "azure_ai/mistral-medium-3-5": { + "input_cost_per_token": 1.5e-06, + "litellm_provider": "azure_ai", + "max_input_tokens": 128000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 7.5e-06, + "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_response_schema": true, + "supports_vision": true + }, "bedrock/us-gov-west-1/nvidia.nemotron-nano-3-30b": { "input_cost_per_token": 7.2e-08, "litellm_provider": "bedrock", @@ -69433,6 +69542,114 @@ "mode": "embedding", "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" }, + "azure/gpt-realtime-2": { + "cache_creation_input_audio_token_cost": 4e-07, + "cache_read_input_audio_token_cost": 4e-07, + "cache_read_input_token_cost": 4e-07, + "input_cost_per_audio_token": 3.2e-05, + "input_cost_per_image_token": 5e-06, + "input_cost_per_token": 4e-06, + "litellm_provider": "azure", + "max_input_tokens": 32000, + "max_output_tokens": 4096, + "max_tokens": 4096, + "mode": "realtime", + "output_cost_per_audio_token": 6.4e-05, + "output_cost_per_token": 2.4e-05, + "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", + "supported_endpoints": [ + "/v1/realtime" + ], + "supported_modalities": [ + "text", + "image", + "audio" + ], + "supported_output_modalities": [ + "text", + "audio" + ], + "supports_audio_input": true, + "supports_audio_output": true, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_system_messages": true, + "supports_tool_choice": true + }, + "azure/gpt-live-1": { + "input_cost_per_second": 0.000833333333333, + "litellm_provider": "azure", + "mode": "realtime", + "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", + "supported_modalities": [ + "text", + "audio" + ], + "supported_output_modalities": [ + "text", + "audio" + ], + "supports_audio_input": true, + "supports_audio_output": true, + "supports_function_calling": true + }, + "azure/gpt-live-transcribe": { + "input_cost_per_second": 0.000283333333333, + "litellm_provider": "azure", + "max_input_tokens": 32000, + "max_output_tokens": 4096, + "max_tokens": 4096, + "mode": "audio_transcription", + "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", + "supported_endpoints": [ + "/v1/realtime", + "/v1/realtime/transcription_sessions" + ], + "supported_modalities": [ + "text", + "audio" + ], + "supported_output_modalities": [ + "text" + ], + "supports_audio_input": true + }, + "azure/gpt-transcribe": { + "input_cost_per_second": 7.5e-05, + "litellm_provider": "azure", + "mode": "audio_transcription", + "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", + "supported_endpoints": [ + "/v1/audio/transcriptions", + "/v1/realtime/transcription_sessions" + ], + "supported_modalities": [ + "text", + "audio" + ], + "supported_output_modalities": [ + "text" + ], + "supports_audio_input": true + }, + "azure/gpt-realtime-translate": { + "input_cost_per_second": 0.000566666666667, + "litellm_provider": "azure", + "max_input_tokens": 32000, + "max_output_tokens": 4096, + "max_tokens": 4096, + "mode": "realtime", + "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", + "supported_modalities": [ + "audio" + ], + "supported_output_modalities": [ + "text", + "audio" + ], + "supports_audio_input": true, + "supports_audio_output": true + }, "aihubmix/agnes-2.5-flash": { "input_cost_per_token": 3e-08, "litellm_provider": "aihubmix", From 11ed3335c8be943da40815bc9bd3ae5990176c80 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 24 Sep 2026 07:18:08 -0700 Subject: [PATCH 041/120] fix(cost-map): sync openrouter deepseek-v4-pro prices (#42956) Co-authored-by: kerry Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/model_prices_and_context_window_backup.json | 6 +++--- model_prices_and_context_window.json | 6 +++--- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index d87a21e6d1f..6257405ed7f 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -41242,21 +41242,21 @@ "supports_web_search": false }, "openrouter/deepseek/deepseek-v4-pro": { - "input_cost_per_token": 9.31074e-07, + "input_cost_per_token": 9.27768e-07, "input_cost_per_token_cache_hit": 4.4e-08, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 384000, "max_tokens": 384000, "mode": "chat", - "output_cost_per_token": 1.862148e-06, + "output_cost_per_token": 1.855536e-06, "source": "https://openrouter.ai/api/v1/models", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "cache_read_input_token_cost": 7.75895e-08, + "cache_read_input_token_cost": 7.7314e-08, "supports_audio_input": false, "supports_pdf_input": false, "supports_vision": false, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index d87a21e6d1f..6257405ed7f 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -41242,21 +41242,21 @@ "supports_web_search": false }, "openrouter/deepseek/deepseek-v4-pro": { - "input_cost_per_token": 9.31074e-07, + "input_cost_per_token": 9.27768e-07, "input_cost_per_token_cache_hit": 4.4e-08, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 384000, "max_tokens": 384000, "mode": "chat", - "output_cost_per_token": 1.862148e-06, + "output_cost_per_token": 1.855536e-06, "source": "https://openrouter.ai/api/v1/models", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "cache_read_input_token_cost": 7.75895e-08, + "cache_read_input_token_cost": 7.7314e-08, "supports_audio_input": false, "supports_pdf_input": false, "supports_vision": false, From b8f3ba03b30b7cb19e21820b9818f32935401ccd Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 24 Sep 2026 07:57:26 -0700 Subject: [PATCH 042/120] fix(cost-map): sync openrouter deepseek-v4-pro prices (#42964) Co-authored-by: kerry Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/model_prices_and_context_window_backup.json | 6 +++--- model_prices_and_context_window.json | 6 +++--- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 6257405ed7f..69a64de6e04 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -41242,21 +41242,21 @@ "supports_web_search": false }, "openrouter/deepseek/deepseek-v4-pro": { - "input_cost_per_token": 9.27768e-07, + "input_cost_per_token": 9.24462e-07, "input_cost_per_token_cache_hit": 4.4e-08, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 384000, "max_tokens": 384000, "mode": "chat", - "output_cost_per_token": 1.855536e-06, + "output_cost_per_token": 1.848924e-06, "source": "https://openrouter.ai/api/v1/models", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "cache_read_input_token_cost": 7.7314e-08, + "cache_read_input_token_cost": 7.70385e-08, "supports_audio_input": false, "supports_pdf_input": false, "supports_vision": false, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 6257405ed7f..69a64de6e04 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -41242,21 +41242,21 @@ "supports_web_search": false }, "openrouter/deepseek/deepseek-v4-pro": { - "input_cost_per_token": 9.27768e-07, + "input_cost_per_token": 9.24462e-07, "input_cost_per_token_cache_hit": 4.4e-08, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 384000, "max_tokens": 384000, "mode": "chat", - "output_cost_per_token": 1.855536e-06, + "output_cost_per_token": 1.848924e-06, "source": "https://openrouter.ai/api/v1/models", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "cache_read_input_token_cost": 7.7314e-08, + "cache_read_input_token_cost": 7.70385e-08, "supports_audio_input": false, "supports_pdf_input": false, "supports_vision": false, From e135a199ad840336f60b61426967444bbd71cea5 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 24 Sep 2026 10:09:49 -0500 Subject: [PATCH 043/120] fix(proxy): fail parked DB lookups at a deadline and flip readiness while they stall (#42654) * fix(proxy): fail parked DB lookups at a deadline and flip readiness while they stall Under a load burst with a slow authentication database every request parked inside the pod with no deadline while /health/readiness kept answering 200 (its own ping gets a fresh connection), so the load balancer kept sending traffic until the pod hit its memory limit, and the parked requests completed against the provider minutes after every client had hung up Every pre-request read (key, team, user, end user, budget, membership, organization, object permission, jwt mapping, project, proxy budget, spend counter reseed) now runs under one deadline, PROXY_DB_LOOKUP_DEADLINE_SECONDS (default 10 s). A lookup that hits it fails the request with the existing 503 "authentication database is temporarily unreachable" answer, honours allow_requests_on_db_unavailable, and never triggers the transport reconnect (the transport is fine, the query is slow), which is what turned the repro's stall into "too many clients". Writes stay unbounded A deadline hit marks the pod stalled for PROXY_DB_LOOKUP_STALL_WINDOW_SECONDS (default 30 s, 0 disables), during which /health/readiness answers 503 with "db": "stalled" behind the same fail-open gate, so the pod leaves rotation before it fills its memory. The existing litellm_in_flight_requests gauge already exposes the parked set on /metrics The deadline is enforced on the wall clock: bounded_db_lookup waits on the lookup task with asyncio.wait and raises DBLookupDeadlineExceeded when the deadline passes even if the lookup absorbs its cancellation, where asyncio.wait_for on 3.12+ would sit on the cancelled task for as long as it takes The failure spend-log row no longer re-runs the key and team lookups when the failure itself is a database connection or deadline error, so a request that hit the deadline is answered after one deadline instead of two * fix(proxy): bound the spend counter gate wait and narrow the stalled lookup shortcut Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): keep the global spend lookup on the prisma client handle 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> Co-authored-by: yassin Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/constants.py | 2 + litellm/proxy/auth/auth_checks.py | 119 ++++- litellm/proxy/auth/user_api_key_auth.py | 6 +- litellm/proxy/db/db_lookup_gate.py | 71 ++- litellm/proxy/db/exception_handler.py | 3 +- litellm/proxy/db/spend_counter_reseed.py | 71 +-- .../health_endpoints/_health_endpoints.py | 33 +- .../proxy/hooks/proxy_track_cost_callback.py | 11 +- .../proxy/auth/test_auth_checks.py | 163 +++++- .../proxy/auth/test_user_api_key_auth.py | 486 +++++++++--------- .../proxy/db/test_db_lookup_gate.py | 111 ++++ .../proxy/db/test_exception_handler.py | 15 + .../proxy/db/test_spend_counter_reseed.py | 27 + .../health_endpoints/test_health_endpoints.py | 128 +++++ .../hooks/test_proxy_track_cost_callback.py | 100 +++- 15 files changed, 1016 insertions(+), 330 deletions(-) create mode 100644 tests/test_litellm/proxy/db/test_db_lookup_gate.py diff --git a/litellm/constants.py b/litellm/constants.py index 3ff80d8b7dd..7b40f432446 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1772,6 +1772,8 @@ RESPONSES_SESSION_LOOKUP_MAX_ATTEMPTS: Final = max(1, int(os.getenv("RESPONSES_S RESPONSES_SESSION_LOOKUP_RETRY_INTERVAL: Final = float(os.getenv("RESPONSES_SESSION_LOOKUP_RETRY_INTERVAL", "0.2")) SPEND_COUNTER_RESEED_LOCKS_MAX_SIZE: Final = int(os.getenv("SPEND_COUNTER_RESEED_LOCKS_MAX_SIZE", 10000)) PROXY_DB_LOOKUP_MAX_CONCURRENCY: Final = max(1, int(os.getenv("PROXY_DB_LOOKUP_MAX_CONCURRENCY", "25"))) +PROXY_DB_LOOKUP_DEADLINE_SECONDS: Final = max(0.1, float(os.getenv("PROXY_DB_LOOKUP_DEADLINE_SECONDS", "10"))) +PROXY_DB_LOOKUP_STALL_WINDOW_SECONDS: Final = max(0.0, float(os.getenv("PROXY_DB_LOOKUP_STALL_WINDOW_SECONDS", "30"))) DEFAULT_CRON_JOB_LOCK_TTL_SECONDS: Final = int(os.getenv("DEFAULT_CRON_JOB_LOCK_TTL_SECONDS", 60)) # 1 minute PROXY_BUDGET_RESCHEDULER_MIN_TIME: Final = int(os.getenv("PROXY_BUDGET_RESCHEDULER_MIN_TIME", 597)) RESET_BUDGET_JOB_BATCH_SIZE: Final = max(1, int(os.getenv("RESET_BUDGET_JOB_BATCH_SIZE", "500"))) diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index f5e98b40ca7..67950e603c0 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -15,11 +15,11 @@ import re import time from collections.abc import Awaitable, Callable, Iterator, Mapping, Sequence from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final, Literal, Optional, Protocol, TypeAlias +from typing import TYPE_CHECKING, Any, Final, Generic, Literal, Optional, Protocol, TypeAlias from fastapi import HTTPException, Request, status from pydantic import BaseModel, TypeAdapter -from typing_extensions import ReadOnly, TypedDict +from typing_extensions import NotRequired, ReadOnly, Required, TypedDict, Unpack import litellm from litellm._logging import verbose_proxy_logger @@ -110,7 +110,7 @@ from litellm.proxy.common_utils.user_api_key_cache import ( team_membership_auth_cache_key, team_membership_reservation_cache_key, ) -from litellm.proxy.db.db_lookup_gate import db_lookup_gate +from litellm.proxy.db.db_lookup_gate import bounded_db_lookup, db_lookup_gate from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler from litellm.proxy.guardrails.tool_name_extraction import ( TOOL_CAPABLE_CALL_TYPES, @@ -223,30 +223,79 @@ class _PrismaTableHolder(Protocol[RowT_co]): def table(self) -> _PrismaAuthTable[RowT_co]: ... -def _dictable_table(repo: _PrismaTableHolder[_PrismaDictableRow]) -> _PrismaAuthTable[_PrismaDictableRow]: - return repo.table +class _FindOneKwargs(TypedDict): + where: ReadOnly[Required[Mapping[str, object]]] + include: ReadOnly[NotRequired[Mapping[str, object] | None]] + + +class _FindManyKwargs(TypedDict): + where: ReadOnly[NotRequired[Mapping[str, object] | None]] + include: ReadOnly[NotRequired[Mapping[str, object] | None]] + take: ReadOnly[NotRequired[int | None]] + + +class _DeadlineBoundedTable(Generic[RowT_co]): + """Every read on the wrapped table fails with ``DBLookupDeadlineExceeded`` once + ``PROXY_DB_LOOKUP_DEADLINE_SECONDS`` passes, so a stalled database fails the + request fast instead of parking it in the pod until it fills its memory.""" + + __slots__ = ("_lookup", "_table") + + def __init__(self, table: _PrismaAuthTable[RowT_co], lookup: str) -> None: + self._table: Final = table + self._lookup: Final = lookup + + async def find_unique( + self, + **kwargs: Unpack[_FindOneKwargs], # kwargs-ok: typed pass-through that forwards exactly what the caller passed + ) -> RowT_co | None: + return await bounded_db_lookup(self._table.find_unique(**kwargs), name=self._lookup) + + async def find_first( + self, + **kwargs: Unpack[_FindOneKwargs], # kwargs-ok: typed pass-through that forwards exactly what the caller passed + ) -> RowT_co | None: + return await bounded_db_lookup(self._table.find_first(**kwargs), name=self._lookup) + + async def find_many( + self, + **kwargs: Unpack[_FindManyKwargs], # kwargs-ok: typed pass-through that forwards exactly what the caller passed + ) -> Sequence[RowT_co]: + return await bounded_db_lookup(self._table.find_many(**kwargs), name=self._lookup) + + async def update(self, *, where: Mapping[str, object], data: Mapping[str, object]) -> RowT_co | None: + return await self._table.update(where=where, data=data) + + async def create(self, *, data: Mapping[str, object], include: Mapping[str, object] | None = None) -> RowT_co: + return await self._table.create(data=data, include=include) + + +def _dictable_table(repo: _PrismaTableHolder[_PrismaDictableRow], lookup: str) -> _PrismaAuthTable[_PrismaDictableRow]: + return _DeadlineBoundedTable(repo.table, lookup) def _jwt_key_mapping_table( repo: _PrismaTableHolder[_PrismaJWTKeyMappingRow], ) -> _PrismaAuthTable[_PrismaJWTKeyMappingRow]: - return repo.table + return _DeadlineBoundedTable(repo.table, "jwt_key_mapping") -def _model_dump_table(repo: _PrismaTableHolder[_PrismaModelDumpRow]) -> _PrismaAuthTable[_PrismaModelDumpRow]: - return repo.table +def _model_dump_table( + repo: _PrismaTableHolder[_PrismaModelDumpRow], lookup: str +) -> _PrismaAuthTable[_PrismaModelDumpRow]: + return _DeadlineBoundedTable(repo.table, lookup) def _team_table(repo: _PrismaTableHolder[_PrismaTeamRow]) -> _PrismaAuthTable[_PrismaTeamRow]: - return repo.table + return _DeadlineBoundedTable(repo.table, "team") def _vector_store_table(repo: _PrismaTableHolder[_PrismaVectorStoreRow]) -> _PrismaAuthTable[_PrismaVectorStoreRow]: - return repo.table + return _DeadlineBoundedTable(repo.table, "vector_store") def _user_table(repo: _PrismaTableHolder[_PrismaUserRow]) -> _PrismaAuthTable[_PrismaUserRow]: - return repo.table + return _DeadlineBoundedTable(repo.table, "user") class _VectorStorePermissionsRow(Protocol): @@ -257,7 +306,7 @@ class _VectorStorePermissionsRow(Protocol): def _object_permission_table( repo: _PrismaTableHolder[_VectorStorePermissionsRow], ) -> _PrismaAuthTable[_VectorStorePermissionsRow]: - return repo.table + return _DeadlineBoundedTable(repo.table, "object_permission") class _PrismaTagRow(Protocol): @@ -1422,7 +1471,7 @@ async def get_default_end_user_budget( # Fetch from database try: - budget_record: Final = await _dictable_table(BudgetRepository(prisma_client)).find_unique( + budget_record: Final = await _dictable_table(BudgetRepository(prisma_client), "budget").find_unique( where={"budget_id": default_budget_id} # mutable-ok: prisma where clause ) @@ -1483,7 +1532,7 @@ async def get_team_member_default_budget( return cached_budget try: - budget_record: Final = await _dictable_table(BudgetRepository(prisma_client)).find_unique( + budget_record: Final = await _dictable_table(BudgetRepository(prisma_client), "budget").find_unique( where={"budget_id": budget_id} ) except Exception: @@ -1877,7 +1926,7 @@ async def get_end_user_object( # Fetch from database try: - response: Final = await _dictable_table(EndUserRepository(prisma_client)).find_unique( + response: Final = await _dictable_table(EndUserRepository(prisma_client), "end_user").find_unique( where={"user_id": end_user_id}, include={"litellm_budget_table": True, "object_permission": True}, ) @@ -2286,7 +2335,7 @@ async def _fetch_team_membership_from_db( proxy_logging_obj: ProxyLogging | None = None, ) -> LiteLLM_TeamMembership | None: _ = parent_otel_span, proxy_logging_obj - response: Final = await _dictable_table(TeamMembershipRepository(prisma_client)).find_unique( + response: Final = await _dictable_table(TeamMembershipRepository(prisma_client), "team_membership").find_unique( where={"user_id_team_id": {"user_id": user_id, "team_id": team_id}}, include={"litellm_budget_table": True}, ) @@ -3290,7 +3339,7 @@ async def get_access_object( # Not in cache - fetch from DB try: - response: Final = await _dictable_table(AccessGroupRepository(prisma_client)).find_unique( + response: Final = await _dictable_table(AccessGroupRepository(prisma_client), "access_group").find_unique( where={"access_group_id": access_group_id} ) @@ -3472,7 +3521,7 @@ async def get_org_object_by_alias( # Query database by organization_alias try: - orgs = await _model_dump_table(OrganizationRepository(prisma_client)).find_many( + orgs = await _model_dump_table(OrganizationRepository(prisma_client), "organization").find_many( where={"organization_alias": org_alias} ) @@ -3650,10 +3699,32 @@ async def _fetch_key_object_from_db_with_reconnect( prisma_client: PrismaClient, parent_otel_span: Span | None, proxy_logging_obj: ProxyLogging | None, + deadline_seconds: float | None = None, ) -> BaseModel | None: """ Fetch key object from DB and retry once if a DB connection error can be healed. + The gate wait, the query, the reconnect, and the retry share one deadline, so a + stalled database fails the request with ``DBLookupDeadlineExceeded`` instead of + parking it. """ + return await bounded_db_lookup( + _fetch_key_object_from_db_unbounded( + hashed_token=hashed_token, + prisma_client=prisma_client, + parent_otel_span=parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + ), + name="key", + deadline_seconds=deadline_seconds, + ) + + +async def _fetch_key_object_from_db_unbounded( + hashed_token: str, + prisma_client: PrismaClient, + parent_otel_span: Span | None, + proxy_logging_obj: ProxyLogging | None, +) -> BaseModel | None: async with db_lookup_gate.current(): try: return await prisma_client.get_data( @@ -3874,9 +3945,9 @@ async def get_object_permission( # else, check db try: - response: Final = await _dictable_table(ObjectPermissionRepository(prisma_client)).find_unique( - where={"object_permission_id": object_permission_id} - ) + response: Final = await _dictable_table( + ObjectPermissionRepository(prisma_client), "object_permission" + ).find_unique(where={"object_permission_id": object_permission_id}) if response is None: return None @@ -4008,7 +4079,9 @@ async def get_org_object( if include_budget_table: query_kwargs["include"] = {"litellm_budget_table": True} - response: Final = await _model_dump_table(OrganizationRepository(prisma_client)).find_unique(**query_kwargs) + response: Final = await _model_dump_table(OrganizationRepository(prisma_client), "organization").find_unique( + **query_kwargs + ) except Exception: # An operational failure (DB down, timeout, cache fault) is NOT the same fact as a confirmed # missing row, and relabelling it as "doesn't exist" made every caller unable to tell them @@ -5948,7 +6021,7 @@ async def get_project_object( return deserialized_project # Fetch from DB - project_row: Final = await _model_dump_table(ProjectRepository(prisma_client)).find_unique( + project_row: Final = await _model_dump_table(ProjectRepository(prisma_client), "project").find_unique( where={"project_id": project_id}, include={"litellm_budget_table": True}, ) diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index a6c0792a86f..ae95e94dd2d 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -120,6 +120,7 @@ from litellm.proxy.common_utils.user_api_key_cache import ( UserApiKeyCache, team_membership_auth_cache_key, ) +from litellm.proxy.db.db_lookup_gate import bounded_db_lookup from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup from litellm.proxy.spend_tracking.carried_budget_state import carry_team_and_user_budget_state @@ -735,8 +736,9 @@ async def _fetch_global_spend_with_event_coordination( """ async def _load_global_spend() -> float | None: - proxy_budget_row: Final = await prisma_client.db.litellm_usertable.find_unique( - where={"user_id": LITELLM_PROXY_BUDGET_NAME} + proxy_budget_row: Final = await bounded_db_lookup( + prisma_client.db.litellm_usertable.find_unique(where={"user_id": LITELLM_PROXY_BUDGET_NAME}), + name="proxy_budget", ) return float(proxy_budget_row.spend) if proxy_budget_row is not None else None diff --git a/litellm/proxy/db/db_lookup_gate.py b/litellm/proxy/db/db_lookup_gate.py index 2fd427687bd..d2aefde929f 100644 --- a/litellm/proxy/db/db_lookup_gate.py +++ b/litellm/proxy/db/db_lookup_gate.py @@ -1,7 +1,14 @@ import asyncio -from typing import Final +import time +from collections.abc import Awaitable, Callable +from typing import Final, TypeVar -from litellm.constants import PROXY_DB_LOOKUP_MAX_CONCURRENCY +from litellm.constants import ( + PROXY_DB_LOOKUP_DEADLINE_SECONDS, + PROXY_DB_LOOKUP_MAX_CONCURRENCY, +) + +LookupT = TypeVar("LookupT") class LoopBoundSemaphore: @@ -20,4 +27,64 @@ class LoopBoundSemaphore: return self._semaphore +class DBLookupDeadlineExceeded(asyncio.TimeoutError): + def __init__(self, lookup: str, deadline_seconds: float) -> None: + super().__init__(f"{lookup} lookup did not answer within {deadline_seconds:g}s") + self.lookup: Final = lookup + self.deadline_seconds: Final = deadline_seconds + + +class DBLookupStallTracker: + __slots__ = ("_clock", "_last_hit") + + def __init__(self, clock: Callable[[], float] = time.monotonic) -> None: + self._clock: Final = clock + self._last_hit: float | None = None + + def record_hit(self) -> None: + self._last_hit = self._clock() + + def clear(self) -> None: + self._last_hit = None + + def stalled_within(self, window_seconds: float) -> bool: + if self._last_hit is None: + return False + return self._clock() - self._last_hit < window_seconds + + db_lookup_gate: Final = LoopBoundSemaphore(PROXY_DB_LOOKUP_MAX_CONCURRENCY) +db_lookup_stall_tracker: Final = DBLookupStallTracker() + + +def _consume_abandoned_lookup(task: asyncio.Future[LookupT]) -> None: + if not task.cancelled(): + task.exception() + + +async def bounded_db_lookup( + lookup: Awaitable[LookupT], + *, + name: str, + deadline_seconds: float | None = None, + tracker: DBLookupStallTracker = db_lookup_stall_tracker, +) -> LookupT: + timeout: Final = PROXY_DB_LOOKUP_DEADLINE_SECONDS if deadline_seconds is None else deadline_seconds + task: Final = asyncio.ensure_future(lookup) + try: + done, _ = await asyncio.wait({task}, timeout=timeout) + except asyncio.CancelledError: + task.cancel() + raise + if task not in done: + task.cancel() + task.add_done_callback(_consume_abandoned_lookup) + tracker.record_hit() + raise DBLookupDeadlineExceeded(name, timeout) + try: + return task.result() + except DBLookupDeadlineExceeded: + raise + except asyncio.TimeoutError as e: + tracker.record_hit() + raise DBLookupDeadlineExceeded(name, timeout) from e diff --git a/litellm/proxy/db/exception_handler.py b/litellm/proxy/db/exception_handler.py index 38d6fb9b99d..022ca2efc9e 100644 --- a/litellm/proxy/db/exception_handler.py +++ b/litellm/proxy/db/exception_handler.py @@ -11,6 +11,7 @@ from litellm.proxy._types import ( ProxyErrorTypes, ProxyException, ) +from litellm.proxy.db.db_lookup_gate import DBLookupDeadlineExceeded from litellm.secret_managers.main import str_to_bool # Bounds the __cause__/__context__ walk in find_database_service_unavailable_error_in_chain. @@ -104,7 +105,7 @@ class PrismaDBExceptionHandler: """ import prisma.engine.errors - if isinstance(e, DB_CONNECTION_ERROR_TYPES): + if isinstance(e, (*DB_CONNECTION_ERROR_TYPES, DBLookupDeadlineExceeded)): return True if isinstance(e, _exception_types(prisma.engine.errors.EngineConnectionError)): return True diff --git a/litellm/proxy/db/spend_counter_reseed.py b/litellm/proxy/db/spend_counter_reseed.py index 2dd028454d6..f8e102d2682 100644 --- a/litellm/proxy/db/spend_counter_reseed.py +++ b/litellm/proxy/db/spend_counter_reseed.py @@ -23,7 +23,7 @@ from litellm._logging import verbose_proxy_logger from litellm.constants import SPEND_COUNTER_RESEED_LOCKS_MAX_SIZE from litellm.litellm_core_utils.duration_parser import duration_in_seconds from litellm.proxy._types import Litellm_EntityType -from litellm.proxy.db.db_lookup_gate import db_lookup_gate +from litellm.proxy.db.db_lookup_gate import bounded_db_lookup, db_lookup_gate from litellm.proxy.spend_tracking.spend_counter_batch import read_batched_spend_counter, record_spend_counter_value from litellm.repositories.organization_repository import OrganizationRepository from litellm.repositories.project_repository import ProjectRepository @@ -134,36 +134,9 @@ class SpendCounterReseed: if SpendCounterReseed._is_key_or_team_window_counter(counter_key): return None try: - async with db_lookup_gate.current(): - if counter_key.startswith("spend:key:"): - token: Final = counter_key[len("spend:key:") :] - row = await VerificationTokenRepository(prisma_client).table.find_unique(where={"token": token}) - elif counter_key.startswith("spend:team_member:"): - suffix: Final = counter_key[len("spend:team_member:") :] - if ":" not in suffix: - return None - user_id, team_id = suffix.rsplit(":", 1) - row = await TeamMembershipRepository(prisma_client).table.find_unique( - where={"user_id_team_id": {"user_id": user_id, "team_id": team_id}} - ) - elif counter_key.startswith("spend:team:"): - team_id = counter_key[len("spend:team:") :] - row = await TeamRepository(prisma_client).table.find_unique(where={"team_id": team_id}) - elif counter_key.startswith("spend:user:"): - user_id = counter_key[len("spend:user:") :] - row = await UserRepository(prisma_client).table.find_unique(where={"user_id": user_id}) - elif counter_key.startswith(END_USER_COUNTER_PREFIX) or counter_key.startswith("spend:tag:"): - return None - elif counter_key.startswith("spend:org:"): - org_id: Final = counter_key[len("spend:org:") :] - row = await OrganizationRepository(prisma_client).table.find_unique( - where={"organization_id": org_id} - ) - elif counter_key.startswith("spend:project:"): - project_id: Final = counter_key[len("spend:project:") :] - row = await ProjectRepository(prisma_client).table.find_unique(where={"project_id": project_id}) - else: - return None + row: Final = await bounded_db_lookup( + SpendCounterReseed._counter_row(prisma_client, counter_key), name="spend_counter" + ) except Exception: verbose_proxy_logger.exception("SpendCounterReseed.from_db: failed for %s", counter_key) return None @@ -171,13 +144,47 @@ class SpendCounterReseed: return None return float(getattr(row, "spend", 0.0) or 0.0) + @staticmethod + async def _counter_row(prisma_client: "PrismaClient", counter_key: str) -> object | None: + async with db_lookup_gate.current(): + if counter_key.startswith("spend:key:"): + token: Final = counter_key[len("spend:key:") :] + return await VerificationTokenRepository(prisma_client).table.find_unique(where={"token": token}) + if counter_key.startswith("spend:team_member:"): + suffix: Final = counter_key[len("spend:team_member:") :] + if ":" not in suffix: + return None + user_id, team_id = suffix.rsplit(":", 1) + return await TeamMembershipRepository(prisma_client).table.find_unique( + where={"user_id_team_id": {"user_id": user_id, "team_id": team_id}} + ) + if counter_key.startswith("spend:team:"): + return await TeamRepository(prisma_client).table.find_unique( + where={"team_id": counter_key[len("spend:team:") :]} + ) + if counter_key.startswith("spend:user:"): + return await UserRepository(prisma_client).table.find_unique( + where={"user_id": counter_key[len("spend:user:") :]} + ) + if counter_key.startswith("spend:org:"): + return await OrganizationRepository(prisma_client).table.find_unique( + where={"organization_id": counter_key[len("spend:org:") :]} + ) + if counter_key.startswith("spend:project:"): + return await ProjectRepository(prisma_client).table.find_unique( + where={"project_id": counter_key[len("spend:project:") :]} + ) + return None + @staticmethod async def end_user_from_db(prisma_client: Optional["PrismaClient"], counter_key: str) -> float | None: if prisma_client is None or not counter_key.startswith(END_USER_COUNTER_PREFIX): return None where: Final[LiteLLM_EndUserTableWhereUniqueInput] = {"user_id": counter_key[len(END_USER_COUNTER_PREFIX) :]} try: - row: Final = await EndUserRepository(prisma_client).table.find_unique(where=where) + row: Final = await bounded_db_lookup( + EndUserRepository(prisma_client).table.find_unique(where=where), name="end_user_spend" + ) except Exception: # noqa: BLE001 # a failed floor read falls back to the cached spend, like from_db verbose_proxy_logger.exception("SpendCounterReseed.end_user_from_db: failed for %s", counter_key) return None diff --git a/litellm/proxy/health_endpoints/_health_endpoints.py b/litellm/proxy/health_endpoints/_health_endpoints.py index 64fd59bbe44..f8801e65c82 100644 --- a/litellm/proxy/health_endpoints/_health_endpoints.py +++ b/litellm/proxy/health_endpoints/_health_endpoints.py @@ -16,7 +16,7 @@ from typing_extensions import ReadOnly import litellm from litellm._logging import verbose_logger, verbose_proxy_logger -from litellm.constants import HEALTH_CHECK_TIMEOUT_SECONDS +from litellm.constants import HEALTH_CHECK_TIMEOUT_SECONDS, PROXY_DB_LOOKUP_STALL_WINDOW_SECONDS from litellm.integrations.SlackAlerting.ms_teams import ( MS_TEAMS_ALERT_HEADERS, build_ms_teams_payload, @@ -44,6 +44,7 @@ from litellm.proxy.auth.auth_utils import ( ) from litellm.proxy.auth.model_checks import get_key_models from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.db.db_lookup_gate import db_lookup_stall_tracker from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler from litellm.proxy.db.health_check_latest import ( LatestHealthCheckRow, @@ -1723,7 +1724,7 @@ async def _get_health_readiness_details( # check DB if prisma_client is not None: # if db passed in, check if it's connected - db_health_status: Final = await _db_health_readiness_check() + db_status: Final = _readiness_db_status(await _db_health_readiness_check()) # A configured DB that is not reachable means the worker cannot # serve requests that depend on persisted state (keys, budgets, # spend logs). Return 503 so orchestrators take this pod out of @@ -1733,13 +1734,13 @@ async def _get_health_readiness_details( # report the DB state through the body instead. if ( response is not None - and db_health_status["status"] != "connected" + and db_status != "connected" and not PrismaDBExceptionHandler.should_allow_request_on_db_unavailable() ): response.status_code = status.HTTP_503_SERVICE_UNAVAILABLE return { "status": "healthy", - "db": db_health_status["status"], + "db": db_status, "cache": cache_type, "litellm_version": version, "success_callbacks": success_callback_names, @@ -1816,24 +1817,32 @@ def _authorize_drain_request(request: Request) -> None: ) +def _readiness_db_status(db_health_status: DBHealthCache) -> str: + """A pod whose pre-request lookups hit their deadline inside the stall window + reports "stalled" even though the ping succeeds: the ping is a fresh + connection, the stalled lookups are the ones requests actually wait on.""" + if db_health_status["status"] != "connected": + return db_health_status["status"] + if db_lookup_stall_tracker.stalled_within(PROXY_DB_LOOKUP_STALL_WINDOW_SECONDS): + return "stalled" + return "connected" + + async def _resolve_public_readiness_db(response: Response) -> str: """ Return the db status string for the public probe and flip the response to - 503 when a configured DB is unreachable. Mirrors the legacy values: - "Not connected" (no DB configured), "connected", "disconnected". + 503 when a configured DB is unreachable or stalled. Mirrors the legacy values: + "Not connected" (no DB configured), "connected", "disconnected", plus "stalled". """ from litellm.proxy.proxy_server import prisma_client if prisma_client is None: return "Not connected" - db_health_status: Final = await _db_health_readiness_check() - if ( - db_health_status["status"] != "connected" - and not PrismaDBExceptionHandler.should_allow_request_on_db_unavailable() - ): + db_status: Final = _readiness_db_status(await _db_health_readiness_check()) + if db_status != "connected" and not PrismaDBExceptionHandler.should_allow_request_on_db_unavailable(): response.status_code = status.HTTP_503_SERVICE_UNAVAILABLE - return db_health_status["status"] + return db_status @router.get( diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index 5dae9e8bb10..4dfea5cd472 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -23,6 +23,7 @@ from litellm.proxy.auth.auth_checks import ( log_db_metrics, ) from litellm.proxy.auth.route_checks import RouteChecks +from litellm.proxy.db.db_lookup_gate import DBLookupDeadlineExceeded from litellm.proxy.db.db_spend_update_writer import ( DBSpendUpdateWriter, debitable_model_access_groups, @@ -186,8 +187,8 @@ class _ProxyDBLogger(CustomLogger): ) _metadata["error_information"] = _error_information - _metadata = await _ProxyDBLogger._enrich_failure_metadata_with_key_info( - metadata=_metadata, + _metadata = await _ProxyDBLogger._enrich_failure_metadata_unless_db_stalled( + metadata=_metadata, original_exception=original_exception ) existing_metadata: Final[dict] = request_data.get("metadata", None) or {} @@ -472,6 +473,12 @@ class _ProxyDBLogger(CustomLogger): spend_log_error("Error in tracking cost callback - %s", str(e), exc=e) + @staticmethod + async def _enrich_failure_metadata_unless_db_stalled(metadata: dict, original_exception: Exception) -> dict: + if isinstance(original_exception, DBLookupDeadlineExceeded): + return metadata + return await _ProxyDBLogger._enrich_failure_metadata_with_key_info(metadata=metadata) + @staticmethod async def _enrich_failure_metadata_with_key_info(metadata: dict, resolve_missing_key_identity: bool = True) -> dict: """ diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index aa4bb2e49d3..30f5abdbb98 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -1,7 +1,7 @@ import asyncio import json import time -from collections.abc import Mapping +from collections.abc import Iterator, Mapping from types import SimpleNamespace from typing import TYPE_CHECKING, Final, Literal, Optional from unittest.mock import AsyncMock, MagicMock, patch @@ -646,6 +646,114 @@ async def test_fetch_key_object_from_db_bounds_in_flight_prisma_requests(): assert prisma.max_in_flight == PROXY_DB_LOOKUP_MAX_CONCURRENCY +@pytest.fixture +def _clear_db_lookup_stall() -> Iterator[None]: + from litellm.proxy.db.db_lookup_gate import db_lookup_stall_tracker + + db_lookup_stall_tracker.clear() + yield + db_lookup_stall_tracker.clear() + + +class _StalledPrisma: + def __init__(self) -> None: + self.attempt_db_reconnect = AsyncMock(return_value=True) + self.db = MagicMock() + self.db.litellm_teamtable.find_unique = AsyncMock(side_effect=_stall_forever) + self.db.litellm_teamtable.update = AsyncMock(side_effect=_answer_slowly) + + async def get_data(self, token: str, table_name: str, parent_otel_span: None, proxy_logging_obj: None) -> None: + await _stall_forever() + + +async def _stall_forever(**kwargs: object) -> None: + await asyncio.Event().wait() + + +async def _answer_slowly(**kwargs: object) -> Mapping[str, object]: + await asyncio.sleep(0.15) + return {"team_id": "slow-write"} + + +@pytest.mark.asyncio +async def test_fetch_key_object_from_db_fails_a_stalled_burst_within_the_deadline_without_reconnecting( + _clear_db_lookup_stall, +): + """The incident: a stalled database parked every request in the pod with liveness + and readiness green until it OOMed. Every lookup in a burst larger than the gate, + the ones queued behind it included, must fail within one deadline, must not try to + reconnect (the transport is fine, the query is slow), and must leave every gate slot + free for the next burst.""" + from litellm.proxy.db.db_lookup_gate import DBLookupDeadlineExceeded + + prisma: Final = _StalledPrisma() + burst: Final = PROXY_DB_LOOKUP_MAX_CONCURRENCY * 3 + started: Final = time.monotonic() + + results: Final = await asyncio.gather( + *( + _fetch_key_object_from_db_with_reconnect( + hashed_token=f"hashed-token-{i}", + prisma_client=prisma, # pyright: ignore[reportArgumentType] # fake stands in for PrismaClient + parent_otel_span=None, + proxy_logging_obj=None, + deadline_seconds=0.2, + ) + for i in range(burst) + ), + return_exceptions=True, + ) + elapsed: Final = time.monotonic() - started + + assert len(results) == burst + assert all(isinstance(result, DBLookupDeadlineExceeded) for result in results) + assert all(PrismaDBExceptionHandler.is_database_service_unavailable_error(result) for result in results) + assert elapsed < 3 + prisma.attempt_db_reconnect.assert_not_awaited() + + recovered: Final = _InFlightCountingPrisma() + after: Final = await asyncio.wait_for( + asyncio.gather( + *( + _fetch_key_object_from_db_with_reconnect( + hashed_token=f"after-{i}", + prisma_client=recovered, # pyright: ignore[reportArgumentType] # fake stands in for PrismaClient + parent_otel_span=None, + proxy_logging_obj=None, + ) + for i in range(PROXY_DB_LOOKUP_MAX_CONCURRENCY) + ) + ), + timeout=5, + ) + assert {r.token for r in after if r is not None} == {f"after-{i}" for i in range(PROXY_DB_LOOKUP_MAX_CONCURRENCY)} + + +@pytest.mark.asyncio +async def test_team_lookup_fails_at_the_db_lookup_deadline_while_writes_stay_unbounded(_clear_db_lookup_stall): + """Team, user, budget, and membership reads share the key lookup's deadline through + the typed table wrappers; writes do not, since a slow write must land rather than + fail the request that already passed auth.""" + from litellm.proxy.auth.auth_checks import _team_table + from litellm.proxy.db.db_lookup_gate import DBLookupDeadlineExceeded + from litellm.repositories.table_repositories import TeamRepository + + prisma: Final = _StalledPrisma() + with patch( # test-quality-ok: lowers the module-level lookup deadline so the stalled-read test finishes fast + "litellm.proxy.db.db_lookup_gate.PROXY_DB_LOOKUP_DEADLINE_SECONDS", 0.05 + ): + started: Final = time.monotonic() + with pytest.raises(DBLookupDeadlineExceeded, match=r"team lookup did not answer within 0\.05s"): + await _get_team_db_check(team_id="stalled-team", prisma_client=prisma) # pyright: ignore[reportArgumentType] # fake stands in for PrismaClient + assert time.monotonic() - started < 2 + + written: Final = await _team_table(TeamRepository(prisma)).update( + where={"team_id": "slow-write"}, data={"spend": 1.0} + ) + + assert written == {"team_id": "slow-write"} + + def _fake_redis_cache(): fake_redis = MagicMock() fake_redis.async_get_cache = AsyncMock(return_value=None) @@ -6184,7 +6292,9 @@ async def test_get_org_object_for_request_serves_last_known_org_through_db_outag proxy_logging_obj=None, ) - with patch("litellm.proxy.proxy_server.general_settings", {}): # test-quality-ok: the outage fallback reads this module global; no dependency injection seam exists + with patch( + "litellm.proxy.proxy_server.general_settings", {} + ): # test-quality-ok: the outage fallback reads this module global; no dependency injection seam exists warm = await _lookup() assert warm is not None and warm.organization_alias == "platform-org" await user_api_key_cache.async_delete_cache("org_id:org-1:with_budget") @@ -8933,20 +9043,34 @@ async def test_access_group_model_fallback_uses_the_injected_database(channel: s reader: Final = AsyncMock(return_value=group) client: Final = MagicMock(db=MagicMock(litellm_accessgrouptable=MagicMock(find_unique=reader))) with ( - patch("litellm.proxy.proxy_server.prisma_client", None), # test-quality-ok: [TQ008] prove reads stay on the injected connection - patch("litellm.proxy.proxy_server.user_api_key_cache", UserApiKeyCache()), # test-quality-ok: [TQ008] isolate the process cache + patch( + "litellm.proxy.proxy_server.prisma_client", None + ), # test-quality-ok: [TQ008] prove reads stay on the injected connection + patch( + "litellm.proxy.proxy_server.user_api_key_cache", UserApiKeyCache() + ), # test-quality-ok: [TQ008] isolate the process cache ): if channel == "team": - assert await can_team_access_model( - model="allowed", team_object=LiteLLM_TeamTable(team_id="team-a", models=["other"], access_group_ids=["group-a"]), - llm_router=None, prisma_client=client, - ) is True + assert ( + await can_team_access_model( + model="allowed", + team_object=LiteLLM_TeamTable(team_id="team-a", models=["other"], access_group_ids=["group-a"]), + llm_router=None, + prisma_client=client, + ) + is True + ) else: - assert await can_key_call_model( - model="allowed", llm_model_list=None, - valid_token=UserAPIKeyAuth(models=["other"], access_group_ids=["group-a"]), - llm_router=None, prisma_client=client, - ) is True + assert ( + await can_key_call_model( + model="allowed", + llm_model_list=None, + valid_token=UserAPIKeyAuth(models=["other"], access_group_ids=["group-a"]), + llm_router=None, + prisma_client=client, + ) + is True + ) reader.assert_awaited_once_with(where={"access_group_id": "group-a"}) @@ -8967,6 +9091,7 @@ def test_jwt_team_role_reaches_the_gateway_token_endpoint_by_default(): litellm_proxy_roles=LiteLLM_JWTAuth(team_allowed_routes=[]), ) + def test_route_skips_budget_checks_marks_only_spend_free_routes() -> None: assert route_skips_budget_checks(route="/v1/models") is True assert route_skips_budget_checks(route="/spend/logs") is True @@ -9083,7 +9208,9 @@ async def test_team_member_budget_check_temp_budget_increase_extends_cap(): return fallback_spend with ( - patch("litellm.proxy.proxy_server.get_current_spend", mock_get_current_spend), # test-quality-ok: [TQ008] no seam on the cross-pod spend counter + patch( + "litellm.proxy.proxy_server.get_current_spend", mock_get_current_spend + ), # test-quality-ok: [TQ008] no seam on the cross-pod spend counter patch( # test-quality-ok: [TQ008] isolates the check from the DB fetch "litellm.proxy.auth.auth_checks.get_team_membership", new_callable=AsyncMock, @@ -9111,7 +9238,9 @@ async def test_team_member_budget_check_temp_budget_increase_extends_cap(): ), ) with ( - patch("litellm.proxy.proxy_server.get_current_spend", mock_get_current_spend), # test-quality-ok: [TQ008] no seam on the cross-pod spend counter + patch( + "litellm.proxy.proxy_server.get_current_spend", mock_get_current_spend + ), # test-quality-ok: [TQ008] no seam on the cross-pod spend counter patch( # test-quality-ok: [TQ008] isolates the check from the DB fetch "litellm.proxy.auth.auth_checks.get_team_membership", new_callable=AsyncMock, @@ -9173,7 +9302,9 @@ async def test_team_member_budget_check_adds_temp_increase_to_live_team_default( return fallback_spend with ( - patch("litellm.proxy.proxy_server.get_current_spend", mock_get_current_spend), # test-quality-ok: [TQ008] no seam on the cross-pod spend counter + patch( + "litellm.proxy.proxy_server.get_current_spend", mock_get_current_spend + ), # test-quality-ok: [TQ008] no seam on the cross-pod spend counter patch( # test-quality-ok: [TQ008] isolates the check from the DB fetch "litellm.proxy.auth.auth_checks.get_team_membership", new_callable=AsyncMock, diff --git a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py index f03abe8f124..de669449f85 100644 --- a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py +++ b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py @@ -4,6 +4,7 @@ import logging import os import subprocess import sys +import time from collections.abc import Mapping from contextlib import contextmanager from datetime import datetime, timedelta, timezone @@ -186,11 +187,7 @@ async def test_disable_budget_reservation_does_not_log_per_request(caplog): general_settings={"disable_budget_reservation": True}, ) - records = [ - record - for record in caplog.records - if "disable_budget_reservation is enabled" in record.message - ] + records = [record for record in caplog.records if "disable_budget_reservation is enabled" in record.message] assert records == [] assert user_api_key_auth_obj.budget_reservation is None @@ -234,9 +231,7 @@ async def test_budget_reservation_runs_when_not_disabled(): ({}, False), ], ) -async def test_fail_closed_budget_enforcement_reaches_reservation( - general_settings, expected_flag -): +async def test_fail_closed_budget_enforcement_reaches_reservation(general_settings, expected_flag): """#33923: the strict flag must be threaded into reserve_budget_for_request so a failed reservation write can reject instead of failing open.""" user_api_key_auth_obj = UserAPIKeyAuth(token="test_token") @@ -259,10 +254,7 @@ async def test_fail_closed_budget_enforcement_reaches_reservation( general_settings=general_settings, ) - assert ( - mock_reserve.await_args.kwargs["fail_closed_budget_enforcement"] - is expected_flag - ) + assert mock_reserve.await_args.kwargs["fail_closed_budget_enforcement"] is expected_flag @pytest.mark.asyncio @@ -274,9 +266,7 @@ async def test_fail_closed_budget_enforcement_reaches_reservation( ({}, False), ], ) -async def test_apply_user_budget_to_team_keys_reaches_reservation( - general_settings, expected_flag -): +async def test_apply_user_budget_to_team_keys_reaches_reservation(general_settings, expected_flag): """The opt-in lives in general_settings but is consumed inside _get_budget_counters, so it has to be threaded through reserve_budget_for_request or the reservation path keeps exempting team keys while the read path enforces.""" @@ -300,9 +290,7 @@ async def test_apply_user_budget_to_team_keys_reaches_reservation( general_settings=general_settings, ) - assert ( - mock_reserve.await_args.kwargs["apply_user_budget_to_team_keys"] is expected_flag - ) + assert mock_reserve.await_args.kwargs["apply_user_budget_to_team_keys"] is expected_flag @pytest.mark.asyncio @@ -402,9 +390,7 @@ async def test_custom_auth_honors_key_level_model_access_restriction_allowed_wit "litellm.proxy.auth.user_api_key_auth.can_key_call_model", new_callable=AsyncMock, ) as mock_can_key, - patch( - "litellm.proxy.auth.user_api_key_auth.common_checks", new_callable=AsyncMock - ), + patch("litellm.proxy.auth.user_api_key_auth.common_checks", new_callable=AsyncMock), patch( "litellm.proxy.proxy_server.general_settings", {"custom_auth_run_common_checks": True}, @@ -435,9 +421,7 @@ async def test_custom_auth_enforces_key_model_access_from_file_route_header_with "litellm.proxy.auth.user_api_key_auth.can_key_call_model", new_callable=AsyncMock, ) as mock_can_key, - patch( - "litellm.proxy.auth.user_api_key_auth.common_checks", new_callable=AsyncMock - ), + patch("litellm.proxy.auth.user_api_key_auth.common_checks", new_callable=AsyncMock), patch( "litellm.proxy.proxy_server.general_settings", {"custom_auth_run_common_checks": True}, @@ -468,9 +452,7 @@ async def test_custom_auth_honors_key_level_model_access_restriction_denied_with "litellm.proxy.auth.user_api_key_auth.can_key_call_model", new_callable=AsyncMock, ) as mock_can_key, - patch( - "litellm.proxy.auth.user_api_key_auth.common_checks", new_callable=AsyncMock - ), + patch("litellm.proxy.auth.user_api_key_auth.common_checks", new_callable=AsyncMock), patch( "litellm.proxy.proxy_server.general_settings", {"custom_auth_run_common_checks": True}, @@ -506,9 +488,7 @@ def _proxy_server_attrs_for_custom_auth(*, user_custom_auth): mock_proxy_logging_obj = MagicMock() mock_proxy_logging_obj.internal_usage_cache = MagicMock() mock_proxy_logging_obj.internal_usage_cache.dual_cache = AsyncMock() - mock_proxy_logging_obj.internal_usage_cache.dual_cache.async_delete_cache = ( - AsyncMock() - ) + mock_proxy_logging_obj.internal_usage_cache.dual_cache.async_delete_cache = AsyncMock() mock_proxy_logging_obj.post_call_failure_hook = AsyncMock(return_value=None) return { @@ -770,9 +750,7 @@ async def test_enterprise_custom_auth_runs_post_custom_auth_checks_when_opt_in() litellm.enable_post_custom_auth_checks = original_flag -def _assert_get_api_key_with_custom_litellm_key_header( - custom_litellm_key_header, api_key, passed_in_key -): +def _assert_get_api_key_with_custom_litellm_key_header(custom_litellm_key_header, api_key, passed_in_key): assert get_api_key( custom_litellm_key_header=custom_litellm_key_header, api_key=None, @@ -829,9 +807,7 @@ def _assert_get_api_key_with_custom_litellm_key_header( ("App:LiteLLM", None, False, False), ], ) -def test_routing_selector_matches_claim_parametrized( - selector_value, claim_value, expected, split_space_delimited -): +def test_routing_selector_matches_claim_parametrized(selector_value, claim_value, expected, split_space_delimited): assert ( _routing_selector_matches_claim( selector_value=selector_value, @@ -925,10 +901,7 @@ def test_routing_selector_matches_claim_parametrized( ], ) def test_matches_routing_override_parametrized(override, token_claims, expected): - assert ( - _matches_routing_override(token_claims=token_claims, override=override) - is expected - ) + assert _matches_routing_override(token_claims=token_claims, override=override) is expected def test_get_api_key_with_custom_litellm_key_header_bearer_prefix(): @@ -1007,12 +980,9 @@ def test_team_metadata_with_tags_flows_through_jwt_auth(): ) # Verify team_metadata is set - assert ( - user_api_key_auth.team_metadata is not None - ), "team_metadata should be populated" + assert user_api_key_auth.team_metadata is not None, "team_metadata should be populated" assert user_api_key_auth.team_metadata == team_object.metadata, ( - f"team_metadata not correctly mapped. " - f"Expected: {team_object.metadata}, Got: {user_api_key_auth.team_metadata}" + f"team_metadata not correctly mapped. Expected: {team_object.metadata}, Got: {user_api_key_auth.team_metadata}" ) # Specifically verify tags are present @@ -1051,9 +1021,7 @@ def test_route_checks_is_llm_api_route(): ] for route in openai_routes: - assert RouteChecks.is_llm_api_route( - route=route - ), f"Route {route} should be identified as LLM API route" + assert RouteChecks.is_llm_api_route(route=route), f"Route {route} should be identified as LLM API route" # Test Anthropic routes anthropic_routes = [ @@ -1062,9 +1030,7 @@ def test_route_checks_is_llm_api_route(): ] for route in anthropic_routes: - assert RouteChecks.is_llm_api_route( - route=route - ), f"Route {route} should be identified as LLM API route" + assert RouteChecks.is_llm_api_route(route=route), f"Route {route} should be identified as LLM API route" # Test passthrough routes (this is the key improvement over the old route checking) passthrough_routes = [ @@ -1084,9 +1050,7 @@ def test_route_checks_is_llm_api_route(): ] for route in passthrough_routes: - assert RouteChecks.is_llm_api_route( - route=route - ), f"Route {route} should be identified as LLM API route" + assert RouteChecks.is_llm_api_route(route=route), f"Route {route} should be identified as LLM API route" # Test MCP routes mcp_routes = [ @@ -1096,9 +1060,7 @@ def test_route_checks_is_llm_api_route(): ] for route in mcp_routes: - assert RouteChecks.is_llm_api_route( - route=route - ), f"Route {route} should be identified as LLM API route" + assert RouteChecks.is_llm_api_route(route=route), f"Route {route} should be identified as LLM API route" # Test LiteLLM native RAG routes rag_routes = [ @@ -1108,9 +1070,7 @@ def test_route_checks_is_llm_api_route(): "/v1/rag/query", ] for route in rag_routes: - assert RouteChecks.is_llm_api_route( - route=route - ), f"Route {route} should be identified as LLM API route" + assert RouteChecks.is_llm_api_route(route=route), f"Route {route} should be identified as LLM API route" # Test routes with placeholders placeholder_routes = [ @@ -1125,9 +1085,7 @@ def test_route_checks_is_llm_api_route(): ] for route in placeholder_routes: - assert RouteChecks.is_llm_api_route( - route=route - ), f"Route {route} should be identified as LLM API route" + assert RouteChecks.is_llm_api_route(route=route), f"Route {route} should be identified as LLM API route" # Test Azure OpenAI routes azure_routes = [ @@ -1138,9 +1096,7 @@ def test_route_checks_is_llm_api_route(): ] for route in azure_routes: - assert RouteChecks.is_llm_api_route( - route=route - ), f"Route {route} should be identified as LLM API route" + assert RouteChecks.is_llm_api_route(route=route), f"Route {route} should be identified as LLM API route" # Test non-LLM routes (should return False) non_llm_routes = [ @@ -1159,9 +1115,7 @@ def test_route_checks_is_llm_api_route(): ] for route in non_llm_routes: - assert not RouteChecks.is_llm_api_route( - route=route - ), f"Route {route} should NOT be identified as LLM API route" + assert not RouteChecks.is_llm_api_route(route=route), f"Route {route} should NOT be identified as LLM API route" # Test invalid inputs invalid_inputs = [ @@ -1173,9 +1127,9 @@ def test_route_checks_is_llm_api_route(): ] for invalid_input in invalid_inputs: - assert not RouteChecks.is_llm_api_route( - route=invalid_input - ), f"Invalid input {invalid_input} should return False" + assert not RouteChecks.is_llm_api_route(route=invalid_input), ( + f"Invalid input {invalid_input} should return False" + ) @pytest.mark.asyncio @@ -1222,9 +1176,7 @@ async def test_proxy_admin_expired_key_from_cache(): mock_proxy_logging_obj = MagicMock() mock_proxy_logging_obj.internal_usage_cache = MagicMock() mock_proxy_logging_obj.internal_usage_cache.dual_cache = AsyncMock() - mock_proxy_logging_obj.internal_usage_cache.dual_cache.async_delete_cache = ( - AsyncMock() - ) + mock_proxy_logging_obj.internal_usage_cache.dual_cache.async_delete_cache = AsyncMock() # Mock post_call_failure_hook as async function returning None (no transformation) mock_proxy_logging_obj.post_call_failure_hook = AsyncMock(return_value=None) @@ -1261,9 +1213,7 @@ async def test_proxy_admin_expired_key_from_cache(): "jwt_handler": None, "litellm_proxy_admin_name": "admin", } - _original_values = { - attr: getattr(_proxy_server_mod, attr, None) for attr in _attrs_to_set - } + _original_values = {attr: getattr(_proxy_server_mod, attr, None) for attr in _attrs_to_set} try: for attr, val in _attrs_to_set.items(): setattr(_proxy_server_mod, attr, val) @@ -1287,36 +1237,30 @@ async def test_proxy_admin_expired_key_from_cache(): ) # Verify that ProxyException was raised with expired_key type - assert hasattr( - exc_info.value, "type" - ), "Exception should have 'type' attribute" - assert ( - exc_info.value.type == ProxyErrorTypes.expired_key - ), f"Expected expired_key error type, got {exc_info.value.type}" + assert hasattr(exc_info.value, "type"), "Exception should have 'type' attribute" + assert exc_info.value.type == ProxyErrorTypes.expired_key, ( + f"Expected expired_key error type, got {exc_info.value.type}" + ) assert int(exc_info.value.code) == status.HTTP_401_UNAUTHORIZED - assert "Expired Key" in str( - exc_info.value.message - ), f"Exception message should mention 'Expired Key', got: {exc_info.value.message}" + assert "Expired Key" in str(exc_info.value.message), ( + f"Exception message should mention 'Expired Key', got: {exc_info.value.message}" + ) # Verify that the param field does NOT leak the full API key (Issue #18731) # The param should be abbreviated like "sk-...XXXX" not the full plaintext key - assert ( - exc_info.value.param is not None - ), "Exception should have 'param' attribute" + assert exc_info.value.param is not None, "Exception should have 'param' attribute" assert exc_info.value.param != api_key, ( f"SECURITY: Full API key should NOT be in param field! " f"Got: {exc_info.value.param}, Expected abbreviated format like 'sk-...XXXX'" ) - assert exc_info.value.param.startswith( - "sk-..." - ), f"Param should be abbreviated to 'sk-...XXXX' format. Got: {exc_info.value.param}" + assert exc_info.value.param.startswith("sk-..."), ( + f"Param should be abbreviated to 'sk-...XXXX' format. Got: {exc_info.value.param}" + ) # Verify that cache deletion was called mock_delete_cache.assert_called_once() call_args = mock_delete_cache.call_args - assert ( - call_args[1]["hashed_token"] == hashed_key - ), "Cache deletion should be called with the hashed key" + assert call_args[1]["hashed_token"] == hashed_key, "Cache deletion should be called with the hashed key" finally: # Restore all module-level attributes so subsequent tests are not affected for attr, val in _original_values.items(): @@ -1354,9 +1298,7 @@ async def test_scim_deactivated_user_key_is_rejected(): mock_proxy_logging_obj = MagicMock() mock_proxy_logging_obj.internal_usage_cache = MagicMock() mock_proxy_logging_obj.internal_usage_cache.dual_cache = AsyncMock() - mock_proxy_logging_obj.internal_usage_cache.dual_cache.async_delete_cache = ( - AsyncMock() - ) + mock_proxy_logging_obj.internal_usage_cache.dual_cache.async_delete_cache = AsyncMock() mock_proxy_logging_obj.post_call_failure_hook = AsyncMock(return_value=None) mock_prisma_client = MagicMock() @@ -1377,9 +1319,7 @@ async def test_scim_deactivated_user_key_is_rejected(): "jwt_handler": None, "litellm_proxy_admin_name": "admin", } - _original_values = { - attr: getattr(_proxy_server_mod, attr, None) for attr in _attrs_to_set - } + _original_values = {attr: getattr(_proxy_server_mod, attr, None) for attr in _attrs_to_set} try: for attr, val in _attrs_to_set.items(): setattr(_proxy_server_mod, attr, val) @@ -1446,9 +1386,7 @@ async def test_cached_proxy_admin_key_sets_via_virtual_key_marker(): mock_proxy_logging_obj = MagicMock() mock_proxy_logging_obj.internal_usage_cache = MagicMock() mock_proxy_logging_obj.internal_usage_cache.dual_cache = AsyncMock() - mock_proxy_logging_obj.internal_usage_cache.dual_cache.async_delete_cache = ( - AsyncMock() - ) + mock_proxy_logging_obj.internal_usage_cache.dual_cache.async_delete_cache = AsyncMock() mock_proxy_logging_obj.post_call_failure_hook = AsyncMock(return_value=None) import litellm.proxy.proxy_server as _proxy_server_mod @@ -1467,9 +1405,7 @@ async def test_cached_proxy_admin_key_sets_via_virtual_key_marker(): "jwt_handler": None, "litellm_proxy_admin_name": "admin", } - _original_values = { - attr: getattr(_proxy_server_mod, attr, None) for attr in _attrs_to_set - } + _original_values = {attr: getattr(_proxy_server_mod, attr, None) for attr in _attrs_to_set} try: for attr, val in _attrs_to_set.items(): setattr(_proxy_server_mod, attr, val) @@ -1521,9 +1457,7 @@ async def test_master_key_auth_sets_via_virtual_key_marker(): mock_proxy_logging_obj = MagicMock() mock_proxy_logging_obj.internal_usage_cache = MagicMock() mock_proxy_logging_obj.internal_usage_cache.dual_cache = AsyncMock() - mock_proxy_logging_obj.internal_usage_cache.dual_cache.async_delete_cache = ( - AsyncMock() - ) + mock_proxy_logging_obj.internal_usage_cache.dual_cache.async_delete_cache = AsyncMock() mock_proxy_logging_obj.post_call_failure_hook = AsyncMock(return_value=None) import litellm.proxy.proxy_server as _proxy_server_mod @@ -1542,9 +1476,7 @@ async def test_master_key_auth_sets_via_virtual_key_marker(): "jwt_handler": None, "litellm_proxy_admin_name": "admin", } - _original_values = { - attr: getattr(_proxy_server_mod, attr, None) for attr in _attrs_to_set - } + _original_values = {attr: getattr(_proxy_server_mod, attr, None) for attr in _attrs_to_set} try: for attr, val in _attrs_to_set.items(): setattr(_proxy_server_mod, attr, val) @@ -1597,9 +1529,7 @@ async def test_db_virtual_key_auth_sets_via_virtual_key_marker(): mock_proxy_logging_obj = MagicMock() mock_proxy_logging_obj.internal_usage_cache = MagicMock() mock_proxy_logging_obj.internal_usage_cache.dual_cache = AsyncMock() - mock_proxy_logging_obj.internal_usage_cache.dual_cache.async_delete_cache = ( - AsyncMock() - ) + mock_proxy_logging_obj.internal_usage_cache.dual_cache.async_delete_cache = AsyncMock() mock_proxy_logging_obj.post_call_failure_hook = AsyncMock(return_value=None) mock_prisma_client = MagicMock() @@ -1620,9 +1550,7 @@ async def test_db_virtual_key_auth_sets_via_virtual_key_marker(): "jwt_handler": None, "litellm_proxy_admin_name": "admin", } - _original_values = { - attr: getattr(_proxy_server_mod, attr, None) for attr in _attrs_to_set - } + _original_values = {attr: getattr(_proxy_server_mod, attr, None) for attr in _attrs_to_set} try: for attr, val in _attrs_to_set.items(): setattr(_proxy_server_mod, attr, val) @@ -2153,7 +2081,10 @@ async def test_auto_register_first_request_propagates_user_email(active: bool) - patch("litellm.proxy.proxy_server.master_key", "sk-master"), patch("litellm.proxy.proxy_server.prisma_client", prisma_client), patch("litellm.proxy.proxy_server.user_api_key_cache", user_api_key_cache), - patch("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock(post_call_failure_hook=AsyncMock(return_value=None))), + patch( + "litellm.proxy.proxy_server.proxy_logging_obj", + MagicMock(post_call_failure_hook=AsyncMock(return_value=None)), + ), patch("litellm.proxy.proxy_server.jwt_handler", jwt_handler), patch( "litellm.proxy.auth.user_api_key_auth._resolve_jwt_to_virtual_key", @@ -2218,7 +2149,9 @@ async def test_auto_register_stamps_new_key_with_jwt_agent_id(): plaintext = "sk-auto-registered-agent" token_hash = hash_token(plaintext) persisted_principal = IdentityStore._principal_from_key( - UserAPIKeyAuth(token=token_hash, user_id="validated-user", team_id="validated-team", agent_id="canonical-agent-id"), + UserAPIKeyAuth( + token=token_hash, user_id="validated-user", team_id="validated-team", agent_id="canonical-agent-id" + ), auth_method=AuthMethod.API_KEY, credential_ref=CredentialRef(token_id=token_hash), ) @@ -2433,10 +2366,7 @@ class TestJWTOAuth2Coexistence: def test_is_jwt_detects_jwt_tokens(self): """JWT tokens have 3 dot-separated parts.""" assert JWTHandler.is_jwt("header.payload.signature") is True - assert ( - JWTHandler.is_jwt("eyJhbGciOiJSUzI1NiJ9.eyJzdWIiOiJ1c2VyMSJ9.sig123") - is True - ) + assert JWTHandler.is_jwt("eyJhbGciOiJSUzI1NiJ9.eyJzdWIiOiJ1c2VyMSJ9.sig123") is True def test_is_jwt_rejects_opaque_tokens(self): """Opaque OAuth2 tokens do not have 3 dot-separated parts.""" @@ -2545,10 +2475,7 @@ class TestJWTOAuth2Coexistence: assert exc_info.value.type == ProxyErrorTypes.auth_error assert exc_info.value.code == "403" - assert ( - "Oauth2 token validation is only available for premium users" - in exc_info.value.message - ) + assert "Oauth2 token validation is only available for premium users" in exc_info.value.message mock_oauth2.assert_not_called() @pytest.mark.asyncio @@ -2740,9 +2667,7 @@ class TestJWTOAuth2Coexistence: assert mock_auto_register.call_args.kwargs["team_id"] == "validated-team" assert mock_auto_register.call_args.kwargs["user_id"] == "validated-user" assert mock_auto_register.call_args.kwargs["org_id"] == "validated-org" - assert ( - mock_auto_register.call_args.kwargs["end_user_id"] == "validated-end-user" - ) + assert mock_auto_register.call_args.kwargs["end_user_id"] == "validated-end-user" assert result.org_id == "validated-org" assert result.user_email == "validated@example.com" @@ -2820,10 +2745,7 @@ class TestJWTOAuth2Coexistence: assert result.user_id == "mapped-user" assert result.user_email == "mapped@example.com" - assert ( - mock_get_user_object.call_args_list[0].kwargs["user_email"] - == "mapped@example.com" - ) + assert mock_get_user_object.call_args_list[0].kwargs["user_email"] == "mapped@example.com" @pytest.mark.asyncio async def test_mapped_virtual_key_does_not_backfill_mismatched_owner(self): @@ -2899,8 +2821,7 @@ class TestJWTOAuth2Coexistence: assert result.user_id == "other-owner" assert result.user_email is None assert all( - call.kwargs.get("user_email") != "principal@example.com" - for call in mock_get_user_object.call_args_list + call.kwargs.get("user_email") != "principal@example.com" for call in mock_get_user_object.call_args_list ) @pytest.mark.asyncio @@ -3705,9 +3626,7 @@ async def test_user_api_key_auth_builder_no_blocking_calls(): mock_proxy_logging_obj = MagicMock() mock_proxy_logging_obj.internal_usage_cache = MagicMock() mock_proxy_logging_obj.internal_usage_cache.dual_cache = AsyncMock() - mock_proxy_logging_obj.internal_usage_cache.dual_cache.async_delete_cache = ( - AsyncMock() - ) + mock_proxy_logging_obj.internal_usage_cache.dual_cache.async_delete_cache = AsyncMock() mock_proxy_logging_obj.post_call_failure_hook = AsyncMock(return_value=None) import litellm.proxy.proxy_server as _proxy_server_mod @@ -3839,9 +3758,7 @@ async def test_team_metadata_refreshed_from_team_object_during_auth(): mock_proxy_logging_obj = MagicMock() mock_proxy_logging_obj.internal_usage_cache = MagicMock() mock_proxy_logging_obj.internal_usage_cache.dual_cache = AsyncMock() - mock_proxy_logging_obj.internal_usage_cache.dual_cache.async_delete_cache = ( - AsyncMock() - ) + mock_proxy_logging_obj.internal_usage_cache.dual_cache.async_delete_cache = AsyncMock() mock_proxy_logging_obj.post_call_failure_hook = AsyncMock(return_value=None) import litellm.proxy.proxy_server as _proxy_server_mod @@ -3891,9 +3808,9 @@ async def test_team_metadata_refreshed_from_team_object_during_auth(): request_data={}, ) - assert result.team_metadata == { - "guardrails": ["test-guardrail-333"] - }, f"team_metadata was not updated from fresh team object. Got: {result.team_metadata}" + assert result.team_metadata == {"guardrails": ["test-guardrail-333"]}, ( + f"team_metadata was not updated from fresh team object. Got: {result.team_metadata}" + ) finally: for k, v in _originals.items(): @@ -4218,9 +4135,7 @@ async def test_auth_flow_fallback_team_object_permission_none_when_unreadable(): # --------------------------------------------------------------------------- -def _proxy_attrs_for_centralized_checks( - user_custom_auth=None, flag=False, master_key="sk-test-master" -): +def _proxy_attrs_for_centralized_checks(user_custom_auth=None, flag=False, master_key="sk-test-master"): """Build the minimal proxy_server module attributes that _run_centralized_common_checks reads. @@ -4430,9 +4345,7 @@ async def _run_centralized_checks_with_key_end_user_budget( request = Request(scope={"type": "http"}) request._url = URL(url="/chat/completions") attrs = { - **_proxy_attrs_for_centralized_checks( - user_custom_auth=AsyncMock() if custom_auth else None, flag=custom_auth - ), + **_proxy_attrs_for_centralized_checks(user_custom_auth=AsyncMock() if custom_auth else None, flag=custom_auth), "prisma_client": prisma_client, "user_api_key_cache": user_api_key_cache if user_api_key_cache is not None else DualCache(), "proxy_logging_obj": proxy_logging_obj, @@ -4623,7 +4536,9 @@ async def test_centralized_common_checks_enforces_team_model_max_budget_from_the for k, v in attrs.items(): setattr(_proxy_server_mod, k, v) with ( - patch("litellm.proxy.auth.user_api_key_auth.common_checks", new_callable=AsyncMock), # test-quality-ok: stubs the sibling check so only the team model-budget gate is under test + patch( + "litellm.proxy.auth.user_api_key_auth.common_checks", new_callable=AsyncMock + ), # test-quality-ok: stubs the sibling check so only the team model-budget gate is under test patch( # test-quality-ok: stubs the budget reservation so only the team model-budget gate is under test "litellm.proxy.auth.user_api_key_auth._reserve_budget_after_common_checks", new_callable=AsyncMock, @@ -4656,9 +4571,7 @@ async def test_centralized_common_checks_skipped_for_custom_auth_without_flag(): request = Request(scope={"type": "http"}) request._url = URL(url="/chat/completions") - attrs = _proxy_attrs_for_centralized_checks( - user_custom_auth=AsyncMock(), flag=False - ) + attrs = _proxy_attrs_for_centralized_checks(user_custom_auth=AsyncMock(), flag=False) originals = {a: getattr(_proxy_server_mod, a, None) for a in attrs} try: for k, v in attrs.items(): @@ -5063,9 +4976,7 @@ async def test_centralized_common_checks_reserves_request_end_user_budget(): "applied_adjustment": 0.0, } ] - assert counter_cache.in_memory_cache.get_cache( - key="spend:end_user:alice" - ) == pytest.approx(0.6) + assert counter_cache.in_memory_cache.get_cache(key="spend:end_user:alice") == pytest.approx(0.6) @pytest.mark.asyncio @@ -5080,9 +4991,7 @@ async def test_centralized_common_checks_short_circuits_when_master_key_unset(): from litellm.proxy._types import LitellmUserRoles - token = UserAPIKeyAuth( - api_key="sk-test", user_id="u", user_role=LitellmUserRoles.INTERNAL_USER - ) + token = UserAPIKeyAuth(api_key="sk-test", user_id="u", user_role=LitellmUserRoles.INTERNAL_USER) request = Request(scope={"type": "http"}) request._url = URL(url="/get/config/callbacks") @@ -5883,9 +5792,7 @@ async def test_centralized_common_checks_user_http_exception_isolates_to_user_on request._url = URL(url="/chat/completions") request._body = json.dumps({"user": "alice", "model": "gpt-4o"}).encode() - fetched_team = LiteLLM_TeamTableCachedObj( - team_id="t1", max_budget=20.0, models=["gpt-4o"] - ) + fetched_team = LiteLLM_TeamTableCachedObj(team_id="t1", max_budget=20.0, models=["gpt-4o"]) fetched_end_user = LiteLLM_EndUserTable(user_id="alice", blocked=False, spend=1.0) fetched_project = LiteLLM_ProjectTableCachedObj( project_id="proj-1", @@ -6014,10 +5921,46 @@ async def test_centralized_common_checks_backfills_org_id_from_team(key_org_id, ("org-pinned", None, None, "preset", None, "success", False, False, "org-pinned", "preset", (None, None, None)), ("org-view", None, None, None, 3, "success", False, False, "org-view", None, (None, None, 3)), ("org-missing", None, None, None, None, "missing", False, False, "org-missing", None, (None, None, None)), - ("org-db-failure-allowed", None, None, None, None, "db_failure", True, False, "org-db-failure-allowed", None, (None, None, None)), - ("org-db-failure-denied", None, None, None, None, "db_failure", False, True, "org-db-failure-denied", None, (None, None, None)), + ( + "org-db-failure-allowed", + None, + None, + None, + None, + "db_failure", + True, + False, + "org-db-failure-allowed", + None, + (None, None, None), + ), + ( + "org-db-failure-denied", + None, + None, + None, + None, + "db_failure", + False, + True, + "org-db-failure-denied", + None, + (None, None, None), + ), ("org-bad-row", None, None, None, None, "bad_row", False, False, "org-bad-row", None, (None, None, None)), - ("org-nobudget", None, None, None, None, "no_budget", False, False, "org-nobudget", "acme-org", (None, None, None)), + ( + "org-nobudget", + None, + None, + None, + None, + "no_budget", + False, + False, + "org-nobudget", + "acme-org", + (None, None, None), + ), ], ) async def test_centralized_common_checks_inherits_org_identity( @@ -6326,9 +6269,7 @@ async def test_user_api_key_auth_sets_end_user_id_when_builder_skips_it(): } ) request._url = URL(url="/chat/completions") - request._body = json.dumps( - {"model": "gpt-4o", "user": "alice@example.com"} - ).encode() + request._body = json.dumps({"model": "gpt-4o", "user": "alice@example.com"}).encode() attrs = _proxy_attrs_for_centralized_checks(user_custom_auth=None) originals = {a: getattr(_proxy_server_mod, a, None) for a in attrs} @@ -6372,9 +6313,7 @@ async def test_user_api_key_auth_does_not_overwrite_end_user_id_set_by_builder() import litellm.proxy.proxy_server as _proxy_server_mod - builder_token = UserAPIKeyAuth( - api_key="sk-test", user_id="u1", end_user_id="builder-resolved-id" - ) + builder_token = UserAPIKeyAuth(api_key="sk-test", user_id="u1", end_user_id="builder-resolved-id") request = Request( scope={ @@ -6384,9 +6323,7 @@ async def test_user_api_key_auth_does_not_overwrite_end_user_id_set_by_builder() } ) request._url = URL(url="/chat/completions") - request._body = json.dumps( - {"model": "gpt-4o", "user": "different-id-from-body"} - ).encode() + request._body = json.dumps({"model": "gpt-4o", "user": "different-id-from-body"}).encode() attrs = _proxy_attrs_for_centralized_checks(user_custom_auth=None) originals = {a: getattr(_proxy_server_mod, a, None) for a in attrs} @@ -6717,6 +6654,83 @@ async def _run_builder_with_key_lookup(get_key_object_mock): setattr(_proxy_server_mod, k, v) +class _StalledKeyLookupPrisma: + """A database whose connection answers the readiness ping but whose key lookups + never return, which is what the incident's locked table looked like.""" + + def __init__(self) -> None: + self.health_check = AsyncMock(return_value=True) + self.attempt_db_reconnect = AsyncMock(return_value=True) + self.db = MagicMock() + + async def get_data(self, token: str, table_name: str, parent_otel_span: None, proxy_logging_obj: None) -> None: + await asyncio.Event().wait() + + +@pytest.mark.asyncio +async def test_burst_against_a_stalled_db_fails_fast_with_503_and_turns_readiness_red(): + """The incident, end to end: N requests into a proxy whose database stalls used to + park in the pod with readiness green until it OOMed. Now every one of them fails + within the lookup deadline as a 503, and the next readiness probe takes the pod out + of rotation.""" + import httpx + from fastapi import Depends, FastAPI + + import litellm.proxy.health_endpoints._health_endpoints as health_endpoints + import litellm.proxy.proxy_server as _proxy_server_mod + from litellm.proxy.db.db_lookup_gate import db_lookup_stall_tracker + + app = FastAPI() + + @app.post("/chat/completions", dependencies=[Depends(user_api_key_auth)]) + async def chat_completions() -> Mapping[str, bool]: + return {"served": True} + + app.include_router(health_endpoints.router) + app.add_exception_handler(ProxyException, _proxy_server_mod.openai_exception_handler) + + attrs = {**_proxy_attrs_for_db_lookup(), "prisma_client": _StalledKeyLookupPrisma()} + originals = {a: getattr(_proxy_server_mod, a, None) for a in attrs} + health_endpoints.db_health_cache = {"status": "unknown", "last_updated": datetime.now() - timedelta(seconds=60)} + db_lookup_stall_tracker.clear() + burst = 60 + try: + for k, v in attrs.items(): + setattr(_proxy_server_mod, k, v) + with ( + patch( # test-quality-ok: lowers the module-level lookup deadline so the stalled burst finishes fast + "litellm.proxy.db.db_lookup_gate.PROXY_DB_LOOKUP_DEADLINE_SECONDS", 0.2 + ), + patch("litellm.proxy.auth.auth_exception_handler.seed_request_identity"), + ): + async with httpx.AsyncClient(transport=httpx.ASGITransport(app=app), base_url="http://t") as client: + started = time.monotonic() + responses = await asyncio.gather( + *( + client.post( + "/chat/completions", + json={"model": "gpt-5.5", "messages": [{"role": "user", "content": "hi"}]}, + headers={"Authorization": f"Bearer sk-stalled-{i}"}, + ) + for i in range(burst) + ) + ) + elapsed = time.monotonic() - started + readiness = await client.get("/health/readiness") + finally: + for k, v in originals.items(): + setattr(_proxy_server_mod, k, v) + db_lookup_stall_tracker.clear() + + assert len(responses) == burst + assert {r.status_code for r in responses} == {status.HTTP_503_SERVICE_UNAVAILABLE} + assert {r.json()["error"]["type"] for r in responses} == {ProxyErrorTypes.no_db_connection.value} + assert all("temporarily unreachable" in r.json()["error"]["message"] for r in responses) + assert elapsed < 5 + assert readiness.status_code == status.HTTP_503_SERVICE_UNAVAILABLE + assert readiness.json()["db"] == "stalled" + + @pytest.mark.asyncio async def test_builder_returns_503_when_db_lookup_raises_infra_error(): """End-to-end: a DB infrastructure failure during the key lookup must @@ -6784,9 +6798,7 @@ def _mint_cli_session_token(monkeypatch, *, user_id="cli-admin"): models=["gpt-3.5-turbo"], max_budget=100.0, ) - return ExperimentalUIJWTToken.get_cli_jwt_auth_token( - user_info, team_id="cli-team", team_alias="cli-team-alias" - ) + return ExperimentalUIJWTToken.get_cli_jwt_auth_token(user_info, team_id="cli-team", team_alias="cli-team-alias") @pytest.mark.asyncio @@ -6836,7 +6848,7 @@ async def test_random_non_sk_token_is_rejected(monkeypatch): patch("litellm.proxy.proxy_server.master_key", "sk-master"), patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), ): - with pytest.raises(Exception, match='LiteLLM Virtual Key expected\\.') as exc_info: + with pytest.raises(Exception, match="LiteLLM Virtual Key expected\\.") as exc_info: await user_api_key_auth( request=mock_request, api_key="Bearer not-a-real-token", @@ -6915,9 +6927,7 @@ async def test_non_admin_cli_session_token_reaches_production_auth_path(monkeypa user_role=LitellmUserRoles.INTERNAL_USER.value, models=[], ) - cli_token = ExperimentalUIJWTToken.get_cli_jwt_auth_token( - user_info, team_id="team-abc", team_alias="my-team" - ) + cli_token = ExperimentalUIJWTToken.get_cli_jwt_auth_token(user_info, team_id="team-abc", team_alias="my-team") import litellm.proxy.proxy_server as _proxy_server_mod from fastapi import Request @@ -7224,7 +7234,7 @@ async def test_real_jwt_still_requires_license_when_jwt_auth_enabled(monkeypatch patch("litellm.proxy.proxy_server.master_key", "sk-master"), patch("litellm.proxy.proxy_server.prisma_client", None), ): - with pytest.raises(Exception, match='JWT Auth is an enterprise only feature\\. You must be a') as exc_info: + with pytest.raises(Exception, match="JWT Auth is an enterprise only feature\\. You must be a") as exc_info: await user_api_key_auth( request=mock_request, api_key=f"Bearer {jwt_token}", @@ -7263,13 +7273,9 @@ async def test_auth_does_not_rewrite_cached_key_object_back_into_cache(): metadata={"model_rpm_limit": {"gpt-5.4-mini": 3}}, last_refreshed_at=1000.0, ) - await key_cache.async_set_cache( - key=hashed_key, value=stale_token, model_type=UserAPIKeyAuth - ) + await key_cache.async_set_cache(key=hashed_key, value=stale_token, model_type=UserAPIKeyAuth) - fetch_from_db = AsyncMock( - side_effect=AssertionError("cache-hit auth must not touch the DB") - ) + fetch_from_db = AsyncMock(side_effect=AssertionError("cache-hit auth must not touch the DB")) proxy_logging_obj = MagicMock() proxy_logging_obj.internal_usage_cache = MagicMock() @@ -7316,9 +7322,7 @@ async def test_auth_does_not_rewrite_cached_key_object_back_into_cache(): assert result.token == hashed_key fetch_from_db.assert_not_called() - cached_after = await key_cache.async_get_cache( - key=hashed_key, model_type=UserAPIKeyAuth - ) + cached_after = await key_cache.async_get_cache(key=hashed_key, model_type=UserAPIKeyAuth) assert cached_after is not None assert cached_after.last_refreshed_at == 1000.0 assert cached_after.metadata == {"model_rpm_limit": {"gpt-5.4-mini": 3}} @@ -7394,7 +7398,9 @@ class TestJWTAuthUserEmail: assert result.user_email == "resolved@example.com" @pytest.mark.asyncio - @pytest.mark.parametrize("route", ["/mcp-rest/tools/list", "/mcp-rest/tools/call", "/v1/chat/completions", "/user/info"]) + @pytest.mark.parametrize( + "route", ["/mcp-rest/tools/list", "/mcp-rest/tools/call", "/v1/chat/completions", "/user/info"] + ) @pytest.mark.parametrize("active", [False, True, None, "false", 0]) @pytest.mark.parametrize("is_admin", [False, True]) async def test_jwt_auth_rejects_deactivated_user( @@ -7469,9 +7475,7 @@ class TestCheckKeyModelBudgetWithFallback: @pytest.mark.asyncio async def test_within_budget_does_not_reroute(self): - valid_token = UserAPIKeyAuth( - token="test-key", budget_fallbacks={"gpt-4o": ["gpt-4o-mini"]} - ) + valid_token = UserAPIKeyAuth(token="test-key", budget_fallbacks={"gpt-4o": ["gpt-4o-mini"]}) limiter = AsyncMock() limiter.is_key_within_model_budget.return_value = True request_data = {"model": "gpt-4o"} @@ -7496,9 +7500,7 @@ class TestCheckKeyModelBudgetWithFallback: budget_fallbacks={"gpt-4o": ["gpt-4o-mini", "claude-haiku"]}, ) limiter = AsyncMock() - limiter.is_key_within_model_budget.side_effect = litellm.BudgetExceededError( - current_cost=10, max_budget=5 - ) + limiter.is_key_within_model_budget.side_effect = litellm.BudgetExceededError(current_cost=10, max_budget=5) limiter.get_fallback_model_within_budget.return_value = "gpt-4o-mini" request_data = {"model": "gpt-4o"} request = self._make_request() @@ -7512,9 +7514,7 @@ class TestCheckKeyModelBudgetWithFallback: ) assert request_data["model"] == "gpt-4o-mini" - limiter.get_fallback_model_within_budget.assert_awaited_once_with( - user_api_key_dict=valid_token, model="gpt-4o" - ) + limiter.get_fallback_model_within_budget.assert_awaited_once_with(user_api_key_dict=valid_token, model="gpt-4o") # the rerouted model must be visible to a later, separate # `_read_request_body` call on the same `request` (route handlers # re-parse the body from this cache instead of reusing the dict). @@ -7523,9 +7523,7 @@ class TestCheckKeyModelBudgetWithFallback: @pytest.mark.asyncio async def test_raises_when_every_fallback_also_exceeded(self): - valid_token = UserAPIKeyAuth( - token="test-key", budget_fallbacks={"gpt-4o": ["gpt-4o-mini"]} - ) + valid_token = UserAPIKeyAuth(token="test-key", budget_fallbacks={"gpt-4o": ["gpt-4o-mini"]}) limiter = AsyncMock() original_error = litellm.BudgetExceededError(current_cost=10, max_budget=5) limiter.is_key_within_model_budget.side_effect = original_error @@ -7595,9 +7593,7 @@ class TestCheckKeyModelBudgetWithFallback: budget_fallbacks={"gpt-4o": ["gpt-4o-mini"]}, ) limiter = AsyncMock() - limiter.is_key_within_model_budget.side_effect = litellm.BudgetExceededError( - current_cost=10, max_budget=5 - ) + limiter.is_key_within_model_budget.side_effect = litellm.BudgetExceededError(current_cost=10, max_budget=5) limiter.get_fallback_model_within_budget.return_value = "gpt-4o-mini" request_data = {"model": "gpt-4o"} request = self._make_request() @@ -7665,9 +7661,7 @@ class TestCheckKeyModelBudgetWithFallback: budget_fallbacks={"gpt-4o": ["gpt-4o-mini"]}, ) limiter = AsyncMock() - limiter.is_key_within_model_budget.side_effect = litellm.BudgetExceededError( - current_cost=10, max_budget=5 - ) + limiter.is_key_within_model_budget.side_effect = litellm.BudgetExceededError(current_cost=10, max_budget=5) limiter.get_fallback_model_within_budget.return_value = "gpt-4o-mini" request_data = {"model": "gpt-4o"} request = self._make_request() @@ -7747,9 +7741,7 @@ async def test_global_proxy_spend_reads_resettable_proxy_budget_row(): ) assert result == 42.5 - prisma_client.db.litellm_usertable.find_unique.assert_awaited_once_with( - where={"user_id": "litellm-proxy-budget"} - ) + prisma_client.db.litellm_usertable.find_unique.assert_awaited_once_with(where={"user_id": "litellm-proxy-budget"}) @pytest.mark.asyncio @@ -8124,9 +8116,7 @@ async def test_jwt_shaped_key_error_names_enable_jwt_auth_when_disabled(): Prometheus invalid-key filter and the admin UI both substring-match it. Keys that are not JWT-shaped must not pick up the hint. """ - jwt_error = await _proxy_exception_for_key( - "eyJhbGciOiJSUzI1NiJ9.eyJzdWIiOiJzdmMtMSJ9.c2lnbmF0dXJl", {}, True - ) + jwt_error = await _proxy_exception_for_key("eyJhbGciOiJSUzI1NiJ9.eyJzdWIiOiJzdmMtMSJ9.c2lnbmF0dXJl", {}, True) assert jwt_error.code == "401" assert "enable_jwt_auth" in jwt_error.message @@ -8136,9 +8126,7 @@ async def test_jwt_shaped_key_error_names_enable_jwt_auth_when_disabled(): assert "is a JWT" not in jwt_error.message opaque_error = await _proxy_exception_for_key("not-a-jwt-at-all", {}, True) - two_segment_error = await _proxy_exception_for_key( - "eyJhbGciOiJSUzI1NiJ9.eyJzdWIiOiJzdmMtMSJ9", {}, True - ) + two_segment_error = await _proxy_exception_for_key("eyJhbGciOiJSUzI1NiJ9.eyJzdWIiOiJzdmMtMSJ9", {}, True) assert "enable_jwt_auth" not in opaque_error.message assert "enable_jwt_auth" not in two_segment_error.message @@ -8167,9 +8155,7 @@ class TestLitellmReceivedAtStamping: on OTEL being configured to see a true request-arrival timestamp.""" def test_stamped_even_when_otel_is_not_configured(self, monkeypatch): - monkeypatch.setattr( - "litellm.proxy.proxy_server.open_telemetry_logger", None - ) + monkeypatch.setattr("litellm.proxy.proxy_server.open_telemetry_logger", None) request = MagicMock() request.state = SimpleNamespace() @@ -8201,7 +8187,7 @@ class TestLitellmReceivedAtStamping: _RECORDING_DDTRACE = dedent( - ''' + """ import functools import inspect @@ -8250,11 +8236,11 @@ _RECORDING_DDTRACE = dedent( tracer = _Tracer() - ''' + """ ) _DDTRACE_AUTH_PROBE = dedent( - ''' + """ import asyncio import json @@ -8293,7 +8279,7 @@ _DDTRACE_AUTH_PROBE = dedent( asyncio.run(main()) - ''' + """ ) @@ -8440,25 +8426,43 @@ async def test_jwt_builder_returns_every_team_grant_the_key_path_gets(is_proxy_a @pytest.mark.asyncio -@pytest.mark.parametrize("route", ["/v1/messages", "/messages", "/v1/chat/completions", "/chat/completions", "/v1/responses", "/responses"]) +@pytest.mark.parametrize( + "route", ["/v1/messages", "/messages", "/v1/chat/completions", "/chat/completions", "/v1/responses", "/responses"] +) async def test_claude_view_normalizes_before_model_access(monkeypatch, route): from starlette.requests import Request from litellm.proxy.auth.user_api_key_auth import _enforce_key_and_fallback_model_access source = "foo[1m]" encoded = "claude-router-" + source.encode().hex() + "[1m]" - router = litellm.Router(model_list=[{"model_name": source, "litellm_params": {"model": "openai/gpt-4o", "api_key": "sk-fake"}}]) + router = litellm.Router( + model_list=[{"model_name": source, "litellm_params": {"model": "openai/gpt-4o", "api_key": "sk-fake"}}] + ) monkeypatch.setattr(litellm.proxy.proxy_server, "llm_router", router) data = {"model": encoded, "messages": [{"role": "user", "content": "hi"}]} request = Request({"type": "http", "method": "POST", "path": route, "headers": [], "query_string": b""}) token = UserAPIKeyAuth(models=[source]) - await _enforce_key_and_fallback_model_access(valid_token=token, request_data=data, route=route, request=request, llm_model_list=router.model_list, llm_router=router) + await _enforce_key_and_fallback_model_access( + valid_token=token, + request_data=data, + route=route, + request=request, + llm_model_list=router.model_list, + llm_router=router, + ) assert data["model"] == source assert (await request.json())["model"] == source assert json.loads(await request.body())["model"] == source assert request.scope["parsed_body"][1]["model"] == source with pytest.raises(ProxyException): - await _enforce_key_and_fallback_model_access(valid_token=UserAPIKeyAuth(models=["other"]), request_data=data, route=route, request=request, llm_model_list=router.model_list, llm_router=router) + await _enforce_key_and_fallback_model_access( + valid_token=UserAPIKeyAuth(models=["other"]), + request_data=data, + route=route, + request=request, + llm_model_list=router.model_list, + llm_router=router, + ) @pytest.mark.asyncio @@ -8470,10 +8474,18 @@ async def test_claude_view_never_reinterprets_explicit_names(monkeypatch, layer) encoded = "claude-router-666f6f" names = ("foo", "other", encoded) if layer == "literal" else ("foo", "other") alias = {encoded: "other"} - router = litellm.Router(model_list=[{"model_name": name, "litellm_params": {"model": "openai/gpt-4o", "api_key": "sk-fake"}} for name in names], model_group_alias=alias if layer == "router" else None) + router = litellm.Router( + model_list=[ + {"model_name": name, "litellm_params": {"model": "openai/gpt-4o", "api_key": "sk-fake"}} for name in names + ], + model_group_alias=alias if layer == "router" else None, + ) monkeypatch.setattr(litellm.proxy.proxy_server, "llm_router", router) monkeypatch.setattr(litellm, "model_alias_map", alias if layer == "global" else {}) - token = UserAPIKeyAuth(aliases=alias if layer == "key" else {}, router_settings={"model_group_alias": alias} if layer == "hierarchical" else None) + token = UserAPIKeyAuth( + aliases=alias if layer == "key" else {}, + router_settings={"model_group_alias": alias} if layer == "hierarchical" else None, + ) data = {"model": encoded} request = Request({"type": "http", "method": "POST", "path": "/v1/messages", "headers": [], "query_string": b""}) await _normalize_claude_model(data, token, request, "/v1/messages") diff --git a/tests/test_litellm/proxy/db/test_db_lookup_gate.py b/tests/test_litellm/proxy/db/test_db_lookup_gate.py new file mode 100644 index 00000000000..68903170840 --- /dev/null +++ b/tests/test_litellm/proxy/db/test_db_lookup_gate.py @@ -0,0 +1,111 @@ +import asyncio +import time +from typing import Final + +import pytest + +from litellm.proxy.db.db_lookup_gate import DBLookupDeadlineExceeded, DBLookupStallTracker, bounded_db_lookup + + +async def _never_answers() -> None: + await asyncio.Event().wait() + + +class _FakeClock: + def __init__(self) -> None: + self.now = 1000.0 + + def __call__(self) -> float: + return self.now + + +@pytest.mark.asyncio +async def test_bounded_db_lookup_fails_a_stalled_lookup_at_the_deadline_and_records_the_hit(): + tracker: Final = DBLookupStallTracker() + started: Final = time.monotonic() + + with pytest.raises(DBLookupDeadlineExceeded) as exc_info: + await bounded_db_lookup(_never_answers(), name="team", deadline_seconds=0.05, tracker=tracker) + + assert time.monotonic() - started < 2 + assert exc_info.value.lookup == "team" + assert exc_info.value.deadline_seconds == 0.05 + assert str(exc_info.value) == "team lookup did not answer within 0.05s" + assert isinstance(exc_info.value, asyncio.TimeoutError) + assert tracker.stalled_within(30) is True + + +@pytest.mark.asyncio +async def test_bounded_db_lookup_returns_a_prompt_answer_without_recording_a_stall(): + tracker: Final = DBLookupStallTracker() + + async def answers() -> str: + return "row" + + assert await bounded_db_lookup(answers(), name="key", deadline_seconds=0.05, tracker=tracker) == "row" + assert tracker.stalled_within(30) is False + + +@pytest.mark.asyncio +async def test_bounded_db_lookup_fails_a_whole_stalled_burst_within_one_deadline(): + tracker: Final = DBLookupStallTracker() + burst: Final = 200 + started: Final = time.monotonic() + + results: Final = await asyncio.gather( + *( + bounded_db_lookup(_never_answers(), name=f"key-{i}", deadline_seconds=0.1, tracker=tracker) + for i in range(burst) + ), + return_exceptions=True, + ) + + assert time.monotonic() - started < 2 + assert len(results) == burst + assert all(isinstance(result, DBLookupDeadlineExceeded) for result in results) + assert tracker.stalled_within(30) is True + + +@pytest.mark.asyncio +async def test_bounded_db_lookup_fails_at_the_deadline_even_when_the_lookup_absorbs_the_cancel(): + tracker: Final = DBLookupStallTracker() + absorbed: Final = asyncio.Event() + let_go: Final = asyncio.Event() + + async def absorbs_the_cancel() -> str: + try: + await asyncio.Event().wait() + except asyncio.CancelledError: + absorbed.set() + await let_go.wait() + return "late row" + + started: Final = time.monotonic() + with pytest.raises(DBLookupDeadlineExceeded): + await asyncio.wait_for( + bounded_db_lookup(absorbs_the_cancel(), name="key", deadline_seconds=0.05, tracker=tracker), + timeout=2, + ) + + assert time.monotonic() - started < 1 + assert tracker.stalled_within(30) is True + await asyncio.wait_for(absorbed.wait(), timeout=1) + let_go.set() + await asyncio.sleep(0) + + +def test_stall_tracker_reports_a_stall_only_inside_the_window(): + clock: Final = _FakeClock() + tracker: Final = DBLookupStallTracker(clock=clock) + + assert tracker.stalled_within(30) is False + tracker.record_hit() + assert tracker.stalled_within(30) is True + assert tracker.stalled_within(0) is False + clock.now += 29.9 + assert tracker.stalled_within(30) is True + clock.now += 0.2 + assert tracker.stalled_within(30) is False + tracker.record_hit() + tracker.clear() + assert tracker.stalled_within(30) is False diff --git a/tests/test_litellm/proxy/db/test_exception_handler.py b/tests/test_litellm/proxy/db/test_exception_handler.py index 613ca847115..09f4d294ad0 100644 --- a/tests/test_litellm/proxy/db/test_exception_handler.py +++ b/tests/test_litellm/proxy/db/test_exception_handler.py @@ -774,3 +774,18 @@ def test_connection_error_answers_when_prisma_is_mocked_after_import(): with patch.dict(sys.modules, {"prisma": MagicMock()}): assert PrismaDBExceptionHandler.is_database_connection_error(Exception("x")) is False assert PrismaDBExceptionHandler.is_database_connection_error(httpx.ConnectError("refused")) is True + + +def test_db_lookup_deadline_is_a_connection_and_unavailability_error_but_never_a_transport_error(): + """A lookup that hit its deadline fails the request as a 503 and counts as a + DB outage for ``allow_requests_on_db_unavailable``, but it must not be read + as a broken transport: that would send every parked request into + ``attempt_db_reconnect`` and turn a slow database into a reconnect storm.""" + from litellm.proxy.db.db_lookup_gate import DBLookupDeadlineExceeded + + deadline: Final = DBLookupDeadlineExceeded("key", 10.0) + + assert PrismaDBExceptionHandler.is_database_connection_error(deadline) is True + assert PrismaDBExceptionHandler.is_database_service_unavailable_error(deadline) is True + assert PrismaDBExceptionHandler.is_database_transport_error(deadline) is False + assert "temporarily unreachable" in PrismaDBExceptionHandler.database_unavailable_message(deadline) diff --git a/tests/test_litellm/proxy/db/test_spend_counter_reseed.py b/tests/test_litellm/proxy/db/test_spend_counter_reseed.py index ff0b67d426b..ab931277313 100644 --- a/tests/test_litellm/proxy/db/test_spend_counter_reseed.py +++ b/tests/test_litellm/proxy/db/test_spend_counter_reseed.py @@ -12,12 +12,14 @@ from collections.abc import Mapping from datetime import datetime, timedelta, timezone from types import SimpleNamespace from typing import Final +from unittest.mock import AsyncMock import pytest from litellm.caching.dual_cache import DualCache from litellm.caching.in_memory_cache import InMemoryCache from litellm.constants import PROXY_DB_LOOKUP_MAX_CONCURRENCY +from litellm.proxy.db.db_lookup_gate import LoopBoundSemaphore, db_lookup_stall_tracker from litellm.proxy.db.spend_counter_reseed import SpendCounterReseed WINDOW_START = datetime(2026, 8, 1, tzinfo=timezone.utc) @@ -445,6 +447,31 @@ async def test_from_db_returns_none_for_a_missing_project_row(): assert await SpendCounterReseed.from_db(prisma_client=prisma, counter_key="spend:project:proj-1") is None +@pytest.mark.asyncio +async def test_from_db_deadline_covers_the_wait_for_a_gate_slot(monkeypatch: pytest.MonkeyPatch) -> None: + """A saturated gate must fail the lookup at the deadline instead of parking + the request on a gate slot outside the bounded window.""" + gate: Final = LoopBoundSemaphore(1) + monkeypatch.setattr("litellm.proxy.db.spend_counter_reseed.db_lookup_gate", gate) + monkeypatch.setattr("litellm.proxy.db.db_lookup_gate.PROXY_DB_LOOKUP_DEADLINE_SECONDS", 0.05) + find_unique: Final = AsyncMock() + prisma: Final = SimpleNamespace( + db=SimpleNamespace(litellm_verificationtoken=SimpleNamespace(find_unique=find_unique)) + ) + db_lookup_stall_tracker.clear() + try: + async with gate.current(): + result: Final = await asyncio.wait_for( + SpendCounterReseed.from_db(prisma_client=prisma, counter_key="spend:key:abc"), + timeout=1.0, + ) + assert result is None + assert db_lookup_stall_tracker.stalled_within(60.0) + find_unique.assert_not_called() + finally: + db_lookup_stall_tracker.clear() + + @pytest.mark.asyncio async def test_from_db_still_never_reads_the_end_user_row(): """A cold end-user counter keeps seeding from the cached end-user object the auth diff --git a/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py b/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py index 761cd0685f2..ee4c468a460 100644 --- a/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py +++ b/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py @@ -2703,6 +2703,134 @@ async def test_health_readiness_details_returns_200_when_db_down_and_allow_reque assert result["db"] == "disconnected" +@pytest.fixture +def _clear_db_lookup_stall() -> Iterator[None]: + from litellm.proxy.db.db_lookup_gate import db_lookup_stall_tracker + + db_lookup_stall_tracker.clear() + yield + db_lookup_stall_tracker.clear() + + +def _connected_prisma() -> MagicMock: + mock_prisma = MagicMock() + mock_prisma.health_check = AsyncMock(return_value=True) + return mock_prisma + + +def _forget_db_health_cache() -> None: + _health_endpoints_module.db_health_cache = { + "status": "unknown", + "last_updated": datetime.now() - timedelta(seconds=60), + } + + +@pytest.mark.asyncio +async def test_health_readiness_returns_503_stalled_after_a_db_lookup_deadline_hit(_clear_db_lookup_stall): + """The incident's readiness stayed green while every request sat parked on the + database: the probe's own ping is a fresh connection that answers fine. A lookup + that hit its deadline inside the stall window must take the pod out of rotation.""" + from fastapi import Response + + from litellm.proxy.db.db_lookup_gate import db_lookup_stall_tracker + from litellm.proxy.health_endpoints._health_endpoints import health_readiness + + _forget_db_health_cache() + db_lookup_stall_tracker.record_hit() + + response = Response() + with patch( # test-quality-ok: the readiness path reads the proxy-global DB client; it has no injection seam + "litellm.proxy.proxy_server.prisma_client", _connected_prisma() + ): + result = await health_readiness(response=response) + + assert response.status_code == 503 + assert result == {"status": "healthy", "db": "stalled"} + + +@pytest.mark.asyncio +async def test_health_readiness_details_returns_503_stalled_after_a_db_lookup_deadline_hit(_clear_db_lookup_stall): + from fastapi import Response + + from litellm.proxy.db.db_lookup_gate import db_lookup_stall_tracker + from litellm.proxy.health_endpoints._health_endpoints import _get_health_readiness_details + + _forget_db_health_cache() + db_lookup_stall_tracker.record_hit() + + response = Response() + with patch( # test-quality-ok: the readiness path reads the proxy-global DB client; it has no injection seam + "litellm.proxy.proxy_server.prisma_client", _connected_prisma() + ): + result = await _get_health_readiness_details(response=response) + + assert response.status_code == 503 + assert result["db"] == "stalled" + + +@pytest.mark.asyncio +async def test_health_readiness_stays_200_with_stalled_body_when_requests_are_allowed_on_db_unavailable( + _clear_db_lookup_stall, +): + """The fail-open deployment keeps serving through a stalled database, so the pod + must stay in rotation and report the stall through the body, exactly as it does + for a disconnected one.""" + from fastapi import Response + + from litellm.proxy.db.db_lookup_gate import db_lookup_stall_tracker + from litellm.proxy.health_endpoints._health_endpoints import health_readiness + + _forget_db_health_cache() + db_lookup_stall_tracker.record_hit() + + response = Response() + with ( + patch( # test-quality-ok: the readiness path reads the proxy-global DB client; it has no injection seam + "litellm.proxy.proxy_server.prisma_client", _connected_prisma() + ), + patch.dict( # test-quality-ok: the fail-open flag lives in the proxy-global general_settings; no injection seam + "litellm.proxy.proxy_server.general_settings", + {"allow_requests_on_db_unavailable": True}, + ), + ): + result = await health_readiness(response=response) + + assert response.status_code == 200 + assert result == {"status": "healthy", "db": "stalled"} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("hit_recorded", [False, True]) +async def test_health_readiness_reports_connected_without_a_stall_inside_the_window( + _clear_db_lookup_stall, hit_recorded: bool +): + """No deadline hit, or a window of 0 (the opt-out), keeps the ordinary connected + answer, so a healthy pod never leaves rotation over the stall check.""" + from fastapi import Response + + from litellm.proxy.db.db_lookup_gate import db_lookup_stall_tracker + from litellm.proxy.health_endpoints._health_endpoints import health_readiness + + _forget_db_health_cache() + if hit_recorded: + db_lookup_stall_tracker.record_hit() + + response = Response() + with ( + patch( # test-quality-ok: the readiness path reads the proxy-global DB client; it has no injection seam + "litellm.proxy.proxy_server.prisma_client", _connected_prisma() + ), + patch( # test-quality-ok: lowers the module-level stall window to its opt-out value for the recorded-hit case + "litellm.proxy.health_endpoints._health_endpoints.PROXY_DB_LOOKUP_STALL_WINDOW_SECONDS", + 0.0 if hit_recorded else 30.0, + ), + ): + result = await health_readiness(response=response) + + assert response.status_code == 200 + assert result == {"status": "healthy", "db": "connected"} + + @pytest.mark.asyncio async def test_db_health_readiness_check_bounds_hung_health_check(): """ diff --git a/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py b/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py index 0e9c336a9eb..b5e594db701 100644 --- a/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py +++ b/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py @@ -5,12 +5,14 @@ from datetime import datetime from typing import Final from unittest.mock import AsyncMock, MagicMock, patch +import httpx import pytest from litellm._logging import verbose_proxy_logger from litellm.litellm_core_utils.internal_call_metadata import MODEL_ACCESS_GROUP_METADATA_KEY from litellm.proxy._types import SpendLogsPayload, UserAPIKeyAuth from litellm.proxy.collector import SpendEventConsumer +from litellm.proxy.db.db_lookup_gate import DBLookupDeadlineExceeded from litellm.proxy.db.db_spend_update_writer import DBSpendUpdateWriter from litellm.proxy.db.spend_log_tool_index import response_tool_call_names from litellm.proxy.hooks.proxy_track_cost_callback import ( @@ -1679,6 +1681,92 @@ async def test_async_post_call_failure_hook_enriches_auth_error_metadata(): assert metadata["user_api_key_team_alias"] == "my-team-alias" +@pytest.mark.asyncio +async def test_async_post_call_failure_hook_skips_the_key_lookup_when_the_failure_is_a_db_stall(): + logger = _ProxyDBLogger() + user_api_key_dict = UserAPIKeyAuth(api_key="hashed_key") + request_data = { + "model": "gpt-5.6", + "messages": [{"role": "user", "content": "Hello"}], + "metadata": {}, + "litellm_params": {}, + } + + with ( + patch( + "litellm.proxy.db.db_spend_update_writer.DBSpendUpdateWriter.update_database", + new_callable=AsyncMock, + ) as mock_update_database, + patch( + "litellm.proxy.hooks.proxy_track_cost_callback.get_key_object", + new_callable=AsyncMock, + ) as mock_get_key_object, + patch( + "litellm.proxy.hooks.proxy_track_cost_callback.get_team_object", + new_callable=AsyncMock, + ) as mock_get_team_object, + ): + await logger.async_post_call_failure_hook( + request_data=request_data, + original_exception=DBLookupDeadlineExceeded("key", 10.0), + user_api_key_dict=user_api_key_dict, + ) + + mock_get_key_object.assert_not_called() + mock_get_team_object.assert_not_called() + mock_update_database.assert_called_once() + metadata = mock_update_database.call_args[1]["kwargs"]["litellm_params"]["metadata"] + assert metadata["status"] == "failure" + assert metadata["user_api_key"] == "hashed_key" + assert metadata["user_api_key_alias"] is None + + +@pytest.mark.asyncio +async def test_async_post_call_failure_hook_still_enriches_metadata_for_a_non_stall_failure(): + """Only a DBLookupDeadlineExceeded skips the key lookup; a transport error + from the provider call must still resolve the key's alias for the failure row.""" + logger = _ProxyDBLogger() + user_api_key_dict = UserAPIKeyAuth(api_key="hashed_key") + request_data = { + "model": "gpt-5.6", + "messages": [{"role": "user", "content": "Hello"}], + "metadata": {}, + "litellm_params": {}, + } + + mock_key_obj = MagicMock() + mock_key_obj.key_alias = "my-key-alias" + mock_key_obj.user_id = "my-user-id" + mock_key_obj.team_id = "my-team-id" + mock_key_obj.org_id = None + mock_key_obj.project_id = None + + with ( + patch( + "litellm.proxy.db.db_spend_update_writer.DBSpendUpdateWriter.update_database", + new_callable=AsyncMock, + ) as mock_update_database, + patch( + "litellm.proxy.hooks.proxy_track_cost_callback.get_key_object", + new_callable=AsyncMock, + return_value=mock_key_obj, + ) as mock_get_key_object, + patch( + "litellm.proxy.hooks.proxy_track_cost_callback.get_team_object", + new_callable=AsyncMock, + ), + ): + await logger.async_post_call_failure_hook( + request_data=request_data, + original_exception=httpx.ConnectError("boom"), + user_api_key_dict=user_api_key_dict, + ) + + mock_get_key_object.assert_called_once() + metadata = mock_update_database.call_args[1]["kwargs"]["litellm_params"]["metadata"] + assert metadata["user_api_key_alias"] == "my-key-alias" + + @pytest.mark.asyncio async def test_async_post_call_failure_hook_enriches_missing_team_alias(): """ @@ -2035,9 +2123,15 @@ async def test_track_cost_callback_keeps_guardrail_cost_on_cache_hit(): } with ( - patch("litellm.proxy.proxy_server.increment_spend_counters", new_callable=AsyncMock) as mock_increment, # test-quality-ok: the callback imports this from proxy_server inside its body, so there is no injection seam - patch("litellm.proxy.proxy_server.update_cache", new_callable=AsyncMock), # test-quality-ok: same function-body import, no injection seam - patch("litellm.proxy.proxy_server.proxy_logging_obj") as mock_proxy_logging, # test-quality-ok: same function-body import, no injection seam + patch( + "litellm.proxy.proxy_server.increment_spend_counters", new_callable=AsyncMock + ) as mock_increment, # test-quality-ok: the callback imports this from proxy_server inside its body, so there is no injection seam + patch( + "litellm.proxy.proxy_server.update_cache", new_callable=AsyncMock + ), # test-quality-ok: same function-body import, no injection seam + patch( + "litellm.proxy.proxy_server.proxy_logging_obj" + ) as mock_proxy_logging, # test-quality-ok: same function-body import, no injection seam ): mock_proxy_logging.db_spend_update_writer.update_database = AsyncMock() mock_proxy_logging.slack_alerting_instance.customer_spend_alert = AsyncMock() From 58d7cafb97007bc4ed9be04ccff40d8ea0f15d26 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 24 Sep 2026 08:24:27 -0700 Subject: [PATCH 044/120] fix(cost-map): sync openrouter deepseek-v4-pro prices (#42969) Co-authored-by: kerry Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/model_prices_and_context_window_backup.json | 6 +++--- model_prices_and_context_window.json | 6 +++--- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 69a64de6e04..93e3a033578 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -41242,21 +41242,21 @@ "supports_web_search": false }, "openrouter/deepseek/deepseek-v4-pro": { - "input_cost_per_token": 9.24462e-07, + "input_cost_per_token": 9.19242e-07, "input_cost_per_token_cache_hit": 4.4e-08, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 384000, "max_tokens": 384000, "mode": "chat", - "output_cost_per_token": 1.848924e-06, + "output_cost_per_token": 1.838484e-06, "source": "https://openrouter.ai/api/v1/models", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "cache_read_input_token_cost": 7.70385e-08, + "cache_read_input_token_cost": 7.66035e-08, "supports_audio_input": false, "supports_pdf_input": false, "supports_vision": false, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 69a64de6e04..93e3a033578 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -41242,21 +41242,21 @@ "supports_web_search": false }, "openrouter/deepseek/deepseek-v4-pro": { - "input_cost_per_token": 9.24462e-07, + "input_cost_per_token": 9.19242e-07, "input_cost_per_token_cache_hit": 4.4e-08, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 384000, "max_tokens": 384000, "mode": "chat", - "output_cost_per_token": 1.848924e-06, + "output_cost_per_token": 1.838484e-06, "source": "https://openrouter.ai/api/v1/models", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "cache_read_input_token_cost": 7.70385e-08, + "cache_read_input_token_cost": 7.66035e-08, "supports_audio_input": false, "supports_pdf_input": false, "supports_vision": false, From 553f0b6ee7d59e709b8adf9c28ac0a55a4f420ae Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 24 Sep 2026 08:27:42 -0700 Subject: [PATCH 045/120] feat(cost-map): sync azure models, add MAI-Image-2.6, deepseek-v4.1-flash, muse-spark-1.3 (#42970) Co-authored-by: kerry Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- ...odel_prices_and_context_window_backup.json | 69 +++++++++++++++++++ model_prices_and_context_window.json | 69 +++++++++++++++++++ 2 files changed, 138 insertions(+) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 93e3a033578..02d7c3c66eb 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -63384,6 +63384,73 @@ "supports_tool_choice": true, "supports_vision": true }, + "azure_ai/deepseek-v4.1-flash": { + "cache_read_input_token_cost": 8e-09, + "input_cost_per_token": 3.75e-07, + "litellm_provider": "azure_ai", + "max_input_tokens": 1000000, + "max_output_tokens": 384000, + "max_tokens": 384000, + "mode": "chat", + "output_cost_per_token": 1.5e-06, + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "deprecation_date": "2026-12-15", + "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" + }, + "azure_ai/muse-spark-1.3": { + "cache_read_input_token_cost": 1.5e-07, + "input_cost_per_token": 1.25e-06, + "litellm_provider": "azure_ai", + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 4.25e-06, + "source": "https://ai.developer.meta.com/docs/pricing-rate-limits", + "supported_endpoints": [ + "/v1/chat/completions" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "azure_ai/MAI-Image-2.6": { + "input_cost_per_image_token": 8e-06, + "input_cost_per_token": 5e-06, + "litellm_provider": "azure_ai", + "mode": "image_generation", + "output_cost_per_image_token": 3.8e-05, + "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", + "supported_endpoints": [ + "/v1/images/generations", + "/v1/images/edits" + ] + }, + "azure_ai/MAI-Image-2.6-Flash": { + "input_cost_per_image_token": 2.5e-06, + "input_cost_per_token": 1.75e-06, + "litellm_provider": "azure_ai", + "mode": "image_generation", + "output_cost_per_image_token": 1.9e-05, + "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", + "supported_endpoints": [ + "/v1/images/generations", + "/v1/images/edits" + ] + }, "azure_ai/FW-DeepSeek-V4.1-Flash": { "cache_read_input_token_cost": 8e-09, "input_cost_per_token": 3.75e-07, @@ -63445,6 +63512,7 @@ "supports_tool_choice": true }, "azure_ai/FW-GPT-OSS-120B": { + "deprecation_date": "2027-07-01", "cache_read_input_token_cost": 8.2e-08, "input_cost_per_token": 1.65e-07, "litellm_provider": "azure_ai", @@ -63462,6 +63530,7 @@ "supports_tool_choice": true }, "azure_ai/Cohere-command-a-plus-05-2026": { + "deprecation_date": "2026-10-16", "input_cost_per_token": 8e-07, "litellm_provider": "azure_ai", "max_input_tokens": 128000, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 93e3a033578..02d7c3c66eb 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -63384,6 +63384,73 @@ "supports_tool_choice": true, "supports_vision": true }, + "azure_ai/deepseek-v4.1-flash": { + "cache_read_input_token_cost": 8e-09, + "input_cost_per_token": 3.75e-07, + "litellm_provider": "azure_ai", + "max_input_tokens": 1000000, + "max_output_tokens": 384000, + "max_tokens": 384000, + "mode": "chat", + "output_cost_per_token": 1.5e-06, + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "deprecation_date": "2026-12-15", + "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" + }, + "azure_ai/muse-spark-1.3": { + "cache_read_input_token_cost": 1.5e-07, + "input_cost_per_token": 1.25e-06, + "litellm_provider": "azure_ai", + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 4.25e-06, + "source": "https://ai.developer.meta.com/docs/pricing-rate-limits", + "supported_endpoints": [ + "/v1/chat/completions" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "azure_ai/MAI-Image-2.6": { + "input_cost_per_image_token": 8e-06, + "input_cost_per_token": 5e-06, + "litellm_provider": "azure_ai", + "mode": "image_generation", + "output_cost_per_image_token": 3.8e-05, + "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", + "supported_endpoints": [ + "/v1/images/generations", + "/v1/images/edits" + ] + }, + "azure_ai/MAI-Image-2.6-Flash": { + "input_cost_per_image_token": 2.5e-06, + "input_cost_per_token": 1.75e-06, + "litellm_provider": "azure_ai", + "mode": "image_generation", + "output_cost_per_image_token": 1.9e-05, + "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", + "supported_endpoints": [ + "/v1/images/generations", + "/v1/images/edits" + ] + }, "azure_ai/FW-DeepSeek-V4.1-Flash": { "cache_read_input_token_cost": 8e-09, "input_cost_per_token": 3.75e-07, @@ -63445,6 +63512,7 @@ "supports_tool_choice": true }, "azure_ai/FW-GPT-OSS-120B": { + "deprecation_date": "2027-07-01", "cache_read_input_token_cost": 8.2e-08, "input_cost_per_token": 1.65e-07, "litellm_provider": "azure_ai", @@ -63462,6 +63530,7 @@ "supports_tool_choice": true }, "azure_ai/Cohere-command-a-plus-05-2026": { + "deprecation_date": "2026-10-16", "input_cost_per_token": 8e-07, "litellm_provider": "azure_ai", "max_input_tokens": 128000, From cfa2830bde4b73c1b5e35d84fae1712490e9dd18 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 24 Sep 2026 08:30:48 -0700 Subject: [PATCH 046/120] fix(bedrock): extrapolate global cris pricing for gpt-5.4 and gpt-5.5 (#42971) Co-authored-by: kerry Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- ...odel_prices_and_context_window_backup.json | 24 +++++++++---------- model_prices_and_context_window.json | 24 +++++++++---------- 2 files changed, 24 insertions(+), 24 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 02d7c3c66eb..14252b727ac 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -57189,12 +57189,12 @@ ] }, "global.openai.gpt-5.4": { - "input_cost_per_token": 2.75e-06, - "input_cost_per_token_above_272k_tokens": 5.5e-06, - "cache_read_input_token_cost": 2.75e-07, - "cache_read_input_token_cost_above_272k_tokens": 5.5e-07, - "output_cost_per_token": 1.65e-05, - "output_cost_per_token_above_272k_tokens": 2.475e-05, + "input_cost_per_token": 2.5e-06, + "input_cost_per_token_above_272k_tokens": 5e-06, + "cache_read_input_token_cost": 2.5e-07, + "cache_read_input_token_cost_above_272k_tokens": 5e-07, + "output_cost_per_token": 1.5e-05, + "output_cost_per_token_above_272k_tokens": 2.25e-05, "litellm_provider": "bedrock_converse", "max_input_tokens": 1000000, "max_output_tokens": 128000, @@ -57251,12 +57251,12 @@ ] }, "global.openai.gpt-5.5": { - "input_cost_per_token": 5.5e-06, - "input_cost_per_token_above_272k_tokens": 1.1e-05, - "cache_read_input_token_cost": 5.5e-07, - "cache_read_input_token_cost_above_272k_tokens": 1.1e-06, - "output_cost_per_token": 3.3e-05, - "output_cost_per_token_above_272k_tokens": 4.95e-05, + "input_cost_per_token": 5e-06, + "input_cost_per_token_above_272k_tokens": 1e-05, + "cache_read_input_token_cost": 5e-07, + "cache_read_input_token_cost_above_272k_tokens": 1e-06, + "output_cost_per_token": 3e-05, + "output_cost_per_token_above_272k_tokens": 4.5e-05, "litellm_provider": "bedrock_converse", "max_input_tokens": 1000000, "max_output_tokens": 128000, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 02d7c3c66eb..14252b727ac 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -57189,12 +57189,12 @@ ] }, "global.openai.gpt-5.4": { - "input_cost_per_token": 2.75e-06, - "input_cost_per_token_above_272k_tokens": 5.5e-06, - "cache_read_input_token_cost": 2.75e-07, - "cache_read_input_token_cost_above_272k_tokens": 5.5e-07, - "output_cost_per_token": 1.65e-05, - "output_cost_per_token_above_272k_tokens": 2.475e-05, + "input_cost_per_token": 2.5e-06, + "input_cost_per_token_above_272k_tokens": 5e-06, + "cache_read_input_token_cost": 2.5e-07, + "cache_read_input_token_cost_above_272k_tokens": 5e-07, + "output_cost_per_token": 1.5e-05, + "output_cost_per_token_above_272k_tokens": 2.25e-05, "litellm_provider": "bedrock_converse", "max_input_tokens": 1000000, "max_output_tokens": 128000, @@ -57251,12 +57251,12 @@ ] }, "global.openai.gpt-5.5": { - "input_cost_per_token": 5.5e-06, - "input_cost_per_token_above_272k_tokens": 1.1e-05, - "cache_read_input_token_cost": 5.5e-07, - "cache_read_input_token_cost_above_272k_tokens": 1.1e-06, - "output_cost_per_token": 3.3e-05, - "output_cost_per_token_above_272k_tokens": 4.95e-05, + "input_cost_per_token": 5e-06, + "input_cost_per_token_above_272k_tokens": 1e-05, + "cache_read_input_token_cost": 5e-07, + "cache_read_input_token_cost_above_272k_tokens": 1e-06, + "output_cost_per_token": 3e-05, + "output_cost_per_token_above_272k_tokens": 4.5e-05, "litellm_provider": "bedrock_converse", "max_input_tokens": 1000000, "max_output_tokens": 128000, From 991c339946ad943633ad7961fe678ae507bed258 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 24 Sep 2026 08:53:31 -0700 Subject: [PATCH 047/120] fix(cost-map): sync openrouter deepseek v4 flash, v4 pro and v4.1 flash prices (#42974) Co-authored-by: kerry Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- ...odel_prices_and_context_window_backup.json | 22 +++++++++---------- model_prices_and_context_window.json | 22 +++++++++---------- 2 files changed, 22 insertions(+), 22 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 14252b727ac..be43affab4e 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -41242,34 +41242,34 @@ "supports_web_search": false }, "openrouter/deepseek/deepseek-v4-pro": { - "input_cost_per_token": 9.19242e-07, + "input_cost_per_token": 9.15936e-07, "input_cost_per_token_cache_hit": 4.4e-08, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 384000, "max_tokens": 384000, "mode": "chat", - "output_cost_per_token": 1.838484e-06, + "output_cost_per_token": 1.831872e-06, "source": "https://openrouter.ai/api/v1/models", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "cache_read_input_token_cost": 7.66035e-08, + "cache_read_input_token_cost": 7.6328e-08, "supports_audio_input": false, "supports_pdf_input": false, "supports_vision": false, "supports_web_search": false }, "openrouter/deepseek/deepseek-v4.1-flash": { - "input_cost_per_token": 1.4e-07, - "output_cost_per_token": 4.2e-07, - "cache_read_input_token_cost": 4.2e-09, + "input_cost_per_token": 3e-07, + "output_cost_per_token": 1.2e-06, + "cache_read_input_token_cost": 6e-09, "litellm_provider": "openrouter", "max_input_tokens": 1048576, - "max_output_tokens": 943718, - "max_tokens": 943718, + "max_output_tokens": 393216, + "max_tokens": 393216, "mode": "chat", "off_peak_pricing": {"windows":[{"weekdays":["saturday","sunday"],"hours_utc":"00:00-00:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"00:00-01:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"04:00-06:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"10:00-00:00"}],"input_cost_per_token":1.5e-7,"output_cost_per_token":6e-7,"cache_read_input_token_cost":3e-9}, "source": "https://openrouter.ai/api/v1/models", @@ -66491,9 +66491,9 @@ "supports_web_search": true }, "openrouter/deepseek/deepseek-v4-flash": { - "input_cost_per_token": 8.554e-08, - "output_cost_per_token": 1.7108e-07, - "cache_read_input_token_cost": 1.7108e-08, + "input_cost_per_token": 8.4e-08, + "output_cost_per_token": 1.68e-07, + "cache_read_input_token_cost": 1.68e-08, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 384000, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 14252b727ac..be43affab4e 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -41242,34 +41242,34 @@ "supports_web_search": false }, "openrouter/deepseek/deepseek-v4-pro": { - "input_cost_per_token": 9.19242e-07, + "input_cost_per_token": 9.15936e-07, "input_cost_per_token_cache_hit": 4.4e-08, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 384000, "max_tokens": 384000, "mode": "chat", - "output_cost_per_token": 1.838484e-06, + "output_cost_per_token": 1.831872e-06, "source": "https://openrouter.ai/api/v1/models", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "cache_read_input_token_cost": 7.66035e-08, + "cache_read_input_token_cost": 7.6328e-08, "supports_audio_input": false, "supports_pdf_input": false, "supports_vision": false, "supports_web_search": false }, "openrouter/deepseek/deepseek-v4.1-flash": { - "input_cost_per_token": 1.4e-07, - "output_cost_per_token": 4.2e-07, - "cache_read_input_token_cost": 4.2e-09, + "input_cost_per_token": 3e-07, + "output_cost_per_token": 1.2e-06, + "cache_read_input_token_cost": 6e-09, "litellm_provider": "openrouter", "max_input_tokens": 1048576, - "max_output_tokens": 943718, - "max_tokens": 943718, + "max_output_tokens": 393216, + "max_tokens": 393216, "mode": "chat", "off_peak_pricing": {"windows":[{"weekdays":["saturday","sunday"],"hours_utc":"00:00-00:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"00:00-01:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"04:00-06:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"10:00-00:00"}],"input_cost_per_token":1.5e-7,"output_cost_per_token":6e-7,"cache_read_input_token_cost":3e-9}, "source": "https://openrouter.ai/api/v1/models", @@ -66491,9 +66491,9 @@ "supports_web_search": true }, "openrouter/deepseek/deepseek-v4-flash": { - "input_cost_per_token": 8.554e-08, - "output_cost_per_token": 1.7108e-07, - "cache_read_input_token_cost": 1.7108e-08, + "input_cost_per_token": 8.4e-08, + "output_cost_per_token": 1.68e-07, + "cache_read_input_token_cost": 1.68e-08, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 384000, From f0e671f7549665b39a562149d5e654682b9e6a51 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 24 Sep 2026 09:02:02 -0700 Subject: [PATCH 048/120] refactor(types): replace Any with proven types in 13 files (#42937) * refactor(types): replace Any with proven types in 13 files Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * refactor(types): drop unused executor import from utils type-checking block Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(realtime): keep reserved-key filtering on azure realtime health params Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(realtime): pin reserved-key filtering in azure realtime health auth params Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(realtime): exercise the real azure header builder in the reserved-key test Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/integrations/opentelemetry.py | 10 +-- litellm/litellm_core_utils/litellm_logging.py | 12 ++- .../prompt_templates/factory.py | 4 +- litellm/llms/custom_httpx/llm_http_handler.py | 58 +++++++++--- litellm/main.py | 7 +- litellm/proxy/common_request_processing.py | 12 +-- .../key_management_endpoints.py | 2 +- litellm/proxy/proxy_server.py | 12 +-- litellm/proxy/utils.py | 2 +- litellm/realtime_api/main.py | 14 +-- litellm/responses/main.py | 2 +- litellm/router.py | 10 +-- litellm/utils.py | 89 +++++++++++++++---- .../test_health_check_helpers.py | 17 ++++ 14 files changed, 177 insertions(+), 74 deletions(-) diff --git a/litellm/integrations/opentelemetry.py b/litellm/integrations/opentelemetry.py index 749f0ce4fcb..c1531f4e4ae 100644 --- a/litellm/integrations/opentelemetry.py +++ b/litellm/integrations/opentelemetry.py @@ -62,10 +62,10 @@ if TYPE_CHECKING: from litellm.proxy.proxy_server import UserAPIKeyAuth as _UserAPIKeyAuth Span = _Span | Any - Tracer = _Tracer | Any - Context = _Context | Any - SpanExporter = _SpanExporter | Any - UserAPIKeyAuth = _UserAPIKeyAuth | Any + Tracer = _Tracer + Context = _Context + SpanExporter = _SpanExporter + UserAPIKeyAuth = _UserAPIKeyAuth ManagementEndpointLoggingPayload = _ManagementEndpointLoggingPayload | Any else: Span = Any @@ -2730,7 +2730,7 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger): self.handle_callback_failure(callback_name=self.callback_name or "opentelemetry") verbose_logger.exception("OpenTelemetry logging error in set_attributes %s", str(e)) - def _cast_as_primitive_value_type(self, value) -> str | bool | int | float: + def _cast_as_primitive_value_type(self, value: object) -> str | bool | int | float: """ Casts the value to a primitive OTEL type if it is not already a primitive type. diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 2cecec729c2..a6391a2ae27 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -1566,14 +1566,14 @@ class Logging(LiteLLMLoggingBaseClass): attr = "debug" if json_logs: - callattr = getattr(verbose_logger, attr) + callattr = verbose_logger.warning if attr == "warning" else verbose_logger.debug callattr( "RAW RESPONSE:\n{}\n\n".format( self.model_call_details.get("original_response", self.model_call_details) ), ) else: - callattr = getattr(verbose_logger, attr) + callattr = verbose_logger.warning if attr == "warning" else verbose_logger.debug callattr( "RAW RESPONSE:\n{}\n\n".format( self.model_call_details.get("original_response", self.model_call_details) @@ -5882,7 +5882,7 @@ class StandardLoggingPayloadSetup: base_model: str | None, custom_pricing: bool | None, custom_llm_provider: str | None, - init_response_obj: Any | BaseModel | dict, + init_response_obj: object, api_base: str | None = None, ) -> StandardLoggingModelInformation: model_cost_name: Final = _select_model_name_for_cost_calc( @@ -5915,9 +5915,7 @@ class StandardLoggingPayloadSetup: return model_cost_information @staticmethod - def get_final_response_obj( - response_obj: dict, init_response_obj: Any | BaseModel | dict, kwargs: dict - ) -> dict | str | list | None: + def get_final_response_obj(response_obj: dict, init_response_obj: object, kwargs: dict) -> dict | str | list | None: """ Get final response object after redacting the message input/output from logging """ @@ -6360,7 +6358,7 @@ def _get_status_fields( def _extract_response_obj_and_hidden_params( - init_response_obj: Any | BaseModel | dict, + init_response_obj: object, original_exception: Exception | None, ) -> tuple[dict, dict | None]: """Extract response_obj and hidden_params from init_response_obj.""" diff --git a/litellm/litellm_core_utils/prompt_templates/factory.py b/litellm/litellm_core_utils/prompt_templates/factory.py index 8424187dcbc..6fc319c26ae 100644 --- a/litellm/litellm_core_utils/prompt_templates/factory.py +++ b/litellm/litellm_core_utils/prompt_templates/factory.py @@ -446,7 +446,7 @@ def _render_chat_template(env, chat_template: str, bos_token: str, eos_token: st async def _afetch_and_extract_template( - model: str, chat_template: Any | None, get_config_fn, get_template_fn + model: str, chat_template: str | None, get_config_fn, get_template_fn ) -> tuple[str, str, str]: """ Async version: Fetch template and tokens from HuggingFace. @@ -500,7 +500,7 @@ async def _afetch_and_extract_template( def _fetch_and_extract_template( - model: str, chat_template: Any | None, get_config_fn, get_template_fn + model: str, chat_template: str | None, get_config_fn, get_template_fn ) -> tuple[str, str, str]: """ Sync version: Fetch template and tokens from HuggingFace. diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 052978c2680..dba0dee38fc 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -12,6 +12,7 @@ from typing import ( Literal, NamedTuple, Optional, + Protocol, TypedDict, TypeVar, Union, @@ -24,6 +25,7 @@ import httpx from httpx import USE_CLIENT_DEFAULT from httpx._types import FileContent from openai.types.file_deleted import FileDeleted +from typing_extensions import ReadOnly import litellm import litellm.litellm_core_utils @@ -206,6 +208,7 @@ if TYPE_CHECKING: FakeAnthropicMessagesStreamIterator, ) from litellm.llms.base_llm.passthrough.transformation import BasePassthroughConfig + from litellm.proxy._types import UserAPIKeyAuth from litellm.types.llms.openai_evals import ( CancelEvalResponse, CancelRunResponse, @@ -221,6 +224,21 @@ if TYPE_CHECKING: else: LiteLLMLoggingObj = Any + +class _RealtimeClientWebSocket(Protocol): + async def send_text(self, data: str) -> None: ... + + async def close(self, code: int = ..., reason: str | None = ...) -> None: ... + + +class _ResponsesClientWebSocket(Protocol): + async def send_text(self, data: str) -> None: ... + + async def receive_text(self) -> str: ... + + async def close(self, code: int = ..., reason: str | None = ...) -> None: ... + + _ResponseT = TypeVar("_ResponseT") @@ -237,6 +255,17 @@ class _MediaUploadKwargs(TypedDict, total=False): timeout: float | httpx.Timeout +class _SignedBodyKwargs(TypedDict, total=False): + data: ReadOnly[bytes] + json: ReadOnly[dict[str, object]] + + +def _signed_body_kwargs(*, signed_body: bytes | None, data: dict[str, object]) -> _SignedBodyKwargs: + if signed_body is not None: + return {"data": signed_body} + return {"json": data} + + def _google_genai_streaming_hidden_params( *, api_base: str, @@ -318,7 +347,9 @@ def _mask_presigned_request_headers(transformed_request: bytes | str | dict) -> } -def _aws_signing_overrides(optional_params: Mapping[str, Any], litellm_params: Mapping[str, Any]) -> Mapping[str, Any]: +def _aws_signing_overrides( + optional_params: Mapping[str, object], litellm_params: Mapping[str, object] +) -> Mapping[str, object]: return MappingProxyType( { key: litellm_params[key] @@ -2739,7 +2770,7 @@ class BaseLLMHTTPHandler: stream=stream, fake_stream=fake_stream, ) - body_kwargs: Final[dict[str, Any]] = {"data": signed_body} if signed_body is not None else {"json": data} + body_kwargs: Final = _signed_body_kwargs(signed_body=signed_body, data=data) ## LOGGING logging_obj.pre_call( @@ -2926,7 +2957,7 @@ class BaseLLMHTTPHandler: stream=stream, fake_stream=fake_stream, ) - body_kwargs: Final[dict[str, Any]] = {"data": signed_body} if signed_body is not None else {"json": data} + body_kwargs: Final = _signed_body_kwargs(signed_body=signed_body, data=data) ## LOGGING logging_obj.pre_call( @@ -4540,7 +4571,7 @@ class BaseLLMHTTPHandler: api_key=litellm_params.api_key, model=model, ) - body_kwargs: Final[dict[str, Any]] = {"data": signed_body} if signed_body is not None else {"json": data} + body_kwargs: Final = _signed_body_kwargs(signed_body=signed_body, data=data) ## LOGGING logging_obj.pre_call( @@ -4634,7 +4665,7 @@ class BaseLLMHTTPHandler: api_key=litellm_params.api_key, model=model, ) - body_kwargs: Final[dict[str, Any]] = {"data": signed_body} if signed_body is not None else {"json": data} + body_kwargs: Final = _signed_body_kwargs(signed_body=signed_body, data=data) ## LOGGING logging_obj.pre_call( @@ -6186,6 +6217,7 @@ class BaseLLMHTTPHandler: "BasePassthroughConfig", "BaseContainerConfig", BaseEvalsAPIConfig, + BaseRealtimeHTTPConfig, ], ): received_status_code: Final = ( @@ -6300,7 +6332,7 @@ class BaseLLMHTTPHandler: async def async_realtime( self, model: str, - websocket: Any, + websocket: _RealtimeClientWebSocket, logging_obj: LiteLLMLoggingObj, provider_config: BaseRealtimeConfig, headers: dict, @@ -6308,7 +6340,7 @@ class BaseLLMHTTPHandler: api_key: str | None = None, client: Any | None = None, timeout: float | None = None, - user_api_key_dict: Any | None = None, + user_api_key_dict: object | None = None, litellm_metadata: dict[str, object] | None = None, query_params: RealtimeQueryParams | None = None, ): @@ -6483,7 +6515,7 @@ class BaseLLMHTTPHandler: request_data: dict[str, object], logging_obj: LiteLLMLoggingObj, timeout: float | httpx.Timeout, - provider_config: Any | None = None, + provider_config: BaseRealtimeHTTPConfig | None = None, model: str | None = None, extra_headers: dict[str, object] | None = None, client: HTTPHandler | AsyncHTTPHandler | None = None, @@ -6555,7 +6587,7 @@ class BaseLLMHTTPHandler: sdp_body: bytes, logging_obj: LiteLLMLoggingObj, timeout: float | httpx.Timeout, - provider_config: Any | None = None, + provider_config: BaseRealtimeHTTPConfig | None = None, model: str | None = None, session_config: dict[str, object] | None = None, extra_headers: dict[str, object] | None = None, @@ -6633,13 +6665,13 @@ class BaseLLMHTTPHandler: async def async_responses_websocket( self, model: str, - websocket: Any, + websocket: _ResponsesClientWebSocket, logging_obj: LiteLLMLoggingObj, responses_api_provider_config: BaseResponsesAPIConfig | None, api_base: str | None = None, api_key: str | None = None, timeout: float | None = None, - user_api_key_dict: Any | None = None, + user_api_key_dict: "UserAPIKeyAuth | None" = None, litellm_metadata: dict[str, object] | None = None, custom_llm_provider: str | None = None, first_message: str | None = None, @@ -7850,7 +7882,7 @@ class BaseLLMHTTPHandler: def video_create_character_handler( self, name: str, - video: Any, + video: FileTypes, video_provider_config: BaseVideoConfig, custom_llm_provider: str, litellm_params, @@ -7934,7 +7966,7 @@ class BaseLLMHTTPHandler: async def async_video_create_character_handler( self, name: str, - video: Any, + video: FileTypes, video_provider_config: BaseVideoConfig, custom_llm_provider: str, litellm_params, diff --git a/litellm/main.py b/litellm/main.py index ceac729d3f0..98bb5126a90 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -57,6 +57,7 @@ from litellm.utils import ( # Logging is imported lazily when needed to avoid loading litellm_logging at import time if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging + from litellm.router import Router from litellm.types.utils import TokenCountResponse from litellm.constants import ( @@ -351,7 +352,7 @@ class LiteLLM: class Chat: - def __init__(self, params, router_obj: Any | None): + def __init__(self, params, router_obj: "Router | None"): self.params = params if self.params.get("acompletion", False) is True: self.params.pop("acompletion") @@ -361,7 +362,7 @@ class Chat: class Completions: - def __init__(self, params, router_obj: Any | None): + def __init__(self, params, router_obj: "Router | None"): self.params = params self.router_obj = router_obj @@ -377,7 +378,7 @@ class Completions: class AsyncCompletions: - def __init__(self, params, router_obj: Any | None): + def __init__(self, params, router_obj: "Router | None"): self.params = params self.router_obj = router_obj diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 904070cfadd..7e1989b0aab 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -218,7 +218,7 @@ ProxyRouteType: TypeAlias = Literal[ from litellm.llms.anthropic.chat.transformation import AnthropicConfig # Type alias for streaming chunk serializer (chunk after hooks + cost injection -> wire format) -StreamChunkSerializer = Callable[[Any], str] +StreamChunkSerializer = Callable[[object], str] # Type alias for streaming error serializer (ProxyException -> wire format) StreamErrorSerializer = Callable[[ProxyException], str] @@ -459,7 +459,7 @@ async def _bill_partial_streamed_spend_on_disconnect(request_data: dict, respons return True -async def _cancel_pending_gather_tasks(tasks: list["asyncio.Task[Any]"]) -> None: +async def _cancel_pending_gather_tasks(tasks: Sequence["asyncio.Task[object]"]) -> None: pending_tasks: Final = [task for task in tasks if not task.done()] for task in pending_tasks: task.cancel() @@ -3145,7 +3145,7 @@ class ProxyBaseLLMRequestProcessing: logging_obj._on_detached_stream_failure = _on_detached_stream_failure - def _is_streaming_response(self, response: Any) -> bool: + def _is_streaming_response(self, response: object) -> bool: """ Check if the response object is actually a streaming response by inspecting its type. @@ -3259,7 +3259,7 @@ class ProxyBaseLLMRequestProcessing: async def _handle_non_streaming_allm_passthrough_route( self, - response: Any, + response: _UpstreamHttpResponse, proxy_logging_obj: "ProxyLogging", user_api_key_dict: "UserAPIKeyAuth", custom_headers: Mapping[str, str], @@ -3852,7 +3852,7 @@ class ProxyBaseLLMRequestProcessing: @staticmethod async def async_streaming_data_generator( - response: Any, + response: object, user_api_key_dict: UserAPIKeyAuth, request_data: dict, proxy_logging_obj: ProxyLogging, @@ -3993,7 +3993,7 @@ class ProxyBaseLLMRequestProcessing: @staticmethod def async_sse_data_generator( - response: Any, + response: object, user_api_key_dict: UserAPIKeyAuth, request_data: dict, proxy_logging_obj: ProxyLogging, diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 306ea90d7f1..2f1dc270cf6 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -453,7 +453,7 @@ def _regenerate_request_as_update_request(key: str, data: RegenerateKeyRequest) ) if not changed_fields: return None - return UpdateKeyRequest(key=key, **changed_fields) + return UpdateKeyRequest.model_validate(MappingProxyType({"key": key, **changed_fields})) class _LegacyDumpable(Protocol): diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 9fb72d2d1b5..c0071aa7c81 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -8700,9 +8700,9 @@ class ProxyConfig: @staticmethod def _merge_config_and_db_search_tools( - config_search_tools: list[SearchToolTypedDict], - db_search_tools: list[dict[str, Any]], - ) -> list[dict[str, Any]]: + config_search_tools: Sequence[SearchToolTypedDict], + db_search_tools: Sequence[dict[str, object]], + ) -> list[dict[str, object]]: db_tool_names: Final = {tool.get("search_tool_name") for tool in db_search_tools} return [ *[ @@ -9273,10 +9273,10 @@ _EMPTY_HEADERS: Final[Mapping[str, str]] = MappingProxyType({}) async def _iter_with_keepalive( - aiter: AsyncIterator[Any], + aiter: AsyncIterator[object], resolve_keepalive_seconds: Callable[[object], float], keepalive_seconds: float, -) -> AsyncGenerator[Any, None]: +) -> AsyncGenerator[object, None]: """Wrap `aiter` with idle-gap heartbeats, re-resolving the interval after each real chunk via `resolve_keepalive_seconds`. A mid-stream router fallback can swap in a deployment with a different keepalive policy, including one that @@ -9287,7 +9287,7 @@ async def _iter_with_keepalive( actually produced it, in both directions. While the interval is <= 0, no task is created and no timeout is awaited: a chunk is forwarded the moment it arrives, at the same cost as a bare `async for`.""" - pending: asyncio.Task[Any] | None = None # rebind-ok: rebound each loop iteration + pending: asyncio.Task[object] | None = None # rebind-ok: rebound each loop iteration current_keepalive_seconds = keepalive_seconds # rebind-ok: re-resolved after each chunk try: while True: diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index fe161f5d50b..c617047fad9 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -3956,7 +3956,7 @@ class ProxyLogging: request_data: dict, # mutable-ok: same request-payload shape the hooks mutate pipelines: "tuple[tuple[str, GuardrailPipeline], ...]", translation: "tuple[str, BaseTranslation]", - ) -> "AsyncGenerator[Any, None]": + ) -> "AsyncGenerator[object, None]": """ Execute post_call policy pipelines against a streamed response. diff --git a/litellm/realtime_api/main.py b/litellm/realtime_api/main.py index acc42c44c04..28814741852 100644 --- a/litellm/realtime_api/main.py +++ b/litellm/realtime_api/main.py @@ -127,15 +127,15 @@ def _get_realtime_http_provider_config( @wrapper_client async def acreate_realtime_client_secret( model: str | None = None, - session: dict[str, Any] | None = None, - expires_after: dict[str, Any] | None = None, + session: Mapping[str, object] | None = None, + expires_after: Mapping[str, object] | None = None, timeout: float | None = None, **kwargs, ): req: Final = RealtimeClientSecretRequest( model=model, - session=RealtimeSessionConfig(**session) if session else None, - expires_after=RealtimeExpiresAfter(**expires_after) if expires_after else None, + session=RealtimeSessionConfig.model_validate(session) if session else None, + expires_after=RealtimeExpiresAfter.model_validate(expires_after) if expires_after else None, ) model_name = (req.session.model if req.session is not None else None) or req.model or "gpt-4o-realtime-preview" litellm_logging_obj: Final[LiteLLMLogging] = kwargs.get("litellm_logging_obj") @@ -614,12 +614,14 @@ def _azure_realtime_health_protocol( def _realtime_health_check_auth_headers( - custom_llm_provider: str, api_key: str | None, model_params: Mapping[str, Any] + custom_llm_provider: str, api_key: str | None, model_params: Mapping[str, object] ) -> Mapping[str, str]: if custom_llm_provider == "azure": return azure_realtime.get_auth_headers( api_key=api_key, - azure_ad_token=(None if api_key else get_azure_ad_token(GenericLiteLLMParams(**model_params))), + azure_ad_token=( + None if api_key else get_azure_ad_token(GenericLiteLLMParams.model_validate(dict(model_params))) + ), ) if api_key is None: return _EMPTY_AUTH_HEADERS diff --git a/litellm/responses/main.py b/litellm/responses/main.py index 884ea7e217d..93c72bc2d3b 100644 --- a/litellm/responses/main.py +++ b/litellm/responses/main.py @@ -549,7 +549,7 @@ def _will_bridge_to_chat_completions( @contextmanager def _prompt_management_sees_a_provisional_message_list( - kwargs: dict[str, Any], # mutable-ok: the signal is read and popped out of the caller's own kwargs + kwargs: dict[str, object], # mutable-ok: the signal is read and popped out of the caller's own kwargs bridged: bool, ) -> Generator[None, None]: """Tell the cache-control hook that this layer's messages are not the ones sent upstream. diff --git a/litellm/router.py b/litellm/router.py index 62042e1c969..d328fbbb12f 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -4231,7 +4231,7 @@ class Router: models: Final = [m.strip() for m in model.split(",")] async def _async_completion_no_exceptions( - model_name: str, messages: list[dict[str, str]], stream: bool, **kwargs: Any + model_name: str, messages: list[dict[str, str]], stream: bool, **kwargs: object ) -> ModelResponse | CustomStreamWrapper | Exception: """ Wrapper around self.acompletion that catches exceptions and returns them as a result @@ -6736,7 +6736,7 @@ class Router: # Handle asynchronous call types async def async_wrapper( custom_llm_provider: str | None = None, - client: Any | None = None, + client: AsyncOpenAI | None = None, **kwargs, ): if call_type == "assistants": @@ -8441,7 +8441,7 @@ class Router: return self._has_content_policy_fallback(model, kwargs) def _should_raise_anthropic_refusal_error( - self, model: str, original_generic_function: Callable, response: object, kwargs: Mapping[str, Any] + self, model: str, original_generic_function: Callable, response: object, kwargs: Mapping[str, object] ) -> bool: """ The /v1/messages twin of _should_raise_content_policy_error: an Anthropic safeguard @@ -10318,7 +10318,7 @@ class Router: ) @staticmethod - def _widest_configured_limit(model_infos: Sequence[Mapping[str, Any]], field: str) -> int | None: + def _widest_configured_limit(model_infos: Sequence[Mapping[str, object]], field: str) -> int | None: """The largest usable value of ``field`` across a group's configured model_info blocks.""" limits: Final = tuple( limit @@ -13409,7 +13409,7 @@ class Router: self, model: str, request_kwargs: dict, - messages: list[dict[str, Any]] | None, + messages: list[dict[str, object]] | None, ) -> RoutingContext: """ Build a RoutingContext for `model`, run it through `self.routing_plugins` diff --git a/litellm/utils.py b/litellm/utils.py index 64097021dff..81258b8ca77 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -369,7 +369,7 @@ if TYPE_CHECKING: ) from litellm.litellm_core_utils.rules import Rules from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper - from litellm.litellm_core_utils.thread_pool_executor import executor + from litellm.litellm_core_utils.thread_pool_executor import BoundedLoggingThreadPoolExecutor from litellm.llms.base_llm.anthropic_messages.transformation import ( BaseAnthropicMessagesConfig, ) @@ -683,12 +683,12 @@ def load_credentials_from_list(kwargs: dict): Updates kwargs with the credentials if credential_name in kwarg """ # Access CredentialAccessor via module to trigger lazy loading if needed - CredentialAccessor: Final = getattr(sys.modules[__name__], "CredentialAccessor") + credential_accessor: Final[type[CredentialAccessor]] = getattr(sys.modules[__name__], "CredentialAccessor") credential_name: Final = kwargs.get("litellm_credential_name") if not credential_name: return - credential: Final = CredentialAccessor.find_credential(credential_name) + credential: Final = credential_accessor.find_credential(credential_name) if credential is None: verbose_logger.warning( "litellm_credential_name=%s matched none of the %d loaded credentials; the request runs without it", @@ -894,6 +894,45 @@ class _NamedFile(Protocol): def name(self) -> object: ... +class _LoggingClassGetter(Protocol): + def __call__(self) -> type[LiteLLMLoggingObject]: ... + + +class _ResponseMetadataUpdater(Protocol): + def __call__( + self, + result: object, + logging_obj: LiteLLMLoggingObject, + model: str | None, + kwargs: dict[str, object], + start_time: datetime.datetime, + end_time: datetime.datetime, + include_overhead: bool = True, + ) -> None: ... + + +class _SupportedOpenAIParamsGetter(Protocol): + def __call__( + self, + model: str, + custom_llm_provider: str | None = None, + request_type: Literal["chat_completion", "embeddings", "transcription"] = "chat_completion", + base_model: str | None = None, + ) -> list[str] | None: ... + + +class _NestedPathChecker(Protocol): + def __call__(self, path: str) -> bool: ... + + +class _NestedValueDeleter(Protocol): + def __call__(self, data: dict[str, object], path: str) -> dict[str, object]: ... + + +class _BaseModelFromMetadataGetter(Protocol): + def __call__(self, metadata: Mapping[str, object] | None) -> str | None: ... + + def _ocr_document_summary(document: object) -> str: if not isinstance(document, Mapping): return "default-message-value" @@ -978,7 +1017,9 @@ def function_setup( len(litellm.input_callback) > 0 or len(litellm.success_callback) > 0 or len(litellm.failure_callback) > 0 ) and len(callback_list) == 0: callback_list = list(set(litellm.input_callback + litellm.success_callback + litellm.failure_callback)) - get_set_callbacks: Final = getattr(sys.modules[__name__], "get_set_callbacks") + get_set_callbacks: Final[Callable[[], Callable[..., None]]] = getattr( + sys.modules[__name__], "get_set_callbacks" + ) get_set_callbacks()(callback_list=callback_list, function_id=function_id) ## ASYNC CALLBACKS - safety net for callbacks added via direct append if len(litellm.input_callback) > 0: @@ -1223,7 +1264,9 @@ def function_setup( call_type=call_type, ): stream = True - get_litellm_logging_class: Final = getattr(sys.modules[__name__], "get_litellm_logging_class") + get_litellm_logging_class: Final[_LoggingClassGetter] = getattr( + sys.modules[__name__], "get_litellm_logging_class" + ) # Victim for object pool logging_obj = get_litellm_logging_class()( # rebind-ok: 2nd assignment to logging_obj (see initial None above) model=model, @@ -1761,7 +1804,9 @@ def client(original_function): return litellm.stream_chunk_builder(chunks, messages=kwargs.get("messages", None)) else: # RETURN RESULT - update_response_metadata = getattr(sys.modules[__name__], "update_response_metadata") + update_response_metadata: _ResponseMetadataUpdater = getattr( + sys.modules[__name__], "update_response_metadata" + ) update_response_metadata( result=result, logging_obj=logging_obj, @@ -1802,7 +1847,9 @@ def client(original_function): kwargs=kwargs, ) - _update_response_metadata: Final = getattr(sys.modules[__name__], "update_response_metadata") + _update_response_metadata: Final[_ResponseMetadataUpdater] = getattr( + sys.modules[__name__], "update_response_metadata" + ) _update_response_metadata( result=result, logging_obj=logging_obj, @@ -1817,7 +1864,7 @@ def client(original_function): # Copy the current context to propagate it to the background thread # This is essential for OpenTelemetry span context propagation ctx: Final = contextvars.copy_context() - executor: Final = getattr(sys.modules[__name__], "executor") + executor: Final[BoundedLoggingThreadPoolExecutor] = getattr(sys.modules[__name__], "executor") executor.submit( ctx.run, logging_obj.success_handler, @@ -1910,7 +1957,9 @@ def client(original_function): print_args_passed_to_litellm(original_function, args, kwargs) start_time: Final = datetime.datetime.now() result = None - _update_response_metadata: Final = getattr(sys.modules[__name__], "update_response_metadata") + _update_response_metadata: Final[_ResponseMetadataUpdater] = getattr( + sys.modules[__name__], "update_response_metadata" + ) logging_obj: LiteLLMLoggingObject | None = kwargs.get("litellm_logging_obj", None) LLMCachingHandler: Final = _get_cached_llm_caching_handler() _llm_caching_handler: Final[LLMCachingHandler] = LLMCachingHandler( @@ -3678,7 +3727,9 @@ def get_optional_params_embeddings( **kwargs, ): # Lazy load get_supported_openai_params - get_supported_openai_params: Final = getattr(sys.modules[__name__], "get_supported_openai_params") + get_supported_openai_params: Final[_SupportedOpenAIParamsGetter] = getattr( + sys.modules[__name__], "get_supported_openai_params" + ) # retrieve all parameters passed to the function passed_params: Final = locals() @@ -4469,7 +4520,9 @@ def get_optional_params( message=f"{custom_llm_provider} does not support parameters: {list(unsupported_params.keys())}, for model={model}. To drop these, set `litellm.drop_params=True` or for proxy:\n\n`litellm_settings:\n drop_params: true`\n. \n If you want to use these params dynamically send allowed_openai_params={list(unsupported_params.keys())} in your request.", ) - get_supported_openai_params: Final = getattr(sys.modules[__name__], "get_supported_openai_params") + get_supported_openai_params: Final[_SupportedOpenAIParamsGetter] = getattr( + sys.modules[__name__], "get_supported_openai_params" + ) supported_params = get_supported_openai_params( model=model, custom_llm_provider=custom_llm_provider, base_model=base_model ) @@ -4640,9 +4693,9 @@ def get_optional_params( drop_params=bool(drop_params), ) elif custom_llm_provider == "bedrock": - BedrockModelInfo: Final = getattr(sys.modules[__name__], "BedrockModelInfo") - bedrock_route: Final = BedrockModelInfo.get_bedrock_route(model) - bedrock_base_model: Final = BedrockModelInfo.get_base_model(model) + bedrock_model_info: Final[type[BedrockModelInfo]] = getattr(sys.modules[__name__], "BedrockModelInfo") + bedrock_route: Final = bedrock_model_info.get_bedrock_route(model) + bedrock_base_model: Final = bedrock_model_info.get_base_model(model) if bedrock_route == "converse" or bedrock_route == "converse_like": optional_params = litellm.AmazonConverseConfig().map_openai_params( model=model, @@ -4680,7 +4733,7 @@ def get_optional_params( drop_params=bool(drop_params), ) if bedrock_route == "claude_platform": - optional_params = BedrockModelInfo.map_claude_platform_auth_params( + optional_params = bedrock_model_info.map_claude_platform_auth_params( passed_params=passed_params, optional_params=optional_params ) elif custom_llm_provider == "cloudflare": @@ -4954,8 +5007,8 @@ def get_optional_params( # Apply nested drops from additional_drop_params if additional_drop_params: - is_nested_path: Final = getattr(sys.modules[__name__], "is_nested_path") - delete_nested_value: Final = getattr(sys.modules[__name__], "delete_nested_value") + is_nested_path: Final[_NestedPathChecker] = getattr(sys.modules[__name__], "is_nested_path") + delete_nested_value: Final[_NestedValueDeleter] = getattr(sys.modules[__name__], "delete_nested_value") nested_paths: Final = [p for p in additional_drop_params if is_nested_path(p)] for path in nested_paths: optional_params = delete_nested_value(optional_params, path) @@ -7852,7 +7905,7 @@ def _get_base_model_from_metadata(model_call_details=None): return _base_model metadata: Final = litellm_params.get("metadata") or {} - _get_base_model_from_litellm_call_metadata: Callable[..., str | None] = getattr( + _get_base_model_from_litellm_call_metadata: _BaseModelFromMetadataGetter = getattr( sys.modules[__name__], "_get_base_model_from_litellm_call_metadata" ) base_model_from_metadata: Final = _get_base_model_from_litellm_call_metadata(metadata=metadata) diff --git a/tests/test_litellm/litellm_core_utils/test_health_check_helpers.py b/tests/test_litellm/litellm_core_utils/test_health_check_helpers.py index 89b377af3a0..12134ba988a 100644 --- a/tests/test_litellm/litellm_core_utils/test_health_check_helpers.py +++ b/tests/test_litellm/litellm_core_utils/test_health_check_helpers.py @@ -2,6 +2,7 @@ import struct import zlib +from types import MappingProxyType from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -482,3 +483,19 @@ async def test_ocr_health_check_sends_the_document_kind_the_provider_config_acce document = mock_aocr.call_args.kwargs["document"] assert document["type"] == expected_document_type assert document[expected_document_type].startswith(expected_uri_prefix) + + +def test_realtime_health_check_azure_ad_params_drop_reserved_keys(): + from litellm.realtime_api import main as realtime_main + + seen = [] + with patch.object(realtime_main, "get_azure_ad_token", lambda params: seen.append(params) or "ad-token"): + headers = realtime_main._realtime_health_check_auth_headers( + "azure", + None, + MappingProxyType({"api_base": "https://x.openai.azure.com", "self": 1, "params": 2, "__class__": 3}), + ) + + assert dict(headers) == {"Authorization": "Bearer ad-token"} + assert seen[0].api_base == "https://x.openai.azure.com" + assert seen[0].model_extra == {} From cc7aaae7482d9cffbae66c34b3462a1df297715d Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 24 Sep 2026 09:27:11 -0700 Subject: [PATCH 049/120] fix(cost-map): sync openrouter deepseek-v4-pro prices (#42978) Co-authored-by: kerry Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/model_prices_and_context_window_backup.json | 6 +++--- model_prices_and_context_window.json | 6 +++--- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index be43affab4e..926672b19a2 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -41242,21 +41242,21 @@ "supports_web_search": false }, "openrouter/deepseek/deepseek-v4-pro": { - "input_cost_per_token": 9.15936e-07, + "input_cost_per_token": 9.1263e-07, "input_cost_per_token_cache_hit": 4.4e-08, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 384000, "max_tokens": 384000, "mode": "chat", - "output_cost_per_token": 1.831872e-06, + "output_cost_per_token": 1.82526e-06, "source": "https://openrouter.ai/api/v1/models", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "cache_read_input_token_cost": 7.6328e-08, + "cache_read_input_token_cost": 7.60525e-08, "supports_audio_input": false, "supports_pdf_input": false, "supports_vision": false, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index be43affab4e..926672b19a2 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -41242,21 +41242,21 @@ "supports_web_search": false }, "openrouter/deepseek/deepseek-v4-pro": { - "input_cost_per_token": 9.15936e-07, + "input_cost_per_token": 9.1263e-07, "input_cost_per_token_cache_hit": 4.4e-08, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 384000, "max_tokens": 384000, "mode": "chat", - "output_cost_per_token": 1.831872e-06, + "output_cost_per_token": 1.82526e-06, "source": "https://openrouter.ai/api/v1/models", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "cache_read_input_token_cost": 7.6328e-08, + "cache_read_input_token_cost": 7.60525e-08, "supports_audio_input": false, "supports_pdf_input": false, "supports_vision": false, From bd55afc0f87cfe88f5afe8eaba8888e1bdf9ac56 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 24 Sep 2026 09:30:37 -0700 Subject: [PATCH 050/120] fix(cost-map): add vertex ai priority audio input prices for gemini flash rows (#42980) Co-authored-by: kerry Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/model_prices_and_context_window_backup.json | 4 ++++ model_prices_and_context_window.json | 4 ++++ 2 files changed, 8 insertions(+) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 926672b19a2..47cdfecfc0f 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -26063,6 +26063,7 @@ "input_cost_per_token_batches": 1.5e-07, "input_cost_per_token_flex": 1.5e-07, "input_cost_per_token_priority": 5.4e-07, + "input_cost_per_audio_token_priority": 1.8e-06, "output_cost_per_token_batches": 1.25e-06, "output_cost_per_token_flex": 1.25e-06, "output_cost_per_token_priority": 4.5e-06, @@ -26400,6 +26401,7 @@ "input_cost_per_token_batches": 1.25e-07, "input_cost_per_token_flex": 1.25e-07, "input_cost_per_token_priority": 4.5e-07, + "input_cost_per_audio_token_priority": 9e-07, "litellm_provider": "vertex_ai-language-models", "max_input_tokens": 1048576, "max_output_tokens": 65536, @@ -26595,6 +26597,7 @@ "input_cost_per_token_batches": 5e-08, "input_cost_per_token_flex": 5e-08, "input_cost_per_token_priority": 1.8e-07, + "input_cost_per_audio_token_priority": 5.4e-07, "output_cost_per_token_batches": 2e-07, "output_cost_per_token_flex": 2e-07, "output_cost_per_token_priority": 7.2e-07, @@ -49499,6 +49502,7 @@ "input_cost_per_token_batches": 1.25e-07, "input_cost_per_token_flex": 1.25e-07, "input_cost_per_token_priority": 4.5e-07, + "input_cost_per_audio_token_priority": 9e-07, "litellm_provider": "vertex_ai-language-models", "max_input_tokens": 1048576, "max_output_tokens": 65536, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 926672b19a2..47cdfecfc0f 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -26063,6 +26063,7 @@ "input_cost_per_token_batches": 1.5e-07, "input_cost_per_token_flex": 1.5e-07, "input_cost_per_token_priority": 5.4e-07, + "input_cost_per_audio_token_priority": 1.8e-06, "output_cost_per_token_batches": 1.25e-06, "output_cost_per_token_flex": 1.25e-06, "output_cost_per_token_priority": 4.5e-06, @@ -26400,6 +26401,7 @@ "input_cost_per_token_batches": 1.25e-07, "input_cost_per_token_flex": 1.25e-07, "input_cost_per_token_priority": 4.5e-07, + "input_cost_per_audio_token_priority": 9e-07, "litellm_provider": "vertex_ai-language-models", "max_input_tokens": 1048576, "max_output_tokens": 65536, @@ -26595,6 +26597,7 @@ "input_cost_per_token_batches": 5e-08, "input_cost_per_token_flex": 5e-08, "input_cost_per_token_priority": 1.8e-07, + "input_cost_per_audio_token_priority": 5.4e-07, "output_cost_per_token_batches": 2e-07, "output_cost_per_token_flex": 2e-07, "output_cost_per_token_priority": 7.2e-07, @@ -49499,6 +49502,7 @@ "input_cost_per_token_batches": 1.25e-07, "input_cost_per_token_flex": 1.25e-07, "input_cost_per_token_priority": 4.5e-07, + "input_cost_per_audio_token_priority": 9e-07, "litellm_provider": "vertex_ai-language-models", "max_input_tokens": 1048576, "max_output_tokens": 65536, From 571ada0b0f09c0356055626160f78c583a9b2998 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 24 Sep 2026 09:36:25 -0700 Subject: [PATCH 051/120] feat(rust_bridge): mark native streams with the x-litellm-rust header (#42758) Non-streaming responses served by the Rust core already carry x-litellm-rust: true through _hidden_params.additional_headers, which the SDK exposes and the gateway renders as a response header. Native streams did not, because the lifecycle Stream and SyncStream objects had nowhere to hold hidden params and the marker writer skips objects without them. Give both stream classes the same _hidden_params bag every other litellm response has, so the existing marker attaches without wrapping the stream or changing its identity. Co-authored-by: Yujong Lee --- litellm/rust_bridge/lifecycle.py | 2 + .../test_litellm/rust_bridge/test_runtime.py | 56 +++++++++++++++++++ .../messages/test_callbacks.py | 3 + 3 files changed, 61 insertions(+) diff --git a/litellm/rust_bridge/lifecycle.py b/litellm/rust_bridge/lifecycle.py index b3b7a1888c3..2f243e8c212 100644 --- a/litellm/rust_bridge/lifecycle.py +++ b/litellm/rust_bridge/lifecycle.py @@ -81,6 +81,7 @@ class Stream(AsyncIterator[object]): def __init__(self, execution: Execution) -> None: self._execution: Final = execution self._done = False + self._hidden_params: dict[str, object] = {} # mutable-ok: header writers mutate _hidden_params in place def __aiter__(self) -> Stream: return self @@ -117,6 +118,7 @@ class SyncStream(Iterator[object]): def __init__(self, execution: Execution) -> None: self._execution: Final = execution self._done = False + self._hidden_params: dict[str, object] = {} # mutable-ok: header writers mutate _hidden_params in place def __iter__(self) -> SyncStream: return self diff --git a/tests/test_litellm/rust_bridge/test_runtime.py b/tests/test_litellm/rust_bridge/test_runtime.py index bff7ded3114..bbe25e0de13 100644 --- a/tests/test_litellm/rust_bridge/test_runtime.py +++ b/tests/test_litellm/rust_bridge/test_runtime.py @@ -12,6 +12,7 @@ from litellm.router_utils.add_retry_fallback_headers import get_hidden_params_di from litellm.rust_bridge import bindings, configuration, runtime from litellm.rust_bridge.catalog import Delivery, Route, RouteContext, RouteRule from litellm.rust_bridge.configuration import Rollout +from litellm.rust_bridge.lifecycle import Complete, Open, Stream, SyncStream, Yield class RustBridgeDeclined(Exception): @@ -264,6 +265,61 @@ async def test_native_response_marker_reaches_caller_with_existing_metadata(shap } +class ScriptedStreamExecution: + def __init__(self, chunks: tuple[bytes, ...]) -> None: + self._steps: Final = iter((*(Yield(chunk) for chunk in chunks), Complete(None))) + self.closed = False + + def start(self) -> Open: + return Open(None) + + def resume_value(self, value: object) -> Yield | Complete: + return next(self._steps) + + def resume_error(self, error: BaseException) -> Complete: + return Complete(None) + + def close(self) -> None: + self.closed = True + + +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", (False, True)) +async def test_native_stream_marker_reaches_caller_without_wrapping_or_consuming_the_stream( + asynchronous: bool, +) -> None: + chunks: Final = (b"event: message_start\n\n", b"event: message_stop\n\n") + execution: Final = ScriptedStreamExecution(chunks) + stream: Final[Stream | SyncStream] = Stream(execution) if asynchronous else SyncStream(execution) + bound: Final[bindings.NativeBinding[Callable[[], object]]] = bindings.NativeBinding( + "messages", validate=lambda _: None + ) + bound.override(lambda: stream) + + def python() -> object: + pytest.fail("native success must not fall back") + + async def anative(fn: Callable[[], object]) -> object: + return fn() + + async def apython() -> object: + return python() + + result: Final = ( + await runtime.arun(CONTEXT, binding=bound, native=anative, python=apython, rules=rules(Rollout.RUST_REQUIRED)) + if asynchronous + else runtime.run( + CONTEXT, binding=bound, native=lambda fn: fn(), python=python, rules=rules(Rollout.RUST_REQUIRED) + ) + ) + assert result is stream + assert get_hidden_params_dict(result) == {"additional_headers": {"x-litellm-rust": "true"}} + assert not execution.closed + delivered: Final = tuple([chunk async for chunk in result]) if isinstance(result, Stream) else tuple(result) + assert delivered == chunks + assert execution.closed + + def test_upstream_error_maps_to_api_error_without_fallback() -> None: calls: Final = recorder(RustUpstreamError(429, "rate limited")) diff --git a/tests/test_litellm_rust/messages/test_callbacks.py b/tests/test_litellm_rust/messages/test_callbacks.py index d5dc5437398..19043780eb6 100644 --- a/tests/test_litellm_rust/messages/test_callbacks.py +++ b/tests/test_litellm_rust/messages/test_callbacks.py @@ -5,6 +5,7 @@ import pytest import litellm from litellm.integrations.custom_logger import CustomLogger +from litellm.router_utils.add_retry_fallback_headers import get_hidden_params_dict from litellm.rust_bridge import catalog from litellm.rust_bridge.catalog import Route, RouteRule from litellm.rust_bridge.configuration import Rollout @@ -126,6 +127,7 @@ async def test_native_messages_stream_relays_provider_events_and_logs_success_on **arguments(messages_server, stream=True, callbacks=[recorder]) ) assert isinstance(stream, AsyncIterator) + assert get_hidden_params_dict(stream) == {"additional_headers": {"x-litellm-rust": "true"}} first: Final = await anext(stream) await drain_logging() assert "async_log_success_event" not in recorder.names @@ -169,6 +171,7 @@ def test_native_sync_messages_stream_relays_provider_events_and_logs_success_onc stream: Final = litellm.anthropic.messages.create(**arguments(messages_server, stream=True, callbacks=[recorder])) assert isinstance(stream, Iterator) + assert get_hidden_params_dict(stream) == {"additional_headers": {"x-litellm-rust": "true"}} assert b"".join(stream) == sse_payload() assert_served_natively(messages_server) From 7a9bfb4f1c7e251b25ee8920d586640a764c8a07 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 24 Sep 2026 09:44:33 -0700 Subject: [PATCH 052/120] feat(azure): add gpt-audio and gpt-realtime alias rows from the Azure model list (#42981) * feat(azure): add gpt-audio and gpt-realtime alias rows from the Azure model list Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(azure): keep existing catalog formatting untouched Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: kerry Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- ...odel_prices_and_context_window_backup.json | 136 ++++++++++++++++++ model_prices_and_context_window.json | 136 ++++++++++++++++++ 2 files changed, 272 insertions(+) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 47cdfecfc0f..e20cd12c9be 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -5603,6 +5603,39 @@ "supports_tool_choice": true, "supports_vision": true }, + "azure/gpt-audio": { + "deprecation_date": "2027-03-02", + "input_cost_per_audio_token": 4e-05, + "input_cost_per_token": 2.5e-06, + "litellm_provider": "azure", + "max_input_tokens": 128000, + "max_output_tokens": 16384, + "max_tokens": 16384, + "mode": "chat", + "output_cost_per_audio_token": 8e-05, + "output_cost_per_token": 1e-05, + "supported_endpoints": [ + "/v1/chat/completions" + ], + "supported_modalities": [ + "text", + "audio" + ], + "supported_output_modalities": [ + "text", + "audio" + ], + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_parallel_function_calling": true, + "supports_prompt_caching": false, + "supports_reasoning": false, + "supports_response_schema": false, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": false, + "source": "https://learn.microsoft.com/en-us/azure/foundry/foundry-models/concepts/models-sold-directly-by-azure" + }, "azure/gpt-audio-2025-08-28": { "deprecation_date": "2027-03-02", "input_cost_per_audio_token": 4e-05, @@ -5635,6 +5668,39 @@ "supports_tool_choice": true, "supports_vision": false }, + "azure/gpt-audio-1.5": { + "deprecation_date": "2027-08-24", + "input_cost_per_audio_token": 4e-05, + "input_cost_per_token": 2.5e-06, + "litellm_provider": "azure", + "max_input_tokens": 128000, + "max_output_tokens": 16384, + "max_tokens": 16384, + "mode": "chat", + "output_cost_per_audio_token": 8e-05, + "output_cost_per_token": 1e-05, + "supported_endpoints": [ + "/v1/chat/completions" + ], + "supported_modalities": [ + "text", + "audio" + ], + "supported_output_modalities": [ + "text", + "audio" + ], + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_parallel_function_calling": true, + "supports_prompt_caching": false, + "supports_reasoning": false, + "supports_response_schema": false, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": false, + "source": "https://learn.microsoft.com/en-us/azure/foundry/foundry-models/concepts/models-sold-directly-by-azure" + }, "azure/gpt-audio-1.5-2026-02-23": { "deprecation_date": "2027-08-24", "input_cost_per_audio_token": 4e-05, @@ -5849,6 +5915,41 @@ "supports_system_messages": true, "supports_tool_choice": true }, + "azure/gpt-realtime": { + "cache_creation_input_audio_token_cost": 4e-06, + "cache_read_input_audio_token_cost": 4e-07, + "cache_read_input_token_cost": 4e-07, + "deprecation_date": "2027-03-02", + "input_cost_per_audio_token": 3.2e-05, + "input_cost_per_image_token": 5e-06, + "input_cost_per_token": 4e-06, + "litellm_provider": "azure", + "max_input_tokens": 32000, + "max_output_tokens": 4096, + "max_tokens": 4096, + "mode": "realtime", + "output_cost_per_audio_token": 6.4e-05, + "output_cost_per_token": 1.6e-05, + "supported_endpoints": [ + "/v1/realtime" + ], + "supported_modalities": [ + "text", + "image", + "audio" + ], + "supported_output_modalities": [ + "text", + "audio" + ], + "supports_audio_input": true, + "supports_audio_output": true, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "source": "https://learn.microsoft.com/en-us/azure/foundry/foundry-models/concepts/models-sold-directly-by-azure" + }, "azure/gpt-realtime-2025-08-28": { "cache_creation_input_audio_token_cost": 4e-06, "cache_read_input_audio_token_cost": 4e-07, @@ -5883,6 +5984,41 @@ "supports_system_messages": true, "supports_tool_choice": true }, + "azure/gpt-realtime-1.5": { + "cache_creation_input_audio_token_cost": 4e-06, + "cache_read_input_audio_token_cost": 4e-07, + "cache_read_input_token_cost": 4e-07, + "deprecation_date": "2027-08-24", + "input_cost_per_audio_token": 3.2e-05, + "input_cost_per_image_token": 5e-06, + "input_cost_per_token": 4e-06, + "litellm_provider": "azure", + "max_input_tokens": 32000, + "max_output_tokens": 4096, + "max_tokens": 4096, + "mode": "realtime", + "output_cost_per_audio_token": 6.4e-05, + "output_cost_per_token": 1.6e-05, + "supported_endpoints": [ + "/v1/realtime" + ], + "supported_modalities": [ + "text", + "image", + "audio" + ], + "supported_output_modalities": [ + "text", + "audio" + ], + "supports_audio_input": true, + "supports_audio_output": true, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "source": "https://learn.microsoft.com/en-us/azure/foundry/foundry-models/concepts/models-sold-directly-by-azure" + }, "azure/gpt-realtime-1.5-2026-02-23": { "cache_creation_input_audio_token_cost": 4e-06, "cache_read_input_audio_token_cost": 4e-07, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 47cdfecfc0f..e20cd12c9be 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -5603,6 +5603,39 @@ "supports_tool_choice": true, "supports_vision": true }, + "azure/gpt-audio": { + "deprecation_date": "2027-03-02", + "input_cost_per_audio_token": 4e-05, + "input_cost_per_token": 2.5e-06, + "litellm_provider": "azure", + "max_input_tokens": 128000, + "max_output_tokens": 16384, + "max_tokens": 16384, + "mode": "chat", + "output_cost_per_audio_token": 8e-05, + "output_cost_per_token": 1e-05, + "supported_endpoints": [ + "/v1/chat/completions" + ], + "supported_modalities": [ + "text", + "audio" + ], + "supported_output_modalities": [ + "text", + "audio" + ], + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_parallel_function_calling": true, + "supports_prompt_caching": false, + "supports_reasoning": false, + "supports_response_schema": false, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": false, + "source": "https://learn.microsoft.com/en-us/azure/foundry/foundry-models/concepts/models-sold-directly-by-azure" + }, "azure/gpt-audio-2025-08-28": { "deprecation_date": "2027-03-02", "input_cost_per_audio_token": 4e-05, @@ -5635,6 +5668,39 @@ "supports_tool_choice": true, "supports_vision": false }, + "azure/gpt-audio-1.5": { + "deprecation_date": "2027-08-24", + "input_cost_per_audio_token": 4e-05, + "input_cost_per_token": 2.5e-06, + "litellm_provider": "azure", + "max_input_tokens": 128000, + "max_output_tokens": 16384, + "max_tokens": 16384, + "mode": "chat", + "output_cost_per_audio_token": 8e-05, + "output_cost_per_token": 1e-05, + "supported_endpoints": [ + "/v1/chat/completions" + ], + "supported_modalities": [ + "text", + "audio" + ], + "supported_output_modalities": [ + "text", + "audio" + ], + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_parallel_function_calling": true, + "supports_prompt_caching": false, + "supports_reasoning": false, + "supports_response_schema": false, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": false, + "source": "https://learn.microsoft.com/en-us/azure/foundry/foundry-models/concepts/models-sold-directly-by-azure" + }, "azure/gpt-audio-1.5-2026-02-23": { "deprecation_date": "2027-08-24", "input_cost_per_audio_token": 4e-05, @@ -5849,6 +5915,41 @@ "supports_system_messages": true, "supports_tool_choice": true }, + "azure/gpt-realtime": { + "cache_creation_input_audio_token_cost": 4e-06, + "cache_read_input_audio_token_cost": 4e-07, + "cache_read_input_token_cost": 4e-07, + "deprecation_date": "2027-03-02", + "input_cost_per_audio_token": 3.2e-05, + "input_cost_per_image_token": 5e-06, + "input_cost_per_token": 4e-06, + "litellm_provider": "azure", + "max_input_tokens": 32000, + "max_output_tokens": 4096, + "max_tokens": 4096, + "mode": "realtime", + "output_cost_per_audio_token": 6.4e-05, + "output_cost_per_token": 1.6e-05, + "supported_endpoints": [ + "/v1/realtime" + ], + "supported_modalities": [ + "text", + "image", + "audio" + ], + "supported_output_modalities": [ + "text", + "audio" + ], + "supports_audio_input": true, + "supports_audio_output": true, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "source": "https://learn.microsoft.com/en-us/azure/foundry/foundry-models/concepts/models-sold-directly-by-azure" + }, "azure/gpt-realtime-2025-08-28": { "cache_creation_input_audio_token_cost": 4e-06, "cache_read_input_audio_token_cost": 4e-07, @@ -5883,6 +5984,41 @@ "supports_system_messages": true, "supports_tool_choice": true }, + "azure/gpt-realtime-1.5": { + "cache_creation_input_audio_token_cost": 4e-06, + "cache_read_input_audio_token_cost": 4e-07, + "cache_read_input_token_cost": 4e-07, + "deprecation_date": "2027-08-24", + "input_cost_per_audio_token": 3.2e-05, + "input_cost_per_image_token": 5e-06, + "input_cost_per_token": 4e-06, + "litellm_provider": "azure", + "max_input_tokens": 32000, + "max_output_tokens": 4096, + "max_tokens": 4096, + "mode": "realtime", + "output_cost_per_audio_token": 6.4e-05, + "output_cost_per_token": 1.6e-05, + "supported_endpoints": [ + "/v1/realtime" + ], + "supported_modalities": [ + "text", + "image", + "audio" + ], + "supported_output_modalities": [ + "text", + "audio" + ], + "supports_audio_input": true, + "supports_audio_output": true, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "source": "https://learn.microsoft.com/en-us/azure/foundry/foundry-models/concepts/models-sold-directly-by-azure" + }, "azure/gpt-realtime-1.5-2026-02-23": { "cache_creation_input_audio_token_cost": 4e-06, "cache_read_input_audio_token_cost": 4e-07, From 1ceb8fb08e778506816766c001d6fc9121c3d46e Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 24 Sep 2026 09:45:09 -0700 Subject: [PATCH 053/120] fix(cost-map): align openrouter deepseek-v4-pro cache hit price with cache read (#42985) Co-authored-by: kerry Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/model_prices_and_context_window_backup.json | 2 +- model_prices_and_context_window.json | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index e20cd12c9be..65f90014151 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -41382,7 +41382,7 @@ }, "openrouter/deepseek/deepseek-v4-pro": { "input_cost_per_token": 9.1263e-07, - "input_cost_per_token_cache_hit": 4.4e-08, + "input_cost_per_token_cache_hit": 7.60525e-08, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 384000, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index e20cd12c9be..65f90014151 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -41382,7 +41382,7 @@ }, "openrouter/deepseek/deepseek-v4-pro": { "input_cost_per_token": 9.1263e-07, - "input_cost_per_token_cache_hit": 4.4e-08, + "input_cost_per_token_cache_hit": 7.60525e-08, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 384000, From de314f8271788593087c3f6065b7d82e677a7a31 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 24 Sep 2026 09:57:05 -0700 Subject: [PATCH 054/120] fix(cost-map): take the later azure deprecation date for gpt-4.1-nano, gpt-4o-transcribe and gpt-realtime-2.1 (#42986) Co-authored-by: kerry Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../model_prices_and_context_window_backup.json | 14 +++++++------- model_prices_and_context_window.json | 14 +++++++------- 2 files changed, 14 insertions(+), 14 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 65f90014151..8fa9740ee84 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -5442,7 +5442,7 @@ "supports_web_search": false }, "azure/gpt-4.1-nano": { - "deprecation_date": "2026-10-14", + "deprecation_date": "2027-04-14", "cache_read_input_token_cost": 2.5e-08, "input_cost_per_token": 1e-07, "input_cost_per_token_batches": 5e-08, @@ -5476,7 +5476,7 @@ "supports_vision": true }, "azure/gpt-4.1-nano-2025-04-14": { - "deprecation_date": "2026-10-14", + "deprecation_date": "2027-04-14", "cache_read_input_token_cost": 2.5e-08, "input_cost_per_token": 1e-07, "input_cost_per_token_batches": 5e-08, @@ -6057,7 +6057,7 @@ "cache_creation_input_audio_token_cost": 4e-07, "cache_read_input_audio_token_cost": 4e-07, "cache_read_input_token_cost": 4e-07, - "deprecation_date": "2027-06-25", + "deprecation_date": "2027-07-31", "input_cost_per_audio_token": 3.2e-05, "input_cost_per_image_token": 5e-06, "input_cost_per_token": 4e-06, @@ -6269,7 +6269,7 @@ "supports_tool_choice": true }, "azure/gpt-4o-transcribe": { - "deprecation_date": "2026-10-15", + "deprecation_date": "2026-12-31", "input_cost_per_audio_token": 2.5e-06, "input_cost_per_token": 2.5e-06, "litellm_provider": "azure", @@ -10796,7 +10796,7 @@ "supports_web_search": false }, "azure/us/gpt-4.1-nano-2025-04-14": { - "deprecation_date": "2026-10-14", + "deprecation_date": "2027-04-14", "cache_read_input_token_cost": 2.8e-08, "input_cost_per_token": 1.1e-07, "input_cost_per_token_batches": 5.5e-08, @@ -69100,7 +69100,7 @@ "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" }, "azure/eu/gpt-4.1-nano": { - "deprecation_date": "2026-10-14", + "deprecation_date": "2027-04-14", "cache_read_input_token_cost": 2.8e-08, "input_cost_per_token": 1.1e-07, "input_cost_per_token_batches": 5.5e-08, @@ -69535,7 +69535,7 @@ "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" }, "azure/us/gpt-4.1-nano": { - "deprecation_date": "2026-10-14", + "deprecation_date": "2027-04-14", "cache_read_input_token_cost": 2.8e-08, "input_cost_per_token": 1.1e-07, "input_cost_per_token_batches": 5.5e-08, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 65f90014151..8fa9740ee84 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -5442,7 +5442,7 @@ "supports_web_search": false }, "azure/gpt-4.1-nano": { - "deprecation_date": "2026-10-14", + "deprecation_date": "2027-04-14", "cache_read_input_token_cost": 2.5e-08, "input_cost_per_token": 1e-07, "input_cost_per_token_batches": 5e-08, @@ -5476,7 +5476,7 @@ "supports_vision": true }, "azure/gpt-4.1-nano-2025-04-14": { - "deprecation_date": "2026-10-14", + "deprecation_date": "2027-04-14", "cache_read_input_token_cost": 2.5e-08, "input_cost_per_token": 1e-07, "input_cost_per_token_batches": 5e-08, @@ -6057,7 +6057,7 @@ "cache_creation_input_audio_token_cost": 4e-07, "cache_read_input_audio_token_cost": 4e-07, "cache_read_input_token_cost": 4e-07, - "deprecation_date": "2027-06-25", + "deprecation_date": "2027-07-31", "input_cost_per_audio_token": 3.2e-05, "input_cost_per_image_token": 5e-06, "input_cost_per_token": 4e-06, @@ -6269,7 +6269,7 @@ "supports_tool_choice": true }, "azure/gpt-4o-transcribe": { - "deprecation_date": "2026-10-15", + "deprecation_date": "2026-12-31", "input_cost_per_audio_token": 2.5e-06, "input_cost_per_token": 2.5e-06, "litellm_provider": "azure", @@ -10796,7 +10796,7 @@ "supports_web_search": false }, "azure/us/gpt-4.1-nano-2025-04-14": { - "deprecation_date": "2026-10-14", + "deprecation_date": "2027-04-14", "cache_read_input_token_cost": 2.8e-08, "input_cost_per_token": 1.1e-07, "input_cost_per_token_batches": 5.5e-08, @@ -69100,7 +69100,7 @@ "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" }, "azure/eu/gpt-4.1-nano": { - "deprecation_date": "2026-10-14", + "deprecation_date": "2027-04-14", "cache_read_input_token_cost": 2.8e-08, "input_cost_per_token": 1.1e-07, "input_cost_per_token_batches": 5.5e-08, @@ -69535,7 +69535,7 @@ "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" }, "azure/us/gpt-4.1-nano": { - "deprecation_date": "2026-10-14", + "deprecation_date": "2027-04-14", "cache_read_input_token_cost": 2.8e-08, "input_cost_per_token": 1.1e-07, "input_cost_per_token_batches": 5.5e-08, From 793627ce83e875209d5705ce544ba52d0371bd62 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 24 Sep 2026 09:57:52 -0700 Subject: [PATCH 055/120] fix(cost-map): add supports_reasoning to azure/eu/gpt-6-astra (#42989) Co-authored-by: kerry Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/model_prices_and_context_window_backup.json | 3 ++- model_prices_and_context_window.json | 3 ++- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 8fa9740ee84..8088fa28b52 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -69288,7 +69288,8 @@ "mode": "chat", "output_cost_per_token": 5.5e-05, "output_cost_per_token_above_272k_tokens": 8.25e-05, - "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" + "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", + "supports_reasoning": true }, "azure/eu/gpt-6-luna": { "deprecation_date": "2028-03-11", diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 8fa9740ee84..8088fa28b52 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -69288,7 +69288,8 @@ "mode": "chat", "output_cost_per_token": 5.5e-05, "output_cost_per_token_above_272k_tokens": 8.25e-05, - "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" + "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", + "supports_reasoning": true }, "azure/eu/gpt-6-luna": { "deprecation_date": "2028-03-11", From 57eb3ff6b66acb87f9c8776292bb10fcbafe1c24 Mon Sep 17 00:00:00 2001 From: joshua-berri Date: Thu, 24 Sep 2026 17:23:30 +0000 Subject: [PATCH 056/120] fix(mcp): reject origins outside the configured allowlist (#42649) Co-authored-by: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> --- .../proxy/_experimental/mcp_server/server.py | 9 +++ .../mcp_server/test_mcp_server.py | 58 +++++++++++++++++++ 2 files changed, 67 insertions(+) diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index fe508fc22c7..6261f36983d 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -111,6 +111,13 @@ _MCP_DESTINATIONS_SCOPE_KEY: Final = "litellm_otel_request_destinations" _MCP_PROTOCOL_VERSION_HEADER: Final = b"mcp-protocol-version" +def reject_disallowed_mcp_origin(request: StarletteRequest) -> None: + from litellm.proxy.proxy_server import origins # noqa: PLC0415 # proxy imports this module during startup + + if "*" not in origins and any(origin not in origins for origin in request.headers.getlist("origin")): + raise HTTPException(status_code=403, detail="Invalid Origin header") + + def unsupported_protocol_version(scope: Scope) -> str | None: """Return the unsupported ``MCP-Protocol-Version`` header value, if any. @@ -1931,6 +1938,7 @@ if MCP_AVAILABLE: async def handle_streamable_http_mcp(scope: Scope, receive: Receive, send: Send) -> None: """Handle MCP requests through StreamableHTTP.""" try: + reject_disallowed_mcp_origin(StarletteRequest(scope)) bad_version: Final = unsupported_protocol_version(scope) if bad_version is not None: supported: Final = ", ".join(sorted(HANDSHAKE_PROTOCOL_VERSIONS)) @@ -2275,6 +2283,7 @@ if MCP_AVAILABLE: async def handle_sse_mcp(scope: Scope, receive: Receive, send: Send) -> None: """Handle MCP requests through SSE.""" try: + reject_disallowed_mcp_origin(StarletteRequest(scope)) bad_version: Final = unsupported_protocol_version(scope) if bad_version is not None: supported: Final = ", ".join(sorted(HANDSHAKE_PROTOCOL_VERSIONS)) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py index ea82c9b069d..6887adf8283 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py @@ -10201,6 +10201,64 @@ async def test_active_request_ctx_var_feeds_get_current_session(_mcp_request_ctx assert _get_current_session() is None +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("method", "path", "session_headers"), + ( + ("POST", "/mcp", ()), + ("GET", "/mcp", (("mcp-session-id", "existing-session"),)), + ("DELETE", "/mcp", (("mcp-session-id", "existing-session"),)), + ("POST", "/server/mcp", ()), + ("GET", "/sse", ()), + ("POST", "/sse/messages", ()), + ), +) +@pytest.mark.parametrize( + ("allowed_origins", "origin_headers", "expected_status"), + ( + (("https://allowed.example",), (("origin", "https://evil.example"),), 403), + (("https://allowed.example",), (("origin", "https://allowed.example.evil.example"),), 403), + (("https://allowed.example",), (("origin", "null"),), 403), + (("https://allowed.example",), (("origin", ""),), 403), + ( + ("https://allowed.example",), + (("origin", "https://allowed.example"), ("origin", "https://evil.example")), + 403, + ), + (("https://allowed.example",), (("origin", "https://allowed.example"),), 401), + (("https://allowed.example",), (), 401), + (("*",), (("origin", "https://another.example"),), 401), + ), +) +async def test_mcp_origin_admission_precedes_authentication( + method: str, + path: str, + session_headers: tuple[tuple[str, str], ...], + allowed_origins: tuple[str, ...], + origin_headers: tuple[tuple[str, str], ...], + expected_status: int, +) -> None: + import httpx + + from litellm.proxy._experimental.mcp_server import server + + authenticate: Final = AsyncMock(side_effect=HTTPException(status_code=401, detail="authentication required")) + with ( + patch("litellm.proxy.proxy_server.origins", allowed_origins), + patch.object(server, "extract_mcp_auth_context", authenticate), + ): + async with httpx.AsyncClient(transport=httpx.ASGITransport(app=server.app), base_url="http://gateway") as client: + response: Final = await client.request(method, path, headers=(*session_headers, *origin_headers)) + + assert response.status_code == expected_status + if expected_status == 403: + assert response.json() == {"detail": "Invalid Origin header"} + authenticate.assert_not_awaited() + else: + assert response.json() == {"detail": "authentication required"} + authenticate.assert_awaited_once() + + @pytest.mark.asyncio async def test_active_request_ctx_var_feeds_auth_resolution_recording(_mcp_request_ctx) -> None: from starlette.requests import Request From 0fa1fe2059513f0bb387959894057d4a12a56657 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 24 Sep 2026 10:37:59 -0700 Subject: [PATCH 057/120] fix(cost-map): drop stale cache_hit field from openrouter deepseek-v4-pro (#42994) Co-authored-by: kerry Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/model_prices_and_context_window_backup.json | 1 - model_prices_and_context_window.json | 1 - 2 files changed, 2 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 8088fa28b52..a338c389ea4 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -41382,7 +41382,6 @@ }, "openrouter/deepseek/deepseek-v4-pro": { "input_cost_per_token": 9.1263e-07, - "input_cost_per_token_cache_hit": 7.60525e-08, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 384000, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 8088fa28b52..a338c389ea4 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -41382,7 +41382,6 @@ }, "openrouter/deepseek/deepseek-v4-pro": { "input_cost_per_token": 9.1263e-07, - "input_cost_per_token_cache_hit": 7.60525e-08, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 384000, From d3462c65b585eb7e51594330a5785669f6b7db51 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 24 Sep 2026 11:08:09 -0700 Subject: [PATCH 058/120] fix(cost-map): add vertex cache read, batch and above 200k prices for gemini image preview rows (#42995) * fix(cost-map): add vertex cache read prices for gemini image preview rows Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(cost-map): add batch cache read and above 200k tiers to vertex gemini image preview rows Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: kerry Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- ...model_prices_and_context_window_backup.json | 18 ++++++++++++++++++ model_prices_and_context_window.json | 18 ++++++++++++++++++ 2 files changed, 36 insertions(+) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index a338c389ea4..2600da22a77 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -26312,7 +26312,11 @@ }, "gemini-3-pro-image-preview": { "input_cost_per_image": 0.0011, + "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "cache_read_input_token_cost_batches": 1e-07, "input_cost_per_token": 2e-06, + "input_cost_per_token_above_200k_tokens": 4e-06, "input_cost_per_token_batches": 1e-06, "litellm_provider": "vertex_ai-language-models", "max_input_tokens": 65536, @@ -26322,6 +26326,7 @@ "output_cost_per_image": 0.134, "output_cost_per_image_token": 0.00012, "output_cost_per_token": 1.2e-05, + "output_cost_per_token_above_200k_tokens": 1.8e-05, "output_cost_per_token_batches": 6e-06, "source": "https://ai.google.dev/gemini-api/docs/pricing", "supported_endpoints": [ @@ -26398,7 +26403,10 @@ }, "gemini-3.1-flash-image-preview": { "input_cost_per_image": 0.00056, + "cache_read_input_token_cost": 5e-08, + "cache_read_input_token_cost_batches": 2.5e-08, "input_cost_per_token": 5e-07, + "input_cost_per_token_batches": 2.5e-07, "litellm_provider": "vertex_ai-language-models", "max_input_tokens": 65536, "max_output_tokens": 32768, @@ -26407,6 +26415,7 @@ "output_cost_per_image": 0.0672, "output_cost_per_image_token": 6e-05, "output_cost_per_token": 3e-06, + "output_cost_per_token_batches": 1.5e-06, "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing#gemini-models", "supported_endpoints": [ "/v1/chat/completions", @@ -49484,7 +49493,11 @@ }, "vertex_ai/gemini-3-pro-image-preview": { "input_cost_per_image": 0.0011, + "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "cache_read_input_token_cost_batches": 1e-07, "input_cost_per_token": 2e-06, + "input_cost_per_token_above_200k_tokens": 4e-06, "input_cost_per_token_batches": 1e-06, "litellm_provider": "vertex_ai-language-models", "max_input_tokens": 65536, @@ -49494,6 +49507,7 @@ "output_cost_per_image": 0.134, "output_cost_per_image_token": 0.00012, "output_cost_per_token": 1.2e-05, + "output_cost_per_token_above_200k_tokens": 1.8e-05, "output_cost_per_token_batches": 6e-06, "supports_reasoning": false, "source": "https://docs.cloud.google.com/vertex-ai/generative-ai/docs/models/gemini/3-pro-image" @@ -49522,7 +49536,10 @@ }, "vertex_ai/gemini-3.1-flash-image-preview": { "input_cost_per_image": 0.00056, + "cache_read_input_token_cost": 5e-08, + "cache_read_input_token_cost_batches": 2.5e-08, "input_cost_per_token": 5e-07, + "input_cost_per_token_batches": 2.5e-07, "litellm_provider": "vertex_ai-language-models", "max_input_tokens": 65536, "max_output_tokens": 32768, @@ -49531,6 +49548,7 @@ "output_cost_per_image": 0.0672, "output_cost_per_image_token": 6e-05, "output_cost_per_token": 3e-06, + "output_cost_per_token_batches": 1.5e-06, "supports_reasoning": false, "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing#gemini-models" }, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index a338c389ea4..2600da22a77 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -26312,7 +26312,11 @@ }, "gemini-3-pro-image-preview": { "input_cost_per_image": 0.0011, + "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "cache_read_input_token_cost_batches": 1e-07, "input_cost_per_token": 2e-06, + "input_cost_per_token_above_200k_tokens": 4e-06, "input_cost_per_token_batches": 1e-06, "litellm_provider": "vertex_ai-language-models", "max_input_tokens": 65536, @@ -26322,6 +26326,7 @@ "output_cost_per_image": 0.134, "output_cost_per_image_token": 0.00012, "output_cost_per_token": 1.2e-05, + "output_cost_per_token_above_200k_tokens": 1.8e-05, "output_cost_per_token_batches": 6e-06, "source": "https://ai.google.dev/gemini-api/docs/pricing", "supported_endpoints": [ @@ -26398,7 +26403,10 @@ }, "gemini-3.1-flash-image-preview": { "input_cost_per_image": 0.00056, + "cache_read_input_token_cost": 5e-08, + "cache_read_input_token_cost_batches": 2.5e-08, "input_cost_per_token": 5e-07, + "input_cost_per_token_batches": 2.5e-07, "litellm_provider": "vertex_ai-language-models", "max_input_tokens": 65536, "max_output_tokens": 32768, @@ -26407,6 +26415,7 @@ "output_cost_per_image": 0.0672, "output_cost_per_image_token": 6e-05, "output_cost_per_token": 3e-06, + "output_cost_per_token_batches": 1.5e-06, "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing#gemini-models", "supported_endpoints": [ "/v1/chat/completions", @@ -49484,7 +49493,11 @@ }, "vertex_ai/gemini-3-pro-image-preview": { "input_cost_per_image": 0.0011, + "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "cache_read_input_token_cost_batches": 1e-07, "input_cost_per_token": 2e-06, + "input_cost_per_token_above_200k_tokens": 4e-06, "input_cost_per_token_batches": 1e-06, "litellm_provider": "vertex_ai-language-models", "max_input_tokens": 65536, @@ -49494,6 +49507,7 @@ "output_cost_per_image": 0.134, "output_cost_per_image_token": 0.00012, "output_cost_per_token": 1.2e-05, + "output_cost_per_token_above_200k_tokens": 1.8e-05, "output_cost_per_token_batches": 6e-06, "supports_reasoning": false, "source": "https://docs.cloud.google.com/vertex-ai/generative-ai/docs/models/gemini/3-pro-image" @@ -49522,7 +49536,10 @@ }, "vertex_ai/gemini-3.1-flash-image-preview": { "input_cost_per_image": 0.00056, + "cache_read_input_token_cost": 5e-08, + "cache_read_input_token_cost_batches": 2.5e-08, "input_cost_per_token": 5e-07, + "input_cost_per_token_batches": 2.5e-07, "litellm_provider": "vertex_ai-language-models", "max_input_tokens": 65536, "max_output_tokens": 32768, @@ -49531,6 +49548,7 @@ "output_cost_per_image": 0.0672, "output_cost_per_image_token": 6e-05, "output_cost_per_token": 3e-06, + "output_cost_per_token_batches": 1.5e-06, "supports_reasoning": false, "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing#gemini-models" }, From 0ac435bd0195ec224c8399f0655ace3fd0d3c68d Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Thu, 24 Sep 2026 11:54:26 -0700 Subject: [PATCH 059/120] test(ollama): assert the images sent to Ollama instead of echoing them through response (#42905) test_ollama_image returned the request's images list as the mocked `response` field and read it back from message content. Since #42838 the completion transform validates `response` as a string, so the list is dropped and the test fails with "string index out of range". Ollama always sends a string there. The mock now returns a real string reply and the test asserts on the images the transform actually sent, which is what it was checking all along. --- tests/local_testing/test_completion.py | 17 ++++++++--------- 1 file changed, 8 insertions(+), 9 deletions(-) diff --git a/tests/local_testing/test_completion.py b/tests/local_testing/test_completion.py index 25c6c50251d..c6dd78c73b4 100644 --- a/tests/local_testing/test_completion.py +++ b/tests/local_testing/test_completion.py @@ -1361,16 +1361,14 @@ def test_ollama_image(): from PIL import Image + sent_images = [] + def mock_post(url, **kwargs): + sent_images.append(json.loads(kwargs["data"])["images"]) mock_response = MagicMock() mock_response.status_code = 200 mock_response.headers = {"Content-Type": "application/json"} - data_json = json.loads(kwargs["data"]) - mock_response.json.return_value = { - # return the image in the response so that it can be tested - # against the original - "response": data_json["images"] - } + mock_response.json.return_value = {"response": "a black pixel"} return mock_response def make_b64image(format): @@ -1399,9 +1397,10 @@ def test_ollama_image(): client = HTTPHandler() for test in tests: + sent_images.clear() try: with patch.object(client, "post", side_effect=mock_post): - response = completion( + completion( model="ollama/llava", messages=[ { @@ -1417,14 +1416,14 @@ def test_ollama_image(): ], client=client, ) + (image_data,) = sent_images[0] if not test[1]: # the conversion process may not always generate the same image, # so just check for a JPEG image when a conversion was done. - image_data = response["choices"][0]["message"]["content"][0] image = Image.open(io.BytesIO(base64.b64decode(image_data))) assert image.format == "JPEG" else: - assert response["choices"][0]["message"]["content"][0] == test[1] + assert image_data == test[1] except Exception as e: pytest.fail(f"Error occurred: {e}") From 12960f3edfb1b3e26cd50fe567f242d9df8c50e4 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 24 Sep 2026 11:54:52 -0700 Subject: [PATCH 060/120] fix(ui): pass is_proxy_admin for proxy admins on the models page team drill-in (#43003) * fix(ui): pass is_proxy_admin for proxy admins on the models page team drill-in Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(ui): drop explanatory comments from the models page team drill-in tests Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(ui): exclude view-only admins from is_proxy_admin on the models page team drill-in Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(ui): add browser integration contract for the team guardrail kill switch on the models page Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: yucheng Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../tests/integrationCritical/expected.json | 1 + .../teamGlobalGuardrailKillSwitch.spec.ts | 77 ++++++ .../page.integration.test.tsx | 243 ++++++++++++++++++ .../models-and-endpoints/page.test.tsx | 23 +- .../(dashboard)/models-and-endpoints/page.tsx | 2 +- 5 files changed, 340 insertions(+), 6 deletions(-) create mode 100644 tests/e2e/ui/tests/integrationCritical/teamGlobalGuardrailKillSwitch.spec.ts create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/page.integration.test.tsx diff --git a/tests/e2e/ui/tests/integrationCritical/expected.json b/tests/e2e/ui/tests/integrationCritical/expected.json index 0354c9f729c..189cdef9a93 100644 --- a/tests/e2e/ui/tests/integrationCritical/expected.json +++ b/tests/e2e/ui/tests/integrationCritical/expected.json @@ -1,5 +1,6 @@ [ "tests/e2e/ui/tests/integrationCritical/projectDetachment.spec.ts::project creation and explicit detachment preserve saved scope and restore serving", + "tests/e2e/ui/tests/integrationCritical/teamGlobalGuardrailKillSwitch.spec.ts::proxy admin can enable the global guardrail kill switch from the models page team drill-in", "tests/e2e/ui/tests/integrationCritical/mcpUserEnvVars.spec.ts::per-user MCP env var stays updatable and clearable from the card after it is set", "tests/e2e/ui/tests/integrationCritical/mcpUserEnvVars.spec.ts::cancelling the clear confirmation keeps the stored value and sends no delete", "tests/e2e/ui/tests/integrationCritical/mcpUserEnvVars.spec.ts::pressing Enter on Update opens the credentials modal instead of the server editor", diff --git a/tests/e2e/ui/tests/integrationCritical/teamGlobalGuardrailKillSwitch.spec.ts b/tests/e2e/ui/tests/integrationCritical/teamGlobalGuardrailKillSwitch.spec.ts new file mode 100644 index 00000000000..ac48dd1a770 --- /dev/null +++ b/tests/e2e/ui/tests/integrationCritical/teamGlobalGuardrailKillSwitch.spec.ts @@ -0,0 +1,77 @@ +import { + test, + expect, + APIRequestContext, + Page as PlaywrightPage, +} from "@playwright/test"; +import { randomUUID } from "node:crypto"; + +const master = process.env.LITELLM_MASTER_KEY ?? "sk-integration-master"; +const headers = { Authorization: `Bearer ${master}` }; + +async function createTeam(request: APIRequestContext): Promise { + const created = await request.post("/team/new", { + headers, + data: { + team_alias: `int_kill_switch_${randomUUID().replace(/-/g, "").slice(0, 12)}`, + }, + }); + expect(created.ok(), await created.text()).toBe(true); + return (await created.json()).team_id as string; +} + +async function loginAsAdmin(page: PlaywrightPage): Promise { + await page.goto("/ui/login"); + await page.getByPlaceholder("Enter your username").fill("admin"); + await page.getByPlaceholder("Enter your password").fill(master); + await page.getByRole("button", { name: "Login", exact: true }).click(); + await expect(page).toHaveURL( + (url) => url.pathname.startsWith("/ui") && !url.pathname.includes("login"), + ); +} + +test("proxy admin can enable the global guardrail kill switch from the models page team drill-in", async ({ + page, + request, +}) => { + const teamId = await createTeam(request); + try { + await loginAsAdmin(page); + await page.goto(`/ui/models-and-endpoints?team=${teamId}`); + await page.getByRole("tab", { name: "Settings" }).click(); + await page.getByRole("button", { name: /edit settings/i }).click(); + await expect(page.getByLabel(/Team Name/)).toBeVisible(); + const killSwitch = page.getByRole("switch", { + name: /disable all global guardrails/i, + }); + await expect(killSwitch).toBeVisible(); + await expect(killSwitch).not.toBeChecked(); + await killSwitch.click(); + await page.getByRole("button", { name: "Save Changes" }).click(); + await expect + .poll(async () => { + const response = await request.get(`/team/info?team_id=${teamId}`, { + headers, + }); + expect(response.ok(), await response.text()).toBe(true); + const json = await response.json(); + return json.team_info?.metadata?.disable_global_guardrails; + }) + .toBe(true); + await page.reload(); + await page.getByRole("tab", { name: "Settings" }).click(); + await page.getByRole("button", { name: /edit settings/i }).click(); + await expect(page.getByLabel(/Team Name/)).toBeVisible(); + await expect( + page.getByRole("switch", { name: /disable all global guardrails/i }), + ).toBeChecked(); + } finally { + const removed = await request.post("/team/delete", { + headers, + data: { team_ids: [teamId] }, + }); + expect(removed.ok() || removed.status() === 404, await removed.text()).toBe( + true, + ); + } +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/page.integration.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/page.integration.test.tsx new file mode 100644 index 00000000000..865cc808627 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/page.integration.test.tsx @@ -0,0 +1,243 @@ +/* @vitest-environment jsdom */ +import { screen } from "@testing-library/react"; +import userEvent from "@testing-library/user-event"; +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; +import * as networking from "@/components/networking"; +import { renderWithProviders } from "../../../../tests/test-utils"; +import ModelsAndEndpointsPage from "./page"; + +vi.mock("./panels/AllModelsPanel", () => ({ default: () =>
})); +vi.mock("./panels/AddModelPanel", () => ({ default: () =>
})); +vi.mock("./panels/AutoRoutersTabPanel", () => ({ default: () =>
})); +vi.mock("./panels/LlmCredentialsPanel", () => ({ default: () =>
})); +vi.mock("./panels/PassThroughPanel", () => ({ default: () =>
})); +vi.mock("./panels/HealthStatusPanel", () => ({ default: () =>
})); +vi.mock("./panels/ModelRetrySettingsPanel", () => ({ default: () =>
})); +vi.mock("./panels/ModelGroupAliasPanel", () => ({ default: () =>
})); +vi.mock("./panels/PriceDataPanel", () => ({ default: () =>
})); +vi.mock("./panels/AccessGroupBudgetsPanel", () => ({ default: () =>
})); +vi.mock("@/components/molecules/cost_optimization_feedback_banner", () => ({ default: () => null })); +vi.mock("@/components/model_info_view", () => ({ + default: ({ modelId }: { modelId: string }) =>
model:{modelId}
, +})); +vi.mock("./useModelDashboardData", () => ({ + useModelDashboardData: () => ({ availableModelAccessGroups: [], allModelsOnProxy: [], availableModelGroups: [] }), +})); + +const authState = vi.hoisted(() => ({ userRole: "Admin", isViewOnly: false })); +vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({ + default: () => ({ + token: "123", + accessToken: "123", + userId: "user-1", + userEmail: "admin@example.com", + userRole: authState.userRole, + premiumUser: true, + isViewOnly: authState.isViewOnly, + disabledPersonalKeyCreation: null, + showSSOBanner: false, + }), +})); + +vi.mock("next/navigation", () => ({ useRouter: () => ({ push: vi.fn() }) })); + +vi.mock("@/components/networking", () => ({ + serverRootPath: "", + teamInfoCall: vi.fn(), + teamMemberDeleteCall: vi.fn(), + teamMemberAddCall: vi.fn(), + teamMemberUpdateCall: vi.fn(), + teamUpdateCall: vi.fn(), + getGuardrailsList: vi.fn(), + getPoliciesList: vi.fn(), + getPolicyInfoWithGuardrails: vi.fn(), + fetchMCPAccessGroups: vi.fn(), + getTeamPermissionsCall: vi.fn(), + organizationInfoCall: vi.fn(), + getRouterSettingsCall: vi.fn().mockResolvedValue({ fields: [] }), + getPassThroughEndpointsCall: vi.fn().mockResolvedValue({ endpoints: [] }), + fetchMCPServers: vi.fn().mockResolvedValue([]), + fetchMCPToolsets: vi.fn().mockResolvedValue([]), + listMCPTools: vi.fn().mockResolvedValue({ tools: [] }), + vectorStoreListCall: vi.fn().mockResolvedValue({ data: [] }), + getAgentsList: vi.fn().mockResolvedValue({ agents: [] }), + getClaudeCodePluginsList: vi.fn().mockResolvedValue({ plugins: [], count: 0 }), +})); + +const can = vi.fn(); +vi.mock("@/app/(dashboard)/hooks/useCan", () => ({ + default: (...args: unknown[]) => can(...args), +})); + +vi.mock("@/components/utils/dataUtils", () => ({ + copyToClipboard: vi.fn().mockResolvedValue(true), + formatNumberWithCommas: vi.fn((value: number) => value.toLocaleString()), +})); + +vi.mock("@/app/(dashboard)/hooks/teams/useTeamMetadataSchema", () => ({ + useTeamMetadataSchema: vi.fn(() => ({ data: [], isLoading: false })), +})); + +vi.mock("@/app/(dashboard)/hooks/uiSettings/useUISettings", () => ({ + useUISettings: vi.fn(() => ({ data: { values: {} }, isLoading: false })), +})); + +vi.mock("@/app/(dashboard)/hooks/models/useModels", () => ({ + useAllProxyModels: vi.fn(() => ({ data: { data: [] }, isLoading: false })), +})); + +vi.mock("@/app/(dashboard)/hooks/teams/useTeams", () => ({ + useTeams: vi.fn(() => ({ data: [], isLoading: false })), + useTeam: vi.fn(() => ({ data: undefined, isLoading: false })), +})); + +vi.mock("@/app/(dashboard)/hooks/organizations/useOrganizations", () => ({ + organizationKeys: { all: ["organizations"] }, + useOrganization: vi.fn(() => ({ data: undefined, isLoading: false })), + useOrganizations: vi.fn().mockReturnValue({ data: [], isLoading: false }), +})); + +vi.mock("@/app/(dashboard)/hooks/users/useCurrentUser", () => ({ + useCurrentUser: vi.fn(() => ({ data: { models: [] }, isLoading: false })), +})); + +vi.mock("@/app/(dashboard)/hooks/mcpServers/useMCPServers", () => ({ + useMCPServers: vi.fn(() => ({ data: [], isLoading: false, isError: false })), +})); + +vi.mock("@/app/(dashboard)/hooks/mcpServers/useMCPToolsets", () => ({ + useMCPToolsets: vi.fn(() => ({ data: [], isLoading: false, isError: false })), +})); + +vi.mock("@/components/mcp_server_management/MCPServerSelector", () => ({ + default: () =>
mcp server selector
, +})); + +vi.mock("@/components/team/TeamMemberTab", () => ({ + default: vi.fn(() =>
member tab
), +})); + +vi.mock("@/components/common_components/user_search_modal", () => ({ + default: vi.fn(() => null), +})); + +vi.mock("@/components/team/EditMembership", () => ({ + default: vi.fn(() => null), +})); + +vi.mock("@/components/common_components/DeleteResourceModal", () => ({ + default: vi.fn(() => null), +})); + +vi.mock("@/components/team/member_permissions", () => ({ + default: vi.fn(() =>
Member Permissions
), +})); + +vi.mock("@/components/common_components/ModelAliasManager", () => ({ + default: vi.fn(() =>
alias manager
), +})); + +vi.mock("@/app/(dashboard)/hooks/accessGroups/useAccessGroups", () => ({ + useAccessGroups: vi.fn().mockReturnValue({ data: [], isLoading: false, isError: false }), +})); + +vi.mock("@/components/common_components/AccessGroupSelector", () => ({ + default: () =>
access group selector
, +})); + +vi.mock("@/app/(dashboard)/hooks/keys/useKeys", () => { + const keysResult = { + data: { keys: [], total_count: 0, current_page: 1, total_pages: 1 }, + isPending: false, + isFetching: false, + refetch: vi.fn(), + }; + return { useKeys: vi.fn(() => keysResult) }; +}); + +vi.mock("@/components/key_team_helpers/filter_helpers", () => ({ + fetchTeamFilterOptions: vi.fn().mockResolvedValue({ keyAliases: [], organizationIds: [], userIds: [] }), + fetchAllKeyAliases: vi.fn().mockResolvedValue([]), + fetchAllOrganizations: vi.fn().mockResolvedValue([]), +})); + +const createMockTeamData = (overrides = {}) => ({ + team_id: "team-a1b2", + team_info: { + team_alias: "Test Team", + team_id: "team-a1b2", + organization_id: null, + admins: ["admin@test.com"], + members: ["user1@test.com"], + members_with_roles: [ + { user_id: "user1@test.com", user_email: "user1@test.com", role: "member", spend: 0, budget_id: "budget1" }, + ], + metadata: { disable_global_guardrails: true }, + tpm_limit: null, + rpm_limit: null, + max_budget: null, + budget_duration: null, + models: [], + blocked: false, + spend: 0, + max_parallel_requests: null, + budget_reset_at: null, + model_id: null, + litellm_model_table: null, + created_at: "2024-01-01T00:00:00Z", + team_member_budget_table: null, + guardrails: [], + policies: [], + object_permission: null, + ...overrides, + }, + keys: [], + team_memberships: [], +}); + +describe("ModelsAndEndpointsPage ?team drill-in", () => { + beforeEach(() => { + authState.userRole = "Admin"; + authState.isViewOnly = false; + can.mockReturnValue(true); + vi.mocked(networking.getGuardrailsList).mockResolvedValue({ guardrails: [] }); + vi.mocked(networking.getPoliciesList).mockResolvedValue({ policies: [] }); + vi.mocked(networking.fetchMCPAccessGroups).mockResolvedValue([]); + vi.mocked(networking.getTeamPermissionsCall).mockResolvedValue({ + all_available_permissions: [], + team_member_permissions: [], + } as never); + vi.mocked(networking.teamInfoCall).mockResolvedValue(createMockTeamData() as never); + // eslint-disable-next-line @typescript-eslint/no-explicit-any -- jsdom has no ResizeObserver global to type against + (global as any).ResizeObserver = class { + observe() {} + unobserve() {} + disconnect() {} + }; + }); + + afterEach(() => { + vi.clearAllMocks(); + }); + + it("shows the Disable all global guardrails switch to a proxy admin session", async () => { + const user = userEvent.setup({ delay: null }); + renderWithProviders(, { searchParams: { team: "team-a1b2" } }); + + await user.click(await screen.findByRole("tab", { name: "Settings" })); + await user.click(await screen.findByRole("button", { name: /edit settings/i })); + await screen.findByLabelText(/Team Name/); + + expect(screen.getByRole("switch", { name: /Disable all global guardrails/i })).toBeChecked(); + }); + + it("keeps the switch hidden from an internal user session on the same team", async () => { + authState.userRole = "Internal User"; + renderWithProviders(, { searchParams: { team: "team-a1b2" } }); + + await screen.findByText("Test Team"); + + expect(screen.queryByRole("button", { name: /edit settings/i })).not.toBeInTheDocument(); + expect(screen.queryByText("Disable all global guardrails")).not.toBeInTheDocument(); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/page.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/page.test.tsx index 105f6ff3043..652a3e804db 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/page.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/page.test.tsx @@ -25,12 +25,16 @@ vi.mock("@/components/molecules/cost_optimization_feedback_banner", () => ({ def vi.mock("@/components/model_info_view", () => ({ default: ({ modelId }: { modelId: string }) =>
model:{modelId}
, })); +const teamInfoProps = vi.hoisted(() => vi.fn()); vi.mock("@/components/team/TeamInfo", () => ({ - default: ({ teamId, is_team_admin }: { teamId: string; is_team_admin: boolean }) => ( -
- team:{teamId} -
- ), + default: (props: { teamId: string; is_team_admin: boolean; is_proxy_admin: boolean }) => { + teamInfoProps(props); + return ( +
+ team:{props.teamId} +
+ ); + }, })); const mockUseAuthorized = vi.fn(); @@ -107,12 +111,21 @@ describe("ModelsAndEndpointsPage", () => { expect(screen.getByTestId("team-info")).toHaveAttribute("data-team-admin", "true"); }); + it("passes is_proxy_admin for an admin session on the ?team drill-in", () => { + detailState.teamId = "team-a1b2"; + renderPage(); + expect(teamInfoProps).toHaveBeenLastCalledWith( + expect.objectContaining({ is_proxy_admin: true, is_team_admin: true }), + ); + }); + it("opens the ?team drill-in without edit rights for a view-only admin", () => { mockUseAuthorized.mockReturnValue(VIEW_ONLY_ADMIN); detailState.teamId = "team-9"; renderPage(); expect(screen.getByTestId("team-info")).toHaveTextContent("team:team-9"); expect(screen.getByTestId("team-info")).toHaveAttribute("data-team-admin", "false"); + expect(teamInfoProps).toHaveBeenLastCalledWith(expect.objectContaining({ is_proxy_admin: false })); }); it("hides admin-only tabs for a non-admin user", () => { diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/page.tsx b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/page.tsx index bbc803af700..d8952b88545 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/page.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/page.tsx @@ -151,7 +151,7 @@ export default function ModelsAndEndpointsPage() { onClose={close} accessToken={accessToken} is_team_admin={userRole === "Admin" && !isViewOnly} - is_proxy_admin={userRole === "Proxy Admin"} + is_proxy_admin={userRole === "Admin" && !isViewOnly} userModels={allModelsOnProxy} editTeam={false} onUpdate={invalidateModels} From 259f954f5642cd1156ce6a5930ea44bd5761655d Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Thu, 24 Sep 2026 11:55:07 -0700 Subject: [PATCH 061/120] test(e2e/ui): give the logout specs their own admin session (#42930) #42463 made Logout revoke the dashboard session key on the server. Both logout specs ran on the shared ADMIN_STORAGE_PATH session that globalSetup mints once, so clicking Logout revoked the key every later admin spec reuses. The CircleCI run is serial, and from the auth/ folder on, every admin-session spec failed with "Invalid proxy server token passed" (80 failures, up from 7) while the internal-user, internal-viewer and team-admin specs kept passing. Each logout spec now starts from an empty storage state and logs in through the login page, so the session it revokes is its own. The login steps live in a shared logInThroughLoginPage helper next to the other onboarding helpers. --- tests/e2e/ui/helpers/userOnboarding.ts | 10 ++++++++++ tests/e2e/ui/tests/auth/logout.spec.ts | 8 ++++++-- tests/e2e/ui/tests/auth/proxyLogoutUrl.spec.ts | 12 ++++++++---- 3 files changed, 24 insertions(+), 6 deletions(-) diff --git a/tests/e2e/ui/helpers/userOnboarding.ts b/tests/e2e/ui/helpers/userOnboarding.ts index a1ea6e5e82b..14e2b0257b2 100644 --- a/tests/e2e/ui/helpers/userOnboarding.ts +++ b/tests/e2e/ui/helpers/userOnboarding.ts @@ -57,3 +57,13 @@ export async function expectUnrestrictedDashboard(page: Page): Promise { expect(info.ok(), `Read own user with dashboard session: HTTP ${info.status()}`).toBe(true); expect((await info.json()).user_id).toBe(session.user_id); } + +export async function logInThroughLoginPage(page: Page, email: string, password: string): Promise { + await page.goto(`${rootPath()}/ui/login`); + await page.getByPlaceholder("Enter your username").fill(email); + await page.getByPlaceholder("Enter your password").fill(password); + await page.getByRole("button", { name: "Login", exact: true }).click(); + await page.waitForURL((url) => url.pathname.startsWith(`${rootPath()}/ui`) && !url.pathname.includes("/login"), { + timeout: 30_000, + }); +} diff --git a/tests/e2e/ui/tests/auth/logout.spec.ts b/tests/e2e/ui/tests/auth/logout.spec.ts index 351ba91e8d7..92c31456353 100644 --- a/tests/e2e/ui/tests/auth/logout.spec.ts +++ b/tests/e2e/ui/tests/auth/logout.spec.ts @@ -1,10 +1,14 @@ import { test, expect } from "@playwright/test"; -import { ADMIN_STORAGE_PATH } from "../../constants"; +import { Role, users } from "../../fixtures/users"; +import { logInThroughLoginPage } from "../../helpers/userOnboarding"; test.describe("Logout", () => { - test.use({ storageState: ADMIN_STORAGE_PATH }); + test.use({ storageState: { cookies: [], origins: [] } }); test("Clicking Logout clears the session and forces re-login on a protected page", async ({ page }) => { + const admin = users[Role.ProxyAdmin]; + await logInThroughLoginPage(page, admin.email, admin.password); + await page.goto("/ui"); // Scope to the sidebar; the top-bar breadcrumb also shows "Virtual Keys". await expect(page.getByRole("complementary").getByText("Virtual Keys")).toBeVisible({ timeout: 10_000 }); diff --git a/tests/e2e/ui/tests/auth/proxyLogoutUrl.spec.ts b/tests/e2e/ui/tests/auth/proxyLogoutUrl.spec.ts index 7f6cc6f2f87..3a79ce82717 100644 --- a/tests/e2e/ui/tests/auth/proxyLogoutUrl.spec.ts +++ b/tests/e2e/ui/tests/auth/proxyLogoutUrl.spec.ts @@ -1,5 +1,6 @@ import { test, expect } from "@playwright/test"; -import { ADMIN_STORAGE_PATH } from "../../constants"; +import { Role, users } from "../../fixtures/users"; +import { logInThroughLoginPage } from "../../helpers/userOnboarding"; /** * Runs as part of the standard e2e suite: both `run_e2e.sh` and the CircleCI @@ -16,9 +17,12 @@ const LOGOUT_URL = process.env.PROXY_LOGOUT_URL ?? ""; test.skip(!LOGOUT_URL, "Requires PROXY_LOGOUT_URL env var"); test.describe("PROXY_LOGOUT_URL redirect", () => { - test.use({ storageState: ADMIN_STORAGE_PATH }); + test.use({ storageState: { cookies: [], origins: [] } }); test("Logout clears the session and redirects to PROXY_LOGOUT_URL", async ({ page }) => { + const admin = users[Role.ProxyAdmin]; + await logInThroughLoginPage(page, admin.email, admin.password); + const target = new URL(LOGOUT_URL); // Stub the external logout destination so the assertion doesn't depend on @@ -46,8 +50,8 @@ test.describe("PROXY_LOGOUT_URL redirect", () => { await expect(page.getByRole("complementary").getByText("Virtual Keys")).toBeVisible({ timeout: 15_000 }); await settingsLoaded; - // Pre-condition: we start authenticated. The admin storage state carries a - // `token` cookie, so a real logout has something to tear down. + // Pre-condition: we start authenticated. The fresh login set a `token` + // cookie, so a real logout has something to tear down. const tokensBefore = (await page.context().cookies()).filter((c) => c.name === "token"); expect(tokensBefore.length, "should start logged in with a token cookie").toBeGreaterThan(0); From 342f9c3bf551a9e8e73dac3ee022c3a245e07fd9 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 24 Sep 2026 11:55:36 -0700 Subject: [PATCH 062/120] fix(azure): update gpt-audio-mini and gpt-5-chat deprecation dates from the retirement schedule (#43017) Co-authored-by: kerry Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/model_prices_and_context_window_backup.json | 4 ++-- model_prices_and_context_window.json | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 2600da22a77..8f38866387a 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -5734,7 +5734,7 @@ "supports_vision": false }, "azure/gpt-audio-mini": { - "deprecation_date": "2027-04-06", + "deprecation_date": "2027-06-15", "input_cost_per_audio_token": 1e-05, "input_cost_per_token": 6e-07, "litellm_provider": "azure", @@ -6768,7 +6768,7 @@ }, "azure/gpt-5-chat": { "cache_read_input_token_cost": 1.25e-07, - "deprecation_date": "2026-05-13", + "deprecation_date": "2026-06-29", "input_cost_per_token": 1.25e-06, "litellm_provider": "azure", "max_input_tokens": 128000, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 2600da22a77..8f38866387a 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -5734,7 +5734,7 @@ "supports_vision": false }, "azure/gpt-audio-mini": { - "deprecation_date": "2027-04-06", + "deprecation_date": "2027-06-15", "input_cost_per_audio_token": 1e-05, "input_cost_per_token": 6e-07, "litellm_provider": "azure", @@ -6768,7 +6768,7 @@ }, "azure/gpt-5-chat": { "cache_read_input_token_cost": 1.25e-07, - "deprecation_date": "2026-05-13", + "deprecation_date": "2026-06-29", "input_cost_per_token": 1.25e-06, "litellm_provider": "azure", "max_input_tokens": 128000, From f4f1a75a9cdea2db8e83287e6f449b89019d4fa0 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Thu, 24 Sep 2026 11:56:42 -0700 Subject: [PATCH 063/120] test(vertex_ai): run the files peak-memory guards without coverage tracing (#42914) The new CircleCI tests pipeline (#42773) runs tests/unit under pytest-cov on CPython 3.12.2, where coverage traces every line through sys.settrace. The two tracemalloc peak comparisons in test_vertex_ai_files_streaming.py drive 8000-row payloads through both pipelines and slow from ~10s to over 3 minutes under that tracer, so both hit the 90s pytest-timeout on every run. Mark them no_cover so pytest-cov pauses tracing for just these two. Their assertions are unchanged and every other test in the file still reports coverage. --- .../unit/llms/vertex_ai/files/test_vertex_ai_files_streaming.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tests/unit/llms/vertex_ai/files/test_vertex_ai_files_streaming.py b/tests/unit/llms/vertex_ai/files/test_vertex_ai_files_streaming.py index 29eaf38b427..6851591f8b5 100644 --- a/tests/unit/llms/vertex_ai/files/test_vertex_ai_files_streaming.py +++ b/tests/unit/llms/vertex_ai/files/test_vertex_ai_files_streaming.py @@ -269,6 +269,7 @@ class TestStreamingPeakMemory: measurement removes any garbage the previous run left behind. """ + @pytest.mark.no_cover def test_streaming_peak_well_below_list_pipeline(self): cfg = VertexAIFilesConfig() raw = _make_openai_jsonl_bytes(8000) @@ -345,6 +346,7 @@ class TestPathSourcedStreaming: first_labels = json.loads(lines[0])["request"]["labels"] assert _get_litellm_batch_custom_id_from_labels(first_labels) == "request-0" + @pytest.mark.no_cover def test_path_source_peak_stays_below_list_pipeline(self, tmp_path): cfg = VertexAIFilesConfig() path, raw = self._write_jsonl(tmp_path, 8000) From 2cbaa4c9b5a4f9a7b6dc62a91ceb7c4dabff9841 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Thu, 24 Sep 2026 11:56:56 -0700 Subject: [PATCH 064/120] test(mcp): stop a comprehension variable from shadowing the body() helper (#42906) test_malformed_bodies_missing_users_and_foreign_servers_are_rejected used `body` as a comprehension variable and then called the module-level `body()` helper a few lines later. CPython 3.12.2, which the CircleCI integration job runs, compiles that later call as a local read, so the test raised UnboundLocalError on every integration-mcp run since #42652. Newer 3.12 patch releases and 3.13 compile it as a global read, which is why it passes locally. Renaming the comprehension variable makes both reads unambiguous. --- tests/integration/mcp/test_mcp_user_env_vars.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/integration/mcp/test_mcp_user_env_vars.py b/tests/integration/mcp/test_mcp_user_env_vars.py index d9cecaccadb..0170eb77491 100644 --- a/tests/integration/mcp/test_mcp_user_env_vars.py +++ b/tests/integration/mcp/test_mcp_user_env_vars.py @@ -195,7 +195,7 @@ def test_malformed_bodies_missing_users_and_foreign_servers_are_rejected(gateway path: Final = f"/v1/mcp/server/{identity}/user-env-vars" payload: Final[dict[str, JsonValue]] = {"values": {TOKEN: "x"}} malformed: Final[tuple[dict[str, JsonValue], ...]] = ({"values": {TOKEN: 7}}, {"values": ["a"]}, {}) - assert [gateway.request("POST", path, body, key=key).status_code for body in malformed] == [422, 422, 422] + assert [gateway.request("POST", path, bad, key=key).status_code for bad in malformed] == [422, 422, 422] assert set_names(env_status(gateway, key, identity)) == {TOKEN: False} assert [gateway.client.request(method, path, json=payload).status_code for method in METHODS] == [401, 401, 401] no_user: Final = tuple(gateway.request(method, path, payload, key=userless) for method in METHODS) From 4c68a97c2585783443fb066fedcee2bd439102f7 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Thu, 24 Sep 2026 11:59:52 -0700 Subject: [PATCH 065/120] test(mcp): patch create_mcp_server_if_identifier_free in the store-model-in-db MCP tests (#42916) #42791 switched the MCP management endpoints from create_mcp_server to create_mcp_server_if_identifier_free, which keeps the same arguments and returns the created row on success. Two tests in tests/store_model_in_db_tests/test_mcp_servers.py still patched the old name, so mock.patch raised AttributeError before the tests ran and proxy_store_model_in_db_tests has been red on main since. --- tests/store_model_in_db_tests/test_mcp_servers.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/store_model_in_db_tests/test_mcp_servers.py b/tests/store_model_in_db_tests/test_mcp_servers.py index 735d5d71ad3..0e20880ede9 100644 --- a/tests/store_model_in_db_tests/test_mcp_servers.py +++ b/tests/store_model_in_db_tests/test_mcp_servers.py @@ -134,7 +134,7 @@ async def test_create_mcp_server_direct(): "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw" ) as mock_get_prisma, mock.patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints.create_mcp_server", + "litellm.proxy.management_endpoints.mcp_management_endpoints.create_mcp_server_if_identifier_free", new_callable=mock.AsyncMock, ) as mock_create, mock.patch( @@ -345,7 +345,7 @@ async def test_create_mcp_server_invalid_alias(): "litellm.proxy.management_endpoints.mcp_management_endpoints.get_mcp_server" ) as mock_get_server, mock.patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints.create_mcp_server" + "litellm.proxy.management_endpoints.mcp_management_endpoints.create_mcp_server_if_identifier_free" ) as mock_create, ): from litellm.proxy.management_endpoints.mcp_management_endpoints import ( From 759c216366dc98c7221e591ecf0665e22a058e1a Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Thu, 24 Sep 2026 12:00:10 -0700 Subject: [PATCH 066/120] test(guardrails): run the cache-hit redis outage test on the shared owned_redis helper (#42925) * test(guardrails): run the cache-hit redis outage test on the shared owned_redis helper test_redis_outage_keeps_serving_in_memory_hits (#42780) spawned `redis-server` straight from PATH. The CircleCI integration machine has no redis-server binary, so the test died with FileNotFoundError before reaching the proxy. Use integration._support.redis_process.owned_redis, which runs the local binary when there is one and otherwise the job's redis-cache container, and take the outage with its stop()/start() pair the way test_redis_recovery already does. The assertions are unchanged. * test(guardrails): describe the owned_redis stop as an outage, not a kill --- .../test_cache_hit_guardrail_metrics_chaos.py | 46 ++++++------------- 1 file changed, 13 insertions(+), 33 deletions(-) diff --git a/tests/integration/observability/test_cache_hit_guardrail_metrics_chaos.py b/tests/integration/observability/test_cache_hit_guardrail_metrics_chaos.py index faaa1fc7325..c77b8eebe33 100644 --- a/tests/integration/observability/test_cache_hit_guardrail_metrics_chaos.py +++ b/tests/integration/observability/test_cache_hit_guardrail_metrics_chaos.py @@ -1,12 +1,9 @@ import json import signal -import socket -import subprocess import threading import uuid -from collections.abc import Callable, Generator +from collections.abc import Callable from concurrent.futures import ThreadPoolExecutor -from contextlib import contextmanager from pathlib import Path from typing import Final @@ -14,6 +11,7 @@ import httpx import psutil from integration._support.client import Gateway, eventually, object_value, string_value from integration._support.process import owned_proxy_process +from integration._support.redis_process import owned_redis from integration._support.wire import Reply, Request, wire_server from prometheus_client.parser import text_string_to_metric_families from test_cache_hit_guardrail_metrics import ( @@ -31,22 +29,6 @@ from test_cache_hit_guardrail_metrics import ( BURST: Final = 10 -@contextmanager -def _redis(port: int) -> Generator[subprocess.Popen[bytes], None, None]: - process: Final = subprocess.Popen(["redis-server", "--port", str(port), "--save", ""], stdout=subprocess.DEVNULL) - try: - yield process - finally: - process.kill() - process.wait(timeout=10) - - -def _free_port() -> int: - with socket.socket() as reserve: - reserve.bind(("127.0.0.1", 0)) - return reserve.getsockname()[1] - - def _deployment_id(candidate: Gateway, model_name: str) -> str: entries: Final = candidate.get("/model/info")["data"] assert isinstance(entries, list) @@ -224,18 +206,16 @@ def test_stalled_guardrail_sink_recovers_and_counts(gateway: Gateway, tmp_path: def test_redis_outage_keeps_serving_in_memory_hits(gateway: Gateway, tmp_path: Path) -> None: - """X2: the redis cache keeps an in-memory shadow, so a redis kill does not stop cache-hit rejects.""" + """X2: the redis cache keeps an in-memory shadow, so a redis outage does not stop cache-hit rejects.""" marker: Final = uuid.uuid4().hex - port: Final = _free_port() - with _redis(port) as redis_one: - with _rig(gateway, tmp_path, marker, env={"REDIS_HOST": "127.0.0.1", "REDIS_PORT": str(port)}) as rig: + with owned_redis(tmp_path) as cache: + with _rig(gateway, tmp_path, marker, env={"REDIS_HOST": cache.host, "REDIS_PORT": str(cache.port)}) as rig: bodies: Final = _burst_bodies(rig, marker, None)[:BURST] _warm(rig, bodies) reject: Final = rig.candidate.request("POST", *bodies[0]) assert reject.status_code == 400, reject.text warmed_hits: Final = rig.provider.received.qsize() - redis_one.kill() - redis_one.wait(timeout=10) + cache.stop() outcomes: Final = _fire(rig, bodies[1:]) assert all(status == 400 for status, _ in outcomes), outcomes assert rig.provider.received.qsize() == warmed_hits, ( @@ -243,13 +223,13 @@ def test_redis_outage_keeps_serving_in_memory_hits(gateway: Gateway, tmp_path: P warmed_hits, rig.provider.received.qsize(), ) - with _redis(port): - recovered: Final = rig.candidate.request( - "POST", - "/v1/chat/completions", - _chat_body(rig.model_name, "x2 rehit " + marker, rig.guardrail_name), - ) - assert recovered.status_code == 400, recovered.text + cache.start() + recovered: Final = rig.candidate.request( + "POST", + "/v1/chat/completions", + _chat_body(rig.model_name, "x2 rehit " + marker, rig.guardrail_name), + ) + assert recovered.status_code == 400, recovered.text _expect_exactly_once(rig, (rig.model_name,), (rig.deployment_id,), 1 + len(bodies)) From bdf854c3ea3ecbfd399c26a33710d2c9644eb616 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 24 Sep 2026 19:08:32 +0000 Subject: [PATCH 067/120] feat(rust): shape Anthropic Messages requests natively (#42982) * test(rust): encode anthropic response serialization shape as rstest cases Co-Authored-By: Claude Opus 5.5 * feat(rust): shape Anthropic Messages requests natively The Rust Messages route only relayed the body. It now runs the request shaping the Python handler does for the direct Anthropic provider: history sanitizers (empty blocks, tool ids, replayed web search results, provider_specific_fields, encrypted reasoning, advisor blocks), reasoning_effort and adaptive/legacy thinking translation against the model's capability flags, the sampling and speed gates under drop_params, the metadata allowlist, additional_drop_params, reasoning auto summary, OAuth and ANTHROPIC_AUTH_TOKEN credentials, provider_specific_header merging and anthropic-beta injection. Capability flags and LiteLLM settings reach Rust through route_host.shaping(). A request the route rejects before the call now maps to BadRequestError instead of APIConnectionError Co-Authored-By: Claude Fable 5.1 * test(rust): port Anthropic Messages shaping tests and pin comment contracts as cases Every Python unit test that exercises the ported shaping for the direct Anthropic provider now has a named rstest counterpart, and every comment that stated a behavior contract is deleted in favor of a case that pins it. Measured with cargo-mutants over the touched files, all viable mutants are caught Porting the tests surfaced parity gaps, fixed here to match Python: every casing of a forwarded anthropic-beta header is merged, replayed web search results are rewritten from their own block (an empty result keeps its slot and a server_tool_use with a non-string query stays), an empty output_config.effort falls back to medium, speed and reasoning effort errors quote values the way Python does, additional_drop_params apply after metadata validation and the auto summary and never touch model or messages, and a non-string metadata.user_id is rejected before the call * fix(rust): resolve Messages credentials through the secret source and scope headers by resolved provider The native Messages route read ANTHROPIC_API_KEY, ANTHROPIC_AUTH_TOKEN and the base URL straight from the process environment, so a key or base held in a configured secret manager was never found. Each provider config now declares its secret names and the route resolves them through the same SecretSource the OCR route uses, with the Python bridge passing in litellm's configured manager provider_specific_header entries were scoped by the explicit custom_llm_provider only, falling back to anthropic, so an azure_ai/ model lost its azure_ai scoped headers. Scoping now happens in the route after the provider is resolved from the model, as Python's handler does The Azure config now adds the same anthropic-beta feature headers Python's Azure route adds, and the metadata allowlist, reasoning auto summary and history sanitizers move from the core route into the llms crate, mirroring their home in Python's messages handler * test(rust): escape the dot in the metadata.user_id match pattern * refactor(rust-bridge): project Messages capabilities without mutable dicts The capability flags and effort tiers were built as dict comprehensions, which the type-discipline gate counts as mutable construction, and the asdict call carried a mutable-ok suppression that suppressed nothing. The flags are now passed one by one and the effort tiers are a frozen dataclass, which asdict projects to the same map the native side reads --------- Co-authored-by: Yujong Lee Co-authored-by: Claude Opus 5.5 --- litellm-rust/Cargo.lock | 1 + .../core-utils/src/dot_notation_indexing.rs | 274 +++ .../src/get_provider_specific_headers.rs | 93 + litellm-rust/crates/core-utils/src/lib.rs | 2 + .../crates/core/src/messages/common_utils.rs | 2 +- .../crates/core/src/messages/error.rs | 28 + litellm-rust/crates/core/src/messages/mod.rs | 8 +- .../crates/core/src/messages/prepare.rs | 501 +++++- .../crates/core/src/messages/route.rs | 52 +- .../crates/core/src/messages/tests.rs | 149 +- .../crates/core/src/messages/types.rs | 104 +- .../crates/llms/src/anthropic/common_utils.rs | 1568 +++++++++++++++++ .../messages/handler.rs | 270 +++ .../messages/headers.rs | 643 +++++++ .../experimental_pass_through/messages/mod.rs | 3 + .../messages/thinking.rs | 1182 +++++++++++++ .../messages/transformation.rs | 808 ++++++++- litellm-rust/crates/llms/src/anthropic/mod.rs | 1 + .../anthropic/messages_transformation.rs | 145 +- .../anthropic_messages/transformation.rs | 232 ++- .../python-bridge/src/routes/messages/host.rs | 175 +- .../python-bridge/src/routes/messages/mod.rs | 3 +- litellm-rust/crates/types/Cargo.toml | 3 + .../anthropic_messages/anthropic_request.rs | 156 ++ .../anthropic_messages/anthropic_response.rs | 60 +- litellm-rust/crates/types/src/utils.rs | 15 + litellm/rust_bridge/messages/route_host.py | 110 +- tests/test_litellm/rust_bridge/AGENTS.md | 2 + .../rust_bridge/messages/test_route_host.py | 112 ++ .../rust_bridge/messages/test_secrets.py | 111 ++ .../messages/test_request_shaping.py | 237 +++ 31 files changed, 6891 insertions(+), 159 deletions(-) create mode 100644 litellm-rust/crates/core-utils/src/dot_notation_indexing.rs create mode 100644 litellm-rust/crates/core-utils/src/get_provider_specific_headers.rs create mode 100644 litellm-rust/crates/llms/src/anthropic/common_utils.rs create mode 100644 litellm-rust/crates/llms/src/anthropic/experimental_pass_through/messages/handler.rs create mode 100644 litellm-rust/crates/llms/src/anthropic/experimental_pass_through/messages/headers.rs create mode 100644 litellm-rust/crates/llms/src/anthropic/experimental_pass_through/messages/thinking.rs create mode 100644 tests/test_litellm/rust_bridge/messages/test_route_host.py create mode 100644 tests/test_litellm/rust_bridge/messages/test_secrets.py create mode 100644 tests/test_litellm_rust/messages/test_request_shaping.py diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index ef032bfe55c..0a91f0759c2 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -3414,6 +3414,7 @@ dependencies = [ name = "litellm-types" version = "0.1.0" dependencies = [ + "rstest", "serde", "serde_json", ] diff --git a/litellm-rust/crates/core-utils/src/dot_notation_indexing.rs b/litellm-rust/crates/core-utils/src/dot_notation_indexing.rs new file mode 100644 index 00000000000..be81d900f57 --- /dev/null +++ b/litellm-rust/crates/core-utils/src/dot_notation_indexing.rs @@ -0,0 +1,274 @@ +use serde_json::Value; + +#[derive(Clone, Debug, PartialEq, Eq)] +enum Segment { + Field(String), + Every, + Index(usize), +} + +fn parse_segments(path: &str) -> Option> { + let mut segments = Vec::new(); + let mut rest = path; + while !rest.is_empty() { + if let Some(after_open) = rest.strip_prefix('[') { + let (inside, after) = after_open.split_once(']')?; + segments.push(match inside { + "*" => Segment::Every, + index => Segment::Index(index.trim().parse().ok()?), + }); + rest = after.strip_prefix('.').unwrap_or(after); + continue; + } + let end = rest.find(['.', '[']).unwrap_or(rest.len()); + let (field, after) = rest.split_at(end); + if !field.is_empty() { + segments.push(Segment::Field(field.to_string())); + } + rest = after.strip_prefix('.').unwrap_or(after); + } + Some(segments) +} + +fn without_path(value: Value, segments: &[Segment]) -> Value { + let Some((segment, tail)) = segments.split_first() else { + return value; + }; + match (segment, value) { + (Segment::Field(name), Value::Object(object)) => Value::Object( + object + .into_iter() + .filter_map(|(key, item)| { + if key != *name { + return Some((key, item)); + } + (!tail.is_empty()).then(|| (key, without_path(item, tail))) + }) + .collect(), + ), + (Segment::Every, Value::Array(items)) => Value::Array( + items + .into_iter() + .map(|item| without_path(item, tail)) + .collect(), + ), + (Segment::Index(index), Value::Array(items)) => Value::Array( + items + .into_iter() + .enumerate() + .map(|(position, item)| { + if position == *index { + without_path(item, tail) + } else { + item + } + }) + .collect(), + ), + (_, value) => value, + } +} + +pub fn delete_nested_value(value: Value, path: &str) -> Value { + match parse_segments(path) { + Some(segments) => without_path(value, &segments), + None => value, + } +} + +#[cfg(test)] +mod tests { + use rstest::{fixture, rstest}; + use serde_json::json; + + use super::*; + + #[fixture] + fn body() -> Value { + json!({ + "tools": [ + {"name": "t0", "examples": ["a"], "arr": [{"f": 1, "k": 1}, {"f": 2, "k": 2}]}, + {"name": "t1", "examples": ["b"], "arr": [{"f": 3, "k": 3}]} + ], + "meta": {"user": "u", "inner": {"drop": 1, "keep": 2}}, + "top": 0.7 + }) + } + + #[rstest] + #[case::top_level_field("top", json!({ + "tools": [ + {"name": "t0", "examples": ["a"], "arr": [{"f": 1, "k": 1}, {"f": 2, "k": 2}]}, + {"name": "t1", "examples": ["b"], "arr": [{"f": 3, "k": 3}]} + ], + "meta": {"user": "u", "inner": {"drop": 1, "keep": 2}} + }))] + #[case::whole_object("meta", json!({ + "tools": [ + {"name": "t0", "examples": ["a"], "arr": [{"f": 1, "k": 1}, {"f": 2, "k": 2}]}, + {"name": "t1", "examples": ["b"], "arr": [{"f": 3, "k": 3}]} + ], + "top": 0.7 + }))] + #[case::nested_field("meta.inner.drop", json!({ + "tools": [ + {"name": "t0", "examples": ["a"], "arr": [{"f": 1, "k": 1}, {"f": 2, "k": 2}]}, + {"name": "t1", "examples": ["b"], "arr": [{"f": 3, "k": 3}]} + ], + "meta": {"user": "u", "inner": {"keep": 2}}, + "top": 0.7 + }))] + #[case::trailing_dot("meta.inner.drop.", json!({ + "tools": [ + {"name": "t0", "examples": ["a"], "arr": [{"f": 1, "k": 1}, {"f": 2, "k": 2}]}, + {"name": "t1", "examples": ["b"], "arr": [{"f": 3, "k": 3}]} + ], + "meta": {"user": "u", "inner": {"keep": 2}}, + "top": 0.7 + }))] + #[case::leading_and_doubled_dots(".meta..inner.drop", json!({ + "tools": [ + {"name": "t0", "examples": ["a"], "arr": [{"f": 1, "k": 1}, {"f": 2, "k": 2}]}, + {"name": "t1", "examples": ["b"], "arr": [{"f": 3, "k": 3}]} + ], + "meta": {"user": "u", "inner": {"keep": 2}}, + "top": 0.7 + }))] + #[case::field_in_every_element("tools[*].examples", json!({ + "tools": [ + {"name": "t0", "arr": [{"f": 1, "k": 1}, {"f": 2, "k": 2}]}, + {"name": "t1", "arr": [{"f": 3, "k": 3}]} + ], + "meta": {"user": "u", "inner": {"drop": 1, "keep": 2}}, + "top": 0.7 + }))] + #[case::whole_array_field_in_every_element("tools[*].arr", json!({ + "tools": [ + {"name": "t0", "examples": ["a"]}, + {"name": "t1", "examples": ["b"]} + ], + "meta": {"user": "u", "inner": {"drop": 1, "keep": 2}}, + "top": 0.7 + }))] + #[case::field_in_indexed_element("tools[1].examples", json!({ + "tools": [ + {"name": "t0", "examples": ["a"], "arr": [{"f": 1, "k": 1}, {"f": 2, "k": 2}]}, + {"name": "t1", "arr": [{"f": 3, "k": 3}]} + ], + "meta": {"user": "u", "inner": {"drop": 1, "keep": 2}}, + "top": 0.7 + }))] + #[case::padded_index("tools[ 1 ].examples", json!({ + "tools": [ + {"name": "t0", "examples": ["a"], "arr": [{"f": 1, "k": 1}, {"f": 2, "k": 2}]}, + {"name": "t1", "arr": [{"f": 3, "k": 3}]} + ], + "meta": {"user": "u", "inner": {"drop": 1, "keep": 2}}, + "top": 0.7 + }))] + #[case::field_right_after_bracket("tools[0]examples", json!({ + "tools": [ + {"name": "t0", "arr": [{"f": 1, "k": 1}, {"f": 2, "k": 2}]}, + {"name": "t1", "examples": ["b"], "arr": [{"f": 3, "k": 3}]} + ], + "meta": {"user": "u", "inner": {"drop": 1, "keep": 2}}, + "top": 0.7 + }))] + #[case::index_then_wildcard("tools[0].arr[*].f", json!({ + "tools": [ + {"name": "t0", "examples": ["a"], "arr": [{"k": 1}, {"k": 2}]}, + {"name": "t1", "examples": ["b"], "arr": [{"f": 3, "k": 3}]} + ], + "meta": {"user": "u", "inner": {"drop": 1, "keep": 2}}, + "top": 0.7 + }))] + #[case::wildcard_then_index_only_where_it_exists("tools[*].arr[1].f", json!({ + "tools": [ + {"name": "t0", "examples": ["a"], "arr": [{"f": 1, "k": 1}, {"k": 2}]}, + {"name": "t1", "examples": ["b"], "arr": [{"f": 3, "k": 3}]} + ], + "meta": {"user": "u", "inner": {"drop": 1, "keep": 2}}, + "top": 0.7 + }))] + #[case::nested_wildcards("tools[*].arr[*].f", json!({ + "tools": [ + {"name": "t0", "examples": ["a"], "arr": [{"k": 1}, {"k": 2}]}, + {"name": "t1", "examples": ["b"], "arr": [{"k": 3}]} + ], + "meta": {"user": "u", "inner": {"drop": 1, "keep": 2}}, + "top": 0.7 + }))] + fn deletes_the_addressed_field(body: Value, #[case] path: &str, #[case] expected: Value) { + assert_eq!(delete_nested_value(body, path), expected); + } + + #[rstest] + #[case::empty_path("")] + #[case::missing_field("missing")] + #[case::missing_parent("missing.field")] + #[case::field_through_a_scalar("top.value")] + #[case::field_on_an_array("tools.name")] + #[case::index_on_an_object("meta[0].user")] + #[case::wildcard_on_an_object("meta[*].user")] + #[case::wildcard_over_scalars("tools[*].examples[*].name")] + #[case::index_out_of_range("tools[5].name")] + #[case::every_element_itself("tools[*]")] + #[case::indexed_element_itself("tools[0]")] + #[case::nested_element_itself("tools[*].arr[0]")] + #[case::negative_index("tools[-1].name")] + #[case::non_numeric_index("tools[x].name")] + #[case::empty_index("tools[].name")] + #[case::unclosed_bracket("top[0")] + fn leaves_the_value_untouched(body: Value, #[case] path: &str) { + assert_eq!(delete_nested_value(body.clone(), path), body); + } + + #[rstest] + #[case::wildcards_indices_and_nesting( + json!({"tools": [ + {"name": "t0", "configs": [{"id": "c0", "remove_me": 1, "keep": 1}, {"id": "c1", "remove_me": 2, "keep": 2}], "metadata": {"drop_this": 1, "preserve": 1}}, + {"name": "t1", "configs": [{"id": "c0", "remove_me": 3, "keep": 3}, {"id": "c1", "remove_me": 4, "keep": 4}], "metadata": {"drop_this": 2, "preserve": 2}}, + {"name": "t2", "configs": [{"id": "c0", "remove_me": 5, "keep": 5}], "metadata": {"drop_this": 3, "preserve": 3}} + ]}), + &["tools[*].configs[1].remove_me", "tools[1].metadata.drop_this", "tools[*].configs[*].id"], + json!({"tools": [ + {"name": "t0", "configs": [{"remove_me": 1, "keep": 1}, {"keep": 2}], "metadata": {"drop_this": 1, "preserve": 1}}, + {"name": "t1", "configs": [{"remove_me": 3, "keep": 3}, {"keep": 4}], "metadata": {"preserve": 2}}, + {"name": "t2", "configs": [{"remove_me": 5, "keep": 5}], "metadata": {"drop_this": 3, "preserve": 3}} + ]}), + )] + #[case::simple_and_wildcard_nesting( + json!({ + "tools": [{"name": "t1", "simple_nested": {"remove": 1, "keep": 2}, "complex": [{"nested": {"remove": 3, "keep": 4}}]}], + "top_level_remove": "should_go", + "top_level_keep": "should_stay" + }), + &["tools[*].simple_nested.remove", "tools[*].complex[*].nested.remove"], + json!({ + "tools": [{"name": "t1", "simple_nested": {"keep": 2}, "complex": [{"nested": {"keep": 4}}]}], + "top_level_remove": "should_go", + "top_level_keep": "should_stay" + }), + )] + #[case::triple_nested_wildcards( + json!({"tools": [{"name": "t1", "arr1": [ + {"arr2": [{"field": 1, "keep": 1}, {"field": 2, "keep": 2}]}, + {"arr2": [{"field": 3, "keep": 3}]} + ]}]}), + &["tools[*].arr1[*].arr2[*].field"], + json!({"tools": [{"name": "t1", "arr1": [ + {"arr2": [{"keep": 1}, {"keep": 2}]}, + {"arr2": [{"keep": 3}]} + ]}]}), + )] + fn applies_paths_in_sequence( + #[case] value: Value, + #[case] paths: &[&str], + #[case] expected: Value, + ) { + let deleted = paths + .iter() + .fold(value, |value, path| delete_nested_value(value, path)); + assert_eq!(deleted, expected); + } +} diff --git a/litellm-rust/crates/core-utils/src/get_provider_specific_headers.rs b/litellm-rust/crates/core-utils/src/get_provider_specific_headers.rs new file mode 100644 index 00000000000..bfcd448e2d8 --- /dev/null +++ b/litellm-rust/crates/core-utils/src/get_provider_specific_headers.rs @@ -0,0 +1,93 @@ +use litellm_types::utils::{ProviderSpecificHeader, ProviderSpecificHeaders}; +use serde_json::{Map, Value}; + +pub fn get_provider_specific_headers( + provider_specific_header: Option<&ProviderSpecificHeaders>, + custom_llm_provider: &str, +) -> Map { + let entries: &[ProviderSpecificHeader] = match provider_specific_header { + None => &[], + Some(ProviderSpecificHeaders::One(entry)) => std::slice::from_ref(entry), + Some(ProviderSpecificHeaders::Many(entries)) => entries, + }; + entries + .iter() + .filter(|entry| { + entry + .custom_llm_provider + .split(',') + .any(|scoped| scoped.trim() == custom_llm_provider) + }) + .flat_map(|entry| entry.extra_headers.clone()) + .collect() +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + use serde_json::json; + + use super::*; + + #[rstest] + #[case::single_entry_for_the_provider( + json!({"custom_llm_provider": "anthropic", "extra_headers": {"Authorization": "Bearer t", "Custom-Header": "v"}}), + json!({"Authorization": "Bearer t", "Custom-Header": "v"}), + )] + #[case::single_entry_for_another_provider( + json!({"custom_llm_provider": "openai", "extra_headers": {"Authorization": "Bearer t"}}), + json!({}), + )] + #[case::provider_in_a_comma_separated_scope( + json!({"custom_llm_provider": "bedrock,anthropic,vertex_ai", "extra_headers": {"anthropic-beta": "context-1m-2025-08-07"}}), + json!({"anthropic-beta": "context-1m-2025-08-07"}), + )] + #[case::provider_missing_from_a_comma_separated_scope( + json!({"custom_llm_provider": "bedrock,vertex_ai", "extra_headers": {"anthropic-beta": "test"}}), + json!({}), + )] + #[case::scope_with_spaces( + json!({"custom_llm_provider": "bedrock, anthropic , vertex_ai", "extra_headers": {"anthropic-beta": "test"}}), + json!({"anthropic-beta": "test"}), + )] + #[case::scope_names_must_match_exactly( + json!({"custom_llm_provider": "anthropic_text", "extra_headers": {"anthropic-beta": "test"}}), + json!({}), + )] + #[case::entries_scope_independently( + json!([ + {"custom_llm_provider": "anthropic,bedrock,vertex_ai", "extra_headers": {"anthropic-beta": "context-1m-2025-08-07"}}, + {"custom_llm_provider": "bedrock", "extra_headers": {"x-bedrock-only": "no"}}, + {"custom_llm_provider": "anthropic", "extra_headers": {"authorization": "Bearer sk-ant-oat01-fake-token"}} + ]), + json!({"anthropic-beta": "context-1m-2025-08-07", "authorization": "Bearer sk-ant-oat01-fake-token"}), + )] + #[case::later_entries_win( + json!([ + {"custom_llm_provider": "anthropic", "extra_headers": {"x-scoped": "first"}}, + {"custom_llm_provider": "anthropic", "extra_headers": {"x-scoped": "second"}} + ]), + json!({"x-scoped": "second"}), + )] + #[case::empty_list(json!([]), json!({}))] + #[case::entry_without_scope(json!({"extra_headers": {"x-scoped": "yes"}}), json!({}))] + #[case::entry_without_headers(json!({"custom_llm_provider": "anthropic"}), json!({}))] + fn provider_specific_headers_match_the_scoped_provider( + #[case] configured: Value, + #[case] expected: Value, + ) { + let configured: ProviderSpecificHeaders = serde_json::from_value(configured).unwrap(); + assert_eq!( + Value::Object(get_provider_specific_headers( + Some(&configured), + "anthropic" + )), + expected + ); + } + + #[test] + fn no_configured_headers_match_nothing() { + assert_eq!(get_provider_specific_headers(None, "anthropic"), Map::new()); + } +} diff --git a/litellm-rust/crates/core-utils/src/lib.rs b/litellm-rust/crates/core-utils/src/lib.rs index ceb0e9eb3f2..a937f55654e 100644 --- a/litellm-rust/crates/core-utils/src/lib.rs +++ b/litellm-rust/crates/core-utils/src/lib.rs @@ -1,7 +1,9 @@ pub mod call_arguments; pub mod core_helpers; +pub mod dot_notation_indexing; pub mod exception_mapping_utils; pub mod get_llm_provider_logic; +pub mod get_provider_specific_headers; pub mod params; pub mod prompt_templates; pub mod secret_redaction; diff --git a/litellm-rust/crates/core/src/messages/common_utils.rs b/litellm-rust/crates/core/src/messages/common_utils.rs index dcefa3ebffc..015e026f6da 100644 --- a/litellm-rust/crates/core/src/messages/common_utils.rs +++ b/litellm-rust/crates/core/src/messages/common_utils.rs @@ -1,5 +1,5 @@ use litellm_http::request::string_headers as shared_string_headers; -pub(super) use litellm_http::request::{has_bearer_auth, has_header, truncate_error_body}; +pub(super) use litellm_http::request::truncate_error_body; use litellm_llms::{ anthropic::experimental_pass_through::messages::transformation::ANTHROPIC_MESSAGES_CONFIG, azure_ai::anthropic::messages_transformation::AZURE_ANTHROPIC_MESSAGES_CONFIG, diff --git a/litellm-rust/crates/core/src/messages/error.rs b/litellm-rust/crates/core/src/messages/error.rs index 51fb764032c..2a9723beb38 100644 --- a/litellm-rust/crates/core/src/messages/error.rs +++ b/litellm-rust/crates/core/src/messages/error.rs @@ -1,3 +1,5 @@ +use std::sync::Arc; + use litellm_llms::base_llm::chat::transformation::Error as LlmError; #[derive(Clone, Debug, PartialEq, Eq, thiserror::Error)] @@ -18,8 +20,34 @@ pub enum Error { Transport(#[from] litellm_http::transport::Error), #[error(transparent)] Headers(#[from] litellm_http::request::HeaderError), + #[error(transparent)] + Secret(#[from] SecretError), } +#[derive(Clone, Debug, thiserror::Error)] +#[error(transparent)] +pub struct SecretError(Arc); + +impl SecretError { + pub fn source_error(&self) -> &litellm_secrets::Error { + &self.0 + } +} + +impl From for Error { + fn from(error: litellm_secrets::Error) -> Self { + Self::Secret(SecretError(Arc::new(error))) + } +} + +impl PartialEq for SecretError { + fn eq(&self, other: &Self) -> bool { + Arc::ptr_eq(&self.0, &other.0) + } +} + +impl Eq for SecretError {} + impl From for Error { fn from(error: LlmError) -> Self { match error { diff --git a/litellm-rust/crates/core/src/messages/mod.rs b/litellm-rust/crates/core/src/messages/mod.rs index 289f79109dd..8795d4f8507 100644 --- a/litellm-rust/crates/core/src/messages/mod.rs +++ b/litellm-rust/crates/core/src/messages/mod.rs @@ -12,6 +12,9 @@ mod common_utils; mod handler; mod prepare; pub mod route; +use std::sync::Arc; + +use litellm_secrets::source::EnvironmentSecrets; use litellm_types::llms::anthropic_messages::anthropic_response::AnthropicMessagesResponse; use route::{LocalMessagesHost, MessagesCall, MessagesOutput, messages_machine}; use serde_json::Value; @@ -31,9 +34,12 @@ pub async fn messages(request: MessagesRequest<'_>) -> Result Ok(*message), MessagesOutput::Streamed => Err(Error::Unsupported( "streamed responses need a streaming host", diff --git a/litellm-rust/crates/core/src/messages/prepare.rs b/litellm-rust/crates/core/src/messages/prepare.rs index 850f9108869..dc4b3562e3f 100644 --- a/litellm-rust/crates/core/src/messages/prepare.rs +++ b/litellm-rust/crates/core/src/messages/prepare.rs @@ -1,51 +1,102 @@ -use litellm_core_utils::get_llm_provider_logic::{CustomLlmProvider, get_custom_llm_provider}; -use litellm_llms::base_llm::anthropic_messages::transformation::{ - BaseAnthropicMessagesConfig, MessagesAuthStrategy, +use litellm_core_utils::{ + dot_notation_indexing::delete_nested_value, + get_llm_provider_logic::{CustomLlmProvider, get_custom_llm_provider}, + get_provider_specific_headers::get_provider_specific_headers, + settings::Lookup, +}; +use litellm_llms::{ + anthropic::experimental_pass_through::messages::handler::shape_anthropic_messages_request, + base_llm::anthropic_messages::transformation::{ + BaseAnthropicMessagesConfig, MessagesTransformContext, + }, }; use litellm_types::llms::anthropic_messages::anthropic_request::AnthropicMessagesRequest; use serde_json::{Map, Value}; use super::{ Error, - common_utils::{has_bearer_auth, has_header, messages_provider_config, string_headers}, + common_utils::{messages_provider_config, string_headers}, }; use crate::messages::types::{MessagesRequest, ProviderMessagesRequest}; -pub(super) fn prepare_provider_request( - request: MessagesRequest<'_>, -) -> Result { - let provider_info = get_custom_llm_provider(request.model, request.custom_llm_provider) +pub(super) struct ResolvedProvider<'a> { + pub(super) model: &'a str, + pub(super) provider: &'a str, + pub(super) config: &'static dyn BaseAnthropicMessagesConfig, +} + +pub(super) fn resolve_provider<'a>( + model: &'a str, + custom_llm_provider: Option<&'a str>, +) -> Result, Error> { + let CustomLlmProvider { + model, + custom_llm_provider: provider, + } = get_custom_llm_provider(model, custom_llm_provider) .or_else(|| { - request - .custom_llm_provider - .map(|provider| CustomLlmProvider { - model: request.model, - custom_llm_provider: provider, - }) + custom_llm_provider.map(|provider| CustomLlmProvider { + model, + custom_llm_provider: provider, + }) }) .ok_or_else(|| { Error::InvalidProvider( "unable to resolve custom_llm_provider for messages request".to_string(), ) })?; - let model = provider_info.model.to_string(); - let provider = provider_info.custom_llm_provider; - let config = messages_provider_config(provider) .ok_or_else(|| Error::InvalidProvider(provider.to_string()))?; - let env_lookup = |key: &str| std::env::var(key).ok(); + Ok(ResolvedProvider { + model, + provider, + config, + }) +} - let headers = - validate_environment(config, request.extra_headers, request.api_key, &env_lookup)?; +pub(super) fn prepare_provider_request( + request: MessagesRequest<'_>, + resolved: ResolvedProvider<'_>, + secrets: &dyn Lookup, +) -> Result { + let ResolvedProvider { + model, + provider, + config, + } = resolved; + let model = model.to_string(); + let env_lookup = |key: &str| secrets.get(key); let typed_request: AnthropicMessagesRequest = - serde_json::from_value(request.body).map_err(|err| { - Error::InvalidRequest(format!("invalid Anthropic messages request: {err}")) - })?; - let transformed = config.transform_anthropic_messages_request(AnthropicMessagesRequest { - model: model.clone(), - ..typed_request - })?; + serde_json::from_value(request.body).map_err(invalid_request)?; + let sanitized = shape_anthropic_messages_request( + AnthropicMessagesRequest { + model: model.clone(), + ..typed_request + }, + request.shaping.reasoning_auto_summary, + )?; + let trimmed = + without_additional_drop_params(sanitized, &request.shaping.additional_drop_params)?; + let transformed = config.transform_anthropic_messages_request( + trimmed, + &MessagesTransformContext::new(request.shaping.capabilities, request.shaping.drop_params), + )?; + + let scoped = get_provider_specific_headers(request.provider_specific_header.as_ref(), provider); + let forwarded = string_headers(Some( + request + .extra_headers + .into_iter() + .flatten() + .chain(scoped) + .collect(), + ))?; + let authenticated = config.authenticate(forwarded, request.api_key, &env_lookup)?; + let headers = config.request_headers( + with_default_headers(authenticated, config.default_headers()), + &transformed, + ); + let body = serde_json::to_value(transformed).map_err(|err| { Error::InvalidRequest(format!( "failed to serialize Anthropic messages request: {err}" @@ -65,33 +116,371 @@ pub(super) fn prepare_provider_request( }) } -fn validate_environment( - config: &dyn BaseAnthropicMessagesConfig, - extra_headers: Option>, - api_key: Option<&str>, - env_lookup: &dyn Fn(&str) -> Option, -) -> Result, Error> { - let mut headers = string_headers(extra_headers)?; - - let auth_strategy = config.auth_strategy(); - let already_authorized = has_header(&headers, auth_strategy.header_name()) - || (config.accepts_bearer_auth() && has_bearer_auth(&headers)); - if !already_authorized { - let api_key = config.resolve_api_key(api_key, env_lookup)?; - let auth_header = match auth_strategy { - MessagesAuthStrategy::Bearer => { - ("authorization".to_string(), format!("Bearer {api_key}")) - } - MessagesAuthStrategy::Header(name) => (name.to_string(), api_key), - }; - headers.push(auth_header); - } - - for (name, value) in config.default_headers() { - if !has_header(&headers, name) { - headers.push((name.to_string(), value.to_string())); - } - } - - Ok(headers) +fn invalid_request(err: serde_json::Error) -> Error { + Error::InvalidRequest(format!("invalid Anthropic messages request: {err}")) +} + +fn without_additional_drop_params( + request: AnthropicMessagesRequest, + paths: &[String], +) -> Result { + if paths.is_empty() { + return Ok(request); + } + let Value::Object(fields) = serde_json::to_value(request).map_err(invalid_request)? else { + return Err(Error::InvalidRequest( + "Anthropic messages request did not serialize to an object".to_string(), + )); + }; + let (required, optional): (Map, Map) = fields + .into_iter() + .partition(|(key, _)| matches!(key.as_str(), "model" | "messages")); + let trimmed = paths.iter().fold(Value::Object(optional), |body, path| { + delete_nested_value(body, path) + }); + let merged: Map = required + .into_iter() + .chain(trimmed.as_object().cloned().unwrap_or_default()) + .collect(); + serde_json::from_value(Value::Object(merged)).map_err(invalid_request) +} + +fn with_default_headers( + headers: Vec<(String, String)>, + defaults: &[(&str, &str)], +) -> Vec<(String, String)> { + let missing: Vec<(String, String)> = defaults + .iter() + .filter(|(name, _)| { + !headers + .iter() + .any(|(header, _)| header.eq_ignore_ascii_case(name)) + }) + .map(|(name, value)| ((*name).to_string(), (*value).to_string())) + .collect(); + headers.into_iter().chain(missing).collect() +} + +#[cfg(test)] +mod tests { + use litellm_types::utils::ProviderSpecificHeaders; + use rstest::{fixture, rstest}; + use serde_json::json; + + use super::*; + use crate::messages::types::MessagesShaping; + + #[fixture] + fn shaping() -> MessagesShaping { + MessagesShaping::default() + } + + fn prepare(request: MessagesRequest<'_>) -> Result { + prepare_with_secrets(request, &|_: &str| None) + } + + fn prepare_with_secrets( + request: MessagesRequest<'_>, + secrets: &dyn Lookup, + ) -> Result { + let resolved = resolve_provider(request.model, request.custom_llm_provider)?; + prepare_provider_request(request, resolved, secrets) + } + + #[rstest] + #[case::api_key( + &[("ANTHROPIC_API_KEY", "sk-secret")], + &[("x-api-key", "sk-secret")], + "https://api.anthropic.com/v1/messages" + )] + #[case::auth_token( + &[("ANTHROPIC_AUTH_TOKEN", "token")], + &[("authorization", "Bearer token")], + "https://api.anthropic.com/v1/messages" + )] + #[case::api_base( + &[("ANTHROPIC_API_KEY", "sk-secret"), ("ANTHROPIC_API_BASE", "https://gateway.test")], + &[("x-api-key", "sk-secret")], + "https://gateway.test/v1/messages" + )] + #[case::sdk_base_url( + &[("ANTHROPIC_API_KEY", "sk-secret"), ("ANTHROPIC_BASE_URL", "https://sdk.test")], + &[("x-api-key", "sk-secret")], + "https://sdk.test/v1/messages" + )] + fn credentials_and_base_come_from_the_resolved_secrets( + shaping: MessagesShaping, + #[case] secrets: &[(&str, &str)], + #[case] expected_auth: &[(&str, &str)], + #[case] expected_url: &str, + ) { + let lookup = |name: &str| { + secrets + .iter() + .find(|(key, _)| *key == name) + .map(|(_, value)| value.to_string()) + }; + let prepared = prepare_with_secrets( + MessagesRequest { + model: "claude-test", + body: json!({"model": "claude-test", "messages": [{"role": "user", "content": "hi"}], "max_tokens": 16}), + api_key: None, + api_base: None, + custom_llm_provider: Some("anthropic"), + extra_headers: None, + provider_specific_header: None, + timeout: None, + shaping, + }, + &lookup, + ) + .unwrap(); + let auth: Vec<(&str, &str)> = prepared + .upstream_headers + .iter() + .filter(|(name, _)| matches!(name.as_str(), "x-api-key" | "authorization")) + .map(|(name, value)| (name.as_str(), value.as_str())) + .collect(); + assert_eq!( + (auth.as_slice(), prepared.url.as_str()), + (expected_auth, expected_url) + ); + } + + fn prepared_body(body: Value, shaping: MessagesShaping) -> Result { + prepare(MessagesRequest { + model: "anthropic/claude-test", + body, + api_key: Some("sk-test"), + api_base: Some("https://anthropic.test"), + custom_llm_provider: Some("anthropic"), + extra_headers: None, + provider_specific_header: None, + timeout: None, + shaping, + }) + .map(|prepared| prepared.body) + } + + #[rstest] + #[case::nothing_forwarded( + &[], + &[("x-version", "1"), ("content-type", "application/json")], + &[("x-version", "1"), ("content-type", "application/json")], + )] + #[case::forwarded_header_wins_in_any_case( + &[("X-Version", "custom"), ("x-api-key", "k")], + &[("x-version", "1"), ("content-type", "application/json")], + &[("X-Version", "custom"), ("x-api-key", "k"), ("content-type", "application/json")], + )] + #[case::no_defaults(&[("x-api-key", "k")], &[], &[("x-api-key", "k")])] + fn default_headers_fill_only_missing_names( + #[case] forwarded: &[(&str, &str)], + #[case] defaults: &[(&str, &str)], + #[case] expected: &[(&str, &str)], + ) { + let owned = |headers: &[(&str, &str)]| -> Vec<(String, String)> { + headers + .iter() + .map(|(name, value)| ((*name).to_string(), (*value).to_string())) + .collect() + }; + assert_eq!( + with_default_headers(owned(forwarded), defaults), + owned(expected) + ); + } + + #[rstest] + #[case::top_level_and_nested_paths( + json!({ + "max_tokens": 1024, + "thinking": {"type": "enabled", "budget_tokens": 2048}, + "context_management": {"edits": [{"type": "clear_thinking_20251015"}]}, + "metadata": {"user_id": "u1"}, + "tools": [{"name": "lookup", "input_schema": {"type": "object"}, "input_examples": [{"q": "x"}]}] + }), + &["thinking", "context_management", "tools[*].input_examples"], + json!({ + "max_tokens": 1024, + "metadata": {"user_id": "u1"}, + "tools": [{"name": "lookup", "input_schema": {"type": "object"}}] + }), + )] + #[case::no_paths( + json!({"max_tokens": 16, "safeguards": [{"type": "dangerous_tool_use"}]}), + &[], + json!({"max_tokens": 16, "safeguards": [{"type": "dangerous_tool_use"}]}), + )] + #[case::model_and_messages_are_never_dropped( + json!({"max_tokens": 16}), + &["model", "messages", "messages[0].content"], + json!({"max_tokens": 16}), + )] + fn prepared_body_drops_configured_paths( + shaping: MessagesShaping, + #[case] fields: Value, + #[case] additional_drop_params: &[&str], + #[case] expected_fields: Value, + ) { + let with_messages = |fields: Value| -> Value { + let Value::Object(fields) = fields else { + unreachable!() + }; + Value::Object( + [ + ("model".to_string(), json!("claude-test")), + ( + "messages".to_string(), + json!([{"role": "user", "content": "hi"}]), + ), + ] + .into_iter() + .chain(fields) + .collect(), + ) + }; + let shaping = MessagesShaping { + additional_drop_params: additional_drop_params + .iter() + .map(ToString::to_string) + .collect(), + ..shaping + }; + assert_eq!( + prepared_body(with_messages(fields), shaping), + Ok(with_messages(expected_fields)) + ); + } + + #[rstest] + #[case::model_prefix_picks_the_provider( + "azure_ai/claude-test", + None, + &[("x-priority", "extra"), ("x-scoped", "azure_ai")] + )] + #[case::explicit_provider( + "claude-test", + Some("anthropic"), + &[("x-priority", "scoped"), ("x-scoped", "anthropic")] + )] + #[case::provider_prefix_on_an_anthropic_model( + "anthropic/claude-test", + None, + &[("x-priority", "scoped"), ("x-scoped", "anthropic")] + )] + fn provider_specific_headers_follow_the_resolved_provider( + shaping: MessagesShaping, + #[case] model: &str, + #[case] custom_llm_provider: Option<&str>, + #[case] expected: &[(&str, &str)], + ) { + let configured: ProviderSpecificHeaders = serde_json::from_value(json!([ + {"custom_llm_provider": "azure_ai", "extra_headers": {"x-scoped": "azure_ai"}}, + {"custom_llm_provider": "anthropic", "extra_headers": {"x-scoped": "anthropic", "x-priority": "scoped"}} + ])) + .unwrap(); + let prepared = prepare(MessagesRequest { + model, + body: json!({"model": model, "messages": [{"role": "user", "content": "hi"}], "max_tokens": 16}), + api_key: Some("sk-test"), + api_base: Some("https://resource.services.ai.azure.com"), + custom_llm_provider, + extra_headers: Some(serde_json::from_value(json!({"x-priority": "extra"})).unwrap()), + provider_specific_header: Some(configured), + timeout: None, + shaping, + }) + .unwrap(); + let caller_headers: Vec<(&str, &str)> = prepared + .upstream_headers + .iter() + .filter(|(name, _)| matches!(name.as_str(), "x-priority" | "x-scoped")) + .map(|(name, value)| (name.as_str(), value.as_str())) + .collect(); + assert_eq!(caller_headers, expected); + } + + #[rstest] + fn prepared_body_carries_the_provider_stripped_model(shaping: MessagesShaping) { + assert_eq!( + prepared_body( + json!({ + "model": "anthropic/claude-test", + "messages": [{"role": "user", "content": "hi"}], + "max_tokens": 16 + }), + shaping, + ), + Ok(json!({ + "model": "claude-test", + "messages": [{"role": "user", "content": "hi"}], + "max_tokens": 16 + })) + ); + } + + #[rstest] + fn dropped_thinking_display_is_not_restored_by_auto_summary(shaping: MessagesShaping) { + let shaping = MessagesShaping { + reasoning_auto_summary: true, + additional_drop_params: vec!["thinking.display".to_string()], + ..shaping + }; + assert_eq!( + prepared_body( + json!({ + "model": "claude-test", + "messages": [{"role": "user", "content": "hi"}], + "max_tokens": 4096, + "thinking": {"type": "enabled", "budget_tokens": 2048} + }), + shaping, + ), + Ok(json!({ + "model": "claude-test", + "messages": [{"role": "user", "content": "hi"}], + "max_tokens": 4096, + "thinking": {"type": "enabled", "budget_tokens": 2048} + })) + ); + } + + #[rstest] + fn dropping_an_invalid_metadata_user_id_does_not_skip_its_validation(shaping: MessagesShaping) { + let shaping = MessagesShaping { + additional_drop_params: vec!["metadata.user_id".to_string()], + ..shaping + }; + assert!(matches!( + prepared_body( + json!({ + "model": "claude-test", + "messages": [{"role": "user", "content": "hi"}], + "max_tokens": 16, + "metadata": {"user_id": 123} + }), + shaping, + ), + Err(Error::InvalidRequest(_)) + )); + } + + #[rstest] + fn prepared_body_rejects_invalid_metadata_before_the_call(shaping: MessagesShaping) { + assert_eq!( + prepared_body( + json!({ + "model": "claude-test", + "messages": [{"role": "user", "content": "hi"}], + "max_tokens": 16, + "metadata": {"user_id": 123} + }), + shaping, + ), + Err(Error::InvalidRequest( + "metadata.user_id must be a string, got 123".to_string() + )) + ); + } } diff --git a/litellm-rust/crates/core/src/messages/route.rs b/litellm-rust/crates/core/src/messages/route.rs index 838b56fcb4b..8cd3eaf3aa3 100644 --- a/litellm-rust/crates/core/src/messages/route.rs +++ b/litellm-rust/crates/core/src/messages/route.rs @@ -1,4 +1,7 @@ -use std::{sync::Mutex, time::Duration}; +use std::{ + sync::{Arc, Mutex}, + time::Duration, +}; use bytes::Bytes; use litellm_auth::SecretValue; @@ -9,15 +12,19 @@ use litellm_host::{ machine::{HostChannel, MachineFault, RouteMachine}, route::Route, }; -use litellm_types::llms::anthropic_messages::anthropic_response::AnthropicMessagesResponse; +use litellm_secrets::source::SecretSource; +use litellm_types::{ + llms::anthropic_messages::anthropic_response::AnthropicMessagesResponse, + utils::ProviderSpecificHeaders, +}; use serde_json::{Map, Value}; use super::{ Error, common_utils::messages_provider_config, handler::{decode_response, network, provider_error, send}, - prepare::prepare_provider_request, - types::MessagesRequest, + prepare::{prepare_provider_request, resolve_provider}, + types::{MessagesRequest, MessagesShaping}, }; use crate::constants::ANTHROPIC_MESSAGES_PROVIDER; @@ -38,7 +45,9 @@ pub struct MessagesCall { pub api_base: Option, pub custom_llm_provider: Option, pub extra_headers: Option>, + pub provider_specific_header: Option, pub timeout: Option, + pub shaping: MessagesShaping, } impl MessagesCall { @@ -120,22 +129,33 @@ impl Host for LocalMessagesHost { } } -pub fn messages_machine() -> MessagesMachine { - RouteMachine::new(|host| Box::pin(execute(host))) +pub fn messages_machine(secrets: Arc) -> MessagesMachine { + RouteMachine::new(move |host| Box::pin(execute(host, secrets.clone()))) } -async fn execute(host: MessagesHost) -> Result { +async fn execute( + host: MessagesHost, + secrets: Arc, +) -> Result { let MessagesOpResult::Request(call) = host.route(MessagesOp::ProjectRequest).await?; let stream = call.streams(); - let request = prepare_provider_request(MessagesRequest { - model: &call.model, - body: Value::Object(call.body.clone()), - api_key: call.api_key.as_deref(), - api_base: call.api_base.as_deref(), - custom_llm_provider: call.custom_llm_provider.as_deref(), - extra_headers: call.extra_headers.clone(), - timeout: call.timeout, - })?; + let resolved = resolve_provider(&call.model, call.custom_llm_provider.as_deref())?; + let secrets = secrets.resolve(resolved.config.secret_names()).await?; + let request = prepare_provider_request( + MessagesRequest { + model: &call.model, + body: Value::Object(call.body.clone()), + api_key: call.api_key.as_deref(), + api_base: call.api_base.as_deref(), + custom_llm_provider: call.custom_llm_provider.as_deref(), + extra_headers: call.extra_headers.clone(), + provider_specific_header: call.provider_specific_header.clone(), + timeout: call.timeout, + shaping: call.shaping.clone(), + }, + resolved, + secrets.as_ref(), + )?; if stream && request.provider != ANTHROPIC_MESSAGES_PROVIDER { return Err(Error::Unsupported("streaming messages for this provider")); } diff --git a/litellm-rust/crates/core/src/messages/tests.rs b/litellm-rust/crates/core/src/messages/tests.rs index 057b42a316c..ce48752864a 100644 --- a/litellm-rust/crates/core/src/messages/tests.rs +++ b/litellm-rust/crates/core/src/messages/tests.rs @@ -1,5 +1,8 @@ -use std::time::Duration; +use std::{sync::Arc, time::Duration}; +use futures_util::future::BoxFuture; +use litellm_http::request::{has_bearer_auth, has_header}; +use litellm_secrets::{SecretValue, source::SecretSource}; use serde_json::{Map, Value, json}; use tokio::{ io::{AsyncReadExt, AsyncWriteExt}, @@ -8,12 +11,132 @@ use tokio::{ use super::{ Error, - common_utils::{ - has_bearer_auth, has_header, messages_provider_config, string_headers, truncate_error_body, - }, + common_utils::{messages_provider_config, string_headers, truncate_error_body}, messages, + route::{LocalMessagesHost, MessagesCall, MessagesOutput, messages_machine}, }; -use crate::messages::types::MessagesRequest; +use crate::messages::types::{MessagesRequest, MessagesShaping}; + +struct RecordingSecrets { + values: Vec<(&'static str, String)>, + fails: bool, + requested: std::sync::Mutex>, +} + +impl RecordingSecrets { + fn new(values: Vec<(&'static str, String)>, fails: bool) -> Self { + Self { + values, + fails, + requested: std::sync::Mutex::new(Vec::new()), + } + } +} + +impl SecretSource for RecordingSecrets { + fn get_secret_str<'a>( + &'a self, + name: &'a str, + ) -> BoxFuture<'a, Result, litellm_secrets::Error>> { + Box::pin(async move { + self.requested.lock().unwrap().push(name.to_string()); + if self.fails { + return Err(litellm_secrets::Error::ManagedSecretMissing); + } + Ok(self + .values + .iter() + .find(|(key, _)| *key == name) + .map(|(_, value)| SecretValue::new(value.clone()))) + }) + } +} + +fn secrets_call() -> MessagesCall { + let Value::Object(body) = json!({ + "model": "claude-sonnet-4-5", + "max_tokens": 16, + "messages": [{"role": "user", "content": "hi"}] + }) else { + unreachable!("literal object") + }; + MessagesCall { + model: "claude-sonnet-4-5".into(), + body, + api_key: None, + api_base: None, + custom_llm_provider: Some("anthropic".into()), + extra_headers: None, + provider_specific_header: None, + timeout: Some(Duration::from_secs(5)), + shaping: MessagesShaping::default(), + } +} + +#[tokio::test] +async fn route_reads_the_provider_credential_and_base_from_the_secret_source() { + let listener = TcpListener::bind("127.0.0.1:0").await.expect("binds"); + let addr = listener.local_addr().expect("addr"); + let server = tokio::spawn(async move { + let (mut socket, _) = listener.accept().await.expect("accepts request"); + let request = read_http_request(&mut socket).await; + let response_body = r#"{"id":"msg_1","type":"message","role":"assistant","content":[],"model":"claude-sonnet-4-5","stop_reason":"end_turn","usage":{"input_tokens":1,"output_tokens":1}}"#; + socket + .write_all(write_response(response_body).as_bytes()) + .await + .expect("writes response"); + request + }); + let secrets = Arc::new(RecordingSecrets::new( + vec![ + ("ANTHROPIC_API_KEY", "sk-from-manager".to_string()), + ("ANTHROPIC_BASE_URL", format!("http://{addr}")), + ], + false, + )); + + let output = litellm_host::run::run( + messages_machine(secrets.clone()), + &LocalMessagesHost::new(secrets_call()), + ) + .await + .expect("messages request succeeds"); + + assert!(matches!(output, MessagesOutput::Message(_))); + let request = server.await.expect("server task completes"); + assert!( + request + .to_ascii_lowercase() + .contains("x-api-key: sk-from-manager"), + "{request}" + ); + let requested = secrets.requested.lock().unwrap().clone(); + assert_eq!( + requested, + messages_provider_config("anthropic") + .unwrap() + .secret_names() + .iter() + .map(ToString::to_string) + .collect::>() + ); +} + +#[tokio::test] +async fn route_surfaces_a_secret_manager_failure_before_the_call() { + let Err(error) = litellm_host::run::run( + messages_machine(Arc::new(RecordingSecrets::new(Vec::new(), true))), + &LocalMessagesHost::new(secrets_call()), + ) + .await + else { + panic!("a secret manager failure fails the call"); + }; + assert!( + matches!(&error, Error::Secret(source) if matches!(source.source_error(), litellm_secrets::Error::ManagedSecretMissing)), + "{error:?}" + ); +} async fn read_http_request(socket: &mut TcpStream) -> String { let mut request = Vec::new(); @@ -159,7 +282,9 @@ async fn messages_round_trip_builds_azure_request_and_passes_response_through() api_base: Some(&format!("http://{addr}")), custom_llm_provider: Some("azure_ai"), extra_headers: None, + provider_specific_header: None, timeout: Some(Duration::from_secs(5)), + shaping: MessagesShaping::default(), }) .await .expect("messages request succeeds"); @@ -215,7 +340,9 @@ async fn messages_round_trip_builds_native_anthropic_request() { api_base: Some(&format!("http://{addr}")), custom_llm_provider: Some("anthropic"), extra_headers: None, + provider_specific_header: None, timeout: Some(Duration::from_secs(5)), + shaping: MessagesShaping::default(), }) .await .expect("messages request succeeds"); @@ -268,7 +395,9 @@ async fn messages_does_not_duplicate_auth_when_x_api_key_supplied() { api_base: Some(&format!("http://{addr}")), custom_llm_provider: Some("azure_ai"), extra_headers: Some(headers), + provider_specific_header: None, timeout: Some(Duration::from_secs(5)), + shaping: MessagesShaping::default(), }) .await .expect("messages request succeeds"); @@ -322,7 +451,9 @@ async fn messages_forwards_entra_id_bearer_without_requiring_api_key() { api_base: Some(&format!("http://{addr}")), custom_llm_provider: Some("azure_ai"), extra_headers: Some(headers), + provider_specific_header: None, timeout: Some(Duration::from_secs(5)), + shaping: MessagesShaping::default(), }) .await .expect("entra id request succeeds without api key"); @@ -346,7 +477,9 @@ async fn messages_requires_auth_when_no_key_and_no_header() { api_base: Some("http://127.0.0.1:1"), custom_llm_provider: Some("azure_ai"), extra_headers: None, + provider_specific_header: None, timeout: Some(Duration::from_millis(50)), + shaping: MessagesShaping::default(), }) .await .expect_err("missing auth errors"); @@ -384,7 +517,9 @@ async fn messages_ignores_malformed_authorization_and_uses_api_key() { api_base: Some(&format!("http://{addr}")), custom_llm_provider: Some("azure_ai"), extra_headers: Some(headers), + provider_specific_header: None, timeout: Some(Duration::from_secs(5)), + shaping: MessagesShaping::default(), }) .await .expect("falls back to api key"); @@ -425,7 +560,9 @@ async fn messages_maps_provider_error_status_to_http_error() { api_base: Some(&format!("http://{addr}")), custom_llm_provider: Some("azure_ai"), extra_headers: None, + provider_specific_header: None, timeout: Some(Duration::from_secs(5)), + shaping: MessagesShaping::default(), }) .await .expect_err("provider error propagates"); @@ -445,7 +582,9 @@ async fn messages_rejects_unsupported_provider() { api_base: Some("http://127.0.0.1:1"), custom_llm_provider: Some("openai"), extra_headers: None, + provider_specific_header: None, timeout: Some(Duration::from_millis(50)), + shaping: MessagesShaping::default(), }) .await .expect_err("unsupported provider errors"); diff --git a/litellm-rust/crates/core/src/messages/types.rs b/litellm-rust/crates/core/src/messages/types.rs index a73ceffad7a..4a5dd2926e0 100644 --- a/litellm-rust/crates/core/src/messages/types.rs +++ b/litellm-rust/crates/core/src/messages/types.rs @@ -1,8 +1,25 @@ use std::time::Duration; -use litellm_llms::base_llm::anthropic_messages::transformation::BaseAnthropicMessagesConfig; +use litellm_llms::{ + anthropic::common_utils::AnthropicModelCapabilities, + base_llm::anthropic_messages::transformation::BaseAnthropicMessagesConfig, +}; +use litellm_types::utils::ProviderSpecificHeaders; +use serde::{Deserialize, Serialize}; use serde_json::{Map, Value}; +#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)] +pub struct MessagesShaping { + #[serde(default)] + pub capabilities: AnthropicModelCapabilities, + #[serde(default)] + pub drop_params: bool, + #[serde(default)] + pub reasoning_auto_summary: bool, + #[serde(default)] + pub additional_drop_params: Vec, +} + pub struct MessagesRequest<'a> { pub model: &'a str, pub body: Value, @@ -10,7 +27,9 @@ pub struct MessagesRequest<'a> { pub api_base: Option<&'a str>, pub custom_llm_provider: Option<&'a str>, pub extra_headers: Option>, + pub provider_specific_header: Option, pub timeout: Option, + pub shaping: MessagesShaping, } pub struct ProviderMessagesRequest { @@ -22,3 +41,86 @@ pub struct ProviderMessagesRequest { pub upstream_headers: Vec<(String, String)>, pub timeout: Option, } + +#[cfg(test)] +mod tests { + use litellm_llms::anthropic::common_utils::SupportedEffortTiers; + use rstest::rstest; + use serde_json::json; + + use super::*; + + #[rstest] + #[case::nothing_projected(json!({}), MessagesShaping::default())] + #[case::only_drop_params( + json!({"drop_params": true}), + MessagesShaping { drop_params: true, ..MessagesShaping::default() }, + )] + #[case::only_reasoning_auto_summary( + json!({"reasoning_auto_summary": true}), + MessagesShaping { reasoning_auto_summary: true, ..MessagesShaping::default() }, + )] + #[case::only_additional_drop_params( + json!({"additional_drop_params": ["tools[*].input_examples"]}), + MessagesShaping { + additional_drop_params: vec!["tools[*].input_examples".to_string()], + ..MessagesShaping::default() + }, + )] + #[case::partial_capabilities( + json!({"capabilities": {"supports_reasoning": true}}), + MessagesShaping { + capabilities: AnthropicModelCapabilities { + supports_reasoning: true, + ..AnthropicModelCapabilities::default() + }, + ..MessagesShaping::default() + }, + )] + #[case::everything_the_python_host_projects( + json!({ + "capabilities": { + "supports_reasoning": true, + "supports_adaptive_thinking": true, + "thinking_always_on": false, + "supports_legacy_thinking": false, + "supports_output_config": true, + "supports_sampling_params": false, + "supports_speed": true, + "effort_tiers": {"minimal": false, "low": true, "medium": true, "high": true, "xhigh": true, "max": false} + }, + "drop_params": true, + "reasoning_auto_summary": true, + "additional_drop_params": ["metadata.user_id", "thinking"] + }), + MessagesShaping { + capabilities: AnthropicModelCapabilities { + supports_reasoning: true, + supports_adaptive_thinking: true, + thinking_always_on: false, + supports_legacy_thinking: false, + supports_output_config: true, + supports_sampling_params: false, + supports_speed: true, + effort_tiers: SupportedEffortTiers { + minimal: false, + low: true, + medium: true, + high: true, + xhigh: true, + max: false, + }, + }, + drop_params: true, + reasoning_auto_summary: true, + additional_drop_params: vec!["metadata.user_id".to_string(), "thinking".to_string()], + }, + )] + fn shaping_deserializes_with_defaults_for_absent_fields( + #[case] projected: Value, + #[case] expected: MessagesShaping, + ) { + let shaping: MessagesShaping = serde_json::from_value(projected).unwrap(); + assert_eq!(shaping, expected); + } +} diff --git a/litellm-rust/crates/llms/src/anthropic/common_utils.rs b/litellm-rust/crates/llms/src/anthropic/common_utils.rs new file mode 100644 index 00000000000..a2234e0df03 --- /dev/null +++ b/litellm-rust/crates/llms/src/anthropic/common_utils.rs @@ -0,0 +1,1568 @@ +use litellm_types::llms::anthropic_messages::anthropic_request::{ + AnthropicMessage, ContentBlock, MessageContent, +}; +use serde::{Deserialize, Serialize}; +use serde_json::Value; + +use crate::anthropic::ANTHROPIC_OAUTH_TOKEN_PREFIX; + +pub const ANTHROPIC_OAUTH_BETA_HEADER: &str = "oauth-2025-04-20"; +pub const ANTHROPIC_ADVISOR_TOOL_TYPE: &str = "advisor_20260301"; +pub const ANTHROPIC_TOOL_SEARCH_TOOL_TYPES: [&str; 2] = [ + "tool_search_tool_regex_20251119", + "tool_search_tool_bm25_20251119", +]; +pub const ENCRYPTED_REASONING_SIGNATURE_PREFIX: &str = "litellm_encrypted_reasoning:"; +const THOUGHT_SIGNATURE_SEPARATOR: &str = "__thought__"; + +pub mod beta { + pub const CONTEXT_MANAGEMENT_2025_06_27: &str = "context-management-2025-06-27"; + pub const COMPACT_2026_01_12: &str = "compact-2026-01-12"; + pub const COMPACT_2026_09_04: &str = "compact-2026-09-04"; + pub const STRUCTURED_OUTPUT: &str = "structured-outputs-2025-11-13"; + pub const ADVANCED_TOOL_USE_2025_11_20: &str = "advanced-tool-use-2025-11-20"; + pub const FAST_MODE_2026_02_01: &str = "fast-mode-2026-02-01"; + pub const ADVISOR_TOOL_2026_03_01: &str = "advisor-tool-2026-03-01"; + pub const PER_TURN_CONTROL_2026_07_01: &str = "per-turn-control-2026-07-01"; +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)] +#[serde(rename_all = "lowercase")] +pub enum EffortLevel { + Low, + Medium, + High, + Xhigh, + Max, +} + +impl EffortLevel { + pub fn as_str(self) -> &'static str { + match self { + Self::Low => "low", + Self::Medium => "medium", + Self::High => "high", + Self::Xhigh => "xhigh", + Self::Max => "max", + } + } + + pub fn parse(value: &str) -> Option { + match value { + "low" => Some(Self::Low), + "medium" => Some(Self::Medium), + "high" => Some(Self::High), + "xhigh" => Some(Self::Xhigh), + "max" => Some(Self::Max), + _ => None, + } + } +} + +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Serialize, Deserialize)] +pub struct SupportedEffortTiers { + #[serde(default)] + pub minimal: bool, + #[serde(default)] + pub low: bool, + #[serde(default)] + pub medium: bool, + #[serde(default)] + pub high: bool, + #[serde(default)] + pub xhigh: bool, + #[serde(default)] + pub max: bool, +} + +impl SupportedEffortTiers { + pub fn any(self) -> bool { + self.minimal || self.low || self.medium || self.high || self.xhigh || self.max + } +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)] +pub struct AnthropicModelCapabilities { + #[serde(default)] + pub supports_reasoning: bool, + #[serde(default)] + pub supports_adaptive_thinking: bool, + #[serde(default)] + pub thinking_always_on: bool, + #[serde(default)] + pub supports_legacy_thinking: bool, + #[serde(default)] + pub supports_output_config: bool, + #[serde(default = "default_true")] + pub supports_sampling_params: bool, + #[serde(default)] + pub supports_speed: bool, + #[serde(default)] + pub effort_tiers: SupportedEffortTiers, +} + +fn default_true() -> bool { + true +} + +impl Default for AnthropicModelCapabilities { + fn default() -> Self { + Self { + supports_reasoning: false, + supports_adaptive_thinking: false, + thinking_always_on: false, + supports_legacy_thinking: false, + supports_output_config: false, + supports_sampling_params: true, + supports_speed: false, + effort_tiers: SupportedEffortTiers::default(), + } + } +} + +impl AnthropicModelCapabilities { + pub fn supports_effort_tier(&self, level: EffortLevel) -> bool { + match level { + EffortLevel::Low => self.effort_tiers.low, + EffortLevel::Medium => self.effort_tiers.medium, + EffortLevel::High => self.effort_tiers.high, + EffortLevel::Xhigh => self.effort_tiers.xhigh, + EffortLevel::Max => self.effort_tiers.max, + } + } + + pub fn supports_effort_param(&self) -> bool { + self.supports_output_config || self.effort_tiers.any() + } + + pub fn effort_level_rejection(&self, effort: &str, model: &str) -> Option { + match effort { + "max" if !(self.supports_adaptive_thinking || self.effort_tiers.max) => Some(format!( + "effort='max' is not supported by this model. Got model: {model}" + )), + "xhigh" if !self.effort_tiers.xhigh => Some(format!( + "effort='xhigh' is not supported by this model. Got model: {model}" + )), + _ => None, + } + } +} + +pub fn is_anthropic_oauth_key(value: &str) -> bool { + value + .strip_prefix("Bearer ") + .unwrap_or(value) + .starts_with(ANTHROPIC_OAUTH_TOKEN_PREFIX) +} + +pub fn split_beta_values(header: Option<&str>) -> impl Iterator + '_ { + header + .into_iter() + .flat_map(|value| value.split(',')) + .map(str::trim) + .filter(|piece| !piece.is_empty()) + .map(str::to_string) +} + +pub fn join_beta_values(values: impl IntoIterator) -> String { + let mut values: Vec = values.into_iter().collect(); + values.sort(); + values.dedup(); + values.join(",") +} + +pub fn is_tool_search_used(tools: Option<&[Value]>) -> bool { + tools.into_iter().flatten().any(|tool| { + tool.get("type") + .and_then(Value::as_str) + .is_some_and(|tool_type| ANTHROPIC_TOOL_SEARCH_TOOL_TYPES.contains(&tool_type)) + }) +} + +pub fn has_advisor_tool(tools: Option<&[Value]>) -> bool { + tools + .into_iter() + .flatten() + .any(|tool| tool.get("type").and_then(Value::as_str) == Some(ANTHROPIC_ADVISOR_TOOL_TYPE)) +} + +pub fn requires_native_compaction_beta( + compaction: Option<&Value>, + messages: &[AnthropicMessage], +) -> bool { + compaction.is_some() + || messages + .iter() + .flat_map(AnthropicMessage::blocks) + .any(|block| { + block.is_type("compaction") + && block.signature.as_deref().is_some_and(|s| !s.is_empty()) + }) +} + +fn is_blank(text: Option<&str>) -> bool { + text.is_none_or(|text| text.trim().is_empty()) +} + +fn is_empty_text_block(block: &ContentBlock) -> bool { + block.is_type("text") && is_blank(block.text.as_deref()) +} + +pub fn is_empty_thinking_block(block: &ContentBlock) -> bool { + block.is_type("thinking") && is_blank(block.thinking.as_deref()) +} + +fn retain_blocks( + messages: Vec, + keep: impl Fn(&ContentBlock) -> bool, +) -> Vec { + messages + .into_iter() + .filter_map(|message| match message.content { + MessageContent::Text(_) => Some(message), + MessageContent::Blocks(ref blocks) => { + let kept: Vec = + blocks.iter().filter(|block| keep(block)).cloned().collect(); + if kept.len() == blocks.len() { + return Some(message); + } + (!kept.is_empty()).then(|| message.with_blocks(kept)) + } + }) + .collect() +} + +pub fn strip_empty_content_blocks(messages: Vec) -> Vec { + retain_blocks(messages, |block| { + !is_empty_text_block(block) && !is_empty_thinking_block(block) + }) +} + +pub fn normalize_anthropic_tool_use_id(raw_id: &str) -> String { + let base = raw_id + .split_once(THOUGHT_SIGNATURE_SEPARATOR) + .map_or(raw_id, |(base, _)| base); + let sanitized: String = base + .chars() + .map(|character| { + if character.is_ascii_alphanumeric() || matches!(character, '_' | '-') { + character + } else { + '_' + } + }) + .collect(); + if sanitized.is_empty() { + "tool_use_id".to_string() + } else { + sanitized + } +} + +fn normalized_if_changed(raw_id: Option<&str>) -> Option { + let raw_id = raw_id?; + let normalized = normalize_anthropic_tool_use_id(raw_id); + (normalized != raw_id).then_some(normalized) +} + +fn sanitize_tool_use_id_block(block: ContentBlock) -> ContentBlock { + match block.block_type.as_deref() { + Some("tool_use" | "server_tool_use") => match normalized_if_changed(block.id.as_deref()) { + Some(id) => ContentBlock { + id: Some(id), + ..block + }, + None => block, + }, + Some("tool_result") => match normalized_if_changed(block.tool_use_id.as_deref()) { + Some(tool_use_id) => ContentBlock { + tool_use_id: Some(tool_use_id), + ..block + }, + None => block, + }, + _ => block, + } +} + +fn map_blocks( + messages: Vec, + rewrite: impl Fn(Vec) -> Vec, +) -> Vec { + messages + .into_iter() + .map(|message| match message.content { + MessageContent::Blocks(blocks) => AnthropicMessage { + content: MessageContent::Blocks(rewrite(blocks)), + ..message + }, + MessageContent::Text(_) => message, + }) + .collect() +} + +pub fn sanitize_tool_use_ids(messages: Vec) -> Vec { + map_blocks(messages, |blocks| { + blocks.into_iter().map(sanitize_tool_use_id_block).collect() + }) +} + +pub fn strip_provider_specific_fields(messages: Vec) -> Vec { + map_blocks(messages, |blocks| { + blocks + .into_iter() + .map(|block| ContentBlock { + provider_specific_fields: None, + ..block + }) + .collect() + }) +} + +pub fn is_encrypted_reasoning_block(block: &ContentBlock) -> bool { + let field = match block.block_type.as_deref() { + Some("thinking") => block.signature.as_deref(), + Some("redacted_thinking") => block.data.as_deref(), + _ => None, + }; + field.is_some_and(|value| value.starts_with(ENCRYPTED_REASONING_SIGNATURE_PREFIX)) +} + +pub fn strip_encrypted_reasoning_blocks(messages: Vec) -> Vec { + retain_blocks(messages, |block| !is_encrypted_reasoning_block(block)) +} + +fn is_advisor_use(block: &ContentBlock) -> bool { + block.is_type("server_tool_use") + && block.name.as_deref() == Some("advisor") + && block.id.as_deref().is_some_and(|id| !id.is_empty()) +} + +pub fn strip_advisor_blocks(messages: Vec) -> Vec { + messages + .into_iter() + .map(|message| { + if message.role != "assistant" { + return message; + } + let MessageContent::Blocks(blocks) = &message.content else { + return message; + }; + let advisor_ids: Vec<&str> = blocks + .iter() + .filter(|block| is_advisor_use(block)) + .filter_map(|block| block.id.as_deref()) + .collect(); + if advisor_ids.is_empty() { + return message; + } + let kept: Vec = blocks + .iter() + .filter(|block| { + let is_result = block.is_type("advisor_tool_result") + && block + .tool_use_id + .as_deref() + .is_some_and(|id| advisor_ids.contains(&id)); + !is_advisor_use(block) && !is_result + }) + .cloned() + .collect(); + message.with_blocks(kept) + }) + .collect() +} + +#[derive(Deserialize)] +struct ReplayedWebSearchResult { + #[serde(default)] + url: String, + #[serde(default)] + title: String, + #[serde(default)] + snippet: String, + #[serde(default)] + encrypted_content: String, +} + +#[derive(Deserialize)] +#[serde(tag = "type")] +enum ReplayedWebSearchContent { + #[serde(rename = "web_search_tool_result_error")] + Error { + #[serde(default)] + error_code: String, + }, +} + +enum WebSearchResults { + Results(Vec), + Error(String), +} + +fn flattenable_web_search_results(block: &ContentBlock) -> Option<(&str, WebSearchResults)> { + if !block.is_type("web_search_tool_result") { + return None; + } + let tool_use_id = block.tool_use_id.as_deref()?; + let results = match block.content.as_ref()? { + Value::Array(items) => { + let results = items + .iter() + .map(|item| { + (item.get("type").and_then(Value::as_str) == Some("web_search_result")) + .then(|| { + serde_json::from_value::(item.clone()).ok() + }) + .flatten() + }) + .collect::>>()?; + if results + .iter() + .any(|result| !result.encrypted_content.is_empty()) + { + return None; + } + WebSearchResults::Results(results) + } + error @ Value::Object(_) => match serde_json::from_value(error.clone()).ok()? { + ReplayedWebSearchContent::Error { error_code } => WebSearchResults::Error(error_code), + }, + _ => return None, + }; + Some((tool_use_id, results)) +} + +fn render_web_search_results(query: &str, results: &WebSearchResults) -> String { + let header = if query.is_empty() { + "Web search results:".to_string() + } else { + format!("Web search results for '{query}':") + }; + match results { + WebSearchResults::Error(code) => { + let code = if code.is_empty() { "unavailable" } else { code }; + format!("{header}\n\nSearch failed: {code}") + } + WebSearchResults::Results(results) if results.is_empty() => { + format!("{header}\n\nNo results were returned.") + } + WebSearchResults::Results(results) => { + let body = results + .iter() + .map(|result| { + [ + (!result.title.is_empty()).then(|| format!("Title: {}", result.title)), + (!result.url.is_empty()).then(|| format!("URL: {}", result.url)), + (!result.snippet.is_empty()) + .then(|| format!("Snippet: {}", result.snippet)), + ] + .into_iter() + .flatten() + .collect::>() + .join("\n") + }) + .collect::>() + .join("\n\n"); + if body.is_empty() { + header + } else { + format!("{header}\n\n{body}") + } + } + } +} + +fn server_tool_use_query(block: &ContentBlock) -> Option<(&str, &str)> { + if !block.is_type("server_tool_use") { + return None; + } + let id = block.id.as_deref()?; + let query = match block.input.as_ref() { + None => "", + Some(Value::Object(input)) => match input.get("query") { + None => "", + Some(query) => query.as_str()?, + }, + Some(_) => return None, + }; + Some((id, query)) +} + +fn flatten_web_search_results_in_blocks(blocks: Vec) -> Vec { + let flattenable_ids: Vec<&str> = blocks + .iter() + .filter_map(flattenable_web_search_results) + .map(|(tool_use_id, _)| tool_use_id) + .collect(); + if flattenable_ids.is_empty() { + return blocks; + } + let queries: Vec<(&str, &str)> = blocks.iter().filter_map(server_tool_use_query).collect(); + blocks + .iter() + .filter_map(|block| { + if let Some((tool_use_id, results)) = flattenable_web_search_results(block) { + let query = queries + .iter() + .rfind(|(id, _)| *id == tool_use_id) + .map_or("", |(_, query)| query); + return Some(ContentBlock::text(render_web_search_results( + query, &results, + ))); + } + if let Some((id, _)) = server_tool_use_query(block) + && flattenable_ids.contains(&id) + { + return None; + } + Some(block.clone()) + }) + .collect() +} + +pub fn flatten_unencrypted_web_search_results( + messages: Vec, +) -> Vec { + map_blocks(messages, flatten_web_search_results_in_blocks) +} + +#[cfg(test)] +mod tests { + use rstest::{fixture, rstest}; + use serde_json::json; + + use super::*; + + const ALL_LEVELS: [EffortLevel; 5] = [ + EffortLevel::Low, + EffortLevel::Medium, + EffortLevel::High, + EffortLevel::Xhigh, + EffortLevel::Max, + ]; + + fn apply( + sanitizer: fn(Vec) -> Vec, + messages: Value, + ) -> Value { + let parsed: Vec = serde_json::from_value(messages).unwrap(); + serde_json::to_value(sanitizer(parsed)).unwrap() + } + + fn block(value: Value) -> ContentBlock { + serde_json::from_value(value).unwrap() + } + + fn history(messages: Value) -> Vec { + serde_json::from_value(messages).unwrap() + } + + fn tools(value: Option) -> Option> { + value.map(|tools| tools.as_array().unwrap().clone()) + } + + fn tagged(encrypted: &str) -> String { + format!("{ENCRYPTED_REASONING_SIGNATURE_PREFIX}{encrypted}") + } + + fn tiers( + minimal: bool, + low: bool, + medium: bool, + high: bool, + xhigh: bool, + max: bool, + ) -> SupportedEffortTiers { + SupportedEffortTiers { + minimal, + low, + medium, + high, + xhigh, + max, + } + } + + fn replayed_search_turn(results: Value) -> Value { + json!([ + {"role": "user", "content": "when was Rome founded?"}, + {"role": "assistant", "content": [ + {"type": "server_tool_use", "id": "srvtoolu_1", "name": "web_search", "input": {"query": "when"}}, + {"type": "web_search_tool_result", "tool_use_id": "srvtoolu_1", "content": results}, + {"type": "text", "text": "753 BC."} + ]} + ]) + } + + #[fixture] + fn unmapped() -> AnthropicModelCapabilities { + AnthropicModelCapabilities::default() + } + + #[rstest] + #[case::empty_text(json!({"type": "thinking", "thinking": ""}), true)] + #[case::whitespace_only(json!({"type": "thinking", "thinking": " \n\t "}), true)] + #[case::null_text(json!({"type": "thinking", "thinking": null}), true)] + #[case::missing_text(json!({"type": "thinking"}), true)] + #[case::empty_text_despite_signature(json!({"type": "thinking", "thinking": "", "signature": "sig_abc"}), true)] + #[case::real_thinking(json!({"type": "thinking", "thinking": "plan", "signature": "sig"}), false)] + #[case::padded_real_thinking(json!({"type": "thinking", "thinking": " plan "}), false)] + #[case::redacted_thinking_is_a_different_type(json!({"type": "redacted_thinking", "data": "opaque"}), false)] + #[case::empty_text_block(json!({"type": "text", "text": ""}), false)] + #[case::untyped_block(json!({"thinking": ""}), false)] + fn empty_thinking_block_detection(#[case] input: Value, #[case] expected: bool) { + assert_eq!(is_empty_thinking_block(&block(input)), expected); + } + + #[rstest] + #[case::empty_text_beside_tool_use( + json!([{"role": "assistant", "content": [ + {"type": "text", "text": ""}, + {"type": "tool_use", "id": "x", "name": "Bash", "input": {}} + ]}]), + json!([{"role": "assistant", "content": [{"type": "tool_use", "id": "x", "name": "Bash", "input": {}}]}]) + )] + #[case::whitespace_text_beside_tool_use( + json!([{"role": "assistant", "content": [ + {"type": "text", "text": " \n "}, + {"type": "tool_use", "id": "x", "name": "Bash", "input": {}} + ]}]), + json!([{"role": "assistant", "content": [{"type": "tool_use", "id": "x", "name": "Bash", "input": {}}]}]) + )] + #[case::null_text( + json!([{"role": "user", "content": [ + {"type": "text", "text": null}, + {"type": "tool_result", "tool_use_id": "x", "content": "y"} + ]}]), + json!([{"role": "user", "content": [{"type": "tool_result", "tool_use_id": "x", "content": "y"}]}]) + )] + #[case::missing_text( + json!([{"role": "user", "content": [ + {"type": "text"}, + {"type": "tool_result", "tool_use_id": "x", "content": "y"} + ]}]), + json!([{"role": "user", "content": [{"type": "tool_result", "tool_use_id": "x", "content": "y"}]}]) + )] + #[case::empty_signed_thinking_beside_tool_use( + json!([ + {"role": "user", "content": "weather?"}, + {"role": "assistant", "content": [ + {"type": "thinking", "thinking": "", "signature": "sig_abc"}, + {"type": "tool_use", "id": "toolu_01A", "name": "get_weather", "input": {"city": "Paris"}} + ]} + ]), + json!([ + {"role": "user", "content": "weather?"}, + {"role": "assistant", "content": [ + {"type": "tool_use", "id": "toolu_01A", "name": "get_weather", "input": {"city": "Paris"}} + ]} + ]) + )] + #[case::whitespace_thinking_beside_real_and_redacted_thinking( + json!([{"role": "assistant", "content": [ + {"type": "thinking", "thinking": " \n "}, + {"type": "thinking", "thinking": "real plan", "signature": "sig"}, + {"type": "redacted_thinking", "data": "opaque"} + ]}]), + json!([{"role": "assistant", "content": [ + {"type": "thinking", "thinking": "real plan", "signature": "sig"}, + {"type": "redacted_thinking", "data": "opaque"} + ]}]) + )] + #[case::blank_text_beside_real_thinking( + json!([{"role": "assistant", "content": [ + {"type": "thinking", "thinking": "plan", "signature": "sig"}, + {"type": "text", "text": ""} + ]}]), + json!([{"role": "assistant", "content": [{"type": "thinking", "thinking": "plan", "signature": "sig"}]}]) + )] + #[case::message_left_without_blocks_is_dropped( + json!([ + {"role": "user", "content": "hello"}, + {"role": "assistant", "content": [{"type": "text", "text": ""}]}, + {"role": "assistant", "content": [{"type": "thinking", "thinking": ""}]} + ]), + json!([{"role": "user", "content": "hello"}]) + )] + fn strip_empty_content_blocks_rewrites(#[case] input: Value, #[case] expected: Value) { + assert_eq!(apply(strip_empty_content_blocks, input), expected); + } + + #[rstest] + #[case::non_empty_text(json!([{"role": "assistant", "content": [{"type": "text", "text": "hi"}]}]))] + #[case::padded_text(json!([{"role": "assistant", "content": [{"type": "text", "text": " hi "}]}]))] + #[case::empty_string_content(json!([{"role": "user", "content": ""}]))] + #[case::textless_non_text_block(json!([{"role": "user", "content": [ + {"type": "image", "source": {"type": "base64", "media_type": "image/png", "data": "AA=="}} + ]}]))] + #[case::encrypted_reasoning_left_for_the_responses_bridge(json!([{"role": "assistant", "content": [ + {"type": "thinking", "thinking": "plan", "signature": tagged("gAAAA_1")}, + {"type": "redacted_thinking", "data": tagged("gAAAA_2")}, + {"type": "text", "text": "The answer."} + ]}]))] + fn strip_empty_content_blocks_leaves_untouched(#[case] input: Value) { + assert_eq!(apply(strip_empty_content_blocks, input.clone()), input); + } + + #[rstest] + #[case::replayed_provider_id("functions.Bash:0", "functions_Bash_0")] + #[case::thought_signature_suffix("call_abc123__thought__CiIBDDnWx+/a==", "call_abc123")] + #[case::splits_at_first_thought_separator("call_1__thought__a__thought__b", "call_1")] + #[case::valid_id("toolu_01-A_b", "toolu_01-A_b")] + #[case::non_ascii_letter("café", "caf_")] + #[case::only_invalid_characters("::", "__")] + #[case::empty("", "tool_use_id")] + #[case::thought_signature_only("__thought__CiIB", "tool_use_id")] + fn normalize_anthropic_tool_use_id_cases(#[case] raw: &str, #[case] expected: &str) { + assert_eq!(normalize_anthropic_tool_use_id(raw), expected); + } + + #[rstest] + #[case::tool_use_and_its_result( + json!([ + {"role": "assistant", "content": [{"type": "tool_use", "id": "functions.Bash:0", "name": "Bash", "input": {}}]}, + {"role": "user", "content": [{"type": "tool_result", "tool_use_id": "functions.Bash:0", "content": "ok"}]} + ]), + json!([ + {"role": "assistant", "content": [{"type": "tool_use", "id": "functions_Bash_0", "name": "Bash", "input": {}}]}, + {"role": "user", "content": [{"type": "tool_result", "tool_use_id": "functions_Bash_0", "content": "ok"}]} + ]) + )] + #[case::server_tool_use( + json!([{"role": "assistant", "content": [ + {"type": "server_tool_use", "id": "srv.1", "name": "web_search", "input": {}} + ]}]), + json!([{"role": "assistant", "content": [ + {"type": "server_tool_use", "id": "srv_1", "name": "web_search", "input": {}} + ]}]) + )] + #[case::tool_use_rewrites_only_its_id( + json!([{"role": "assistant", "content": [ + {"type": "tool_use", "id": "a.b", "tool_use_id": "c.d", "name": "Bash", "input": {}} + ]}]), + json!([{"role": "assistant", "content": [ + {"type": "tool_use", "id": "a_b", "tool_use_id": "c.d", "name": "Bash", "input": {}} + ]}]) + )] + #[case::tool_result_rewrites_only_its_tool_use_id( + json!([{"role": "user", "content": [ + {"type": "tool_result", "id": "a.b", "tool_use_id": "c.d", "content": "ok"} + ]}]), + json!([{"role": "user", "content": [ + {"type": "tool_result", "id": "a.b", "tool_use_id": "c_d", "content": "ok"} + ]}]) + )] + fn sanitize_tool_use_ids_rewrites(#[case] input: Value, #[case] expected: Value) { + assert_eq!(apply(sanitize_tool_use_ids, input), expected); + } + + #[rstest] + #[case::valid_ids(json!([ + {"role": "assistant", "content": [{"type": "tool_use", "id": "toolu_01", "name": "Bash", "input": {}}]}, + {"role": "user", "content": [{"type": "tool_result", "tool_use_id": "toolu_01", "content": "ok"}]} + ]))] + #[case::id_mentioned_in_text(json!([{"role": "user", "content": [{"type": "text", "text": "id: functions.Bash:0"}]}]))] + #[case::tool_use_without_id(json!([{"role": "assistant", "content": [{"type": "tool_use", "name": "Bash", "input": {}}]}]))] + #[case::tool_result_without_tool_use_id(json!([{"role": "user", "content": [{"type": "tool_result", "content": "ok"}]}]))] + #[case::string_content(json!([{"role": "user", "content": "functions.Bash:0"}]))] + fn sanitize_tool_use_ids_leaves_untouched(#[case] input: Value) { + assert_eq!(apply(sanitize_tool_use_ids, input.clone()), input); + } + + #[rstest] + #[case::thinking_block( + json!([{"role": "assistant", "content": [ + {"type": "thinking", "thinking": "hm", "signature": "s", "provider_specific_fields": {"a": 1}} + ]}]), + json!([{"role": "assistant", "content": [{"type": "thinking", "thinking": "hm", "signature": "s"}]}]) + )] + #[case::every_block_of_every_message( + json!([ + {"role": "assistant", "content": [ + {"type": "text", "text": "a", "provider_specific_fields": {"x": 1}}, + {"type": "tool_use", "id": "t1", "name": "f", "input": {}, "provider_specific_fields": {"y": 2}} + ]}, + {"role": "user", "content": [ + {"type": "tool_result", "tool_use_id": "t1", "content": "ok", "provider_specific_fields": {}} + ]} + ]), + json!([ + {"role": "assistant", "content": [ + {"type": "text", "text": "a"}, + {"type": "tool_use", "id": "t1", "name": "f", "input": {}} + ]}, + {"role": "user", "content": [{"type": "tool_result", "tool_use_id": "t1", "content": "ok"}]} + ]) + )] + fn strip_provider_specific_fields_rewrites(#[case] input: Value, #[case] expected: Value) { + assert_eq!(apply(strip_provider_specific_fields, input), expected); + } + + #[rstest] + #[case::string_content(json!([{"role": "user", "content": "provider_specific_fields"}]))] + #[case::blocks_without_the_field(json!([{"role": "assistant", "content": [{"type": "text", "text": "a"}]}]))] + fn strip_provider_specific_fields_leaves_untouched(#[case] input: Value) { + assert_eq!(apply(strip_provider_specific_fields, input.clone()), input); + } + + #[rstest] + #[case::tagged_thinking_signature(json!({"type": "thinking", "thinking": "x", "signature": tagged("g")}), true)] + #[case::tagged_redacted_data(json!({"type": "redacted_thinking", "data": tagged("g")}), true)] + #[case::bare_tag_signature(json!({"type": "thinking", "thinking": "x", "signature": tagged("")}), true)] + #[case::bare_tag_data(json!({"type": "redacted_thinking", "data": tagged("")}), true)] + #[case::anthropic_signature(json!({"type": "thinking", "thinking": "x", "signature": "ErcBCkgIValid"}), false)] + #[case::anthropic_data(json!({"type": "redacted_thinking", "data": "EmwKAhgBEgy"}), false)] + #[case::unsigned_thinking(json!({"type": "thinking", "thinking": "x"}), false)] + #[case::tag_in_text_block(json!({"type": "text", "text": tagged("g")}), false)] + #[case::tag_in_thinking_data(json!({"type": "thinking", "thinking": "x", "data": tagged("g")}), false)] + #[case::tag_in_redacted_signature( + json!({"type": "redacted_thinking", "data": "EmwKAhgBEgy", "signature": tagged("g")}), + false + )] + #[case::tag_not_at_start(json!({"type": "thinking", "thinking": "x", "signature": format!("x{}", tagged("g"))}), false)] + fn encrypted_reasoning_block_detection(#[case] input: Value, #[case] expected: bool) { + assert_eq!(is_encrypted_reasoning_block(&block(input)), expected); + } + + #[rstest] + #[case::only_the_bridge_tagged_blocks( + json!([ + {"role": "user", "content": "Solve it."}, + {"role": "assistant", "content": [ + {"type": "thinking", "thinking": "plan", "signature": tagged("gAAAA_1")}, + {"type": "redacted_thinking", "data": tagged("gAAAA_2")} + ]}, + {"role": "assistant", "content": [ + {"type": "thinking", "thinking": "plan", "signature": tagged("gAAAA_3")}, + {"type": "thinking", "thinking": "native", "signature": "EqQBCkYIAxgCIkA_anthropic_signed"}, + {"type": "redacted_thinking", "data": "EmwKAhgBEgy_anthropic_minted"}, + {"type": "text", "text": "The answer."} + ]} + ]), + json!([ + {"role": "user", "content": "Solve it."}, + {"role": "assistant", "content": [ + {"type": "thinking", "thinking": "native", "signature": "EqQBCkYIAxgCIkA_anthropic_signed"}, + {"type": "redacted_thinking", "data": "EmwKAhgBEgy_anthropic_minted"}, + {"type": "text", "text": "The answer."} + ]} + ]) + )] + #[case::bridge_turn_keeps_its_text( + json!([ + {"role": "user", "content": "Solve it."}, + {"role": "assistant", "content": [ + {"type": "thinking", "thinking": "plan", "signature": tagged("gAAAA_1")}, + {"type": "redacted_thinking", "data": tagged("gAAAA_2")}, + {"type": "text", "text": "The answer."} + ]}, + {"role": "user", "content": "And the next one?"} + ]), + json!([ + {"role": "user", "content": "Solve it."}, + {"role": "assistant", "content": [{"type": "text", "text": "The answer."}]}, + {"role": "user", "content": "And the next one?"} + ]) + )] + fn strip_encrypted_reasoning_blocks_rewrites(#[case] input: Value, #[case] expected: Value) { + assert_eq!(apply(strip_encrypted_reasoning_blocks, input), expected); + } + + #[rstest] + #[case::anthropic_signed_blocks(json!([ + {"role": "user", "content": "Solve it."}, + {"role": "assistant", "content": [ + {"type": "thinking", "thinking": "plan", "signature": "EqQBCkYIAxgCIkA_anthropic_signed"}, + {"type": "redacted_thinking", "data": "EmwKAhgBEgy_anthropic_minted"}, + {"type": "text", "text": "The answer."} + ]} + ]))] + #[case::string_content(json!([{"role": "user", "content": tagged("g")}]))] + fn strip_encrypted_reasoning_blocks_leaves_untouched(#[case] input: Value) { + assert_eq!( + apply(strip_encrypted_reasoning_blocks, input.clone()), + input + ); + } + + #[rstest] + #[case::advisor_exchange_between_texts( + json!([ + {"role": "user", "content": "Build a worker pool."}, + {"role": "assistant", "content": [ + {"type": "text", "text": "Let me consult the advisor."}, + {"type": "server_tool_use", "id": "srvtoolu_abc123", "name": "advisor", "input": {}}, + {"type": "advisor_tool_result", "tool_use_id": "srvtoolu_abc123", + "content": {"type": "advisor_result", "text": "Use channels."}}, + {"type": "text", "text": "Here is the implementation."} + ]} + ]), + json!([ + {"role": "user", "content": "Build a worker pool."}, + {"role": "assistant", "content": [ + {"type": "text", "text": "Let me consult the advisor."}, + {"type": "text", "text": "Here is the implementation."} + ]} + ]) + )] + #[case::only_results_of_this_turns_advisor_calls( + json!([{"role": "assistant", "content": [ + {"type": "server_tool_use", "id": "adv_1", "name": "advisor", "input": {}}, + {"type": "advisor_tool_result", "tool_use_id": "adv_1", "content": "advice"}, + {"type": "advisor_tool_result", "tool_use_id": "other", "content": "kept"}, + {"type": "tool_result", "tool_use_id": "adv_1", "content": "kept"}, + {"type": "text", "text": "answer"} + ]}]), + json!([{"role": "assistant", "content": [ + {"type": "advisor_tool_result", "tool_use_id": "other", "content": "kept"}, + {"type": "tool_result", "tool_use_id": "adv_1", "content": "kept"}, + {"type": "text", "text": "answer"} + ]}]) + )] + #[case::advisor_call_without_result( + json!([{"role": "assistant", "content": [ + {"type": "server_tool_use", "id": "adv_1", "name": "advisor", "input": {}}, + {"type": "text", "text": "answer"} + ]}]), + json!([{"role": "assistant", "content": [{"type": "text", "text": "answer"}]}]) + )] + #[case::advisor_only_turn_keeps_an_empty_block_list( + json!([{"role": "assistant", "content": [ + {"type": "server_tool_use", "id": "adv_1", "name": "advisor", "input": {}}, + {"type": "advisor_tool_result", "tool_use_id": "adv_1", "content": "advice"} + ]}]), + json!([{"role": "assistant", "content": []}]) + )] + fn strip_advisor_blocks_rewrites(#[case] input: Value, #[case] expected: Value) { + assert_eq!(apply(strip_advisor_blocks, input), expected); + } + + #[rstest] + #[case::no_advisor_blocks(json!([ + {"role": "user", "content": "Hello"}, + {"role": "assistant", "content": [ + {"type": "text", "text": "Hi there"}, + {"type": "tool_use", "id": "toolu_abc", "name": "get_weather", "input": {"location": "SF"}} + ]} + ]))] + #[case::user_turn(json!([{"role": "user", "content": [ + {"type": "server_tool_use", "id": "adv_2", "name": "advisor", "input": {}}, + {"type": "advisor_tool_result", "tool_use_id": "adv_2", "content": "advice"} + ]}]))] + #[case::other_server_tool(json!([{"role": "assistant", "content": [ + {"type": "server_tool_use", "id": "s1", "name": "web_search", "input": {}}, + {"type": "advisor_tool_result", "tool_use_id": "s1", "content": "advice"} + ]}]))] + #[case::client_tool_named_advisor(json!([{"role": "assistant", "content": [ + {"type": "tool_use", "id": "t1", "name": "advisor", "input": {}}, + {"type": "advisor_tool_result", "tool_use_id": "t1", "content": "advice"} + ]}]))] + #[case::advisor_call_with_empty_id(json!([{"role": "assistant", "content": [ + {"type": "server_tool_use", "id": "", "name": "advisor", "input": {}}, + {"type": "advisor_tool_result", "tool_use_id": "", "content": "advice"} + ]}]))] + #[case::advisor_call_without_id(json!([{"role": "assistant", "content": [ + {"type": "server_tool_use", "name": "advisor", "input": {}} + ]}]))] + #[case::string_content(json!([{"role": "assistant", "content": "advisor"}]))] + fn strip_advisor_blocks_leaves_untouched(#[case] input: Value) { + assert_eq!(apply(strip_advisor_blocks, input.clone()), input); + } + + #[rstest] + #[case::results_keep_their_evidence( + json!([ + {"role": "user", "content": "latest version?"}, + {"role": "assistant", "content": [ + {"type": "server_tool_use", "id": "srvtoolu_1", "name": "web_search", "input": {"query": "latest version"}}, + {"type": "web_search_tool_result", "tool_use_id": "srvtoolu_1", "content": [ + {"type": "web_search_result", "url": "https://example.com/releases", "title": "Releases", + "page_age": null, "encrypted_content": "", "snippet": "Latest release v1.95.0"} + ]}, + {"type": "text", "text": "v1.95.0"} + ]} + ]), + json!([ + {"role": "user", "content": "latest version?"}, + {"role": "assistant", "content": [ + {"type": "text", "text": "Web search results for 'latest version':\n\nTitle: Releases\nURL: https://example.com/releases\nSnippet: Latest release v1.95.0"}, + {"type": "text", "text": "v1.95.0"} + ]} + ]) + )] + #[case::each_result_lists_only_its_present_fields( + json!([{"role": "assistant", "content": [ + {"type": "server_tool_use", "id": "s1", "name": "web_search", "input": {"query": "q"}}, + {"type": "web_search_tool_result", "tool_use_id": "s1", "content": [ + {"type": "web_search_result", "title": "A"}, + {"type": "web_search_result", "snippet": "b"}, + {"type": "web_search_result", "url": "https://c"} + ]} + ]}]), + json!([{"role": "assistant", "content": [ + {"type": "text", "text": "Web search results for 'q':\n\nTitle: A\n\nSnippet: b\n\nURL: https://c"} + ]}]) + )] + #[case::result_without_fields_renders_the_header_only( + json!([{"role": "assistant", "content": [ + {"type": "server_tool_use", "id": "s1", "name": "web_search", "input": {"query": "q"}}, + {"type": "web_search_tool_result", "tool_use_id": "s1", "content": [{"type": "web_search_result"}]} + ]}]), + json!([{"role": "assistant", "content": [{"type": "text", "text": "Web search results for 'q':"}]}]) + )] + #[case::resultless_search( + json!([{"role": "assistant", "content": [ + {"type": "server_tool_use", "id": "srvtoolu_1", "name": "web_search", "input": {"query": "who won"}}, + {"type": "web_search_tool_result", "tool_use_id": "srvtoolu_1", "content": []}, + {"type": "text", "text": "I could not find that."} + ]}]), + json!([{"role": "assistant", "content": [ + {"type": "text", "text": "Web search results for 'who won':\n\nNo results were returned."}, + {"type": "text", "text": "I could not find that."} + ]}]) + )] + #[case::failed_search( + json!([{"role": "assistant", "content": [ + {"type": "server_tool_use", "id": "srvtoolu_1", "name": "web_search", "input": {"query": "q"}}, + {"type": "web_search_tool_result", "tool_use_id": "srvtoolu_1", + "content": {"type": "web_search_tool_result_error", "error_code": "max_uses_exceeded"}} + ]}]), + json!([{"role": "assistant", "content": [ + {"type": "text", "text": "Web search results for 'q':\n\nSearch failed: max_uses_exceeded"} + ]}]) + )] + #[case::failed_search_without_error_code( + json!([{"role": "assistant", "content": [ + {"type": "web_search_tool_result", "tool_use_id": "e1", "content": {"type": "web_search_tool_result_error"}} + ]}]), + json!([{"role": "assistant", "content": [ + {"type": "text", "text": "Web search results:\n\nSearch failed: unavailable"} + ]}]) + )] + #[case::server_tool_use_without_query( + json!([{"role": "assistant", "content": [ + {"type": "server_tool_use", "id": "s1", "name": "web_search"}, + {"type": "server_tool_use", "id": "s2", "name": "web_search", "input": {}}, + {"type": "web_search_tool_result", "tool_use_id": "s1", "content": []}, + {"type": "web_search_tool_result", "tool_use_id": "s2", "content": []} + ]}]), + json!([{"role": "assistant", "content": [ + {"type": "text", "text": "Web search results:\n\nNo results were returned."}, + {"type": "text", "text": "Web search results:\n\nNo results were returned."} + ]}]) + )] + #[case::genuine_results_in_the_same_turn_stay( + json!([{"role": "assistant", "content": [ + {"type": "server_tool_use", "id": "s1", "name": "web_search", "input": {"query": "rust"}}, + {"type": "web_search_tool_result", "tool_use_id": "s1", "content": [ + {"type": "web_search_result", "url": "https://r", "title": "Rust", "snippet": "fast"} + ]}, + {"type": "server_tool_use", "id": "s2", "name": "web_search", "input": {"query": "real"}}, + {"type": "web_search_tool_result", "tool_use_id": "s2", "content": [ + {"type": "web_search_result", "url": "https://a", "title": "A", "snippet": "b", "encrypted_content": "enc"} + ]}, + {"type": "text", "text": "done"} + ]}]), + json!([{"role": "assistant", "content": [ + {"type": "text", "text": "Web search results for 'rust':\n\nTitle: Rust\nURL: https://r\nSnippet: fast"}, + {"type": "server_tool_use", "id": "s2", "name": "web_search", "input": {"query": "real"}}, + {"type": "web_search_tool_result", "tool_use_id": "s2", "content": [ + {"type": "web_search_result", "url": "https://a", "title": "A", "snippet": "b", "encrypted_content": "enc"} + ]}, + {"type": "text", "text": "done"} + ]}]) + )] + #[case::other_blocks_sharing_the_tool_use_id_stay( + json!([{"role": "assistant", "content": [ + {"type": "server_tool_use", "id": "s1", "name": "web_search", "input": {"query": "q"}}, + {"type": "web_search_tool_result", "tool_use_id": "s1", "content": []}, + {"type": "tool_result", "tool_use_id": "s1", "content": "x"} + ]}]), + json!([{"role": "assistant", "content": [ + {"type": "text", "text": "Web search results for 'q':\n\nNo results were returned."}, + {"type": "tool_result", "tool_use_id": "s1", "content": "x"} + ]}]) + )] + #[case::query_lookup_stays_within_the_message( + json!([ + {"role": "assistant", "content": [ + {"type": "server_tool_use", "id": "s1", "name": "web_search", "input": {"query": "q"}} + ]}, + {"role": "assistant", "content": [ + {"type": "web_search_tool_result", "tool_use_id": "s1", "content": []} + ]} + ]), + json!([ + {"role": "assistant", "content": [ + {"type": "server_tool_use", "id": "s1", "name": "web_search", "input": {"query": "q"}} + ]}, + {"role": "assistant", "content": [ + {"type": "text", "text": "Web search results:\n\nNo results were returned."} + ]} + ]) + )] + #[case::result_without_any_field_keeps_its_slot( + json!([{"role": "assistant", "content": [ + {"type": "web_search_tool_result", "tool_use_id": "s1", "content": [ + {"type": "web_search_result", "url": "https://a", "title": "A"}, + {"type": "web_search_result"}, + {"type": "web_search_result", "title": "B"} + ]} + ]}]), + json!([{"role": "assistant", "content": [ + {"type": "text", "text": "Web search results:\n\nTitle: A\nURL: https://a\n\n\n\nTitle: B"} + ]}]) + )] + #[case::non_string_query_keeps_its_server_tool_use( + json!([{"role": "assistant", "content": [ + {"type": "server_tool_use", "id": "s1", "name": "web_search", "input": {"query": 123}}, + {"type": "web_search_tool_result", "tool_use_id": "s1", "content": []} + ]}]), + json!([{"role": "assistant", "content": [ + {"type": "server_tool_use", "id": "s1", "name": "web_search", "input": {"query": 123}}, + {"type": "text", "text": "Web search results:\n\nNo results were returned."} + ]}]) + )] + #[case::non_object_input_keeps_its_server_tool_use( + json!([{"role": "assistant", "content": [ + {"type": "server_tool_use", "id": "s1", "name": "web_search", "input": "q"}, + {"type": "web_search_tool_result", "tool_use_id": "s1", "content": []} + ]}]), + json!([{"role": "assistant", "content": [ + {"type": "server_tool_use", "id": "s1", "name": "web_search", "input": "q"}, + {"type": "text", "text": "Web search results:\n\nNo results were returned."} + ]}]) + )] + #[case::repeated_tool_use_id_renders_each_block_from_its_own_results( + json!([{"role": "assistant", "content": [ + {"type": "web_search_tool_result", "tool_use_id": "s1", "content": []}, + {"type": "web_search_tool_result", "tool_use_id": "s1", "content": {"type": "web_search_tool_result_error", "error_code": "max_uses"}} + ]}]), + json!([{"role": "assistant", "content": [ + {"type": "text", "text": "Web search results:\n\nNo results were returned."}, + {"type": "text", "text": "Web search results:\n\nSearch failed: max_uses"} + ]}]) + )] + #[case::encrypted_block_sharing_a_replayed_id_stays( + json!([{"role": "assistant", "content": [ + {"type": "server_tool_use", "id": "s1", "name": "web_search", "input": {"query": "q"}}, + {"type": "web_search_tool_result", "tool_use_id": "s1", "content": []}, + {"type": "web_search_tool_result", "tool_use_id": "s1", "content": [ + {"type": "web_search_result", "url": "https://a", "encrypted_content": "enc"} + ]} + ]}]), + json!([{"role": "assistant", "content": [ + {"type": "text", "text": "Web search results for 'q':\n\nNo results were returned."}, + {"type": "web_search_tool_result", "tool_use_id": "s1", "content": [ + {"type": "web_search_result", "url": "https://a", "encrypted_content": "enc"} + ]} + ]}]) + )] + #[case::last_query_wins_for_a_repeated_server_tool_use_id( + json!([{"role": "assistant", "content": [ + {"type": "server_tool_use", "id": "s1", "name": "web_search", "input": {"query": "first"}}, + {"type": "server_tool_use", "id": "s1", "name": "web_search", "input": {"query": "second"}}, + {"type": "web_search_tool_result", "tool_use_id": "s1", "content": []} + ]}]), + json!([{"role": "assistant", "content": [ + {"type": "text", "text": "Web search results for 'second':\n\nNo results were returned."} + ]}]) + )] + fn flatten_unencrypted_web_search_results_rewrites( + #[case] input: Value, + #[case] expected: Value, + ) { + assert_eq!( + apply(flatten_unencrypted_web_search_results, input), + expected + ); + } + + #[rstest] + #[case::anthropic_issued_results(json!([{"role": "assistant", "content": [ + {"type": "server_tool_use", "id": "srvtoolu_1", "name": "web_search", "input": {"query": "q"}}, + {"type": "web_search_tool_result", "tool_use_id": "srvtoolu_1", "content": [ + {"type": "web_search_result", "url": "https://example.com", "title": "Example", + "page_age": null, "encrypted_content": "EqgfCioIARgBIiQ4"} + ]} + ]}]))] + #[case::any_encrypted_result_marks_the_block_genuine(json!([{"role": "assistant", "content": [ + {"type": "web_search_tool_result", "tool_use_id": "s1", "content": [ + {"type": "web_search_result", "url": "https://a", "encrypted_content": ""}, + {"type": "web_search_result", "url": "https://b", "encrypted_content": "enc"} + ]} + ]}]))] + #[case::result_without_tool_use_id(json!([{"role": "assistant", "content": [ + {"type": "web_search_tool_result", "content": []} + ]}]))] + #[case::foreign_item_in_results(json!([{"role": "assistant", "content": [ + {"type": "web_search_tool_result", "tool_use_id": "s1", "content": [ + {"type": "web_search_result", "url": "https://a"}, + {"type": "text", "text": "x"} + ]} + ]}]))] + #[case::result_with_null_url(json!([{"role": "assistant", "content": [ + {"type": "web_search_tool_result", "tool_use_id": "s1", "content": [ + {"type": "web_search_result", "url": null, "title": "A"} + ]} + ]}]))] + #[case::string_result_content(json!([{"role": "assistant", "content": [ + {"type": "web_search_tool_result", "tool_use_id": "s1", "content": "oops"} + ]}]))] + #[case::object_content_that_is_not_an_error(json!([{"role": "assistant", "content": [ + {"type": "web_search_tool_result", "tool_use_id": "s1", "content": {"type": "web_search_result", "url": "https://a"}} + ]}]))] + #[case::string_content(json!([{"role": "assistant", "content": "web_search_tool_result"}]))] + fn flatten_unencrypted_web_search_results_leaves_untouched(#[case] input: Value) { + assert_eq!( + apply(flatten_unencrypted_web_search_results, input.clone()), + input + ); + } + + #[rstest] + #[case::with_results(json!([{"type": "web_search_result", "url": "u", "title": "Rome", "snippet": "s", "page_age": null}]))] + #[case::without_results(json!([]))] + fn flatten_unencrypted_web_search_results_is_idempotent(#[case] results: Value) { + let input = replayed_search_turn(results); + let once = apply(flatten_unencrypted_web_search_results, input.clone()); + let twice = apply(flatten_unencrypted_web_search_results, once.clone()); + assert_ne!(once, input); + assert_eq!(twice, once); + } + + #[rstest] + #[case::no_existing_header(None, "b", "b")] + #[case::empty_existing_header(Some(""), "b", "b")] + #[case::whitespace_existing_header(Some(" "), "b", "b")] + #[case::sorted_after_merge(Some("c,a"), "b", "a,b,c")] + #[case::already_present(Some("a,b"), "a", "a,b")] + #[case::trimmed_and_deduplicated(Some("b, a ,b"), "c", "a,b,c")] + #[case::blank_pieces_skipped(Some("a,,b"), "c", "a,b,c")] + fn beta_values_merge_sorted_and_deduplicated( + #[case] existing: Option<&str>, + #[case] new_beta: &str, + #[case] expected: &str, + ) { + assert_eq!( + join_beta_values(split_beta_values(existing).chain([new_beta.to_string()])), + expected + ); + } + + #[rstest] + #[case::raw_token("sk-ant-oat01-abc123", true)] + #[case::bearer_token("Bearer sk-ant-oat02-xyz789", true)] + #[case::bare_prefix(ANTHROPIC_OAUTH_TOKEN_PREFIX, true)] + #[case::api_key("sk-ant-api01-abc123", false)] + #[case::bearer_api_key("Bearer sk-ant-api01-abc123", false)] + #[case::empty("", false)] + #[case::uppercase_prefix("sk-ant-OAT01-abc123", false)] + #[case::shouting_prefix("SK-ANT-OAT01-abc123", false)] + #[case::lowercase_bearer("bearer sk-ant-oat01-abc123", false)] + #[case::bearer_stripped_once("Bearer Bearer sk-ant-oat01-abc123", false)] + #[case::prefix_not_at_start(" sk-ant-oat01-abc123", false)] + fn anthropic_oauth_key_detection(#[case] value: &str, #[case] expected: bool) { + assert_eq!(is_anthropic_oauth_key(value), expected); + } + + #[rstest] + #[case::regex_tool(Some(json!([{"type": ANTHROPIC_TOOL_SEARCH_TOOL_TYPES[0], "name": "tool_search_tool_regex"}])), true)] + #[case::bm25_tool(Some(json!([{"type": ANTHROPIC_TOOL_SEARCH_TOOL_TYPES[1], "name": "tool_search_tool_bm25"}])), true)] + #[case::after_other_tools( + Some(json!([{"name": "get_weather", "input_schema": {}}, {"type": ANTHROPIC_TOOL_SEARCH_TOOL_TYPES[1]}])), + true + )] + #[case::function_tool(Some(json!([{"type": "function", "function": {"name": "get_weather"}}])), false)] + #[case::name_without_type(Some(json!([{"name": ANTHROPIC_TOOL_SEARCH_TOOL_TYPES[0]}])), false)] + #[case::empty_tools(Some(json!([])), false)] + #[case::no_tools(None, false)] + fn tool_search_detection(#[case] input: Option, #[case] expected: bool) { + assert_eq!(is_tool_search_used(tools(input).as_deref()), expected); + } + + #[rstest] + #[case::advisor_tool(Some(json!([{"type": ANTHROPIC_ADVISOR_TOOL_TYPE, "name": "advisor"}])), true)] + #[case::after_other_tools(Some(json!([{"name": "f", "input_schema": {}}, {"type": ANTHROPIC_ADVISOR_TOOL_TYPE}])), true)] + #[case::tool_named_advisor(Some(json!([{"name": "advisor", "input_schema": {}}])), false)] + #[case::other_server_tool(Some(json!([{"type": "web_search_20250305", "name": "web_search"}])), false)] + #[case::empty_tools(Some(json!([])), false)] + #[case::no_tools(None, false)] + fn advisor_tool_detection(#[case] input: Option, #[case] expected: bool) { + assert_eq!(has_advisor_tool(tools(input).as_deref()), expected); + } + + #[rstest] + #[case::param_without_history(Some(json!({})), json!([]), true)] + #[case::param_with_unsigned_history( + Some(json!({"trigger": 1})), + json!([{"role": "assistant", "content": [{"type": "compaction", "content": "c"}]}]), + true + )] + #[case::signed_block(None, json!([{"role": "assistant", "content": [{"type": "compaction", "content": "c", "signature": "s"}]}]), true)] + #[case::signed_block_later_in_history( + None, + json!([ + {"role": "user", "content": "hi"}, + {"role": "assistant", "content": [{"type": "text", "text": "a"}, {"type": "compaction", "content": "c", "signature": "s"}]} + ]), + true + )] + #[case::unsigned_block(None, json!([{"role": "assistant", "content": [{"type": "compaction", "content": "c"}]}]), false)] + #[case::empty_signature(None, json!([{"role": "assistant", "content": [{"type": "compaction", "content": "c", "signature": ""}]}]), false)] + #[case::signed_non_compaction_block( + None, + json!([{"role": "assistant", "content": [{"type": "thinking", "thinking": "t", "signature": "s"}]}]), + false + )] + #[case::string_content(None, json!([{"role": "user", "content": "compaction"}]), false)] + #[case::neither(None, json!([]), false)] + fn native_compaction_beta_requirement( + #[case] compaction: Option, + #[case] messages: Value, + #[case] expected: bool, + ) { + assert_eq!( + requires_native_compaction_beta(compaction.as_ref(), &history(messages)), + expected + ); + } + + #[rstest] + #[case::low(EffortLevel::Low, "low")] + #[case::medium(EffortLevel::Medium, "medium")] + #[case::high(EffortLevel::High, "high")] + #[case::xhigh(EffortLevel::Xhigh, "xhigh")] + #[case::max(EffortLevel::Max, "max")] + fn effort_level_names_agree_across_str_parse_and_serde( + #[case] level: EffortLevel, + #[case] name: &str, + ) { + assert_eq!(level.as_str(), name); + assert_eq!(EffortLevel::parse(name), Some(level)); + assert_eq!(serde_json::to_value(level).unwrap(), json!(name)); + assert_eq!( + serde_json::from_value::(json!(name)).unwrap(), + level + ); + } + + #[rstest] + #[case::unknown("ultra")] + #[case::minimal_is_not_an_output_config_level("minimal")] + #[case::uppercase("HIGH")] + #[case::empty("")] + fn effort_level_parse_rejects(#[case] value: &str) { + assert_eq!(EffortLevel::parse(value), None); + } + + #[rstest] + #[case::minimal_only(tiers(true, false, false, false, false, false), [false, false, false, false, false])] + #[case::low_only(tiers(false, true, false, false, false, false), [true, false, false, false, false])] + #[case::medium_only(tiers(false, false, true, false, false, false), [false, true, false, false, false])] + #[case::high_only(tiers(false, false, false, true, false, false), [false, false, true, false, false])] + #[case::xhigh_only(tiers(false, false, false, false, true, false), [false, false, false, true, false])] + #[case::max_only(tiers(false, false, false, false, false, true), [false, false, false, false, true])] + fn supports_effort_tier_reads_the_matching_flag( + #[case] effort_tiers: SupportedEffortTiers, + #[case] expected: [bool; 5], + unmapped: AnthropicModelCapabilities, + ) { + let capabilities = AnthropicModelCapabilities { + effort_tiers, + ..unmapped + }; + assert_eq!( + ALL_LEVELS.map(|level| capabilities.supports_effort_tier(level)), + expected + ); + } + + #[rstest] + #[case::unmapped(false, false, false, SupportedEffortTiers::default(), false)] + #[case::reasoning_and_adaptive_thinking_alone( + true, + true, + false, + SupportedEffortTiers::default(), + false + )] + #[case::output_config_without_tiers(false, false, true, SupportedEffortTiers::default(), true)] + #[case::minimal_tier( + false, + false, + false, + tiers(true, false, false, false, false, false), + true + )] + #[case::low_tier( + false, + false, + false, + tiers(false, true, false, false, false, false), + true + )] + #[case::medium_tier( + false, + false, + false, + tiers(false, false, true, false, false, false), + true + )] + #[case::high_tier( + false, + false, + false, + tiers(false, false, false, true, false, false), + true + )] + #[case::xhigh_tier( + false, + false, + false, + tiers(false, false, false, false, true, false), + true + )] + #[case::max_tier( + false, + false, + false, + tiers(false, false, false, false, false, true), + true + )] + fn supports_effort_param_cases( + #[case] supports_reasoning: bool, + #[case] supports_adaptive_thinking: bool, + #[case] supports_output_config: bool, + #[case] effort_tiers: SupportedEffortTiers, + #[case] expected: bool, + unmapped: AnthropicModelCapabilities, + ) { + let capabilities = AnthropicModelCapabilities { + supports_reasoning, + supports_adaptive_thinking, + supports_output_config, + effort_tiers, + ..unmapped + }; + assert_eq!(capabilities.supports_effort_param(), expected); + } + + #[rstest] + #[case::max_on_adaptive_thinking_model(true, SupportedEffortTiers::default(), "max", None)] + #[case::max_on_max_tier_model( + false, + tiers(false, false, false, false, false, true), + "max", + None + )] + #[case::max_on_output_config_only_model( + false, + SupportedEffortTiers::default(), + "max", + Some("effort='max' is not supported by this model. Got model: claude-test") + )] + #[case::max_on_xhigh_tier_model( + false, + tiers(false, false, false, false, true, false), + "max", + Some("effort='max' is not supported by this model. Got model: claude-test") + )] + #[case::xhigh_on_xhigh_tier_model( + false, + tiers(false, false, false, false, true, false), + "xhigh", + None + )] + #[case::xhigh_on_adaptive_thinking_model( + true, + SupportedEffortTiers::default(), + "xhigh", + Some("effort='xhigh' is not supported by this model. Got model: claude-test") + )] + #[case::xhigh_on_max_tier_model( + false, + tiers(false, false, false, false, false, true), + "xhigh", + Some("effort='xhigh' is not supported by this model. Got model: claude-test") + )] + #[case::high_on_unmapped_model(false, SupportedEffortTiers::default(), "high", None)] + #[case::low_on_unmapped_model(false, SupportedEffortTiers::default(), "low", None)] + #[case::unknown_level_is_left_to_other_validation( + false, + SupportedEffortTiers::default(), + "ultra", + None + )] + fn effort_level_rejection_cases( + #[case] supports_adaptive_thinking: bool, + #[case] effort_tiers: SupportedEffortTiers, + #[case] effort: &str, + #[case] expected: Option<&str>, + unmapped: AnthropicModelCapabilities, + ) { + let capabilities = AnthropicModelCapabilities { + supports_output_config: true, + supports_adaptive_thinking, + effort_tiers, + ..unmapped + }; + assert_eq!( + capabilities + .effort_level_rejection(effort, "claude-test") + .as_deref(), + expected + ); + } + + #[rstest] + fn unmapped_model_has_no_reasoning_features_but_accepts_sampling_params( + unmapped: AnthropicModelCapabilities, + ) { + assert_eq!( + unmapped, + AnthropicModelCapabilities { + supports_reasoning: false, + supports_adaptive_thinking: false, + thinking_always_on: false, + supports_legacy_thinking: false, + supports_output_config: false, + supports_sampling_params: true, + supports_speed: false, + effort_tiers: tiers(false, false, false, false, false, false), + } + ); + assert_eq!( + serde_json::from_value::(json!({})).unwrap(), + unmapped + ); + } + + #[rstest] + #[case::sampling_params_removed( + json!({"supports_sampling_params": false}), + AnthropicModelCapabilities { supports_sampling_params: false, ..AnthropicModelCapabilities::default() } + )] + #[case::fast_mode( + json!({"supports_speed": true}), + AnthropicModelCapabilities { supports_speed: true, ..AnthropicModelCapabilities::default() } + )] + #[case::partial_effort_tiers( + json!({"supports_reasoning": true, "effort_tiers": {"xhigh": true}}), + AnthropicModelCapabilities { + supports_reasoning: true, + effort_tiers: tiers(false, false, false, false, true, false), + ..AnthropicModelCapabilities::default() + } + )] + fn capabilities_fill_missing_flags_with_unmapped_defaults( + #[case] input: Value, + #[case] expected: AnthropicModelCapabilities, + ) { + assert_eq!( + serde_json::from_value::(input).unwrap(), + expected + ); + } +} diff --git a/litellm-rust/crates/llms/src/anthropic/experimental_pass_through/messages/handler.rs b/litellm-rust/crates/llms/src/anthropic/experimental_pass_through/messages/handler.rs new file mode 100644 index 00000000000..0e2ab97956a --- /dev/null +++ b/litellm-rust/crates/llms/src/anthropic/experimental_pass_through/messages/handler.rs @@ -0,0 +1,270 @@ +use litellm_types::llms::anthropic_messages::anthropic_request::{ + AnthropicMessage, AnthropicMessagesRequest, +}; +use serde_json::{Value, json}; + +use crate::{ + anthropic::common_utils::{ + flatten_unencrypted_web_search_results, sanitize_tool_use_ids, strip_empty_content_blocks, + strip_provider_specific_fields, + }, + base_llm::chat::transformation::Error, +}; + +pub fn shape_anthropic_messages_request( + request: AnthropicMessagesRequest, + reasoning_auto_summary: bool, +) -> Result { + Ok(AnthropicMessagesRequest { + messages: sanitize_anthropic_messages(request.messages), + metadata: request + .metadata + .as_ref() + .map(validate_anthropic_api_metadata) + .transpose()?, + thinking: with_reasoning_auto_summary(request.thinking, reasoning_auto_summary), + ..request + }) +} + +fn sanitize_anthropic_messages(messages: Vec) -> Vec { + strip_provider_specific_fields(flatten_unencrypted_web_search_results( + sanitize_tool_use_ids(strip_empty_content_blocks(messages)), + )) +} + +fn validate_anthropic_api_metadata(metadata: &Value) -> Result { + let Value::Object(fields) = metadata else { + return Err(Error::InvalidRequest(format!( + "metadata must be an object, got {metadata}" + ))); + }; + match fields.get("user_id") { + None | Some(Value::Null) => Ok(json!({})), + Some(Value::String(user_id)) => Ok(json!({"user_id": user_id})), + Some(other) => Err(Error::InvalidRequest(format!( + "metadata.user_id must be a string, got {other}" + ))), + } +} + +fn with_reasoning_auto_summary(thinking: Option, enabled: bool) -> Option { + let Some(Value::Object(thinking)) = thinking else { + return thinking; + }; + if !enabled || thinking.get("type").and_then(Value::as_str) == Some("disabled") { + return Some(Value::Object(thinking)); + } + Some(Value::Object( + thinking + .into_iter() + .filter(|(key, _)| key != "display") + .chain([("display".to_string(), json!("summarized"))]) + .collect(), + )) +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + + use super::*; + + fn messages(value: Value) -> Vec { + serde_json::from_value(value).unwrap() + } + + fn request(body: Value) -> AnthropicMessagesRequest { + serde_json::from_value(body).unwrap() + } + + #[rstest] + #[case::empty_text_next_to_a_tool_use( + json!([{"role": "assistant", "content": [ + {"type": "text", "text": " "}, + {"type": "tool_use", "id": "t", "name": "B", "input": {}} + ]}]), + json!([{"role": "assistant", "content": [ + {"type": "tool_use", "id": "t", "name": "B", "input": {}} + ]}]), + )] + #[case::cross_provider_tool_ids( + json!([ + {"role": "assistant", "content": [{"type": "tool_use", "id": "functions.Bash:0", "name": "Bash", "input": {}}]}, + {"role": "user", "content": [{"type": "tool_result", "tool_use_id": "functions.Bash:0", "content": "ok"}]} + ]), + json!([ + {"role": "assistant", "content": [{"type": "tool_use", "id": "functions_Bash_0", "name": "Bash", "input": {}}]}, + {"role": "user", "content": [{"type": "tool_result", "tool_use_id": "functions_Bash_0", "content": "ok"}]} + ]), + )] + #[case::replayed_unencrypted_web_search_results( + json!([ + {"role": "user", "content": "latest litellm version?"}, + {"role": "assistant", "content": [ + {"type": "server_tool_use", "id": "srvtoolu_1", "name": "web_search", "input": {"query": "latest litellm version"}}, + {"type": "web_search_tool_result", "tool_use_id": "srvtoolu_1", "content": [{ + "type": "web_search_result", + "url": "https://github.com/BerriAI/litellm/releases", + "title": "Releases", + "page_age": null, + "encrypted_content": "", + "snippet": "Latest release v1.95.0" + }]} + ]}, + {"role": "user", "content": "which version?"} + ]), + json!([ + {"role": "user", "content": "latest litellm version?"}, + {"role": "assistant", "content": [{ + "type": "text", + "text": "Web search results for 'latest litellm version':\n\nTitle: Releases\nURL: https://github.com/BerriAI/litellm/releases\nSnippet: Latest release v1.95.0" + }]}, + {"role": "user", "content": "which version?"} + ]), + )] + #[case::replayed_provider_specific_fields( + json!([ + {"role": "assistant", "content": [{ + "type": "tool_use", "id": "toolu_01", "name": "get_weather", "input": {"city": "Paris"}, + "provider_specific_fields": {"signature": "sig_abc"} + }]}, + {"role": "user", "content": [{"type": "tool_result", "tool_use_id": "toolu_01", "content": "Sunny"}]} + ]), + json!([ + {"role": "assistant", "content": [{"type": "tool_use", "id": "toolu_01", "name": "get_weather", "input": {"city": "Paris"}}]}, + {"role": "user", "content": [{"type": "tool_result", "tool_use_id": "toolu_01", "content": "Sunny"}]} + ]), + )] + #[case::ids_are_normalized_before_web_search_results_flatten( + json!([ + {"role": "user", "content": "run it"}, + {"role": "assistant", "content": [ + {"type": "thinking", "thinking": "", "signature": "sig"}, + {"type": "text", "text": ""}, + {"type": "tool_use", "id": "functions.Bash:0", "name": "Bash", "input": {}, "provider_specific_fields": {"x": 1}}, + {"type": "server_tool_use", "id": "srv.1", "name": "web_search", "input": {"query": "q"}, "provider_specific_fields": {"x": 2}}, + {"type": "web_search_tool_result", "tool_use_id": "srv.1", "provider_specific_fields": {"x": 3}, "content": [ + {"type": "web_search_result", "url": "u", "title": "", "encrypted_content": "", "provider_specific_fields": {"x": 4}} + ]} + ]}, + {"role": "user", "content": [{"type": "tool_result", "tool_use_id": "functions.Bash:0", "content": "ok"}]}, + {"role": "assistant", "content": [{"type": "text", "text": " "}]} + ]), + json!([ + {"role": "user", "content": "run it"}, + {"role": "assistant", "content": [ + {"type": "tool_use", "id": "functions_Bash_0", "name": "Bash", "input": {}}, + {"type": "server_tool_use", "id": "srv_1", "name": "web_search", "input": {"query": "q"}}, + {"type": "text", "text": "Web search results:\n\nURL: u"} + ]}, + {"role": "user", "content": [{"type": "tool_result", "tool_use_id": "functions_Bash_0", "content": "ok"}]} + ]), + )] + fn sanitize_anthropic_messages_cleans_replayed_history( + #[case] history: Value, + #[case] expected: Value, + ) { + assert_eq!( + serde_json::to_value(sanitize_anthropic_messages(messages(history))).unwrap(), + expected + ); + } + + #[rstest] + #[case::keeps_only_user_id(json!({"user_id": "u-1", "trace_id": "internal"}), Ok(json!({"user_id": "u-1"})))] + #[case::null_user_id(json!({"user_id": null, "trace_id": "internal"}), Ok(json!({})))] + #[case::no_user_id(json!({"trace_id": "internal"}), Ok(json!({})))] + #[case::empty(json!({}), Ok(json!({})))] + #[case::numeric_user_id( + json!({"user_id": 123}), + Err(Error::InvalidRequest("metadata.user_id must be a string, got 123".to_string())), + )] + #[case::boolean_user_id( + json!({"user_id": true}), + Err(Error::InvalidRequest("metadata.user_id must be a string, got true".to_string())), + )] + #[case::not_an_object( + json!(["u-1"]), + Err(Error::InvalidRequest(r#"metadata must be an object, got ["u-1"]"#.to_string())), + )] + fn validate_anthropic_api_metadata_passes_only_a_string_user_id( + #[case] metadata: Value, + #[case] expected: Result, + ) { + assert_eq!(validate_anthropic_api_metadata(&metadata), expected); + } + + #[rstest] + #[case::adaptive( + Some(json!({"type": "adaptive", "budget_tokens": 5000})), + true, + Some(json!({"type": "adaptive", "budget_tokens": 5000, "display": "summarized"})), + )] + #[case::enabled( + Some(json!({"type": "enabled", "budget_tokens": 10000})), + true, + Some(json!({"type": "enabled", "budget_tokens": 10000, "display": "summarized"})), + )] + #[case::no_type(Some(json!({})), true, Some(json!({"display": "summarized"})))] + #[case::display_omitted_is_overridden( + Some(json!({"type": "enabled", "budget_tokens": 10000, "display": "omitted"})), + true, + Some(json!({"type": "enabled", "budget_tokens": 10000, "display": "summarized"})), + )] + #[case::display_summarized_is_kept( + Some(json!({"type": "enabled", "display": "summarized"})), + true, + Some(json!({"type": "enabled", "display": "summarized"})), + )] + #[case::disabled_thinking(Some(json!({"type": "disabled"})), true, Some(json!({"type": "disabled"})))] + #[case::flag_off( + Some(json!({"type": "enabled", "budget_tokens": 10000})), + false, + Some(json!({"type": "enabled", "budget_tokens": 10000})), + )] + #[case::flag_off_keeps_callers_display( + Some(json!({"type": "enabled", "display": "omitted"})), + false, + Some(json!({"type": "enabled", "display": "omitted"})), + )] + #[case::no_thinking(None, true, None)] + #[case::non_object_thinking(Some(json!("enabled")), true, Some(json!("enabled")))] + fn reasoning_auto_summary_marks_active_thinking_as_summarized( + #[case] thinking: Option, + #[case] enabled: bool, + #[case] expected: Option, + ) { + assert_eq!(with_reasoning_auto_summary(thinking, enabled), expected); + } + + #[test] + fn shaping_cleans_messages_metadata_and_thinking() { + let sanitized = shape_anthropic_messages_request( + request(json!({ + "model": "m", + "messages": [{"role": "assistant", "content": [ + {"type": "text", "text": ""}, + {"type": "tool_use", "id": "functions.Bash:0", "name": "Bash", "input": {}} + ]}], + "metadata": {"user_id": "u", "trace_id": "t"}, + "thinking": {"type": "enabled", "budget_tokens": 1024}, + "safeguards": [{"type": "dangerous_tool_use"}] + })), + true, + ) + .unwrap(); + assert_eq!( + serde_json::to_value(sanitized).unwrap(), + json!({ + "model": "m", + "messages": [{"role": "assistant", "content": [ + {"type": "tool_use", "id": "functions_Bash_0", "name": "Bash", "input": {}} + ]}], + "metadata": {"user_id": "u"}, + "thinking": {"type": "enabled", "budget_tokens": 1024, "display": "summarized"}, + "safeguards": [{"type": "dangerous_tool_use"}] + }) + ); + } +} diff --git a/litellm-rust/crates/llms/src/anthropic/experimental_pass_through/messages/headers.rs b/litellm-rust/crates/llms/src/anthropic/experimental_pass_through/messages/headers.rs new file mode 100644 index 00000000000..8d48d7a0f5c --- /dev/null +++ b/litellm-rust/crates/llms/src/anthropic/experimental_pass_through/messages/headers.rs @@ -0,0 +1,643 @@ +use litellm_types::llms::anthropic_messages::anthropic_request::AnthropicMessagesRequest; +use serde_json::Value; + +use crate::{ + anthropic::{ + ANTHROPIC_OAUTH_TOKEN_PREFIX, + common_utils::{ + ANTHROPIC_OAUTH_BETA_HEADER, beta, has_advisor_tool, is_anthropic_oauth_key, + is_tool_search_used, join_beta_values, requires_native_compaction_beta, + split_beta_values, + }, + }, + base_llm::anthropic_messages::transformation::Headers, +}; + +const ANTHROPIC_API_KEY_ENV: &str = "ANTHROPIC_API_KEY"; +const ANTHROPIC_AUTH_TOKEN_ENV: &str = "ANTHROPIC_AUTH_TOKEN"; +const BETA_HEADER: &str = "anthropic-beta"; +const AUTHORIZATION: &str = "authorization"; +const API_KEY_HEADER: &str = "x-api-key"; +const DIRECT_BROWSER_ACCESS_HEADER: &str = "anthropic-dangerous-direct-browser-access"; + +fn header_value<'a>(headers: &'a [(String, String)], name: &str) -> Option<&'a str> { + headers + .iter() + .find(|(header, _)| header.eq_ignore_ascii_case(name)) + .map(|(_, value)| value.as_str()) +} + +fn without(headers: Headers, names: &[&str]) -> Headers { + headers + .into_iter() + .filter(|(header, _)| !names.iter().any(|name| header.eq_ignore_ascii_case(name))) + .collect() +} + +fn existing_betas(headers: &[(String, String)]) -> impl Iterator + '_ { + headers + .iter() + .filter(|(header, _)| header.eq_ignore_ascii_case(BETA_HEADER)) + .flat_map(|(_, value)| split_beta_values(Some(value))) +} + +fn with_oauth_bearer(headers: Headers, bearer: String) -> Headers { + let beta = + join_beta_values(existing_betas(&headers).chain([ANTHROPIC_OAUTH_BETA_HEADER.to_string()])); + without(headers, &[API_KEY_HEADER, AUTHORIZATION, BETA_HEADER]) + .into_iter() + .chain([ + (AUTHORIZATION.to_string(), bearer), + (BETA_HEADER.to_string(), beta), + (DIRECT_BROWSER_ACCESS_HEADER.to_string(), "true".to_string()), + ]) + .collect() +} + +fn non_empty(value: Option<&str>) -> Option<&str> { + value.map(str::trim).filter(|value| !value.is_empty()) +} + +pub fn authenticate( + headers: Headers, + api_key: Option<&str>, + env_lookup: &dyn Fn(&str) -> Option, +) -> Result { + if let Some(forwarded) = header_value(&headers, AUTHORIZATION) + && forwarded + .strip_prefix("Bearer ") + .is_some_and(|token| token.starts_with(ANTHROPIC_OAUTH_TOKEN_PREFIX)) + { + let bearer = forwarded.to_string(); + return Ok(with_oauth_bearer(headers, bearer)); + } + if let Some(key) = api_key.filter(|key| key.starts_with(ANTHROPIC_OAUTH_TOKEN_PREFIX)) { + return Ok(with_oauth_bearer(headers, format!("Bearer {key}"))); + } + if header_value(&headers, API_KEY_HEADER).is_some() + || header_value(&headers, AUTHORIZATION).is_some() + { + return Ok(headers); + } + let resolved_key = non_empty(api_key) + .map(str::to_string) + .or_else(|| env_lookup(ANTHROPIC_API_KEY_ENV).filter(|value| !value.trim().is_empty())); + let auth = match resolved_key { + Some(key) if is_anthropic_oauth_key(&key) => { + (AUTHORIZATION.to_string(), format!("Bearer {key}")) + } + Some(key) => (API_KEY_HEADER.to_string(), key), + None => match env_lookup(ANTHROPIC_AUTH_TOKEN_ENV).filter(|value| !value.trim().is_empty()) + { + Some(token) => (AUTHORIZATION.to_string(), format!("Bearer {token}")), + None => { + return Err(litellm_auth::Error::MissingApiKey { + provider: "Anthropic", + environment_variable: ANTHROPIC_API_KEY_ENV, + }); + } + }, + }; + Ok(headers.into_iter().chain([auth]).collect()) +} + +fn context_management_betas( + context_management: Option<&Value>, +) -> impl Iterator { + let edits = context_management + .and_then(|value| value.get("edits")) + .and_then(Value::as_array) + .map(Vec::as_slice) + .unwrap_or(&[]); + let (compact, other) = edits.iter().fold((false, false), |(compact, other), edit| { + match edit.get("type").and_then(Value::as_str) { + Some("compact_20260112") => (true, other), + _ => (compact, true), + } + }); + compact + .then_some(beta::COMPACT_2026_01_12) + .into_iter() + .chain(other.then_some(beta::CONTEXT_MANAGEMENT_2025_06_27)) +} + +fn uses_structured_output(request: &AnthropicMessagesRequest) -> bool { + request.output_format.is_some() + || request + .output_config + .as_ref() + .and_then(|config| config.get("format")) + .is_some_and(|format| !format.is_null()) +} + +fn messages_carry_output_config(request: &AnthropicMessagesRequest) -> bool { + request + .messages + .iter() + .any(|message| message.extra.contains_key("output_config")) +} + +pub fn feature_betas(request: &AnthropicMessagesRequest) -> Vec<&'static str> { + let tools = request.tools.as_deref(); + [ + requires_native_compaction_beta(request.compaction.as_ref(), &request.messages) + .then_some(beta::COMPACT_2026_09_04), + uses_structured_output(request).then_some(beta::STRUCTURED_OUTPUT), + (request.speed.as_deref() == Some("fast")).then_some(beta::FAST_MODE_2026_02_01), + messages_carry_output_config(request).then_some(beta::PER_TURN_CONTROL_2026_07_01), + has_advisor_tool(tools).then_some(beta::ADVISOR_TOOL_2026_03_01), + is_tool_search_used(tools).then_some(beta::ADVANCED_TOOL_USE_2025_11_20), + ] + .into_iter() + .flatten() + .chain(context_management_betas( + request.context_management.as_ref(), + )) + .collect() +} + +pub fn with_feature_betas(headers: Headers, request: &AnthropicMessagesRequest) -> Headers { + let existing = existing_betas(&headers).collect::>(); + let features = feature_betas(request); + if existing.is_empty() && features.is_empty() { + return headers; + } + let merged = join_beta_values( + existing + .into_iter() + .chain(features.into_iter().map(str::to_string)), + ); + without(headers, &[BETA_HEADER]) + .into_iter() + .chain([(BETA_HEADER.to_string(), merged)]) + .collect() +} + +#[cfg(test)] +mod tests { + use rstest::{fixture, rstest}; + use serde_json::json; + + use super::*; + + const OAUTH_TOKEN: &str = "sk-ant-oat01-token"; + const OAUTH_BEARER: &str = "Bearer sk-ant-oat01-token"; + const REGULAR_KEY: &str = "sk-ant-api03-regular"; + const BROWSER_ACCESS: (&str, &str) = ("anthropic-dangerous-direct-browser-access", "true"); + + type Env = &'static [(&'static str, &'static str)]; + + fn request(fields: Value) -> AnthropicMessagesRequest { + let mut body = + json!({"model": "claude", "messages": [{"role": "user", "content": "Hello"}]}); + body.as_object_mut() + .unwrap() + .extend(fields.as_object().unwrap().clone()); + serde_json::from_value(body).unwrap() + } + + fn headers(pairs: &[(&str, &str)]) -> Headers { + pairs + .iter() + .map(|(name, value)| (name.to_string(), value.to_string())) + .collect() + } + + fn betas(values: &[&str]) -> String { + values.join(",") + } + + #[fixture] + fn no_env() -> Env { + &[] + } + + #[fixture] + fn full_env() -> Env { + &[ + ("ANTHROPIC_API_KEY", "sk-env"), + ("ANTHROPIC_AUTH_TOKEN", "env-token"), + ] + } + + fn authenticate_with( + forwarded: &[(&str, &str)], + api_key: Option<&str>, + env: Env, + ) -> Result { + let lookup = |name: &str| { + env.iter() + .find(|(key, _)| *key == name) + .map(|(_, value)| value.to_string()) + }; + authenticate(headers(forwarded), api_key, &lookup) + } + + #[rstest] + #[case::forwarded_bearer_drops_forwarded_and_deployment_keys( + &[("X-Api-Key", REGULAR_KEY), ("Authorization", OAUTH_BEARER)], + Some(REGULAR_KEY), + OAUTH_BEARER, + &[], + )] + #[case::forwarded_bearer_in_uppercase_authorization_header( + &[("AUTHORIZATION", OAUTH_BEARER)], + None, + OAUTH_BEARER, + &[], + )] + #[case::forwarded_bearer_keeps_unrelated_headers_in_place( + &[("anthropic-version", "2023-06-01"), ("authorization", OAUTH_BEARER)], + None, + OAUTH_BEARER, + &[("anthropic-version", "2023-06-01")], + )] + #[case::forwarded_bearer_wins_over_an_oauth_api_key( + &[("authorization", OAUTH_BEARER)], + Some("sk-ant-oat01-deployment"), + OAUTH_BEARER, + &[], + )] + #[case::api_key_authenticates_as_a_bearer(&[], Some(OAUTH_TOKEN), OAUTH_BEARER, &[])] + #[case::api_key_removes_a_forwarded_x_api_key( + &[("x-api-key", OAUTH_TOKEN)], + Some(OAUTH_TOKEN), + OAUTH_BEARER, + &[], + )] + #[case::api_key_replaces_a_forwarded_non_oauth_bearer( + &[("Authorization", "Bearer some-proxy-token")], + Some(OAUTH_TOKEN), + OAUTH_BEARER, + &[], + )] + fn oauth_token_is_the_whole_credential( + #[case] forwarded: &[(&str, &str)], + #[case] api_key: Option<&str>, + #[case] expected_bearer: &str, + #[case] kept: &[(&str, &str)], + full_env: Env, + ) { + let expected = kept + .iter() + .copied() + .chain([ + ("authorization", expected_bearer), + ("anthropic-beta", ANTHROPIC_OAUTH_BETA_HEADER), + BROWSER_ACCESS, + ]) + .collect::>(); + assert_eq!( + authenticate_with(forwarded, api_key, full_env).unwrap(), + headers(&expected) + ); + } + + #[rstest] + #[case::forwarded_bearer_merges_a_differently_cased_beta_header( + &[("Anthropic-Beta", "web-search-2025-03-05"), ("authorization", OAUTH_BEARER)], + None, + )] + #[case::forwarded_bearer_dedupes_an_existing_oauth_beta( + &[("anthropic-beta", "web-search-2025-03-05, oauth-2025-04-20"), ("authorization", OAUTH_BEARER)], + None, + )] + #[case::api_key_merges_the_existing_beta_header( + &[("anthropic-beta", " web-search-2025-03-05 ,")], + Some(OAUTH_TOKEN), + )] + #[case::forwarded_bearer_unions_every_beta_header_casing( + &[("anthropic-beta", "oauth-2025-04-20"), ("ANTHROPIC-BETA", "web-search-2025-03-05"), ("authorization", OAUTH_BEARER)], + None, + )] + fn oauth_beta_merges_into_existing_betas( + #[case] forwarded: &[(&str, &str)], + #[case] api_key: Option<&str>, + no_env: Env, + ) { + assert_eq!( + authenticate_with(forwarded, api_key, no_env).unwrap(), + headers(&[ + ("authorization", OAUTH_BEARER), + ( + "anthropic-beta", + &betas(&[ANTHROPIC_OAUTH_BETA_HEADER, "web-search-2025-03-05"]) + ), + BROWSER_ACCESS, + ]) + ); + } + + #[rstest] + #[case::x_api_key_over_the_deployment_key(&[("x-api-key", "caller-key")], Some("sk-other"))] + #[case::uppercase_x_api_key(&[("X-API-KEY", "caller-key")], None)] + #[case::non_oauth_bearer(&[("Authorization", "Bearer some-proxy-token")], None)] + #[case::non_oauth_bearer_over_a_regular_api_key( + &[("authorization", "Bearer sk-ant-api03-forwarded")], + Some(REGULAR_KEY), + )] + #[case::oauth_token_without_the_bearer_scheme(&[("authorization", OAUTH_TOKEN)], None)] + #[case::oauth_token_behind_a_lowercase_bearer_scheme( + &[("authorization", "bearer sk-ant-oat01-token")], + None, + )] + fn forwarded_auth_header_is_kept_untouched( + #[case] forwarded: &[(&str, &str)], + #[case] api_key: Option<&str>, + full_env: Env, + ) { + assert_eq!( + authenticate_with(forwarded, api_key, full_env).unwrap(), + headers(forwarded) + ); + } + + #[rstest] + #[case::api_key_param(Some("sk-param"), &[], ("x-api-key", "sk-param"))] + #[case::api_key_param_over_env_key_and_auth_token( + Some("sk-param"), + &[("ANTHROPIC_API_KEY", "sk-env"), ("ANTHROPIC_AUTH_TOKEN", "env-token")], + ("x-api-key", "sk-param"), + )] + #[case::env_key_without_a_param(None, &[("ANTHROPIC_API_KEY", "sk-env")], ("x-api-key", "sk-env"))] + #[case::env_key_when_the_param_is_empty(Some(""), &[("ANTHROPIC_API_KEY", "sk-env")], ("x-api-key", "sk-env"))] + #[case::env_key_when_the_param_is_whitespace( + Some(" "), + &[("ANTHROPIC_API_KEY", "sk-env")], + ("x-api-key", "sk-env"), + )] + #[case::env_key_over_auth_token( + None, + &[("ANTHROPIC_API_KEY", "sk-env"), ("ANTHROPIC_AUTH_TOKEN", "env-token")], + ("x-api-key", "sk-env"), + )] + #[case::auth_token_as_a_bearer( + None, + &[("ANTHROPIC_AUTH_TOKEN", "env-token")], + ("authorization", "Bearer env-token"), + )] + #[case::auth_token_when_the_env_key_is_whitespace( + None, + &[("ANTHROPIC_API_KEY", " \t"), ("ANTHROPIC_AUTH_TOKEN", "env-token")], + ("authorization", "Bearer env-token"), + )] + #[case::oauth_env_key_as_a_plain_bearer( + None, + &[("ANTHROPIC_API_KEY", "sk-ant-oat01-env")], + ("authorization", "Bearer sk-ant-oat01-env"), + )] + fn credential_is_resolved_after_the_existing_headers( + #[case] api_key: Option<&str>, + #[case] env: Env, + #[case] expected: (&str, &str), + ) { + let forwarded = [("anthropic-beta", "web-search-2025-03-05")]; + assert_eq!( + authenticate_with(&forwarded, api_key, env).unwrap(), + headers(&[forwarded[0], expected]) + ); + } + + #[rstest] + #[case::no_credentials(&[], None, &[])] + #[case::empty_api_key(&[], Some(""), &[])] + #[case::whitespace_only_env_values( + &[], + None, + &[("ANTHROPIC_API_KEY", " "), ("ANTHROPIC_AUTH_TOKEN", " \t")], + )] + #[case::unrelated_forwarded_headers(&[("anthropic-beta", "web-search-2025-03-05")], None, &[])] + fn missing_credentials_are_an_auth_error( + #[case] forwarded: &[(&str, &str)], + #[case] api_key: Option<&str>, + #[case] env: Env, + ) { + assert!(matches!( + authenticate_with(forwarded, api_key, env), + Err(litellm_auth::Error::MissingApiKey { + provider: "Anthropic", + environment_variable: "ANTHROPIC_API_KEY", + }) + )); + } + + #[rstest] + #[case::no_features(json!({}), &[])] + #[case::output_format(json!({"output_format": {"type": "json_schema"}}), &[beta::STRUCTURED_OUTPUT])] + #[case::null_output_format(json!({"output_format": null}), &[])] + #[case::output_config_format( + json!({"output_config": {"format": {"type": "json_schema"}, "effort": "xhigh"}}), + &[beta::STRUCTURED_OUTPUT] + )] + #[case::null_output_config_format(json!({"output_config": {"format": null}}), &[])] + #[case::top_level_output_config_without_format(json!({"output_config": {"effort": "high"}}), &[])] + #[case::fast_speed(json!({"speed": "fast"}), &[beta::FAST_MODE_2026_02_01])] + #[case::standard_speed(json!({"speed": "standard"}), &[])] + #[case::compaction_param(json!({"compaction": {"enabled": true}}), &[beta::COMPACT_2026_09_04])] + #[case::empty_compaction_param(json!({"compaction": {}}), &[beta::COMPACT_2026_09_04])] + #[case::signed_compaction_block_in_history( + json!({"messages": [ + {"role": "assistant", "content": [{"type": "compaction", "content": "summary", "signature": "sig"}]}, + {"role": "user", "content": "Continue"}, + ]}), + &[beta::COMPACT_2026_09_04] + )] + #[case::unsigned_compaction_block_in_history( + json!({"messages": [ + {"role": "assistant", "content": [{"type": "compaction", "content": "summary", "signature": ""}]}, + {"role": "user", "content": "Continue"}, + ]}), + &[] + )] + #[case::advisor_tool( + json!({"tools": [{"type": "advisor_20260301", "name": "advisor", "model": "claude-opus-4-6"}]}), + &[beta::ADVISOR_TOOL_2026_03_01] + )] + #[case::no_tools(json!({"tools": []}), &[])] + #[case::regex_tool_search( + json!({"tools": [{"type": "tool_search_tool_regex_20251119"}]}), + &[beta::ADVANCED_TOOL_USE_2025_11_20] + )] + #[case::bm25_tool_search( + json!({"tools": [{"type": "tool_search_tool_bm25_20251119"}]}), + &[beta::ADVANCED_TOOL_USE_2025_11_20] + )] + #[case::unrelated_server_tool(json!({"tools": [{"type": "web_search_20250305", "name": "web_search"}]}), &[])] + #[case::only_compact_edits( + json!({"context_management": {"edits": [{"type": "compact_20260112"}]}}), + &[beta::COMPACT_2026_01_12] + )] + #[case::only_other_edits( + json!({"context_management": {"edits": [{"type": "clear_tool_uses_20250919", "keep": {"type": "tool_uses", "value": 3}}]}}), + &[beta::CONTEXT_MANAGEMENT_2025_06_27] + )] + #[case::compact_and_other_edits( + json!({"context_management": {"edits": [{"type": "compact_20260112"}, {"type": "clear_tool_uses_20250919"}]}}), + &[beta::COMPACT_2026_01_12, beta::CONTEXT_MANAGEMENT_2025_06_27] + )] + #[case::edit_without_a_type(json!({"context_management": {"edits": [{}]}}), &[beta::CONTEXT_MANAGEMENT_2025_06_27])] + #[case::empty_edits(json!({"context_management": {"edits": []}}), &[])] + #[case::context_management_without_edits(json!({"context_management": {}}), &[])] + #[case::per_message_output_config( + json!({"messages": [{"role": "user", "content": "hi", "output_config": {"effort": "low"}}]}), + &[beta::PER_TURN_CONTROL_2026_07_01] + )] + #[case::per_message_null_output_config( + json!({"messages": [{"role": "user", "content": "hi", "output_config": null}]}), + &[beta::PER_TURN_CONTROL_2026_07_01] + )] + fn feature_betas_follow_the_request(#[case] fields: Value, #[case] expected: &[&str]) { + assert_eq!(feature_betas(&request(fields)), expected); + } + + #[rstest] + #[case::no_betas(&[("x-api-key", "k"), ("anthropic-version", "2023-06-01")], json!({}))] + #[case::blank_beta_header(&[("Anthropic-Beta", " , "), ("x-api-key", "k")], json!({}))] + fn headers_without_any_beta_value_are_untouched( + #[case] input: &[(&str, &str)], + #[case] fields: Value, + ) { + assert_eq!( + with_feature_betas(headers(input), &request(fields)), + headers(input) + ); + } + + #[rstest] + #[case::feature_beta_is_appended( + &[("x-api-key", "k")], + json!({"speed": "fast"}), + &[("x-api-key", "k"), ("anthropic-beta", beta::FAST_MODE_2026_02_01)], + )] + #[case::existing_betas_are_normalized_without_features( + &[("Anthropic-Beta", "web-search-2025-03-05, interleaved-thinking-2025-05-14 ,web-search-2025-03-05"), ("x-api-key", "k")], + json!({}), + &[("x-api-key", "k"), ("anthropic-beta", "interleaved-thinking-2025-05-14,web-search-2025-03-05")], + )] + #[case::existing_advisor_beta_is_kept_without_an_advisor_tool( + &[("anthropic-beta", beta::ADVISOR_TOOL_2026_03_01)], + json!({"tools": []}), + &[("anthropic-beta", beta::ADVISOR_TOOL_2026_03_01)], + )] + #[case::feature_already_sent_is_not_duplicated( + &[("anthropic-beta", beta::FAST_MODE_2026_02_01)], + json!({"speed": "fast"}), + &[("anthropic-beta", beta::FAST_MODE_2026_02_01)], + )] + fn feature_betas_merge_into_the_headers( + #[case] input: &[(&str, &str)], + #[case] fields: Value, + #[case] expected: &[(&str, &str)], + ) { + assert_eq!( + with_feature_betas(headers(input), &request(fields)), + headers(expected) + ); + } + + #[test] + fn differently_cased_beta_header_is_replaced_by_one_sorted_header() { + let merged = with_feature_betas( + headers(&[("Anthropic-Beta", "interleaved-thinking-2025-05-14")]), + &request( + json!({"messages": [{"role": "system", "content": "env", "output_config": {"effort": "low"}}]}), + ), + ); + assert_eq!( + merged, + headers(&[( + "anthropic-beta", + &betas(&[ + "interleaved-thinking-2025-05-14", + beta::PER_TURN_CONTROL_2026_07_01 + ]) + )]) + ); + } + + #[test] + fn every_beta_header_casing_is_unioned_into_one_header() { + let merged = with_feature_betas( + headers(&[ + ("anthropic-beta", "interleaved-thinking-2025-05-14"), + ("Anthropic-Beta", "web-search-2025-03-05"), + ]), + &request(json!({"speed": "fast"})), + ); + assert_eq!( + merged, + headers(&[( + "anthropic-beta", + &betas(&[ + beta::FAST_MODE_2026_02_01, + "interleaved-thinking-2025-05-14", + "web-search-2025-03-05" + ]) + )]) + ); + } + + #[test] + fn unknown_client_betas_survive_alongside_the_added_one() { + let client_betas = [ + "claude-code-20250219", + "interleaved-thinking-2025-05-14", + beta::CONTEXT_MANAGEMENT_2025_06_27, + beta::PER_TURN_CONTROL_2026_07_01, + "effort-2025-11-24", + ]; + let merged = with_feature_betas( + headers(&[("anthropic-beta", &betas(&client_betas))]), + &request( + json!({"messages": [{"role": "user", "content": "hi", "output_config": {"effort": "low"}}]}), + ), + ); + assert_eq!( + merged, + headers(&[( + "anthropic-beta", + &betas(&[ + "claude-code-20250219", + beta::CONTEXT_MANAGEMENT_2025_06_27, + "effort-2025-11-24", + "interleaved-thinking-2025-05-14", + beta::PER_TURN_CONTROL_2026_07_01, + ]) + )]) + ); + } + + #[test] + fn every_feature_merges_with_the_oauth_beta_sorted_and_last() { + let oauth_headers = authenticate_with(&[], Some(OAUTH_TOKEN), &[]).unwrap(); + let all_features = request(json!({ + "compaction": {"enabled": true}, + "output_format": {"type": "json_schema"}, + "speed": "fast", + "tools": [{"type": "advisor_20260301"}, {"type": "tool_search_tool_bm25_20251119"}], + "context_management": {"edits": [{"type": "compact_20260112"}, {"type": "clear_thinking_20251015"}]}, + "messages": [{"role": "user", "content": "hi", "output_config": {"effort": "low"}}], + })); + assert_eq!( + with_feature_betas(oauth_headers, &all_features), + headers(&[ + ("authorization", OAUTH_BEARER), + BROWSER_ACCESS, + ( + "anthropic-beta", + &betas(&[ + beta::ADVANCED_TOOL_USE_2025_11_20, + beta::ADVISOR_TOOL_2026_03_01, + beta::COMPACT_2026_01_12, + beta::COMPACT_2026_09_04, + beta::CONTEXT_MANAGEMENT_2025_06_27, + beta::FAST_MODE_2026_02_01, + ANTHROPIC_OAUTH_BETA_HEADER, + beta::PER_TURN_CONTROL_2026_07_01, + beta::STRUCTURED_OUTPUT, + ]) + ), + ]) + ); + } +} diff --git a/litellm-rust/crates/llms/src/anthropic/experimental_pass_through/messages/mod.rs b/litellm-rust/crates/llms/src/anthropic/experimental_pass_through/messages/mod.rs index 481d98c4e9d..5adf5fda16f 100644 --- a/litellm-rust/crates/llms/src/anthropic/experimental_pass_through/messages/mod.rs +++ b/litellm-rust/crates/llms/src/anthropic/experimental_pass_through/messages/mod.rs @@ -1,2 +1,5 @@ +pub mod handler; +pub mod headers; pub mod streaming_iterator; +pub mod thinking; pub mod transformation; diff --git a/litellm-rust/crates/llms/src/anthropic/experimental_pass_through/messages/thinking.rs b/litellm-rust/crates/llms/src/anthropic/experimental_pass_through/messages/thinking.rs new file mode 100644 index 00000000000..ffa4c8ffeb8 --- /dev/null +++ b/litellm-rust/crates/llms/src/anthropic/experimental_pass_through/messages/thinking.rs @@ -0,0 +1,1182 @@ +use litellm_core_utils::settings::Lookup; +use litellm_types::llms::anthropic_messages::anthropic_request::AnthropicMessagesRequest; +use serde_json::{Map, Value, json}; + +use crate::{ + anthropic::common_utils::AnthropicModelCapabilities, base_llm::chat::transformation::Error, +}; + +pub const ANTHROPIC_MIN_THINKING_BUDGET_TOKENS: u64 = 1024; + +const EFFORT_NAMES: &str = "'minimal', 'low', 'medium', 'high', 'xhigh', 'max', 'none'"; + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct ThinkingBudgets { + pub minimal: u64, + pub low: u64, + pub medium: u64, + pub high: u64, + pub xhigh: u64, + pub max: u64, +} + +impl Default for ThinkingBudgets { + fn default() -> Self { + Self { + minimal: 128, + low: 1024, + medium: 2048, + high: 4096, + xhigh: 8192, + max: 16384, + } + } +} + +impl ThinkingBudgets { + pub fn from_lookup(env: &impl Lookup) -> Self { + let defaults = Self::default(); + let tier = |name: &str, default: u64| { + env.parsed::(&format!("DEFAULT_REASONING_EFFORT_{name}_THINKING_BUDGET")) + .unwrap_or(default) + }; + Self { + minimal: tier("MINIMAL", defaults.minimal), + low: tier("LOW", defaults.low), + medium: tier("MEDIUM", defaults.medium), + high: tier("HIGH", defaults.high), + xhigh: tier("XHIGH", defaults.xhigh), + max: tier("MAX", defaults.max), + } + } + + fn for_effort(&self, reasoning_effort: &str) -> Option { + match reasoning_effort { + "low" => Some(self.low), + "medium" => Some(self.medium), + "high" => Some(self.high), + "xhigh" => Some(self.xhigh), + "max" => Some(self.max), + "minimal" => Some(self.minimal.max(ANTHROPIC_MIN_THINKING_BUDGET_TOKENS)), + _ => None, + } + } + + fn effort_for_budget( + &self, + budget_tokens: u64, + capabilities: &AnthropicModelCapabilities, + ) -> &'static str { + if budget_tokens >= self.xhigh && capabilities.effort_tiers.xhigh { + return "xhigh"; + } + if budget_tokens >= self.high { + return "high"; + } + if budget_tokens >= self.medium { + return "medium"; + } + "low" + } +} + +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +pub struct ThinkingContext { + pub capabilities: AnthropicModelCapabilities, + pub budgets: ThinkingBudgets, +} + +fn bad_request(message: String) -> Error { + Error::InvalidRequest(message) +} + +fn thinking_type(thinking: Option<&Value>) -> Option<&str> { + thinking?.get("type")?.as_str() +} + +fn output_config_effort(output_config: Option<&Value>) -> Option<&str> { + output_config?.get("effort")?.as_str() +} + +fn enabled_thinking(budget_tokens: u64) -> Value { + json!({"type": "enabled", "budget_tokens": budget_tokens}) +} + +fn map_reasoning_effort( + reasoning_effort: &str, + context: &ThinkingContext, +) -> Result, Error> { + if reasoning_effort == "none" { + return Ok(None); + } + if context.capabilities.supports_adaptive_thinking { + return Ok(Some(json!({"type": "adaptive", "display": "summarized"}))); + } + context + .budgets + .for_effort(reasoning_effort) + .map(|budget| Some(enabled_thinking(budget))) + .ok_or_else(|| { + bad_request(format!( + "Unmapped reasoning effort: '{reasoning_effort}'. Must be one of: {EFFORT_NAMES}." + )) + }) +} + +fn cap_thinking_budget_to_max_tokens(thinking: Value, max_tokens: Option) -> Option { + let (Some(max_tokens), Some(budget)) = ( + max_tokens, + thinking.get("budget_tokens").and_then(Value::as_u64), + ) else { + return Some(thinking); + }; + if max_tokens <= ANTHROPIC_MIN_THINKING_BUDGET_TOKENS { + return None; + } + if budget < max_tokens { + return Some(thinking); + } + Some(enabled_thinking(max_tokens - 1)) +} + +fn reasoning_effort_to_output_config_effort(reasoning_effort: &str) -> Option<&'static str> { + match reasoning_effort { + "low" | "minimal" => Some("low"), + "medium" => Some("medium"), + "high" => Some("high"), + "xhigh" => Some("xhigh"), + "max" => Some("max"), + _ => None, + } +} + +fn with_default_effort(output_config: Option, effort: &str) -> Value { + let mut config = match output_config { + Some(Value::Object(config)) => config, + _ => Map::new(), + }; + if !config.contains_key("effort") { + config.insert("effort".to_string(), Value::String(effort.to_string())); + } + Value::Object(config) +} + +fn translate_reasoning_effort( + request: AnthropicMessagesRequest, + context: &ThinkingContext, +) -> Result { + let Some(reasoning_effort) = request.reasoning_effort.clone() else { + return Ok(request); + }; + let request = AnthropicMessagesRequest { + reasoning_effort: None, + ..request + }; + let Some(mapped) = map_reasoning_effort(&reasoning_effort, context)? else { + return Ok(AnthropicMessagesRequest { + thinking: None, + output_config: None, + ..request + }); + }; + let Some(fitted) = cap_thinking_budget_to_max_tokens(mapped, request.max_tokens) else { + return Ok(request); + }; + let thinking = Some(request.thinking.clone().unwrap_or(fitted)); + if !context.capabilities.supports_adaptive_thinking { + return Ok(AnthropicMessagesRequest { + thinking, + ..request + }); + } + let effort = reasoning_effort_to_output_config_effort(&reasoning_effort).ok_or_else(|| { + bad_request(format!( + "Invalid reasoning_effort: '{reasoning_effort}'. Must be one of: {EFFORT_NAMES}" + )) + })?; + if let Some(rejection) = context + .capabilities + .effort_level_rejection(effort, &request.model) + { + return Err(bad_request(rejection)); + } + Ok(AnthropicMessagesRequest { + thinking, + output_config: Some(with_default_effort(request.output_config.clone(), effort)), + ..request + }) +} + +fn drop_disabled_thinking( + request: AnthropicMessagesRequest, + context: &ThinkingContext, +) -> AnthropicMessagesRequest { + if !context.capabilities.thinking_always_on + || thinking_type(request.thinking.as_ref()) != Some("disabled") + { + return request; + } + AnthropicMessagesRequest { + thinking: None, + ..request + } +} + +fn translate_legacy_thinking_for_adaptive_model( + request: AnthropicMessagesRequest, + context: &ThinkingContext, +) -> AnthropicMessagesRequest { + let capabilities = &context.capabilities; + if !capabilities.supports_adaptive_thinking + || capabilities.supports_legacy_thinking + || thinking_type(request.thinking.as_ref()) != Some("enabled") + { + return request; + } + let budget = request + .thinking + .as_ref() + .and_then(|thinking| thinking.get("budget_tokens")) + .and_then(Value::as_u64) + .unwrap_or(0); + let effort = context.budgets.effort_for_budget(budget, capabilities); + AnthropicMessagesRequest { + thinking: Some(json!({"type": "adaptive"})), + output_config: Some(with_default_effort(request.output_config.clone(), effort)), + ..request + } +} + +fn output_config_without_effort(output_config: Option) -> Option { + let Some(Value::Object(config)) = output_config else { + return output_config; + }; + if !config.contains_key("effort") { + return Some(Value::Object(config)); + } + let residual: Map = config + .into_iter() + .filter(|(key, _)| key != "effort") + .collect(); + (!residual.is_empty()).then_some(Value::Object(residual)) +} + +fn translate_adaptive_effort_for_non_adaptive_model( + request: AnthropicMessagesRequest, + context: &ThinkingContext, +) -> Result { + let capabilities = &context.capabilities; + if capabilities.supports_adaptive_thinking { + return Ok(request); + } + let effort = output_config_effort(request.output_config.as_ref()).map(str::to_string); + let adaptive_thinking = thinking_type(request.thinking.as_ref()) == Some("adaptive"); + if effort.is_none() && !adaptive_thinking { + return Ok(request); + } + let level_supported = effort.as_deref().is_none_or(|effort| { + capabilities + .effort_level_rejection(effort, &request.model) + .is_none() + }); + if capabilities.supports_effort_param() && (!adaptive_thinking || level_supported) { + return Ok(AnthropicMessagesRequest { + thinking: if adaptive_thinking { + None + } else { + request.thinking.clone() + }, + ..request + }); + } + let legacy = if capabilities.supports_reasoning { + map_reasoning_effort( + effort + .as_deref() + .filter(|effort| !effort.is_empty()) + .unwrap_or("medium"), + context, + )? + } else { + None + }; + let capped = + legacy.and_then(|thinking| cap_thinking_budget_to_max_tokens(thinking, request.max_tokens)); + Ok(AnthropicMessagesRequest { + thinking: capped, + output_config: output_config_without_effort(request.output_config.clone()), + ..request + }) +} + +fn drop_incompatible_temperature_for_thinking( + request: AnthropicMessagesRequest, + context: &ThinkingContext, +) -> AnthropicMessagesRequest { + if context.capabilities.supports_adaptive_thinking { + return request; + } + let pinned = request + .temperature + .is_some_and(|temperature| temperature != 1.0); + let thinking_enabled = thinking_type(request.thinking.as_ref()) == Some("enabled"); + let effort_enabled = output_config_effort(request.output_config.as_ref()).is_some(); + if !pinned || !(thinking_enabled || effort_enabled) { + return request; + } + AnthropicMessagesRequest { + temperature: None, + ..request + } +} + +pub fn translate_thinking( + request: AnthropicMessagesRequest, + context: &ThinkingContext, +) -> Result { + let request = translate_reasoning_effort(request, context)?; + let request = drop_disabled_thinking(request, context); + let request = translate_legacy_thinking_for_adaptive_model(request, context); + let request = translate_adaptive_effort_for_non_adaptive_model(request, context)?; + Ok(drop_incompatible_temperature_for_thinking(request, context)) +} + +#[cfg(test)] +mod tests { + use rstest::{fixture, rstest}; + + use super::*; + use crate::anthropic::common_utils::SupportedEffortTiers; + + const EFFORT_CHOICES: &str = "'minimal', 'low', 'medium', 'high', 'xhigh', 'max', 'none'"; + + fn request(fields: Value) -> AnthropicMessagesRequest { + let mut body = + json!({"model": "claude", "messages": [{"role": "user", "content": "Hello"}]}); + body.as_object_mut() + .unwrap() + .extend(fields.as_object().unwrap().clone()); + serde_json::from_value(body).unwrap() + } + + fn context(capabilities: AnthropicModelCapabilities) -> ThinkingContext { + ThinkingContext { + capabilities, + budgets: ThinkingBudgets::default(), + } + } + + fn translate( + capabilities: AnthropicModelCapabilities, + fields: Value, + ) -> Result { + translate_thinking(request(fields), &context(capabilities)) + } + + fn overridden_budgets(overrides: &[(&str, &str)]) -> ThinkingBudgets { + let env = |name: &str| { + overrides + .iter() + .find(|(tier, _)| { + name == format!("DEFAULT_REASONING_EFFORT_{tier}_THINKING_BUDGET") + }) + .map(|(_, value)| value.to_string()) + }; + ThinkingBudgets::from_lookup(&env) + } + + fn claude_code_payload(effort: &str, max_tokens: u64) -> Value { + json!({"max_tokens": max_tokens, "thinking": {"type": "adaptive"}, "output_config": {"effort": effort}}) + } + + fn with_temperature(fields: Value, temperature: f64) -> Value { + let mut fields = fields; + fields + .as_object_mut() + .unwrap() + .insert("temperature".to_string(), json!(temperature)); + fields + } + + #[fixture] + fn haiku_3_5() -> AnthropicModelCapabilities { + AnthropicModelCapabilities::default() + } + + #[fixture] + fn haiku_4_5() -> AnthropicModelCapabilities { + AnthropicModelCapabilities { + supports_reasoning: true, + ..Default::default() + } + } + + #[fixture] + fn opus_4_5() -> AnthropicModelCapabilities { + AnthropicModelCapabilities { + supports_reasoning: true, + supports_output_config: true, + ..Default::default() + } + } + + #[fixture] + fn sonnet_4_6() -> AnthropicModelCapabilities { + AnthropicModelCapabilities { + supports_reasoning: true, + supports_adaptive_thinking: true, + supports_legacy_thinking: true, + supports_output_config: true, + effort_tiers: SupportedEffortTiers { + max: true, + ..Default::default() + }, + ..Default::default() + } + } + + #[fixture] + fn opus_4_7() -> AnthropicModelCapabilities { + AnthropicModelCapabilities { + supports_reasoning: true, + supports_adaptive_thinking: true, + supports_output_config: true, + effort_tiers: SupportedEffortTiers { + xhigh: true, + max: true, + ..Default::default() + }, + ..Default::default() + } + } + + #[fixture] + fn fable_5_1() -> AnthropicModelCapabilities { + AnthropicModelCapabilities { + thinking_always_on: true, + ..opus_4_7() + } + } + + #[fixture] + fn newfamily_6() -> AnthropicModelCapabilities { + AnthropicModelCapabilities { + supports_reasoning: true, + supports_adaptive_thinking: true, + ..Default::default() + } + } + + #[rstest] + #[case::minimal_maps_to_low(opus_4_7(), "minimal", "low")] + #[case::low(opus_4_7(), "low", "low")] + #[case::medium(opus_4_7(), "medium", "medium")] + #[case::high(opus_4_7(), "high", "high")] + #[case::xhigh_with_xhigh_tier(opus_4_7(), "xhigh", "xhigh")] + #[case::max(opus_4_7(), "max", "max")] + #[case::minimal_maps_to_low_on_4_6(sonnet_4_6(), "minimal", "low")] + #[case::low_on_4_6(sonnet_4_6(), "low", "low")] + #[case::max_without_max_tier_is_allowed_on_adaptive_models(newfamily_6(), "max", "max")] + fn reasoning_effort_on_adaptive_model_becomes_summarized_adaptive_thinking_and_effort( + #[case] capabilities: AnthropicModelCapabilities, + #[case] reasoning_effort: &str, + #[case] expected_effort: &str, + ) { + assert_eq!( + translate( + capabilities, + json!({"max_tokens": 1024, "reasoning_effort": reasoning_effort}) + ), + Ok(request(json!({ + "max_tokens": 1024, + "thinking": {"type": "adaptive", "display": "summarized"}, + "output_config": {"effort": expected_effort} + }))) + ); + } + + #[rstest] + #[case::adaptive_shape_is_not_dropped_for_small_max_tokens( + opus_4_7(), + json!({"max_tokens": 64, "reasoning_effort": "high"}), + json!({"max_tokens": 64, "thinking": {"type": "adaptive", "display": "summarized"}, "output_config": {"effort": "high"}}) + )] + #[case::caller_output_config_effort_wins( + opus_4_7(), + json!({"max_tokens": 1024, "reasoning_effort": "low", "output_config": {"effort": "max"}}), + json!({"max_tokens": 1024, "thinking": {"type": "adaptive", "display": "summarized"}, "output_config": {"effort": "max"}}) + )] + #[case::effort_merges_into_caller_output_config( + opus_4_7(), + json!({"max_tokens": 1024, "reasoning_effort": "high", "output_config": {"format": {"type": "json_schema"}}}), + json!({ + "max_tokens": 1024, + "thinking": {"type": "adaptive", "display": "summarized"}, + "output_config": {"format": {"type": "json_schema"}, "effort": "high"} + }) + )] + #[case::non_object_output_config_is_replaced( + opus_4_7(), + json!({"max_tokens": 1024, "reasoning_effort": "high", "output_config": "bogus"}), + json!({"max_tokens": 1024, "thinking": {"type": "adaptive", "display": "summarized"}, "output_config": {"effort": "high"}}) + )] + #[case::caller_thinking_and_output_config_win( + sonnet_4_6(), + json!({ + "max_tokens": 16000, + "reasoning_effort": "low", + "thinking": {"type": "enabled", "budget_tokens": 8000}, + "output_config": {"effort": "high"} + }), + json!({ + "max_tokens": 16000, + "thinking": {"type": "enabled", "budget_tokens": 8000}, + "output_config": {"effort": "high"} + }) + )] + #[case::caller_legacy_thinking_is_then_translated_while_reasoning_effort_level_stays( + opus_4_7(), + json!({"max_tokens": 16000, "reasoning_effort": "low", "thinking": {"type": "enabled", "budget_tokens": 8000}}), + json!({"max_tokens": 16000, "thinking": {"type": "adaptive"}, "output_config": {"effort": "low"}}) + )] + #[case::caller_disabled_thinking_is_kept_then_omitted_on_always_on_model( + fable_5_1(), + json!({"max_tokens": 1024, "reasoning_effort": "high", "thinking": {"type": "disabled"}}), + json!({"max_tokens": 1024, "output_config": {"effort": "high"}}) + )] + #[case::non_adaptive_model_gets_no_output_config( + opus_4_5(), + json!({"max_tokens": 8192, "reasoning_effort": "high"}), + json!({"max_tokens": 8192, "thinking": {"type": "enabled", "budget_tokens": 4096}}) + )] + #[case::caller_thinking_wins_on_non_adaptive_model( + opus_4_5(), + json!({"max_tokens": 16000, "reasoning_effort": "low", "thinking": {"type": "enabled", "budget_tokens": 8000}}), + json!({"max_tokens": 16000, "thinking": {"type": "enabled", "budget_tokens": 8000}}) + )] + #[case::caller_thinking_survives_when_mapped_budget_cannot_fit( + opus_4_5(), + json!({"max_tokens": 1024, "reasoning_effort": "low", "thinking": {"type": "enabled", "budget_tokens": 8000}}), + json!({"max_tokens": 1024, "thinking": {"type": "enabled", "budget_tokens": 8000}}) + )] + #[case::missing_max_tokens_leaves_budget_uncapped( + haiku_4_5(), + json!({"reasoning_effort": "high"}), + json!({"thinking": {"type": "enabled", "budget_tokens": 4096}}) + )] + #[case::budget_below_max_tokens_is_kept( + haiku_4_5(), + json!({"max_tokens": 4097, "reasoning_effort": "high"}), + json!({"max_tokens": 4097, "thinking": {"type": "enabled", "budget_tokens": 4096}}) + )] + #[case::budget_equal_to_max_tokens_is_capped( + haiku_4_5(), + json!({"max_tokens": 4096, "reasoning_effort": "high"}), + json!({"max_tokens": 4096, "thinking": {"type": "enabled", "budget_tokens": 4095}}) + )] + #[case::budget_above_max_tokens_is_capped( + haiku_4_5(), + json!({"max_tokens": 4000, "reasoning_effort": "xhigh"}), + json!({"max_tokens": 4000, "thinking": {"type": "enabled", "budget_tokens": 3999}}) + )] + #[case::max_tokens_just_above_min_budget_caps_to_min_budget( + haiku_4_5(), + json!({"max_tokens": 1025, "reasoning_effort": "xhigh"}), + json!({"max_tokens": 1025, "thinking": {"type": "enabled", "budget_tokens": 1024}}) + )] + #[case::max_tokens_at_min_budget_drops_thinking( + haiku_4_5(), + json!({"max_tokens": 1024, "reasoning_effort": "xhigh"}), + json!({"max_tokens": 1024}) + )] + #[case::pinned_temperature_is_dropped_after_thinking_is_synthesized( + haiku_4_5(), + json!({"max_tokens": 8192, "reasoning_effort": "low", "temperature": 0}), + json!({"max_tokens": 8192, "thinking": {"type": "enabled", "budget_tokens": 1024}}) + )] + fn reasoning_effort_is_translated( + #[case] capabilities: AnthropicModelCapabilities, + #[case] input: Value, + #[case] expected: Value, + ) { + assert_eq!(translate(capabilities, input), Ok(request(expected))); + } + + #[rstest] + #[case::minimal_floors_at_min_budget("minimal", 1024)] + #[case::low("low", 1024)] + #[case::medium("medium", 2048)] + #[case::high("high", 4096)] + #[case::xhigh("xhigh", 8192)] + #[case::max("max", 16384)] + fn reasoning_effort_on_non_adaptive_model_uses_the_tier_budget( + haiku_4_5: AnthropicModelCapabilities, + #[case] reasoning_effort: &str, + #[case] expected_budget: u64, + ) { + assert_eq!( + translate( + haiku_4_5, + json!({"max_tokens": 32000, "reasoning_effort": reasoning_effort}) + ), + Ok(request(json!({ + "max_tokens": 32000, + "thinking": {"type": "enabled", "budget_tokens": expected_budget} + }))) + ); + } + + #[rstest] + #[case::adaptive_model(opus_4_7())] + #[case::effort_capable_model(opus_4_5())] + #[case::budget_model(haiku_4_5())] + fn reasoning_effort_none_clears_thinking_and_output_config( + #[case] capabilities: AnthropicModelCapabilities, + ) { + assert_eq!( + translate( + capabilities, + json!({ + "max_tokens": 1024, + "reasoning_effort": "none", + "thinking": {"type": "adaptive"}, + "output_config": {"effort": "high"} + }) + ), + Ok(request(json!({"max_tokens": 1024}))) + ); + } + + #[rstest] + #[case::bogus_on_budget_model( + opus_4_5(), + json!({"max_tokens": 1024, "reasoning_effort": "bogus"}), + format!("Unmapped reasoning effort: 'bogus'. Must be one of: {EFFORT_CHOICES}.") + )] + #[case::disabled_on_budget_model( + haiku_4_5(), + json!({"max_tokens": 1024, "reasoning_effort": "disabled"}), + format!("Unmapped reasoning effort: 'disabled'. Must be one of: {EFFORT_CHOICES}.") + )] + #[case::empty_on_budget_model( + haiku_4_5(), + json!({"max_tokens": 1024, "reasoning_effort": ""}), + format!("Unmapped reasoning effort: ''. Must be one of: {EFFORT_CHOICES}.") + )] + #[case::invalid_on_adaptive_model( + opus_4_7(), + json!({"max_tokens": 1024, "reasoning_effort": "invalid"}), + format!("Invalid reasoning_effort: 'invalid'. Must be one of: {EFFORT_CHOICES}") + )] + #[case::disabled_on_adaptive_model( + opus_4_7(), + json!({"max_tokens": 1024, "reasoning_effort": "disabled"}), + format!("Invalid reasoning_effort: 'disabled'. Must be one of: {EFFORT_CHOICES}") + )] + #[case::empty_on_adaptive_model( + opus_4_7(), + json!({"max_tokens": 1024, "reasoning_effort": ""}), + format!("Invalid reasoning_effort: ''. Must be one of: {EFFORT_CHOICES}") + )] + #[case::xhigh_without_xhigh_tier_on_4_6( + sonnet_4_6(), + json!({"max_tokens": 1024, "reasoning_effort": "xhigh"}), + "effort='xhigh' is not supported by this model. Got model: claude".to_string() + )] + #[case::xhigh_without_xhigh_tier_on_unmapped_adaptive_model( + newfamily_6(), + json!({"max_tokens": 1024, "reasoning_effort": "xhigh"}), + "effort='xhigh' is not supported by this model. Got model: claude".to_string() + )] + #[case::unrecognized_adaptive_effort_on_budget_model( + haiku_4_5(), + claude_code_payload("turbo", 8192), + format!("Unmapped reasoning effort: 'turbo'. Must be one of: {EFFORT_CHOICES}.") + )] + fn unsupported_effort_is_a_request_error( + #[case] capabilities: AnthropicModelCapabilities, + #[case] input: Value, + #[case] expected_message: String, + ) { + assert_eq!( + translate(capabilities, input), + Err(Error::InvalidRequest(expected_message)) + ); + } + + #[rstest] + #[case::omitted_on_always_on_model(fable_5_1(), json!({"type": "disabled"}), None)] + #[case::kept_on_adaptive_model(opus_4_7(), json!({"type": "disabled"}), Some(json!({"type": "disabled"})))] + #[case::kept_on_budget_model(haiku_4_5(), json!({"type": "disabled"}), Some(json!({"type": "disabled"})))] + #[case::adaptive_kept_on_always_on_model( + fable_5_1(), + json!({"type": "adaptive"}), + Some(json!({"type": "adaptive"})) + )] + fn disabled_thinking_is_omitted_only_for_always_on_models( + #[case] capabilities: AnthropicModelCapabilities, + #[case] thinking: Value, + #[case] expected_thinking: Option, + ) { + let expected = match expected_thinking { + Some(thinking) => json!({"max_tokens": 64, "thinking": thinking}), + None => json!({"max_tokens": 64}), + }; + assert_eq!( + translate( + capabilities, + json!({"max_tokens": 64, "thinking": thinking}) + ), + Ok(request(expected)) + ); + } + + #[rstest] + #[case::far_above_xhigh_budget(opus_4_7(), json!(16384), "xhigh")] + #[case::at_xhigh_budget(opus_4_7(), json!(8192), "xhigh")] + #[case::below_xhigh_budget(opus_4_7(), json!(8191), "high")] + #[case::xhigh_budget_without_xhigh_tier(newfamily_6(), json!(8192), "high")] + #[case::large_budget_without_xhigh_tier(newfamily_6(), json!(31999), "high")] + #[case::at_high_budget(opus_4_7(), json!(4096), "high")] + #[case::below_high_budget(opus_4_7(), json!(4095), "medium")] + #[case::at_medium_budget(opus_4_7(), json!(2048), "medium")] + #[case::below_medium_budget(opus_4_7(), json!(2047), "low")] + #[case::tiny_budget(opus_4_7(), json!(1), "low")] + #[case::missing_budget(opus_4_7(), Value::Null, "low")] + #[case::always_on_model(fable_5_1(), json!(24000), "xhigh")] + fn legacy_thinking_is_bucketed_into_adaptive_effort_on_adaptive_only_models( + #[case] capabilities: AnthropicModelCapabilities, + #[case] budget_tokens: Value, + #[case] expected_effort: &str, + ) { + let thinking = match budget_tokens { + Value::Null => json!({"type": "enabled"}), + budget_tokens => json!({"type": "enabled", "budget_tokens": budget_tokens}), + }; + assert_eq!( + translate( + capabilities, + json!({"max_tokens": 1024, "thinking": thinking}) + ), + Ok(request(json!({ + "max_tokens": 1024, + "thinking": {"type": "adaptive"}, + "output_config": {"effort": expected_effort} + }))) + ); + } + + #[rstest] + #[case::verbatim_on_model_accepting_legacy_thinking( + sonnet_4_6(), + json!({"max_tokens": 1024, "thinking": {"type": "enabled", "budget_tokens": 31999}}), + json!({"max_tokens": 1024, "thinking": {"type": "enabled", "budget_tokens": 31999}}) + )] + #[case::verbatim_with_explicit_output_config_on_model_accepting_legacy_thinking( + sonnet_4_6(), + json!({"max_tokens": 1024, "thinking": {"type": "enabled", "budget_tokens": 31999}, "output_config": {"effort": "low"}}), + json!({"max_tokens": 1024, "thinking": {"type": "enabled", "budget_tokens": 31999}, "output_config": {"effort": "low"}}) + )] + #[case::verbatim_on_non_adaptive_model( + opus_4_5(), + json!({"max_tokens": 1024, "thinking": {"type": "enabled", "budget_tokens": 31999}}), + json!({"max_tokens": 1024, "thinking": {"type": "enabled", "budget_tokens": 31999}}) + )] + #[case::caller_output_config_effort_wins( + opus_4_7(), + json!({ + "max_tokens": 32000, + "thinking": {"type": "enabled", "budget_tokens": 31999}, + "output_config": {"effort": "low", "format": {"type": "json_schema"}} + }), + json!({ + "max_tokens": 32000, + "thinking": {"type": "adaptive"}, + "output_config": {"effort": "low", "format": {"type": "json_schema"}} + }) + )] + #[case::effort_merges_into_caller_output_config( + opus_4_7(), + json!({ + "max_tokens": 32000, + "thinking": {"type": "enabled", "budget_tokens": 4096}, + "output_config": {"format": {"type": "json_schema"}} + }), + json!({ + "max_tokens": 32000, + "thinking": {"type": "adaptive"}, + "output_config": {"effort": "high", "format": {"type": "json_schema"}} + }) + )] + #[case::adaptive_thinking_is_left_alone( + opus_4_7(), + json!({"max_tokens": 8192, "thinking": {"type": "adaptive", "display": "summarized"}}), + json!({"max_tokens": 8192, "thinking": {"type": "adaptive", "display": "summarized"}}) + )] + fn legacy_thinking_on_adaptive_capable_models( + #[case] capabilities: AnthropicModelCapabilities, + #[case] input: Value, + #[case] expected: Value, + ) { + assert_eq!(translate(capabilities, input), Ok(request(expected))); + } + + #[rstest] + #[case::bare_adaptive_becomes_medium_budget_on_budget_model( + haiku_4_5(), + json!({"max_tokens": 8192, "thinking": {"type": "adaptive"}}), + json!({"max_tokens": 8192, "thinking": {"type": "enabled", "budget_tokens": 2048}}) + )] + #[case::medium_effort_becomes_medium_budget_on_budget_model( + haiku_4_5(), + claude_code_payload("medium", 8192), + json!({"max_tokens": 8192, "thinking": {"type": "enabled", "budget_tokens": 2048}}) + )] + #[case::empty_effort_becomes_medium_budget_on_budget_model( + haiku_4_5(), + claude_code_payload("", 8192), + json!({"max_tokens": 8192, "thinking": {"type": "enabled", "budget_tokens": 2048}}) + )] + #[case::high_effort_becomes_high_budget_on_budget_model( + haiku_4_5(), + claude_code_payload("high", 8192), + json!({"max_tokens": 8192, "thinking": {"type": "enabled", "budget_tokens": 4096}}) + )] + #[case::effort_only_becomes_budget_on_budget_model( + haiku_4_5(), + json!({"max_tokens": 8192, "output_config": {"effort": "high"}}), + json!({"max_tokens": 8192, "thinking": {"type": "enabled", "budget_tokens": 4096}}) + )] + #[case::effort_replaces_caller_legacy_budget_on_budget_model( + haiku_4_5(), + json!({"max_tokens": 8192, "thinking": {"type": "enabled", "budget_tokens": 3000}, "output_config": {"effort": "high"}}), + json!({"max_tokens": 8192, "thinking": {"type": "enabled", "budget_tokens": 4096}}) + )] + #[case::residual_output_config_survives_effort_translation( + haiku_4_5(), + json!({ + "max_tokens": 8192, + "thinking": {"type": "adaptive"}, + "output_config": {"effort": "medium", "format": {"type": "json_schema"}} + }), + json!({ + "max_tokens": 8192, + "thinking": {"type": "enabled", "budget_tokens": 2048}, + "output_config": {"format": {"type": "json_schema"}} + }) + )] + #[case::effortless_output_config_is_kept( + haiku_4_5(), + json!({"max_tokens": 8192, "thinking": {"type": "adaptive"}, "output_config": {"format": {"type": "json_schema"}}}), + json!({ + "max_tokens": 8192, + "thinking": {"type": "enabled", "budget_tokens": 2048}, + "output_config": {"format": {"type": "json_schema"}} + }) + )] + #[case::empty_output_config_is_kept( + haiku_4_5(), + json!({"max_tokens": 8192, "thinking": {"type": "adaptive"}, "output_config": {}}), + json!({"max_tokens": 8192, "thinking": {"type": "enabled", "budget_tokens": 2048}, "output_config": {}}) + )] + #[case::missing_max_tokens_leaves_budget_uncapped( + haiku_4_5(), + json!({"thinking": {"type": "adaptive"}}), + json!({"thinking": {"type": "enabled", "budget_tokens": 2048}}) + )] + #[case::budget_is_capped_below_max_tokens( + haiku_4_5(), + claude_code_payload("high", 3000), + json!({"max_tokens": 3000, "thinking": {"type": "enabled", "budget_tokens": 2999}}) + )] + #[case::max_tokens_just_above_min_budget_caps_to_min_budget( + haiku_4_5(), + claude_code_payload("medium", 1025), + json!({"max_tokens": 1025, "thinking": {"type": "enabled", "budget_tokens": 1024}}) + )] + #[case::max_tokens_at_min_budget_drops_thinking_and_effort( + haiku_4_5(), + claude_code_payload("medium", 1024), + json!({"max_tokens": 1024}) + )] + #[case::max_tokens_below_min_budget_drops_thinking_and_effort( + haiku_4_5(), + claude_code_payload("medium", 512), + json!({"max_tokens": 512}) + )] + #[case::bare_adaptive_is_dropped_on_non_reasoning_model( + haiku_3_5(), + json!({"max_tokens": 8192, "thinking": {"type": "adaptive"}}), + json!({"max_tokens": 8192}) + )] + #[case::adaptive_and_effort_are_dropped_on_non_reasoning_model( + haiku_3_5(), + claude_code_payload("medium", 8192), + json!({"max_tokens": 8192}) + )] + #[case::effort_only_is_dropped_on_non_reasoning_model( + haiku_3_5(), + json!({"max_tokens": 8192, "output_config": {"effort": "high", "format": {"type": "json_schema"}}}), + json!({"max_tokens": 8192, "output_config": {"format": {"type": "json_schema"}}}) + )] + #[case::supported_effort_is_kept_and_adaptive_thinking_dropped_on_effort_model( + opus_4_5(), + claude_code_payload("medium", 8192), + json!({"max_tokens": 8192, "output_config": {"effort": "medium"}}) + )] + #[case::bare_adaptive_is_dropped_on_effort_model( + opus_4_5(), + json!({"max_tokens": 8192, "thinking": {"type": "adaptive"}}), + json!({"max_tokens": 8192}) + )] + #[case::effort_only_is_left_alone_on_effort_model( + opus_4_5(), + json!({"max_tokens": 8192, "output_config": {"effort": "high"}}), + json!({"max_tokens": 8192, "output_config": {"effort": "high"}}) + )] + #[case::unsupported_effort_only_is_left_for_provider_normalization( + opus_4_5(), + json!({"max_tokens": 4096, "output_config": {"effort": "xhigh"}}), + json!({"max_tokens": 4096, "output_config": {"effort": "xhigh"}}) + )] + #[case::legacy_thinking_is_kept_beside_native_effort_on_effort_model( + opus_4_5(), + json!({"max_tokens": 8192, "thinking": {"type": "enabled", "budget_tokens": 4096}, "output_config": {"effort": "high"}}), + json!({"max_tokens": 8192, "thinking": {"type": "enabled", "budget_tokens": 4096}, "output_config": {"effort": "high"}}) + )] + #[case::unsupported_xhigh_with_adaptive_thinking_falls_back_to_budget( + opus_4_5(), + claude_code_payload("xhigh", 64000), + json!({"max_tokens": 64000, "thinking": {"type": "enabled", "budget_tokens": 8192}}) + )] + #[case::unsupported_max_with_adaptive_thinking_falls_back_to_budget( + opus_4_5(), + claude_code_payload("max", 64000), + json!({"max_tokens": 64000, "thinking": {"type": "enabled", "budget_tokens": 16384}}) + )] + #[case::bare_adaptive_is_native_on_4_6( + sonnet_4_6(), + json!({"max_tokens": 8192, "thinking": {"type": "adaptive"}}), + json!({"max_tokens": 8192, "thinking": {"type": "adaptive"}}) + )] + #[case::adaptive_payload_is_native_on_4_6( + sonnet_4_6(), + claude_code_payload("high", 8192), + claude_code_payload("high", 8192) + )] + #[case::request_without_adaptive_interface_is_left_alone( + haiku_4_5(), + json!({"max_tokens": 1024}), + json!({"max_tokens": 1024}) + )] + fn adaptive_interface_is_reshaped_for_non_adaptive_models( + #[case] capabilities: AnthropicModelCapabilities, + #[case] input: Value, + #[case] expected: Value, + ) { + assert_eq!(translate(capabilities, input), Ok(request(expected))); + } + + #[rstest] + #[case::adaptive_downgraded_to_enabled_thinking( + haiku_4_5(), + claude_code_payload("medium", 8192), + 0.0, + json!({"max_tokens": 8192, "thinking": {"type": "enabled", "budget_tokens": 2048}}) + )] + #[case::bare_adaptive_downgraded_to_enabled_thinking( + haiku_4_5(), + json!({"max_tokens": 8192, "thinking": {"type": "adaptive"}}), + 0.0, + json!({"max_tokens": 8192, "thinking": {"type": "enabled", "budget_tokens": 2048}}) + )] + #[case::reasoning_effort_synthesized_enabled_thinking( + haiku_4_5(), + json!({"max_tokens": 8192, "reasoning_effort": "high"}), + 0.2, + json!({"max_tokens": 8192, "thinking": {"type": "enabled", "budget_tokens": 4096}}) + )] + #[case::above_one_with_enabled_thinking( + haiku_4_5(), + json!({"max_tokens": 8192, "thinking": {"type": "enabled", "budget_tokens": 2048}}), + 1.5, + json!({"max_tokens": 8192, "thinking": {"type": "enabled", "budget_tokens": 2048}}) + )] + #[case::native_effort_kept_on_effort_model( + opus_4_5(), + claude_code_payload("medium", 8192), + 0.0, + json!({"max_tokens": 8192, "output_config": {"effort": "medium"}}) + )] + #[case::effort_only_on_effort_model( + opus_4_5(), + json!({"max_tokens": 8192, "output_config": {"effort": "high"}}), + 0.0, + json!({"max_tokens": 8192, "output_config": {"effort": "high"}}) + )] + fn pinned_temperature_is_dropped_when_thinking_or_effort_survives_on_non_adaptive_model( + #[case] capabilities: AnthropicModelCapabilities, + #[case] input: Value, + #[case] temperature: f64, + #[case] expected: Value, + ) { + assert_eq!( + translate(capabilities, with_temperature(input, temperature)), + Ok(request(expected)) + ); + } + + #[rstest] + #[case::temperature_one_with_enabled_thinking( + haiku_4_5(), + claude_code_payload("medium", 8192), + 1.0, + json!({"max_tokens": 8192, "thinking": {"type": "enabled", "budget_tokens": 2048}}) + )] + #[case::thinking_dropped_for_small_max_tokens( + haiku_4_5(), + claude_code_payload("medium", 512), + 0.0, + json!({"max_tokens": 512}) + )] + #[case::thinking_dropped_on_non_reasoning_model( + haiku_3_5(), + claude_code_payload("medium", 8192), + 0.0, + json!({"max_tokens": 8192}) + )] + #[case::disabled_thinking( + haiku_4_5(), + json!({"max_tokens": 8192, "thinking": {"type": "disabled"}}), + 0.0, + json!({"max_tokens": 8192, "thinking": {"type": "disabled"}}) + )] + #[case::no_thinking(haiku_4_5(), json!({"max_tokens": 8192}), 0.0, json!({"max_tokens": 8192}))] + #[case::output_config_without_effort( + haiku_4_5(), + json!({"max_tokens": 8192, "output_config": {"format": {"type": "json_schema"}}}), + 0.0, + json!({"max_tokens": 8192, "output_config": {"format": {"type": "json_schema"}}}) + )] + #[case::adaptive_model( + opus_4_7(), + claude_code_payload("medium", 8192), + 0.0, + claude_code_payload("medium", 8192) + )] + #[case::legacy_thinking_on_adaptive_model( + sonnet_4_6(), + json!({"max_tokens": 8192, "thinking": {"type": "enabled", "budget_tokens": 4096}}), + 0.0, + json!({"max_tokens": 8192, "thinking": {"type": "enabled", "budget_tokens": 4096}}) + )] + fn temperature_is_kept( + #[case] capabilities: AnthropicModelCapabilities, + #[case] input: Value, + #[case] temperature: f64, + #[case] expected: Value, + ) { + assert_eq!( + translate(capabilities, with_temperature(input, temperature)), + Ok(request(with_temperature(expected, temperature))) + ); + } + + #[rstest] + #[case::minimal("MINIMAL", ThinkingBudgets { minimal: 5000, ..ThinkingBudgets::default() })] + #[case::low("LOW", ThinkingBudgets { low: 5000, ..ThinkingBudgets::default() })] + #[case::medium("MEDIUM", ThinkingBudgets { medium: 5000, ..ThinkingBudgets::default() })] + #[case::high("HIGH", ThinkingBudgets { high: 5000, ..ThinkingBudgets::default() })] + #[case::xhigh("XHIGH", ThinkingBudgets { xhigh: 5000, ..ThinkingBudgets::default() })] + #[case::max("MAX", ThinkingBudgets { max: 5000, ..ThinkingBudgets::default() })] + fn each_tier_budget_reads_only_its_own_environment_override( + #[case] tier: &str, + #[case] expected: ThinkingBudgets, + ) { + assert_eq!(overridden_budgets(&[(tier, "5000")]), expected); + } + + #[rstest] + #[case::whitespace_is_trimmed(" 6000 ", 6000)] + #[case::unparseable_value_keeps_default("lots", 4096)] + fn environment_override_parsing(#[case] raw: &str, #[case] expected_high: u64) { + assert_eq!( + overridden_budgets(&[("HIGH", raw)]), + ThinkingBudgets { + high: expected_high, + ..ThinkingBudgets::default() + } + ); + } + + #[rstest] + #[case::reasoning_effort_uses_overridden_budget( + &[("HIGH", "6000")], + haiku_4_5(), + json!({"max_tokens": 32000, "reasoning_effort": "high"}), + json!({"max_tokens": 32000, "thinking": {"type": "enabled", "budget_tokens": 6000}}) + )] + #[case::minimal_override_below_min_budget_is_floored( + &[("MINIMAL", "512")], + haiku_4_5(), + json!({"max_tokens": 32000, "reasoning_effort": "minimal"}), + json!({"max_tokens": 32000, "thinking": {"type": "enabled", "budget_tokens": 1024}}) + )] + #[case::minimal_override_above_min_budget_is_used( + &[("MINIMAL", "2000")], + haiku_4_5(), + json!({"max_tokens": 32000, "reasoning_effort": "minimal"}), + json!({"max_tokens": 32000, "thinking": {"type": "enabled", "budget_tokens": 2000}}) + )] + #[case::adaptive_fallback_uses_overridden_medium_budget( + &[("MEDIUM", "3000")], + haiku_4_5(), + json!({"max_tokens": 32000, "thinking": {"type": "adaptive"}}), + json!({"max_tokens": 32000, "thinking": {"type": "enabled", "budget_tokens": 3000}}) + )] + #[case::legacy_bucket_below_overridden_high_budget( + &[("HIGH", "6000")], + opus_4_7(), + json!({"max_tokens": 32000, "thinking": {"type": "enabled", "budget_tokens": 5999}}), + json!({"max_tokens": 32000, "thinking": {"type": "adaptive"}, "output_config": {"effort": "medium"}}) + )] + #[case::legacy_bucket_at_overridden_high_budget( + &[("HIGH", "6000")], + opus_4_7(), + json!({"max_tokens": 32000, "thinking": {"type": "enabled", "budget_tokens": 6000}}), + json!({"max_tokens": 32000, "thinking": {"type": "adaptive"}, "output_config": {"effort": "high"}}) + )] + #[case::legacy_bucket_below_overridden_xhigh_budget( + &[("XHIGH", "20000")], + opus_4_7(), + json!({"max_tokens": 32000, "thinking": {"type": "enabled", "budget_tokens": 19999}}), + json!({"max_tokens": 32000, "thinking": {"type": "adaptive"}, "output_config": {"effort": "high"}}) + )] + #[case::legacy_bucket_at_overridden_medium_budget( + &[("MEDIUM", "3000")], + opus_4_7(), + json!({"max_tokens": 32000, "thinking": {"type": "enabled", "budget_tokens": 3000}}), + json!({"max_tokens": 32000, "thinking": {"type": "adaptive"}, "output_config": {"effort": "medium"}}) + )] + #[case::legacy_bucket_below_overridden_medium_budget( + &[("MEDIUM", "3000")], + opus_4_7(), + json!({"max_tokens": 32000, "thinking": {"type": "enabled", "budget_tokens": 2999}}), + json!({"max_tokens": 32000, "thinking": {"type": "adaptive"}, "output_config": {"effort": "low"}}) + )] + fn translation_honors_budget_overrides( + #[case] overrides: &[(&str, &str)], + #[case] capabilities: AnthropicModelCapabilities, + #[case] input: Value, + #[case] expected: Value, + ) { + let context = ThinkingContext { + capabilities, + budgets: overridden_budgets(overrides), + }; + assert_eq!( + translate_thinking(request(input), &context), + Ok(request(expected)) + ); + } +} diff --git a/litellm-rust/crates/llms/src/anthropic/experimental_pass_through/messages/transformation.rs b/litellm-rust/crates/llms/src/anthropic/experimental_pass_through/messages/transformation.rs index c791749ac6d..59280c04a70 100644 --- a/litellm-rust/crates/llms/src/anthropic/experimental_pass_through/messages/transformation.rs +++ b/litellm-rust/crates/llms/src/anthropic/experimental_pass_through/messages/transformation.rs @@ -1,9 +1,28 @@ -use crate::base_llm::{ - anthropic_messages::transformation::BaseAnthropicMessagesConfig, chat::transformation::Error, +use litellm_core_utils::settings::{Lookup, ProcessEnvironment}; +use litellm_types::llms::anthropic_messages::anthropic_request::AnthropicMessagesRequest; +use serde_json::{Map, Value, json}; + +use super::{ + headers::{authenticate, with_feature_betas}, + thinking::{ThinkingBudgets, ThinkingContext, translate_thinking}, +}; +use crate::{ + anthropic::common_utils::{ + AnthropicModelCapabilities, has_advisor_tool, strip_advisor_blocks, + strip_encrypted_reasoning_blocks, + }, + base_llm::{ + anthropic_messages::transformation::{ + BaseAnthropicMessagesConfig, Headers, MessagesTransformContext, + }, + chat::transformation::Error, + }, }; const ANTHROPIC_API_KEY_ENV: &str = "ANTHROPIC_API_KEY"; +const ANTHROPIC_AUTH_TOKEN_ENV: &str = "ANTHROPIC_AUTH_TOKEN"; const ANTHROPIC_API_BASE_ENV: &str = "ANTHROPIC_API_BASE"; +const ANTHROPIC_BASE_URL_ENV: &str = "ANTHROPIC_BASE_URL"; const DEFAULT_ANTHROPIC_API_BASE: &str = "https://api.anthropic.com"; const MESSAGES_PATH_SUFFIX: &str = "/v1/messages"; @@ -11,6 +30,26 @@ pub struct AnthropicMessagesConfig; pub const ANTHROPIC_MESSAGES_CONFIG: AnthropicMessagesConfig = AnthropicMessagesConfig; +impl MessagesTransformContext { + pub fn new(capabilities: AnthropicModelCapabilities, drop_params: bool) -> Self { + Self::with_lookup(capabilities, drop_params, &ProcessEnvironment) + } + + pub fn with_lookup( + capabilities: AnthropicModelCapabilities, + drop_params: bool, + env: &impl Lookup, + ) -> Self { + Self { + thinking: ThinkingContext { + capabilities, + budgets: ThinkingBudgets::from_lookup(env), + }, + drop_params, + } + } +} + impl BaseAnthropicMessagesConfig for AnthropicMessagesConfig { fn get_complete_url( &self, @@ -21,6 +60,35 @@ impl BaseAnthropicMessagesConfig for AnthropicMessagesConfig { Ok(complete_anthropic_url(api_base, env_lookup)) } + fn transform_anthropic_messages_request( + &self, + request: AnthropicMessagesRequest, + context: &MessagesTransformContext, + ) -> Result { + if request.max_tokens.is_none() { + return Err(Error::InvalidRequest( + "max_tokens is required for Anthropic /v1/messages API".to_string(), + )); + } + let request = drop_unsupported_params(request, context)?; + let request = translate_thinking(request, &context.thinking)?; + let context_management = request + .context_management + .as_ref() + .and_then(map_openai_context_management_to_anthropic) + .or_else(|| request.context_management.clone()); + let messages = if has_advisor_tool(request.tools.as_deref()) { + request.messages + } else { + strip_advisor_blocks(request.messages) + }; + Ok(AnthropicMessagesRequest { + messages: strip_encrypted_reasoning_blocks(messages), + context_management, + ..request + }) + } + fn resolve_api_key( &self, api_key: Option<&str>, @@ -28,6 +96,113 @@ impl BaseAnthropicMessagesConfig for AnthropicMessagesConfig { ) -> Result { resolve_anthropic_api_key(api_key, env_lookup).map_err(Error::from) } + + fn secret_names(&self) -> &'static [&'static str] { + &[ + ANTHROPIC_API_KEY_ENV, + ANTHROPIC_AUTH_TOKEN_ENV, + ANTHROPIC_API_BASE_ENV, + ANTHROPIC_BASE_URL_ENV, + ] + } + + fn authenticate( + &self, + headers: Headers, + api_key: Option<&str>, + env_lookup: &dyn Fn(&str) -> Option, + ) -> Result { + authenticate(headers, api_key, env_lookup).map_err(Error::from) + } + + fn request_headers(&self, headers: Headers, request: &AnthropicMessagesRequest) -> Headers { + with_feature_betas(headers, request) + } +} + +fn unsupported_param(model: &str, param: &str, value: &str, hint: &str) -> Error { + Error::InvalidRequest(format!( + "{model} does not support {param}={value}. {hint}To drop unsupported params, set `litellm.drop_params = True`." + )) +} + +fn drop_unsupported_params( + request: AnthropicMessagesRequest, + context: &MessagesTransformContext, +) -> Result { + let capabilities = &context.thinking.capabilities; + let model = request.model.clone(); + let reject = |param: &str, value: String, hint: &str| -> Result<(), Error> { + if context.drop_params { + return Ok(()); + } + Err(unsupported_param(&model, param, &value, hint)) + }; + let speed = match request.speed.as_deref() { + Some(speed) if !capabilities.supports_speed => { + reject("speed", format!("'{speed}'"), "")?; + None + } + _ => request.speed.clone(), + }; + if capabilities.supports_sampling_params { + return Ok(AnthropicMessagesRequest { speed, ..request }); + } + let temperature = match request.temperature { + Some(temperature) if temperature != 1.0 => { + reject( + "temperature", + json!(temperature).to_string(), + "Only temperature=1 is supported. ", + )?; + None + } + temperature => temperature, + }; + if let Some(top_p) = request.top_p { + reject("top_p", json!(top_p).to_string(), "")?; + } + if let Some(top_k) = request.top_k { + reject("top_k", json!(top_k).to_string(), "")?; + } + Ok(AnthropicMessagesRequest { + speed, + temperature, + top_p: None, + top_k: None, + ..request + }) +} + +pub fn map_openai_context_management_to_anthropic(context_management: &Value) -> Option { + match context_management { + Value::Object(edits) if edits.contains_key("edits") => Some(context_management.clone()), + Value::Array(entries) => { + let edits: Vec = entries + .iter() + .filter_map(Value::as_object) + .filter(|entry| entry.get("type").and_then(Value::as_str) == Some("compaction")) + .map(|entry| { + let trigger = entry.get("compact_threshold").and_then(Value::as_f64).map( + |threshold| json!({"type": "input_tokens", "value": threshold as i64}), + ); + let passthrough = entry + .iter() + .filter(|(key, _)| !matches!(key.as_str(), "type" | "compact_threshold")) + .map(|(key, value)| (key.clone(), value.clone())); + Value::Object( + [("type".to_string(), json!("compact_20260112"))] + .into_iter() + .chain(trigger.map(|trigger| ("trigger".to_string(), trigger))) + .chain(passthrough) + .collect::>(), + ) + }) + .collect(); + (!edits.is_empty()).then(|| json!({"edits": edits})) + } + _ => None, + } } pub fn non_empty(value: Option<&str>) -> Option<&str> { @@ -64,70 +239,619 @@ pub fn resolve_anthropic_api_base( api_base: Option<&str>, env_lookup: &dyn Fn(&str) -> Option, ) -> String { + let env = |name: &str| env_lookup(name).filter(|value| !value.trim().is_empty()); non_empty(api_base) .map(str::to_string) - .or_else(|| env_lookup(ANTHROPIC_API_BASE_ENV).filter(|value| !value.trim().is_empty())) + .or_else(|| env(ANTHROPIC_API_BASE_ENV)) + .or_else(|| env(ANTHROPIC_BASE_URL_ENV)) .unwrap_or_else(|| DEFAULT_ANTHROPIC_API_BASE.to_string()) } #[cfg(test)] mod tests { + use std::process::Command; + + use rstest::{fixture, rstest}; + use super::*; + use crate::anthropic::common_utils::{ENCRYPTED_REASONING_SIGNATURE_PREFIX, beta}; - #[test] - fn url_defaults_to_public_anthropic_endpoint() { + type Env = &'static [(&'static str, &'static str)]; + + const BOTH_BASE_ENVS: Env = &[ + (ANTHROPIC_API_BASE_ENV, "https://api-base.example.com"), + (ANTHROPIC_BASE_URL_ENV, "https://base-url.example.com"), + ]; + const API_KEY_ENV: Env = &[(ANTHROPIC_API_KEY_ENV, "sk-env")]; + const MISSING_API_KEY: &str = + "Missing Anthropic API Key - Set `api_key` or the ANTHROPIC_API_KEY environment variable"; + const LOW_BUDGET_ENV: &str = "DEFAULT_REASONING_EFFORT_LOW_THINKING_BUDGET"; + const PROCESS_ENV_PROBE: &str = "LITELLM_MESSAGES_TRANSFORM_CONTEXT_PROBE"; + + fn merged(base: Value, fields: Value) -> Value { + Value::Object( + base.as_object() + .unwrap() + .clone() + .into_iter() + .chain(fields.as_object().unwrap().clone()) + .collect(), + ) + } + + fn body(fields: Value) -> Value { + merged( + json!({ + "model": "claude", + "max_tokens": 1024, + "messages": [{"role": "user", "content": "Hello"}] + }), + fields, + ) + } + + fn request(fields: Value) -> AnthropicMessagesRequest { + serde_json::from_value(body(fields)).unwrap() + } + + fn no_env(_: &str) -> Option { + None + } + + fn env(vars: Env) -> impl Fn(&str) -> Option { + move |name| { + vars.iter() + .find(|(key, _)| *key == name) + .map(|(_, value)| value.to_string()) + } + } + + fn headers(pairs: &[(&str, &str)]) -> Headers { + pairs + .iter() + .map(|(name, value)| (name.to_string(), value.to_string())) + .collect() + } + + fn transform( + fields: Value, + capabilities: AnthropicModelCapabilities, + drop_params: bool, + ) -> Result { + ANTHROPIC_MESSAGES_CONFIG + .transform_anthropic_messages_request( + request(fields), + &MessagesTransformContext::with_lookup(capabilities, drop_params, &no_env), + ) + .map(|transformed| serde_json::to_value(transformed).unwrap()) + } + + fn invalid(message: &str) -> Result { + Err(Error::InvalidRequest(message.to_string())) + } + + fn advisor_history() -> Value { + json!([ + {"role": "user", "content": "Build a worker pool."}, + {"role": "assistant", "content": [ + {"type": "text", "text": "Let me consult the advisor."}, + {"type": "server_tool_use", "id": "srvtoolu_abc123", "name": "advisor", "input": {}}, + {"type": "advisor_tool_result", "tool_use_id": "srvtoolu_abc123", "content": {"type": "advisor_result", "text": "Use channels."}}, + {"type": "text", "text": "Here is the implementation."} + ]} + ]) + } + + #[fixture] + fn unmapped() -> AnthropicModelCapabilities { + AnthropicModelCapabilities::default() + } + + #[fixture] + fn sampling_removed() -> AnthropicModelCapabilities { + AnthropicModelCapabilities { + supports_sampling_params: false, + ..Default::default() + } + } + + #[fixture] + fn fast_mode() -> AnthropicModelCapabilities { + AnthropicModelCapabilities { + supports_speed: true, + ..Default::default() + } + } + + #[rstest] + #[case::alone(json!({"max_tokens": null}))] + #[case::ahead_of_the_param_gate(json!({"max_tokens": null, "speed": "fast"}))] + fn missing_max_tokens_is_rejected(#[case] fields: Value, unmapped: AnthropicModelCapabilities) { assert_eq!( - complete_anthropic_url(None, &|_| None), - "https://api.anthropic.com/v1/messages" + transform(fields, unmapped, false), + invalid("max_tokens is required for Anthropic /v1/messages API") + ); + } + + #[rstest] + #[case::sampling_params_on_a_sampling_model( + unmapped(), + false, + json!({"temperature": 0.3, "top_p": 0.9, "top_k": 40}) + )] + #[case::sampling_params_on_a_sampling_model_under_drop_params( + unmapped(), + true, + json!({"temperature": 0.3, "top_p": 0.9, "top_k": 40}) + )] + #[case::unit_temperature_on_a_sampling_removed_model( + sampling_removed(), + false, + json!({"temperature": 1.0}) + )] + #[case::unit_temperature_on_a_sampling_removed_model_under_drop_params( + sampling_removed(), + true, + json!({"temperature": 1.0}) + )] + #[case::speed_on_a_fast_mode_model(fast_mode(), false, json!({"speed": "fast"}))] + #[case::speed_on_a_fast_mode_model_under_drop_params(fast_mode(), true, json!({"speed": "fast"}))] + #[case::native_context_management_edits(unmapped(), false, json!({"context_management": {"edits": [{ + "type": "clear_tool_uses_20250919", + "trigger": {"type": "input_tokens", "value": 30000}, + "keep": {"type": "tool_uses", "value": 3}, + "clear_at_least": {"type": "input_tokens", "value": 5000}, + "exclude_tools": ["web_search"], + "clear_tool_inputs": false + }]}}))] + #[case::first_party_billing_header_system_block(unmapped(), false, json!({"system": [ + {"type": "text", "text": "x-anthropic-billing-header: cc_version=1"}, + {"type": "text", "text": "real system prompt"} + ]}))] + #[case::anthropic_signed_reasoning_history(unmapped(), false, json!({"messages": [ + {"role": "user", "content": "Solve it."}, + {"role": "assistant", "content": [ + {"type": "thinking", "thinking": "plan", "signature": "EqQBCkYIAxgCIkA_anthropic_signed"}, + {"type": "redacted_thinking", "data": "EmwKAhgBEgy_anthropic_minted"}, + {"type": "text", "text": "The answer."} + ]} + ]}))] + #[case::advisor_history_alongside_the_advisor_tool(unmapped(), false, json!({ + "messages": advisor_history(), + "tools": [{"type": "advisor_20260301", "name": "advisor"}] + }))] + fn request_is_forwarded_unchanged( + #[case] capabilities: AnthropicModelCapabilities, + #[case] drop_params: bool, + #[case] fields: Value, + ) { + assert_eq!( + transform(fields.clone(), capabilities, drop_params), + Ok(body(fields)) + ); + } + + #[rstest] + #[case::temperature(sampling_removed(), json!({"temperature": 0.3}), json!({}))] + #[case::top_p(sampling_removed(), json!({"top_p": 0.9}), json!({}))] + #[case::top_k(sampling_removed(), json!({"top_k": 40}), json!({}))] + #[case::every_sampling_param_keeping_the_rest( + sampling_removed(), + json!({"temperature": 0.3, "top_p": 0.9, "top_k": 40, "stream": true}), + json!({"stream": true}) + )] + #[case::speed_on_a_sampling_model( + unmapped(), + json!({"speed": "fast", "temperature": 0.5}), + json!({"temperature": 0.5}) + )] + #[case::speed_on_a_sampling_removed_model( + sampling_removed(), + json!({"speed": "fast", "temperature": 1.0}), + json!({"temperature": 1.0}) + )] + fn removed_params_are_dropped_under_drop_params( + #[case] capabilities: AnthropicModelCapabilities, + #[case] fields: Value, + #[case] expected: Value, + ) { + assert_eq!(transform(fields, capabilities, true), Ok(body(expected))); + } + + #[rstest] + #[case::temperature( + sampling_removed(), + json!({"temperature": 0.3}), + "claude does not support temperature=0.3. Only temperature=1 is supported. To drop unsupported params, set `litellm.drop_params = True`." + )] + #[case::temperature_just_below_one( + sampling_removed(), + json!({"temperature": 0.99}), + "claude does not support temperature=0.99. Only temperature=1 is supported. To drop unsupported params, set `litellm.drop_params = True`." + )] + #[case::whole_number_temperature_keeps_its_decimal( + sampling_removed(), + json!({"temperature": 2.0}), + "claude does not support temperature=2.0. Only temperature=1 is supported. To drop unsupported params, set `litellm.drop_params = True`." + )] + #[case::top_p( + sampling_removed(), + json!({"top_p": 0.9}), + "claude does not support top_p=0.9. To drop unsupported params, set `litellm.drop_params = True`." + )] + #[case::top_k( + sampling_removed(), + json!({"top_k": 5}), + "claude does not support top_k=5. To drop unsupported params, set `litellm.drop_params = True`." + )] + #[case::top_k_next_to_unit_temperature( + sampling_removed(), + json!({"temperature": 1.0, "top_k": 5}), + "claude does not support top_k=5. To drop unsupported params, set `litellm.drop_params = True`." + )] + #[case::temperature_ahead_of_top_k( + sampling_removed(), + json!({"temperature": 0.5, "top_k": 5}), + "claude does not support temperature=0.5. Only temperature=1 is supported. To drop unsupported params, set `litellm.drop_params = True`." + )] + #[case::top_p_ahead_of_top_k( + sampling_removed(), + json!({"top_p": 0.9, "top_k": 5}), + "claude does not support top_p=0.9. To drop unsupported params, set `litellm.drop_params = True`." + )] + #[case::speed( + unmapped(), + json!({"speed": "fast"}), + "claude does not support speed='fast'. To drop unsupported params, set `litellm.drop_params = True`." + )] + #[case::speed_ahead_of_sampling_params( + sampling_removed(), + json!({"speed": "fast", "temperature": 0.5}), + "claude does not support speed='fast'. To drop unsupported params, set `litellm.drop_params = True`." + )] + fn removed_params_are_rejected_without_drop_params( + #[case] capabilities: AnthropicModelCapabilities, + #[case] fields: Value, + #[case] message: &str, + ) { + assert_eq!(transform(fields, capabilities, false), invalid(message)); + } + + #[rstest] + #[case::compaction_threshold( + json!([{"type": "compaction", "compact_threshold": 200000}]), + Some(json!({"edits": [{"type": "compact_20260112", "trigger": {"type": "input_tokens", "value": 200000}}]})) + )] + #[case::other_keys_pass_through( + json!([{"type": "compaction", "compact_threshold": 150000, "instructions": "Focus on preserving code snippets"}]), + Some(json!({"edits": [{ + "type": "compact_20260112", + "trigger": {"type": "input_tokens", "value": 150000}, + "instructions": "Focus on preserving code snippets" + }]})) + )] + #[case::float_threshold_is_truncated( + json!([{"type": "compaction", "compact_threshold": 150000.9}]), + Some(json!({"edits": [{"type": "compact_20260112", "trigger": {"type": "input_tokens", "value": 150000}}]})) + )] + #[case::compaction_without_threshold( + json!([{"type": "compaction"}]), + Some(json!({"edits": [{"type": "compact_20260112"}]})) + )] + #[case::non_numeric_threshold_is_dropped( + json!([{"type": "compaction", "compact_threshold": "150000"}]), + Some(json!({"edits": [{"type": "compact_20260112"}]})) + )] + #[case::non_object_entries_are_skipped( + json!([42, "compaction", null, [], {"type": "compaction", "compact_threshold": 1000}]), + Some(json!({"edits": [{"type": "compact_20260112", "trigger": {"type": "input_tokens", "value": 1000}}]})) + )] + #[case::only_compaction_entries_are_mapped_in_order( + json!([ + {"type": "compaction", "compact_threshold": 1000}, + {"type": "other", "compact_threshold": 5}, + {"type": "compaction", "instructions": "second"} + ]), + Some(json!({"edits": [ + {"type": "compact_20260112", "trigger": {"type": "input_tokens", "value": 1000}}, + {"type": "compact_20260112", "instructions": "second"} + ]})) + )] + #[case::list_without_compaction(json!([{"type": "other"}]), None)] + #[case::empty_list(json!([]), None)] + #[case::anthropic_edits_pass_through( + json!({"edits": [{"type": "compact_20260112", "trigger": {"type": "input_tokens", "value": 150000}}]}), + Some(json!({"edits": [{"type": "compact_20260112", "trigger": {"type": "input_tokens", "value": 150000}}]})) + )] + #[case::object_without_edits(json!({"type": "compaction"}), None)] + #[case::scalar(json!("compaction"), None)] + fn openai_context_management_maps_to_anthropic_edits( + #[case] context_management: Value, + #[case] expected: Option, + ) { + assert_eq!( + map_openai_context_management_to_anthropic(&context_management), + expected + ); + } + + #[rstest] + #[case::openai_list_is_mapped( + json!([{"type": "compaction", "compact_threshold": 200000}]), + json!({"edits": [{"type": "compact_20260112", "trigger": {"type": "input_tokens", "value": 200000}}]}) + )] + #[case::unmappable_list_is_kept(json!([{"type": "other"}]), json!([{"type": "other"}]))] + #[case::unmappable_object_is_kept(json!({"type": "other"}), json!({"type": "other"}))] + fn context_management_reaches_the_wire( + #[case] context_management: Value, + #[case] expected: Value, + unmapped: AnthropicModelCapabilities, + ) { + assert_eq!( + transform( + json!({"context_management": context_management}), + unmapped, + false + ), + Ok(body(json!({"context_management": expected}))) + ); + } + + #[rstest] + #[case::without_tools(json!({}))] + #[case::with_only_other_tools(json!({"tools": [{"name": "get_weather", "input_schema": {"type": "object"}}]}))] + fn advisor_history_is_stripped_without_the_advisor_tool( + #[case] tools: Value, + unmapped: AnthropicModelCapabilities, + ) { + let stripped = json!([ + {"role": "user", "content": "Build a worker pool."}, + {"role": "assistant", "content": [ + {"type": "text", "text": "Let me consult the advisor."}, + {"type": "text", "text": "Here is the implementation."} + ]} + ]); + assert_eq!( + transform( + merged(tools.clone(), json!({"messages": advisor_history()})), + unmapped, + false + ), + Ok(body(merged(tools, json!({"messages": stripped})))) + ); + } + + #[rstest] + fn bridge_minted_reasoning_is_stripped_from_the_wire(unmapped: AnthropicModelCapabilities) { + let messages = json!([ + {"role": "user", "content": "Solve it."}, + {"role": "assistant", "content": [ + {"type": "thinking", "thinking": "plan", "signature": format!("{ENCRYPTED_REASONING_SIGNATURE_PREFIX}gAAAA_1")}, + {"type": "redacted_thinking", "data": format!("{ENCRYPTED_REASONING_SIGNATURE_PREFIX}gAAAA_2")}, + {"type": "text", "text": "The answer."} + ]}, + {"role": "user", "content": "And the next one?"} + ]); + assert_eq!( + transform(json!({"messages": messages}), unmapped, false), + Ok(body(json!({"messages": [ + {"role": "user", "content": "Solve it."}, + {"role": "assistant", "content": [{"type": "text", "text": "The answer."}]}, + {"role": "user", "content": "And the next one?"} + ]}))) ); } #[test] - fn url_appends_messages_suffix_to_custom_base() { + fn thinking_is_translated_with_the_context_budgets() { + let context = MessagesTransformContext::with_lookup( + AnthropicModelCapabilities { + supports_reasoning: true, + ..Default::default() + }, + false, + &env(&[(LOW_BUDGET_ENV, "2000")]), + ); + let transformed = ANTHROPIC_MESSAGES_CONFIG + .transform_anthropic_messages_request( + request(json!({"max_tokens": 4096, "reasoning_effort": "low"})), + &context, + ) + .map(|transformed| serde_json::to_value(transformed).unwrap()); assert_eq!( - complete_anthropic_url(Some("https://proxy.internal"), &|_| None), - "https://proxy.internal/v1/messages" + transformed, + Ok(body(json!({ + "max_tokens": 4096, + "thinking": {"type": "enabled", "budget_tokens": 2000} + }))) ); } #[test] - fn url_leaves_complete_messages_endpoint_untouched() { + fn new_reads_thinking_budgets_from_the_process_environment() { + if std::env::var_os(PROCESS_ENV_PROBE).is_some() { + assert_eq!( + MessagesTransformContext::new(sampling_removed(), true), + MessagesTransformContext { + thinking: ThinkingContext { + capabilities: sampling_removed(), + budgets: ThinkingBudgets { + low: 2000, + ..ThinkingBudgets::default() + }, + }, + drop_params: true, + } + ); + return; + } + let (_, test_path) = concat!( + module_path!(), + "::new_reads_thinking_budgets_from_the_process_environment" + ) + .split_once("::") + .unwrap(); + let other_tiers = ["MINIMAL", "MEDIUM", "HIGH", "XHIGH", "MAX"] + .map(|tier| format!("DEFAULT_REASONING_EFFORT_{tier}_THINKING_BUDGET")); + let output = other_tiers + .iter() + .fold( + Command::new(std::env::current_exe().unwrap()), + |mut command, name| { + command.env_remove(name); + command + }, + ) + .args([test_path, "--exact"]) + .env(PROCESS_ENV_PROBE, "1") + .env(LOW_BUDGET_ENV, "2000") + .output() + .unwrap(); + let stdout = String::from_utf8_lossy(&output.stdout); + assert!( + output.status.success() && stdout.contains("1 passed"), + "{stdout}{}", + String::from_utf8_lossy(&output.stderr) + ); + } + + #[rstest] + #[case::public_endpoint_by_default(None, &[], "https://api.anthropic.com")] + #[case::explicit_api_base_beats_env( + Some("https://explicit.example.com"), + BOTH_BASE_ENVS, + "https://explicit.example.com" + )] + #[case::explicit_api_base_is_trimmed( + Some(" https://explicit.example.com "), + &[], + "https://explicit.example.com" + )] + #[case::blank_api_base_falls_back_to_env( + Some(" "), + BOTH_BASE_ENVS, + "https://api-base.example.com" + )] + #[case::api_base_env_beats_base_url_env(None, BOTH_BASE_ENVS, "https://api-base.example.com")] + #[case::base_url_env_without_api_base_env( + None, + &[(ANTHROPIC_BASE_URL_ENV, "https://base-url.example.com")], + "https://base-url.example.com" + )] + #[case::blank_api_base_env_falls_back_to_base_url_env( + None, + &[(ANTHROPIC_API_BASE_ENV, " \t "), (ANTHROPIC_BASE_URL_ENV, "https://base-url.example.com")], + "https://base-url.example.com" + )] + #[case::blank_envs_fall_back_to_public_endpoint( + None, + &[(ANTHROPIC_API_BASE_ENV, ""), (ANTHROPIC_BASE_URL_ENV, " ")], + "https://api.anthropic.com" + )] + fn api_base_resolution( + #[case] api_base: Option<&str>, + #[case] vars: Env, + #[case] expected: &str, + ) { + assert_eq!(resolve_anthropic_api_base(api_base, &env(vars)), expected); + } + + #[rstest] + #[case::public_endpoint(None, &[], "https://api.anthropic.com/v1/messages")] + #[case::base_url_env( + None, + &[(ANTHROPIC_BASE_URL_ENV, "https://custom.example.com")], + "https://custom.example.com/v1/messages" + )] + #[case::custom_base(Some("https://proxy.internal"), &[], "https://proxy.internal/v1/messages")] + #[case::trailing_slash(Some("https://proxy.internal/"), &[], "https://proxy.internal/v1/messages")] + #[case::complete_endpoint( + Some("https://proxy.internal/v1/messages"), + &[], + "https://proxy.internal/v1/messages" + )] + #[case::complete_endpoint_with_trailing_slash( + Some("https://proxy.internal/v1/messages/"), + &[], + "https://proxy.internal/v1/messages" + )] + fn complete_url_ends_in_the_messages_path( + #[case] api_base: Option<&str>, + #[case] vars: Env, + #[case] expected: &str, + ) { assert_eq!( - complete_anthropic_url(Some("https://proxy.internal/v1/messages"), &|_| None), - "https://proxy.internal/v1/messages" + ANTHROPIC_MESSAGES_CONFIG.get_complete_url(api_base, "claude", &env(vars)), + Ok(expected.to_string()) + ); + } + + #[rstest] + #[case::param_beats_env(Some("sk-param"), API_KEY_ENV, Ok("sk-param"))] + #[case::param_is_trimmed(Some(" sk-param "), &[], Ok("sk-param"))] + #[case::blank_param_falls_back_to_env(Some(" "), API_KEY_ENV, Ok("sk-env"))] + #[case::env_without_param(None, API_KEY_ENV, Ok("sk-env"))] + #[case::blank_env_is_missing(None, &[(ANTHROPIC_API_KEY_ENV, " ")], Err(MISSING_API_KEY))] + #[case::nothing_is_missing(None, &[], Err(MISSING_API_KEY))] + fn api_key_resolution( + #[case] api_key: Option<&str>, + #[case] vars: Env, + #[case] expected: Result<&str, &str>, + ) { + assert_eq!( + resolve_anthropic_api_key(api_key, &env(vars)).map_err(|error| error.to_string()), + expected.map(str::to_string).map_err(str::to_string) ); } #[test] - fn url_falls_back_to_env_base() { - let with_env = |key: &str| { - (key == ANTHROPIC_API_BASE_ENV).then(|| "https://env.anthropic".to_string()) - }; + fn config_reports_a_missing_key_as_an_auth_error() { assert_eq!( - complete_anthropic_url(Some(" "), &with_env), - "https://env.anthropic/v1/messages" + ANTHROPIC_MESSAGES_CONFIG.resolve_api_key(None, &no_env), + Err(Error::Auth(litellm_auth::Error::MissingApiKey { + provider: "Anthropic", + environment_variable: ANTHROPIC_API_KEY_ENV, + })) ); } #[test] - fn api_key_prefers_param_then_env_then_errors() { + fn config_authenticates_with_the_anthropic_auth_token() { assert_eq!( - resolve_anthropic_api_key(Some("sk-param"), &|_| None).unwrap(), - "sk-param" + ANTHROPIC_MESSAGES_CONFIG.authenticate( + vec![], + None, + &env(&[("ANTHROPIC_AUTH_TOKEN", "auth-token")]) + ), + Ok(headers(&[("authorization", "Bearer auth-token")])) ); - let with_env = |key: &str| (key == ANTHROPIC_API_KEY_ENV).then(|| "sk-env".to_string()); + } + + #[test] + fn config_requests_the_betas_the_request_features_need() { assert_eq!( - resolve_anthropic_api_key(Some(" "), &with_env).unwrap(), - "sk-env" - ); - assert_eq!( - resolve_anthropic_api_key(None, &|_| None) - .expect_err("missing key") - .to_string(), - "Missing Anthropic API Key - Set `api_key` or the ANTHROPIC_API_KEY environment variable" + ANTHROPIC_MESSAGES_CONFIG.request_headers( + headers(&[("x-api-key", "sk")]), + &request(json!({"speed": "fast"})) + ), + headers(&[ + ("x-api-key", "sk"), + ("anthropic-beta", beta::FAST_MODE_2026_02_01) + ]) ); } + #[rstest] + #[case::absent(None, None)] + #[case::blank(Some(" \t "), None)] + #[case::padded(Some(" value "), Some("value"))] + fn non_empty_trims_and_drops_blank_values( + #[case] value: Option<&str>, + #[case] expected: Option<&str>, + ) { + assert_eq!(non_empty(value), expected); + } + #[test] fn auth_strategy_and_default_headers_match_anthropic() { assert_eq!( @@ -142,4 +866,26 @@ mod tests { ] ); } + + #[test] + fn secret_names_cover_every_credential_and_base_lookup() { + let requested = std::cell::RefCell::new(Vec::::new()); + let record = |name: &str| -> Option { + requested.borrow_mut().push(name.to_string()); + None + }; + let _ = ANTHROPIC_MESSAGES_CONFIG.authenticate(Vec::new(), None, &record); + let _ = ANTHROPIC_MESSAGES_CONFIG.get_complete_url(None, "claude", &record); + let requested = requested.into_inner(); + assert!(!requested.is_empty()); + let undeclared: Vec<&String> = requested + .iter() + .filter(|name| { + !ANTHROPIC_MESSAGES_CONFIG + .secret_names() + .contains(&name.as_str()) + }) + .collect(); + assert_eq!(undeclared, Vec::<&String>::new()); + } } diff --git a/litellm-rust/crates/llms/src/anthropic/mod.rs b/litellm-rust/crates/llms/src/anthropic/mod.rs index d181ceaca3c..755bc7d1907 100644 --- a/litellm-rust/crates/llms/src/anthropic/mod.rs +++ b/litellm-rust/crates/llms/src/anthropic/mod.rs @@ -1,5 +1,6 @@ pub mod batches; pub mod chat; +pub mod common_utils; pub mod count_tokens; pub mod experimental_pass_through; diff --git a/litellm-rust/crates/llms/src/azure_ai/anthropic/messages_transformation.rs b/litellm-rust/crates/llms/src/azure_ai/anthropic/messages_transformation.rs index 99f55f18afc..c409f7f687e 100644 --- a/litellm-rust/crates/llms/src/azure_ai/anthropic/messages_transformation.rs +++ b/litellm-rust/crates/llms/src/azure_ai/anthropic/messages_transformation.rs @@ -4,14 +4,15 @@ use litellm_types::llms::anthropic_messages::{ }, anthropic_response::AnthropicMessagesResponse, }; -use serde_json::{Map, Value}; use crate::{ anthropic::experimental_pass_through::messages::transformation::{ ANTHROPIC_MESSAGES_CONFIG, AnthropicMessagesConfig, non_empty, }, base_llm::{ - anthropic_messages::transformation::{BaseAnthropicMessagesConfig, MessagesAuthStrategy}, + anthropic_messages::transformation::{ + BaseAnthropicMessagesConfig, Headers, MessagesAuthStrategy, MessagesTransformContext, + }, chat::transformation::Error, }, }; @@ -21,7 +22,6 @@ const AZURE_API_BASE_ENV: &str = "AZURE_API_BASE"; const ANTHROPIC_PATH_SEGMENT: &str = "/anthropic"; const MESSAGES_PATH_SUFFIX: &str = "/v1/messages"; const SYSTEM_ROLE: &str = "system"; -const TEXT_BLOCK_TYPE: &str = "text"; pub struct AzureAnthropicMessagesConfig { anthropic: AnthropicMessagesConfig, @@ -45,6 +45,7 @@ impl BaseAnthropicMessagesConfig for AzureAnthropicMessagesConfig { fn transform_anthropic_messages_request( &self, request: AnthropicMessagesRequest, + context: &MessagesTransformContext, ) -> Result { let mut request = fold_system_role_messages(request); if let Some(system) = request.system.as_mut() { @@ -54,7 +55,8 @@ impl BaseAnthropicMessagesConfig for AzureAnthropicMessagesConfig { .messages .iter_mut() .for_each(strip_scope_from_message); - self.anthropic.transform_anthropic_messages_request(request) + self.anthropic + .transform_anthropic_messages_request(request, context) } fn transform_anthropic_messages_response( @@ -74,6 +76,10 @@ impl BaseAnthropicMessagesConfig for AzureAnthropicMessagesConfig { resolve_azure_api_key(api_key, env_lookup) } + fn secret_names(&self) -> &'static [&'static str] { + &[AZURE_API_KEY_ENV, AZURE_API_BASE_ENV] + } + fn auth_strategy(&self) -> MessagesAuthStrategy { self.anthropic.auth_strategy() } @@ -85,6 +91,10 @@ impl BaseAnthropicMessagesConfig for AzureAnthropicMessagesConfig { fn default_headers(&self) -> &'static [(&'static str, &'static str)] { self.anthropic.default_headers() } + + fn request_headers(&self, headers: Headers, request: &AnthropicMessagesRequest) -> Headers { + self.anthropic.request_headers(headers, request) + } } pub fn resolve_azure_api_key( @@ -143,17 +153,7 @@ fn strip_scope_from_message(message: &mut AnthropicMessage) { } fn text_content_block(text: String) -> ContentBlock { - let extra = Map::from_iter([ - ( - "type".to_string(), - Value::String(TEXT_BLOCK_TYPE.to_string()), - ), - ("text".to_string(), Value::String(text)), - ]); - ContentBlock { - cache_control: None, - extra, - } + ContentBlock::text(text) } fn content_into_blocks(content: MessageContent) -> Vec { @@ -202,6 +202,7 @@ mod tests { use serde_json::json; use super::*; + use crate::anthropic::common_utils::AnthropicModelCapabilities; fn request_from(value: serde_json::Value) -> AnthropicMessagesRequest { serde_json::from_value(value).expect("valid request") @@ -346,7 +347,7 @@ mod tests { let transformed = to_value( AZURE_ANTHROPIC_MESSAGES_CONFIG - .transform_anthropic_messages_request(request) + .transform_anthropic_messages_request(request, &MessagesTransformContext::default()) .expect("request transforms"), ); @@ -373,10 +374,13 @@ mod tests { "messages": [{"role": "user", "content": "hi"}] })); let once = AZURE_ANTHROPIC_MESSAGES_CONFIG - .transform_anthropic_messages_request(request) + .transform_anthropic_messages_request(request, &MessagesTransformContext::default()) .expect("request transforms"); let twice = AZURE_ANTHROPIC_MESSAGES_CONFIG - .transform_anthropic_messages_request(once.clone()) + .transform_anthropic_messages_request( + once.clone(), + &MessagesTransformContext::default(), + ) .expect("request transforms"); assert_eq!(once, twice); assert_eq!(to_value(once)["system"], json!("plain string system")); @@ -408,9 +412,21 @@ mod tests { "inference_geo": "us", "litellm_metadata": {"trace": "abc"} }); + let context = MessagesTransformContext::with_lookup( + AnthropicModelCapabilities { + supports_reasoning: true, + supports_adaptive_thinking: true, + supports_legacy_thinking: true, + supports_output_config: true, + supports_speed: true, + ..Default::default() + }, + false, + &|_: &str| None, + ); let transformed = to_value( AZURE_ANTHROPIC_MESSAGES_CONFIG - .transform_anthropic_messages_request(request_from(body.clone())) + .transform_anthropic_messages_request(request_from(body.clone()), &context) .expect("request transforms"), ); assert_eq!(transformed, body); @@ -430,7 +446,7 @@ mod tests { let transformed = to_value( AZURE_ANTHROPIC_MESSAGES_CONFIG - .transform_anthropic_messages_request(request) + .transform_anthropic_messages_request(request, &MessagesTransformContext::default()) .expect("request transforms"), ); @@ -460,7 +476,7 @@ mod tests { let transformed = to_value( AZURE_ANTHROPIC_MESSAGES_CONFIG - .transform_anthropic_messages_request(request) + .transform_anthropic_messages_request(request, &MessagesTransformContext::default()) .expect("request transforms"), ); @@ -485,9 +501,21 @@ mod tests { {"role": "assistant", "content": "hello"} ] }); + let context = MessagesTransformContext::with_lookup( + AnthropicModelCapabilities { + supports_reasoning: true, + supports_adaptive_thinking: true, + supports_legacy_thinking: true, + supports_output_config: true, + supports_speed: true, + ..Default::default() + }, + false, + &|_: &str| None, + ); let transformed = to_value( AZURE_ANTHROPIC_MESSAGES_CONFIG - .transform_anthropic_messages_request(request_from(body.clone())) + .transform_anthropic_messages_request(request_from(body.clone()), &context) .expect("request transforms"), ); assert_eq!(transformed, body); @@ -500,6 +528,57 @@ mod tests { assert!(err.is_data()); } + #[rstest::rstest] + #[case::compact_context_management_edit( + json!({"context_management": {"edits": [{"type": "compact_20260112"}]}}), + &[], + &[("x-api-key", "k"), ("anthropic-beta", "compact-2026-01-12")] + )] + #[case::forwarded_beta_merged_with_structured_output( + json!({"output_config": {"format": {"type": "json_schema"}}}), + &[("anthropic-beta", "web-search-2025-03-05")], + &[("x-api-key", "k"), ("anthropic-beta", "structured-outputs-2025-11-13,web-search-2025-03-05")] + )] + #[case::no_feature_needs_a_beta(json!({}), &[], &[("x-api-key", "k")])] + fn request_headers_carry_the_anthropic_feature_betas( + #[case] fields: serde_json::Value, + #[case] forwarded: &[(&str, &str)], + #[case] expected: &[(&str, &str)], + ) { + let pairs = |pairs: &[(&str, &str)]| -> Vec<(String, String)> { + pairs + .iter() + .map(|(name, value)| (name.to_string(), value.to_string())) + .collect() + }; + let serde_json::Value::Object(fields) = fields else { + panic!("case fields are an object") + }; + let request = request_from(serde_json::Value::Object( + [ + ("model".to_string(), json!("claude-sonnet")), + ("max_tokens".to_string(), json!(16)), + ( + "messages".to_string(), + json!([{"role": "user", "content": "hi"}]), + ), + ] + .into_iter() + .chain(fields) + .collect(), + )); + assert_eq!( + AZURE_ANTHROPIC_MESSAGES_CONFIG.request_headers( + pairs(&[("x-api-key", "k")]) + .into_iter() + .chain(pairs(forwarded)) + .collect(), + &request + ), + pairs(expected) + ); + } + #[test] fn transform_response_passes_through() { let response: AnthropicMessagesResponse = serde_json::from_value(json!({ @@ -521,4 +600,26 @@ mod tests { assert_eq!(value["stop_sequence"], json!(null)); assert_eq!(value["content"][0]["text"], json!("hello")); } + + #[test] + fn secret_names_cover_every_credential_and_base_lookup() { + let requested = std::cell::RefCell::new(Vec::::new()); + let record = |name: &str| -> Option { + requested.borrow_mut().push(name.to_string()); + None + }; + let _ = AZURE_ANTHROPIC_MESSAGES_CONFIG.authenticate(Vec::new(), None, &record); + let _ = AZURE_ANTHROPIC_MESSAGES_CONFIG.get_complete_url(None, "claude", &record); + let requested = requested.into_inner(); + assert!(!requested.is_empty()); + let undeclared: Vec<&String> = requested + .iter() + .filter(|name| { + !AZURE_ANTHROPIC_MESSAGES_CONFIG + .secret_names() + .contains(&name.as_str()) + }) + .collect(); + assert_eq!(undeclared, Vec::<&String>::new()); + } } diff --git a/litellm-rust/crates/llms/src/base_llm/anthropic_messages/transformation.rs b/litellm-rust/crates/llms/src/base_llm/anthropic_messages/transformation.rs index 5b4afb601d2..8db14687214 100644 --- a/litellm-rust/crates/llms/src/base_llm/anthropic_messages/transformation.rs +++ b/litellm-rust/crates/llms/src/base_llm/anthropic_messages/transformation.rs @@ -1,8 +1,14 @@ +use litellm_http::request::{has_bearer_auth, has_header}; use litellm_types::llms::anthropic_messages::{ anthropic_request::AnthropicMessagesRequest, anthropic_response::AnthropicMessagesResponse, }; -use crate::base_llm::chat::transformation::Error; +use crate::{ + anthropic::experimental_pass_through::messages::thinking::ThinkingContext, + base_llm::chat::transformation::Error, +}; + +pub type Headers = Vec<(String, String)>; #[derive(Clone, Copy, Debug, PartialEq, Eq)] pub enum MessagesAuthStrategy { @@ -19,6 +25,12 @@ impl MessagesAuthStrategy { } } +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +pub struct MessagesTransformContext { + pub thinking: ThinkingContext, + pub drop_params: bool, +} + pub trait BaseAnthropicMessagesConfig: Sync { fn get_complete_url( &self, @@ -30,6 +42,7 @@ pub trait BaseAnthropicMessagesConfig: Sync { fn transform_anthropic_messages_request( &self, request: AnthropicMessagesRequest, + _context: &MessagesTransformContext, ) -> Result { Ok(request) } @@ -48,6 +61,8 @@ pub trait BaseAnthropicMessagesConfig: Sync { env_lookup: &dyn Fn(&str) -> Option, ) -> Result; + fn secret_names(&self) -> &'static [&'static str]; + fn auth_strategy(&self) -> MessagesAuthStrategy { MessagesAuthStrategy::Header("x-api-key") } @@ -56,10 +71,225 @@ pub trait BaseAnthropicMessagesConfig: Sync { false } + fn authenticate( + &self, + headers: Headers, + api_key: Option<&str>, + env_lookup: &dyn Fn(&str) -> Option, + ) -> Result { + let strategy = self.auth_strategy(); + if has_header(&headers, strategy.header_name()) + || (self.accepts_bearer_auth() && has_bearer_auth(&headers)) + { + return Ok(headers); + } + let api_key = self.resolve_api_key(api_key, env_lookup)?; + let auth_header = match strategy { + MessagesAuthStrategy::Bearer => { + ("authorization".to_string(), format!("Bearer {api_key}")) + } + MessagesAuthStrategy::Header(name) => (name.to_string(), api_key), + }; + Ok(headers.into_iter().chain([auth_header]).collect()) + } + fn default_headers(&self) -> &'static [(&'static str, &'static str)] { &[ ("anthropic-version", "2023-06-01"), ("content-type", "application/json"), ] } + + fn request_headers(&self, headers: Headers, _request: &AnthropicMessagesRequest) -> Headers { + headers + } +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + + use super::*; + + const X_API_KEY: MessagesAuthStrategy = MessagesAuthStrategy::Header("x-api-key"); + + struct StubConfig { + strategy: MessagesAuthStrategy, + accepts_bearer: bool, + } + + impl BaseAnthropicMessagesConfig for StubConfig { + fn secret_names(&self) -> &'static [&'static str] { + &[] + } + + fn get_complete_url( + &self, + _api_base: Option<&str>, + _model: &str, + _env_lookup: &dyn Fn(&str) -> Option, + ) -> Result { + Ok(String::new()) + } + + fn resolve_api_key( + &self, + api_key: Option<&str>, + _env_lookup: &dyn Fn(&str) -> Option, + ) -> Result { + api_key + .map(str::to_string) + .ok_or(Error::MissingField("api_key")) + } + + fn auth_strategy(&self) -> MessagesAuthStrategy { + self.strategy + } + + fn accepts_bearer_auth(&self) -> bool { + self.accepts_bearer + } + } + + struct DefaultsConfig; + + impl BaseAnthropicMessagesConfig for DefaultsConfig { + fn secret_names(&self) -> &'static [&'static str] { + &[] + } + + fn get_complete_url( + &self, + _api_base: Option<&str>, + _model: &str, + _env_lookup: &dyn Fn(&str) -> Option, + ) -> Result { + Ok(String::new()) + } + + fn resolve_api_key( + &self, + api_key: Option<&str>, + _env_lookup: &dyn Fn(&str) -> Option, + ) -> Result { + api_key + .map(str::to_string) + .ok_or(Error::MissingField("api_key")) + } + } + + #[test] + fn default_config_adds_its_key_next_to_a_forwarded_bearer() { + assert_eq!( + DefaultsConfig.authenticate( + headers(&[("authorization", "Bearer forwarded")]), + Some("sk"), + &|_| None + ), + Ok(headers(&[ + ("authorization", "Bearer forwarded"), + ("x-api-key", "sk") + ])) + ); + } + + #[test] + fn default_request_headers_are_the_given_headers() { + let request: AnthropicMessagesRequest = serde_json::from_value(serde_json::json!({ + "model": "claude", + "max_tokens": 16, + "speed": "fast", + "messages": [{"role": "user", "content": "hi"}] + })) + .unwrap(); + assert_eq!( + DefaultsConfig.request_headers(headers(&[("x-api-key", "sk")]), &request), + headers(&[("x-api-key", "sk")]) + ); + } + + fn headers(pairs: &[(&str, &str)]) -> Headers { + pairs + .iter() + .map(|(name, value)| (name.to_string(), value.to_string())) + .collect() + } + + #[rstest] + #[case::own_header_is_kept( + X_API_KEY, + false, + headers(&[("x-api-key", "forwarded")]), + None, + Ok(headers(&[("x-api-key", "forwarded")])) + )] + #[case::own_header_in_any_casing_is_kept( + X_API_KEY, + false, + headers(&[("X-Api-Key", "forwarded")]), + None, + Ok(headers(&[("X-Api-Key", "forwarded")])) + )] + #[case::accepted_bearer_is_kept( + X_API_KEY, + true, + headers(&[("authorization", "Bearer forwarded")]), + None, + Ok(headers(&[("authorization", "Bearer forwarded")])) + )] + #[case::bearer_the_provider_does_not_accept_gets_the_key_too( + X_API_KEY, + false, + headers(&[("authorization", "Bearer forwarded")]), + Some("sk"), + Ok(headers(&[("authorization", "Bearer forwarded"), ("x-api-key", "sk")])) + )] + #[case::blank_bearer_gets_the_key( + X_API_KEY, + true, + headers(&[("authorization", "Bearer ")]), + Some("sk"), + Ok(headers(&[("authorization", "Bearer "), ("x-api-key", "sk")])) + )] + #[case::key_goes_in_the_provider_header( + X_API_KEY, + false, + headers(&[("content-type", "application/json")]), + Some("sk"), + Ok(headers(&[("content-type", "application/json"), ("x-api-key", "sk")])) + )] + #[case::key_goes_in_a_bearer( + MessagesAuthStrategy::Bearer, + false, + headers(&[]), + Some("sk"), + Ok(headers(&[("authorization", "Bearer sk")])) + )] + #[case::bearer_strategy_keeps_a_forwarded_authorization( + MessagesAuthStrategy::Bearer, + false, + headers(&[("authorization", "Bearer forwarded")]), + None, + Ok(headers(&[("authorization", "Bearer forwarded")])) + )] + #[case::missing_key_is_an_error( + X_API_KEY, + false, + headers(&[]), + None, + Err(Error::MissingField("api_key")) + )] + fn default_authenticate_applies_the_key_unless_a_credential_is_forwarded( + #[case] strategy: MessagesAuthStrategy, + #[case] accepts_bearer: bool, + #[case] forwarded: Headers, + #[case] api_key: Option<&str>, + #[case] expected: Result, + ) { + let config = StubConfig { + strategy, + accepts_bearer, + }; + assert_eq!(config.authenticate(forwarded, api_key, &|_| None), expected); + } } diff --git a/litellm-rust/crates/python-bridge/src/routes/messages/host.rs b/litellm-rust/crates/python-bridge/src/routes/messages/host.rs index 1a9b170f661..9d97094aeda 100644 --- a/litellm-rust/crates/python-bridge/src/routes/messages/host.rs +++ b/litellm-rust/crates/python-bridge/src/routes/messages/host.rs @@ -2,9 +2,11 @@ use bytes::Bytes; use litellm_core::messages::{ Error, route::{Messages, MessagesCall, MessagesOp, MessagesOpResult, MessagesOutput}, + types::MessagesShaping, }; use litellm_host_python::{InvokeError, RouteHost, from_py, lookup, to_py}; use litellm_http::transport::Error as TransportError; +use litellm_types::utils::ProviderSpecificHeaders; use pyo3::{ exceptions::{PyException, PyValueError}, gc::{PyTraverseError, PyVisit}, @@ -18,9 +20,10 @@ use crate::{ marshal::{optional_timeout, python_timeout_seconds}, }; -/// The Anthropic Messages body fields a caller may pass besides `model` and `messages`, -/// as `AnthropicMessagesRequestOptionalParams` declares them. -const BODY_FIELDS: [&str; 20] = [ +const ROUTE_HOST_MODULE: &str = "litellm.rust_bridge.messages.route_host"; +const REQUEST_ERROR_MARKER: &str = "messages_request_error"; + +const BODY_FIELDS: [&str; 22] = [ "max_tokens", "metadata", "stop_sequences", @@ -35,14 +38,46 @@ const BODY_FIELDS: [&str; 20] = [ "top_p", "mcp_servers", "context_management", + "compaction", "container", "output_format", "speed", "output_config", "cache_control", "reasoning_effort", + "safeguards", ]; +fn merge_headers( + forwarded: Option>, + extra_headers: Option>, +) -> Option> { + let merged: Map = forwarded + .into_iter() + .flatten() + .chain(extra_headers.into_iter().flatten()) + .collect(); + (!merged.is_empty()).then_some(merged) +} + +fn native_error(py: Python<'_>, error: Error) -> PyResult { + match error { + Error::Transport(TransportError::Http { status, body }) => { + let error = RustUpstreamError::new_err((status, body)); + error + .value(py) + .setattr("headers", Vec::<(String, String)>::new())?; + Ok(error) + } + Error::InvalidRequest(message) => { + let error = PyValueError::new_err(message); + error.value(py).setattr(REQUEST_ERROR_MARKER, true)?; + Ok(error) + } + other => Ok(messages_error_to_pyerr(other)), + } +} + /// The Python side of the Messages route: projects the prepared arguments and builds the /// public response, chunks and exceptions. pub(super) struct MessagesRouteHost { @@ -84,19 +119,65 @@ impl MessagesRouteHost { .map(|value| python_timeout_seconds(py, value.unbind())) .transpose()? .flatten(); + let custom_llm_provider = string("custom_llm_provider")?; + let shaping = self.shaping(py, &model, custom_llm_provider.as_deref(), arguments)?; Ok(MessagesCall { model, body, api_key: string("api_key")?, api_base: string("api_base")?, - custom_llm_provider: string("custom_llm_provider")?, - extra_headers: argument("extra_headers")? - .map(|value| from_py(&value)) - .transpose()?, + extra_headers: self.merged_headers(py, arguments)?, + provider_specific_header: self.provider_specific_header(py, arguments)?, + custom_llm_provider, timeout: optional_timeout(timeout), + shaping, }) } + fn merged_headers( + &self, + py: Python<'_>, + arguments: &Bound<'_, PyDict>, + ) -> PyResult>> { + let request = self.request.bind(py); + let mapping = |name: &str| -> PyResult>> { + lookup(arguments, request, name)? + .filter(|value| !value.is_none()) + .map(|value| from_py(&value)) + .transpose() + }; + Ok(merge_headers( + mapping("headers")?, + mapping("extra_headers")?, + )) + } + + fn provider_specific_header( + &self, + py: Python<'_>, + arguments: &Bound<'_, PyDict>, + ) -> PyResult> { + lookup(arguments, self.request.bind(py), "provider_specific_header")? + .filter(|value| !value.is_none()) + .map(|value| from_py(&value)) + .transpose() + } + + fn shaping( + &self, + py: Python<'_>, + model: &str, + custom_llm_provider: Option<&str>, + arguments: &Bound<'_, PyDict>, + ) -> PyResult { + let projected = py.import(ROUTE_HOST_MODULE)?.getattr("shaping")?.call1(( + model, + custom_llm_provider, + arguments, + ))?; + from_py(&projected) + } + fn provider(&self, py: Python<'_>) -> String { self.request .bind(py) @@ -112,7 +193,7 @@ impl MessagesRouteHost { return error; } let mapped = py - .import("litellm.rust_bridge.messages.route_host") + .import(ROUTE_HOST_MODULE) .and_then(|module| module.getattr("map_failure")) .and_then(|map| map.call1((error.value(py), self.request.bind(py), self.provider(py)))) .and_then(|mapped| { @@ -148,7 +229,7 @@ impl RouteHost for MessagesRouteHost { fn complete(&mut self, py: Python<'_>, response: MessagesOutput) -> PyResult> { match response { MessagesOutput::Message(message) => py - .import("litellm.rust_bridge.messages.route_host")? + .import(ROUTE_HOST_MODULE)? .getattr("response")? .call1((to_py(py, message.as_ref())?,)) .map(Bound::unbind), @@ -161,17 +242,12 @@ impl RouteHost for MessagesRouteHost { } fn classify(&self, py: Python<'_>, error: Error) -> PyResult { - let native = match error { - Error::Transport(TransportError::Http { status, body }) => { - let error = RustUpstreamError::new_err((status, body)); - error - .value(py) - .setattr("headers", Vec::<(String, String)>::new())?; - error - } - other => messages_error_to_pyerr(other), - }; - Ok(self.map_failure(py, native)) + if let Error::Secret(source) = &error + && let Some(original) = crate::secrets::python_error(py, source.source_error()) + { + return Ok(original); + } + Ok(self.map_failure(py, native_error(py, error)?)) } fn host_error(error: &PyErr) -> Error { @@ -184,3 +260,62 @@ impl RouteHost for MessagesRouteHost { visit.call(&self.request) } } + +#[cfg(test)] +mod tests { + use rstest::rstest; + use serde_json::json; + + use super::*; + + fn map(value: Value) -> Map { + serde_json::from_value(value).unwrap() + } + + #[rstest] + #[case::extra_over_forwarded( + Some(json!({"X-Priority": "forwarded", "X-Forwarded-Only": "keep"})), + Some(json!({"X-Priority": "extra", "X-Extra-Only": "also-keep"})), + Some(json!({"X-Priority": "extra", "X-Forwarded-Only": "keep", "X-Extra-Only": "also-keep"})), + )] + #[case::only_forwarded(Some(json!({"X-Forwarded": "yes"})), None, Some(json!({"X-Forwarded": "yes"})))] + #[case::only_extra_headers( + None, + Some(json!({"X-Custom-Header": "from-kwargs", "X-Auth-Token": "token123"})), + Some(json!({"X-Custom-Header": "from-kwargs", "X-Auth-Token": "token123"})), + )] + #[case::nothing(None, Some(json!({})), None)] + fn headers_merge_forwarded_then_extra( + #[case] forwarded: Option, + #[case] extra_headers: Option, + #[case] expected: Option, + ) { + assert_eq!( + merge_headers(forwarded.map(map), extra_headers.map(map)), + expected.map(map) + ); + } + + #[rstest] + #[case::rejected_request(Error::InvalidRequest("does not support top_k=5".into()), true)] + #[case::unresolvable_provider(Error::InvalidProvider("openai".into()), false)] + #[case::upstream_failure( + Error::Transport(TransportError::Http { status: 400, body: "bad".into() }), + false, + )] + fn only_request_rejections_carry_the_request_error_marker( + #[case] error: Error, + #[case] marked: bool, + ) { + Python::initialize(); + Python::attach(|py| { + let native = native_error(py, error).unwrap(); + let marker = native + .value(py) + .getattr_opt(REQUEST_ERROR_MARKER) + .unwrap() + .map(|value| value.extract::().unwrap()); + assert_eq!(marker.unwrap_or(false), marked); + }); + } +} diff --git a/litellm-rust/crates/python-bridge/src/routes/messages/mod.rs b/litellm-rust/crates/python-bridge/src/routes/messages/mod.rs index 804d883e9ae..fd474e6b2d4 100644 --- a/litellm-rust/crates/python-bridge/src/routes/messages/mod.rs +++ b/litellm-rust/crates/python-bridge/src/routes/messages/mod.rs @@ -39,11 +39,12 @@ fn run_messages( "the Rust Messages route does not serve this provider", )); } + let secrets = crate::secrets::source(py)?; run_legacy_call( py, SURFACE, PublicCall::capture(&request, &args, &kwargs)?, - crate::logger::LoggedMachine::new(messages_machine()), + crate::logger::LoggedMachine::new(messages_machine(secrets)), MessagesRouteHost::new(request.unbind()), asynchronous, ) diff --git a/litellm-rust/crates/types/Cargo.toml b/litellm-rust/crates/types/Cargo.toml index 6a2efa90ab4..0a0927386f0 100644 --- a/litellm-rust/crates/types/Cargo.toml +++ b/litellm-rust/crates/types/Cargo.toml @@ -8,3 +8,6 @@ repository.workspace = true [dependencies] serde.workspace = true serde_json.workspace = true + +[dev-dependencies] +rstest.workspace = true diff --git a/litellm-rust/crates/types/src/llms/anthropic_messages/anthropic_request.rs b/litellm-rust/crates/types/src/llms/anthropic_messages/anthropic_request.rs index 50eedf7ba09..2f7a75ba517 100644 --- a/litellm-rust/crates/types/src/llms/anthropic_messages/anthropic_request.rs +++ b/litellm-rust/crates/types/src/llms/anthropic_messages/anthropic_request.rs @@ -17,12 +17,48 @@ pub enum MessageContent { #[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)] pub struct ContentBlock { + #[serde(rename = "type", default, skip_serializing_if = "Option::is_none")] + pub block_type: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub text: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub thinking: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub signature: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub data: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub id: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub name: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub tool_use_id: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub input: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub content: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub provider_specific_fields: Option, #[serde(skip_serializing_if = "Option::is_none")] pub cache_control: Option, #[serde(flatten)] pub extra: Map, } +impl ContentBlock { + pub fn text(text: impl Into) -> Self { + Self { + block_type: Some("text".to_string()), + text: Some(text.into()), + ..Self::default() + } + } + + pub fn is_type(&self, block_type: &str) -> bool { + self.block_type.as_deref() == Some(block_type) + } +} + #[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)] pub struct CacheControl { #[serde(rename = "type", skip_serializing_if = "Option::is_none")] @@ -85,6 +121,126 @@ pub struct AnthropicMessagesRequest { pub speed: Option, #[serde(skip_serializing_if = "Option::is_none")] pub inference_geo: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub reasoning_effort: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub compaction: Option, #[serde(flatten)] pub extra: Map, } + +impl AnthropicMessage { + pub fn blocks(&self) -> &[ContentBlock] { + match &self.content { + MessageContent::Blocks(blocks) => blocks, + MessageContent::Text(_) => &[], + } + } + + pub fn with_blocks(self, blocks: Vec) -> Self { + Self { + content: MessageContent::Blocks(blocks), + ..self + } + } +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + use serde_json::json; + + use super::*; + + fn round_trip(value: &Value) -> Value { + let parsed: T = serde_json::from_value(value.clone()).unwrap(); + serde_json::to_value(parsed).unwrap() + } + + #[rstest] + #[case::text(json!({"type": "text", "text": "hi"}))] + #[case::text_with_citations_and_cache_control(json!({ + "type": "text", + "text": "hi", + "citations": [{"type": "char_location", "cited_text": "x"}], + "cache_control": {"type": "ephemeral", "ttl": "1h", "scope": "global", "future": 1} + }))] + #[case::image(json!({"type": "image", "source": {"type": "base64", "media_type": "image/png", "data": "AA=="}}))] + #[case::thinking(json!({"type": "thinking", "thinking": "hmm", "signature": "sig"}))] + #[case::redacted_thinking(json!({"type": "redacted_thinking", "data": "opaque"}))] + #[case::tool_use(json!({"type": "tool_use", "id": "toolu_1", "name": "f", "input": {"q": [1, null]}}))] + #[case::tool_result_with_text(json!({"type": "tool_result", "tool_use_id": "toolu_1", "content": "ok", "is_error": false}))] + #[case::tool_result_with_blocks(json!({"type": "tool_result", "tool_use_id": "toolu_1", "content": [{"type": "text", "text": "ok"}]}))] + #[case::web_search_result_with_nulls(json!({ + "type": "web_search_tool_result", + "tool_use_id": "srvtoolu_1", + "content": [{"type": "web_search_result", "url": "u", "page_age": null, "encrypted_content": ""}] + }))] + #[case::provider_specific_fields(json!({"type": "tool_use", "id": "t", "name": "f", "input": {}, "provider_specific_fields": {"x": 1}}))] + #[case::untyped(json!({"unknown": {"nested": true}}))] + fn content_block_round_trips_unchanged(#[case] block: Value) { + assert_eq!(round_trip::(&block), block); + } + + #[test] + fn text_constructor_serializes_as_a_text_block() { + assert_eq!( + serde_json::to_value(ContentBlock::text("hello")).unwrap(), + json!({"type": "text", "text": "hello"}) + ); + } + + #[rstest] + #[case::same_type(json!({"type": "tool_use"}), "tool_use", true)] + #[case::other_type(json!({"type": "tool_result"}), "tool_use", false)] + #[case::prefix_of_type(json!({"type": "tool_use"}), "tool", false)] + #[case::no_type(json!({"text": "x"}), "text", false)] + fn is_type_matches_the_exact_block_type( + #[case] block: Value, + #[case] block_type: &str, + #[case] expected: bool, + ) { + let block: ContentBlock = serde_json::from_value(block).unwrap(); + assert_eq!(block.is_type(block_type), expected); + } + + #[rstest] + #[case::string_content(json!({"role": "user", "content": "hi"}), vec![])] + #[case::block_content( + json!({"role": "user", "content": [{"type": "text", "text": "a"}, {"type": "text", "text": "b"}]}), + vec![ContentBlock::text("a"), ContentBlock::text("b")], + )] + fn message_blocks_list_only_block_content( + #[case] message: Value, + #[case] expected: Vec, + ) { + let message: AnthropicMessage = serde_json::from_value(message).unwrap(); + assert_eq!(message.blocks(), expected.as_slice()); + } + + #[rstest] + #[case::replaces_string_content(json!({"role": "assistant", "content": "old", "name": "kept"}))] + #[case::replaces_block_content(json!({"role": "assistant", "content": [{"type": "text", "text": "old"}], "name": "kept"}))] + fn with_blocks_replaces_content_and_keeps_the_rest(#[case] message: Value) { + let message: AnthropicMessage = serde_json::from_value(message).unwrap(); + assert_eq!( + serde_json::to_value(message.with_blocks(vec![ContentBlock::text("new")])).unwrap(), + json!({"role": "assistant", "content": [{"type": "text", "text": "new"}], "name": "kept"}) + ); + } + + #[rstest] + #[case::minimal(json!({"model": "m", "messages": [{"role": "user", "content": "hi"}]}))] + #[case::reasoning_effort_compaction_and_unknown_fields(json!({ + "model": "m", + "messages": [{"role": "user", "content": [{"type": "text", "text": "hi"}]}], + "max_tokens": 8, + "reasoning_effort": "high", + "compaction": {"type": "auto"}, + "safeguards": [{"type": "dangerous_tool_use", "classifier_context": {"v": 1}}], + "metadata": {"user_id": "u"} + }))] + fn request_round_trips_unchanged(#[case] request: Value) { + assert_eq!(round_trip::(&request), request); + } +} diff --git a/litellm-rust/crates/types/src/llms/anthropic_messages/anthropic_response.rs b/litellm-rust/crates/types/src/llms/anthropic_messages/anthropic_response.rs index 0c3876aac59..0a2653f352f 100644 --- a/litellm-rust/crates/types/src/llms/anthropic_messages/anthropic_response.rs +++ b/litellm-rust/crates/types/src/llms/anthropic_messages/anthropic_response.rs @@ -9,8 +9,6 @@ pub struct AnthropicMessagesResponse { pub role: String, pub model: String, pub content: Vec, - // Anthropic always includes stop_reason / stop_sequence, null until the turn - // ends; serialize them even when None so callers see the same shape as Python. pub stop_reason: Option, pub stop_sequence: Option, #[serde(skip_serializing_if = "Option::is_none")] @@ -20,3 +18,61 @@ pub struct AnthropicMessagesResponse { #[serde(flatten)] pub extra: Map, } + +#[cfg(test)] +mod tests { + use rstest::rstest; + use serde_json::json; + + use super::*; + + fn response( + stop_reason: Option<&str>, + stop_sequence: Option<&str>, + usage: Option, + container: Option, + ) -> AnthropicMessagesResponse { + AnthropicMessagesResponse { + id: "msg_1".to_string(), + message_type: "message".to_string(), + role: "assistant".to_string(), + model: "claude".to_string(), + content: vec![], + stop_reason: stop_reason.map(str::to_string), + stop_sequence: stop_sequence.map(str::to_string), + usage, + container, + extra: Map::new(), + } + } + + #[rstest] + #[case::turn_in_progress(None, None, json!(null), json!(null))] + #[case::ended_on_end_turn(Some("end_turn"), None, json!("end_turn"), json!(null))] + #[case::ended_on_stop_sequence(Some("stop_sequence"), Some("###"), json!("stop_sequence"), json!("###"))] + fn stop_fields_are_always_serialized( + #[case] stop_reason: Option<&str>, + #[case] stop_sequence: Option<&str>, + #[case] expected_reason: Value, + #[case] expected_sequence: Value, + ) { + let body: Value = serde_json::to_value(response(stop_reason, stop_sequence, None, None)) + .expect("serializable"); + assert_eq!(body.get("stop_reason"), Some(&expected_reason)); + assert_eq!(body.get("stop_sequence"), Some(&expected_sequence)); + } + + #[rstest] + #[case::absent(None, None)] + #[case::present(Some(json!({"input_tokens": 1})), Some(json!({"id": "c_1"})))] + fn usage_and_container_are_omitted_only_when_none( + #[case] usage: Option, + #[case] container: Option, + ) { + let body: Value = + serde_json::to_value(response(None, None, usage.clone(), container.clone())) + .expect("serializable"); + assert_eq!(body.get("usage").cloned(), usage); + assert_eq!(body.get("container").cloned(), container); + } +} diff --git a/litellm-rust/crates/types/src/utils.rs b/litellm-rust/crates/types/src/utils.rs index 7f0c18f9f2c..5ca56ec9e49 100644 --- a/litellm-rust/crates/types/src/utils.rs +++ b/litellm-rust/crates/types/src/utils.rs @@ -3,6 +3,21 @@ use serde_json::{Map, Value}; use crate::llms::openai::{ChatCompletionThinkingBlock, ChatCompletionToolCallChunk}; +#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)] +pub struct ProviderSpecificHeader { + #[serde(default)] + pub custom_llm_provider: String, + #[serde(default)] + pub extra_headers: Map, +} + +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +#[serde(untagged)] +pub enum ProviderSpecificHeaders { + One(ProviderSpecificHeader), + Many(Vec), +} + /// OpenAI `usage`, including the `prompt_tokens_details` split LiteLLM's Python /// path reports so cost tracking sees the same numbers on either path. #[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)] diff --git a/litellm/rust_bridge/messages/route_host.py b/litellm/rust_bridge/messages/route_host.py index beef0f81eca..d49d7b75a6f 100644 --- a/litellm/rust_bridge/messages/route_host.py +++ b/litellm/rust_bridge/messages/route_host.py @@ -1,12 +1,50 @@ from __future__ import annotations -from collections.abc import Mapping -from typing import cast # noqa: TID251 # narrows the normalized native payload to the public TypedDict +from collections.abc import Mapping, Sequence +from dataclasses import asdict, dataclass +from typing import Final, cast # noqa: TID251 # narrows the normalized native payload to the public TypedDict +from pydantic import TypeAdapter, ValidationError + +import litellm +from litellm.litellm_core_utils.core_helpers import normalize_drop_params +from litellm.llms.anthropic.experimental_pass_through.utils import is_reasoning_auto_summary_enabled from litellm.rust_bridge import failures from litellm.rust_bridge.messages.entrypoints import LiteLLMMessagesRequest from litellm.types.llms.anthropic_messages.anthropic_response import AnthropicMessagesResponse +_DROP_PATHS: Final = TypeAdapter(list[object]) + + +@dataclass(frozen=True, slots=True) +class EffortTiers: + minimal: bool + low: bool + medium: bool + high: bool + xhigh: bool + max: bool + + +@dataclass(frozen=True, slots=True) +class ModelCapabilities: + supports_reasoning: bool + supports_adaptive_thinking: bool + thinking_always_on: bool + supports_legacy_thinking: bool + supports_output_config: bool + supports_sampling_params: bool + supports_speed: bool + effort_tiers: EffortTiers + + +@dataclass(frozen=True, slots=True) +class MessagesShaping: + capabilities: ModelCapabilities + drop_params: bool + reasoning_auto_summary: bool + additional_drop_params: Sequence[str] + def response(value: Mapping[str, object]) -> AnthropicMessagesResponse: return cast( # cast-ok: AnthropicMessagesResponse is a TypedDict over the normalized native payload @@ -20,4 +58,72 @@ def arguments(request: LiteLLMMessagesRequest) -> Mapping[str, object]: def map_failure(error: Exception, request: LiteLLMMessagesRequest, request_provider: str) -> Exception: + if getattr(error, "messages_request_error", False): + return litellm.BadRequestError( + message=str(error), + model=request.model.removeprefix(f"{request_provider}/"), + llm_provider=request_provider, + ) return failures.map_native_failure(error, request.model, request_provider, arguments(request), request.api_base) + + +def _resolved_provider(model: str, custom_llm_provider: str | None) -> tuple[str, str]: + try: + resolved_model, provider, _, _ = litellm.get_llm_provider(model=model, custom_llm_provider=custom_llm_provider) + except Exception: # noqa: BLE001 # an unroutable model still shapes as a bare Anthropic id + return model, custom_llm_provider or "anthropic" + return resolved_model, provider + + +def model_capabilities(model: str, custom_llm_provider: str | None) -> ModelCapabilities: + from litellm.llms.anthropic.chat.transformation import AnthropicConfig + from litellm.llms.anthropic.common_utils import AnthropicModelInfo + + resolved_model, provider = _resolved_provider(model, custom_llm_provider) + + def supports(flag: str) -> bool: + return AnthropicModelInfo._supports_model_capability(model, flag, provider) # pyright: ignore[reportPrivateUsage] # same probes the Python transform runs; forking them would drift + + def tier(level: str) -> bool: + return AnthropicConfig._supports_effort_level(model, level, provider) # pyright: ignore[reportPrivateUsage] # same probe the Python transform runs + + return ModelCapabilities( + supports_reasoning=supports("supports_reasoning"), + supports_adaptive_thinking=supports("supports_adaptive_thinking"), + thinking_always_on=supports("thinking_always_on"), + supports_legacy_thinking=supports("supports_legacy_thinking"), + supports_output_config=supports("supports_output_config"), + supports_sampling_params=AnthropicModelInfo._supports_sampling_params(resolved_model), # pyright: ignore[reportPrivateUsage] # same gate the handler applies + supports_speed=AnthropicConfig._model_supports_speed_param(resolved_model, provider), # pyright: ignore[reportPrivateUsage] # same gate the handler applies + effort_tiers=EffortTiers( + minimal=tier("minimal"), + low=tier("low"), + medium=tier("medium"), + high=tier("high"), + xhigh=tier("xhigh"), + max=tier("max"), + ), + ) + + +def _drop_params(kwargs: Mapping[str, object]) -> bool: + return bool(litellm.drop_params) or normalize_drop_params(kwargs.get("drop_params")) is True + + +def _additional_drop_params(kwargs: Mapping[str, object]) -> tuple[str, ...]: + try: + configured: Final = _DROP_PATHS.validate_python(kwargs.get("additional_drop_params")) + except ValidationError: + return () + return tuple(path for path in configured if isinstance(path, str)) + + +def shaping(model: str, custom_llm_provider: str | None, kwargs: Mapping[str, object]) -> dict[str, object]: + return asdict( + MessagesShaping( + capabilities=model_capabilities(model, custom_llm_provider), + drop_params=_drop_params(kwargs), + reasoning_auto_summary=is_reasoning_auto_summary_enabled(), + additional_drop_params=_additional_drop_params(kwargs), + ) + ) diff --git a/tests/test_litellm/rust_bridge/AGENTS.md b/tests/test_litellm/rust_bridge/AGENTS.md index 351bd582ce7..994d24112e8 100644 --- a/tests/test_litellm/rust_bridge/AGENTS.md +++ b/tests/test_litellm/rust_bridge/AGENTS.md @@ -5,3 +5,5 @@ Test what each side of the bridge does, not the rollout policy that picks a side Call each path directly with an explicit decision instead. The Python path is the implementation the dispatcher falls back to, e.g. `litellm.ocr.main.ocr`. The Rust path is the native binding, e.g. `NATIVE_OCR.load()` from `litellm/rust_bridge/ocr/entrypoints.py`, called with the request, args and kwargs that dispatch would hand it. When the native side reads a policy-derived setting such as `settings.secret_manager().native`, pin that field in the test instead of deriving it from the catalog. `ocr/test_secrets.py` shows the pattern Rollout policy itself, meaning which rule matches and what `LITELLM_RUST` changes, belongs in `test_catalog.py`, `test_configuration.py` and `test_dispatch.py`, tested against rules the test builds rather than the shipped `catalog.RULES` + +Before adding a test here, ask whether it checks something Rust cannot. A `route_host.py` module is the Python half of a native route: it projects Python-only state (the cost map, `litellm.*` settings, request kwargs) into the plain values the Rust side consumes, and maps native failures back onto public exceptions. Those projections are what belongs here, because a wrong key or an ignored provider prefix ships the wrong value to Rust and no Rust test sees it. `messages/test_route_host.py` shows the shape. Behavior that lives in Rust (a request transform given its inputs, header assembly, stream relay) is tested in the crate, and the route end to end is tested against a recording server in `tests/test_litellm_rust/`. A test that only re-checks a Python helper the route host happens to call is a duplicate of that helper's own test and should not be added diff --git a/tests/test_litellm/rust_bridge/messages/test_route_host.py b/tests/test_litellm/rust_bridge/messages/test_route_host.py new file mode 100644 index 00000000000..f47333a45d9 --- /dev/null +++ b/tests/test_litellm/rust_bridge/messages/test_route_host.py @@ -0,0 +1,112 @@ +from dataclasses import astuple +from typing import Final + +import pytest + +import litellm +from litellm.rust_bridge.messages import route_host + +pytestmark = pytest.mark.usefixtures("local_model_cost_map") + + +def _flag_model(monkeypatch: pytest.MonkeyPatch, name: str, **flags: bool) -> None: + monkeypatch.setitem( + litellm.model_cost, + name, + { + "litellm_provider": "anthropic", + "mode": "chat", + "input_cost_per_token": 0, + "output_cost_per_token": 0, + **flags, + }, + ) + + +def test_capabilities_come_from_the_model_map_under_the_callers_provider(monkeypatch: pytest.MonkeyPatch) -> None: + _flag_model( + monkeypatch, + "claude-test-adaptive", + supports_reasoning=True, + supports_adaptive_thinking=True, + supports_output_config=True, + supports_xhigh_reasoning_effort=True, + supports_sampling_params=False, + ) + + capabilities: Final = route_host.model_capabilities("anthropic/claude-test-adaptive", None) + + assert capabilities.supports_adaptive_thinking + assert capabilities.supports_output_config + assert not capabilities.supports_legacy_thinking + assert not capabilities.supports_sampling_params + assert capabilities.effort_tiers.xhigh + assert not capabilities.effort_tiers.max + + +def test_unmapped_model_keeps_sampling_params_and_no_reasoning_features() -> None: + capabilities: Final = route_host.model_capabilities("anthropic/not-a-real-model", None) + + assert capabilities.supports_sampling_params + assert not capabilities.supports_reasoning + assert not capabilities.supports_adaptive_thinking + assert not any(astuple(capabilities.effort_tiers)) + + +@pytest.mark.parametrize( + ("global_flag", "kwargs", "expected"), + [ + (False, {}, False), + (True, {}, True), + (False, {"drop_params": "true"}, True), + (False, {"drop_params": "nonsense"}, False), + (False, {"drop_params": False}, False), + ], +) +def test_drop_params_merges_the_global_flag_with_the_request( + monkeypatch: pytest.MonkeyPatch, global_flag: bool, kwargs: dict[str, object], expected: bool +) -> None: + monkeypatch.setattr(litellm, "drop_params", global_flag) + + assert route_host.shaping("anthropic/not-a-real-model", None, kwargs)["drop_params"] is expected + + +@pytest.mark.parametrize( + ("configured", "expected"), + [ + (["tools[*].input_examples", 3, "metadata.user_id"], ("tools[*].input_examples", "metadata.user_id")), + ("tools", ()), + (None, ()), + ], +) +def test_additional_drop_params_keep_only_string_paths(configured: object, expected: tuple[str, ...]) -> None: + shaping: Final = route_host.shaping("anthropic/not-a-real-model", None, {"additional_drop_params": configured}) + + assert shaping["additional_drop_params"] == expected + + +def test_native_request_rejections_map_to_the_public_400() -> None: + from types import MappingProxyType + + from litellm.rust_bridge.messages.entrypoints import LiteLLMMessagesRequest + + request: Final = LiteLLMMessagesRequest( + model="anthropic/claude-sonnet-5", + messages=(), + max_tokens=8, + stream=None, + api_key=None, + api_base=None, + custom_llm_provider=None, + kwargs=MappingProxyType({}), + ) + rejected: Final = ValueError("claude-sonnet-5 does not support top_k=5") + rejected.messages_request_error = True # pyright: ignore[reportAttributeAccessIssue] # marker the native host sets + + mapped: Final = route_host.map_failure(rejected, request, "anthropic") + + assert isinstance(mapped, litellm.BadRequestError) + assert mapped.status_code == 400 + assert "does not support top_k=5" in mapped.message + assert mapped.model == "claude-sonnet-5" + assert not isinstance(route_host.map_failure(ValueError("plain"), request, "anthropic"), litellm.BadRequestError) diff --git a/tests/test_litellm/rust_bridge/messages/test_secrets.py b/tests/test_litellm/rust_bridge/messages/test_secrets.py new file mode 100644 index 00000000000..cf37ed0830b --- /dev/null +++ b/tests/test_litellm/rust_bridge/messages/test_secrets.py @@ -0,0 +1,111 @@ +from __future__ import annotations + +from collections.abc import Awaitable, Mapping +from dataclasses import replace +from types import MappingProxyType +from typing import Final, Protocol, cast # noqa: TID251 # narrows the parametrized path to its protocol + +import httpx +import pytest + +import litellm +from litellm.integrations.custom_secret_manager import CustomSecretManager +from litellm.llms.anthropic.experimental_pass_through.messages.handler import anthropic_messages +from litellm.rust_bridge import settings +from litellm.rust_bridge.messages.entrypoints import NATIVE_AMESSAGES, NATIVE_MESSAGES, LiteLLMMessagesRequest +from litellm.types.secret_managers.main import KeyManagementSettings, KeyManagementSystem +from tests.test_litellm_rust.support.recording_server import ResponseSpec, recording_service +from tests.test_litellm_rust.support.requests import MESSAGES, MESSAGES_MODEL, MESSAGES_RESPONSE + +pytest.importorskip("litellm.rust_bridge._native") + +pytestmark = pytest.mark.usefixtures("local_model_cost_map") + + +class Messages(Protocol): + def __call__(self) -> Awaitable[object]: ... + + +class _ManagedSecrets(CustomSecretManager): + def __init__(self, values: Mapping[str, str]) -> None: + super().__init__(secret_manager_name="rust_bridge_messages_test") + self.values: Final = values + + async def async_read_secret( + self, + secret_name: str, + optional_params: dict[str, object] | None = None, + timeout: float | httpx.Timeout | None = None, + ) -> str | None: + raise AssertionError("get_secret reads custom managers synchronously") + + def sync_read_secret( + self, + secret_name: str, + optional_params: dict[str, object] | None = None, + timeout: float | httpx.Timeout | None = None, + ) -> str | None: + return self.values.get(secret_name) + + +def _native_request() -> LiteLLMMessagesRequest: + return LiteLLMMessagesRequest( + model=MESSAGES_MODEL, + messages=MESSAGES, + max_tokens=8, + stream=None, + api_key=None, + api_base=None, + custom_llm_provider=None, + kwargs=MappingProxyType({}), + ) + + +def _public_kwargs() -> dict[str, object]: + return {"model": MESSAGES_MODEL, "messages": [dict(message) for message in MESSAGES], "max_tokens": 8} + + +async def _python_messages() -> object: + return await anthropic_messages(**_public_kwargs()) + + +async def _rust_messages() -> object: + route: Final = NATIVE_MESSAGES.load() + assert route is not None + return route(_native_request(), (), _public_kwargs()) + + +async def _rust_amessages() -> object: + route: Final = NATIVE_AMESSAGES.load() + assert route is not None + return await route(_native_request(), (), _public_kwargs()) + + +@pytest.fixture( + params=(_python_messages, _rust_messages, _rust_amessages), ids=("python-async", "rust-sync", "rust-async") +) +def messages(request: pytest.FixtureRequest) -> Messages: + return cast(Messages, request.param) + + +async def test_secret_manager_supplies_the_anthropic_key_and_base( + monkeypatch: pytest.MonkeyPatch, messages: Messages +) -> None: + for name in ("ANTHROPIC_API_KEY", "ANTHROPIC_AUTH_TOKEN", "ANTHROPIC_API_BASE", "ANTHROPIC_BASE_URL"): + monkeypatch.delenv(name, raising=False) + with recording_service() as server: + server.default_response = ResponseSpec(body=MESSAGES_RESPONSE) + monkeypatch.setattr( + litellm, + "secret_manager_client", + _ManagedSecrets({"ANTHROPIC_API_KEY": "vault-key", "ANTHROPIC_BASE_URL": server.base_url}), + ) + monkeypatch.setattr(litellm, "_key_management_system", KeyManagementSystem.CUSTOM) + monkeypatch.setattr(litellm, "_key_management_settings", KeyManagementSettings(access_mode="read_only")) + configured: Final = settings.secret_manager + monkeypatch.setattr(settings, "secret_manager", lambda: replace(configured(), native=True)) + + await messages() + + assert len(server.requests) == 1 + assert server.requests[0].headers["x-api-key"] == "vault-key" diff --git a/tests/test_litellm_rust/messages/test_request_shaping.py b/tests/test_litellm_rust/messages/test_request_shaping.py new file mode 100644 index 00000000000..f885fda5f42 --- /dev/null +++ b/tests/test_litellm_rust/messages/test_request_shaping.py @@ -0,0 +1,237 @@ +"""The native Messages route shapes the wire request the way the Python handler does. + +Model capability expectations come from model_prices_and_context_window.json (Claude Sonnet 5 is an +adaptive-thinking model without sampling params; Claude Haiku 4.5 is a legacy-thinking model), read at +2026-09-24; the cost map is LiteLLM's own file. +""" + +from collections.abc import Iterator +from typing import Final + +import pytest + +import litellm +from litellm.rust_bridge import catalog +from litellm.rust_bridge.catalog import Route, RouteRule +from litellm.rust_bridge.configuration import Rollout +from tests.test_litellm_rust.support.isolation import rebound +from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec +from tests.test_litellm_rust.support.requests import MESSAGES, MESSAGES_RESPONSE + +pytestmark = pytest.mark.requires_rust_extension + +ADAPTIVE_MODEL: Final = "anthropic/claude-sonnet-5" +LEGACY_THINKING_MODEL: Final = "anthropic/claude-haiku-4-5" + + +@pytest.fixture(autouse=True) +def opt_messages_into_rust() -> Iterator[None]: + with rebound(catalog, "RULES", (RouteRule(Route.MESSAGES, Rollout.RUST_OPT_IN), *catalog.RULES)): + yield + + +@pytest.fixture +def messages_server(recording_server: RecordingServer) -> RecordingServer: + recording_server.default_response = ResponseSpec(body=MESSAGES_RESPONSE) + return recording_server + + +def arguments(server: RecordingServer, **kwargs: object) -> dict[str, object]: + return { + "model": ADAPTIVE_MODEL, + "messages": [dict(message) for message in MESSAGES], + "max_tokens": 8192, + "api_key": "test-key", + "api_base": server.base_url, + **kwargs, + } + + +def sent(server: RecordingServer) -> tuple[dict[str, object], dict[str, str]]: + assert len(server.requests) == 1 + request: Final = server.requests[0] + assert not request.headers.get("user-agent", "").startswith("python-httpx") + assert isinstance(request.body, dict) + return request.body, request.headers + + +@pytest.mark.asyncio +async def test_reasoning_effort_becomes_adaptive_thinking_and_effort_on_the_wire( + messages_server: RecordingServer, +) -> None: + await litellm.anthropic.messages.acreate(**arguments(messages_server, reasoning_effort="high")) + + body, _ = sent(messages_server) + assert "reasoning_effort" not in body + assert body["thinking"] == {"type": "adaptive", "display": "summarized"} + assert body["output_config"] == {"effort": "high"} + + +@pytest.mark.asyncio +async def test_claude_code_adaptive_payload_is_downgraded_to_a_capped_budget_for_a_legacy_model( + messages_server: RecordingServer, +) -> None: + await litellm.anthropic.messages.acreate( + **arguments( + messages_server, + model=LEGACY_THINKING_MODEL, + max_tokens=3000, + thinking={"type": "adaptive"}, + output_config={"effort": "high"}, + temperature=0, + ) + ) + + body, _ = sent(messages_server) + assert body["thinking"] == {"type": "enabled", "budget_tokens": 2999} + assert "output_config" not in body + assert "temperature" not in body + + +@pytest.mark.asyncio +async def test_removed_sampling_params_are_dropped_under_drop_params(messages_server: RecordingServer) -> None: + await litellm.anthropic.messages.acreate( + **arguments(messages_server, temperature=0.2, top_p=0.9, top_k=5, drop_params=True) + ) + + body, _ = sent(messages_server) + assert not {"temperature", "top_p", "top_k"} & body.keys() + + +@pytest.mark.asyncio +async def test_removed_sampling_params_are_rejected_without_drop_params(messages_server: RecordingServer) -> None: + messages_server.expected_requests = 0 + + with pytest.raises(litellm.BadRequestError, match="does not support top_k=5"): + await litellm.anthropic.messages.acreate(**arguments(messages_server, top_k=5)) + + assert messages_server.requests == [] + + +@pytest.mark.asyncio +async def test_replayed_history_is_sanitized_before_it_reaches_the_provider( + messages_server: RecordingServer, +) -> None: + history: Final = [ + {"role": "user", "content": "run it"}, + { + "role": "assistant", + "content": [ + {"type": "text", "text": ""}, + { + "type": "tool_use", + "id": "functions.Bash:0", + "name": "Bash", + "input": {}, + "provider_specific_fields": {"x": 1}, + }, + ], + }, + {"role": "user", "content": [{"type": "tool_result", "tool_use_id": "functions.Bash:0", "content": "ok"}]}, + ] + + await litellm.anthropic.messages.acreate(**arguments(messages_server, messages=history)) + + body, _ = sent(messages_server) + assert body["messages"] == [ + {"role": "user", "content": "run it"}, + {"role": "assistant", "content": [{"type": "tool_use", "id": "functions_Bash_0", "name": "Bash", "input": {}}]}, + {"role": "user", "content": [{"type": "tool_result", "tool_use_id": "functions_Bash_0", "content": "ok"}]}, + ] + + +@pytest.mark.asyncio +async def test_feature_betas_merge_into_the_forwarded_beta_header(messages_server: RecordingServer) -> None: + await litellm.anthropic.messages.acreate( + **arguments( + messages_server, + output_format={"type": "json_schema", "schema": {"type": "object"}}, + extra_headers={"anthropic-beta": "web-search-2025-03-05"}, + ) + ) + + _, headers = sent(messages_server) + assert headers["anthropic-beta"] == "structured-outputs-2025-11-13,web-search-2025-03-05" + + +@pytest.mark.asyncio +async def test_oauth_token_authenticates_as_a_bearer_with_the_oauth_beta(messages_server: RecordingServer) -> None: + await litellm.anthropic.messages.acreate(**arguments(messages_server, api_key="sk-ant-oat01-token")) + + _, headers = sent(messages_server) + assert "x-api-key" not in headers + assert headers["authorization"] == "Bearer sk-ant-oat01-token" + assert headers["anthropic-beta"] == "oauth-2025-04-20" + assert headers["anthropic-dangerous-direct-browser-access"] == "true" + + +@pytest.mark.asyncio +async def test_metadata_is_reduced_to_the_fields_anthropic_accepts(messages_server: RecordingServer) -> None: + await litellm.anthropic.messages.acreate( + **arguments(messages_server, metadata={"user_id": "u-1", "trace_id": "internal"}) + ) + + body, _ = sent(messages_server) + assert body["metadata"] == {"user_id": "u-1"} + + +@pytest.mark.asyncio +async def test_additional_drop_params_remove_nested_fields_from_the_wire(messages_server: RecordingServer) -> None: + tools: Final = [{"name": "lookup", "input_schema": {"type": "object"}, "input_examples": [{"q": "x"}]}] + + await litellm.anthropic.messages.acreate( + **arguments(messages_server, tools=tools, additional_drop_params=["tools[*].input_examples"]) + ) + + body, _ = sent(messages_server) + assert body["tools"] == [{"name": "lookup", "input_schema": {"type": "object"}}] + + +@pytest.mark.asyncio +async def test_provider_specific_headers_scoped_to_anthropic_reach_the_wire(messages_server: RecordingServer) -> None: + await litellm.anthropic.messages.acreate( + **arguments( + messages_server, + provider_specific_header=[ + {"custom_llm_provider": "anthropic, azure_ai", "extra_headers": {"x-scoped": "yes"}}, + {"custom_llm_provider": "openai", "extra_headers": {"x-other": "no"}}, + ], + ) + ) + + _, headers = sent(messages_server) + assert headers["x-scoped"] == "yes" + assert "x-other" not in headers + + +@pytest.mark.asyncio +async def test_scoped_headers_override_extra_headers_which_override_forwarded_headers( + messages_server: RecordingServer, +) -> None: + await litellm.anthropic.messages.acreate( + **arguments( + messages_server, + headers={"x-priority": "forwarded", "x-forwarded-only": "kept"}, + extra_headers={"x-priority": "extra", "x-extra-only": "kept"}, + provider_specific_header={"custom_llm_provider": "anthropic", "extra_headers": {"x-priority": "scoped"}}, + ) + ) + + _, headers = sent(messages_server) + assert {name: headers.get(name) for name in ("x-priority", "x-forwarded-only", "x-extra-only")} == { + "x-priority": "scoped", + "x-forwarded-only": "kept", + "x-extra-only": "kept", + } + + +@pytest.mark.asyncio +async def test_non_string_metadata_user_id_is_rejected_before_the_provider_call( + messages_server: RecordingServer, +) -> None: + messages_server.expected_requests = 0 + + with pytest.raises(litellm.BadRequestError, match=r"metadata\.user_id must be a string"): + await litellm.anthropic.messages.acreate(**arguments(messages_server, metadata={"user_id": 123})) + + assert messages_server.requests == [] From f39a56b004c83f4242f7f95c48504235f1af2180 Mon Sep 17 00:00:00 2001 From: ahamedshaik16 Date: Fri, 25 Sep 2026 00:45:01 +0530 Subject: [PATCH 068/120] fix(prometheus): add model_group label to deployment request and rate limit metrics (#42966) * fix(prometheus): add model_group label to deployment request and rate limit metrics litellm_deployment_total_requests, litellm_deployment_success_responses, litellm_deployment_failure_responses, litellm_deployment_tpm_limit and litellm_deployment_rpm_limit had no way to identify which model_group a pooled deployment belongs to, only requested_model, litellm_model_name and model_id, none of which name the alias a model_name resolves through when it fans out to more than one deployment. model_group was already resolved onto enum_values for every request in async_log_success_event, so this is a label-list addition for the metrics built directly from that enum_values (the two request counters). The failure counter builds its own UserAPIKeyLabelValues locally and had a model_group variable already in scope that it never passed through, and the tpm/rpm limit gauges are set from a helper that took no model_group parameter at all even though its only caller already had it on enum_values. Both now thread the value through. * test(prometheus): expect model_group in deployment success/total request labels test_set_llm_deployment_success_metrics asserts the exact label set passed to litellm_deployment_success_responses.labels() and litellm_deployment_total_requests.labels(), which now includes model_group since it was added to those metrics' label list. * fix(prometheus): bound model_group on deployment failure metrics On a pre-routing reject (no deployment selected), model_group is caller-supplied via litellm_params.metadata and was passed through unbounded, letting an unrecognized value mint unlimited label series on litellm_deployment_failure_responses / litellm_deployment_total_requests. Bound it with the same _bounded_requested_model_label used for requested_model on this path. When a deployment is actually selected, model_group is router-resolved and passed through as-is. Also documents the model_group parameter on _set_deployment_tpm_rpm_limit_metrics and the bounding behavior on set_llm_deployment_failure_metrics. --------- Co-authored-by: ahamedshaik16 <24526479+ahamedshaik16@users.noreply.github.com> --- litellm/integrations/prometheus.py | 25 +++ litellm/types/integrations/prometheus.py | 3 + .../test_prometheus_logging_callbacks.py | 2 + .../integrations/test_prometheus_labels.py | 170 ++++++++++++++++++ 4 files changed, 200 insertions(+) diff --git a/litellm/integrations/prometheus.py b/litellm/integrations/prometheus.py index 2fdcb8ef745..fb010ab5886 100644 --- a/litellm/integrations/prometheus.py +++ b/litellm/integrations/prometheus.py @@ -2778,6 +2778,13 @@ class PrometheusLogger(CustomLogger): - increment deployment failure responses metric - increment deployment total requests metric + Both counters also carry a model_group label. When a deployment was + actually selected, model_group is the router-resolved value and is + trusted as-is. On a pre-routing reject (no deployment selected), it + is caller-supplied via litellm_params.metadata and is bounded with + _bounded_requested_model_label the same way requested_model is, so an + unrecognized value cannot mint unbounded label series. + Args: request_kwargs: dict @@ -2844,6 +2851,7 @@ class PrometheusLogger(CustomLogger): label_api_base = api_base label_api_provider = llm_provider label_requested_model = model_group or litellm_model_name + label_model_group = model_group else: label_litellm_model_name = "" label_model_id = "" @@ -2852,6 +2860,7 @@ class PrometheusLogger(CustomLogger): label_requested_model = ( _bounded_requested_model_label(litellm_model_name or model_group, router_originated=True) or "" ) + label_model_group = _bounded_requested_model_label(model_group, router_originated=True) enum_values: Final = UserAPIKeyLabelValues( litellm_model_name=label_litellm_model_name, @@ -2861,6 +2870,7 @@ class PrometheusLogger(CustomLogger): exception_status=exception_status, exception_class=(self._get_exception_class_name(exception) if exception else None), requested_model=label_requested_model, + model_group=label_model_group, hashed_api_key=hashed_api_key, api_key_alias=api_key_alias, user_email=user_email, @@ -2912,9 +2922,21 @@ class PrometheusLogger(CustomLogger): model_id: str | None, api_base: str | None, llm_provider: str | None, + model_group: str | None, ): """ Set the deployment TPM and RPM limits metrics + + Args: + model_info: the deployment's static model_info config (id, tpm, rpm, etc.) + litellm_params: the deployment's litellm_params, as a tpm/rpm fallback source + litellm_model_name: the resolved deployment model name + model_id: the deployment's model_id + api_base: the deployment's api_base + llm_provider: the deployment's custom_llm_provider + model_group: the router-resolved model_group the deployment belongs to, + from the caller's already-resolved enum_values.model_group (trusted, + not caller-supplied at this call site) """ tpm: Final = model_info.get("tpm") or litellm_params.get("tpm") rpm: Final = model_info.get("rpm") or litellm_params.get("rpm") @@ -2927,6 +2949,7 @@ class PrometheusLogger(CustomLogger): model_id=model_id, api_base=api_base, api_provider=llm_provider, + model_group=model_group, ), ) self.litellm_deployment_tpm_limit.labels(**_labels).set(tpm) @@ -2939,6 +2962,7 @@ class PrometheusLogger(CustomLogger): model_id=model_id, api_base=api_base, api_provider=llm_provider, + model_group=model_group, ), ) self.litellm_deployment_rpm_limit.labels(**_labels).set(rpm) @@ -3058,6 +3082,7 @@ class PrometheusLogger(CustomLogger): model_id=model_id, api_base=api_base, llm_provider=llm_provider, + model_group=enum_values.model_group, ) remaining_requests: int | None = None diff --git a/litellm/types/integrations/prometheus.py b/litellm/types/integrations/prometheus.py index c929ee2ee79..239fc7f2779 100644 --- a/litellm/types/integrations/prometheus.py +++ b/litellm/types/integrations/prometheus.py @@ -664,6 +664,7 @@ class PrometheusMetricLabels: ] litellm_deployment_tpm_limit = [ + UserAPIKeyLabelNames.MODEL_GROUP.value, UserAPIKeyLabelNames.v2_LITELLM_MODEL_NAME.value, UserAPIKeyLabelNames.MODEL_ID.value, UserAPIKeyLabelNames.API_BASE.value, @@ -770,6 +771,7 @@ class PrometheusMetricLabels: # Add deployment metrics litellm_deployment_failure_responses = [ + UserAPIKeyLabelNames.MODEL_GROUP.value, UserAPIKeyLabelNames.REQUESTED_MODEL.value, UserAPIKeyLabelNames.v2_LITELLM_MODEL_NAME.value, UserAPIKeyLabelNames.MODEL_ID.value, @@ -786,6 +788,7 @@ class PrometheusMetricLabels: ] litellm_deployment_total_requests = [ + UserAPIKeyLabelNames.MODEL_GROUP.value, UserAPIKeyLabelNames.REQUESTED_MODEL.value, UserAPIKeyLabelNames.v2_LITELLM_MODEL_NAME.value, UserAPIKeyLabelNames.MODEL_ID.value, diff --git a/tests/enterprise/litellm_enterprise/enterprise_callbacks/test_prometheus_logging_callbacks.py b/tests/enterprise/litellm_enterprise/enterprise_callbacks/test_prometheus_logging_callbacks.py index 58cde4c8103..92ff3d5813c 100644 --- a/tests/enterprise/litellm_enterprise/enterprise_callbacks/test_prometheus_logging_callbacks.py +++ b/tests/enterprise/litellm_enterprise/enterprise_callbacks/test_prometheus_logging_callbacks.py @@ -1031,6 +1031,7 @@ def test_set_llm_deployment_success_metrics(prometheus_logger): api_base="https://api.openai.com", api_provider="openai", requested_model="my_custom_model_group", + model_group="my_custom_model_group", hashed_api_key=standard_logging_payload["metadata"]["user_api_key_hash"], api_key_alias=standard_logging_payload["metadata"]["user_api_key_alias"], team=standard_logging_payload["metadata"]["user_api_key_team_id"], @@ -1047,6 +1048,7 @@ def test_set_llm_deployment_success_metrics(prometheus_logger): api_base="https://api.openai.com", api_provider="openai", requested_model="my_custom_model_group", + model_group="my_custom_model_group", hashed_api_key=standard_logging_payload["metadata"]["user_api_key_hash"], api_key_alias=standard_logging_payload["metadata"]["user_api_key_alias"], team=standard_logging_payload["metadata"]["user_api_key_team_id"], diff --git a/tests/test_litellm/integrations/test_prometheus_labels.py b/tests/test_litellm/integrations/test_prometheus_labels.py index 200f8e65add..41d0c44ff89 100644 --- a/tests/test_litellm/integrations/test_prometheus_labels.py +++ b/tests/test_litellm/integrations/test_prometheus_labels.py @@ -787,6 +787,176 @@ async def test_failure_hook_prefers_request_data_provider_over_exception_provide ) == ["azure"] +def test_model_group_in_deployment_metrics(): + """ + Test that model_group label is present on the deployment-scoped metrics + needed to build model-group dashboards (request counts, success/failure + counts, tpm/rpm limits). These metrics previously only carried + requested_model, litellm_model_name and model_id, none of which identify + the model_group a pooled deployment belongs to. + """ + model_group_label = UserAPIKeyLabelNames.MODEL_GROUP.value + + metrics_with_model_group = [ + "litellm_deployment_total_requests", + "litellm_deployment_success_responses", + "litellm_deployment_failure_responses", + "litellm_deployment_tpm_limit", + "litellm_deployment_rpm_limit", + ] + + for metric_name in metrics_with_model_group: + labels = PrometheusMetricLabels.get_labels(metric_name) + assert ( + model_group_label in labels + ), f"Metric {metric_name} should contain model_group label" + print(f"✅ {metric_name} contains model_group label") + + +def test_model_group_value_flows_through_deployment_metrics_label_factory(): + """ + The label being in the allow-list is necessary but not sufficient: the + factory must also carry the value from the enum through to the emitted + label. This would fail if the label were dropped from a metric's list or + if the value plumbing regressed, which the allow-list assertion above + cannot catch on its own. + """ + from unittest.mock import MagicMock + + from litellm.integrations.prometheus import ( + PrometheusLogger, + UserAPIKeyLabelValues, + prometheus_label_factory, + ) + + prometheus_logger = MagicMock() + prometheus_logger._cached_metric_labels = {} + prometheus_logger.label_filters = {} + prometheus_logger.get_labels_for_metric = ( + PrometheusLogger.get_labels_for_metric.__get__(prometheus_logger) + ) + + enum_values = UserAPIKeyLabelValues( + model_group="example-model-group", + litellm_model_name="gpt-4o-mini", + requested_model="example-model-group", + status_code="200", + ) + + for metric_name in [ + "litellm_deployment_total_requests", + "litellm_deployment_success_responses", + "litellm_deployment_failure_responses", + "litellm_deployment_tpm_limit", + "litellm_deployment_rpm_limit", + ]: + labels = prometheus_label_factory( + supported_enum_labels=prometheus_logger.get_labels_for_metric( + metric_name=metric_name + ), + enum_values=enum_values, + ) + assert ( + labels.get("model_group") == "example-model-group" + ), f"{metric_name} should emit model_group=example-model-group, got {labels.get('model_group')!r}" + + +def test_deployment_failure_metrics_emit_model_group_from_standard_logging_payload(): + """ + End-to-end emit wiring for the failure path. + + The label-list and factory tests above prove the label exists and that + the factory carries a value handed to it, but neither drives the real + set_llm_deployment_failure_metrics code path, so deleting the production + model_group=model_group assignment there would still pass them. This + calls it directly with a standard_logging_object carrying model_group and + asserts the real litellm_deployment_failure_responses / _total_requests + Counter series actually carry it. + """ + from litellm.integrations.prometheus import PrometheusLogger + + _clear_prometheus_registry() + try: + logger = PrometheusLogger() + logger.set_llm_deployment_failure_metrics( + request_kwargs={ + "model": "gpt-4o-mini", + "litellm_params": {"metadata": {}}, + "standard_logging_object": { + "model_group": "example-model-group", + "model_id": "model-123", + "api_base": "https://api.openai.com", + "request_tags": [], + }, + "exception": Exception("boom"), + } + ) + + for metric in ( + logger.litellm_deployment_failure_responses, + logger.litellm_deployment_total_requests, + ): + index = metric._labelnames.index("model_group") + values = {sample_key[index] for sample_key in metric._metrics} + assert values == {"example-model-group"}, ( + f"expected model_group=example-model-group on {metric._name}, got {values}" + ) + finally: + _clear_prometheus_registry() + + +def test_deployment_tpm_rpm_limit_metrics_emit_model_group_from_enum_values(): + """ + End-to-end emit wiring for the tpm/rpm limit gauges. + + _set_deployment_tpm_rpm_limit_metrics used to build its own + UserAPIKeyLabelValues with no model_group parameter at all, dropping the + value even though its only caller (set_llm_deployment_success_metrics) + already had it on enum_values. This drives set_llm_deployment_success_metrics + directly with a deployment that has tpm/rpm configured and asserts the real + litellm_deployment_tpm_limit / litellm_deployment_rpm_limit Gauge series + carry model_group; it fails if that plumbing is removed. + """ + import datetime + + from litellm.integrations.prometheus import PrometheusLogger, UserAPIKeyLabelValues + + _clear_prometheus_registry() + try: + logger = PrometheusLogger() + now = datetime.datetime.now() + enum_values = UserAPIKeyLabelValues( + model_group="example-model-group", + litellm_model_name="gpt-4o-mini", + requested_model="example-model-group", + status_code="200", + ) + logger.set_llm_deployment_success_metrics( + request_kwargs={ + "model": "gpt-4o-mini", + "litellm_params": {"metadata": {"model_info": {"id": "model-123", "tpm": 1000, "rpm": 10}}}, + "standard_logging_object": { + "model_group": "example-model-group", + "model_id": "model-123", + "api_base": "https://api.openai.com", + "hidden_params": {"additional_headers": None, "litellm_overhead_time_ms": None}, + }, + }, + start_time=now, + end_time=now, + enum_values=enum_values, + ) + + for metric in (logger.litellm_deployment_tpm_limit, logger.litellm_deployment_rpm_limit): + index = metric._labelnames.index("model_group") + values = {sample_key[index] for sample_key in metric._metrics} + assert values == {"example-model-group"}, ( + f"expected model_group=example-model-group on {metric._name}, got {values}" + ) + finally: + _clear_prometheus_registry() + + if __name__ == "__main__": test_user_email_in_required_metrics() test_user_email_label_exists() From 7abaf4edb2458b46ac4014beecd0e778ab5a9c89 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 24 Sep 2026 12:30:42 -0700 Subject: [PATCH 069/120] fix(cost-map): add the Vertex shutdown date to gemini-2.5-flash-native-audio (#43024) Co-authored-by: kerry Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/model_prices_and_context_window_backup.json | 1 + model_prices_and_context_window.json | 1 + 2 files changed, 2 insertions(+) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 8f38866387a..dad89fc58a8 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -68565,6 +68565,7 @@ "source": "https://api.together.ai/v1/models" }, "vertex_ai/gemini-2.5-flash-native-audio": { + "deprecation_date": "2026-12-13", "input_cost_per_audio_token": 3e-06, "input_cost_per_token": 5e-07, "litellm_provider": "vertex_ai", diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 8f38866387a..dad89fc58a8 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -68565,6 +68565,7 @@ "source": "https://api.together.ai/v1/models" }, "vertex_ai/gemini-2.5-flash-native-audio": { + "deprecation_date": "2026-12-13", "input_cost_per_audio_token": 3e-06, "input_cost_per_token": 5e-07, "litellm_provider": "vertex_ai", From c9e8a04139af587d3280ea5fb248744e3785f500 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 24 Sep 2026 12:35:34 -0700 Subject: [PATCH 070/120] feat(vertex): native batch JSONL passthrough with cost tracking (#42810) * feat(vertex): native batch JSONL passthrough with cost tracking Add a per-request `passthrough=true` multipart field on `POST /v1/files` (and the same kwarg on `litellm.create_file`) that uploads a native Vertex AI batch JSONL to the deployment's GCS bucket unchanged, so rows using `googleSearch` and other Gemini-only features run as written and the output, `groundingMetadata` included, comes back untouched. Passthrough is sticky through the GCS object path (`litellm-vertex-files/passthrough/...`), so batch create and output retrieval inherit it without new state. Native output rows are costed from their `usageMetadata` with the deployment's model and model_info, in the polling and retrieve paths and for the existing global `disable_vertex_batch_output_transformation` flag, which billed $0 before. The proxy requires the target to resolve to vertex_ai deployments only, refuses `passthrough` with a non-batch purpose, a non-default `target_storage`, or pre-call guardrails, and validates native rows on `request` instead of the OpenAI batch keys. * refactor(vertex): keep native batch row pricing inside the Vertex adapter Moves native Vertex batch row detection, response parsing, and per-row pricing from litellm/batches/batch_utils.py into litellm/llms/vertex_ai/batches/transformation.py, so batch_utils only aggregates the rows it gets back. Adds tests/test_litellm/files to the misc unit shard so the new test directory is claimed by a shard. * fix(files): say what a passthrough batch upload takes when a row is not native The missing-key 400 listed bare key names, so an OpenAI-shaped row under passthrough=true read "Each line must be a JSON object with keys request". The batch line shape now carries its own hint, and the passthrough one says a passthrough upload takes native Vertex batch rows with a request key * fix(batches): bill native Vertex embedding batch rows on the native cost path A native Vertex output row whose response holds an embedding was validated as a generateContent response, so the documented tokenCount-only shape counted as a failed row. Price embedding rows from their own usage (promptTokenCount, else tokenCount) with the helper the transformed embeddings path already used, and drop the prompt-details helper nothing calls anymore. * fix(batches): keep modality batch rates on native Vertex embedding rows An embedding row that carries usageMetadata was billed from promptTokenCount alone, so its promptTokensDetails no longer reached the audio, image, and video batch rates the way it did before the native cost path. Run every row with usageMetadata through the Gemini usage parser and keep the flat tokenCount fallback for embedding rows without it. * fix(batches): price native Vertex batch rows by modelVersion under a wildcard deployment A `vertex_ai/*` deployment hands the batch cost path `*` as the deployment model, which no cost map resolves, so every native (passthrough or flag-on) row was billed at $0. A wildcard deployment model now defers to the row's own `modelVersion`, the way the transformed path already prices by the row's `model`. Also moves the native passthrough tests under tests/test_litellm, the tree codecov reads, and covers the raw upload chunking, the embedding output translation, the unpriceable-row path, and the flag-on dispatch. * fix(batches): keep explicit deployment prices for native Vertex rows without a modelVersion Under a wildcard deployment a native batch row that carries no modelVersion (an embedding row, or a generateContent row Vertex returned without one) was billed at $0 even when the deployment's model_info sets explicit batch prices, because the cost calculator was never called. The row now falls back to the wildcard name, which the cost calculator prices from the explicit model_info, and only a row with neither a modelVersion nor a deployment model is billed at $0 with the warning --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- .github/workflows/test-unit.yml | 1 + litellm/batches/batch_utils.py | 131 +++--- litellm/files/main.py | 16 + .../llms/vertex_ai/batches/transformation.py | 147 +++++-- .../llms/vertex_ai/files/transformation.py | 127 ++++-- .../batch_file_validation.py | 35 +- .../openai_files_endpoints/files_endpoints.py | 109 ++++- litellm/router.py | 1 + litellm/router_utils/batch_utils.py | 6 +- tests/e2e/batches/test_batches_e2e.py | 112 +++++ .../llm_nonconversational.yaml | 3 + tests/e2e/coverage_registry/schema.py | 1 + tests/e2e/e2e_http.py | 7 + .../test_router_batch_utils.py | 1 + .../test_litellm/batches/test_batch_utils.py | 387 ++++++++++++++++++ tests/test_litellm/files/__init__.py | 0 tests/test_litellm/files/test_main.py | 71 ++++ .../vertex_ai/batches/test_transformation.py | 38 +- .../llms/vertex_ai/files/__init__.py | 0 .../vertex_ai/files/test_transformation.py | 310 ++++++++++++++ .../test_files_batch_file_validation.py | 38 +- .../test_files_endpoint.py | 238 +++++++++++ tests/test_litellm/test_router.py | 33 ++ tests/unit/batches/test_batch_utils.py | 17 +- ui/litellm-dashboard/src/lib/http/schema.d.ts | 15 + 25 files changed, 1677 insertions(+), 167 deletions(-) create mode 100644 tests/test_litellm/batches/test_batch_utils.py create mode 100644 tests/test_litellm/files/__init__.py create mode 100644 tests/test_litellm/files/test_main.py create mode 100644 tests/test_litellm/llms/vertex_ai/files/__init__.py create mode 100644 tests/test_litellm/llms/vertex_ai/files/test_transformation.py diff --git a/.github/workflows/test-unit.yml b/.github/workflows/test-unit.yml index 4580ad17a19..686bbc89467 100644 --- a/.github/workflows/test-unit.yml +++ b/.github/workflows/test-unit.yml @@ -108,6 +108,7 @@ jobs: tests/test_litellm/completion_extras tests/test_litellm/containers tests/test_litellm/endpoints + tests/test_litellm/files tests/test_litellm/images tests/test_litellm/interactions tests/test_litellm/messages diff --git a/litellm/batches/batch_utils.py b/litellm/batches/batch_utils.py index 7209ac6a1e7..819a279a43c 100644 --- a/litellm/batches/batch_utils.py +++ b/litellm/batches/batch_utils.py @@ -11,7 +11,10 @@ from litellm.litellm_core_utils.get_litellm_params import AWS_CREDENTIAL_KWARGS_ from litellm.litellm_core_utils.llm_cost_calc.utils import parse_prompt_tokens_details from litellm.llms.base_llm.ocr.transformation import OCRUsageInfo from litellm.llms.bedrock.batches.transformation import titan_embedding_usage_from_batch_output -from litellm.llms.vertex_ai.batches.transformation import vertex_prompt_tokens_details +from litellm.llms.vertex_ai.batches.transformation import ( + is_native_vertex_batch_output_row, + native_vertex_batch_row_stats, +) from litellm.types.llms.openai import Batch from litellm.types.utils import ModelInfo, Usage from litellm.utils import token_counter @@ -31,6 +34,20 @@ class BatchCostUsageResult: _COMPLETED_BATCH_STATUSES: Final = frozenset({"completed", "complete"}) + + +def _uses_native_vertex_output( + custom_llm_provider: str, + model_name: str | None, + first_row: Mapping[str, object] | None, +) -> bool: + if custom_llm_provider != "vertex_ai": + return False + if model_name and getattr(litellm, "disable_vertex_batch_output_transformation", False): + return True + return first_row is not None and is_native_vertex_batch_output_row(first_row) + + _TERMINAL_BATCH_STATUSES: Final = _COMPLETED_BATCH_STATUSES | frozenset({"failed", "cancelled", "expired"}) @@ -66,12 +83,9 @@ async def calculate_batch_cost_and_usage( deployment-specific pricing (e.g. input_cost_per_token_batches) is used instead of the global cost map. """ - if ( - custom_llm_provider == "vertex_ai" - and model_name - and getattr(litellm, "disable_vertex_batch_output_transformation", False) - ): - return calculate_vertex_ai_batch_cost_and_usage(file_content_dictionary, model_name) + first_row: Final = file_content_dictionary[0] if file_content_dictionary else None + if _uses_native_vertex_output(custom_llm_provider, model_name, first_row): + return calculate_vertex_ai_batch_cost_and_usage(file_content_dictionary, model_name, model_info=model_info) return _aggregate_batch_cost_usage_models( entries=file_content_dictionary, @@ -126,11 +140,11 @@ async def _handle_completed_batch( ) output_file_result: Final = ( - calculate_vertex_ai_batch_cost_and_usage(_get_file_content_as_dictionary(file_content), model_name) - if ( - custom_llm_provider == "vertex_ai" - and model_name - and getattr(litellm, "disable_vertex_batch_output_transformation", False) + calculate_vertex_ai_batch_cost_and_usage( + _iter_batch_output_entries(file_content), model_name, model_info=model_info + ) + if _uses_native_vertex_output( + custom_llm_provider, model_name, next(_iter_batch_output_entries(file_content), None) ) else _aggregate_batch_cost_usage_models( entries=_iter_batch_output_entries(file_content), @@ -332,69 +346,36 @@ def _aggregate_batch_cost_usage_models( def calculate_vertex_ai_batch_cost_and_usage( - vertex_ai_batch_responses: list[dict], + vertex_ai_batch_responses: Iterable[dict], model_name: str | None = None, + model_info: ModelInfo | None = None, ) -> BatchCostUsageResult: """ - Calculate both cost and usage from raw Vertex AI batch responses. - - Used only when ``litellm.disable_vertex_batch_output_transformation = True``. - In that case the GCS predictions.jsonl is returned as-is, with each line in - the native Vertex format: - - {"request": ..., "response": {"candidates": [...], "usageMetadata": {...}}} - - usageMetadata contains promptTokenCount, candidatesTokenCount, totalTokenCount. - - A row with no ``response`` is counted as failed - the same signal already - used to skip it from cost/usage aggregation, since Vertex batch prediction - output doesn't establish a distinct error shape in this (non-default) path. + Cost and usage of a native Vertex predictions.jsonl, one + `{"request": ..., "response": {"candidates": [...], "usageMetadata": {...}, "modelVersion": ...}}` + generateContent row or `{"request": ..., "response": {"embedding": {...}, "usageMetadata": {...}}}` + embedding row per line. `model_name` (the deployment model) prices every row, else each row's own + `modelVersion` does; a row without a usable response counts as failed. """ from litellm.cost_calculator import batch_cost_calculator + from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import VertexGeminiConfig - total_prompt_cost = 0.0 # rebind-ok: loop accumulator, matches total_tokens below - total_completion_cost = 0.0 # rebind-ok: loop accumulator, matches total_tokens below - total_tokens = 0 - prompt_tokens = 0 - completion_tokens = 0 - successful_requests = 0 # rebind-ok: loop accumulator, matches total_cost/total_tokens above - failed_requests = 0 # rebind-ok: loop accumulator, matches total_cost/total_tokens above - actual_model_name: Final = model_name or "gemini-2.0-flash-001" - - for response in vertex_ai_batch_responses: - response_body = response.get("response") - if response_body is None: - failed_requests += 1 - continue - successful_requests += 1 - - usage_metadata = response_body.get("usageMetadata", {}) - _prompt = usage_metadata.get("promptTokenCount", 0) or 0 - _completion = usage_metadata.get("candidatesTokenCount", 0) or 0 - _total = usage_metadata.get("totalTokenCount", 0) or (_prompt + _completion) - - line_usage = Usage( - prompt_tokens=_prompt, - completion_tokens=_completion, - total_tokens=_total, - prompt_tokens_details=vertex_prompt_tokens_details(usage_metadata), + row_stats: Final = tuple( + native_vertex_batch_row_stats( + row, + model_name, + model_info=model_info, + calculate_usage=VertexGeminiConfig._calculate_usage, + cost_calculator=batch_cost_calculator, ) - - try: - p_cost, c_cost = batch_cost_calculator( - usage=line_usage, - model=actual_model_name, - custom_llm_provider="vertex_ai", - ) - total_prompt_cost += p_cost - total_completion_cost += c_cost - except Exception as e: - verbose_logger.debug("vertex_ai batch cost calculation error for line: %s", str(e)) - - prompt_tokens += _prompt - completion_tokens += _completion - total_tokens += _total - + for row in vertex_ai_batch_responses + ) + priced: Final = tuple(stats for stats in row_stats if stats is not None) + total_prompt_cost: Final = sum(stats.prompt_cost for stats in priced) + total_completion_cost: Final = sum(stats.completion_cost for stats in priced) + prompt_tokens: Final = sum(stats.usage.prompt_tokens for stats in priced) + completion_tokens: Final = sum(stats.usage.completion_tokens for stats in priced) + total_tokens: Final = sum(stats.total_tokens for stats in priced) total_cost: Final = total_prompt_cost + total_completion_cost verbose_logger.info( "vertex_ai batch cost: cost=%s, prompt=%d, completion=%d, total=%d, successful=%d, failed=%d", @@ -402,8 +383,8 @@ def calculate_vertex_ai_batch_cost_and_usage( prompt_tokens, completion_tokens, total_tokens, - successful_requests, - failed_requests, + len(priced), + len(row_stats) - len(priced), ) return BatchCostUsageResult( @@ -413,9 +394,13 @@ def calculate_vertex_ai_batch_cost_and_usage( prompt_tokens=prompt_tokens, completion_tokens=completion_tokens, ), - models=[actual_model_name], - successful_requests=successful_requests, - failed_requests=failed_requests, + models=( + [model_name] + if model_name + else list(dict.fromkeys(stats.model for stats in priced if stats.model is not None)) + ), + successful_requests=len(priced), + failed_requests=len(row_stats) - len(priced), prompt_cost=total_prompt_cost, completion_cost=total_completion_cost, ) diff --git a/litellm/files/main.py b/litellm/files/main.py index e0804244ff7..72832aeccc9 100644 --- a/litellm/files/main.py +++ b/litellm/files/main.py @@ -176,6 +176,22 @@ def create_file( if logging_obj is None: raise ValueError("logging_obj is required") client: Final = kwargs.get("client") + if litellm_params_dict.get("passthrough") is True and ( + custom_llm_provider != "vertex_ai" or purpose != "batch" + ): + raise litellm.exceptions.BadRequestError( + message=( + "`passthrough=True` uploads the file bytes unchanged for a native Vertex AI batch, so it needs " + f"custom_llm_provider='vertex_ai' and purpose='batch', got '{custom_llm_provider}' and '{purpose}'." + ), + model="n/a", + llm_provider=custom_llm_provider or "n/a", + response=httpx.Response( + status_code=400, + content="passthrough needs a vertex_ai batch", + request=httpx.Request(method="create_file", url="https://github.com/BerriAI/litellm"), + ), + ) ### TIMEOUT LOGIC ### timeout = optional_params.timeout or kwargs.get("request_timeout", 600) or 600 diff --git a/litellm/llms/vertex_ai/batches/transformation.py b/litellm/llms/vertex_ai/batches/transformation.py index f5f1ab2068a..a7dbb058465 100644 --- a/litellm/llms/vertex_ai/batches/transformation.py +++ b/litellm/llms/vertex_ai/batches/transformation.py @@ -1,7 +1,11 @@ -from collections.abc import Mapping -from typing import Any, Final +from collections.abc import Callable, Mapping +from dataclasses import dataclass +from typing import Any, Final, Protocol from urllib.parse import unquote +from pydantic import TypeAdapter, ValidationError + +from litellm._logging import verbose_logger from litellm._uuid import uuid from litellm.llms.vertex_ai.common_utils import ( VertexAIError, @@ -9,35 +13,128 @@ from litellm.llms.vertex_ai.common_utils import ( ) from litellm.types.llms.openai import BatchJobStatus, CreateBatchRequest from litellm.types.llms.vertex_ai import * -from litellm.types.utils import LiteLLMBatch, PromptTokensDetailsWrapper +from litellm.types.llms.vertex_ai import GenerateContentResponseBody +from litellm.types.utils import LiteLLMBatch, ModelInfo, Usage + +_NATIVE_VERTEX_RESPONSE: Final = TypeAdapter(GenerateContentResponseBody) -def vertex_prompt_tokens_details( - usage_metadata: Mapping[str, object], -) -> PromptTokensDetailsWrapper | None: - raw_details: Final = usage_metadata.get("promptTokensDetails") - if not isinstance(raw_details, list): - return None +def _int_field(mapping: Mapping[str, object], key: str) -> int: + value: Final = mapping.get(key) + if isinstance(value, int): + return value + return int(value) if isinstance(value, str) and value.isdigit() else 0 - def _normalize(detail: object) -> tuple[str, int] | None: - if not isinstance(detail, Mapping): + +def vertex_embedding_prompt_token_count(vertex_response: Mapping[str, object]) -> int: + """ + Prompt tokens billed for one Vertex Gemini Embedding batch row. + + Live rows report usage under `usageMetadata`; the documented `tokenCount` is kept as + a fallback. + """ + usage_metadata: Final = vertex_response.get("usageMetadata") + if isinstance(usage_metadata, Mapping): + return _int_field(usage_metadata, "promptTokenCount") + return _int_field(vertex_response, "tokenCount") + + +def is_vertex_embedding_batch_output_response(response_body: Mapping[str, object]) -> bool: + return isinstance(response_body.get("embedding"), dict) + + +def is_native_vertex_batch_output_row(row: Mapping[str, object]) -> bool: + return isinstance(row.get("request"), dict) + + +class NativeVertexBatchCostCalculator(Protocol): + def __call__( + self, + usage: Usage, + model: str, + custom_llm_provider: str | None = None, + model_info: ModelInfo | None = None, + ) -> tuple[float, float]: ... + + +@dataclass(frozen=True, slots=True) +class NativeVertexBatchRowStats: + usage: Usage + total_tokens: int + model: str | None + prompt_cost: float + completion_cost: float + + +def _native_vertex_row_usage( + response_body: Mapping[str, object], + calculate_usage: Callable[[GenerateContentResponseBody], Usage], +) -> Usage | None: + if "usageMetadata" not in response_body: + if not is_vertex_embedding_batch_output_response(response_body): return None - modality: Final = detail.get("modality") - token_count: Final = detail.get("tokenCount") - if not isinstance(modality, str) or not isinstance(token_count, int): - return None - return modality.upper(), token_count - - parsed_details: Final = tuple(_normalize(detail) for detail in raw_details) - normalized: Final = tuple(detail for detail in parsed_details if detail is not None) - if len(normalized) != len(parsed_details): + prompt_tokens: Final = vertex_embedding_prompt_token_count(response_body) + return Usage(prompt_tokens=prompt_tokens, completion_tokens=0, total_tokens=prompt_tokens) + try: + completion_response: Final = _NATIVE_VERTEX_RESPONSE.validate_python(response_body) + except ValidationError as e: + verbose_logger.debug("vertex_ai batch row response is not a GenerateContentResponse: %s", str(e)) return None + return calculate_usage(completion_response) - return PromptTokensDetailsWrapper( - text_tokens=sum(token_count for modality, token_count in normalized if modality in ("TEXT", "DOCUMENT")), - audio_tokens=sum(token_count for modality, token_count in normalized if modality == "AUDIO"), - image_tokens=sum(token_count for modality, token_count in normalized if modality == "IMAGE"), - video_tokens=sum(token_count for modality, token_count in normalized if modality == "VIDEO"), + +def native_vertex_batch_row_stats( + row: Mapping[str, object], + model_name: str | None, + *, + model_info: ModelInfo | None, + calculate_usage: Callable[[GenerateContentResponseBody], Usage], + cost_calculator: NativeVertexBatchCostCalculator, +) -> NativeVertexBatchRowStats | None: + """ + Usage and cost of one native Vertex predictions.jsonl row, a + `{"request": ..., "response": {"candidates": [...], "usageMetadata": {...}, "modelVersion": ...}}` + generateContent object or a `{"request": ..., "response": {"embedding": {...}, "usageMetadata": {...}}}` + embedding object (an embedding row without `usageMetadata` is billed from its documented `tokenCount`). + `model_name` (the deployment model) prices the row unless it is a wildcard, else its own `modelVersion` + does, else the wildcard name so explicit deployment prices still apply; a row without a response, a + generateContent row without `response.usageMetadata`, and a row whose response fails validation are + None (failed). + """ + response_body: Final = row.get("response") + if not isinstance(response_body, dict): + return None + usage: Final = _native_vertex_row_usage(response_body, calculate_usage) + if usage is None: + return None + total_tokens: Final = usage.total_tokens or (usage.prompt_tokens + usage.completion_tokens) + model_version: Final = response_body.get("modelVersion") + deployment_model: Final = model_name if model_name and "*" not in model_name else None + model: Final = deployment_model or (model_version if isinstance(model_version, str) else model_name) + if model is None: + verbose_logger.warning( + "vertex_ai batch output row could not be costed, so it is billed at $0 and the rest of the batch " + "is still billed: the row has no modelVersion and the batch has no deployment model" + ) + return NativeVertexBatchRowStats( + usage=usage, total_tokens=total_tokens, model=None, prompt_cost=0.0, completion_cost=0.0 + ) + try: + prompt_cost, completion_cost = cost_calculator( + usage=usage, model=model, custom_llm_provider="vertex_ai", model_info=model_info + ) + except Exception as e: # noqa: BLE001 # one unpriceable row must not abort the batch's cost accounting + verbose_logger.warning( + "vertex_ai batch output row could not be costed, so it is billed at $0 and the rest of the batch " + "is still billed. model=%s error=%s", + model, + str(e), + ) + return NativeVertexBatchRowStats( + usage=usage, total_tokens=total_tokens, model=model, prompt_cost=0.0, completion_cost=0.0 + ) + return NativeVertexBatchRowStats( + usage=usage, total_tokens=total_tokens, model=model, prompt_cost=prompt_cost, completion_cost=completion_cost ) diff --git a/litellm/llms/vertex_ai/files/transformation.py b/litellm/llms/vertex_ai/files/transformation.py index 80d32289c94..789b36ef3d0 100644 --- a/litellm/llms/vertex_ai/files/transformation.py +++ b/litellm/llms/vertex_ai/files/transformation.py @@ -9,7 +9,7 @@ from collections.abc import AsyncGenerator, Callable, Iterable, Iterator, Mappin from contextlib import aclosing from dataclasses import dataclass from types import MappingProxyType -from typing import Any, Final, TypedDict +from typing import IO, Any, Final, TypedDict from urllib.parse import quote, unquote import httpx @@ -41,6 +41,7 @@ from litellm.llms.base_llm.files.transformation import ( BaseFileUploadStream, LiteLLMLoggingObj, ) +from litellm.llms.vertex_ai.batches.transformation import vertex_embedding_prompt_token_count from litellm.llms.vertex_ai.common_utils import ( _convert_vertex_datetime_to_openai_datetime, get_vertex_ai_fine_tuned_endpoint_id, @@ -56,6 +57,7 @@ from litellm.types.files import StreamingMediaUploadConfig from litellm.types.llms.openai import ( AllMessageValues, CreateFileRequest, + FileContent, FileTypes, HttpxBinaryResponseContent, OpenAICreateFileRequestOptionalParams, @@ -87,6 +89,8 @@ _EMBED_REQUEST_FIELD_BY_GEMINI_PARAM: Final = ( _VERTEX_BATCH_FANNED_OUT_KEY_PATTERN: Final = re.compile(r"(?P[^#]*)#(?P\d+)/(?P\d+)") _JSONL_NEWLINE: Final = b"\n" _BATCH_OUTPUT_FIRST_ROW_PEEK_LIMIT_BYTES: Final = 32 * 1024 * 1024 +_PASSTHROUGH_MANAGED_GCS_PREFIX: Final = f"{VERTEX_AI_MANAGED_GCS_PREFIX}passthrough/" +_RAW_UPLOAD_CHUNK_BYTES: Final = 1024 * 1024 class _GcsObjectMetadataJson(TypedDict, total=False): @@ -418,19 +422,6 @@ def _split_vertex_batch_key(vertex_output_row: Mapping[str, object]) -> tuple[st return unquote(match["custom_id"]), int(match["index"]), int(match["total"]) -def _embedding_prompt_token_count(vertex_response: _VertexEmbeddingResponse) -> int: - """ - Prompt tokens billed for one Vertex Gemini Embedding batch row. - - Live rows report usage under `usageMetadata`; the documented `tokenCount` is kept as - a fallback. - """ - usage_metadata = vertex_response.get("usageMetadata") - if isinstance(usage_metadata, Mapping): - return int(usage_metadata.get("promptTokenCount") or 0) - return int(vertex_response.get("tokenCount") or 0) - - def _vertex_embeddings_rows_to_openai_batch_output_row( custom_id: str, vertex_output_rows: tuple[_VertexEmbeddingBatchRow, ...], @@ -471,7 +462,7 @@ def _vertex_embeddings_rows_to_openai_batch_output_row( ) responses = tuple(row["response"] for row in vertex_output_rows) - token_count = sum(_embedding_prompt_token_count(response) for response in responses) + token_count = sum(vertex_embedding_prompt_token_count(response) for response in responses) body = EmbeddingResponse( model=model or "", data=[ @@ -528,6 +519,16 @@ def _model_from_managed_gcs_url(url: str) -> str | None: return match.group(1) if match else None +def is_passthrough_managed_gcs_url(url: str) -> bool: + decoded_url: Final = unquote(url) + managed_prefix_start: Final = decoded_url.find(VERTEX_AI_MANAGED_GCS_PREFIX) + return managed_prefix_start >= 0 and decoded_url.startswith(_PASSTHROUGH_MANAGED_GCS_PREFIX, managed_prefix_start) + + +def is_passthrough_batch_upload(create_file_data: Mapping[str, object], litellm_params: Mapping[str, object]) -> bool: + return create_file_data.get("purpose") == "batch" and litellm_params.get("passthrough") is True + + def _is_embeddings_batch_entry(openai_entry: Mapping[str, object]) -> bool: """ Whether an OpenAI batch JSONL line targets the embeddings endpoint. @@ -791,6 +792,58 @@ class _OpenAIToVertexBatchUploadStream(BaseFileUploadStream): return self._iter_vertex_jsonl_chunks() +def _read_chunk_as_bytes(handle: IO[bytes]) -> bytes: + chunk: Final[bytes | str] = handle.read(_RAW_UPLOAD_CHUNK_BYTES) + return chunk.encode("utf-8") if isinstance(chunk, str) else bytes(chunk) + + +def _iter_raw_file_chunks(file_content: FileTypes) -> Iterator[bytes]: + content: Final[FileContent | str] = file_content[1] if isinstance(file_content, tuple) else file_content + if isinstance(content, (bytes, bytearray)): + yield from ( + bytes(content[offset : offset + _RAW_UPLOAD_CHUNK_BYTES]) + for offset in range(0, len(content), _RAW_UPLOAD_CHUNK_BYTES) + ) + return + if isinstance(content, str): + yield content.encode("utf-8") + return + if isinstance(content, PathLike): + with open(str(content), "rb") as handle: + yield from iter(lambda: handle.read(_RAW_UPLOAD_CHUNK_BYTES), b"") + return + if not hasattr(content, "read"): + raise ValueError("Unsupported file content type") + seek: Final = getattr(content, "seek", None) + if seek is None: + raise ValueError( + "Batch upload file handle must be seekable; got a non-seekable " + "stream. Pass bytes, a path, or a seekable handle." + ) + seek(0) + yield from iter(lambda: _read_chunk_as_bytes(content), b"") + + +class _RawFileUploadStream(BaseFileUploadStream): + def __init__(self, file_content: FileTypes) -> None: + self._file_content = file_content + + def iter_bytes(self) -> Iterator[bytes]: + return _iter_raw_file_chunks(self._file_content) + + +def _managed_batch_object_name(raw_model: str, *, passthrough: bool) -> str: + endpoint_id: Final = get_vertex_ai_fine_tuned_endpoint_id(raw_model) + model_path: Final = ( + f"endpoints/{endpoint_id}" + if endpoint_id is not None + else (raw_model if "publishers/google/models" in raw_model else f"publishers/google/models/{raw_model}") + ) + safe_model_path: Final = sanitize_cloud_object_path(model_path, fallback="model") + prefix: Final = _PASSTHROUGH_MANAGED_GCS_PREFIX if passthrough else VERTEX_AI_MANAGED_GCS_PREFIX + return f"{prefix}{safe_model_path}/{uuid.uuid4()}" + + class VertexAIFilesConfig(VertexBase, BaseFilesConfig): """ Config for VertexAI Files @@ -848,23 +901,34 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): if deployment_model else openai_jsonl_content[0].get("body", {}).get("model", "") ) - endpoint_id: Final = get_vertex_ai_fine_tuned_endpoint_id(raw_model) - model_path: Final = ( - f"endpoints/{endpoint_id}" - if endpoint_id is not None - else (raw_model if "publishers/google/models" in raw_model else f"publishers/google/models/{raw_model}") - ) - safe_model_path: Final = sanitize_cloud_object_path(model_path, fallback="model") - object_name: Final = f"{VERTEX_AI_MANAGED_GCS_PREFIX}{safe_model_path}/{uuid.uuid4()}" - return object_name + return _managed_batch_object_name(raw_model, passthrough=False) - def get_object_name(self, file_data: FileTypes, purpose: str, deployment_model: str | None = None) -> str: + def _get_passthrough_gcs_object_name(self, deployment_model: str | None) -> str: + if not deployment_model: + raise VertexAIError( + status_code=400, + message=( + "Native Vertex batch passthrough uploads need the deployment model to name the GCS object, " + "since native rows carry no model: pass `target_model_names` (proxy) or `model` (SDK)." + ), + ) + return _managed_batch_object_name(deployment_model.removeprefix("vertex_ai/"), passthrough=True) + + def get_object_name( + self, + file_data: FileTypes, + purpose: str, + deployment_model: str | None = None, + passthrough: bool = False, + ) -> str: """ Get the object name for the request. Reads only the first JSONL entry (streamed) for batch files, so a large upload is never materialized just to derive the GCS object name. """ + if purpose == "batch" and passthrough: + return self._get_passthrough_gcs_object_name(deployment_model) if purpose == "batch": ## 1. If jsonl, derive the object name from the deployment model (or the first entry's) first_entry: Final = next(_iter_openai_jsonl_entries(file_data), None) @@ -922,6 +986,7 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): file_data, purpose, deployment_model=configured_model if isinstance(configured_model, str) else None, + passthrough=is_passthrough_batch_upload(data, litellm_params), ) if object_prefix: object_name = f"{object_prefix}/{object_name}" @@ -984,6 +1049,14 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): if file_data is None: raise ValueError("file is required") + if is_passthrough_batch_upload(create_file_data, litellm_params): + return { + "streaming_media_upload": StreamingMediaUploadConfig( + body_stream=_RawFileUploadStream(file_data), + content_type="application/json", + ) + } + _, content_type = extract_file_metadata(file_data) if FilesAPIUtils.is_batch_jsonl_request( create_file_data=create_file_data, @@ -1164,6 +1237,8 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): # transformation, e.g. if they consume raw `predictions.jsonl` directly. if getattr(litellm, "disable_vertex_batch_output_transformation", False): return HttpxBinaryResponseContent(response=raw_response) + if is_passthrough_managed_gcs_url(str(raw_response.request.url)): + return HttpxBinaryResponseContent(response=raw_response) # Try to transform batch output if it's a JSONL file content: Final = raw_response.content @@ -1209,7 +1284,7 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): Everything else is passed through unchanged, including a row that fails to transform mid-stream. """ - if litellm.disable_vertex_batch_output_transformation: + if litellm.disable_vertex_batch_output_transformation or is_passthrough_managed_gcs_url(request_url): return FileContentStreamingResult(stream_iterator=stream_iterator, headers=headers) first_line, buffered = await _peek_first_jsonl_line( diff --git a/litellm/proxy/openai_files_endpoints/batch_file_validation.py b/litellm/proxy/openai_files_endpoints/batch_file_validation.py index a41bd36d510..fcd3be56ae7 100644 --- a/litellm/proxy/openai_files_endpoints/batch_file_validation.py +++ b/litellm/proxy/openai_files_endpoints/batch_file_validation.py @@ -8,10 +8,25 @@ from typing_extensions import assert_never from litellm.proxy._types import ProxyException -BATCH_LINE_REQUIRED_KEYS: Final = ("custom_id", "method", "url", "body") _MB: Final = 1024 * 1024 +@dataclass(frozen=True, slots=True) +class BatchLineShape: + required_keys: tuple[str, ...] + hint: str + + +BATCH_LINE_SHAPE: Final = BatchLineShape( + required_keys=("custom_id", "method", "url", "body"), + hint="Each line must be a JSON object with keys custom_id, method, url, body", +) +PASSTHROUGH_BATCH_LINE_SHAPE: Final = BatchLineShape( + required_keys=("request",), + hint="A passthrough upload takes native Vertex batch rows, so each line must be a JSON object with a request key", +) + + @dataclass(frozen=True, slots=True) class BatchFileTooLarge: size_bytes: int @@ -42,6 +57,7 @@ class BatchFileLineNotObject: class BatchFileMissingLineKey: line_number: int key: str + line_shape: BatchLineShape = BATCH_LINE_SHAPE BatchFileValidationFailure = ( @@ -70,20 +86,20 @@ def _iter_lines(file_source: bytes | BinaryIO) -> Iterator[bytes]: return iter(file_source) -def _check_line(line_number: int, raw_line: bytes) -> BatchFileValidationFailure | None: +def _check_line(line_number: int, raw_line: bytes, line_shape: BatchLineShape) -> BatchFileValidationFailure | None: try: parsed: Final = json.loads(raw_line) except (json.JSONDecodeError, UnicodeDecodeError): return BatchFileInvalidJsonLine(line_number=line_number) if not isinstance(parsed, dict): return BatchFileLineNotObject(line_number=line_number) - missing: Final = next((key for key in BATCH_LINE_REQUIRED_KEYS if key not in parsed), None) + missing: Final = next((key for key in line_shape.required_keys if key not in parsed), None) if missing is None: return None - return BatchFileMissingLineKey(line_number=line_number, key=missing) + return BatchFileMissingLineKey(line_number=line_number, key=missing, line_shape=line_shape) -def _scan_lines(file_source: bytes | BinaryIO) -> BatchFileValidationFailure | None: +def _scan_lines(file_source: bytes | BinaryIO, line_shape: BatchLineShape) -> BatchFileValidationFailure | None: content_lines: Final = ( (line_number, raw_line) for line_number, raw_line in enumerate(_iter_lines(file_source), start=1) @@ -96,7 +112,7 @@ def _scan_lines(file_source: bytes | BinaryIO) -> BatchFileValidationFailure | N ( failure for line_number, raw_line in chain((first_line,), content_lines) - for failure in (_check_line(line_number, raw_line),) + for failure in (_check_line(line_number, raw_line, line_shape),) if failure is not None ), None, @@ -107,6 +123,7 @@ def check_batch_file_upload( filename: str | None, file_source: bytes | BinaryIO, max_batch_file_size_mb: int | None, + line_shape: BatchLineShape = BATCH_LINE_SHAPE, ) -> BatchFileValidationFailure | None: if filename is None or not filename.lower().endswith(".jsonl"): return BatchFileWrongExtension(filename=filename or "") @@ -114,7 +131,7 @@ def check_batch_file_upload( size_bytes: Final = _file_size_bytes(file_source) if size_bytes > max_batch_file_size_mb * _MB: return BatchFileTooLarge(size_bytes=size_bytes, limit_mb=max_batch_file_size_mb) - scan_failure: Final = _scan_lines(file_source) + scan_failure: Final = _scan_lines(file_source, line_shape) if not isinstance(file_source, bytes): file_source.seek(0) return scan_failure @@ -169,11 +186,11 @@ def raise_batch_file_validation_failure(failure: BatchFileValidationFailure) -> param="file", code=400, ) - case BatchFileMissingLineKey(line_number=line_number, key=key): + case BatchFileMissingLineKey(line_number=line_number, key=key, line_shape=line_shape): raise ProxyException( message=( f"Missing required parameter: '{key}' (batch input file line {line_number}). " - f"Each line must be a JSON object with keys {', '.join(BATCH_LINE_REQUIRED_KEYS)}. " + f"{line_shape.hint}. " "The file was not forwarded to the provider." ), type="invalid_request_error", diff --git a/litellm/proxy/openai_files_endpoints/files_endpoints.py b/litellm/proxy/openai_files_endpoints/files_endpoints.py index ea2cad558c2..f5e62da5962 100644 --- a/litellm/proxy/openai_files_endpoints/files_endpoints.py +++ b/litellm/proxy/openai_files_endpoints/files_endpoints.py @@ -57,6 +57,8 @@ from litellm.proxy.common_utils.openai_error_payload import ( openai_error_type, ) from litellm.proxy.openai_files_endpoints.batch_file_validation import ( + BATCH_LINE_SHAPE, + PASSTHROUGH_BATCH_LINE_SHAPE, check_batch_file_upload, raise_batch_file_validation_failure, ) @@ -207,10 +209,91 @@ def get_files_provider_config( return None +def _deployment_provider(llm_router: Router, model_id: str, team_id: str | None) -> str | None: + credentials: Final = llm_router.get_deployment_credentials_with_provider(model_id=model_id, team_id=team_id) + return None if credentials is None else credentials.get("custom_llm_provider") + + +def _resolves_to_vertex_deployments_only(llm_router: Router | None, model_name: str, team_id: str | None) -> bool: + if llm_router is None or _deployment_provider(llm_router, model_name, team_id) != "vertex_ai": + return False + return all( + _deployment_provider(llm_router, str(deployment["model_info"]["id"]), team_id) == "vertex_ai" + for deployment in llm_router.get_model_list(model_name=model_name, team_id=team_id) or () + if "id" in deployment.get("model_info", {}) + ) + + +def _validate_passthrough_upload( + *, + purpose: str, + target_model_names: Sequence[str], + model: str | None, + target_storage: str | None, + llm_router: Router | None, + team_id: str | None, +) -> None: + if purpose != "batch": + raise ProxyException( + message=( + "`passthrough` uploads the file bytes unchanged for a native Vertex batch, " + f"so purpose must be 'batch', got '{purpose}'." + ), + type="invalid_request_error", + param="passthrough", + code=400, + ) + if target_storage and target_storage != "default": + raise ProxyException( + message=( + "`passthrough` writes the native batch file to the Vertex AI deployment's GCS bucket, " + f"so it cannot be combined with target_storage='{target_storage}'." + ), + type="invalid_request_error", + param="target_storage", + code=400, + ) + named_deployments: Final = ( + *(("target_model_names", name) for name in target_model_names), + *((("model", model),) if model else ()), + ) + if not named_deployments: + raise ProxyException( + message=( + "`passthrough` needs the Vertex AI deployment that will run the batch, " + "since native rows carry no model: pass `target_model_names` or `model`." + ), + type="invalid_request_error", + param="target_model_names", + code=400, + ) + offending: Final = next( + ( + (param, name) + for param, name in named_deployments + if not _resolves_to_vertex_deployments_only(llm_router, name, team_id) + ), + None, + ) + if offending is None: + return + param, name = offending + raise ProxyException( + message=( + f"`passthrough` is only supported for Vertex AI deployments; '{name}' does not resolve " + "to vertex_ai deployments only." + ), + type="invalid_request_error", + param=param, + code=400, + ) + + async def _scan_batch_upload( *, file_source: bytes | BinaryIO, purpose: str, + passthrough: bool, request_metadata: Mapping[str, object], user_api_key_dict: UserAPIKeyAuth, proxy_logging_obj: ProxyLogging, @@ -222,6 +305,17 @@ async def _scan_batch_upload( or not proxy_logging_obj.has_pre_call_guardrails(request_metadata) ): return None + if passthrough: + raise ProxyException( + message=( + "Batch guardrails cannot scan native Vertex batch rows, so a `passthrough` upload is refused " + "when the key, team, or request has pre-call guardrails configured. " + "The file was not forwarded to the provider." + ), + type="invalid_request_error", + param="passthrough", + code=400, + ) outcome: Final = await scan_batch_input_file( file_source=file_source, request_metadata=request_metadata, @@ -458,6 +552,7 @@ async def create_file( custom_llm_provider: str = Form(default="openai"), file: UploadFile = File(...), litellm_metadata: str | None = Form(default=None), + passthrough: bool = Form(default=False), user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ): """ @@ -560,17 +655,28 @@ async def create_file( if blocked_extension_failure is not None: raise_upload_validation_failure(blocked_extension_failure) + if passthrough: + _validate_passthrough_upload( + purpose=purpose, + target_model_names=target_model_names_list, + model=model_param, + target_storage=target_storage, + llm_router=llm_router, + team_id=user_api_key_dict.team_id, + ) + if purpose == "batch": batch_file_failure: Final = await asyncio.to_thread( check_batch_file_upload, file.filename, file_source, _MAX_BATCH_FILE_SIZE_MB_ADAPTER.validate_python(general_settings.get("max_batch_file_size_mb")), + PASSTHROUGH_BATCH_LINE_SHAPE if passthrough else BATCH_LINE_SHAPE, ) if batch_file_failure is not None: raise_batch_file_validation_failure(batch_file_failure) - data = {} + data = {"passthrough": True} if passthrough else {} # Parse expires_after if provided expires_after: FileExpiresAfter | None = None @@ -673,6 +779,7 @@ async def create_file( scan_result: Final = await _scan_batch_upload( file_source=file_source, purpose=purpose, + passthrough=passthrough, request_metadata=request_metadata, user_api_key_dict=user_api_key_dict, proxy_logging_obj=proxy_logging_obj, diff --git a/litellm/router.py b/litellm/router.py index d328fbbb12f..8960cd92cd8 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -5980,6 +5980,7 @@ class Router: replace_model_in_jsonl_bool: Final = should_replace_model_in_jsonl( purpose=purpose, + passthrough=kwargs.get("passthrough") is True, ) if replace_model_in_jsonl_bool: file = replace_model_in_jsonl( diff --git a/litellm/router_utils/batch_utils.py b/litellm/router_utils/batch_utils.py index be20c358202..386a2135239 100644 --- a/litellm/router_utils/batch_utils.py +++ b/litellm/router_utils/batch_utils.py @@ -62,15 +62,15 @@ def parse_jsonl_with_embedded_newlines(content: str) -> list[dict]: def should_replace_model_in_jsonl( purpose: OpenAIFilesPurpose, + passthrough: bool = False, ) -> bool: """ Check if the model name should be replaced in the JSONL file for the deployment model name. Azure raises an error on create batch if the model name for deployment is not in the .jsonl. + A passthrough upload keeps the caller's bytes untouched, so its rows are never rewritten. """ - if purpose == "batch": - return True - return False + return purpose == "batch" and not passthrough def replace_model_in_jsonl(file_content: FileTypes, new_model_name: str) -> FileTypes: diff --git a/tests/e2e/batches/test_batches_e2e.py b/tests/e2e/batches/test_batches_e2e.py index 9bb6d05bec8..6e9cf45e787 100644 --- a/tests/e2e/batches/test_batches_e2e.py +++ b/tests/e2e/batches/test_batches_e2e.py @@ -61,6 +61,7 @@ from e2e_http import ( StreamingResponse, Success, UnknownApiError, + proxy_error, require_successful_call, unwrap, ) @@ -1752,3 +1753,114 @@ class TestBatchTerminalState: assert (cost_row.total_tokens or 0) > 0, ( f"batch cost row has no token usage: {cost_row.total_tokens!r}" ) + + +NATIVE_VERTEX_BATCH_ROWS: Final = b"".join( + json.dumps( + { + "request": { + "contents": [{"role": "user", "parts": [{"text": text}]}], + "tools": [{"googleSearch": {"excludeDomains": ["example.com"]}}], + } + } + ).encode() + + b"\n" + for text in ("What is the tallest building in the world?", "Who won the last FIFA World Cup?") +) +VERTEX_BATCH_PROVIDER: Final = next(p for p in PROVIDERS if p.name == "vertex_ai") + + +class TestVertexNativePassthrough: + """`passthrough=true` on POST /v1/files uploads native Vertex batch JSONL byte for + byte (no OpenAI-to-Vertex translation, so `googleSearch` tools and the grounding + metadata they produce survive), and a batch created from that file is accepted. + + Terminal-state assertions (native output rows with groundingMetadata, the spend + row) are deliberately not here: retrieving a non-terminal batch books a $0 spend + row that blocks the real-cost row, the same reason TestBatchTerminalState polls + the list endpoint only. Those are proven by the PR's live curl proof instead. + """ + + @pytest.mark.covers( + "llm.files.vertex.native_passthrough.nonstream.works", + "llm.batches.vertex.native_passthrough.nonstream.works", + exercised_on=["files", "batches"], + ) + def test_native_jsonl_round_trips_untouched_and_starts_a_batch( + self, client: BatchClient, resources: ResourceManager, batch_deployments: None + ) -> None: + key = resources.key() + file = unwrap( + client.upload_file( + content=NATIVE_VERTEX_BATCH_ROWS, + form=FileUploadForm( + purpose="batch", target_model_names=VERTEX_BATCH_PROVIDER.model, passthrough=True + ), + key=key, + ) + ) + resources.defer(lambda: cleanup_file(client, file.id, key=key)) + assert_file_object(file, provider="vertex_ai") + assert is_managed_id(file.id), f"passthrough upload must return a managed file id, got {file.id!r}" + assert file.bytes == len(NATIVE_VERTEX_BATCH_ROWS), ( + f"passthrough upload must report the caller's byte count, got {file.bytes}" + ) + + downloaded = client.proxy.transport.download( + f"/v1/files/{file.id}/content", headers=client.proxy.transport.bearer(key) + ) + assert downloaded.status_code == 200, ( + f"file content must be 200, got {downloaded.status_code}: {downloaded.body[:300]}" + ) + assert downloaded.body.encode() == NATIVE_VERTEX_BATCH_ROWS, ( + "passthrough file content must be the uploaded native rows byte for byte" + ) + + created = client.create_batch(body=BatchCreateBody(input_file_id=file.id), key=key) + require_successful_call(created) + batch = BatchObject.model_validate_json(created.body) + resources.defer(lambda: cleanup_batch(client, batch.id, key=key, delete_output_files=True)) + assert is_managed_id(batch.id), f"passthrough batch must be LiteLLM-managed, got {batch.id!r}" + assert batch.status in CREATED_BATCH_STATUSES, f"passthrough batch has non-transitional status {batch.status!r}" + assert batch.input_file_id == file.id + + @pytest.mark.covers("llm.files.vertex.native_passthrough_validation.nonstream.works", exercised_on=["files"]) + @pytest.mark.parametrize( + "content, form, expected_param", + [ + pytest.param( + NATIVE_VERTEX_BATCH_ROWS, + FileUploadForm(purpose="batch", passthrough=True), + "target_model_names", + id="no-target-model", + ), + pytest.param( + NATIVE_VERTEX_BATCH_ROWS, + FileUploadForm(purpose="batch", target_model_names=OPENAI_BATCH_MODEL, passthrough=True), + "target_model_names", + id="non-vertex-target-model", + ), + pytest.param( + render_jsonl(VERTEX_BATCH_PROVIDER.raw_model), + FileUploadForm(purpose="batch", target_model_names=VERTEX_BATCH_PROVIDER.model, passthrough=True), + "request", + id="openai-shaped-rows", + ), + ], + ) + def test_passthrough_upload_is_rejected_outside_a_native_vertex_batch( + self, + content: bytes, + form: FileUploadForm, + expected_param: str, + client: BatchClient, + resources: ResourceManager, + batch_deployments: None, + ) -> None: + key = resources.key() + result = client.upload_file(content=content, form=form, key=key) + assert isinstance(result, UnknownApiError), f"expected a 400, got {result!r}" + assert result.status_code == 400, f"expected 400, got {result.status_code}: {result.body[:300]}" + error = proxy_error(result.body) + assert error.param == expected_param, f"unexpected error param in {error!r}" + assert "passthrough" in error.message diff --git a/tests/e2e/coverage_registry/llm_nonconversational.yaml b/tests/e2e/coverage_registry/llm_nonconversational.yaml index 3e389acc2a9..7d334ed41ff 100644 --- a/tests/e2e/coverage_registry/llm_nonconversational.yaml +++ b/tests/e2e/coverage_registry/llm_nonconversational.yaml @@ -21,6 +21,7 @@ - {id: llm.batches.openai_provider_fallback.basic.nonstream.works, module: llm, tier: P0, subject_endpoint: batches, route: openai, capability: basic, streaming: nonstream, assertions: [works], source: "batches/capabilities.py", rationale: "Provider-fallback raw-id scenario"} - {id: llm.batches.azure_openai.basic.nonstream.works, module: llm, tier: P0, subject_endpoint: batches, route: azure_openai, capability: basic, streaming: nonstream, assertions: [works], source: "batches/capabilities.py:98", rationale: "Azure batches all scenarios"} - {id: llm.batches.vertex.basic.nonstream.works, module: llm, tier: P0, subject_endpoint: batches, route: vertex, capability: basic, streaming: nonstream, assertions: [works], source: "batches/capabilities.py:98", rationale: "Vertex batches"} +- {id: llm.batches.vertex.native_passthrough.nonstream.works, module: llm, tier: P1, subject_endpoint: batches, route: vertex, capability: native_passthrough, streaming: nonstream, assertions: [works], source: "test_batches_e2e.py / LIT-4790", rationale: "A batch created from a passthrough-uploaded native Vertex JSONL file is accepted and starts on the deployment named at upload"} - {id: llm.batches.bedrock.basic.nonstream.works, module: llm, tier: P0, subject_endpoint: batches, route: bedrock_converse, capability: basic, streaming: nonstream, assertions: [works], source: "batches/capabilities.py:98", rationale: "Bedrock batches (encoded/unified only)"} - {id: llm.batches.bedrock.assume_role.nonstream.works, module: llm, tier: P0, subject_endpoint: batches, route: bedrock_converse, capability: assume_role, streaming: nonstream, assertions: [works], source: "test_batches_e2e.py", rationale: "Bedrock batch create under STS assume-role credentials"} - {id: llm.batches.bedrock.govcloud_partition.nonstream.works, module: llm, tier: P0, subject_endpoint: batches, route: bedrock_converse, capability: govcloud_partition, streaming: nonstream, assertions: [works], source: "test_batches_e2e.py", rationale: "Bedrock batch create in the us-gov-west-1 partition"} @@ -46,6 +47,8 @@ - {id: llm.files.openai.passthrough.nonstream.works, module: llm, tier: P1, subject_endpoint: files, route: openai, capability: basic, streaming: nonstream, assertions: [works], source: "test_passthrough_e2e.py", rationale: "POST/DELETE /openai_passthrough/v1/files relay OpenAI's own file object; the dedicated prefix must not bind as a provider name on the /{provider}/v1/files route (GitHub issue #36086)"} - {id: llm.files.azure_openai.upload.nonstream.works, module: llm, tier: P0, subject_endpoint: files, route: azure_openai, capability: basic, streaming: nonstream, assertions: [works], source: "batches/capabilities.py:45", rationale: "Azure file upload managed backend"} - {id: llm.files.vertex.upload.nonstream.works, module: llm, tier: P0, subject_endpoint: files, route: vertex, capability: basic, streaming: nonstream, assertions: [works], source: "batches/capabilities.py:52", rationale: "Vertex file upload to GCS"} +- {id: llm.files.vertex.native_passthrough.nonstream.works, module: llm, tier: P1, subject_endpoint: files, route: vertex, capability: native_passthrough, streaming: nonstream, assertions: [works], source: "test_batches_e2e.py / LIT-4790", rationale: "POST /v1/files with passthrough=true ships native Vertex batch JSONL (googleSearch tools and all) to GCS untouched and GET /v1/files/{id}/content returns the same bytes"} +- {id: llm.files.vertex.native_passthrough_validation.nonstream.works, module: llm, tier: P1, subject_endpoint: files, route: vertex, capability: input_validation, streaming: nonstream, assertions: [works], source: "test_batches_e2e.py / LIT-4790", rationale: "passthrough=true without a Vertex target_model_names, or with OpenAI-shaped rows, is a 400 naming the offending field and nothing is uploaded"} - {id: llm.files.bedrock.upload.nonstream.works, module: llm, tier: P0, subject_endpoint: files, route: bedrock_converse, capability: basic, streaming: nonstream, assertions: [works], source: "batches/capabilities.py:59", rationale: "Bedrock file upload to S3"} - {id: llm.files.bedrock.govcloud_partition.nonstream.works, module: llm, tier: P0, subject_endpoint: files, route: bedrock_converse, capability: govcloud_partition, streaming: nonstream, assertions: [works], source: "test_batches_e2e.py", rationale: "Bedrock file upload to an S3 bucket in the us-gov-west-1 partition"} - {id: llm.files.bedrock.split_s3_credentials.nonstream.works, module: llm, tier: P0, subject_endpoint: files, route: bedrock_converse, capability: split_s3_credentials, streaming: nonstream, assertions: [works], source: "test_batches_e2e.py / LIT-8297", rationale: "Bedrock file upload, content and delete sign S3 with s3_access_key_id / s3_secret_access_key when they differ from the aws_* identity"} diff --git a/tests/e2e/coverage_registry/schema.py b/tests/e2e/coverage_registry/schema.py index e009b02b69c..e5626144fad 100644 --- a/tests/e2e/coverage_registry/schema.py +++ b/tests/e2e/coverage_registry/schema.py @@ -73,6 +73,7 @@ LlmCapability = Literal[ "long_context_1m", "mid_conversation_system", "multi_turn", + "native_passthrough", "pdf_input", "prompt_cache_1h", "prompt_cache_5m", diff --git a/tests/e2e/e2e_http.py b/tests/e2e/e2e_http.py index 022caddd42a..e5d50d05c87 100644 --- a/tests/e2e/e2e_http.py +++ b/tests/e2e/e2e_http.py @@ -69,6 +69,7 @@ class FileUploadForm(BaseModel): purpose: str = "batch" target_model_names: str | None = None custom_llm_provider: str | None = None + passthrough: bool | None = None # ---------- Result types ---------- @@ -376,12 +377,18 @@ class ProxyErrorDetail(BaseModel): message: str type: str code: str + param: str | None = None class _ProxyErrorBody(BaseModel): error: ProxyErrorDetail +def proxy_error(body: str) -> ProxyErrorDetail: + """The proxy's own error envelope (`{"error": {message, type, param, code}}`) parsed off a rejected call.""" + return _ProxyErrorBody.model_validate_json(body).error + + def relayed_provider_rate_limit(outcome: RateLimitedError) -> ProxyErrorDetail | None: """The provider's own 429 as the proxy relayed it, or None when the 429 is the proxy's own.""" if PROVIDER_RATE_LIMIT_MARKER not in outcome.body: diff --git a/tests/router_unit_tests/test_router_batch_utils.py b/tests/router_unit_tests/test_router_batch_utils.py index e274ac61a01..4336185a07f 100644 --- a/tests/router_unit_tests/test_router_batch_utils.py +++ b/tests/router_unit_tests/test_router_batch_utils.py @@ -141,6 +141,7 @@ def test_should_replace_model_in_jsonl(): from litellm.router_utils.batch_utils import should_replace_model_in_jsonl assert should_replace_model_in_jsonl(purpose="batch") is True + assert should_replace_model_in_jsonl(purpose="batch", passthrough=True) is False assert should_replace_model_in_jsonl(purpose="test") is False assert should_replace_model_in_jsonl(purpose="user_data") is False diff --git a/tests/test_litellm/batches/test_batch_utils.py b/tests/test_litellm/batches/test_batch_utils.py new file mode 100644 index 00000000000..0b2bfe9d266 --- /dev/null +++ b/tests/test_litellm/batches/test_batch_utils.py @@ -0,0 +1,387 @@ +import json + +import pytest + +import litellm +import litellm.batches.batch_utils as bu +from litellm.types.llms.openai import Batch + +GROUNDED_USAGE_METADATA = { + "promptTokenCount": 19, + "candidatesTokenCount": 59, + "thoughtsTokenCount": 406, + "toolUsePromptTokenCount": 73, + "totalTokenCount": 557, + "promptTokensDetails": [{"modality": "TEXT", "tokenCount": 19}], + "candidatesTokensDetails": [{"modality": "TEXT", "tokenCount": 59}], + "toolUsePromptTokensDetails": [{"modality": "TEXT", "tokenCount": 73}], + "trafficType": "ON_DEMAND", +} +PASSTHROUGH_OUTPUT_URI = ( + "gs://litellm-bucket/litellm-vertex-files/passthrough/publishers/google/models/gemini-2.5-flash/u/" + "predictions.jsonl" +) +UNGROUNDED_USAGE_METADATA = { + "promptTokenCount": 20, + "candidatesTokenCount": 48, + "thoughtsTokenCount": 195, + "toolUsePromptTokenCount": 73, + "totalTokenCount": 336, + "promptTokensDetails": [{"modality": "TEXT", "tokenCount": 20}], + "trafficType": "ON_DEMAND", +} + + +def _batch(output_file_id: str) -> Batch: + return Batch( + id="b", + completion_window="24h", + created_at=1, + endpoint="/v1/chat/completions", + input_file_id="f", + object="batch", + status="completed", + output_file_id=output_file_id, + ) + + +def _vertex_jsonl(rows: list[dict]) -> bytes: + return "\n".join(json.dumps(row) for row in rows).encode() + + +def _vertex_openai_row(custom_id: str, model: str, prompt_tokens: int, completion_tokens: int) -> dict: + return { + "id": f"batch_req_{custom_id}", + "custom_id": custom_id, + "response": { + "status_code": 200, + "request_id": custom_id, + "body": { + "id": f"chatcmpl-{custom_id}", + "object": "chat.completion", + "model": model, + "choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}], + "usage": { + "prompt_tokens": prompt_tokens, + "completion_tokens": completion_tokens, + "total_tokens": prompt_tokens + completion_tokens, + }, + }, + }, + "error": None, + } + + +def _native_vertex_row(usage_metadata: dict, *, grounded: bool, model_version: str | None = "gemini-2.5-flash"): + candidate = {"content": {"role": "model", "parts": [{"text": "ok"}]}, "finishReason": "STOP"} + grounding = {"groundingMetadata": {"webSearchQueries": ["q"]}} if grounded else {} + response = {"candidates": [{**candidate, **grounding}], "usageMetadata": usage_metadata} + return { + "request": {"contents": [{"role": "user", "parts": [{"text": "q"}]}], "tools": [{"googleSearch": {}}]}, + "status": "", + "response": {**response, **({"modelVersion": model_version} if model_version else {})}, + "processed_time": "2026-09-23T19:02:00.000+00:00", + } + + +def _capture_cost_calls(monkeypatch, prompt_cost=0.5, completion_cost=0.25) -> list: + import litellm.cost_calculator as cc + + calls: list = [] + + def _calc(**kw): + calls.append(kw) + return (prompt_cost, completion_cost) + + monkeypatch.setattr(cc, "batch_cost_calculator", _calc) + return calls + + +def test_vertex_native_cost_bills_embedding_rows(monkeypatch): + monkeypatch.setitem(litellm.model_cost, "vertex_ai/gemini-embedding-2", {"input_cost_per_token_batches": 1e-7}) + rows = [ + { + "key": "id_1", + "status": "", + "request": {"content": {"parts": [{"text": "hello world"}]}}, + "response": {"embedding": {"values": [0.1, 0.2]}, "usageMetadata": {"promptTokenCount": 2}}, + }, + { + "key": "id_2", + "status": "", + "request": {"content": {"parts": [{"text": "hello"}]}}, + "response": {"embedding": {"values": [0.3]}, "tokenCount": "3"}, + }, + {"key": "id_3", "status": "INVALID_ARGUMENT", "request": {"content": {"parts": [{"text": ""}]}}}, + ] + + result = bu.calculate_vertex_ai_batch_cost_and_usage(rows, "gemini-embedding-2") + + assert (result.successful_requests, result.failed_requests) == (2, 1) + assert (result.usage.prompt_tokens, result.usage.completion_tokens, result.usage.total_tokens) == (5, 0, 5) + assert result.cost == pytest.approx(5 * 1e-7) + assert result.models == ["gemini-embedding-2"] + + +@pytest.mark.asyncio +async def test_native_vertex_rows_route_to_vertex_cost_path_without_flag(monkeypatch): + monkeypatch.setattr(litellm, "disable_vertex_batch_output_transformation", False, raising=False) + monkeypatch.setattr( + bu, "_aggregate_batch_cost_usage_models", lambda **kw: pytest.fail("generic path should not run") + ) + calls = _capture_cost_calls(monkeypatch) + rows = [ + _native_vertex_row(GROUNDED_USAGE_METADATA, grounded=True), + _native_vertex_row(UNGROUNDED_USAGE_METADATA, grounded=False), + ] + + result = await bu.calculate_batch_cost_and_usage( + file_content_dictionary=rows, custom_llm_provider="vertex_ai", model_name="gemini-2.5-flash" + ) + + assert result.cost == pytest.approx(1.5) + assert (result.successful_requests, result.failed_requests) == (2, 0) + assert result.models == ["gemini-2.5-flash"] + assert {(call["model"], call["custom_llm_provider"]) for call in calls} == {("gemini-2.5-flash", "vertex_ai")} + + +@pytest.mark.asyncio +async def test_openai_shaped_vertex_rows_keep_the_generic_path_without_flag(monkeypatch): + monkeypatch.setattr(litellm, "disable_vertex_batch_output_transformation", False, raising=False) + monkeypatch.setattr( + bu, "calculate_vertex_ai_batch_cost_and_usage", lambda *a, **kw: pytest.fail("native path should not run") + ) + _capture_cost_calls(monkeypatch) + rows = [_vertex_openai_row("request-1", "gemini-2.5-flash", 10, 5)] + + result = await bu.calculate_batch_cost_and_usage( + file_content_dictionary=rows, custom_llm_provider="vertex_ai", model_name="gemini-2.5-flash" + ) + + assert result.successful_requests == 1 + + +@pytest.mark.asyncio +async def test_native_vertex_rows_on_another_provider_keep_the_generic_path(monkeypatch): + monkeypatch.setattr( + bu, "calculate_vertex_ai_batch_cost_and_usage", lambda *a, **kw: pytest.fail("native path should not run") + ) + _capture_cost_calls(monkeypatch) + + result = await bu.calculate_batch_cost_and_usage( + file_content_dictionary=[_native_vertex_row(GROUNDED_USAGE_METADATA, grounded=True)], + custom_llm_provider="openai", + ) + + assert result.successful_requests == 0 + + +@pytest.mark.asyncio +async def test_handle_completed_batch_routes_native_rows_without_flag(monkeypatch): + monkeypatch.setattr(litellm, "disable_vertex_batch_output_transformation", False, raising=False) + raw_rows = [_native_vertex_row(GROUNDED_USAGE_METADATA, grounded=True)] + + async def fake_fetch(batch, custom_llm_provider, litellm_params=None): + return _vertex_jsonl(raw_rows) + + monkeypatch.setattr(bu, "_fetch_batch_output_file_content", fake_fetch) + monkeypatch.setattr( + bu, "_aggregate_batch_cost_usage_models", lambda **kw: pytest.fail("generic path should not run") + ) + calls = _capture_cost_calls(monkeypatch, prompt_cost=0.7, completion_cost=0.3) + deployment_model_info = {"input_cost_per_token_batches": 1e-6, "output_cost_per_token_batches": 2e-6} + + result = await bu._handle_completed_batch( + _batch(PASSTHROUGH_OUTPUT_URI), + custom_llm_provider="vertex_ai", + model_name="gemini-2.5-flash", + model_info=deployment_model_info, + ) + + assert result.cost == pytest.approx(1.0) + assert result.usage.total_tokens == 557 + assert [call["model_info"] for call in calls] == [deployment_model_info] + + +def test_native_vertex_usage_is_billed_like_the_online_path(monkeypatch): + calls = _capture_cost_calls(monkeypatch) + grounded = _native_vertex_row(GROUNDED_USAGE_METADATA, grounded=True) + ungrounded = _native_vertex_row(UNGROUNDED_USAGE_METADATA, grounded=False) + + result = bu.calculate_vertex_ai_batch_cost_and_usage([grounded, ungrounded], "gemini-2.5-flash") + + grounded_usage, ungrounded_usage = (call["usage"] for call in calls) + assert grounded_usage.prompt_tokens == 19 + assert grounded_usage.completion_tokens == 59 + 406 + assert grounded_usage.completion_tokens_details.reasoning_tokens == 406 + assert ungrounded_usage.prompt_tokens == 20 + 73 + assert ungrounded_usage.completion_tokens == 48 + 195 + assert (result.usage.prompt_tokens, result.usage.completion_tokens, result.usage.total_tokens) == ( + 19 + 93, + 465 + 243, + 557 + 336, + ) + + +def test_native_vertex_rows_are_priced_by_model_version_without_a_model_name(monkeypatch): + calls = _capture_cost_calls(monkeypatch) + rows = [ + _native_vertex_row(GROUNDED_USAGE_METADATA, grounded=True, model_version="gemini-2.5-flash"), + _native_vertex_row(UNGROUNDED_USAGE_METADATA, grounded=False, model_version="gemini-2.5-pro"), + _native_vertex_row(UNGROUNDED_USAGE_METADATA, grounded=False, model_version=None), + ] + + result = bu.calculate_vertex_ai_batch_cost_and_usage(rows) + + assert [call["model"] for call in calls] == ["gemini-2.5-flash", "gemini-2.5-pro"] + assert result.models == ["gemini-2.5-flash", "gemini-2.5-pro"] + assert result.cost == pytest.approx(1.5) + assert result.successful_requests == 3 + assert result.usage.total_tokens == 557 + 336 + 336 + + +def test_native_vertex_rows_without_usage_metadata_count_as_failed(monkeypatch): + _capture_cost_calls(monkeypatch) + rows = [ + {"request": {"contents": []}, "status": "Error: bad request", "processed_time": "t"}, + {"request": {"contents": []}, "response": {"candidates": []}}, + _native_vertex_row(GROUNDED_USAGE_METADATA, grounded=True), + ] + + result = bu.calculate_vertex_ai_batch_cost_and_usage(rows, "gemini-2.5-flash") + + assert (result.successful_requests, result.failed_requests) == (1, 2) + assert result.usage.total_tokens == 557 + + +def test_native_vertex_batch_whose_rows_all_failed_still_names_the_deployment_model(monkeypatch): + calls = _capture_cost_calls(monkeypatch) + rows = [{"request": {"contents": []}, "status": "Error: quota exceeded", "processed_time": "t"}] * 2 + + result = bu.calculate_vertex_ai_batch_cost_and_usage(rows, "gemini-2.5-flash") + + assert result.models == ["gemini-2.5-flash"] + assert (result.successful_requests, result.failed_requests, result.cost) == (0, 2, 0.0) + assert calls == [] + + +def test_native_vertex_rows_are_priced_with_the_deployment_model_info(monkeypatch): + calls = _capture_cost_calls(monkeypatch) + deployment_model_info = {"input_cost_per_token_batches": 1e-6, "output_cost_per_token_batches": 2e-6} + + bu.calculate_vertex_ai_batch_cost_and_usage( + [_native_vertex_row(GROUNDED_USAGE_METADATA, grounded=True)], + "gemini-2.5-flash", + model_info=deployment_model_info, + ) + + assert [call["model_info"] for call in calls] == [deployment_model_info] + + +@pytest.mark.asyncio +async def test_native_vertex_rows_keep_the_deployment_model_info_through_the_batch_entrypoint(monkeypatch): + calls = _capture_cost_calls(monkeypatch) + deployment_model_info = {"input_cost_per_token_batches": 1e-6} + + await bu.calculate_batch_cost_and_usage( + file_content_dictionary=[_native_vertex_row(GROUNDED_USAGE_METADATA, grounded=True)], + custom_llm_provider="vertex_ai", + model_name="gemini-2.5-flash", + model_info=deployment_model_info, + ) + + assert [call["model_info"] for call in calls] == [deployment_model_info] + + +def test_native_vertex_rows_are_priced_by_the_deployment_model_over_model_version(monkeypatch): + calls = _capture_cost_calls(monkeypatch) + rows = [_native_vertex_row(GROUNDED_USAGE_METADATA, grounded=True, model_version="gemini-2.5-pro")] + + result = bu.calculate_vertex_ai_batch_cost_and_usage(rows, "gemini-2.5-flash") + + assert [call["model"] for call in calls] == ["gemini-2.5-flash"] + assert result.models == ["gemini-2.5-flash"] + + +def test_native_vertex_rows_that_fail_response_validation_count_as_failed(monkeypatch): + calls = _capture_cost_calls(monkeypatch) + rows = [ + {"request": {"contents": []}, "response": {"candidates": "nope", "usageMetadata": GROUNDED_USAGE_METADATA}}, + _native_vertex_row(GROUNDED_USAGE_METADATA, grounded=True), + ] + + result = bu.calculate_vertex_ai_batch_cost_and_usage(rows, "gemini-2.5-flash") + + assert (result.successful_requests, result.failed_requests) == (1, 1) + assert result.usage.total_tokens == 557 + assert len(calls) == 1 + + +@pytest.mark.parametrize("wildcard_model", ["*", "vertex_ai/*"]) +def test_native_vertex_rows_under_a_wildcard_deployment_are_priced_by_model_version(monkeypatch, wildcard_model): + calls = _capture_cost_calls(monkeypatch) + rows = [ + _native_vertex_row(GROUNDED_USAGE_METADATA, grounded=True, model_version="gemini-2.5-flash"), + _native_vertex_row(UNGROUNDED_USAGE_METADATA, grounded=False, model_version=None), + ] + + result = bu.calculate_vertex_ai_batch_cost_and_usage(rows, wildcard_model) + + assert [call["model"] for call in calls] == ["gemini-2.5-flash", wildcard_model] + assert result.cost == pytest.approx(1.5) + assert (result.successful_requests, result.failed_requests) == (2, 0) + assert result.usage.total_tokens == 557 + 336 + + +def test_native_vertex_row_without_model_version_under_a_wildcard_deployment_bills_its_explicit_prices(): + deployment_model_info = {"input_cost_per_token_batches": 1e-6, "output_cost_per_token_batches": 2e-6} + with_version = _native_vertex_row(GROUNDED_USAGE_METADATA, grounded=True, model_version="gemini-2.5-flash") + without_version = _native_vertex_row(GROUNDED_USAGE_METADATA, grounded=True, model_version=None) + + twin = bu.calculate_vertex_ai_batch_cost_and_usage([with_version], "vertex_ai/*", model_info=deployment_model_info) + both = bu.calculate_vertex_ai_batch_cost_and_usage( + [with_version, without_version], "vertex_ai/*", model_info=deployment_model_info + ) + + assert twin.cost > 0 + assert both.cost == pytest.approx(2 * twin.cost) + assert (both.successful_requests, both.failed_requests) == (2, 0) + + +def test_native_vertex_row_the_cost_map_cannot_price_is_billed_at_zero_and_the_rest_still_bills(monkeypatch): + import litellm.cost_calculator as cc + + def _calc(**kw): + if kw["model"] == "gemini-unpriced": + raise ValueError("no pricing") + return (0.5, 0.25) + + monkeypatch.setattr(cc, "batch_cost_calculator", _calc) + rows = [ + _native_vertex_row(GROUNDED_USAGE_METADATA, grounded=True, model_version="gemini-unpriced"), + _native_vertex_row(UNGROUNDED_USAGE_METADATA, grounded=False, model_version="gemini-2.5-flash"), + ] + + result = bu.calculate_vertex_ai_batch_cost_and_usage(rows) + + assert result.cost == pytest.approx(0.75) + assert (result.successful_requests, result.failed_requests) == (2, 0) + assert result.usage.total_tokens == 557 + 336 + assert result.models == ["gemini-unpriced", "gemini-2.5-flash"] + + +@pytest.mark.asyncio +async def test_flag_sends_every_vertex_row_down_the_native_path_when_a_model_is_known(monkeypatch): + monkeypatch.setattr(litellm, "disable_vertex_batch_output_transformation", True, raising=False) + monkeypatch.setattr( + bu, "_aggregate_batch_cost_usage_models", lambda **kw: pytest.fail("generic path should not run") + ) + calls = _capture_cost_calls(monkeypatch) + rows = [_vertex_openai_row("request-1", "gemini-2.5-flash", 10, 5)] + + result = await bu.calculate_batch_cost_and_usage( + file_content_dictionary=rows, custom_llm_provider="vertex_ai", model_name="gemini-2.5-flash" + ) + + assert calls == [] + assert (result.successful_requests, result.failed_requests) == (0, 1) diff --git a/tests/test_litellm/files/__init__.py b/tests/test_litellm/files/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/files/test_main.py b/tests/test_litellm/files/test_main.py new file mode 100644 index 00000000000..2704bdfb6ff --- /dev/null +++ b/tests/test_litellm/files/test_main.py @@ -0,0 +1,71 @@ +from typing import Final +from urllib.parse import parse_qs, urlparse + +import httpx +import pytest + +import litellm +from litellm.llms.custom_httpx.http_handler import HTTPHandler + +NATIVE_VERTEX_ROWS: Final = ( + b'{"request": {"contents": [{"role": "user", "parts": [{"text": "Who won the 2024 Tour de France?"}]}],' + b' "tools": [{"googleSearch": {"excludeDomains": ["example.com"]}}]}}\n' + b'{"request": {"contents": [{"role": "user", "parts": [{"text": "What is the tallest building in Tokyo?"}]}],' + b' "tools": [{"googleSearch": {}}]}}\n' +) + + +@pytest.mark.parametrize( + "custom_llm_provider, purpose", + [("openai", "batch"), ("vertex_ai", "assistants")], + ids=["non-vertex-provider", "non-batch-purpose"], +) +def test_create_file_passthrough_is_rejected_outside_a_vertex_batch(custom_llm_provider, purpose): + with pytest.raises(litellm.BadRequestError) as exc_info: + litellm.create_file( + file=("batch.jsonl", b'{"request": {"contents": []}}\n', "application/jsonl"), + purpose=purpose, + custom_llm_provider=custom_llm_provider, + passthrough=True, + api_key="sk-test", + api_base="http://127.0.0.1:9", + ) + + assert "vertex_ai" in str(exc_info.value) + assert "batch" in str(exc_info.value) + + +def _gcs_upload_transport(uploads: list[httpx.Request]) -> httpx.MockTransport: + def respond(request: httpx.Request) -> httpx.Response: + uploads.append(request) + object_name: Final = parse_qs(urlparse(str(request.url)).query)["name"][0] + return httpx.Response( + 200, + json={ + "id": f"my-bucket/{object_name}/1758585600000000", + "name": object_name, + "size": str(len(request.read())), + "timeCreated": "2026-09-23T00:00:00.000Z", + }, + ) + + return httpx.MockTransport(respond) + + +def test_create_file_passthrough_kwarg_ships_native_rows_byte_for_byte_under_the_passthrough_prefix(): + uploads: Final[list[httpx.Request]] = [] + file_object = litellm.create_file( + file=("batch.jsonl", NATIVE_VERTEX_ROWS, "application/jsonl"), + purpose="batch", + custom_llm_provider="vertex_ai", + passthrough=True, + model="vertex_ai/gemini-2.5-flash", + gcs_bucket_name="my-bucket", + api_key="test-token", + client=HTTPHandler(client=httpx.Client(transport=_gcs_upload_transport(uploads))), + ) + (upload,) = uploads + object_name: Final = parse_qs(urlparse(str(upload.url)).query)["name"][0] + assert upload.read() == NATIVE_VERTEX_ROWS + assert object_name.startswith("litellm-vertex-files/passthrough/publishers/google/models/gemini-2.5-flash/") + assert file_object.id == f"gs://my-bucket/{object_name}" diff --git a/tests/test_litellm/llms/vertex_ai/batches/test_transformation.py b/tests/test_litellm/llms/vertex_ai/batches/test_transformation.py index e6126b02790..ae045edbec5 100644 --- a/tests/test_litellm/llms/vertex_ai/batches/test_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/batches/test_transformation.py @@ -19,7 +19,6 @@ import pytest from litellm.llms.vertex_ai.batches.transformation import ( # noqa: E402 VertexAIBatchTransformation, - vertex_prompt_tokens_details, ) from litellm.llms.vertex_ai.common_utils import ( # noqa: E402 VertexAIError, @@ -41,27 +40,6 @@ ENDPOINT_INPUT_FILE = ( ) -def test_vertex_prompt_tokens_details_rejects_malformed_details(): - assert vertex_prompt_tokens_details({"promptTokensDetails": [1]}) is None - assert vertex_prompt_tokens_details({"promptTokensDetails": [{"modality": "AUDIO"}]}) is None - assert ( - vertex_prompt_tokens_details( - { - "promptTokensDetails": [ - {"modality": "AUDIO", "tokenCount": 1}, - "malformed", - ] - } - ) - is None - ) - - -# =========================================================================== # -# transform_openai_batch_request_to_vertex_ai_batch_request -# =========================================================================== # - - def test_transform_openai_request_builds_full_vertex_job(): with patch( "litellm.llms.vertex_ai.batches.transformation.uuid.uuid4", @@ -477,3 +455,19 @@ def test_list_response_none_jobs_treated_as_empty(): out = T.transform_vertex_ai_batch_list_response_to_openai_list_response({"batchPredictionJobs": None}) assert out["data"] == [] assert out["first_id"] is None + + +PASSTHROUGH_INPUT_FILE = ( + "gs://litellm-testing-bucket/litellm-vertex-files/passthrough/publishers/google/models/gemini-2.5-flash/uuid-1" +) + + +def test_get_model_from_passthrough_gcs_file(): + assert T._get_model_from_gcs_file(PASSTHROUGH_INPUT_FILE) == "publishers/google/models/gemini-2.5-flash" + + +def test_get_gcs_uri_prefix_keeps_passthrough_segment_so_output_lands_beside_input(): + assert ( + T._get_gcs_uri_prefix_from_file(PASSTHROUGH_INPUT_FILE) + == "gs://litellm-testing-bucket/litellm-vertex-files/passthrough/publishers/google/models/gemini-2.5-flash" + ) diff --git a/tests/test_litellm/llms/vertex_ai/files/__init__.py b/tests/test_litellm/llms/vertex_ai/files/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/llms/vertex_ai/files/test_transformation.py b/tests/test_litellm/llms/vertex_ai/files/test_transformation.py new file mode 100644 index 00000000000..6958576b0c8 --- /dev/null +++ b/tests/test_litellm/llms/vertex_ai/files/test_transformation.py @@ -0,0 +1,310 @@ +import io +import json +import urllib.parse +from pathlib import Path +from unittest.mock import MagicMock +from urllib.parse import parse_qs, urlparse + +import httpx +import pytest + +from litellm.llms.vertex_ai.common_utils import VertexAIError +from litellm.llms.vertex_ai.files.transformation import VertexAIFilesConfig, is_passthrough_managed_gcs_url + +NATIVE_VERTEX_ROW = json.dumps( + { + "request": { + "contents": [{"role": "user", "parts": [{"text": "What is the tallest building in the world?"}]}], + "tools": [{"googleSearch": {"excludeDomains": ["example.com"]}}], + } + } +).encode() +NATIVE_VERTEX_JSONL = NATIVE_VERTEX_ROW + b"\n" + NATIVE_VERTEX_ROW + b"\n" +OPENAI_BATCH_JSONL = ( + b'{"custom_id": "r1", "method": "POST", "url": "/v1/chat/completions",' + b' "body": {"model": "gemini-2.5-flash", "messages": [{"role": "user", "content": "hi"}]}}\n' +) +PASSTHROUGH_OBJECT = ( + "litellm-vertex-files/passthrough/publishers/google/models/gemini-2.5-flash/uuid-1/predictions.jsonl" +) +TRANSFORMED_OBJECT = "litellm-vertex-files/publishers/google/models/gemini-2.5-flash/uuid-1/predictions.jsonl" +UPLOAD_CHUNK_BYTES = 1024 * 1024 + + +@pytest.fixture +def config() -> VertexAIFilesConfig: + return VertexAIFilesConfig() + + +def _gcs_media_url(object_name: str) -> str: + return ( + f"https://storage.googleapis.com/storage/v1/b/my-bucket/o/{urllib.parse.quote(object_name, safe='')}?alt=media" + ) + + +def _native_output_jsonl() -> bytes: + return ( + json.dumps( + { + "request": json.loads(NATIVE_VERTEX_ROW)["request"], + "status": "", + "response": { + "candidates": [ + { + "content": {"role": "model", "parts": [{"text": "The Burj Khalifa."}]}, + "finishReason": "STOP", + "groundingMetadata": {"webSearchQueries": ["tallest building in the world"]}, + } + ], + "modelVersion": "gemini-2.5-flash", + "usageMetadata": {"promptTokenCount": 20, "candidatesTokenCount": 48, "totalTokenCount": 68}, + }, + "processed_time": "2026-09-23T19:02:00.000+00:00", + } + ).encode() + + b"\n" + ) + + +def _upload_chunks(config: VertexAIFilesConfig, file: object, litellm_params: dict) -> list[bytes]: + body = config.transform_create_file_request( + model="", + create_file_data={"file": file, "purpose": "batch"}, + optional_params={}, + litellm_params=litellm_params, + ) + return list(body["streaming_media_upload"]["body_stream"].iter_bytes()) + + +def _upload_body_bytes(config: VertexAIFilesConfig, file: object, litellm_params: dict) -> bytes: + return b"".join(_upload_chunks(config, file, litellm_params)) + + +class TestPassthroughBatchUpload: + """`passthrough=True` on a batch upload ships the caller's native Vertex JSONL + to GCS byte for byte, filed under a `passthrough/` object path so the batch + output that lands beside it is recognized and returned untouched as well.""" + + def _upload_url(self, config, litellm_params, file, purpose="batch") -> str: + return config.get_complete_file_url( + api_base=None, + api_key=None, + model="", + optional_params={}, + litellm_params=litellm_params, + data={"file": file, "purpose": purpose}, + ) + + def test_passthrough_object_is_filed_under_passthrough_prefix_named_by_deployment_model(self, config): + url = self._upload_url( + config, + {"gcs_bucket_name": "my-bucket", "model": "vertex_ai/gemini-2.5-flash", "passthrough": True}, + ("batch.jsonl", NATIVE_VERTEX_JSONL, "application/jsonl"), + ) + object_name = parse_qs(urlparse(url).query)["name"][0] + assert object_name.startswith("litellm-vertex-files/passthrough/publishers/google/models/gemini-2.5-flash/") + + def test_passthrough_upload_without_deployment_model_is_rejected(self, config): + with pytest.raises(VertexAIError) as exc_info: + self._upload_url( + config, + {"gcs_bucket_name": "my-bucket", "passthrough": True}, + ("batch.jsonl", NATIVE_VERTEX_JSONL, "application/jsonl"), + ) + assert exc_info.value.status_code == 400 + assert "target_model_names" in exc_info.value.message + + def test_passthrough_flag_does_not_ship_a_non_batch_upload_raw(self, config): + result = config.transform_create_file_request( + model="", + create_file_data={"file": ("notes.txt", b"plain text", "text/plain"), "purpose": "user_data"}, + optional_params={}, + litellm_params={"gcs_bucket_name": "my-bucket", "passthrough": True}, + ) + assert result == b"plain text" + + def test_passthrough_flag_is_ignored_for_non_batch_purposes(self, config): + url = self._upload_url( + config, + {"gcs_bucket_name": "my-bucket", "model": "vertex_ai/gemini-2.5-flash", "passthrough": True}, + ("notes.txt", b"plain text", "text/plain"), + purpose="user_data", + ) + object_name = parse_qs(urlparse(url).query)["name"][0] + assert object_name.startswith("litellm-vertex-files/uploads/") + assert "passthrough" not in object_name + + @pytest.mark.parametrize( + "file", + [ + ("batch.jsonl", NATIVE_VERTEX_JSONL, "application/jsonl"), + NATIVE_VERTEX_JSONL, + ("batch.jsonl", io.BytesIO(NATIVE_VERTEX_JSONL), "application/jsonl"), + ("batch.jsonl", NATIVE_VERTEX_JSONL.decode(), "application/jsonl"), + ], + ids=["bytes-tuple", "bare-bytes", "handle-tuple", "text-tuple"], + ) + def test_passthrough_upload_body_is_the_callers_bytes(self, config, file): + body = config.transform_create_file_request( + model="", + create_file_data={"file": file, "purpose": "batch"}, + optional_params={}, + litellm_params={"passthrough": True}, + ) + stream = body["streaming_media_upload"]["body_stream"] + assert b"".join(stream.iter_bytes()) == NATIVE_VERTEX_JSONL + assert b"".join(stream.iter_bytes()) == NATIVE_VERTEX_JSONL + assert body["streaming_media_upload"]["content_type"] == "application/json" + + def test_passthrough_upload_streams_a_large_handle_in_bounded_chunks(self, config): + content = NATIVE_VERTEX_ROW * (3 * UPLOAD_CHUNK_BYTES // len(NATIVE_VERTEX_ROW) + 1) + chunks = _upload_chunks( + config, ("batch.jsonl", io.BytesIO(content), "application/jsonl"), {"passthrough": True} + ) + assert len(chunks) >= 3 + assert max(len(chunk) for chunk in chunks) <= UPLOAD_CHUNK_BYTES + assert b"".join(chunks) == content + + def test_passthrough_upload_streams_a_path_in_bounded_chunks(self, config, tmp_path: Path): + content = NATIVE_VERTEX_ROW * (2 * UPLOAD_CHUNK_BYTES // len(NATIVE_VERTEX_ROW) + 1) + batch_path = tmp_path / "batch.jsonl" + batch_path.write_bytes(content) + chunks = _upload_chunks(config, ("batch.jsonl", batch_path, "application/jsonl"), {"passthrough": True}) + assert len(chunks) >= 2 + assert max(len(chunk) for chunk in chunks) <= UPLOAD_CHUNK_BYTES + assert b"".join(chunks) == content + + def test_passthrough_upload_rejects_a_non_seekable_handle(self, config): + class _Pipe: + def read(self, size=-1): + return b"" + + with pytest.raises(ValueError, match="seekable"): + _upload_body_bytes(config, ("batch.jsonl", _Pipe(), "application/jsonl"), {"passthrough": True}) + + def test_passthrough_upload_rejects_content_that_is_neither_bytes_path_nor_handle(self, config): + with pytest.raises(ValueError, match="Unsupported file content type"): + _upload_body_bytes(config, ("batch.jsonl", 42, "application/jsonl"), {"passthrough": True}) + + def test_openai_rows_are_translated_unless_passthrough_is_set(self, config): + file = ("batch.jsonl", OPENAI_BATCH_JSONL, "application/jsonl") + translated = _upload_body_bytes(config, file, {}) + untouched = _upload_body_bytes(config, file, {"passthrough": True}) + assert untouched == OPENAI_BATCH_JSONL + assert translated != OPENAI_BATCH_JSONL + assert b'"contents"' in translated + + def test_passthrough_output_content_is_returned_untouched(self, config): + raw_jsonl = _native_output_jsonl() + + def _download(object_name: str) -> bytes: + raw_response = httpx.Response( + status_code=200, + content=raw_jsonl, + headers={"content-type": "application/octet-stream"}, + request=httpx.Request("GET", _gcs_media_url(object_name)), + ) + result = config.transform_file_content_response( + raw_response=raw_response, logging_obj=MagicMock(), litellm_params={} + ) + return result.response.content + + assert _download(PASSTHROUGH_OBJECT) == raw_jsonl + assert _download(f"team-a/{PASSTHROUGH_OBJECT}") == raw_jsonl + transformed = _download(TRANSFORMED_OBJECT) + assert transformed != raw_jsonl + assert json.loads(transformed.splitlines()[0])["response"]["body"]["choices"] + nested = _download(f"litellm-vertex-files/{PASSTHROUGH_OBJECT}") + assert nested != raw_jsonl + assert json.loads(nested.splitlines()[0])["response"]["body"]["choices"] + + def test_output_of_an_upload_whose_model_smuggles_the_passthrough_segment_is_still_transformed(self, config): + smuggled_model = b"litellm-vertex-files/passthrough/publishers/google/models/gemini-2.5-flash" + upload_url = self._upload_url( + config, + {"gcs_bucket_name": "my-bucket"}, + ("batch.jsonl", OPENAI_BATCH_JSONL.replace(b"gemini-2.5-flash", smuggled_model), "application/jsonl"), + ) + object_name = parse_qs(urlparse(upload_url).query)["name"][0] + raw_jsonl = _native_output_jsonl() + raw_response = httpx.Response( + status_code=200, + content=raw_jsonl, + headers={"content-type": "application/octet-stream"}, + request=httpx.Request("GET", _gcs_media_url(f"{object_name}/predictions.jsonl")), + ) + + result = config.transform_file_content_response( + raw_response=raw_response, logging_obj=MagicMock(), litellm_params={} + ) + + assert object_name.startswith("litellm-vertex-files/litellm-vertex-files/passthrough/") + assert json.loads(result.response.content.splitlines()[0])["response"]["body"]["choices"] + + @pytest.mark.parametrize( + "url, expected", + [ + (f"gs://my-bucket/{PASSTHROUGH_OBJECT}", True), + (f"gs://my-bucket/team-a/{PASSTHROUGH_OBJECT}", True), + (f"gs://my-bucket/litellm-vertex-files/{PASSTHROUGH_OBJECT}", False), + (_gcs_media_url(f"team-a/sub/{PASSTHROUGH_OBJECT}"), True), + (_gcs_media_url(f"litellm-vertex-files/publishers/google/models/x/{PASSTHROUGH_OBJECT}"), False), + (_gcs_media_url(TRANSFORMED_OBJECT), False), + ], + ids=["gs", "gs-prefixed", "gs-smuggled", "https-prefixed", "https-model-path-smuggled", "https-transformed"], + ) + def test_passthrough_detection_anchors_on_the_first_managed_segment(self, url, expected): + assert is_passthrough_managed_gcs_url(url) is expected + + @pytest.mark.asyncio + async def test_passthrough_output_stream_is_returned_untouched(self, config): + stream_iterator = object() + headers = {"content-type": "application/octet-stream"} + result = await config.transform_file_content_stream( + stream_iterator=stream_iterator, + headers=headers, + request_url=_gcs_media_url(f"team-a/{PASSTHROUGH_OBJECT}"), + logging_obj=MagicMock(), + litellm_params={}, + ) + assert result.stream_iterator is stream_iterator + assert result.headers == headers + + +class TestEmbeddingOutputTranslation: + EMBEDDING_OBJECT = ( + "litellm-vertex-files/publishers/google/models/gemini-embedding-2/prediction-model-1/predictions.jsonl" + ) + + def _transform(self, config: VertexAIFilesConfig, rows: list[dict]) -> list[dict]: + raw_response = httpx.Response( + status_code=200, + content="\n".join(json.dumps(row) for row in rows).encode(), + headers={"content-type": "application/octet-stream"}, + request=httpx.Request("GET", _gcs_media_url(self.EMBEDDING_OBJECT)), + ) + result = config.transform_file_content_response( + raw_response=raw_response, logging_obj=MagicMock(), litellm_params={} + ) + return [json.loads(line) for line in result.response.content.decode().splitlines()] + + def test_embedding_rows_become_openai_batch_rows_billed_by_their_prompt_tokens(self, config): + live_row = { + "key": "request-1", + "request": {"content": {"parts": [{"text": "hello world"}]}}, + "response": {"embedding": {"values": [-0.015, 0.024]}, "usageMetadata": {"promptTokenCount": 2}}, + } + documented_row = { + "key": "request-2", + "request": {"content": {"parts": [{"text": "hello"}]}}, + "response": {"embedding": {"values": [0.5]}, "tokenCount": "3"}, + } + + live, documented = self._transform(config, [live_row, documented_row]) + + assert (live["custom_id"], live["error"], live["response"]["status_code"]) == ("request-1", None, 200) + assert live["response"]["body"]["model"] == "gemini-embedding-2" + assert live["response"]["body"]["data"] == [{"embedding": [-0.015, 0.024], "index": 0, "object": "embedding"}] + live_usage, documented_usage = (row["response"]["body"]["usage"] for row in (live, documented)) + assert (live_usage["prompt_tokens"], live_usage["total_tokens"]) == (2, 2) + assert (documented_usage["prompt_tokens"], documented_usage["total_tokens"]) == (3, 3) diff --git a/tests/test_litellm/proxy/openai_files_endpoint/test_files_batch_file_validation.py b/tests/test_litellm/proxy/openai_files_endpoint/test_files_batch_file_validation.py index f5542fc0446..3a73c39e177 100644 --- a/tests/test_litellm/proxy/openai_files_endpoint/test_files_batch_file_validation.py +++ b/tests/test_litellm/proxy/openai_files_endpoint/test_files_batch_file_validation.py @@ -4,7 +4,8 @@ import pytest from litellm.proxy._types import ProxyException from litellm.proxy.openai_files_endpoints.batch_file_validation import ( - BATCH_LINE_REQUIRED_KEYS, + BATCH_LINE_SHAPE, + PASSTHROUGH_BATCH_LINE_SHAPE, BatchFileEmpty, BatchFileInvalidJsonLine, BatchFileLineNotObject, @@ -96,7 +97,7 @@ def test_non_object_line_rejected(): assert check_batch_file_upload("batch.jsonl", content, None) == BatchFileLineNotObject(line_number=2) -@pytest.mark.parametrize("missing_key", BATCH_LINE_REQUIRED_KEYS) +@pytest.mark.parametrize("missing_key", BATCH_LINE_SHAPE.required_keys) def test_missing_required_key_rejected(missing_key): import json @@ -174,3 +175,36 @@ def test_failures_map_to_openai_shaped_proxy_exceptions(failure, expected_code, assert exc_info.value.param == expected_param for fragment in expected_fragments: assert fragment in exc_info.value.message + + +NATIVE_VERTEX_LINE = b'{"request": {"contents": [{"role": "user", "parts": [{"text": "hi"}]}]}}' + + +def test_passthrough_keys_accept_native_vertex_rows(): + content = NATIVE_VERTEX_LINE + b"\n" + NATIVE_VERTEX_LINE + b"\n" + assert check_batch_file_upload("batch.jsonl", content, None, PASSTHROUGH_BATCH_LINE_SHAPE) is None + + +def test_passthrough_keys_reject_openai_rows(): + content = NATIVE_VERTEX_LINE + b"\n" + VALID_LINE + b"\n" + assert check_batch_file_upload( + "batch.jsonl", content, None, PASSTHROUGH_BATCH_LINE_SHAPE + ) == BatchFileMissingLineKey(line_number=2, key="request", line_shape=PASSTHROUGH_BATCH_LINE_SHAPE) + + +def test_default_keys_still_reject_native_vertex_rows(): + assert check_batch_file_upload("batch.jsonl", NATIVE_VERTEX_LINE, None) == BatchFileMissingLineKey( + line_number=1, key="custom_id" + ) + + +def test_passthrough_missing_key_message_says_what_a_passthrough_upload_takes(): + with pytest.raises(ProxyException) as exc_info: + raise_batch_file_validation_failure( + BatchFileMissingLineKey(line_number=3, key="request", line_shape=PASSTHROUGH_BATCH_LINE_SHAPE) + ) + assert exc_info.value.param == "request" + assert "line 3" in exc_info.value.message + assert "passthrough upload takes native Vertex batch rows" in exc_info.value.message + assert "with a request key." in exc_info.value.message + assert "custom_id" not in exc_info.value.message diff --git a/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py b/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py index 48699b47e7f..84cf4ea7c32 100644 --- a/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py +++ b/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py @@ -5669,3 +5669,241 @@ def test_model_routed_file_retrieve_allows_key_with_model_grant(mocker: MockerFi assert response.status_code == 200, response.text assert captured_kwargs["api_key"] == "mistral-key" assert captured_kwargs["custom_llm_provider"] == "mistral" + + +NATIVE_VERTEX_BATCH_LINE = ( + b'{"request": {"contents": [{"role": "user", "parts": [{"text": "What is the tallest building?"}]}],' + b' "tools": [{"googleSearch": {"excludeDomains": ["example.com"]}}]}}\n' +) + + +def _passthrough_router() -> Router: + return Router( + model_list=[ + { + "model_name": "vertex-batch", + "litellm_params": { + "model": "vertex_ai/gemini-2.5-flash", + "vertex_project": "proj", + "vertex_location": "us-central1", + }, + "model_info": {"id": "vertex-batch-id"}, + }, + { + "model_name": "gpt-3.5-turbo", + "litellm_params": {"model": "openai/gpt-3.5-turbo", "api_key": "openai_api_key"}, + "model_info": {"id": "gpt-3.5-turbo-id"}, + }, + ] + ) + + +def _setup_passthrough_upload_endpoint(monkeypatch, llm_router: Router) -> list: + """Like _setup_batch_upload_endpoint, but reads the forwarded file bytes while the spool is open.""" + from litellm.proxy.openai_files_endpoints import files_endpoints as fe + + forwarded_calls = _setup_batch_upload_endpoint(monkeypatch, llm_router) + + async def fake_route_create_file(**kwargs): + upload_source = kwargs["_create_file_request"]["file"][1] + upload_source.seek(0) + forwarded_calls.append({**kwargs, "file_bytes": upload_source.read()}) + return OpenAIFileObject( + id="dummy-id", + object="file", + bytes=0, + created_at=1234567890, + filename="batch.jsonl", + purpose="batch", + status="uploaded", + ) + + monkeypatch.setattr(fe, "route_create_file", fake_route_create_file) + return forwarded_calls + + +def _upload(content: bytes, form: dict): + return client.post( + "/v1/files", + files={"file": ("batch.jsonl", content, "application/jsonl")}, + data=form, + headers={"Authorization": "Bearer test-key"}, + ) + + +def test_create_file_passthrough_forwards_native_vertex_rows_untouched(monkeypatch): + forwarded_calls = _setup_passthrough_upload_endpoint(monkeypatch, _passthrough_router()) + content = NATIVE_VERTEX_BATCH_LINE * 2 + + try: + response = _upload(content, {"purpose": "batch", "target_model_names": "vertex-batch", "passthrough": "true"}) + finally: + _teardown_batch_upload_endpoint() + + assert response.status_code == 200, response.text + (call,) = forwarded_calls + assert call["_create_file_request"]["passthrough"] is True + assert call["file_bytes"] == content + assert call["target_model_names_list"] == ["vertex-batch"] + + +def test_create_file_passthrough_rejects_rows_without_a_request(monkeypatch): + forwarded_calls = _setup_passthrough_upload_endpoint(monkeypatch, _passthrough_router()) + + try: + response = _upload( + NATIVE_VERTEX_BATCH_LINE + VALID_BATCH_LINE, + {"purpose": "batch", "target_model_names": "vertex-batch", "passthrough": "true"}, + ) + finally: + _teardown_batch_upload_endpoint() + + assert response.status_code == 400, response.text + error = response.json()["error"] + assert error["param"] == "request" + assert "line 2" in error["message"] + assert forwarded_calls == [] + + +def test_create_file_without_passthrough_still_rejects_native_vertex_rows(monkeypatch): + forwarded_calls = _setup_passthrough_upload_endpoint(monkeypatch, _passthrough_router()) + + try: + response = _upload(NATIVE_VERTEX_BATCH_LINE, {"purpose": "batch", "target_model_names": "vertex-batch"}) + finally: + _teardown_batch_upload_endpoint() + + assert response.status_code == 400, response.text + assert response.json()["error"]["param"] == "custom_id" + assert forwarded_calls == [] + + +@pytest.mark.parametrize( + "form, expected_param, expected_fragment", + [ + ({"purpose": "batch", "passthrough": "true"}, "target_model_names", "target_model_names"), + ( + {"purpose": "batch", "target_model_names": "gpt-3.5-turbo", "passthrough": "true"}, + "target_model_names", + "'gpt-3.5-turbo'", + ), + ( + {"purpose": "batch", "target_model_names": "vertex-batch,gpt-3.5-turbo", "passthrough": "true"}, + "target_model_names", + "'gpt-3.5-turbo'", + ), + ({"purpose": "user_data", "target_model_names": "vertex-batch", "passthrough": "true"}, "passthrough", "batch"), + ( + {"purpose": "batch", "target_model_names": "vertex-batch", "passthrough": "true", "target_storage": "s3"}, + "target_storage", + "'s3'", + ), + ( + {"purpose": "batch", "model": "gpt-3.5-turbo", "passthrough": "true"}, + "model", + "'gpt-3.5-turbo'", + ), + ], + ids=[ + "no-model", + "non-vertex-model", + "mixed-models", + "non-batch-purpose", + "target-storage", + "non-vertex-model-param", + ], +) +def test_create_file_passthrough_rejected_outside_a_vertex_batch(monkeypatch, form, expected_param, expected_fragment): + forwarded_calls = _setup_passthrough_upload_endpoint(monkeypatch, _passthrough_router()) + + try: + response = _upload(NATIVE_VERTEX_BATCH_LINE, form) + finally: + _teardown_batch_upload_endpoint() + + assert response.status_code == 400, response.text + error = response.json()["error"] + assert error["type"] == "invalid_request_error" + assert error["param"] == expected_param + assert expected_fragment in error["message"] + assert forwarded_calls == [] + + +def test_create_file_passthrough_accepts_the_model_param_as_the_deployment(monkeypatch): + forwarded_calls = _setup_passthrough_upload_endpoint(monkeypatch, _passthrough_router()) + + try: + response = _upload( + NATIVE_VERTEX_BATCH_LINE, {"purpose": "batch", "model": "vertex-batch", "passthrough": "true"} + ) + finally: + _teardown_batch_upload_endpoint() + + assert response.status_code == 200, response.text + (call,) = forwarded_calls + assert call["model"] == "vertex-batch" + assert call["_create_file_request"]["passthrough"] is True + + +def test_create_file_passthrough_rejects_a_model_group_with_a_non_vertex_deployment(monkeypatch): + mixed_router = Router( + model_list=[ + { + "model_name": "vertex-batch", + "litellm_params": { + "model": "vertex_ai/gemini-2.5-flash", + "vertex_project": "proj", + "vertex_location": "us-central1", + }, + "model_info": {"id": "vertex-batch-id"}, + }, + { + "model_name": "vertex-batch", + "litellm_params": {"model": "openai/gpt-4.1-mini", "api_key": "openai_api_key"}, + "model_info": {"id": "vertex-batch-openai-id"}, + }, + ] + ) + forwarded_calls = _setup_passthrough_upload_endpoint(monkeypatch, mixed_router) + + try: + response = _upload( + NATIVE_VERTEX_BATCH_LINE, {"purpose": "batch", "target_model_names": "vertex-batch", "passthrough": "true"} + ) + finally: + _teardown_batch_upload_endpoint() + + assert response.status_code == 400, response.text + error = response.json()["error"] + assert error["param"] == "target_model_names" + assert "'vertex-batch'" in error["message"] + assert forwarded_calls == [] + + +def test_create_file_passthrough_fails_closed_when_guardrails_would_scan_the_batch(monkeypatch): + """Batch guardrails read OpenAI-shaped rows, so a passthrough upload on a guardrailed + key is refused rather than forwarded unscanned.""" + from litellm.integrations.custom_guardrail import CustomGuardrail + from litellm.proxy.utils import ProxyLogging + + class _Redactor(CustomGuardrail): + async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type): + return data + + forwarded_calls = _setup_passthrough_upload_endpoint(monkeypatch, _passthrough_router()) + monkeypatch.setattr(litellm, "callbacks", [_Redactor(guardrail_name="g", default_on=True)]) + ProxyLogging._callback_capabilities_cache.clear() + + try: + response = _upload( + NATIVE_VERTEX_BATCH_LINE, {"purpose": "batch", "target_model_names": "vertex-batch", "passthrough": "true"} + ) + finally: + _teardown_batch_upload_endpoint() + ProxyLogging._callback_capabilities_cache.clear() + + assert response.status_code == 400, response.text + error = response.json()["error"] + assert error["param"] == "passthrough" + assert "guardrails" in error["message"] + assert forwarded_calls == [] diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 411377f29cf..f985b212b01 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -528,6 +528,39 @@ async def test_async_router_acreate_file_with_jsonl(): assert first_call_content == non_jsonl_content +@pytest.mark.asyncio +async def test_async_router_acreate_file_passthrough_keeps_the_file_and_forwards_the_flag(): + """A passthrough batch upload must reach the provider byte for byte: the router + neither rewrites body.model to the deployment model nor drops the flag.""" + from io import BytesIO + from unittest.mock import MagicMock, patch + + jsonl_content = b'{"custom_id": "r1", "method": "POST", "url": "/v1/chat/completions", "body": {"model": "vertex-batch"}}\n' + router = litellm.Router( + model_list=[ + { + "model_name": "vertex-batch", + "litellm_params": {"model": "vertex_ai/gemini-2.5-flash", "vertex_project": "p"}, + } + ], + ) + + with patch("litellm.acreate_file", return_value=MagicMock()) as mock_acreate_file: + await router.acreate_file( + model="vertex-batch", purpose="batch", file=BytesIO(jsonl_content), passthrough=True + ) + forwarded = mock_acreate_file.call_args.kwargs + assert forwarded["passthrough"] is True + forwarded["file"].seek(0) + assert forwarded["file"].read() == jsonl_content + + mock_acreate_file.reset_mock() + await router.acreate_file(model="vertex-batch", purpose="batch", file=BytesIO(jsonl_content)) + rewritten = mock_acreate_file.call_args.kwargs["file"] + rewritten.seek(0) + assert b'"gemini-2.5-flash"' in rewritten.read() + + @pytest.mark.asyncio async def test_async_router_acreate_file_does_not_fall_back_across_model_groups(): """A file created for batches only exists under the credentials of the model group diff --git a/tests/unit/batches/test_batch_utils.py b/tests/unit/batches/test_batch_utils.py index a4de8eee23c..d1572f4a7c9 100644 --- a/tests/unit/batches/test_batch_utils.py +++ b/tests/unit/batches/test_batch_utils.py @@ -640,7 +640,7 @@ async def test_calculate_vertex_disable_transform_path(monkeypatch): monkeypatch.setattr( bu, "calculate_vertex_ai_batch_cost_and_usage", - lambda content, model: bu.BatchCostUsageResult( + lambda content, model, model_info=None: bu.BatchCostUsageResult( cost=9.9, usage=Usage(prompt_tokens=1, completion_tokens=2, total_tokens=3), models=["gemini-2.0-flash-001"], @@ -671,7 +671,7 @@ async def test_calculate_vertex_disable_transform_needs_model_name(monkeypatch): monkeypatch.setattr( bu, "calculate_vertex_ai_batch_cost_and_usage", - lambda content, model: pytest.fail("raw vertex path should not run"), + lambda content, model, model_info=None: pytest.fail("raw vertex path should not run"), ) result = await bu.calculate_batch_cost_and_usage(file_content_dictionary=[], custom_llm_provider="vertex_ai") @@ -735,7 +735,11 @@ def test_vertex_batch_usage_preserves_modality_token_details(monkeypatch): ) responses = [ { + "key": "id_1", + "status": "", + "request": {"content": {"parts": [{"text": "hello"}, {"fileData": {"mimeType": "audio/wav"}}]}}, "response": { + "embedding": {"values": [0.1, 0.2]}, "usageMetadata": { "promptTokenCount": 84, "candidatesTokenCount": 0, @@ -744,13 +748,14 @@ def test_vertex_batch_usage_preserves_modality_token_details(monkeypatch): {"modality": "AUDIO", "tokenCount": 64}, {"modality": "TEXT", "tokenCount": 20}, ], - } - } + }, + }, } ] result = bu.calculate_vertex_ai_batch_cost_and_usage(responses, "gemini-embedding-2") + assert (result.successful_requests, result.usage.prompt_tokens) == (1, 84) assert result.prompt_cost == pytest.approx(64 * 3.25e-6 + 20 * 1e-7) @@ -1336,7 +1341,7 @@ async def test_handle_completed_batch_vertex_disable_transform_path(monkeypatch) monkeypatch.setattr(litellm, "disable_vertex_batch_output_transformation", True, raising=False) seen: dict = {} - def fake_vertex_calc(content, model): + def fake_vertex_calc(content, model, model_info=None): seen["content"] = content seen["model"] = model return bu.BatchCostUsageResult( @@ -1358,7 +1363,7 @@ async def test_handle_completed_batch_vertex_disable_transform_path(monkeypatch) assert result.cost == 7.7 assert result.usage.total_tokens == 3 assert result.models == ["gemini-x"] - assert seen["content"] == raw_rows + assert list(seen["content"]) == raw_rows assert seen["model"] == "gemini-x" diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 2938e65cfde..052cc693f8b 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -25526,6 +25526,11 @@ export interface components { file: string; /** Litellm Metadata */ litellm_metadata?: string | null; + /** + * Passthrough + * @default false + */ + passthrough: boolean; /** Purpose */ purpose: string; /** @@ -25550,6 +25555,11 @@ export interface components { file: string; /** Litellm Metadata */ litellm_metadata?: string | null; + /** + * Passthrough + * @default false + */ + passthrough: boolean; /** Purpose */ purpose: string; /** @@ -25574,6 +25584,11 @@ export interface components { file: string; /** Litellm Metadata */ litellm_metadata?: string | null; + /** + * Passthrough + * @default false + */ + passthrough: boolean; /** Purpose */ purpose: string; /** From 1edc4ba580091d3e3d344b63d3f92d08b9b21a4d Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 24 Sep 2026 13:01:12 -0700 Subject: [PATCH 071/120] fix(logging): pass provider response headers to callbacks on every endpoint (#42824) * fix(logging): pass provider response headers to callbacks on every endpoint Custom callbacks only received kwargs["response_headers"] for chat completions. Responses, image generation and edit, speech, and transcription calls either never recorded the provider's headers or recorded them in one place and not the other. Every handler now records the provider's httpx headers on the response's hidden params as "headers" (raw) and "additional_headers" (processed, with LiteLLM's own entries winning on a clash), and the logging object derives model_call_details["response_headers"] from those hidden params before cost calculation on the non-stream and both streaming success paths, keeping a handler-set value authoritative. Binary speech responses expose their hidden params to the standard logging payload, and the sync OpenAI transcription request always fetches the raw response. * test(images): point the legacy image and speech fakes at the raw response surface Image generation now goes through the SDK's raw response so the provider headers can be read, and the speech binary response now carries hidden params. The unit fakes in the image generation, xinference, proxy provider, image edit, Vertex speech, and otel suites still pinned the old call surface and the old "no hidden params" assertion, so they read an uncalled mock or a fake response without headers. * test(images): drop the rewritten mock comments and the generated edit PNGs * test(images): move the llm-span test's image fake to the raw response surface --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- litellm/litellm_core_utils/core_helpers.py | 27 +++ litellm/litellm_core_utils/litellm_logging.py | 20 +- litellm/llms/custom_httpx/llm_http_handler.py | 20 +- litellm/llms/openai/openai.py | 29 ++- litellm/llms/openai/transcriptions/handler.py | 33 +-- tests/image_gen_tests/test_image_edits.py | 4 + tests/image_gen_tests/test_xinference.py | 30 ++- .../test_litellm_proxy_provider.py | 16 +- tests/llm_translation/test_openai.py | 13 +- .../otel/test_otel_v2_sources_of_truth.py | 8 +- .../litellm_core_utils/test_core_helpers.py | 68 ++++++ .../test_litellm_logging.py | 87 ++++++++ .../custom_httpx/test_llm_http_handler.py | 203 +++++++++++++++++- tests/test_litellm/llms/openai/test_openai.py | 99 ++++++++- .../test_openai_transcriptions_handler.py | 71 ++++++ .../test_non_chat_routes_open_llm_spans.py | 14 +- ...t_openai_image_generation_extra_headers.py | 50 +++-- .../text_to_speech/test_transformation.py | 1 + 18 files changed, 714 insertions(+), 79 deletions(-) create mode 100644 tests/test_litellm/llms/openai/transcriptions/test_openai_transcriptions_handler.py diff --git a/litellm/litellm_core_utils/core_helpers.py b/litellm/litellm_core_utils/core_helpers.py index 3afa6a913b5..b095b4b12c6 100644 --- a/litellm/litellm_core_utils/core_helpers.py +++ b/litellm/litellm_core_utils/core_helpers.py @@ -765,3 +765,30 @@ def set_response_cost_in_hidden_params(response: _CarriesHiddenParams, cost: flo RESPONSE_COST_HEADER: cost, } hidden_params["additional_headers"] = merged + + +_HIDDEN_PARAMS_ADAPTER: Final = TypeAdapter(Mapping[str, object]) +_PROVIDER_HEADERS_ADAPTER: Final = TypeAdapter(Mapping[str, str]) + + +def set_provider_response_headers_in_hidden_params( + response: _CarriesHiddenParams, headers: httpx.Headers | Mapping[str, str] +) -> None: + hidden_params: Final = response._hidden_params # pyright: ignore[reportPrivateUsage] # no public accessor + existing_additional_headers: Final[object] = hidden_params.get("additional_headers") + raw_headers: Final[dict[str, str]] = dict(headers) # mutable-ok: stored as the plain-dict hidden param + additional_headers: Final[dict[str, object]] = { # mutable-ok: assigned into the plain-dict hidden params + **process_response_headers(raw_headers), + **(existing_additional_headers if isinstance(existing_additional_headers, Mapping) else _NO_HEADERS), + } + hidden_params["headers"] = raw_headers + hidden_params["additional_headers"] = additional_headers + + +def get_provider_response_headers_from_hidden_params(response: object) -> Mapping[str, str] | None: + hidden_params: Final[object] = getattr(response, "_hidden_params", None) + try: + validated: Final = _HIDDEN_PARAMS_ADAPTER.validate_python(hidden_params) + return _PROVIDER_HEADERS_ADAPTER.validate_python(validated.get("headers")) + except ValidationError: + return None diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index a6391a2ae27..3b9419d483b 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -72,6 +72,7 @@ from litellm.litellm_core_utils.classifier_logging import ( is_classifier_call, ) from litellm.litellm_core_utils.core_helpers import ( + get_provider_response_headers_from_hidden_params, is_expected_client_error, reconstruct_model_name, set_response_cost_in_hidden_params, @@ -2353,6 +2354,15 @@ class Logging(LiteLLMLoggingBaseClass): ) return logging_result + def _surface_response_headers_from_result(self, logging_result: object) -> None: + existing: Final[object] = self.model_call_details.get("response_headers") + if existing is not None: + return + headers: Final = get_provider_response_headers_from_hidden_params(logging_result) + if headers is None: + return + self.model_call_details["response_headers"] = headers + def _merge_hidden_params_from_response_into_metadata(self, logging_result: object) -> None: """ Copy response._hidden_params into litellm_params.metadata['hidden_params']. @@ -2386,6 +2396,7 @@ class Logging(LiteLLMLoggingBaseClass): build_logging_payload: bool = True, ): """Resolve hidden params, compute response cost, and emit the standard logging payload.""" + self._surface_response_headers_from_result(logging_result) hidden_params: Final = getattr(logging_result, "_hidden_params", {}) if hidden_params: if self.model_call_details.get("litellm_params") is not None: @@ -2788,6 +2799,7 @@ class Logging(LiteLLMLoggingBaseClass): if complete_streaming_response is not None: verbose_logger.debug("Logging Details LiteLLM-Success Call streaming complete") self.model_call_details["complete_streaming_response"] = complete_streaming_response + self._surface_response_headers_from_result(complete_streaming_response) self.model_call_details["response_cost"] = self._response_cost_calculator( result=complete_streaming_response ) @@ -3302,6 +3314,7 @@ class Logging(LiteLLMLoggingBaseClass): print_verbose("Async success callbacks: Got a complete streaming response") self.model_call_details["async_complete_streaming_response"] = complete_streaming_response + self._surface_response_headers_from_result(complete_streaming_response) try: if self.model_call_details.get("cache_hit", False) is True: @@ -6362,12 +6375,15 @@ def _extract_response_obj_and_hidden_params( original_exception: Exception | None, ) -> tuple[dict, dict | None]: """Extract response_obj and hidden_params from init_response_obj.""" - hidden_params: dict | None = None + hidden_params: dict | None = ( + getattr(init_response_obj, "_hidden_params", None) + if isinstance(init_response_obj, BaseModel | HttpxBinaryResponseContent) + else None + ) if init_response_obj is None: response_obj = {} elif isinstance(init_response_obj, BaseModel): response_obj = init_response_obj.model_dump() - hidden_params = getattr(init_response_obj, "_hidden_params", None) elif isinstance(init_response_obj, dict): response_obj = init_response_obj elif isinstance(init_response_obj, HttpxBinaryResponseContent): diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index dba0dee38fc..4f31742aaaa 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -45,6 +45,7 @@ from litellm.litellm_core_utils.audio_utils.subtitle_utils import ( SUBTITLE_RESPONSE_FORMATS, synthesize_subtitle_document, ) +from litellm.litellm_core_utils.core_helpers import set_provider_response_headers_in_hidden_params from litellm.litellm_core_utils.get_litellm_params import AWS_CREDENTIAL_KWARGS_KEYS from litellm.litellm_core_utils.llm_request_utils import serialize_multipart_form_fields from litellm.litellm_core_utils.realtime_errors import ( @@ -1461,6 +1462,7 @@ class BaseLLMHTTPHandler: transformed: Final = provider_config.transform_audio_transcription_response( raw_response=response, ) + set_provider_response_headers_in_hidden_params(transformed, response.headers) if not provider_config.supports_subtitle_synthesis: return transformed requested_format: Final = optional_params.get("response_format") @@ -6960,11 +6962,13 @@ class BaseLLMHTTPHandler: provider_config=image_edit_provider_config, ) - return image_edit_provider_config.transform_image_edit_response( + image_edit_response: Final = image_edit_provider_config.transform_image_edit_response( model=model, raw_response=response, logging_obj=logging_obj, ) + set_provider_response_headers_in_hidden_params(image_edit_response, response.headers) + return image_edit_response async def async_image_edit_handler( self, @@ -7059,11 +7063,13 @@ class BaseLLMHTTPHandler: provider_config=image_edit_provider_config, ) - return image_edit_provider_config.transform_image_edit_response( + image_edit_response: Final = image_edit_provider_config.transform_image_edit_response( model=model, raw_response=response, logging_obj=logging_obj, ) + set_provider_response_headers_in_hidden_params(image_edit_response, response.headers) + return image_edit_response def image_generation_handler( self, @@ -7186,6 +7192,7 @@ class BaseLLMHTTPHandler: litellm_params=dict(litellm_params), encoding=None, ) + set_provider_response_headers_in_hidden_params(model_response, response.headers) return model_response @@ -7293,6 +7300,7 @@ class BaseLLMHTTPHandler: litellm_params=dict(litellm_params), encoding=None, ) + set_provider_response_headers_in_hidden_params(model_response, response.headers) return model_response @@ -12077,11 +12085,13 @@ class BaseLLMHTTPHandler: provider_config=text_to_speech_provider_config, ) - return text_to_speech_provider_config.transform_text_to_speech_response( + speech_response: Final = text_to_speech_provider_config.transform_text_to_speech_response( model=model, raw_response=response, logging_obj=logging_obj, ) + set_provider_response_headers_in_hidden_params(speech_response, response.headers) + return speech_response async def async_text_to_speech_handler( self, @@ -12176,11 +12186,13 @@ class BaseLLMHTTPHandler: provider_config=text_to_speech_provider_config, ) - return text_to_speech_provider_config.transform_text_to_speech_response( + speech_response: Final = text_to_speech_provider_config.transform_text_to_speech_response( model=model, raw_response=response, logging_obj=logging_obj, ) + set_provider_response_headers_in_hidden_params(speech_response, response.headers) + return speech_response ######################################################### ########## SKILLS API HANDLERS ########################## diff --git a/litellm/llms/openai/openai.py b/litellm/llms/openai/openai.py index 63874ca9619..d6340d182ae 100644 --- a/litellm/llms/openai/openai.py +++ b/litellm/llms/openai/openai.py @@ -27,6 +27,7 @@ from litellm import LlmProviders from litellm._logging import verbose_logger from litellm.constants import DEFAULT_MAX_RETRIES from litellm.files.types import FileContentStreamingResult +from litellm.litellm_core_utils.core_helpers import set_provider_response_headers_in_hidden_params from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.logging_utils import speech_request_body, track_llm_api_timing from litellm.llms.base_llm.base_model_iterator import BaseModelResponseIterator @@ -1404,7 +1405,6 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): organization: str | None = None, headers: dict | None = None, ): - response = None try: openai_aclient: Final = self._get_openai_client( is_async=True, @@ -1428,8 +1428,10 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): ) request_data: Final = {**data, "extra_headers": headers} if headers else data - response = await openai_aclient.images.generate(**request_data, timeout=timeout) - stringified_response: Final = response.model_dump() + raw_response: Final = await openai_aclient.images.with_raw_response.generate( + **request_data, timeout=timeout + ) + stringified_response: Final = raw_response.parse().model_dump() ## LOGGING logging_obj.post_call( input=prompt, @@ -1437,11 +1439,13 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): additional_args={"complete_input_dict": data}, original_response=stringified_response, ) - return convert_to_model_response_object( + image_response: Final[ImageResponse] = convert_to_model_response_object( response_object=stringified_response, model_response_object=model_response, response_type="image_generation", ) + set_provider_response_headers_in_hidden_params(image_response, raw_response.headers) + return image_response except Exception as e: ## LOGGING logging_obj.post_call( @@ -1512,9 +1516,9 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): ## COMPLETION CALL request_data: Final = {**data, "extra_headers": headers} if headers else data - _response: Final = openai_client.images.generate(**request_data, timeout=timeout) + raw_response: Final = openai_client.images.with_raw_response.generate(**request_data, timeout=timeout) - response: Final = _response.model_dump() + response: Final = raw_response.parse().model_dump() ## LOGGING logging_obj.post_call( input=prompt, @@ -1522,11 +1526,13 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): additional_args={"complete_input_dict": data}, original_response=response, ) - return convert_to_model_response_object( + image_response: Final[ImageResponse] = convert_to_model_response_object( response_object=response, model_response_object=model_response, response_type="image_generation", ) + set_provider_response_headers_in_hidden_params(image_response, raw_response.headers) + return image_response except OpenAIError as e: ## LOGGING logging_obj.post_call( @@ -1609,7 +1615,9 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): input=input, **optional_params, ) - return HttpxBinaryResponseContent(response=response.response) + speech_response: Final = HttpxBinaryResponseContent(response=response.response) + set_provider_response_headers_in_hidden_params(speech_response, response.response.headers) + return speech_response async def async_audio_speech( self, @@ -1655,8 +1663,9 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): input=input, **optional_params, ) - - return HttpxBinaryResponseContent(response=response.response) + speech_response: Final = HttpxBinaryResponseContent(response=response.response) + set_provider_response_headers_in_hidden_params(speech_response, response.response.headers) + return speech_response class OpenAIFilesAPI(BaseLLM): diff --git a/litellm/llms/openai/transcriptions/handler.py b/litellm/llms/openai/transcriptions/handler.py index 701b3d30362..014251db821 100644 --- a/litellm/llms/openai/transcriptions/handler.py +++ b/litellm/llms/openai/transcriptions/handler.py @@ -4,11 +4,10 @@ import httpx from openai import AsyncOpenAI, OpenAI from pydantic import BaseModel -import litellm - if TYPE_CHECKING: from aiohttp import ClientSession from litellm.litellm_core_utils.audio_utils.utils import get_audio_file_name +from litellm.litellm_core_utils.core_helpers import set_provider_response_headers_in_hidden_params from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.base_llm.audio_transcription.transformation import ( BaseAudioTranscriptionConfig, @@ -31,11 +30,6 @@ class OpenAIAudioTranscription(OpenAIChatCompletion): data: dict, timeout: float | httpx.Timeout, ): - """ - Helper to: - - call openai_aclient.audio.transcriptions.with_raw_response when litellm.return_response_headers is True - - call openai_aclient.audio.transcriptions.create by default - """ try: raw_response = await openai_aclient.audio.transcriptions.with_raw_response.create(**data, timeout=timeout) headers: Final = dict(raw_response.headers) @@ -51,20 +45,11 @@ class OpenAIAudioTranscription(OpenAIChatCompletion): data: dict, timeout: float | httpx.Timeout, ): - """ - Helper to: - - call openai_aclient.audio.transcriptions.with_raw_response when litellm.return_response_headers is True - - call openai_aclient.audio.transcriptions.create by default - """ try: - if litellm.return_response_headers is True: - raw_response = openai_client.audio.transcriptions.with_raw_response.create(**data, timeout=timeout) - headers: Final = dict(raw_response.headers) - response = raw_response.parse() - return headers, response - else: - response = openai_client.audio.transcriptions.create(**data, timeout=timeout) - return None, response + raw_response: Final = openai_client.audio.transcriptions.with_raw_response.create(**data, timeout=timeout) + headers: Final = dict(raw_response.headers) + response: Final = raw_response.parse() + return headers, response except Exception as e: raise e @@ -133,11 +118,12 @@ class OpenAIAudioTranscription(OpenAIChatCompletion): "complete_input_dict": data, }, ) - _, response = self.make_sync_openai_audio_transcriptions_request( + headers, response = self.make_sync_openai_audio_transcriptions_request( openai_client=openai_client, data=data, timeout=timeout, ) + logging_obj.model_call_details["response_headers"] = headers if isinstance(response, BaseModel): stringified_response = response.model_dump() @@ -158,6 +144,7 @@ class OpenAIAudioTranscription(OpenAIChatCompletion): hidden_params=hidden_params, response_type="audio_transcription", ) + set_provider_response_headers_in_hidden_params(final_response, headers) return final_response async def async_audio_transcriptions( @@ -217,12 +204,14 @@ class OpenAIAudioTranscription(OpenAIChatCompletion): actual_model: Final = data.get("model", "whisper-1") hidden_params: Final = {"model": actual_model, "custom_llm_provider": "openai"} - return convert_to_model_response_object( + final_response: Final[TranscriptionResponse] = convert_to_model_response_object( response_object=stringified_response, model_response_object=model_response, hidden_params=hidden_params, response_type="audio_transcription", ) + set_provider_response_headers_in_hidden_params(final_response, headers) + return final_response except Exception as e: ## LOGGING logging_obj.post_call( diff --git a/tests/image_gen_tests/test_image_edits.py b/tests/image_gen_tests/test_image_edits.py index 0c2f57066e8..36fd65ba71b 100644 --- a/tests/image_gen_tests/test_image_edits.py +++ b/tests/image_gen_tests/test_image_edits.py @@ -250,6 +250,7 @@ async def test_azure_image_edit_litellm_sdk(): self._json_data = json_data self.status_code = status_code self.text = json.dumps(json_data) + self.headers = {} def json(self): return self._json_data @@ -370,6 +371,7 @@ async def test_openai_image_edit_cost_tracking(): self._json_data = json_data self.status_code = status_code self.text = json.dumps(json_data) + self.headers = {} def json(self): return self._json_data @@ -460,6 +462,7 @@ async def test_azure_image_edit_cost_tracking(): self._json_data = json_data self.status_code = status_code self.text = json.dumps(json_data) + self.headers = {} def json(self): return self._json_data @@ -737,6 +740,7 @@ async def test_image_edit_array_handling(): self._json_data = json_data self.status_code = status_code self.text = json.dumps(json_data) + self.headers = {} def json(self): return self._json_data diff --git a/tests/image_gen_tests/test_xinference.py b/tests/image_gen_tests/test_xinference.py index 3dc4fee85da..76cae593e41 100644 --- a/tests/image_gen_tests/test_xinference.py +++ b/tests/image_gen_tests/test_xinference.py @@ -24,9 +24,14 @@ async def test_xinference_image_generation(): def model_dump(self): return mock_openai_response - # Create a mock client with the images.generate method + class MockRawResponse: + headers = {} + + def parse(self): + return MockResponse() + mock_client = AsyncMock() - mock_client.images.generate = AsyncMock(return_value=MockResponse()) + mock_client.images.with_raw_response.generate = AsyncMock(return_value=MockRawResponse()) # Capture the actual arguments sent to OpenAI client captured_args = None @@ -36,9 +41,9 @@ async def test_xinference_image_generation(): nonlocal captured_args, captured_kwargs captured_args = args captured_kwargs = kwargs - return MockResponse() + return MockRawResponse() - mock_client.images.generate.side_effect = capture_generate_call + mock_client.images.with_raw_response.generate.side_effect = capture_generate_call # Mock the _get_openai_client method to return our mock client with patch.object( @@ -65,7 +70,7 @@ async def test_xinference_image_generation(): assert response.data[0].url == "https://example.com/image.png" # Validate that the OpenAI client was called with correct parameters - mock_client.images.generate.assert_called_once() + mock_client.images.with_raw_response.generate.assert_called_once() assert captured_kwargs is not None assert ( captured_kwargs["model"] == "stabilityai/stable-diffusion-3.5-large" @@ -97,9 +102,14 @@ async def test_xinference_image_generation_with_response_format(): def model_dump(self): return mock_openai_response - # Create a mock client with the images.generate method + class MockRawResponse: + headers = {} + + def parse(self): + return MockResponse() + mock_client = AsyncMock() - mock_client.images.generate = AsyncMock(return_value=MockResponse()) + mock_client.images.with_raw_response.generate = AsyncMock(return_value=MockRawResponse()) # Capture the actual arguments sent to OpenAI client captured_args = None @@ -109,9 +119,9 @@ async def test_xinference_image_generation_with_response_format(): nonlocal captured_args, captured_kwargs captured_args = args captured_kwargs = kwargs - return MockResponse() + return MockRawResponse() - mock_client.images.generate.side_effect = capture_generate_call + mock_client.images.with_raw_response.generate.side_effect = capture_generate_call # Mock the _get_openai_client method to return our mock client with patch.object( @@ -141,7 +151,7 @@ async def test_xinference_image_generation_with_response_format(): assert response.data[0].b64_json is not None # Validate that the OpenAI client was called with correct parameters - mock_client.images.generate.assert_called_once() + mock_client.images.with_raw_response.generate.assert_called_once() assert captured_kwargs is not None assert ( captured_kwargs["model"] == "stabilityai/stable-diffusion-3.5-large" diff --git a/tests/llm_translation/test_litellm_proxy_provider.py b/tests/llm_translation/test_litellm_proxy_provider.py index a10fc55ecc5..a7a2848a514 100644 --- a/tests/llm_translation/test_litellm_proxy_provider.py +++ b/tests/llm_translation/test_litellm_proxy_provider.py @@ -210,11 +210,14 @@ async def test_litellm_gateway_image_generation_direct(is_async): "created": 1, "data": [{"url": "https://example.com/image.png"}], } + mock_raw_response = MagicMock() + mock_raw_response.parse.return_value = mock_openai_response + mock_raw_response.headers = {} if is_async: # Mock the AsyncOpenAI client that gets created inside _get_openai_client mock_async_client = AsyncMock() - mock_async_client.images.generate = AsyncMock(return_value=mock_openai_response) + mock_async_client.images.with_raw_response.generate = AsyncMock(return_value=mock_raw_response) with patch( "litellm.llms.openai.openai.AsyncOpenAI", return_value=mock_async_client @@ -234,14 +237,14 @@ async def test_litellm_gateway_image_generation_direct(is_async): assert constructor_kwargs["base_url"] == "http://my-proxy" # Verify the AsyncOpenAI client was called correctly - mock_async_client.images.generate.assert_awaited_once() - call_kwargs = mock_async_client.images.generate.call_args.kwargs + mock_async_client.images.with_raw_response.generate.assert_awaited_once() + call_kwargs = mock_async_client.images.with_raw_response.generate.call_args.kwargs assert call_kwargs["model"] == "dall-e-3" assert call_kwargs["prompt"] == "A beautiful sunset over mountains" else: # Mock the sync OpenAI client that gets created inside _get_openai_client mock_sync_client = MagicMock() - mock_sync_client.images.generate.return_value = mock_openai_response + mock_sync_client.images.with_raw_response.generate.return_value = mock_raw_response with patch( "litellm.llms.openai.openai.OpenAI", return_value=mock_sync_client @@ -260,8 +263,8 @@ async def test_litellm_gateway_image_generation_direct(is_async): assert constructor_kwargs["base_url"] == "http://my-proxy" # Verify the OpenAI client was called correctly - mock_sync_client.images.generate.assert_called_once() - call_kwargs = mock_sync_client.images.generate.call_args.kwargs + mock_sync_client.images.with_raw_response.generate.assert_called_once() + call_kwargs = mock_sync_client.images.with_raw_response.generate.call_args.kwargs assert call_kwargs["model"] == "dall-e-3" assert call_kwargs["prompt"] == "A beautiful sunset over mountains" @@ -285,6 +288,7 @@ async def test_litellm_gateway_from_sdk_image_edit(is_async): self._json_data = json_data self.status_code = status_code self.text = json.dumps(json_data) + self.headers = {} def json(self): return self._json_data diff --git a/tests/llm_translation/test_openai.py b/tests/llm_translation/test_openai.py index af4ba85d58e..0488c4c68e6 100644 --- a/tests/llm_translation/test_openai.py +++ b/tests/llm_translation/test_openai.py @@ -313,7 +313,7 @@ def test_openai_max_retries_0(mock_get_openai_client): def test_openai_image_generation_forwards_organization(mock_get_openai_client): """Ensure organization flows to OpenAI client for image generation.""" - class _DummyImages: + class _DummyRawImages: def generate(self, **kwargs): # type: ignore class _Resp: def model_dump(self_inner): # minimal OpenAI ImagesResponse shape @@ -327,7 +327,16 @@ def test_openai_image_generation_forwards_organization(mock_get_openai_client): }, } - return _Resp() + class _RawResp: + headers = {} + + def parse(self_inner): + return _Resp() + + return _RawResp() + + class _DummyImages: + with_raw_response = _DummyRawImages() class _DummyClient: def __init__(self): diff --git a/tests/test_litellm/integrations/otel/test_otel_v2_sources_of_truth.py b/tests/test_litellm/integrations/otel/test_otel_v2_sources_of_truth.py index 7e93d3d67a7..57f4557c6f7 100644 --- a/tests/test_litellm/integrations/otel/test_otel_v2_sources_of_truth.py +++ b/tests/test_litellm/integrations/otel/test_otel_v2_sources_of_truth.py @@ -1207,14 +1207,18 @@ def test_speech_response_without_a_byte_count_produces_no_output() -> None: def test_speech_binary_response_is_logged_as_its_summary_not_dropped() -> None: import httpx + from litellm.litellm_core_utils.core_helpers import set_provider_response_headers_in_hidden_params from litellm.litellm_core_utils.litellm_logging import _extract_response_obj_and_hidden_params from litellm.types.llms.openai import HttpxBinaryResponseContent raw: Final = httpx.Response(200, headers={"content-type": "audio/mpeg"}, content=b"\x00" * 1234) - response_obj, hidden_params = _extract_response_obj_and_hidden_params(HttpxBinaryResponseContent(raw), None) + speech: Final = HttpxBinaryResponseContent(raw) + set_provider_response_headers_in_hidden_params(speech, raw.headers) + response_obj, hidden_params = _extract_response_obj_and_hidden_params(speech, None) assert response_obj == {"object": "binary", "content_type": "audio/mpeg", "num_bytes": 1234} - assert hidden_params is None + assert hidden_params is not None + assert hidden_params["headers"]["content-type"] == "audio/mpeg" def test_speech_binary_response_still_streaming_reports_the_bytes_downloaded_so_far() -> None: diff --git a/tests/test_litellm/litellm_core_utils/test_core_helpers.py b/tests/test_litellm/litellm_core_utils/test_core_helpers.py index 6eeea271127..2a6dd347d5f 100644 --- a/tests/test_litellm/litellm_core_utils/test_core_helpers.py +++ b/tests/test_litellm/litellm_core_utils/test_core_helpers.py @@ -2,22 +2,27 @@ import logging +import httpx import pytest from litellm.litellm_core_utils.core_helpers import ( _FINISH_REASON_MAP, + RESPONSE_COST_HEADER, bind_budget_reservation_to_callbacks, budget_reservation_from_metadata, drop_params_env_flag, drop_params_flag, get_or_create_metadata_bucket, + get_provider_response_headers_from_hidden_params, map_finish_reason, normalize_drop_params, reconstruct_model_name, redact_nested_match_and_regex_keys, + set_provider_response_headers_in_hidden_params, unbind_budget_reservation_from_callbacks, ) from litellm.proxy._types import UserAPIKeyAuth +from litellm.types.utils import ImageResponse, TranscriptionResponse class TestBudgetReservationBinding: @@ -489,3 +494,66 @@ class TestIsExpectedClientError: category=RateLimitErrorCategory.VENDOR_RATE_LIMIT, ) assert is_expected_client_error(vendor_limit) is False + + +class TestProviderResponseHeadersInHiddenParams: + def test_records_raw_headers_and_the_processed_additional_headers(self): + response = ImageResponse() + response._hidden_params = {"additional_headers": {RESPONSE_COST_HEADER: 0.04}} + + set_provider_response_headers_in_hidden_params( + response, httpx.Headers({"X-Request-Id": "req_img", "x-ratelimit-remaining-requests": "41"}) + ) + + assert response._hidden_params["headers"] == { + "x-request-id": "req_img", + "x-ratelimit-remaining-requests": "41", + } + additional_headers = response._hidden_params["additional_headers"] + assert additional_headers["llm_provider-x-request-id"] == "req_img" + assert additional_headers["x-ratelimit-remaining-requests"] == "41" + assert additional_headers[RESPONSE_COST_HEADER] == 0.04 + + def test_litellm_owned_additional_headers_win_over_provider_headers(self): + response = TranscriptionResponse(text="hi") + response._hidden_params = {"additional_headers": {"llm_provider-x-request-id": "kept"}} + + set_provider_response_headers_in_hidden_params(response, {"x-request-id": "provider"}) + + assert response._hidden_params["additional_headers"]["llm_provider-x-request-id"] == "kept" + assert response._hidden_params["headers"] == {"x-request-id": "provider"} + + def test_getter_returns_the_recorded_headers(self): + response = ImageResponse() + + set_provider_response_headers_in_hidden_params(response, {"x-request-id": "req_img"}) + + assert get_provider_response_headers_from_hidden_params(response) == {"x-request-id": "req_img"} + + @pytest.mark.parametrize( + "hidden_params", + [ + None, + "headers", + {"additional_headers": {}}, + {"headers": "x-request-id: req_img"}, + {"headers": {"x-request-id": 7}}, + ], + ) + def test_getter_returns_none_without_a_string_header_mapping(self, hidden_params): + response = ImageResponse() + response._hidden_params = hidden_params + + assert get_provider_response_headers_from_hidden_params(response) is None + + def test_getter_returns_none_for_an_object_without_hidden_params(self): + assert get_provider_response_headers_from_hidden_params(object()) is None + + def test_headers_never_leak_into_a_sibling_response(self): + recorded = TranscriptionResponse() + sibling = TranscriptionResponse() + + set_provider_response_headers_in_hidden_params(recorded, {"x-request-id": "req_stt"}) + + assert get_provider_response_headers_from_hidden_params(sibling) is None + assert "additional_headers" not in sibling._hidden_params diff --git a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py index 23c01841b1b..bb8098e6e23 100644 --- a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py +++ b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py @@ -24,6 +24,7 @@ from litellm.cost_calculator import ocr_batch_cost from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.litellm_logging import Logging as LitellmLogging from litellm.litellm_core_utils.litellm_logging import ( + _extract_response_obj_and_hidden_params, _get_status_fields, set_callbacks, ) @@ -32,6 +33,7 @@ from litellm.proxy._types import UserAPIKeyAuth from litellm.types.llms.openai import ResponseAPIUsage, ResponseCompletedEvent, ResponsesAPIResponse from litellm.types.utils import ( CallTypes, + ImageResponse, LiteLLMRealtimeStreamLoggingObject, ModelResponse, TextCompletionResponse, @@ -8694,3 +8696,88 @@ async def test_async_failure_handler_delivers_failure_payload_to_custom_logger() assert "smoke-failure" in payload["error_str"] assert payload["model"] == "openai/gpt-5.6" assert events.empty() + + +def _image_logging_obj() -> LitellmLogging: + logging_obj = LitellmLogging( + model="gpt-image-2", + messages="a cat", + stream=False, + call_type="aimage_generation", + start_time=time.time(), + litellm_call_id="response-headers-test", + function_id="response-headers-test", + ) + logging_obj.model_call_details["litellm_params"] = {"metadata": {}} + logging_obj.optional_params = {} + return logging_obj + + +def _image_result_with_headers(request_id: str) -> ImageResponse: + result = ImageResponse(created=1, data=[]) + result._hidden_params = {"headers": {"x-request-id": request_id}} + return result + + +def test_process_hidden_params_surfaces_response_headers_from_the_result(): + logging_obj = _image_logging_obj() + + logging_obj._process_hidden_params_and_response_cost( + _image_result_with_headers("req_img"), datetime.datetime.now(), datetime.datetime.now() + ) + + assert logging_obj.model_call_details["response_headers"] == {"x-request-id": "req_img"} + + +def test_process_hidden_params_keeps_handler_set_response_headers(): + logging_obj = _image_logging_obj() + logging_obj.model_call_details["response_headers"] = {"x-request-id": "from-handler"} + + logging_obj._process_hidden_params_and_response_cost( + _image_result_with_headers("from-result"), datetime.datetime.now(), datetime.datetime.now() + ) + + assert logging_obj.model_call_details["response_headers"] == {"x-request-id": "from-handler"} + + +def _assembled_stream_result_with_headers() -> ModelResponse: + result = _assembled_stream_result() + result._hidden_params = {"headers": {"x-request-id": "req_stream"}} + return result + + +@pytest.mark.asyncio +async def test_async_streaming_success_passes_result_headers_to_callback_kwargs(): + releasing = CustomLogger() + releasing.async_log_success_event = AsyncMock() + patcher, logging_obj = _streaming_logging_obj_with_callbacks([releasing]) + + with patcher: + await logging_obj.async_success_handler(result=_assembled_stream_result_with_headers()) + + kwargs = releasing.async_log_success_event.await_args.kwargs["kwargs"] + assert kwargs["response_headers"] == {"x-request-id": "req_stream"} + + +def test_sync_streaming_success_passes_result_headers_to_callback_kwargs(): + releasing = CustomLogger() + releasing.log_success_event = MagicMock() + patcher, logging_obj = _streaming_logging_obj_with_callbacks([releasing]) + + with patcher: + logging_obj.success_handler(result=_assembled_stream_result_with_headers()) + + kwargs = releasing.log_success_event.call_args.kwargs["kwargs"] + assert kwargs["response_headers"] == {"x-request-id": "req_stream"} + + +def test_extract_response_obj_and_hidden_params_reads_binary_content_hidden_params(): + from litellm.types.llms.openai import HttpxBinaryResponseContent as LiteLLMBinaryResponseContent + + result = LiteLLMBinaryResponseContent(response=httpx.Response(status_code=200, content=b"audio bytes")) + result._hidden_params = {"headers": {"x-request-id": "req_tts"}} + + response_obj, hidden_params = _extract_response_obj_and_hidden_params(result, None) + + assert hidden_params == {"headers": {"x-request-id": "req_tts"}} + assert response_obj["object"] == "binary" diff --git a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py b/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py index 68f37c8ffcc..75b6ce4c626 100644 --- a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py +++ b/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py @@ -26,6 +26,8 @@ from litellm.llms.base_llm.search.transformation import BaseSearchConfig, Search from litellm.llms.bedrock.base_aws_llm import SignsRequestsWithAWS from litellm.llms.brave.search.transformation import BraveSearchConfig from litellm.llms.base_llm.image_edit.transformation import BaseImageEditConfig +from litellm.llms.base_llm.image_generation.transformation import BaseImageGenerationConfig +from litellm.llms.base_llm.text_to_speech.transformation import BaseTextToSpeechConfig from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.llms.custom_httpx.llm_http_handler import ( BaseLLMHTTPHandler, @@ -41,7 +43,7 @@ from litellm.llms.bedrock.messages.invoke_transformations.anthropic_claude3_tran from litellm.llms.mistral.ocr.transformation import MistralOCRConfig from litellm.llms.openai.videos.transformation import OpenAIVideoConfig from litellm.llms.tinyfish.search.transformation import TinyfishSearchConfig -from litellm.types.llms.openai import ResponsesAPIResponse +from litellm.types.llms.openai import HttpxBinaryResponseContent, ResponsesAPIResponse from litellm.types.router import GenericLiteLLMParams from litellm.types.utils import ImageObject, ImageResponse, ModelResponse, TranscriptionResponse from tests.test_litellm.llms.bedrock.event_loop_probe import EventLoopProbe @@ -4166,3 +4168,202 @@ async def test_chat_completion_agentic_followup_does_not_repeat_request_params_f assert followup_calls[0]["temperature"] == 0.2 assert followup_calls[0]["api_base"] == "https://a" assert followup_calls[0]["model"] == "openai/gpt-5" + + +_UPSTREAM_HEADERS: Final = {"x-request-id": "req_upstream", "x-ratelimit-remaining-requests": "41"} + + +def _assert_upstream_headers_recorded(response) -> None: + assert response._hidden_params["headers"]["x-request-id"] == "req_upstream" + assert response._hidden_params["additional_headers"]["llm_provider-x-request-id"] == "req_upstream" + assert response._hidden_params["additional_headers"]["x-ratelimit-remaining-requests"] == "41" + + +def _json_with_upstream_headers(payload: dict) -> httpx.MockTransport: + return httpx.MockTransport(lambda request: httpx.Response(200, json=payload, headers=_UPSTREAM_HEADERS)) + + +def _binary_with_upstream_headers() -> httpx.MockTransport: + return httpx.MockTransport( + lambda request: httpx.Response( + 200, content=b"audio-bytes", headers={**_UPSTREAM_HEADERS, "content-type": "audio/mpeg"} + ) + ) + + +def test_audio_transcriptions_records_upstream_response_headers(): + client = HTTPHandler(client=httpx.Client(transport=_json_with_upstream_headers({"text": "transcribed"}))) + + response = BaseLLMHTTPHandler().audio_transcriptions( + client=client, + atranscription=False, + **_json_transcription_call_kwargs(_JSONBodyAudioTranscriptionConfig()), + ) + + _assert_upstream_headers_recorded(response) + + +@pytest.mark.asyncio +async def test_async_audio_transcriptions_records_upstream_response_headers(): + client = AsyncHTTPHandler() + client.client = httpx.AsyncClient(transport=_json_with_upstream_headers({"text": "transcribed"})) + + response = await BaseLLMHTTPHandler().async_audio_transcriptions( + client=client, + **_json_transcription_call_kwargs(_JSONBodyAudioTranscriptionConfig()), + ) + + _assert_upstream_headers_recorded(response) + + +def _image_edit_call_kwargs() -> dict: + return { + "model": "edit-model", + "image": b"raw-image", + "prompt": "add a hat", + "image_edit_provider_config": _ImageEditRecordingConfig(), + "image_edit_optional_request_params": {}, + "custom_llm_provider": "openai", + "litellm_params": GenericLiteLLMParams(), + "logging_obj": Mock(), + "timeout": 10.0, + } + + +def test_image_edit_handler_records_upstream_response_headers(): + client = HTTPHandler() + client.client = httpx.Client(transport=_json_with_upstream_headers({"transformed_by": "sync"})) + + response = BaseLLMHTTPHandler().image_edit_handler(client=client, **_image_edit_call_kwargs()) + + _assert_upstream_headers_recorded(response) + + +@pytest.mark.asyncio +async def test_async_image_edit_handler_records_upstream_response_headers(): + client = AsyncHTTPHandler() + client.client = httpx.AsyncClient(transport=_json_with_upstream_headers({"transformed_by": "async"})) + + response = await BaseLLMHTTPHandler().async_image_edit_handler(client=client, **_image_edit_call_kwargs()) + + _assert_upstream_headers_recorded(response) + + +class _HeaderImageGenerationConfig(BaseImageGenerationConfig): + def get_supported_openai_params(self, model): + return [] + + def map_openai_params(self, non_default_params, optional_params, model, drop_params): + return optional_params + + def get_complete_url(self, api_base, api_key, model, optional_params, litellm_params, stream=None): + return "https://images.example/v1/generations" + + def transform_image_generation_request(self, model, prompt, optional_params, litellm_params, headers): + return {"prompt": prompt} + + def transform_image_generation_response( + self, + model, + raw_response, + model_response, + logging_obj, + request_data, + optional_params, + litellm_params, + encoding, + api_key=None, + json_mode=None, + ): + return ImageResponse(data=[ImageObject(b64_json=raw_response.json()["b64_json"])]) + + +def _image_generation_call_kwargs() -> dict: + return { + "model": "image-model", + "prompt": "a cat", + "image_generation_provider_config": _HeaderImageGenerationConfig(), + "image_generation_optional_request_params": {}, + "custom_llm_provider": "openai", + "litellm_params": {}, + "logging_obj": Mock(), + "timeout": 10.0, + } + + +def test_image_generation_handler_records_upstream_response_headers(): + client = HTTPHandler() + client.client = httpx.Client(transport=_json_with_upstream_headers({"b64_json": "abc"})) + + response = BaseLLMHTTPHandler().image_generation_handler(client=client, **_image_generation_call_kwargs()) + + assert response.data[0].b64_json == "abc" + _assert_upstream_headers_recorded(response) + + +@pytest.mark.asyncio +async def test_async_image_generation_handler_records_upstream_response_headers(): + client = AsyncHTTPHandler() + client.client = httpx.AsyncClient(transport=_json_with_upstream_headers({"b64_json": "abc"})) + + response = await BaseLLMHTTPHandler().async_image_generation_handler( + client=client, **_image_generation_call_kwargs() + ) + + assert response.data[0].b64_json == "abc" + _assert_upstream_headers_recorded(response) + + +class _HeaderTextToSpeechConfig(BaseTextToSpeechConfig): + def get_supported_openai_params(self, model): + return [] + + def map_openai_params(self, model, optional_params, voice=None, drop_params=False, kwargs=None): + return voice, optional_params + + def validate_environment(self, headers, model, api_key=None, api_base=None): + return {} + + def get_complete_url(self, model, api_base, litellm_params): + return "https://tts.example/v1/speech" + + def transform_text_to_speech_request(self, model, input, voice, optional_params, litellm_params, headers): + return {"dict_body": {"input": input}} + + def transform_text_to_speech_response(self, model, raw_response, logging_obj): + return HttpxBinaryResponseContent(response=raw_response) + + +def _text_to_speech_call_kwargs() -> dict: + return { + "model": "tts-model", + "input": "hello", + "voice": "alloy", + "text_to_speech_provider_config": _HeaderTextToSpeechConfig(), + "text_to_speech_optional_params": {}, + "custom_llm_provider": "openai", + "litellm_params": {}, + "logging_obj": Mock(), + "timeout": 10.0, + } + + +def test_text_to_speech_handler_records_upstream_response_headers(): + client = HTTPHandler() + client.client = httpx.Client(transport=_binary_with_upstream_headers()) + + response = BaseLLMHTTPHandler().text_to_speech_handler(client=client, **_text_to_speech_call_kwargs()) + + assert response.content == b"audio-bytes" + _assert_upstream_headers_recorded(response) + + +@pytest.mark.asyncio +async def test_async_text_to_speech_handler_records_upstream_response_headers(): + client = AsyncHTTPHandler() + client.client = httpx.AsyncClient(transport=_binary_with_upstream_headers()) + + response = await BaseLLMHTTPHandler().async_text_to_speech_handler(client=client, **_text_to_speech_call_kwargs()) + + assert response.content == b"audio-bytes" + _assert_upstream_headers_recorded(response) diff --git a/tests/test_litellm/llms/openai/test_openai.py b/tests/test_litellm/llms/openai/test_openai.py index 9539e13a802..2be691fa65b 100644 --- a/tests/test_litellm/llms/openai/test_openai.py +++ b/tests/test_litellm/llms/openai/test_openai.py @@ -1,13 +1,15 @@ import asyncio import json from typing import Final +from unittest.mock import Mock import httpx import pytest -from openai import AsyncOpenAI +from openai import AsyncOpenAI, OpenAI import litellm from litellm.llms.openai.openai import OpenAIChatCompletion +from litellm.types.utils import ImageResponse @pytest.mark.parametrize( @@ -253,3 +255,98 @@ async def test_acompletion_streams_tool_call_arguments_over_injected_transport() assert tool_call.function.name == "get_weather" assert json.loads(tool_call.function.arguments) == {"city": "Paris"} assert rebuilt.choices[0].finish_reason == "tool_calls" + + +_PROVIDER_HEADERS: Final = {"x-request-id": "req_openai", "x-ratelimit-remaining-requests": "41"} + + +def _image_generation_transport() -> httpx.MockTransport: + return httpx.MockTransport( + lambda request: httpx.Response( + 200, json={"created": 1, "data": [{"b64_json": "abc"}]}, headers=_PROVIDER_HEADERS + ) + ) + + +def _speech_transport() -> httpx.MockTransport: + return httpx.MockTransport( + lambda request: httpx.Response( + 200, content=b"audio-bytes", headers={**_PROVIDER_HEADERS, "content-type": "audio/mpeg"} + ) + ) + + +def _assert_provider_headers_recorded(response) -> None: + assert response._hidden_params["headers"]["x-request-id"] == "req_openai" + assert response._hidden_params["additional_headers"]["llm_provider-x-request-id"] == "req_openai" + assert response._hidden_params["additional_headers"]["x-ratelimit-remaining-requests"] == "41" + + +def _image_generation_kwargs() -> dict: + return { + "model": "gpt-image-2", + "prompt": "a cat", + "timeout": 10, + "optional_params": {}, + "logging_obj": Mock(), + "api_key": "transport-only", + "model_response": ImageResponse(), + } + + +def test_image_generation_records_provider_response_headers(): + with httpx.Client(transport=_image_generation_transport()) as http_client: + response = OpenAIChatCompletion().image_generation( + client=OpenAI(api_key="transport-only", http_client=http_client), **_image_generation_kwargs() + ) + + _assert_provider_headers_recorded(response) + + +@pytest.mark.asyncio +async def test_aimage_generation_records_provider_response_headers(): + async with httpx.AsyncClient(transport=_image_generation_transport()) as http_client: + response = await OpenAIChatCompletion().image_generation( + client=AsyncOpenAI(api_key="transport-only", http_client=http_client), + aimg_generation=True, + **_image_generation_kwargs(), + ) + + _assert_provider_headers_recorded(response) + + +def _audio_speech_kwargs() -> dict: + return { + "model": "gpt-4o-mini-tts", + "input": "hello", + "voice": "alloy", + "optional_params": {}, + "api_key": "transport-only", + "api_base": None, + "organization": None, + "project": None, + "max_retries": 0, + "timeout": 10, + "logging_obj": Mock(), + } + + +def test_audio_speech_records_provider_response_headers(): + with httpx.Client(transport=_speech_transport()) as http_client: + response = OpenAIChatCompletion().audio_speech( + client=OpenAI(api_key="transport-only", http_client=http_client), **_audio_speech_kwargs() + ) + + _assert_provider_headers_recorded(response) + + +@pytest.mark.asyncio +async def test_async_audio_speech_records_provider_response_headers(): + async with httpx.AsyncClient(transport=_speech_transport()) as http_client: + response = await OpenAIChatCompletion().audio_speech( + client=AsyncOpenAI(api_key="transport-only", http_client=http_client), + aspeech=True, + **_audio_speech_kwargs(), + ) + + _assert_provider_headers_recorded(response) diff --git a/tests/test_litellm/llms/openai/transcriptions/test_openai_transcriptions_handler.py b/tests/test_litellm/llms/openai/transcriptions/test_openai_transcriptions_handler.py new file mode 100644 index 00000000000..f2dbb71fea1 --- /dev/null +++ b/tests/test_litellm/llms/openai/transcriptions/test_openai_transcriptions_handler.py @@ -0,0 +1,71 @@ +from typing import Final +from unittest.mock import Mock + +import httpx +import pytest +from openai import AsyncOpenAI, OpenAI + +from litellm.llms.openai.transcriptions.handler import OpenAIAudioTranscription +from litellm.types.utils import TranscriptionResponse + +_PROVIDER_HEADERS: Final = {"x-request-id": "req_stt", "x-ratelimit-remaining-requests": "41"} + + +def _transcription_transport() -> httpx.MockTransport: + return httpx.MockTransport(lambda request: httpx.Response(200, json={"text": "hello"}, headers=_PROVIDER_HEADERS)) + + +def _logging_obj() -> Mock: + logging_obj = Mock() + logging_obj.model_call_details = {} + return logging_obj + + +def _call_kwargs(logging_obj: Mock) -> dict: + return { + "model": "gpt-4o-mini-transcribe", + "audio_file": ("audio.wav", b"riff-bytes", "audio/wav"), + "optional_params": {}, + "litellm_params": {}, + "model_response": TranscriptionResponse(), + "timeout": 10.0, + "max_retries": 0, + "logging_obj": logging_obj, + "api_key": "transport-only", + "api_base": None, + } + + +def _assert_headers_recorded(response: TranscriptionResponse, logging_obj: Mock) -> None: + assert response.text == "hello" + assert response._hidden_params["headers"]["x-request-id"] == "req_stt" + assert response._hidden_params["additional_headers"]["llm_provider-x-request-id"] == "req_stt" + assert response._hidden_params["additional_headers"]["x-ratelimit-remaining-requests"] == "41" + assert logging_obj.model_call_details["response_headers"]["x-request-id"] == "req_stt" + + +def test_audio_transcriptions_records_provider_response_headers(): + logging_obj = _logging_obj() + + with httpx.Client(transport=_transcription_transport()) as http_client: + response = OpenAIAudioTranscription().audio_transcriptions( + client=OpenAI(api_key="transport-only", http_client=http_client), + atranscription=False, + **_call_kwargs(logging_obj), + ) + + _assert_headers_recorded(response, logging_obj) + + +@pytest.mark.asyncio +async def test_async_audio_transcriptions_records_provider_response_headers(): + logging_obj = _logging_obj() + + async with httpx.AsyncClient(transport=_transcription_transport()) as http_client: + response = await OpenAIAudioTranscription().audio_transcriptions( + client=AsyncOpenAI(api_key="transport-only", http_client=http_client), + atranscription=True, + **_call_kwargs(logging_obj), + ) + + _assert_headers_recorded(response, logging_obj) diff --git a/tests/test_litellm/test_non_chat_routes_open_llm_spans.py b/tests/test_litellm/test_non_chat_routes_open_llm_spans.py index d62959ccd43..02c95c4bb2a 100644 --- a/tests/test_litellm/test_non_chat_routes_open_llm_spans.py +++ b/tests/test_litellm/test_non_chat_routes_open_llm_spans.py @@ -46,9 +46,9 @@ class _FakeSpeech: )() -class _FakeImages: +class _FakeRawImages: async def generate(self, **kwargs: Any) -> Any: - return type( + parsed: Final = type( "_Images", (), { @@ -58,6 +58,16 @@ class _FakeImages: } }, )() + return type( + "_RawImages", + (), + {"parse": lambda self: parsed, "headers": httpx.Headers({"x-request-id": "req-image"})}, + )() + + +class _FakeImages: + def __init__(self) -> None: + self.with_raw_response = _FakeRawImages() class _FakeModerations: diff --git a/tests/unit/llms/openai/image_generation/test_openai_image_generation_extra_headers.py b/tests/unit/llms/openai/image_generation/test_openai_image_generation_extra_headers.py index 55ef74abd7b..11df07f5fea 100644 --- a/tests/unit/llms/openai/image_generation/test_openai_image_generation_extra_headers.py +++ b/tests/unit/llms/openai/image_generation/test_openai_image_generation_extra_headers.py @@ -12,6 +12,14 @@ import pytest from litellm.llms.openai.openai import OpenAIChatCompletion +from litellm.types.utils import ImageResponse + + +def _raw_image_response(mock_image_data): + raw_response = MagicMock() + raw_response.parse.return_value = mock_image_data + raw_response.headers = {"x-request-id": "req-image"} + return raw_response @pytest.fixture @@ -41,7 +49,7 @@ class TestImageGenerationExtraHeaders: } mock_openai_client = MagicMock() - mock_openai_client.images.generate.return_value = mock_image_data + mock_openai_client.images.with_raw_response.generate.return_value = _raw_image_response(mock_image_data) mock_openai_client.api_key = "test-key" mock_openai_client._base_url._uri_reference = "https://api.openai.com" @@ -58,7 +66,7 @@ class TestImageGenerationExtraHeaders: client=mock_openai_client, ) - _, kwargs = mock_openai_client.images.generate.call_args + _, kwargs = mock_openai_client.images.with_raw_response.generate.call_args assert kwargs.get("extra_headers") == test_headers def test_sync_image_generation_without_headers( @@ -72,7 +80,7 @@ class TestImageGenerationExtraHeaders: } mock_openai_client = MagicMock() - mock_openai_client.images.generate.return_value = mock_image_data + mock_openai_client.images.with_raw_response.generate.return_value = _raw_image_response(mock_image_data) mock_openai_client.api_key = "test-key" mock_openai_client._base_url._uri_reference = "https://api.openai.com" @@ -86,7 +94,7 @@ class TestImageGenerationExtraHeaders: client=mock_openai_client, ) - _, kwargs = mock_openai_client.images.generate.call_args + _, kwargs = mock_openai_client.images.with_raw_response.generate.call_args assert "extra_headers" not in kwargs @pytest.mark.asyncio @@ -101,7 +109,9 @@ class TestImageGenerationExtraHeaders: } mock_openai_client = MagicMock() - mock_openai_client.images.generate = AsyncMock(return_value=mock_image_data) + mock_openai_client.images.with_raw_response.generate = AsyncMock( + return_value=_raw_image_response(mock_image_data) + ) mock_openai_client.api_key = "test-key" test_headers = {"cf-aig-authorization": "Bearer custom-token"} @@ -109,7 +119,7 @@ class TestImageGenerationExtraHeaders: await openai_chat_completions.aimage_generation( prompt="A white cat", data={"model": "dall-e-3", "prompt": "A white cat"}, - model_response=MagicMock(), + model_response=ImageResponse(), timeout=60.0, logging_obj=mock_logging_obj, api_key="test-key", @@ -117,7 +127,7 @@ class TestImageGenerationExtraHeaders: client=mock_openai_client, ) - _, kwargs = mock_openai_client.images.generate.call_args + _, kwargs = mock_openai_client.images.with_raw_response.generate.call_args assert kwargs.get("extra_headers") == test_headers @pytest.mark.asyncio @@ -132,20 +142,22 @@ class TestImageGenerationExtraHeaders: } mock_openai_client = MagicMock() - mock_openai_client.images.generate = AsyncMock(return_value=mock_image_data) + mock_openai_client.images.with_raw_response.generate = AsyncMock( + return_value=_raw_image_response(mock_image_data) + ) mock_openai_client.api_key = "test-key" await openai_chat_completions.aimage_generation( prompt="A white cat", data={"model": "dall-e-3", "prompt": "A white cat"}, - model_response=MagicMock(), + model_response=ImageResponse(), timeout=60.0, logging_obj=mock_logging_obj, api_key="test-key", client=mock_openai_client, ) - _, kwargs = mock_openai_client.images.generate.call_args + _, kwargs = mock_openai_client.images.with_raw_response.generate.call_args assert "extra_headers" not in kwargs @pytest.mark.parametrize("is_async", [False, True]) @@ -169,11 +181,13 @@ class TestImageGenerationExtraHeaders: test_headers = {"cf-aig-authorization": "Bearer custom-token"} if is_async: - mock_openai_client.images.generate = AsyncMock(return_value=mock_image_data) + mock_openai_client.images.with_raw_response.generate = AsyncMock( + return_value=_raw_image_response(mock_image_data) + ) await openai_chat_completions.aimage_generation( prompt="A white cat", data={"model": "dall-e-3", "prompt": "A white cat"}, - model_response=MagicMock(), + model_response=ImageResponse(), timeout=60.0, logging_obj=mock_logging_obj, api_key="test-key", @@ -181,7 +195,7 @@ class TestImageGenerationExtraHeaders: client=mock_openai_client, ) else: - mock_openai_client.images.generate.return_value = mock_image_data + mock_openai_client.images.with_raw_response.generate.return_value = _raw_image_response(mock_image_data) openai_chat_completions.image_generation( model="dall-e-3", prompt="A white cat", @@ -197,7 +211,7 @@ class TestImageGenerationExtraHeaders: "complete_input_dict" ] assert "extra_headers" not in logged_body - _, kwargs = mock_openai_client.images.generate.call_args + _, kwargs = mock_openai_client.images.with_raw_response.generate.call_args assert kwargs.get("extra_headers") == test_headers def test_sync_image_generation_forwards_headers_to_async( @@ -242,7 +256,9 @@ class TestImageGenerationEntryPointHeaders: } mock_openai_client = MagicMock() - mock_openai_client.images.generate = AsyncMock(return_value=mock_image_data) + mock_openai_client.images.with_raw_response.generate = AsyncMock( + return_value=_raw_image_response(mock_image_data) + ) mock_openai_client.api_key = "test-key" mock_openai_client._base_url._uri_reference = "https://api.openai.com" @@ -256,6 +272,6 @@ class TestImageGenerationEntryPointHeaders: api_key="test-key", ) - mock_openai_client.images.generate.assert_called_once() - _, kwargs = mock_openai_client.images.generate.call_args + mock_openai_client.images.with_raw_response.generate.assert_called_once() + _, kwargs = mock_openai_client.images.with_raw_response.generate.call_args assert kwargs.get("extra_headers") == test_headers diff --git a/tests/unit/llms/vertex_ai/text_to_speech/test_transformation.py b/tests/unit/llms/vertex_ai/text_to_speech/test_transformation.py index b5eec42b569..ee7bdebe745 100644 --- a/tests/unit/llms/vertex_ai/text_to_speech/test_transformation.py +++ b/tests/unit/llms/vertex_ai/text_to_speech/test_transformation.py @@ -526,6 +526,7 @@ class TestVertexAILyriaTextToSpeechConfig: ): mock_response = Mock(spec=httpx.Response) mock_response.status_code = 200 + mock_response.headers = {"content-type": "application/json"} mock_response.json.return_value = response_json with ( patch.object( # test-quality-ok: litellm.speech has no seam for Vertex token minting From 5a8ec1378611e6620f5cfe291f0e971ca403293e Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 24 Sep 2026 15:04:03 -0500 Subject: [PATCH 072/120] feat(usage): search keys beyond the top-N usage subset (#42827) --- litellm/proxy/_types.py | 2 + .../internal_user_endpoints.py | 153 ++++++++++++++-- .../internal_user_endpoints.py | 12 +- .../endpointaudit/coverage_allowlist.txt | 1 + .../proxy/auth/test_route_checks.py | 1 + .../test_internal_user_endpoints.py | 166 ++++++++++++++++++ .../_components/components/UsagePageView.tsx | 16 +- .../components/KeyActivityPanel.test.tsx | 64 +++++++ .../UsagePage/components/KeyActivityPanel.tsx | 73 +++++++- .../src/components/networking.tsx | 30 ++++ ui/litellm-dashboard/src/lib/http/schema.d.ts | 65 +++++++ 11 files changed, 560 insertions(+), 23 deletions(-) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 4affa55f903..51a3eb03067 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -699,6 +699,7 @@ class LiteLLMRoutes(enum.Enum): "/user/list", "/user/daily/activity", "/user/daily/activity/aggregated", + "/user/daily/activity/aggregated/search", # team "/team/new", "/team/update", @@ -901,6 +902,7 @@ class LiteLLMRoutes(enum.Enum): "/model/delete", "/user/daily/activity", "/user/daily/activity/aggregated", + "/user/daily/activity/aggregated/search", # Endpoint restricts results to organizations the caller is ORG_ADMIN # of; a caller who administers none gets an empty result set. "/organization/daily/activity", diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index 59d8dd821d8..7b25348aa53 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -28,6 +28,7 @@ from typing_extensions import ReadOnly, TypedDict import litellm from litellm._logging import verbose_proxy_logger from litellm._uuid import uuid +from litellm.constants import USAGE_TOP_API_KEYS_LIMIT from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.proxy._types import * from litellm.proxy.auth.auth_checks import ( @@ -87,11 +88,13 @@ from litellm.repositories.verification_token_repository import ( VerificationTokenRepository, ) from litellm.types.proxy.management_endpoints.common_daily_activity import ( + DailySpendMetadata, SpendAnalyticsPaginatedResponse, ) from litellm.types.proxy.management_endpoints.internal_user_endpoints import ( BulkUpdateUserRequest, BulkUpdateUserResponse, + KeyActivitySearchWhere, UserListResponse, UserSearchWhere, UserUpdateResult, @@ -2991,6 +2994,27 @@ async def get_user_daily_activity( ) +def _resolve_user_daily_activity_entity_id( + user_api_key_dict: UserAPIKeyAuth, + user_id: str | None, +) -> str | None: + is_admin: Final = _user_has_admin_view(user_api_key_dict) + + if is_admin: + return user_id + + caller_user_id: Final = require_caller_user_id_for_non_admin(user_api_key_dict) + effective_user_id: Final = user_id if user_id is not None else caller_user_id + if effective_user_id != caller_user_id: + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail={ # mutable-ok: FastAPI detail payload shape + "error": "Non-admin users can only view their own spend data." + }, + ) + return effective_user_id + + @router.get( "/user/daily/activity/aggregated", tags=["Budget & Spend Tracking", "Internal User management"], @@ -3057,20 +3081,7 @@ async def get_user_daily_activity_aggregated( ) try: - is_admin: Final = _user_has_admin_view(user_api_key_dict) - - if is_admin: - entity_id = user_id # None means global view, otherwise filter by user - else: - caller_user_id: Final = require_caller_user_id_for_non_admin(user_api_key_dict) - if user_id is None: - user_id = caller_user_id - if user_id != caller_user_id: - raise HTTPException( - status_code=status.HTTP_403_FORBIDDEN, - detail={"error": "Non-admin users can only view their own spend data."}, - ) - entity_id = user_id + entity_id: Final = _resolve_user_daily_activity_entity_id(user_api_key_dict, user_id) return await get_daily_activity_aggregated( prisma_client=prisma_client, @@ -3094,3 +3105,117 @@ async def get_user_daily_activity_aggregated( status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail={"error": f"Failed to fetch analytics: {e}"}, ) + + +@router.get( + "/user/daily/activity/aggregated/search", + tags=["Budget & Spend Tracking", "Internal User management"], # mutable-ok: FastAPI route tags shape + dependencies=[Depends(user_api_key_auth)], # mutable-ok: FastAPI route dependencies shape + response_model=SpendAnalyticsPaginatedResponse, +) +@management_endpoint_wrapper +async def search_user_daily_activity_keys( + search: str = fastapi.Query( + ..., + min_length=1, + description="Matches keys whose hash equals the value, or whose key alias or user ID contains it (case-insensitive)", + ), + start_date: str | None = fastapi.Query( + default=None, + description="Start date in YYYY-MM-DD format", + ), + end_date: str | None = fastapi.Query( + default=None, + description="End date in YYYY-MM-DD format", + ), + user_id: str | None = fastapi.Query( + default=None, + description="Filter by specific user ID. Admins can filter by any user or omit for global view. Non-admins must provide their own user_id.", + ), + timezone: int | None = fastapi.Query( + default=None, + description="Timezone offset in minutes from UTC (e.g., 480 for PST). " + "Matches JavaScript's Date.getTimezoneOffset() convention.", + ), + include_current_utc_day: bool = fastapi.Query( + default=False, + description="When the range ends on the caller's current local day, extend it to " + "today's UTC bucket so spend written after the caller's local midnight (in UTC " + "terms) is included. Requires the timezone parameter. Historical ranges are " + "never extended.", + ), + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), # noqa: B008 # FastAPI dependency injection +) -> SpendAnalyticsPaginatedResponse: + """ + Search verification tokens by exact token hash or by a case-insensitive substring of + the key alias or owning user ID, then return the aggregated daily activity for the + matches. Lets the Usage page surface keys that fell outside the top-spend subset + the aggregated endpoint loads. + """ + from litellm.proxy.proxy_server import prisma_client + + if prisma_client is None: + raise HTTPException( + status_code=500, + detail={ # mutable-ok: FastAPI detail payload shape + "error": CommonProxyErrors.db_not_connected_error.value + }, + ) + + if start_date is None or end_date is None: + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail={"error": "Please provide start_date and end_date"}, # mutable-ok: FastAPI detail payload shape + ) + + try: + entity_id: Final = _resolve_user_daily_activity_entity_id(user_api_key_dict, user_id) + + search_or: Final = ( + {"token": search}, # mutable-ok: prisma serializes where clauses, keep plain dicts + {"key_alias": {"contains": search, "mode": "insensitive"}}, # mutable-ok: prisma where clause leaf + {"user_id": {"contains": search, "mode": "insensitive"}}, # mutable-ok: prisma where clause leaf + ) + where: Final[KeyActivitySearchWhere] = ( + {"OR": search_or} # mutable-ok: prisma where clause root + if entity_id is None + else {"user_id": entity_id, "OR": search_or} # mutable-ok: prisma where clause root + ) + matched_keys: Final = await VerificationTokenRepository(prisma_client).table.find_many( + where=where, + take=USAGE_TOP_API_KEYS_LIMIT, + order={"spend": "desc"}, # mutable-ok: prisma serializes order, keep it a plain dict + ) + tokens: Final = [key.token for key in matched_keys] # mutable-ok: api_key filter union expects a list + + if not tokens: + return SpendAnalyticsPaginatedResponse( + results=[], # mutable-ok: response model field shape + metadata=DailySpendMetadata( + api_key_limit=USAGE_TOP_API_KEYS_LIMIT, + total_api_keys=0, + ), + ) + + return await get_daily_activity_aggregated( + prisma_client=prisma_client, + table_name="litellm_dailyuserspend", + entity_id_field="user_id", + entity_id=entity_id, + entity_metadata_field=None, + start_date=start_date, + end_date=end_date, + model=None, + api_key=tokens, + timezone_offset_minutes=timezone, + include_current_utc_day=include_current_utc_day, + ) + + except HTTPException: + raise + except Exception as e: + verbose_proxy_logger.exception("/user/daily/activity/aggregated/search: Exception occured - %s", e) + raise HTTPException( + status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, + detail={"error": f"Failed to fetch analytics: {e}"}, # mutable-ok: FastAPI detail payload shape + ) diff --git a/litellm/types/proxy/management_endpoints/internal_user_endpoints.py b/litellm/types/proxy/management_endpoints/internal_user_endpoints.py index 43e3899d523..05cbd4507a2 100644 --- a/litellm/types/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/types/proxy/management_endpoints/internal_user_endpoints.py @@ -2,7 +2,7 @@ from collections.abc import Mapping, Sequence from typing import Any, Final, Literal from pydantic import BaseModel, ConfigDict, Field, field_validator -from typing_extensions import ReadOnly, TypedDict +from typing_extensions import NotRequired, ReadOnly, TypedDict from litellm.proxy._types import ( LiteLLM_UserTableWithKeyCount, @@ -28,6 +28,16 @@ class UserSearchWhere(TypedDict): OR: ReadOnly[tuple[Mapping[Literal["user_id", "user_email"], InsensitiveContains], ...]] +class KeyActivitySearchWhere(TypedDict): + """Prisma filter behind `/user/daily/activity/aggregated/search`: exact token hash, or key alias + or user id containing the term, case-insensitive.""" + + user_id: NotRequired[ReadOnly[str]] + OR: ReadOnly[ + tuple[Mapping[Literal["token"], str] | Mapping[Literal["key_alias", "user_id"], InsensitiveContains], ...] + ] + + class UserListResponse(BaseModel): """ Response model for the user list endpoint diff --git a/terraform/provider/tools/endpointaudit/coverage_allowlist.txt b/terraform/provider/tools/endpointaudit/coverage_allowlist.txt index f8277c83a64..d0f9c31ecaf 100644 --- a/terraform/provider/tools/endpointaudit/coverage_allowlist.txt +++ b/terraform/provider/tools/endpointaudit/coverage_allowlist.txt @@ -32,6 +32,7 @@ GET /team/spend/by_user GET /team/spend/report GET /user/daily/activity GET /user/daily/activity/aggregated +GET /user/daily/activity/aggregated/search GET /user/spend/report # Admin UI helper endpoints; serve UI forms and caller-scoped views, not desired state diff --git a/tests/test_litellm/proxy/auth/test_route_checks.py b/tests/test_litellm/proxy/auth/test_route_checks.py index e5179387f82..571e066e947 100644 --- a/tests/test_litellm/proxy/auth/test_route_checks.py +++ b/tests/test_litellm/proxy/auth/test_route_checks.py @@ -3517,6 +3517,7 @@ def test_internal_user_still_blocked_from_another_users_info(): [ "/user/daily/activity", "/user/daily/activity/aggregated", + "/user/daily/activity/aggregated/search", ], ) @pytest.mark.parametrize( diff --git a/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py index c663e63414c..8260aec9326 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py @@ -2659,6 +2659,172 @@ async def test_get_user_daily_activity_aggregated_non_admin_cannot_view_other_us assert mock_get_daily_agg.call_args.kwargs["entity_id"] == "regular-user-123" +@pytest.mark.asyncio +async def test_search_user_daily_activity_keys_passes_matched_tokens_to_aggregation(monkeypatch): + """The search endpoint resolves matching verification tokens by hash, alias, or + user id, then aggregates daily spend for exactly those tokens. This is what lets + the Usage page find keys outside the top-spend subset the aggregated endpoint caps.""" + from types import SimpleNamespace + from unittest.mock import AsyncMock, MagicMock + + from litellm.constants import USAGE_TOP_API_KEYS_LIMIT + from litellm.proxy.management_endpoints.internal_user_endpoints import ( + search_user_daily_activity_keys, + ) + + mock_prisma_client = MagicMock() + mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock( + return_value=[SimpleNamespace(token="tok-a"), SimpleNamespace(token="tok-b")] + ) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + + mock_response = MagicMock() + mock_get_daily_agg = AsyncMock(return_value=mock_response) + monkeypatch.setattr( + "litellm.proxy.management_endpoints.internal_user_endpoints.get_daily_activity_aggregated", + mock_get_daily_agg, + ) + + admin_key_dict = UserAPIKeyAuth( + user_id="admin-user-001", + user_role=LitellmUserRoles.PROXY_ADMIN, + ) + + result = await search_user_daily_activity_keys( + search="gamma", + start_date="2025-02-01", + end_date="2025-02-28", + user_id=None, + timezone=480, + include_current_utc_day=False, + user_api_key_dict=admin_key_dict, + ) + + assert result is mock_response + + find_many_kwargs = mock_prisma_client.db.litellm_verificationtoken.find_many.call_args.kwargs + assert find_many_kwargs["take"] == USAGE_TOP_API_KEYS_LIMIT + assert find_many_kwargs["where"]["OR"] == ( + {"token": "gamma"}, + {"key_alias": {"contains": "gamma", "mode": "insensitive"}}, + {"user_id": {"contains": "gamma", "mode": "insensitive"}}, + ) + assert "user_id" not in find_many_kwargs["where"] + + mock_get_daily_agg.assert_called_once_with( + prisma_client=mock_prisma_client, + table_name="litellm_dailyuserspend", + entity_id_field="user_id", + entity_id=None, + entity_metadata_field=None, + start_date="2025-02-01", + end_date="2025-02-28", + model=None, + api_key=["tok-a", "tok-b"], + timezone_offset_minutes=480, + include_current_utc_day=False, + ) + + +@pytest.mark.asyncio +async def test_search_user_daily_activity_keys_no_match_returns_empty_without_aggregating(monkeypatch): + from unittest.mock import AsyncMock, MagicMock + + from litellm.constants import USAGE_TOP_API_KEYS_LIMIT + from litellm.proxy.management_endpoints.internal_user_endpoints import ( + search_user_daily_activity_keys, + ) + + mock_prisma_client = MagicMock() + mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[]) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + + mock_get_daily_agg = AsyncMock() + monkeypatch.setattr( + "litellm.proxy.management_endpoints.internal_user_endpoints.get_daily_activity_aggregated", + mock_get_daily_agg, + ) + + admin_key_dict = UserAPIKeyAuth( + user_id="admin-user-001", + user_role=LitellmUserRoles.PROXY_ADMIN, + ) + + result = await search_user_daily_activity_keys( + search="nothing-matches", + start_date="2025-02-01", + end_date="2025-02-28", + user_id=None, + timezone=None, + include_current_utc_day=False, + user_api_key_dict=admin_key_dict, + ) + + assert result.results == [] + assert result.metadata.api_key_limit == USAGE_TOP_API_KEYS_LIMIT + assert result.metadata.total_api_keys == 0 + mock_get_daily_agg.assert_not_called() + + +@pytest.mark.asyncio +async def test_search_user_daily_activity_keys_non_admin_scoped_to_caller(monkeypatch): + """Same scoping contract as the aggregated route: a non-admin with no user_id + is scoped to their own rows, and any other user_id is a 403.""" + from types import SimpleNamespace + from unittest.mock import AsyncMock, MagicMock + + from fastapi import HTTPException + + from litellm.proxy.management_endpoints.internal_user_endpoints import ( + search_user_daily_activity_keys, + ) + + mock_prisma_client = MagicMock() + mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[SimpleNamespace(token="tok-a")]) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + + non_admin_key_dict = UserAPIKeyAuth( + user_id="user-1", + user_role=LitellmUserRoles.INTERNAL_USER, + ) + + mock_response = MagicMock() + mock_get_daily_agg = AsyncMock(return_value=mock_response) + monkeypatch.setattr( + "litellm.proxy.management_endpoints.internal_user_endpoints.get_daily_activity_aggregated", + mock_get_daily_agg, + ) + + result = await search_user_daily_activity_keys( + search="gamma", + start_date="2025-02-01", + end_date="2025-02-28", + user_id=None, + timezone=None, + include_current_utc_day=False, + user_api_key_dict=non_admin_key_dict, + ) + + assert result is mock_response + assert mock_get_daily_agg.call_args.kwargs["entity_id"] == "user-1" + find_many_kwargs = mock_prisma_client.db.litellm_verificationtoken.find_many.call_args.kwargs + assert find_many_kwargs["where"]["user_id"] == "user-1" + + with pytest.raises(HTTPException) as exc_info: + await search_user_daily_activity_keys( + search="gamma", + start_date="2025-02-01", + end_date="2025-02-28", + user_id="user-2", + timezone=None, + include_current_utc_day=False, + user_api_key_dict=non_admin_key_dict, + ) + + assert exc_info.value.status_code == 403 + assert "Non-admin users can only view their own spend data" in str(exc_info.value.detail) + + @pytest.mark.asyncio async def test_delete_user_cleans_up_created_by_invitation_links(mocker): """ diff --git a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsagePageView.tsx b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsagePageView.tsx index 228d8acf146..e0171d97423 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsagePageView.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsagePageView.tsx @@ -39,6 +39,7 @@ import { tagListCall, userDailyActivityAggregatedCall, userDailyActivityCall, + userDailyActivityKeySearchCall, } from "@/components/networking"; import AdvancedDatePicker from "@/components/shared/advanced_date_picker"; import { ChartLoader } from "@/components/shared/chart_loader"; @@ -437,6 +438,15 @@ const UsagePage: React.FC = ({ teams, organizations }) => { [userSpendData, modelViewType, teams], ); const keyMetrics = useMemo(() => processActivityData(userSpendData, "api_keys", teams), [userSpendData, teams]); + const searchKeys = useCallback( + (q: string) => { + if (!accessToken || !startTime || !endTime) return Promise.resolve({}); + return userDailyActivityKeySearchCall(accessToken, startTime, endTime, q, effectiveUserId).then((data) => + processActivityData(data, "api_keys", teams), + ); + }, + [accessToken, startTime, endTime, effectiveUserId, teams], + ); const mcpServerMetrics = useMemo( () => processActivityData(userSpendData, "mcp_servers", teams), [userSpendData, teams], @@ -865,7 +875,11 @@ const UsagePage: React.FC = ({ teams, organizations }) => { - + diff --git a/ui/litellm-dashboard/src/components/UsagePage/components/KeyActivityPanel.test.tsx b/ui/litellm-dashboard/src/components/UsagePage/components/KeyActivityPanel.test.tsx index 830139143e9..f8a7d07633b 100644 --- a/ui/litellm-dashboard/src/components/UsagePage/components/KeyActivityPanel.test.tsx +++ b/ui/litellm-dashboard/src/components/UsagePage/components/KeyActivityPanel.test.tsx @@ -78,4 +78,68 @@ describe("KeyActivityPanel", () => { render(); expect(screen.queryByRole("note")).not.toBeInTheDocument(); }); + + it("finds keys outside the loaded top-spend subset via server search", async () => { + const searchKeys = vi + .fn<(query: string) => Promise>>() + .mockResolvedValue({ "hash-gamma": activity("gamma-low-key", "gamma@example.com", "user-gamma") }); + render( + , + ); + + fireEvent.change(screen.getByLabelText("Search keys"), { target: { value: "gamma" } }); + + expect(await screen.findByText("hash-gamma")).toBeInTheDocument(); + expect(searchKeys).toHaveBeenCalledWith("gamma"); + expect(screen.getByText("Showing 1 of 3 keys")).toBeInTheDocument(); + }); + + it("never calls the server search when every key is already loaded", async () => { + const searchKeys = vi + .fn<(query: string) => Promise>>() + .mockResolvedValue({ "hash-gamma": activity("gamma-low-key", "gamma@example.com", "user-gamma") }); + render(); + + fireEvent.change(screen.getByLabelText("Search keys"), { target: { value: "gamma" } }); + + expect(await screen.findByText('No keys match "gamma" in this date range')).toBeInTheDocument(); + await new Promise((resolve) => setTimeout(resolve, 400)); + expect(searchKeys).not.toHaveBeenCalled(); + }); + + it("drops stale server results as soon as the search callback is rebuilt", async () => { + const searchKeysA = vi + .fn<(query: string) => Promise>>() + .mockResolvedValue({ "hash-gamma": activity("gamma-low-key", "gamma@example.com", "user-gamma") }); + const searchKeysB = vi + .fn<(query: string) => Promise>>() + .mockReturnValue(new Promise(() => {})); + const { rerender } = render( + , + ); + + fireEvent.change(screen.getByLabelText("Search keys"), { target: { value: "gamma" } }); + expect(await screen.findByText("hash-gamma")).toBeInTheDocument(); + + rerender( + , + ); + + expect(screen.getByRole("status")).toHaveTextContent("Searching all keys"); + expect(screen.queryByText("hash-gamma")).not.toBeInTheDocument(); + }); + + it("reports a failed server search but keeps the local matches", async () => { + const searchKeys = vi + .fn<(query: string) => Promise>>() + .mockRejectedValue(new Error("boom")); + render( + , + ); + + fireEvent.change(screen.getByLabelText("Search keys"), { target: { value: "alice" } }); + + expect(await screen.findByRole("alert")).toHaveTextContent("Key search failed"); + expect(screen.getByTestId("rendered-keys")).toHaveTextContent("hash-alice"); + }); }); diff --git a/ui/litellm-dashboard/src/components/UsagePage/components/KeyActivityPanel.tsx b/ui/litellm-dashboard/src/components/UsagePage/components/KeyActivityPanel.tsx index 8b2141f8528..3467ba61b22 100644 --- a/ui/litellm-dashboard/src/components/UsagePage/components/KeyActivityPanel.tsx +++ b/ui/litellm-dashboard/src/components/UsagePage/components/KeyActivityPanel.tsx @@ -1,5 +1,5 @@ import { Search, X } from "lucide-react"; -import React, { useMemo, useState } from "react"; +import React, { useEffect, useMemo, useState } from "react"; import { ActivityMetrics } from "@/components/activity_metrics"; import type { ApiKeyTruncation } from "@/components/EntityUsageExport/exportBlockedReason"; @@ -12,18 +12,67 @@ interface KeyActivityPanelProps { keyMetrics: Record; hidePromptCachingMetrics?: boolean; apiKeyTruncation?: ApiKeyTruncation; + searchKeys?: SearchKeys; } +type SearchKeys = (query: string) => Promise>; + +type RemoteSearch = + | { status: "idle" } + | { status: "loading"; query: string; searchKeys: SearchKeys } + | { status: "done"; query: string; searchKeys: SearchKeys; keys: Record } + | { status: "error"; query: string; searchKeys: SearchKeys }; + +const REMOTE_SEARCH_DEBOUNCE_MS = 300; + const KeyActivityPanel: React.FC = ({ keyMetrics, hidePromptCachingMetrics = false, apiKeyTruncation, + searchKeys, }) => { const [query, setQuery] = useState(""); + const [remote, setRemote] = useState({ status: "idle" }); const filtered = useMemo(() => filterKeyActivity(keyMetrics, query), [keyMetrics, query]); + const trimmedQuery = query.trim(); + const remoteEnabled = searchKeys !== undefined && apiKeyTruncation !== undefined && trimmedQuery !== ""; + + useEffect(() => { + if (!remoteEnabled) return; + let cancelled = false; + const timer = setTimeout(() => { + setRemote({ status: "loading", query: trimmedQuery, searchKeys }); + searchKeys(trimmedQuery) + .then((keys) => { + if (!cancelled) setRemote({ status: "done", query: trimmedQuery, searchKeys, keys }); + }) + .catch(() => { + if (!cancelled) setRemote({ status: "error", query: trimmedQuery, searchKeys }); + }); + }, REMOTE_SEARCH_DEBOUNCE_MS); + return () => { + cancelled = true; + clearTimeout(timer); + }; + }, [remoteEnabled, trimmedQuery, searchKeys]); + + const remoteMatchesSearch = + "searchKeys" in remote && remote.searchKeys === searchKeys && remote.query === trimmedQuery; + const remoteCurrent = remoteEnabled && remoteMatchesSearch; + const remoteLoading = remoteEnabled && (remote.status === "loading" || !remoteCurrent); + const remoteFailed = remoteCurrent && remote.status === "error"; + + const extraRemoteKeys = useMemo(() => { + const remoteKeys = remoteCurrent && remote.status === "done" ? remote.keys : {}; + return Object.fromEntries(Object.entries(remoteKeys).filter(([hash]) => !(hash in keyMetrics))); + }, [remoteCurrent, remote, keyMetrics]); + const displayed = useMemo(() => ({ ...extraRemoteKeys, ...filtered }), [extraRemoteKeys, filtered]); + const totalKeys = Object.keys(keyMetrics).length; - const shownKeys = Object.keys(filtered).length; - const isFiltering = query.trim() !== ""; + const shownKeys = Object.keys(displayed).length; + const totalShown = totalKeys + Object.keys(extraRemoteKeys).length; + const isFiltering = trimmedQuery !== ""; + const noMatches = isFiltering && !remoteLoading && totalKeys > 0 && shownKeys === 0; return (
@@ -47,8 +96,18 @@ const KeyActivityPanel: React.FC = ({ )} - Showing {shownKeys.toLocaleString()} of {totalKeys.toLocaleString()} keys + Showing {shownKeys.toLocaleString()} of {totalShown.toLocaleString()} keys + {remoteLoading && ( + + Searching all keys... + + )} + {remoteFailed && ( + + Key search failed + + )} {apiKeyTruncation !== undefined && ( Only the {apiKeyTruncation.limit.toLocaleString()} highest-spend keys of{" "} @@ -56,12 +115,12 @@ const KeyActivityPanel: React.FC = ({ )}
- {isFiltering && totalKeys > 0 && shownKeys === 0 ? ( + {noMatches ? (

- No keys match "{query.trim()}" in this date range + No keys match "{trimmedQuery}" in this date range

) : ( - + )}
); diff --git a/ui/litellm-dashboard/src/components/networking.tsx b/ui/litellm-dashboard/src/components/networking.tsx index c2f4fa80634..e6923a26d4c 100644 --- a/ui/litellm-dashboard/src/components/networking.tsx +++ b/ui/litellm-dashboard/src/components/networking.tsx @@ -2556,6 +2556,36 @@ export const userDailyActivityAggregatedCall = async ( } }; +export const userDailyActivityKeySearchCall = async ( + accessToken: string, + startTime: Date, + endTime: Date, + ...options: [search: string, userId?: string | null] +) => { + const [search, userId = null] = options; + try { + const formatDate = (date: Date) => { + const year = date.getFullYear(); + const month = String(date.getMonth() + 1).padStart(2, "0"); + const day = String(date.getDate()).padStart(2, "0"); + return `${year}-${month}-${day}`; + }; + return await apiClient.get(`/user/daily/activity/aggregated/search`, { + accessToken, + query: { + start_date: formatDate(startTime), + end_date: formatDate(endTime), + timezone: new Date().getTimezoneOffset().toString(), + search, + user_id: userId || undefined, + }, + }); + } catch (error) { + console.error("Failed to search user daily activity keys:", error); + throw error; + } +}; + export const gatewayDailyActivityCall = async (accessToken: string, startTime: Date, endTime: Date) => { /** * Get gateway request counts (SGR) recorded by the proxy middleware. diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 052cc693f8b..79d008c1b48 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -17326,6 +17326,29 @@ export interface paths { patch?: never; trace?: never; }; + "/user/daily/activity/aggregated/search": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + /** + * Search User Daily Activity Keys + * @description Search verification tokens by exact token hash or by a case-insensitive substring of + * the key alias or owning user ID, then return the aggregated daily activity for the + * matches. Lets the Usage page surface keys that fell outside the top-spend subset + * the aggregated endpoint loads. + */ + get: operations["search_user_daily_activity_keys_user_daily_activity_aggregated_search_get"]; + put?: never; + post?: never; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; "/user/delete": { parameters: { query?: never; @@ -68731,6 +68754,48 @@ export interface operations { }; }; }; + search_user_daily_activity_keys_user_daily_activity_aggregated_search_get: { + parameters: { + query: { + /** @description Matches keys whose hash equals the value, or whose key alias or user ID contains it (case-insensitive) */ + search: string; + /** @description Start date in YYYY-MM-DD format */ + start_date?: string | null; + /** @description End date in YYYY-MM-DD format */ + end_date?: string | null; + /** @description Filter by specific user ID. Admins can filter by any user or omit for global view. Non-admins must provide their own user_id. */ + user_id?: string | null; + /** @description Timezone offset in minutes from UTC (e.g., 480 for PST). Matches JavaScript's Date.getTimezoneOffset() convention. */ + timezone?: number | null; + /** @description When the range ends on the caller's current local day, extend it to today's UTC bucket so spend written after the caller's local midnight (in UTC terms) is included. Requires the timezone parameter. Historical ranges are never extended. */ + include_current_utc_day?: boolean; + }; + header?: never; + path?: never; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["SpendAnalyticsPaginatedResponse"]; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; delete_user_user_delete_post: { parameters: { query?: never; From c2eb549ee685b442b74bc7c1cf9c4c1a7698a5ec Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 24 Sep 2026 15:04:12 -0500 Subject: [PATCH 073/120] feat(usage): search team keys beyond the top-N in the Team usage view (#42857) --- litellm/proxy/_types.py | 5 + .../management_endpoints/team_endpoints.py | 109 +++++++++ .../management_endpoints/team_endpoints.py | 24 ++ .../endpointaudit/coverage_allowlist.txt | 1 + .../test_team_daily_activity_key_search.py | 128 +++++++++++ .../management/test_team_daily_activity.py | 15 +- .../proxy/auth/test_route_checks.py | 49 ++++ .../test_team_endpoints.py | 212 ++++++++++++++++++ .../components/EntityUsage/EntityUsage.tsx | 14 +- .../src/components/networking.tsx | 25 +++ ui/litellm-dashboard/src/lib/http/schema.d.ts | 60 ++++- 11 files changed, 635 insertions(+), 7 deletions(-) create mode 100644 tests/integration/spend/test_team_daily_activity_key_search.py diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 51a3eb03067..73e3e0e6ee0 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -307,6 +307,7 @@ class KeyManagementRoutes(str, enum.Enum): # team usage routes TEAM_DAILY_ACTIVITY = "/team/daily/activity" TEAM_DAILY_ACTIVITY_AGGREGATED = "/team/daily/activity/aggregated" + TEAM_DAILY_ACTIVITY_AGGREGATED_SEARCH = "/team/daily/activity/aggregated/search" # team spend-log viewing SPEND_LOGS = "/spend/logs" @@ -673,6 +674,7 @@ class LiteLLMRoutes(enum.Enum): KeyManagementRoutes.TEAM_KEY_BULK_UPDATE.value, KeyManagementRoutes.TEAM_DAILY_ACTIVITY.value, KeyManagementRoutes.TEAM_DAILY_ACTIVITY_AGGREGATED.value, + KeyManagementRoutes.TEAM_DAILY_ACTIVITY_AGGREGATED_SEARCH.value, KeyManagementRoutes.SPEND_LOGS.value, KeyManagementRoutes.SPEND_LOGS_V2.value, KeyManagementRoutes.KEY_RESET_SPEND.value, @@ -717,6 +719,7 @@ class LiteLLMRoutes(enum.Enum): "/team/permissions_bulk_update", "/team/daily/activity", "/team/daily/activity/aggregated", + "/team/daily/activity/aggregated/search", "/team/spend/by_user", # gateway request counts (SGR); deployment-wide, admin-only "/gateway/daily/activity", @@ -887,6 +890,7 @@ class LiteLLMRoutes(enum.Enum): "/team/permissions_update", "/team/daily/activity", "/team/daily/activity/aggregated", + "/team/daily/activity/aggregated/search", "/team/spend/by_user", "/team/{team_id}/members/me", # POST/GET the team's logging callbacks, and DELETE one of them. Every @@ -986,6 +990,7 @@ class LiteLLMRoutes(enum.Enum): "/user/daily/activity", "/team/daily/activity", "/team/daily/activity/aggregated", + "/team/daily/activity/aggregated/search", "/tag/daily/activity", "/tag/list", "/audit", diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 6cec3e714ec..493c83c730a 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -40,6 +40,7 @@ from typing_extensions import ReadOnly, TypedDict, assert_never import litellm from litellm._logging import verbose_proxy_logger from litellm._uuid import uuid +from litellm.constants import USAGE_TOP_API_KEYS_LIMIT from litellm.integrations.prometheus import PrometheusLogger from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.proxy._types import ( @@ -196,6 +197,7 @@ from litellm.repositories.verification_token_repository import ( from litellm.router import Router from litellm.types.proxy.auth.auth_checks import UserNotFoundError from litellm.types.proxy.management_endpoints.common_daily_activity import ( + DailySpendMetadata, SpendAnalyticsPaginatedResponse, ) from litellm.types.proxy.management_endpoints.team_endpoints import ( @@ -204,7 +206,9 @@ from litellm.types.proxy.management_endpoints.team_endpoints import ( BulkUpdateTeamMemberPermissionsRequest, BulkUpdateTeamMemberPermissionsResponse, GetTeamMemberPermissionsResponse, + TeamIdSearchFilter, TeamIdSearchMatch, + TeamKeyActivitySearchWhere, TeamListItem, TeamListResponse, TeamMemberAddResult, @@ -6805,6 +6809,111 @@ async def get_team_daily_activity_aggregated( ) +def _team_key_search_where(*, search: str, scope: _TeamDailyActivityScope) -> TeamKeyActivitySearchWhere: + """Caller scoping lives inside the same Prisma where as the search term so `take` + never trims visible matches in favour of keys the caller is not allowed to see.""" + search_or: Final = ( + {"token": search}, # mutable-ok: prisma where clause leaf + {"key_alias": {"contains": search, "mode": "insensitive"}}, # mutable-ok: prisma where clause leaf + {"user_id": {"contains": search, "mode": "insensitive"}}, # mutable-ok: prisma where clause leaf + ) + own_keys: Final = tuple(scope.api_key_filter) if isinstance(scope.api_key_filter, list) else None + team_filter: Final[TeamIdSearchFilter | None] = ( + { # mutable-ok: prisma where clause leaf + "in": tuple(scope.team_ids), + "notIn": tuple(scope.exclude_team_ids), + } + if scope.team_ids is not None and scope.exclude_team_ids is not None + else {"in": tuple(scope.team_ids)} # mutable-ok: prisma where clause leaf + if scope.team_ids is not None + else {"notIn": tuple(scope.exclude_team_ids)} # mutable-ok: prisma where clause leaf + if scope.exclude_team_ids is not None + else None + ) + if team_filter is None and own_keys is None: + return {"OR": search_or} # mutable-ok: prisma where clause root + if team_filter is None and own_keys is not None: + return {"token": {"in": own_keys}, "OR": search_or} # mutable-ok: prisma where clause root + if team_filter is not None and own_keys is None: + return {"team_id": team_filter, "OR": search_or} # mutable-ok: prisma where clause root + assert team_filter is not None and own_keys is not None + return { # mutable-ok: prisma where clause root + "team_id": team_filter, + "token": {"in": own_keys}, # mutable-ok: prisma where clause leaf + "OR": search_or, + } + + +@router.get( + "/team/daily/activity/aggregated/search", + response_model=SpendAnalyticsPaginatedResponse, + tags=["team management"], # mutable-ok: FastAPI route tags shape +) +async def search_team_daily_activity_keys( + user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)], + search: str = fastapi.Query( + ..., + min_length=1, + description="Exact token hash, or a case-insensitive substring of the key alias or owning user id", + ), + team_ids: str | None = None, + start_date: str | None = None, + end_date: str | None = None, + exclude_team_ids: str | None = None, + timezone: int | None = None, +) -> SpendAnalyticsPaginatedResponse: + """Aggregated daily team activity for the keys matching `search`, across every key the caller may + see rather than only the top USAGE_TOP_API_KEYS_LIMIT keys by spend.""" + from litellm.proxy.proxy_server import ( + prisma_client, + proxy_logging_obj, + user_api_key_cache, + ) + + if prisma_client is None: + raise _daily_activity_error(status_code=500, message=CommonProxyErrors.db_not_connected_error.value) + + range_error: Final = _aggregated_date_range_error(start_date, end_date) + if range_error is not None: + raise _daily_activity_error(status_code=400, message=range_error) + + scope: Final = await _resolve_team_daily_activity_scope( + team_ids=team_ids, + exclude_team_ids=exclude_team_ids, + api_key=None, + user_api_key_dict=user_api_key_dict, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + matched_keys: Final = await _tokens_db(prisma_client).find_many( + where=_team_key_search_where(search=search, scope=scope), + take=USAGE_TOP_API_KEYS_LIMIT, + order={"spend": "desc"}, # mutable-ok: prisma serializes order, keep it a plain dict + ) + tokens: Final = [key.token for key in matched_keys] # mutable-ok: get_daily_activity_aggregated takes list[str] + if not tokens: + return SpendAnalyticsPaginatedResponse( + results=[], # mutable-ok: response model field shape + metadata=DailySpendMetadata(api_key_limit=USAGE_TOP_API_KEYS_LIMIT, total_api_keys=0), + ) + + return await get_daily_activity_aggregated( + prisma_client=prisma_client, + table_name="litellm_dailyteamspend", + entity_id_field="team_id", + entity_id=scope.team_ids, + entity_metadata_field=scope.team_alias_metadata, + start_date=start_date, + end_date=end_date, + model=None, + api_key=tokens, + exclude_entity_ids=scope.exclude_team_ids, + timezone_offset_minutes=timezone, + include_entity_breakdown=True, + ) + + def _team_user_spend_sql(*, team_count: int, restrict_to_user: bool) -> str: team_placeholders: Final = ", ".join(f"${i}" for i in range(3, 3 + team_count)) user_clause: Final = f' AND sl."user" = ${3 + team_count}' if restrict_to_user else "" diff --git a/litellm/types/proxy/management_endpoints/team_endpoints.py b/litellm/types/proxy/management_endpoints/team_endpoints.py index 4524c47ec38..aac2703e918 100644 --- a/litellm/types/proxy/management_endpoints/team_endpoints.py +++ b/litellm/types/proxy/management_endpoints/team_endpoints.py @@ -1,6 +1,8 @@ +from collections.abc import Mapping, Sequence from typing import Any, Final, Literal from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator +from typing_extensions import NotRequired, ReadOnly, TypedDict from litellm.proxy._types import ( KeyManagementRoutes, @@ -11,10 +13,32 @@ from litellm.proxy._types import ( MemberDeleteRequest, ) from litellm.proxy.common_utils.timezone_utils import budget_duration_error +from litellm.types.proxy.management_endpoints.internal_user_endpoints import InsensitiveContains from litellm.types.proxy.management_endpoints.management_v1 import ResourceResponse TeamIdSearchMatch = Literal["exact", "prefix"] + +TeamIdSearchFilter = TypedDict( + "TeamIdSearchFilter", + { # mutable-ok: functional TypedDict field map + "in": NotRequired[ReadOnly[Sequence[str]]], + "notIn": NotRequired[ReadOnly[Sequence[str]]], + }, +) + + +class TeamKeyActivitySearchWhere(TypedDict): + """Prisma filter behind `/team/daily/activity/aggregated/search`: exact token hash, or key alias + or user id containing the term, case-insensitive, narrowed to the teams and keys the caller may see.""" + + team_id: NotRequired[ReadOnly[TeamIdSearchFilter]] + token: NotRequired[ReadOnly[Mapping[Literal["in"], Sequence[str]]]] + OR: ReadOnly[ + tuple[Mapping[Literal["token"], str] | Mapping[Literal["key_alias", "user_id"], InsensitiveContains], ...] + ] + + MAX_BULK_TEAM_MEMBER_DELETES: Final = 500 MAX_BULK_TEAM_MEMBER_BUDGET_UPDATES: Final = 500 diff --git a/terraform/provider/tools/endpointaudit/coverage_allowlist.txt b/terraform/provider/tools/endpointaudit/coverage_allowlist.txt index d0f9c31ecaf..b0f28a9c740 100644 --- a/terraform/provider/tools/endpointaudit/coverage_allowlist.txt +++ b/terraform/provider/tools/endpointaudit/coverage_allowlist.txt @@ -28,6 +28,7 @@ GET /tag/user-agent/per-user-analytics GET /tag/wau GET /team/daily/activity GET /team/daily/activity/aggregated +GET /team/daily/activity/aggregated/search GET /team/spend/by_user GET /team/spend/report GET /user/daily/activity diff --git a/tests/integration/spend/test_team_daily_activity_key_search.py b/tests/integration/spend/test_team_daily_activity_key_search.py new file mode 100644 index 00000000000..2b0395cf4ba --- /dev/null +++ b/tests/integration/spend/test_team_daily_activity_key_search.py @@ -0,0 +1,128 @@ +import uuid +from datetime import datetime, timedelta, timezone +from hashlib import sha256 +from typing import Final + +import pytest +from integration._support.client import Gateway, eventually, object_value +from integration._support.database import read_rows +from pydantic import JsonValue + +_SEARCH_PATH: Final = "/team/daily/activity/aggregated/search" + + +def _range_around_today() -> dict[str, str]: + today: Final = datetime.now(timezone.utc) + return { + "start_date": (today - timedelta(days=1)).strftime("%Y-%m-%d"), + "end_date": (today + timedelta(days=1)).strftime("%Y-%m-%d"), + "timezone": "0", + } + + +def _team_key_breakdown(body: dict[str, JsonValue], team: str) -> dict[str, JsonValue]: + results: Final = body["results"] + assert isinstance(results, list) and len(results) == 1, body + entities: Final = object_value(object_value(object_value(results[0])["breakdown"])["entities"]) + return object_value(object_value(entities[team])["api_key_breakdown"]) + + +def test_team_key_search_returns_only_the_matching_key_spend_by_alias_and_by_hash(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + model: Final = scenario.model(input_cost_per_token=0.001, output_cost_per_token=0.002) + team: Final = scenario.team(models=[model]) + needle_alias: Final = f"needle-{uuid.uuid4().hex}" + needle: Final = scenario.key(team_id=team, models=[model], key_alias=needle_alias) + other: Final = scenario.key(team_id=team, models=[model], key_alias=f"other-{uuid.uuid4().hex}") + needle_digest: Final = sha256(needle.encode()).hexdigest() + other_digest: Final = sha256(other.encode()).hexdigest() + for key in (needle, other): + reply: Final = gateway.chat(model, key=key, text=f"key search {uuid.uuid4().hex}") + assert object_value(reply["usage"])["total_tokens"] == 40, reply + daily: Final = eventually( + lambda: read_rows('SELECT api_key, spend FROM "LiteLLM_DailyTeamSpend" WHERE team_id=%s', (team,)), + lambda values: sorted(row["api_key"] for row in values) == sorted((needle_digest, other_digest)), + seconds=70, + ) + assert all(float(row["spend"]) == pytest.approx(0.06) for row in daily), daily + for search in (needle_alias.upper(), needle_digest): + response: Final = gateway.request( + "GET", _SEARCH_PATH, params={"team_ids": team, "search": search, **_range_around_today()} + ) + assert response.status_code == 200, response.text + body: Final = object_value(response.json()) + assert object_value(body["metadata"])["total_spend"] == pytest.approx(0.06), response.text + per_key: Final = _team_key_breakdown(body, team) + assert set(per_key) == {needle_digest}, response.text + assert object_value(object_value(per_key[needle_digest])["metrics"])["spend"] == pytest.approx(0.06) + + +def test_team_key_search_is_scoped_to_the_teams_the_caller_belongs_to(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + model: Final = scenario.model(input_cost_per_token=0.001, output_cost_per_token=0.002) + team: Final = scenario.team(models=[model]) + needle_alias: Final = f"needle-{uuid.uuid4().hex}" + needle: Final = scenario.key(team_id=team, models=[model], key_alias=needle_alias) + needle_digest: Final = sha256(needle.encode()).hexdigest() + reply: Final = gateway.chat(model, key=needle, text=f"key search {uuid.uuid4().hex}") + assert object_value(reply["usage"])["total_tokens"] == 40, reply + eventually( + lambda: read_rows('SELECT api_key FROM "LiteLLM_DailyTeamSpend" WHERE team_id=%s', (team,)), + lambda values: [row["api_key"] for row in values] == [needle_digest], + seconds=70, + ) + outsider: Final = scenario.user(user_role="internal_user") + outsider_team: Final = scenario.team(models=[model], members_with_roles=[{"user_id": outsider, "role": "user"}]) + outsider_key: Final = scenario.key(user_id=outsider, team_id=outsider_team, models=[model]) + params: Final = {"search": needle_alias, **_range_around_today()} + admin_view: Final = gateway.request("GET", _SEARCH_PATH, params={"team_ids": team, **params}) + assert admin_view.status_code == 200, admin_view.text + assert set(_team_key_breakdown(object_value(admin_view.json()), team)) == {needle_digest}, admin_view.text + own_teams_view: Final = gateway.request("GET", _SEARCH_PATH, params=params, key=outsider_key) + assert own_teams_view.status_code == 200, own_teams_view.text + own_teams_body: Final = object_value(own_teams_view.json()) + assert own_teams_body["results"] == [], own_teams_view.text + assert object_value(own_teams_body["metadata"])["total_api_keys"] == 0, own_teams_view.text + foreign_team_view: Final = gateway.request( + "GET", _SEARCH_PATH, params={"team_ids": team, **params}, key=outsider_key + ) + assert foreign_team_view.status_code == 404, foreign_team_view.text + + +def test_team_key_search_excludes_teams_inside_the_where(gateway: Gateway) -> None: + """The dashboard always sends exclude_team_ids; a matching key in an excluded + team with higher spend must not consume a take slot nor appear in the result.""" + with gateway.scenario() as scenario: + model: Final = scenario.model(input_cost_per_token=0.001, output_cost_per_token=0.002) + team_keep: Final = scenario.team(models=[model]) + team_drop: Final = scenario.team(models=[model]) + shared_alias: Final = f"needle-{uuid.uuid4().hex}" + keep: Final = scenario.key(team_id=team_keep, models=[model], key_alias=f"{shared_alias}-keep") + drop: Final = scenario.key(team_id=team_drop, models=[model], key_alias=f"{shared_alias}-drop") + keep_digest: Final = sha256(keep.encode()).hexdigest() + drop_digest: Final = sha256(drop.encode()).hexdigest() + for _ in range(2): + reply: Final = gateway.chat(model, key=drop, text=f"key search {uuid.uuid4().hex}") + assert object_value(reply["usage"])["total_tokens"] == 40, reply + reply = gateway.chat(model, key=keep, text=f"key search {uuid.uuid4().hex}") + assert object_value(reply["usage"])["total_tokens"] == 40, reply + eventually( + lambda: read_rows( + 'SELECT api_key, spend FROM "LiteLLM_DailyTeamSpend" WHERE team_id IN (%s, %s)', + (team_keep, team_drop), + ), + lambda values: sorted(row["api_key"] for row in values) == sorted((keep_digest, drop_digest)), + seconds=70, + ) + response: Final = gateway.request( + "GET", + _SEARCH_PATH, + params={"search": shared_alias, "exclude_team_ids": team_drop, **_range_around_today()}, + ) + assert response.status_code == 200, response.text + body: Final = object_value(response.json()) + results: Final = body["results"] + assert isinstance(results, list) and len(results) == 1, body + entities: Final = object_value(object_value(object_value(results[0])["breakdown"])["entities"]) + assert set(entities) == {team_keep}, response.text + assert set(_team_key_breakdown(body, team_keep)) == {keep_digest}, response.text diff --git a/tests/proxy_behavior/management/test_team_daily_activity.py b/tests/proxy_behavior/management/test_team_daily_activity.py index d84cc4c94af..9bbc8fdde29 100644 --- a/tests/proxy_behavior/management/test_team_daily_activity.py +++ b/tests/proxy_behavior/management/test_team_daily_activity.py @@ -5,8 +5,9 @@ from .actors import Actor pytestmark = pytest.mark.asyncio(loop_scope="session") -# GET /team/daily/activity and its /aggregated variant (same shared scope -# resolver, so the matrix must hold for both). A proxy admin (admin view) sees +# GET /team/daily/activity, its /aggregated variant, and the key-search +# variant (same shared scope resolver, so the matrix must hold for all +# three). A proxy admin (admin view) sees # activity for any team. A non-admin is scoped to user_info.teams: a bare query # defaults to its own teams (200), and an explicit team_ids filter naming a # team it does not belong to is 404 (the VERIA-43 fix). Org admins have no @@ -43,8 +44,12 @@ _DATES = "start_date=2024-01-01&end_date=2024-12-31" @pytest.mark.parametrize( "endpoint", - ("/team/daily/activity", "/team/daily/activity/aggregated"), - ids=("paginated", "aggregated"), + ( + "/team/daily/activity", + "/team/daily/activity/aggregated", + "/team/daily/activity/aggregated/search", + ), + ids=("paginated", "aggregated", "search"), ) @pytest.mark.parametrize( "actor,team,expected_status", @@ -54,7 +59,7 @@ _DATES = "start_date=2024-01-01&end_date=2024-12-31" async def test_team_daily_activity_matrix( actor: Actor, team: str, expected_status: int, endpoint: str, proxy_client, world ): - query = _DATES + query = _DATES + ("&search=x" if endpoint.endswith("/search") else "") if team == "alpha": query += f"&team_ids={world.team_alpha_id}" elif team == "beta": diff --git a/tests/test_litellm/proxy/auth/test_route_checks.py b/tests/test_litellm/proxy/auth/test_route_checks.py index 571e066e947..7bb79a115dd 100644 --- a/tests/test_litellm/proxy/auth/test_route_checks.py +++ b/tests/test_litellm/proxy/auth/test_route_checks.py @@ -3600,6 +3600,55 @@ def test_user_daily_activity_aggregated_not_covered_by_prefix_match(): ) +@pytest.mark.parametrize( + "route", + [ + "/team/daily/activity", + "/team/daily/activity/aggregated", + "/team/daily/activity/aggregated/search", + ], +) +@pytest.mark.parametrize( + "user_role", + [ + LitellmUserRoles.INTERNAL_USER.value, + LitellmUserRoles.INTERNAL_USER_VIEW_ONLY.value, + ], +) +def test_team_daily_activity_routes_reachable_by_non_admin(route, user_role): + """The Team Usage dashboard calls all three team daily-activity routes, and + each handler self-scopes to the caller's teams and own keys + (_resolve_team_daily_activity_scope). self_managed_routes is the only list + granting them to a non-admin, and check_route_access is exact-match, so each + sub-path needs its own entry: dropping one 401s the dashboard before the + handler ever runs. + """ + user_obj = LiteLLM_UserTable( + user_id="test_user", + user_email="test@example.com", + user_role=user_role, + ) + valid_token = UserAPIKeyAuth(user_id="test_user", user_role=user_role) + request = MagicMock(spec=Request) + request.query_params = {} + + def outcome() -> str: + try: + RouteChecks.non_proxy_admin_allowed_routes_check( + user_obj=user_obj, + _user_role=user_role, + route=route, + request=request, + valid_token=valid_token, + request_data={}, + ) + except Exception as exc: + return f"denied: {exc}" + return "allowed" + + assert outcome() == "allowed" + + @pytest.mark.parametrize( "user_role", [ diff --git a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py index b066b3b80e6..38241926f8e 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -14646,6 +14646,218 @@ async def test_get_team_daily_activity_aggregated_rejects_bad_ranges( mock_aggregated.assert_not_called() +def _key_search_team_setup(mock_db_client, user_id: str, team_id: str): + mock_user_info = LiteLLM_UserTable( + user_id=user_id, + teams=[team_id], + max_budget=1000.0, + spend=0.0, + user_email="test@example.com", + user_role="internal_user", + ) + mock_team = MagicMock(spec=LiteLLM_TeamTable) + mock_team.team_id = team_id + mock_team.team_alias = "Test Team" + mock_team.members_with_roles = [Member(user_id=user_id, role="user")] + mock_team.model_dump.return_value = { + "team_id": team_id, + "team_alias": "Test Team", + "members_with_roles": [{"user_id": user_id, "role": "user"}], + } + mock_db_client.db.litellm_teamtable.find_many = AsyncMock(return_value=[mock_team]) + return mock_user_info + + +@pytest.mark.asyncio +async def test_search_team_daily_activity_keys_scopes_where_before_take(mock_db_client): + """A member's search must put the team and own-key scoping inside the same + Prisma where as the term, because `take` trims rows before Python sees them: + scoped outside the where, the top-N slice could be spent entirely on keys + the caller is not allowed to see.""" + from litellm.constants import USAGE_TOP_API_KEYS_LIMIT + from litellm.proxy.management_endpoints.team_endpoints import ( + search_team_daily_activity_keys, + ) + + user_id = "test_user_123" + team_id = "test_team_456" + user_api_key_dict = UserAPIKeyAuth(user_id=user_id, user_role=LitellmUserRoles.INTERNAL_USER) + mock_user_info = _key_search_team_setup(mock_db_client, user_id, team_id) + + user_key_1 = MagicMock() + user_key_1.token = "user_key_1" + matched = MagicMock() + matched.token = "user_key_1" + mock_db_client.db.litellm_verificationtoken.find_many = AsyncMock(side_effect=[[user_key_1], [matched]]) + + with patch( + "litellm.proxy.management_endpoints.team_endpoints.get_user_object", + new_callable=AsyncMock, + ) as mock_get_user_object: + mock_get_user_object.return_value = mock_user_info + + with patch( + "litellm.proxy.management_endpoints.team_endpoints.get_daily_activity_aggregated", + new_callable=AsyncMock, + ) as mock_aggregated: + mock_aggregated.return_value = MagicMock() + + await search_team_daily_activity_keys( + user_api_key_dict=user_api_key_dict, + search="Needle", + team_ids=team_id, + start_date="2024-01-01", + end_date="2024-01-31", + exclude_team_ids=None, + timezone=480, + ) + + token_calls = mock_db_client.db.litellm_verificationtoken.find_many.call_args_list + assert len(token_calls) == 2 + search_kwargs = token_calls[1][1] + assert search_kwargs["where"] == { + "team_id": {"in": (team_id,)}, + "token": {"in": ("user_key_1",)}, + "OR": ( + {"token": "Needle"}, + {"key_alias": {"contains": "Needle", "mode": "insensitive"}}, + {"user_id": {"contains": "Needle", "mode": "insensitive"}}, + ), + } + assert search_kwargs["take"] == USAGE_TOP_API_KEYS_LIMIT + assert search_kwargs["order"] == {"spend": "desc"} + + call_kwargs = mock_aggregated.call_args[1] + assert call_kwargs["api_key"] == ["user_key_1"] + assert call_kwargs["entity_id"] == [team_id] + assert call_kwargs["table_name"] == "litellm_dailyteamspend" + assert call_kwargs["include_entity_breakdown"] is True + assert call_kwargs["timezone_offset_minutes"] == 480 + assert call_kwargs["model"] is None + assert call_kwargs["entity_metadata_field"] == {team_id: {"team_alias": "Test Team"}} + + +@pytest.mark.asyncio +async def test_search_team_daily_activity_keys_admin_unscoped_where(mock_db_client): + """An admin's search has no caller scoping, so the where is the bare OR over + token, key alias and user id; every matched hash is passed through to the + aggregation.""" + from litellm.constants import USAGE_TOP_API_KEYS_LIMIT + from litellm.proxy.management_endpoints.team_endpoints import ( + search_team_daily_activity_keys, + ) + + match_1 = MagicMock() + match_1.token = "h1" + match_2 = MagicMock() + match_2.token = "h2" + mock_db_client.db.litellm_teamtable.find_many = AsyncMock(return_value=[]) + mock_db_client.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[match_1, match_2]) + + with patch( + "litellm.proxy.management_endpoints.team_endpoints.get_daily_activity_aggregated", + new_callable=AsyncMock, + ) as mock_aggregated: + mock_aggregated.return_value = MagicMock() + + await search_team_daily_activity_keys( + user_api_key_dict=UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN), + search="Needle", + team_ids=None, + start_date="2024-01-01", + end_date="2024-01-31", + exclude_team_ids=None, + timezone=None, + ) + + search_kwargs = mock_db_client.db.litellm_verificationtoken.find_many.call_args[1] + assert search_kwargs["where"] == { + "OR": ( + {"token": "Needle"}, + {"key_alias": {"contains": "Needle", "mode": "insensitive"}}, + {"user_id": {"contains": "Needle", "mode": "insensitive"}}, + ) + } + assert search_kwargs["take"] == USAGE_TOP_API_KEYS_LIMIT + assert mock_aggregated.call_args[1]["api_key"] == ["h1", "h2"] + + +@pytest.mark.asyncio +async def test_search_team_daily_activity_keys_no_match_returns_empty_without_aggregating( + mock_db_client, +): + """A term matching no key still owes the caller the standard metadata shape + (api_key_limit, total_api_keys), and the aggregated query must not run.""" + from litellm.constants import USAGE_TOP_API_KEYS_LIMIT + from litellm.proxy.management_endpoints.team_endpoints import ( + search_team_daily_activity_keys, + ) + + mock_db_client.db.litellm_teamtable.find_many = AsyncMock(return_value=[]) + mock_db_client.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[]) + + with patch( + "litellm.proxy.management_endpoints.team_endpoints.get_daily_activity_aggregated", + new_callable=AsyncMock, + ) as mock_aggregated: + result = await search_team_daily_activity_keys( + user_api_key_dict=UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN), + search="Needle", + team_ids=None, + start_date="2024-01-01", + end_date="2024-01-31", + exclude_team_ids=None, + timezone=None, + ) + + assert result.results == [] + assert result.metadata.total_api_keys == 0 + assert result.metadata.api_key_limit == USAGE_TOP_API_KEYS_LIMIT + mock_aggregated.assert_not_called() + + +@pytest.mark.asyncio +async def test_search_team_daily_activity_keys_excludes_teams_in_where(mock_db_client): + """The dashboard always sends exclude_team_ids=litellm-dashboard; if that + filter stayed out of the where, matching keys in excluded teams could fill + the take=N slice and push visible matches out.""" + from litellm.proxy.management_endpoints.team_endpoints import ( + search_team_daily_activity_keys, + ) + + matched = MagicMock() + matched.token = "h1" + mock_db_client.db.litellm_teamtable.find_many = AsyncMock(return_value=[]) + mock_db_client.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[matched]) + + with patch( + "litellm.proxy.management_endpoints.team_endpoints.get_daily_activity_aggregated", + new_callable=AsyncMock, + ) as mock_aggregated: + mock_aggregated.return_value = MagicMock() + + await search_team_daily_activity_keys( + user_api_key_dict=UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN), + search="Needle", + team_ids=None, + start_date="2024-01-01", + end_date="2024-01-31", + exclude_team_ids="litellm-dashboard", + timezone=None, + ) + + search_kwargs = mock_db_client.db.litellm_verificationtoken.find_many.call_args[1] + assert search_kwargs["where"] == { + "team_id": {"notIn": ("litellm-dashboard",)}, + "OR": ( + {"token": "Needle"}, + {"key_alias": {"contains": "Needle", "mode": "insensitive"}}, + {"user_id": {"contains": "Needle", "mode": "insensitive"}}, + ), + } + assert mock_aggregated.call_args[1]["exclude_entity_ids"] == ["litellm-dashboard"] + + def _wire_new_team_prisma(mock_db_client): mock_db_client.jsonify_team_object = lambda db_data: db_data mock_db_client.get_data = AsyncMock(return_value=None) diff --git a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/EntityUsage.tsx b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/EntityUsage.tsx index 27460b21108..16dc41c3ba8 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/EntityUsage.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/EntityUsage.tsx @@ -20,7 +20,7 @@ import type { ColumnDef } from "@tanstack/react-table"; import PaginationStatusAlerts from "@/components/shared/PaginationStatusAlerts"; import { Tabs, TabsContent, TabsList, TabsTrigger } from "@/components/ui/tabs"; import { Tooltip, TooltipContent, TooltipTrigger } from "@/components/ui/tooltip"; -import React, { type ReactNode, useMemo, useState } from "react"; +import React, { type ReactNode, useCallback, useMemo, useState } from "react"; import TeamMultiSelect from "@/components/common_components/team_multi_select"; import UserDropdown from "@/components/common_components/UserDropdown"; import { ActivityMetrics, processActivityData } from "@/components/activity_metrics"; @@ -34,6 +34,7 @@ import { tagDailyActivityCall, teamDailyActivityAggregatedCall, teamDailyActivityCall, + teamDailyActivityKeySearchCall, userDailyActivityCall, } from "@/components/networking"; import { Logo } from "@/components/molecules/logo/Logo"; @@ -182,6 +183,16 @@ const EntityUsage: React.FC = ({ const modelBreakdownKey = modelViewType === "groups" ? "model_groups" : "models"; const modelMetrics = processActivityData(spendData, modelBreakdownKey, teams || []); const keyMetrics = processActivityData(spendData, "api_keys", teams || []); + const searchTeamKeys = useCallback( + (query: string) => { + if (!accessToken || !startTime || !endTime) return Promise.resolve({}); + const teamIds = Array.isArray(entityFilterArg) ? entityFilterArg : null; + return teamDailyActivityKeySearchCall(accessToken, startTime, endTime, query, teamIds).then((data) => + processActivityData(data, "api_keys", teams || []), + ); + }, + [accessToken, startTime, endTime, entityFilterArg, teams], + ); const agentMetrics = showAgentBreakdown ? processActivityData(agentSpendData, "entities", teams || []) : {}; const getAllTags = () => { @@ -667,6 +678,7 @@ const EntityUsage: React.FC = ({ keyMetrics={keyMetrics} hidePromptCachingMetrics={entityType === "agent"} apiKeyTruncation={apiKeyTruncation} + searchKeys={entityType === "team" ? searchTeamKeys : undefined} /> ), }, diff --git a/ui/litellm-dashboard/src/components/networking.tsx b/ui/litellm-dashboard/src/components/networking.tsx index e6923a26d4c..edaf56a8a16 100644 --- a/ui/litellm-dashboard/src/components/networking.tsx +++ b/ui/litellm-dashboard/src/components/networking.tsx @@ -1467,6 +1467,31 @@ export const teamDailyActivityAggregatedCall = async ( } }; +export const teamDailyActivityKeySearchCall = async ( + accessToken: string, + startTime: Date, + endTime: Date, + ...options: [search: string, teamIds?: string[] | null] +) => { + const [search, teamIds = null] = options; + try { + return await apiClient.get(`/team/daily/activity/aggregated/search`, { + accessToken, + query: { + start_date: formatDate(startTime), + end_date: formatDate(endTime), + timezone: new Date().getTimezoneOffset().toString(), + search, + team_ids: teamIds && teamIds.length > 0 ? teamIds.join(",") : undefined, + exclude_team_ids: "litellm-dashboard", + }, + }); + } catch (error) { + console.error("Failed to search team daily activity keys:", error); + throw error; + } +}; + export type TeamUserSpendResponse = components["schemas"]["TeamUserSpendResponse"]; export const teamSpendByUserCall = async ( diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 79d008c1b48..a398b2db93c 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -15645,6 +15645,27 @@ export interface paths { patch?: never; trace?: never; }; + "/team/daily/activity/aggregated/search": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + /** + * Search Team Daily Activity Keys + * @description Aggregated daily team activity for the keys matching `search`, across every key the caller may + * see rather than only the top USAGE_TOP_API_KEYS_LIMIT keys by spend. + */ + get: operations["search_team_daily_activity_keys_team_daily_activity_aggregated_search_get"]; + put?: never; + post?: never; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; "/team/delete": { parameters: { query?: never; @@ -31378,7 +31399,7 @@ export interface components { * @description Enum for key management routes * @enum {string} */ - KeyManagementRoutes: "/key/generate" | "/key/update" | "/key/delete" | "/key/regenerate" | "/key/service-account/generate" | "/key/{key_id}/regenerate" | "/key/block" | "/key/unblock" | "/key/bulk_update" | "/team/key/bulk_update" | "/key/{key_id}/reset_spend" | "/key/access_group_assignment" | "/auto_router/manage" | "/key/info" | "/key/health" | "/key/list" | "/key/aliases" | "/team/daily/activity" | "/team/daily/activity/aggregated" | "/spend/logs" | "/spend/logs/v2"; + KeyManagementRoutes: "/key/generate" | "/key/update" | "/key/delete" | "/key/regenerate" | "/key/service-account/generate" | "/key/{key_id}/regenerate" | "/key/block" | "/key/unblock" | "/key/bulk_update" | "/team/key/bulk_update" | "/key/{key_id}/reset_spend" | "/key/access_group_assignment" | "/auto_router/manage" | "/key/info" | "/key/health" | "/key/list" | "/key/aliases" | "/team/daily/activity" | "/team/daily/activity/aggregated" | "/team/daily/activity/aggregated/search" | "/spend/logs" | "/spend/logs/v2"; /** * KeyManagementSystem * @enum {string} @@ -66624,6 +66645,43 @@ export interface operations { }; }; }; + search_team_daily_activity_keys_team_daily_activity_aggregated_search_get: { + parameters: { + query: { + /** @description Exact token hash, or a case-insensitive substring of the key alias or owning user id */ + search: string; + team_ids?: string | null; + start_date?: string | null; + end_date?: string | null; + exclude_team_ids?: string | null; + timezone?: number | null; + }; + header?: never; + path?: never; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["SpendAnalyticsPaginatedResponse"]; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; delete_team_team_delete_post: { parameters: { query?: never; From 1fbd1e9ce90580f801612d2016e9e2cc4ab553b5 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 24 Sep 2026 13:08:27 -0700 Subject: [PATCH 074/120] fix(bedrock): route unmapped openai family model ids to converse (#42713) * fix(bedrock): route unmapped openai family model ids to converse Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(bedrock): rename the e2e openai family backend constant global.openai.gpt-6-sol has a cost-map row now, so the constant no longer names an unmapped model --------- 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> --- litellm/llms/bedrock/common_utils.py | 2 ++ .../test_bedrock_provider_matrix_e2e.py | 19 ++++++++++++++++++- .../llms/bedrock/test_bedrock_common_utils.py | 19 ++++++++++++++++++- 3 files changed, 38 insertions(+), 2 deletions(-) diff --git a/litellm/llms/bedrock/common_utils.py b/litellm/llms/bedrock/common_utils.py index 9b52f531cbb..5f044897b2c 100644 --- a/litellm/llms/bedrock/common_utils.py +++ b/litellm/llms/bedrock/common_utils.py @@ -1218,6 +1218,8 @@ class BedrockModelInfo(BaseLLMModelInfo): alt_model: Final = BedrockModelInfo.get_non_litellm_routing_model_name(model=model) if base_model in litellm.bedrock_converse_models or alt_model in litellm.bedrock_converse_models: return "converse" + if _OPENAI_FAMILY_MODEL_RE.search(base_model): + return "converse" return "invoke" @staticmethod diff --git a/tests/e2e/llm_translation/test_bedrock_provider_matrix_e2e.py b/tests/e2e/llm_translation/test_bedrock_provider_matrix_e2e.py index 5f0a931109c..21333d39849 100644 --- a/tests/e2e/llm_translation/test_bedrock_provider_matrix_e2e.py +++ b/tests/e2e/llm_translation/test_bedrock_provider_matrix_e2e.py @@ -8,7 +8,8 @@ caller can hand AWS support the request id behind a completion. Regional inference-profile ids are the deployment shape most Bedrock customers run; a v1.90.0 regression timed them out, and the Converse route keeps them covered in test_chat_completions_regression_e2e.py, so the invoke route carries its own -rows here. +rows here. The file also covers Bedrock-native OpenAI model ids taking the +default (Converse) route with max_tokens. """ from __future__ import annotations @@ -26,6 +27,7 @@ pytestmark = pytest.mark.e2e CONVERSE_REGIONAL_BACKEND = "bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0" INVOKE_REGIONAL_BACKEND = "bedrock/invoke/us.anthropic.claude-haiku-4-5-20251001-v1:0" +OPENAI_FAMILY_BACKEND = "bedrock/global.openai.gpt-6-sol" PROVIDER_HEADER_PREFIX = "llm_provider-" BEDROCK_REQUEST_ID_HEADER = "llm_provider-x-amzn-requestid" @@ -199,3 +201,18 @@ class TestBedrockInvokeRegionalModelIds: ) _assert_streamed_completion(result) + + +class TestBedrockOpenAIFamilyDefaultRoute: + @pytest.mark.covers("llm.chat_completions.bedrock_converse.basic.nonstream.works", exercised_on=[]) + def test_openai_family_model_id_completes_with_max_tokens( + self, client: PassthroughClient, resources: ResourceManager + ) -> None: + model = _register_bedrock_model( + client, resources, "e2e-bedrock-openai-family", OPENAI_FAMILY_BACKEND + ) + key = resources.key() + + response = unwrap(client.proxy.chat(key, ChatBody(model=model, messages=_prompt(), max_tokens=64))) + + _assert_completion(response) diff --git a/tests/test_litellm/llms/bedrock/test_bedrock_common_utils.py b/tests/test_litellm/llms/bedrock/test_bedrock_common_utils.py index 117814a41ff..e5118f90e44 100644 --- a/tests/test_litellm/llms/bedrock/test_bedrock_common_utils.py +++ b/tests/test_litellm/llms/bedrock/test_bedrock_common_utils.py @@ -346,7 +346,7 @@ def test_route_prefix_matched_as_path_segment_not_substring(): BedrockModelInfo.get_bedrock_route("bedrock_mantle/openai.gpt-5.5") != "mantle" ) assert ( - BedrockModelInfo.get_bedrock_route("bedrock_mantle/openai.gpt-5.4") == "invoke" + BedrockModelInfo.get_bedrock_route("bedrock_mantle/openai.gpt-5.4") == "converse" ) assert ( BedrockModelInfo._explicit_mantle_route("bedrock_mantle/openai.gpt-5.5") @@ -964,3 +964,20 @@ def test_s3_static_key_pair_is_none_without_a_full_pair(partial_s3_pair): from litellm.llms.bedrock.common_utils import s3_static_key_pair assert s3_static_key_pair({"aws_access_key_id": "bedrock-key", **partial_s3_pair}) is None + + +def test_unmapped_openai_family_model_routes_to_converse(): + """A Bedrock-native OpenAI model that is not in the cost map yet must not fall to the invoke route. + + The invoke ``openai`` provider is the imported-model path and sends ``max_tokens``, which Bedrock + rejects for these models; Converse maps it to ``inferenceConfig.maxTokens``. + """ + from typing import Final + + import litellm + + unmapped: Final = "bedrock/global.openai.gpt-99-unmapped" + assert unmapped.removeprefix("bedrock/") not in litellm.bedrock_converse_models + assert BedrockModelInfo.get_bedrock_route(unmapped) == "converse" + imported: Final = "bedrock/openai/arn:aws:bedrock:us-east-1:123456789012:imported-model/abc123" + assert BedrockModelInfo.get_bedrock_route(imported) == "openai" From e2781c47132a129b67667ed824d668899b4c3661 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 24 Sep 2026 20:19:00 +0000 Subject: [PATCH 075/120] refactor(rust): move tests.rs files inline or under tests/ and drop autotests = false (#43028) * refactor(rust): move tests.rs files inline or under tests/ Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * refactor(rust): move tests.rs files inline or under tests/ Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * refactor(rust): inline path-included test files into their owning src files Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * chore(rust): drop stray proptest regression file Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(rust): cover lowercase, empty and non-authorization headers in bearer detection Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: Yujong Lee Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm-rust/AGENTS.md | 10 + .../crates/callbacks-legacy-python/Cargo.toml | 1 - .../callbacks-legacy-python/src/adapter.rs | 1116 ++++- .../callbacks-legacy-python/src/deferred.rs | 150 +- .../crates/callbacks-legacy-python/src/lib.rs | 213 +- .../callbacks-legacy-python/tests/deferred.rs | 146 - .../tests/deployment_hooks.rs | 282 -- .../callbacks-legacy-python/tests/payload.rs | 523 --- .../callbacks-legacy-python/tests/support.rs | 205 - .../callbacks-legacy-python/tests/terminal.rs | 291 -- litellm-rust/crates/core/Cargo.toml | 1 - .../core/src/audio_transcription/mod.rs | 3 - .../crates/core/src/chat_completions/mod.rs | 3 - .../core/src/chat_completions/prepare.rs | 838 ++++ .../crates/core/src/chat_completions/tests.rs | 833 ---- .../crates/core/src/messages/common_utils.rs | 183 + litellm-rust/crates/core/src/messages/mod.rs | 3 - litellm-rust/crates/core/src/ocr/document.rs | 157 + litellm-rust/crates/core/src/ocr/mod.rs | 238 +- litellm-rust/crates/core/src/ocr/route.rs | 3586 +++++++++++++++++ .../tests.rs => tests/audio_transcription.rs} | 4 +- .../crates/core/tests/aws_textract_ocr.rs | 193 - .../crates/core/tests/azure_ai_ocr.rs | 293 -- .../tests/azure_document_intelligence_ocr.rs | 712 ---- litellm-rust/crates/core/tests/cohere_ocr.rs | 136 - .../crates/core/tests/deepseek_ocr.rs | 133 - .../messages/tests.rs => tests/messages.rs} | 132 +- litellm-rust/crates/core/tests/ocr.rs | 1033 ----- .../crates/core/tests/ocr/document.rs | 152 - litellm-rust/crates/core/tests/ocr/support.rs | 203 - litellm-rust/crates/core/tests/reducto_ocr.rs | 584 --- .../core/tests/vertex_ai_deepseek_ocr.rs | 143 - .../crates/core/tests/vertex_ai_ocr.rs | 293 -- litellm-rust/crates/http/src/request.rs | 12 + .../llms/src/anthropic/chat/transformation.rs | 4 - .../bedrock/chat/converse_transformation.rs | 4 - .../anthropic_chat_transformation.rs} | 18 +- .../bedrock_converse_transformation.rs} | 12 +- litellm-rust/crates/types/src/utils.rs | 2 +- 39 files changed, 6486 insertions(+), 6359 deletions(-) create mode 100644 litellm-rust/AGENTS.md delete mode 100644 litellm-rust/crates/callbacks-legacy-python/tests/deferred.rs delete mode 100644 litellm-rust/crates/callbacks-legacy-python/tests/deployment_hooks.rs delete mode 100644 litellm-rust/crates/callbacks-legacy-python/tests/payload.rs delete mode 100644 litellm-rust/crates/callbacks-legacy-python/tests/support.rs delete mode 100644 litellm-rust/crates/callbacks-legacy-python/tests/terminal.rs delete mode 100644 litellm-rust/crates/core/src/chat_completions/tests.rs rename litellm-rust/crates/core/{src/audio_transcription/tests.rs => tests/audio_transcription.rs} (95%) delete mode 100644 litellm-rust/crates/core/tests/aws_textract_ocr.rs delete mode 100644 litellm-rust/crates/core/tests/azure_ai_ocr.rs delete mode 100644 litellm-rust/crates/core/tests/azure_document_intelligence_ocr.rs delete mode 100644 litellm-rust/crates/core/tests/cohere_ocr.rs delete mode 100644 litellm-rust/crates/core/tests/deepseek_ocr.rs rename litellm-rust/crates/core/{src/messages/tests.rs => tests/messages.rs} (79%) delete mode 100644 litellm-rust/crates/core/tests/ocr.rs delete mode 100644 litellm-rust/crates/core/tests/ocr/document.rs delete mode 100644 litellm-rust/crates/core/tests/ocr/support.rs delete mode 100644 litellm-rust/crates/core/tests/reducto_ocr.rs delete mode 100644 litellm-rust/crates/core/tests/vertex_ai_deepseek_ocr.rs delete mode 100644 litellm-rust/crates/core/tests/vertex_ai_ocr.rs rename litellm-rust/crates/llms/{src/anthropic/chat/tests.rs => tests/anthropic_chat_transformation.rs} (96%) rename litellm-rust/crates/llms/{src/bedrock/chat/tests.rs => tests/bedrock_converse_transformation.rs} (98%) diff --git a/litellm-rust/AGENTS.md b/litellm-rust/AGENTS.md new file mode 100644 index 00000000000..e5ffcd1c57a --- /dev/null +++ b/litellm-rust/AGENTS.md @@ -0,0 +1,10 @@ +# Rust workspace rules + +## Test placement + +- Never create a `tests.rs` (or `test.rs`) file under `src/`, and never `#[path = "tests.rs"] mod tests;` +- A test that reaches private items lives inline, in a `#[cfg(test)] mod tests { ... }` at the bottom of the file that owns those items +- A test that only uses the crate's public API lives in `crates//tests/.rs`, next to `src/` +- Split a mixed test file along that line instead of widening visibility to move it +- A test for another crate's item belongs in that crate, not in a downstream one +- Never set `autotests = false` or hand-list `[[test]]` targets; every file directly under `tests/` is discovered by cargo, and a shared helper goes in `tests//mod.rs` or `tests//support.rs` so it is not picked up as a test crate of its own diff --git a/litellm-rust/crates/callbacks-legacy-python/Cargo.toml b/litellm-rust/crates/callbacks-legacy-python/Cargo.toml index fe19578e04d..ed5e0fb9691 100644 --- a/litellm-rust/crates/callbacks-legacy-python/Cargo.toml +++ b/litellm-rust/crates/callbacks-legacy-python/Cargo.toml @@ -4,7 +4,6 @@ version = "0.1.0" edition.workspace = true license.workspace = true repository.workspace = true -autotests = false [dependencies] litellm-host.workspace = true diff --git a/litellm-rust/crates/callbacks-legacy-python/src/adapter.rs b/litellm-rust/crates/callbacks-legacy-python/src/adapter.rs index e1742190205..75a635e9c63 100644 --- a/litellm-rust/crates/callbacks-legacy-python/src/adapter.rs +++ b/litellm-rust/crates/callbacks-legacy-python/src/adapter.rs @@ -487,11 +487,1115 @@ impl PythonLifecycle for LegacyLogging { } #[cfg(test)] -#[path = "../tests/deployment_hooks.rs"] -mod deployment_hooks_tests; +mod deployment_hooks_tests { + use std::ffi::CStr; + + use litellm_host::event::{FailureOrigin, Timing}; + use litellm_host_python::{LifecycleEvent, LifecycleStep, PythonLifecycle}; + use pyo3::exceptions::asyncio::CancelledError; + use pyo3::prelude::*; + use pyo3::types::PyDict; + use rstest::rstest; + + use super::LegacyLogging; + use crate::test_support::{legacy_call, local, namespace, run}; + + const CALL: &CStr = c" +document = {'type': 'document_url', 'document_url': 'data:application/pdf;base64,YWJj'} +kwargs = {'logger': logger, 'document': document} +"; + + const TIMING: Timing = Timing { + start_time: 0.0, + end_time: 1.0, + }; + + fn begin<'py>( + py: Python<'py>, + locals: &Bound<'py, PyDict>, + asynchronous: bool, + ) -> (LegacyLogging, LifecycleStep) { + let mut logging = legacy_call(py, locals, asynchronous); + let kwargs = local(locals, "kwargs") + .cast_into::() + .unwrap() + .unbind(); + let step = logging.begin(py, kwargs, 0.0).unwrap(); + (logging, step) + } + + fn arguments<'py>(py: Python<'py>, step: LifecycleStep) -> Bound<'py, PyDict> { + let LifecycleStep::Arguments(arguments) = step else { + panic!("expected the prepared arguments"); + }; + arguments.into_bound(py) + } + + fn awaits_deployment_hook(step: &LifecycleStep) -> bool { + matches!(step, LifecycleStep::Await(_)) + } + + #[rstest] + #[case::synchronous(false)] + #[case::asynchronous(true)] + fn deployment_pre_call_hook_runs_only_for_asynchronous_calls(#[case] asynchronous: bool) { + Python::initialize(); + Python::attach(|py| { + let locals = namespace(py, CALL); + let (_, step) = begin(py, &locals, asynchronous); + assert_eq!(awaits_deployment_hook(&step), asynchronous); + let names: Vec = local(&locals, "logger") + .call_method0("names") + .unwrap() + .extract() + .unwrap(); + assert_eq!(names.contains(&"pre_hook".to_string()), asynchronous); + }); + } + + #[test] + fn kwargs_returned_by_the_pre_call_hook_are_what_the_call_prepares() { + Python::initialize(); + Python::attach(|py| { + let locals = namespace( + py, + c" +document = {'type': 'document_url', 'document_url': 'data:application/pdf;base64,YWJj'} +replacement = {'type': 'document_url', 'document_url': 'data:application/pdf;base64,ZWRpdGVk'} +kwargs = {'logger': logger, 'document': document} +replaced_kwargs = {'logger': logger, 'document': replacement, 'pages': [0]} +", + ); + let (mut logging, step) = begin(py, &locals, true); + assert!(awaits_deployment_hook(&step)); + let step = logging + .resume(py, Ok(local(&locals, "replaced_kwargs").unbind())) + .unwrap(); + locals.set_item("prepared", arguments(py, step)).unwrap(); + run( + py, + &locals, + c" +assert prepared['document'] is replacement +assert prepared['pages'] is replaced_kwargs['pages'] +assert prepared['litellm_logging_obj'] is logger +assert 'litellm_logging_obj' not in replaced_kwargs +[checked] = [value for name, value in logger.calls if name == 'check_limits'] +assert checked is prepared +", + ); + }); + } + + #[rstest] + #[case::synchronous(false)] + #[case::asynchronous(true)] + fn a_keyword_the_bridge_never_reads_reaches_every_reader_as_the_callers_object( + #[case] asynchronous: bool, + ) { + Python::initialize(); + Python::attach(|py| { + let locals = namespace( + py, + c" +opaque = object() +hooked = [] +logger.hooks = {'pre': lambda kwargs: hooked.append(kwargs['vendor_extension']) or kwargs} +kwargs = {'logger': logger, 'vendor_extension': opaque} +", + ); + let (mut logging, step) = begin(py, &locals, asynchronous); + let step = match step { + LifecycleStep::Await(hook_result) => logging.resume(py, Ok(hook_result)).unwrap(), + step => step, + }; + locals.set_item("prepared", arguments(py, step)).unwrap(); + locals.set_item("asynchronous", asynchronous).unwrap(); + run( + py, + &locals, + c" +assert prepared['vendor_extension'] is opaque +[checked] = [value for name, value in logger.calls if name == 'check_limits'] +assert checked['vendor_extension'] is opaque +assert hooked == ([opaque] if asynchronous else []), hooked +", + ); + }); + } + + #[test] + fn response_returned_by_the_post_call_hook_is_finalized_and_returned() { + Python::initialize(); + Python::attach(|py| { + let locals = namespace( + py, + c" +kwargs = {'logger': logger} +response = object() +replacement = object() +logger.hooks = {'pre': lambda kwargs: kwargs} +", + ); + let (mut logging, _) = begin(py, &locals, true); + logging + .resume(py, Ok(local(&locals, "kwargs").unbind())) + .unwrap(); + let step = logging + .after_success(py, local(&locals, "response").unbind(), TIMING) + .unwrap(); + assert!(awaits_deployment_hook(&step)); + let step = logging + .resume(py, Ok(local(&locals, "replacement").unbind())) + .unwrap(); + let LifecycleStep::Response(returned) = step else { + panic!("expected the finalized response"); + }; + assert!(returned.bind(py).is(local(&locals, "replacement"))); + run( + py, + &locals, + c" +[finalized] = [value for name, value in logger.calls if name == 'finalize'] +assert finalized is replacement +", + ); + }); + } + + #[rstest] + #[case::pre_call(false)] + #[case::post_call(true)] + fn cancelling_a_deployment_hook_ends_the_call_with_that_cancellation(#[case] post_call: bool) { + Python::initialize(); + Python::attach(|py| { + let locals = namespace(py, c"kwargs = {'logger': logger}\nresponse = object()"); + let (mut logging, _) = begin(py, &locals, true); + if post_call { + logging + .resume(py, Ok(local(&locals, "kwargs").unbind())) + .unwrap(); + logging + .after_success(py, local(&locals, "response").unbind(), TIMING) + .unwrap(); + } + let cancellation = CancelledError::new_err("cancelled"); + let cancelled = cancellation.value(py).clone(); + let error = logging.resume(py, Err(cancellation)).err().unwrap(); + assert!(error.value(py).is(&cancelled)); + let names: Vec = local(&locals, "logger") + .call_method0("names") + .unwrap() + .extract() + .unwrap(); + assert!(!names.iter().any(|name| name.contains("handler"))); + }); + } + + #[rstest] + #[case::hook_completed(false)] + #[case::hook_cancelled(true)] + fn failure_callbacks_run_after_the_failure_hook_however_it_ends(#[case] cancelled: bool) { + Python::initialize(); + Python::attach(|py| { + let locals = namespace( + py, + c"kwargs = {'logger': logger}\nfailure = ValueError('provider')", + ); + let (mut logging, _) = begin(py, &locals, true); + logging + .resume(py, Ok(local(&locals, "kwargs").unbind())) + .unwrap(); + let failure = PyErr::from_value(local(&locals, "failure")); + let failed = LifecycleEvent::Failed { + timing: TIMING, + origin: FailureOrigin::Call, + error: &failure, + }; + let step = logging.emit(py, failed).unwrap(); + assert!(awaits_deployment_hook(&step)); + let hook_result = if cancelled { + Err(CancelledError::new_err("cancelled")) + } else { + Ok(py.None()) + }; + assert!(matches!( + logging.resume(py, hook_result).unwrap(), + LifecycleStep::Await(_) + )); + run( + py, + &locals, + c" +assert logger.names()[-3:] == ['failure_hook', 'failure_handler', 'async_failure_handler'], logger.calls +assert all(value is failure for name, value in logger.calls if name.endswith('_handler')) +", + ); + }); + } + + #[rstest] + #[case::synchronous(false)] + #[case::asynchronous(true)] + fn a_limit_rejected_before_the_call_surfaces_as_the_callers_error(#[case] asynchronous: bool) { + Python::initialize(); + Python::attach(|py| { + let locals = namespace( + py, + c" +class BudgetExceeded(Exception): + pass + +rejection = BudgetExceeded('over budget') + +class LimitedLogger(StubLogger): + def check_limits(self, arguments): + raise rejection + +logger = LimitedLogger() +logger.hooks = {'pre': lambda kwargs: kwargs} +kwargs = {'logger': logger} +", + ); + let mut logging = legacy_call(py, &locals, asynchronous); + let kwargs = local(&locals, "kwargs") + .cast_into::() + .unwrap() + .unbind(); + let result = logging.begin(py, kwargs, 0.0).and_then(|step| match step { + LifecycleStep::Await(_) => { + logging.resume(py, Ok(local(&locals, "kwargs").unbind())) + } + step => Ok(step), + }); + let error = result.err().unwrap(); + assert!(error.value(py).is(local(&locals, "rejection"))); + }); + } +} + #[cfg(test)] -#[path = "../tests/payload.rs"] -mod payload_tests; +mod payload_tests { + use std::ffi::CStr; + + use litellm_auth::SecretValue; + use litellm_host::event::{MachineEvent, RawResponse, RequestContext, WireRequest}; + use litellm_host_python::{LifecycleEvent, LifecycleStep, PythonLifecycle, to_py}; + use proptest::prelude::*; + use pyo3::prelude::*; + use rstest::rstest; + use serde_json::{Map, Value, json}; + + use super::LegacyLogging; + use crate::PythonLogger; + use crate::test_support::{legacy_call, local, namespace, run}; + + /// The payload phases of `Logging` on top of `StubLogger`, with `pre_call` handing the + /// payload to the case's `on_pre_call`. + const PAYLOAD_LOGGER: &CStr = c" +class Request: + pass + +class PayloadLogger(StubLogger): + def update_from_kwargs(self, **update): + self.update = update + + def pre_call(self, input, api_key, additional_args): + self.record('pre_call', None) + self.pre = additional_args + self.pre_api_key = api_key + on_pre_call(additional_args) + + def post_call(self, original_response, api_key, additional_args): + self.record('post_call', None) + self.post = (original_response, api_key, additional_args) + +request = Request() +kwargs = {} +logger = PayloadLogger() +on_pre_call = lambda additional_args: None +check = lambda: None +"; + + const DOCUMENT: &str = "data:application/pdf;base64,YWJj"; + const EDITED: &str = "data:application/pdf;base64,ZWRpdGVk"; + + fn document(source: &str) -> Value { + json!({"type": "document_url", "document_url": source}) + } + + fn before_send(script: &CStr, body: Value) -> WireRequest { + before_send_with_secrets(script, json!({}), body, &[]) + } + + /// Runs `before_send` over `body` for a route whose parameters are `optional_params`, with + /// the Python objects `script` binds, then delivers the provider's raw response the way the + /// driver does and runs the script's `check()`. + fn before_send_with_secrets( + script: &CStr, + optional_params: Value, + body: Value, + secret_fields: &[&str], + ) -> WireRequest { + before_send_bound(&[], script, optional_params, body, secret_fields) + } + + /// [`before_send_with_secrets`] with `bindings` placed in the namespace before `script` runs. + fn before_send_bound( + bindings: &[(&str, &Value)], + script: &CStr, + optional_params: Value, + body: Value, + secret_fields: &[&str], + ) -> WireRequest { + Python::initialize(); + Python::attach(|py| { + let locals = namespace(py, PAYLOAD_LOGGER); + for &(name, value) in bindings { + locals.set_item(name, to_py(py, value).unwrap()).unwrap(); + } + run(py, &locals, script); + let mut logging = LegacyLogging { + logger: Some(PythonLogger::new(local(&locals, "logger").unbind())), + ..legacy_call(py, &locals, false) + }; + let context = RequestContext { + model: "model".into(), + custom_llm_provider: "provider".into(), + optional_params, + secret_fields: secret_fields.iter().map(|name| name.to_string()).collect(), + api_key: Some(SecretValue::new("route-key")), + }; + let wire = WireRequest { + url: "https://provider.invalid/ocr".into(), + headers: vec![("x-route".into(), "route".into())], + body, + }; + let step = logging.before_send(py, Box::new(wire), &context).unwrap(); + let raw = MachineEvent::ResponseReceived { + raw: RawResponse { + body: "raw response".into(), + }, + }; + assert!(matches!( + logging.emit(py, LifecycleEvent::Machine(&raw)).unwrap(), + LifecycleStep::Done + )); + run(py, &locals, c"check()"); + let LifecycleStep::Wire(wire) = step else { + panic!("before_send did not hand back the wire request"); + }; + *wire + }) + } + + #[rstest] + #[case::caller_keyword(c" +document = {'type': 'document_url', 'document_url': 'data:application/pdf;base64,YWJj'} +pages = [0] +kwargs = {'document': document, 'pages': pages} +observed = [] +on_pre_call = lambda args: observed.append( + (args['complete_input_dict']['document'] is document, args['complete_input_dict']['pages'] is pages) +) +def check(): + assert observed == [(True, True)], observed +")] + #[case::request_attribute_behind_an_omitted_keyword(c" +document = {'type': 'document_url', 'document_url': 'data:application/pdf;base64,YWJj'} +pages = [0] +request.document = document +kwargs = {'pages': pages} +observed = [] +on_pre_call = lambda args: observed.append( + (args['complete_input_dict']['document'] is document, args['complete_input_dict']['pages'] is pages) +) +def check(): + assert observed == [(True, True)], observed +")] + fn passthrough_keys_reach_pre_call_as_the_callers_own_objects(#[case] script: &CStr) { + let body = json!({"model": "model", "document": document(DOCUMENT), "pages": [0]}); + let wire = before_send(script, body.clone()); + assert_eq!(wire.body, body); + } + + #[test] + fn pre_call_edit_of_a_passthrough_object_reaches_the_caller_and_the_wire() { + let wire = before_send( + c" +document = {'type': 'document_url', 'document_url': 'data:application/pdf;base64,YWJj'} +kwargs = {'document': document} +def on_pre_call(args): + args['complete_input_dict']['document']['document_url'] = 'data:application/pdf;base64,ZWRpdGVk' +def check(): + assert document['document_url'] == 'data:application/pdf;base64,ZWRpdGVk' +", + json!({"document": document(DOCUMENT)}), + ); + assert_eq!(wire.body["document"], document(EDITED)); + } + + #[test] + fn a_body_key_the_route_rewrote_is_not_the_callers_object() { + let wire = before_send( + c" +document = {'type': 'document_url', 'document_url': 'https://example.invalid/scan.pdf'} +kwargs = {'document': document} +observed = [] +def on_pre_call(args): + observed.append(args['complete_input_dict']['document'] is document) + args['complete_input_dict']['document']['document_name'] = 'edited.pdf' +def check(): + assert observed == [False], observed + assert document == {'type': 'document_url', 'document_url': 'https://example.invalid/scan.pdf'} +", + json!({"document": document(DOCUMENT)}), + ); + assert_eq!( + wire.body["document"], + json!({"type": "document_url", "document_url": DOCUMENT, "document_name": "edited.pdf"}) + ); + } + + #[test] + fn a_caller_value_with_no_json_form_is_left_out_of_realiasing() { + let body = json!({"pages": [0]}); + let wire = before_send( + c" +opaque = object() +kwargs = {'pages': opaque} +observed = [] +on_pre_call = lambda args: observed.append(args['complete_input_dict']['pages']) +def check(): + assert observed == [[0]], observed +", + body.clone(), + ); + assert_eq!(wire.body, body); + } + + #[rstest] + #[case::body( + c" +def on_pre_call(args): + args['complete_input_dict'] = {'replacement': True} +" + )] + #[case::headers( + c" +def on_pre_call(args): + args['headers'] = {'x-replacement': 'yes'} +" + )] + fn rebinding_the_payload_envelope_does_not_reach_the_wire(#[case] script: &CStr) { + let body = json!({"document": document(DOCUMENT)}); + let wire = before_send(script, body.clone()); + assert_eq!(wire.body, body); + assert_eq!(wire.headers, [("x-route".to_string(), "route".to_string())]); + } + + #[test] + fn pre_call_header_edit_reaches_the_wire() { + let wire = before_send( + c" +def on_pre_call(args): + args['headers']['x-callback'] = 'edited' +", + json!({}), + ); + assert_eq!( + wire.headers, + [ + ("x-route".to_string(), "route".to_string()), + ("x-callback".to_string(), "edited".to_string()), + ] + ); + } + + #[test] + fn pre_call_receives_the_wire_request_and_the_logger_its_redacted_request() { + let body = json!({"model": "model", "document": document(DOCUMENT)}); + before_send_with_secrets( + c" +logger_fn = lambda *args: None +kwargs = { + 'litellm_call_id': 'call-1', + 'client_secret': 'shh', + 'proxy_server_request': {'body': {}}, + 'logger_fn': logger_fn, + 'litellm_request_debug': True, + 'ocr_cost_per_page': 0.05, +} +observed = [] +on_pre_call = observed.append +def check(): + [args] = observed + assert args['api_base'] == 'https://provider.invalid/ocr', args + assert args['complete_input_dict'] == { + 'model': 'model', + 'document': {'type': 'document_url', 'document_url': 'data:application/pdf;base64,YWJj'}, + }, args + update = logger.update + assert update['model'] == 'model' and update['custom_llm_provider'] == 'provider', update + assert update['litellm_params']['litellm_call_id'] == 'call-1', update + assert update['litellm_params']['api_base'] == 'https://provider.invalid/ocr', update + assert update['litellm_params']['logger_fn'] is logger_fn, update + assert update['litellm_params']['litellm_request_debug'] is True, update + assert update['litellm_params']['ocr_cost_per_page'] == 0.05, update + assert update['kwargs']['client_secret'] == '****', update + assert 'proxy_server_request' not in update['kwargs'], update + assert update['optional_params']['client_secret'] == '****', update +", + json!({"client_secret": "shh"}), + body, + &["client_secret"], + ); + } + + #[rstest] + #[case::added_key( + c" +def on_pre_call(args): + args['complete_input_dict']['include_image_base64'] = True +", + json!({"document": document(DOCUMENT), "include_image_base64": true}) + )] + #[case::replaced_document( + c" +document = {'type': 'document_url', 'document_url': 'data:application/pdf;base64,YWJj'} +kwargs = {'document': document} +def on_pre_call(args): + args['complete_input_dict']['document'] = { + 'type': 'document_url', 'document_url': 'data:application/pdf;base64,ZWRpdGVk' + } +def check(): + assert document['document_url'] == 'data:application/pdf;base64,YWJj', document +", + json!({"document": document(EDITED)}) + )] + #[case::retained_body_edited_after_rebinding( + c" +def on_pre_call(args): + retained = args['complete_input_dict'] + args['complete_input_dict'] = {'rebound': True} + retained['include_image_base64'] = True +", + json!({"document": document(DOCUMENT), "include_image_base64": true}) + )] + fn pre_call_body_edits_reach_the_wire(#[case] script: &CStr, #[case] expected: Value) { + let body = json!({"document": document(DOCUMENT)}); + let wire = before_send(script, body); + assert_eq!(wire.body, expected); + } + + #[test] + fn retained_headers_edited_after_rebinding_reach_the_wire() { + let wire = before_send( + c" +def on_pre_call(args): + retained = args['headers'] + args['headers'] = {'x-rebound': 'rebound'} + retained['x-retained'] = 'sent' +", + json!({}), + ); + assert_eq!( + wire.headers, + [ + ("x-route".to_string(), "route".to_string()), + ("x-retained".to_string(), "sent".to_string()), + ] + ); + } + + #[test] + fn post_call_receives_the_raw_response_the_route_key_and_the_body_and_headers_pre_call_saw() { + before_send( + c" +def check(): + original_response, api_key, additional_args = logger.post + assert original_response == 'raw response', original_response + assert api_key == logger.pre_api_key == 'route-key', (api_key, logger.pre_api_key) + assert additional_args == { + 'complete_input_dict': logger.pre['complete_input_dict'], + 'headers': logger.pre['headers'], + }, additional_args + assert additional_args['complete_input_dict'] is logger.pre['complete_input_dict'] + assert additional_args['headers'] is logger.pre['headers'] +", + json!({"document": document(DOCUMENT)}), + ); + } + + #[test] + fn every_request_runs_the_full_pre_call_and_post_call() { + let wire = before_send( + c" +def on_pre_call(args): + args['complete_input_dict']['include_image_base64'] = True +def check(): + assert logger.names() == ['pre_call', 'post_call'], logger.calls +", + json!({"document": document(DOCUMENT)}), + ); + assert_eq!( + wire.body, + json!({"document": document(DOCUMENT), "include_image_base64": true}) + ); + } + + /// What one pre-call callback does to the payload it is handed. + #[derive(Clone, Debug)] + enum Edit { + Nothing, + Set(String, Value), + Remove(String), + Rebind(Value), + RebindThenSetRetained(String, Value), + } + + impl Edit { + fn script(&self) -> Value { + match self { + Self::Nothing => json!({"kind": "nothing"}), + Self::Set(key, value) => json!({"kind": "set", "key": key, "value": value}), + Self::Remove(key) => json!({"kind": "remove", "key": key}), + Self::Rebind(value) => json!({"kind": "rebind", "value": value}), + Self::RebindThenSetRetained(key, value) => { + json!({"kind": "rebind_then_set_retained", "key": key, "value": value}) + } + } + } + + /// The legacy contract: the provider is sent the body object `pre_call` received, as + /// the callback left it. Rebinding the envelope's key points the envelope elsewhere and + /// leaves that object alone. + fn sent(&self, body: &Map) -> Value { + let mut sent = body.clone(); + match self { + Self::Nothing | Self::Rebind(_) => {} + Self::Set(key, value) | Self::RebindThenSetRetained(key, value) => { + sent.insert(key.clone(), value.clone()); + } + Self::Remove(key) => { + sent.remove(key); + } + } + Value::Object(sent) + } + } + + /// How the caller's keyword for a body key relates to what the route sends under it. + #[derive(Clone, Copy, Debug, PartialEq, Eq)] + enum Caller { + PassedUnchanged, + RewrittenByTheRoute, + NotPassed, + } + + const MODEL: &CStr = c" +aliased = {} +def on_pre_call(args): + body = args['complete_input_dict'] + aliased.update({name: body[name] is kwargs[name] for name in unchanged}) + kind = edit['kind'] + if kind == 'set': + body[edit['key']] = edit['value'] + elif kind == 'remove': + body.pop(edit['key'], None) + elif kind == 'rebind': + args['complete_input_dict'] = edit['value'] + elif kind == 'rebind_then_set_retained': + args['complete_input_dict'] = {} + body[edit['key']] = edit['value'] +def check(): + assert aliased == {name: True for name in unchanged}, aliased + assert logger.names() == ['pre_call', 'post_call'], logger.calls +"; + + fn json_value() -> impl Strategy { + let leaf = prop_oneof![ + Just(Value::Null), + any::().prop_map(Value::from), + any::().prop_map(Value::from), + any::() + .prop_filter("JSON has no NaN or infinity", |number| number.is_finite()) + .prop_map(Value::from), + ".{0,8}".prop_map(Value::from), + ]; + leaf.prop_recursive(3, 24, 4, |inner| { + prop_oneof![ + prop::collection::vec(inner.clone(), 0..4).prop_map(Value::from), + prop::collection::btree_map(key(), inner, 0..4) + .prop_map(|fields| Value::Object(fields.into_iter().collect())), + ] + }) + } + + fn key() -> impl Strategy { + "[a-z]{1,6}" + } + + fn caller() -> impl Strategy { + prop_oneof![ + Just(Caller::PassedUnchanged), + Just(Caller::RewrittenByTheRoute), + Just(Caller::NotPassed), + ] + } + + fn edit() -> impl Strategy { + prop_oneof![ + Just(Edit::Nothing), + (key(), json_value()).prop_map(|(key, value)| Edit::Set(key, value)), + key().prop_map(Edit::Remove), + json_value().prop_map(Edit::Rebind), + (key(), json_value()).prop_map(|(key, value)| Edit::RebindThenSetRetained(key, value)), + ] + } + + proptest! { + #![proptest_config(ProptestConfig::with_cases(128))] + + /// For any body, any caller keywords and any callback edit: every keyword the route + /// sends unchanged reaches `pre_call` as the caller's own object, and the provider is + /// sent exactly what the model says, so a callback that edits nothing changes nothing. + #[test] + fn the_wire_is_the_body_pre_call_received_as_the_callback_left_it( + fields in prop::collection::btree_map(key(), (json_value(), caller()), 0..5), + edit in edit(), + ) { + let body: Map = fields + .iter() + .map(|(name, (value, _))| (name.clone(), value.clone())) + .collect(); + let kwargs: Map = fields + .iter() + .filter_map(|(name, (value, caller))| match caller { + Caller::PassedUnchanged => Some((name.clone(), value.clone())), + Caller::RewrittenByTheRoute => Some((name.clone(), json!([value]))), + Caller::NotPassed => None, + }) + .collect(); + let unchanged: Value = fields + .iter() + .filter(|(_, (_, caller))| *caller == Caller::PassedUnchanged) + .map(|(name, _)| Value::from(name.clone())) + .collect(); + + let wire = before_send_bound( + &[ + ("kwargs", &Value::Object(kwargs)), + ("unchanged", &unchanged), + ("edit", &edit.script()), + ], + MODEL, + json!({}), + Value::Object(body.clone()), + &[], + ); + + prop_assert_eq!(wire.body, edit.sent(&body)); + prop_assert_eq!(wire.headers, [("x-route".to_string(), "route".to_string())]); + } + } +} + #[cfg(test)] -#[path = "../tests/terminal.rs"] -mod terminal_tests; +mod terminal_tests { + use std::ffi::CStr; + + use litellm_host::event::{FailureOrigin, Timing}; + use litellm_host_python::{LifecycleEvent, LifecycleStep, PythonLifecycle}; + use pyo3::exceptions::PyRuntimeError; + use pyo3::exceptions::asyncio::CancelledError; + use pyo3::prelude::*; + use pyo3::types::PyDict; + use rstest::rstest; + + use super::LegacyLogging; + use crate::PythonLogger; + use crate::test_support::{legacy_call, local, namespace, run}; + + const TIMING: Timing = Timing { + start_time: 0.0, + end_time: 1.0, + }; + + fn logged(py: Python<'_>, locals: &Bound<'_, PyDict>, asynchronous: bool) -> LegacyLogging { + LegacyLogging { + logger: Some(PythonLogger::new(local(locals, "logger").unbind())), + ..legacy_call(py, locals, asynchronous) + } + } + + fn succeed( + py: Python<'_>, + locals: &Bound<'_, PyDict>, + logging: &mut LegacyLogging, + ) -> LifecycleStep { + let response = local(locals, "response").unbind(); + logging + .emit( + py, + LifecycleEvent::Succeeded { + timing: TIMING, + response: &response, + }, + ) + .unwrap() + } + + fn fail( + py: Python<'_>, + locals: &Bound<'_, PyDict>, + logging: &mut LegacyLogging, + ) -> LifecycleStep { + let failure = PyErr::from_value(local(locals, "failure")); + logging + .emit( + py, + LifecycleEvent::Failed { + timing: TIMING, + origin: FailureOrigin::Host, + error: &failure, + }, + ) + .unwrap() + } + + #[rstest] + #[case::sync_listened(false, c"", &["submit"])] + #[case::async_listened( + true, + c"", + &["async_success_handler", "enqueued", "sync_success_for_async_call"] + )] + #[case::async_deferred(true, c"logger._defer_async_logging = True", &["sync_success_for_async_call"])] + #[case::async_with_fallbacks(true, c"kwargs = {'fallbacks': ['other']}", &["sync_success_for_async_call"])] + fn success_reaches_the_logging_handlers( + #[case] asynchronous: bool, + #[case] script: &CStr, + #[case] expected: &[&str], + ) { + Python::initialize(); + Python::attach(|py| { + let locals = namespace(py, c"response = object()"); + run(py, &locals, script); + let mut logging = logged(py, &locals, asynchronous); + assert!(matches!( + succeed(py, &locals, &mut logging), + LifecycleStep::Done + )); + let names: Vec = local(&locals, "logger") + .call_method0("names") + .unwrap() + .extract() + .unwrap(); + assert_eq!(names, expected); + run( + py, + &locals, + c" +assert all(value is response for name, value in logger.calls if name.endswith('_handler')) +assert hasattr(logger, '_native_pending_logging') == getattr(logger, '_defer_async_logging', False) +", + ); + }); + } + + #[rstest] + #[case::synchronous(false, &["failure_handler"])] + #[case::asynchronous(true, &[])] + fn internal_calls_skip_failure_callbacks_only_when_asynchronous( + #[case] asynchronous: bool, + #[case] expected: &[&str], + ) { + Python::initialize(); + Python::attach(|py| { + let locals = namespace(py, c"failure = ValueError('provider')"); + let mut logging = LegacyLogging { + internal: true, + ..logged(py, &locals, asynchronous) + }; + assert!(matches!( + fail(py, &locals, &mut logging), + LifecycleStep::Done + )); + let names: Vec = local(&locals, "logger") + .call_method0("names") + .unwrap() + .extract() + .unwrap(); + assert_eq!(names, expected); + }); + } + + #[test] + fn internal_async_calls_skip_the_async_success_fan_out() { + Python::initialize(); + Python::attach(|py| { + let locals = namespace(py, c"response = object()"); + let mut logging = LegacyLogging { + internal: true, + ..logged(py, &locals, true) + }; + succeed(py, &locals, &mut logging); + run( + py, + &locals, + c"assert logger.names() == ['sync_success_for_async_call'], logger.calls", + ); + }); + } + + #[test] + fn a_failing_success_callback_is_reported_without_replacing_the_response() { + Python::initialize(); + Python::attach(|py| { + let locals = namespace( + py, + c" +response = object() +failure = ValueError('terminal diagnostic') + +class FailingLogger(StubLogger): + def handle_sync_success_callbacks_for_async_calls(self, *args): + raise failure + +logger = FailingLogger() +", + ); + let mut logging = logged(py, &locals, true); + assert!(matches!( + succeed(py, &locals, &mut logging), + LifecycleStep::Done + )); + assert!( + logging + .response + .as_ref() + .unwrap() + .bind(py) + .is(local(&locals, "response")) + ); + run(py, &locals, c"assert unraisable_from(logger) == [failure]"); + }); + } + + #[rstest] + #[case::sync_listened(false, c"", &["failure_handler"])] + #[case::async_listened(true, c"", &["failure_handler", "async_failure_handler"])] + fn failure_reaches_the_logging_handlers( + #[case] asynchronous: bool, + #[case] script: &CStr, + #[case] expected: &[&str], + ) { + Python::initialize(); + Python::attach(|py| { + let locals = namespace(py, c"failure = ValueError('provider')"); + run(py, &locals, script); + let mut logging = logged(py, &locals, asynchronous); + let step = fail(py, &locals, &mut logging); + let awaits_async_handler = expected.contains(&"async_failure_handler"); + assert_eq!( + matches!(step, LifecycleStep::Await(_)), + awaits_async_handler + ); + let names: Vec = local(&locals, "logger") + .call_method0("names") + .unwrap() + .extract() + .unwrap(); + assert_eq!(names, expected); + run( + py, + &locals, + c"assert all(value is failure for name, value in logger.calls if name.endswith('_handler'))", + ); + }); + } + + #[test] + fn a_failing_sync_failure_callback_keeps_the_error_and_still_runs_the_async_family() { + Python::initialize(); + Python::attach(|py| { + let locals = namespace( + py, + c" +failure = ValueError('selected') + +class FailingLogger(StubLogger): + def failure_handler(self, error, trace, start, end): + self.record('failure_handler', error) + raise RuntimeError('handler failed') + +logger = FailingLogger() +", + ); + let mut logging = logged(py, &locals, true); + assert!(matches!( + fail(py, &locals, &mut logging), + LifecycleStep::Await(_) + )); + assert!( + logging + .error + .as_ref() + .unwrap() + .bind(py) + .is(local(&locals, "failure")) + ); + run( + py, + &locals, + c"assert logger.names() == ['failure_handler', 'async_failure_handler'], logger.calls", + ); + }); + } + + #[rstest] + #[case::completed(None, true)] + #[case::handler_error(Some(false), true)] + #[case::cancelled(Some(true), false)] + fn the_async_failure_handler_ends_the_call_unless_it_was_cancelled( + #[case] error: Option, + #[case] done: bool, + ) { + Python::initialize(); + Python::attach(|py| { + let locals = namespace(py, c"failure = ValueError('provider')"); + let mut logging = logged(py, &locals, true); + fail(py, &locals, &mut logging); + let result = match error { + None => Ok(py.None()), + Some(false) => Err(PyRuntimeError::new_err("handler failed")), + Some(true) => Err(CancelledError::new_err("cancelled")), + }; + let expected = result.as_ref().err().map(|error| error.value(py).clone()); + match logging.resume(py, result) { + Ok(step) => assert!(done && matches!(step, LifecycleStep::Done)), + Err(propagated) => { + assert!(!done); + assert!(propagated.value(py).is(expected.unwrap())); + } + } + }); + } + + #[test] + fn closing_restores_the_correlation_context_once() { + Python::initialize(); + Python::attach(|py| { + let locals = namespace(py, c""); + let mut logging = logged(py, &locals, true); + logging.close(py); + logging.close(py); + run( + py, + &locals, + c"assert logger.names() == ['restore'], logger.calls", + ); + }); + } +} diff --git a/litellm-rust/crates/callbacks-legacy-python/src/deferred.rs b/litellm-rust/crates/callbacks-legacy-python/src/deferred.rs index b18012f926e..292648ba69c 100644 --- a/litellm-rust/crates/callbacks-legacy-python/src/deferred.rs +++ b/litellm-rust/crates/callbacks-legacy-python/src/deferred.rs @@ -63,5 +63,151 @@ impl PendingLogging { } #[cfg(test)] -#[path = "../tests/deferred.rs"] -mod tests; +mod tests { + use std::ffi::CStr; + + use pyo3::prelude::*; + use pyo3::types::PyDict; + use rstest::rstest; + + use super::{PendingLogging, PendingSuccess}; + use crate::PythonLogger; + use crate::test_support::{local, namespace, run}; + + /// A deferred success for the namespace's `logger` and `response`, bound as `pending`. + fn defer<'py>(py: Python<'py>, script: &CStr) -> Bound<'py, PyDict> { + let locals = namespace(py, c"response = object()"); + run(py, &locals, script); + let pending = Py::new( + py, + PendingLogging { + pending: Some(PendingSuccess { + logger: PythonLogger::new(local(&locals, "logger").unbind()), + response: Some(local(&locals, "response").unbind()), + start: py.None(), + end: Some(py.None()), + }), + }, + ) + .unwrap(); + locals.set_item("pending", pending).unwrap(); + locals + } + + #[test] + fn release_enqueues_the_success_once_in_the_releasing_context() { + Python::initialize(); + Python::attach(|py| { + let locals = defer( + py, + c" +from contextvars import ContextVar + +marker = ContextVar('marker', default='unset') +observed = [] + +def on_enqueue(coroutine): + observed.append(marker.get()) + pending.release(True) + +logger.on_enqueue = on_enqueue +", + ); + run( + py, + &locals, + c" +marker.set('release') +pending.release(True) +pending.release(True) +assert observed == ['release'], observed +assert logger.names() == ['async_success_handler', 'enqueued'], logger.calls +assert logger.calls[0][1] is response +", + ); + }); + } + + #[test] + fn a_blocked_release_drops_the_success_for_good() { + Python::initialize(); + Python::attach(|py| { + let locals = defer(py, c""); + run( + py, + &locals, + c" +pending.release(False) +pending.release(True) +assert logger.calls == [], logger.calls +", + ); + }); + } + + #[rstest] + #[case::ordinary_error(c"RuntimeError('queue full')", false)] + #[case::cancellation(c"asyncio.CancelledError()", true)] + fn a_failed_enqueue_closes_the_coroutine_and_is_never_replayed( + #[case] failure: &CStr, + #[case] propagates: bool, + ) { + Python::initialize(); + Python::attach(|py| { + let locals = defer( + py, + c" +import asyncio + +def on_enqueue(coroutine): + raise failure + +logger.on_enqueue = on_enqueue +", + ); + locals + .set_item("failure", py.eval(failure, None, Some(&locals)).unwrap()) + .unwrap(); + let released = local(&locals, "pending").call_method1("release", (true,)); + match released { + Ok(_) => assert!(!propagates), + Err(error) => { + assert!(propagates); + assert!(error.value(py).is(local(&locals, "failure"))); + } + } + locals.set_item("propagates", propagates).unwrap(); + run( + py, + &locals, + c" +pending.release(True) +assert logger.names() == ['async_success_handler', 'enqueued', 'closed'], logger.calls +assert unraisable_from(logger) == ([] if propagates else [failure]) +", + ); + }); + } + + #[test] + fn an_unreleased_success_does_not_keep_its_logger_alive() { + Python::initialize(); + Python::attach(|py| { + let locals = defer(py, c""); + run( + py, + &locals, + c" +import gc +import weakref + +logger.pending = pending +reference = weakref.ref(logger) +del logger, pending +gc.collect() +assert reference() is None +", + ); + }); + } +} diff --git a/litellm-rust/crates/callbacks-legacy-python/src/lib.rs b/litellm-rust/crates/callbacks-legacy-python/src/lib.rs index 44393792d1f..030bf03d4ba 100644 --- a/litellm-rust/crates/callbacks-legacy-python/src/lib.rs +++ b/litellm-rust/crates/callbacks-legacy-python/src/lib.rs @@ -16,13 +16,218 @@ mod deferred; mod logger; mod preparation; mod python; -#[cfg(test)] -#[path = "../tests/support.rs"] -mod test_support; - pub(crate) use adapter::LegacyLogging; pub use adapter::{LegacySurface, PassThroughStream}; pub use call::{PublicCall, run_legacy_call}; pub(crate) use callbacks::{LegacyCallbacks, is_internal_call}; pub(crate) use logger::{DeploymentHooks, PythonLogger, finalize, setup}; pub(crate) use preparation::prepare; + +#[cfg(test)] +mod test_support { + use std::ffi::CStr; + + use pyo3::prelude::*; + use pyo3::types::{PyDict, PyTuple}; + + use crate::{LegacyLogging, LegacySurface, PublicCall}; + + /// The parameters of every `callbacks_legacy_python` function, as the real module declares them. + /// `tests/test_litellm/rust_bridge/test_callbacks_legacy_python.py` pins this file to the Python + /// signatures, and [`namespace`] binds every fake call against it. + pub(crate) const PYTHON_CONTRACT: &str = include_str!("../python_contract.json"); + + /// Stand-ins for `callbacks_legacy_python`, the only Python module the crate calls. Tests + /// share one interpreter and run concurrently, so each fake is installed idempotently and + /// forwards to the per-test `StubLogger` it is handed (directly, or as `kwargs['logger']`). + /// Every fake is bound against the contract first, so a call the real module would reject + /// fails here too. + const STUBS: &CStr = c" +import contextvars +import inspect +import json +import sys +import traceback +import types + +for name in ('litellm', 'litellm.rust_bridge', 'litellm.rust_bridge.callbacks_legacy_python'): + sys.modules.setdefault(name, types.ModuleType(name)) + +legacy = sys.modules['litellm.rust_bridge.callbacks_legacy_python'] +CONTRACT = json.loads(python_contract) + + +def contracted(name, fake): + signature = inspect.Signature( + [inspect.Parameter(parameter, inspect.Parameter.POSITIONAL_OR_KEYWORD) for parameter in CONTRACT[name]] + ) + + def checked(*args, **kwargs): + signature.bind(*args, **kwargs) + return fake(*args, **kwargs) + + return checked + + +if not hasattr(legacy, 'is_internal'): + legacy.is_internal = contextvars.ContextVar('is_internal_call', default=False) + +FAKES = { + 'setup': lambda call_type, args, kwargs, start, asynchronous: types.SimpleNamespace( + logger=kwargs['logger_factory'](kwargs) if 'logger_factory' in kwargs else kwargs['logger'], + kwargs=kwargs, + ), + 'check_limits': lambda arguments: arguments['logger'].check_limits(arguments), + 'finalize': lambda response, logger, kwargs, start, end: logger.record('finalize', response), + 'update_logging': lambda logger, kwargs, model, optional_params, litellm_params, provider: logger.update_from_kwargs( + kwargs=kwargs, + model=model, + optional_params=optional_params, + litellm_params=litellm_params, + custom_llm_provider=provider, + ), + 'pre_call': lambda logger, input, api_key, additional_args: logger.pre_call(input, api_key, additional_args), + 'post_call': lambda logger, original_response, api_key, additional_args: logger.post_call( + original_response, api_key, additional_args + ), + 'defers_async_logging': lambda logger: bool(getattr(logger, '_defer_async_logging', False)), + 'defer_success': lambda logger, pending: setattr(logger, '_native_pending_logging', pending), + 'sync_success_for_async_call': lambda logger, response, start, end: logger.handle_sync_success_callbacks_for_async_calls( + response, start, end + ), + 'failure_handler': lambda logger, error, start, end, asynchronous: ( + logger.async_failure_handler if asynchronous else logger.failure_handler + )(error, ''.join(traceback.format_exception(error)), start, end), + 'submit_success': lambda logger, response, start, end: logger.record('submit', (response, start, end)), + 'async_success_handler': lambda logger, response, start, end: logger.async_success_handler(response, start, end), + 'enqueue_logging': lambda coroutine: coroutine.enqueue(), + 'restore_context': lambda logger: logger.record('restore', None), + 'custom_pricing_fields': lambda: ('ocr_cost_per_page',), + 'is_internal_call': lambda: legacy.is_internal.get(), + 'credential_list': lambda: [], + 'warn_unknown_credential': lambda name, loaded: None, + 'before_deployment_call': lambda kwargs, call_type: kwargs['logger'].hook('pre', kwargs, call_type), + 'after_deployment_success': lambda kwargs, response, call_type: kwargs['logger'].hook( + 'success', response, call_type + ), + 'after_deployment_failure': lambda kwargs, error, call_type: kwargs['logger'].hook('failure', error, call_type), + 'stream_opened': lambda logger: logger.record('stream_opened', None), + 'stream_success': lambda logger, request_body, chunks, start, end, first_chunk: logger.record( + 'stream_success', list(chunks) + ), + 'stream_failure': lambda logger, request_body, chunks, error: logger.record('stream_failure', error), +} +assert FAKES.keys() == CONTRACT.keys(), sorted(FAKES.keys() ^ CONTRACT.keys()) +for name, fake in FAKES.items(): + setattr(legacy, name, contracted(name, fake)) + + +unraisable = sys.modules.setdefault( + 'litellm_test_unraisable', types.ModuleType('litellm_test_unraisable') +) +if not hasattr(unraisable, 'events'): + unraisable.events = [] + sys.unraisablehook = lambda event: unraisable.events.append((event.object, event.exc_value)) + + +def unraisable_from(owner): + return [error for source, error in unraisable.events if source is owner] + + +class StubCoroutine: + def __init__(self, logger): + self.logger = logger + + def enqueue(self): + self.logger.record('enqueued', None) + self.logger.on_enqueue(self) + + def close(self): + self.logger.record('closed', None) + + +class StubLogger: + def __init__(self): + self.calls = [] + self.hooks = {} + self.on_enqueue = lambda coroutine: None + + def record(self, name, value): + self.calls.append((name, value)) + + def names(self): + return [name for name, _ in self.calls] + + def hook(self, phase, value, call_type): + self.record(phase + '_hook', call_type) + return self.hooks.get(phase, lambda value: 'awaitable')(value) + + def check_limits(self, arguments): + self.record('check_limits', arguments) + + def failure_handler(self, error, trace, start, end): + self.record('failure_handler', error) + + def async_failure_handler(self, error, trace, start, end): + self.record('async_failure_handler', error) + return 'awaitable' + + def success_handler(self, response, start, end): + self.record('success_handler', response) + + def async_success_handler(self, response, start, end): + self.record('async_success_handler', response) + return StubCoroutine(self) + + def handle_sync_success_callbacks_for_async_calls(self, response, start, end): + self.record('sync_success_for_async_call', response) + + +logger = StubLogger() +"; + + /// A namespace with the stubs, `StubLogger` and a fresh `logger`, after `script` ran in it. + pub(crate) fn namespace<'py>(py: Python<'py>, script: &CStr) -> Bound<'py, PyDict> { + let locals = PyDict::new(py); + locals.set_item("python_contract", PYTHON_CONTRACT).unwrap(); + py.run(STUBS, Some(&locals), Some(&locals)).unwrap(); + py.run(script, Some(&locals), Some(&locals)).unwrap(); + locals + } + + pub(crate) fn run(py: Python<'_>, locals: &Bound<'_, PyDict>, code: &CStr) { + py.run(code, Some(locals), Some(locals)).unwrap(); + } + + pub(crate) fn local<'py>(locals: &Bound<'py, PyDict>, name: &str) -> Bound<'py, PyAny> { + locals.get_item(name).unwrap().unwrap() + } + + /// A legacy call over the namespace's `kwargs` (or none) and `request` (or `None`). + pub(crate) fn legacy_call( + py: Python<'_>, + locals: &Bound<'_, PyDict>, + asynchronous: bool, + ) -> LegacyLogging { + let request = locals + .get_item("request") + .unwrap() + .unwrap_or_else(|| py.None().into_bound(py)); + let kwargs = locals + .get_item("kwargs") + .unwrap() + .map(|kwargs| kwargs.cast_into::().unwrap()) + .unwrap_or_else(|| PyDict::new(py)); + let call = PublicCall::capture(&request, &PyTuple::empty(py), &kwargs).unwrap(); + LegacyLogging::new( + py, + LegacySurface { + call_type: "test", + input_description: "test input", + stream: None, + }, + call, + asynchronous, + ) + } +} diff --git a/litellm-rust/crates/callbacks-legacy-python/tests/deferred.rs b/litellm-rust/crates/callbacks-legacy-python/tests/deferred.rs deleted file mode 100644 index 289ea1b2e7f..00000000000 --- a/litellm-rust/crates/callbacks-legacy-python/tests/deferred.rs +++ /dev/null @@ -1,146 +0,0 @@ -use std::ffi::CStr; - -use pyo3::prelude::*; -use pyo3::types::PyDict; -use rstest::rstest; - -use super::{PendingLogging, PendingSuccess}; -use crate::PythonLogger; -use crate::test_support::{local, namespace, run}; - -/// A deferred success for the namespace's `logger` and `response`, bound as `pending`. -fn defer<'py>(py: Python<'py>, script: &CStr) -> Bound<'py, PyDict> { - let locals = namespace(py, c"response = object()"); - run(py, &locals, script); - let pending = Py::new( - py, - PendingLogging { - pending: Some(PendingSuccess { - logger: PythonLogger::new(local(&locals, "logger").unbind()), - response: Some(local(&locals, "response").unbind()), - start: py.None(), - end: Some(py.None()), - }), - }, - ) - .unwrap(); - locals.set_item("pending", pending).unwrap(); - locals -} - -#[test] -fn release_enqueues_the_success_once_in_the_releasing_context() { - Python::initialize(); - Python::attach(|py| { - let locals = defer( - py, - c" -from contextvars import ContextVar - -marker = ContextVar('marker', default='unset') -observed = [] - -def on_enqueue(coroutine): - observed.append(marker.get()) - pending.release(True) - -logger.on_enqueue = on_enqueue -", - ); - run( - py, - &locals, - c" -marker.set('release') -pending.release(True) -pending.release(True) -assert observed == ['release'], observed -assert logger.names() == ['async_success_handler', 'enqueued'], logger.calls -assert logger.calls[0][1] is response -", - ); - }); -} - -#[test] -fn a_blocked_release_drops_the_success_for_good() { - Python::initialize(); - Python::attach(|py| { - let locals = defer(py, c""); - run( - py, - &locals, - c" -pending.release(False) -pending.release(True) -assert logger.calls == [], logger.calls -", - ); - }); -} - -#[rstest] -#[case::ordinary_error(c"RuntimeError('queue full')", false)] -#[case::cancellation(c"asyncio.CancelledError()", true)] -fn a_failed_enqueue_closes_the_coroutine_and_is_never_replayed( - #[case] failure: &CStr, - #[case] propagates: bool, -) { - Python::initialize(); - Python::attach(|py| { - let locals = defer( - py, - c" -import asyncio - -def on_enqueue(coroutine): - raise failure - -logger.on_enqueue = on_enqueue -", - ); - locals - .set_item("failure", py.eval(failure, None, Some(&locals)).unwrap()) - .unwrap(); - let released = local(&locals, "pending").call_method1("release", (true,)); - match released { - Ok(_) => assert!(!propagates), - Err(error) => { - assert!(propagates); - assert!(error.value(py).is(local(&locals, "failure"))); - } - } - locals.set_item("propagates", propagates).unwrap(); - run( - py, - &locals, - c" -pending.release(True) -assert logger.names() == ['async_success_handler', 'enqueued', 'closed'], logger.calls -assert unraisable_from(logger) == ([] if propagates else [failure]) -", - ); - }); -} - -#[test] -fn an_unreleased_success_does_not_keep_its_logger_alive() { - Python::initialize(); - Python::attach(|py| { - let locals = defer(py, c""); - run( - py, - &locals, - c" -import gc -import weakref - -logger.pending = pending -reference = weakref.ref(logger) -del logger, pending -gc.collect() -assert reference() is None -", - ); - }); -} diff --git a/litellm-rust/crates/callbacks-legacy-python/tests/deployment_hooks.rs b/litellm-rust/crates/callbacks-legacy-python/tests/deployment_hooks.rs deleted file mode 100644 index 52c5e47f83f..00000000000 --- a/litellm-rust/crates/callbacks-legacy-python/tests/deployment_hooks.rs +++ /dev/null @@ -1,282 +0,0 @@ -use std::ffi::CStr; - -use litellm_host::event::{FailureOrigin, Timing}; -use litellm_host_python::{LifecycleEvent, LifecycleStep, PythonLifecycle}; -use pyo3::exceptions::asyncio::CancelledError; -use pyo3::prelude::*; -use pyo3::types::PyDict; -use rstest::rstest; - -use super::LegacyLogging; -use crate::test_support::{legacy_call, local, namespace, run}; - -const CALL: &CStr = c" -document = {'type': 'document_url', 'document_url': 'data:application/pdf;base64,YWJj'} -kwargs = {'logger': logger, 'document': document} -"; - -const TIMING: Timing = Timing { - start_time: 0.0, - end_time: 1.0, -}; - -fn begin<'py>( - py: Python<'py>, - locals: &Bound<'py, PyDict>, - asynchronous: bool, -) -> (LegacyLogging, LifecycleStep) { - let mut logging = legacy_call(py, locals, asynchronous); - let kwargs = local(locals, "kwargs") - .cast_into::() - .unwrap() - .unbind(); - let step = logging.begin(py, kwargs, 0.0).unwrap(); - (logging, step) -} - -fn arguments<'py>(py: Python<'py>, step: LifecycleStep) -> Bound<'py, PyDict> { - let LifecycleStep::Arguments(arguments) = step else { - panic!("expected the prepared arguments"); - }; - arguments.into_bound(py) -} - -fn awaits_deployment_hook(step: &LifecycleStep) -> bool { - matches!(step, LifecycleStep::Await(_)) -} - -#[rstest] -#[case::synchronous(false)] -#[case::asynchronous(true)] -fn deployment_pre_call_hook_runs_only_for_asynchronous_calls(#[case] asynchronous: bool) { - Python::initialize(); - Python::attach(|py| { - let locals = namespace(py, CALL); - let (_, step) = begin(py, &locals, asynchronous); - assert_eq!(awaits_deployment_hook(&step), asynchronous); - let names: Vec = local(&locals, "logger") - .call_method0("names") - .unwrap() - .extract() - .unwrap(); - assert_eq!(names.contains(&"pre_hook".to_string()), asynchronous); - }); -} - -#[test] -fn kwargs_returned_by_the_pre_call_hook_are_what_the_call_prepares() { - Python::initialize(); - Python::attach(|py| { - let locals = namespace( - py, - c" -document = {'type': 'document_url', 'document_url': 'data:application/pdf;base64,YWJj'} -replacement = {'type': 'document_url', 'document_url': 'data:application/pdf;base64,ZWRpdGVk'} -kwargs = {'logger': logger, 'document': document} -replaced_kwargs = {'logger': logger, 'document': replacement, 'pages': [0]} -", - ); - let (mut logging, step) = begin(py, &locals, true); - assert!(awaits_deployment_hook(&step)); - let step = logging - .resume(py, Ok(local(&locals, "replaced_kwargs").unbind())) - .unwrap(); - locals.set_item("prepared", arguments(py, step)).unwrap(); - run( - py, - &locals, - c" -assert prepared['document'] is replacement -assert prepared['pages'] is replaced_kwargs['pages'] -assert prepared['litellm_logging_obj'] is logger -assert 'litellm_logging_obj' not in replaced_kwargs -[checked] = [value for name, value in logger.calls if name == 'check_limits'] -assert checked is prepared -", - ); - }); -} - -#[rstest] -#[case::synchronous(false)] -#[case::asynchronous(true)] -fn a_keyword_the_bridge_never_reads_reaches_every_reader_as_the_callers_object( - #[case] asynchronous: bool, -) { - Python::initialize(); - Python::attach(|py| { - let locals = namespace( - py, - c" -opaque = object() -hooked = [] -logger.hooks = {'pre': lambda kwargs: hooked.append(kwargs['vendor_extension']) or kwargs} -kwargs = {'logger': logger, 'vendor_extension': opaque} -", - ); - let (mut logging, step) = begin(py, &locals, asynchronous); - let step = match step { - LifecycleStep::Await(hook_result) => logging.resume(py, Ok(hook_result)).unwrap(), - step => step, - }; - locals.set_item("prepared", arguments(py, step)).unwrap(); - locals.set_item("asynchronous", asynchronous).unwrap(); - run( - py, - &locals, - c" -assert prepared['vendor_extension'] is opaque -[checked] = [value for name, value in logger.calls if name == 'check_limits'] -assert checked['vendor_extension'] is opaque -assert hooked == ([opaque] if asynchronous else []), hooked -", - ); - }); -} - -#[test] -fn response_returned_by_the_post_call_hook_is_finalized_and_returned() { - Python::initialize(); - Python::attach(|py| { - let locals = namespace( - py, - c" -kwargs = {'logger': logger} -response = object() -replacement = object() -logger.hooks = {'pre': lambda kwargs: kwargs} -", - ); - let (mut logging, _) = begin(py, &locals, true); - logging - .resume(py, Ok(local(&locals, "kwargs").unbind())) - .unwrap(); - let step = logging - .after_success(py, local(&locals, "response").unbind(), TIMING) - .unwrap(); - assert!(awaits_deployment_hook(&step)); - let step = logging - .resume(py, Ok(local(&locals, "replacement").unbind())) - .unwrap(); - let LifecycleStep::Response(returned) = step else { - panic!("expected the finalized response"); - }; - assert!(returned.bind(py).is(local(&locals, "replacement"))); - run( - py, - &locals, - c" -[finalized] = [value for name, value in logger.calls if name == 'finalize'] -assert finalized is replacement -", - ); - }); -} - -#[rstest] -#[case::pre_call(false)] -#[case::post_call(true)] -fn cancelling_a_deployment_hook_ends_the_call_with_that_cancellation(#[case] post_call: bool) { - Python::initialize(); - Python::attach(|py| { - let locals = namespace(py, c"kwargs = {'logger': logger}\nresponse = object()"); - let (mut logging, _) = begin(py, &locals, true); - if post_call { - logging - .resume(py, Ok(local(&locals, "kwargs").unbind())) - .unwrap(); - logging - .after_success(py, local(&locals, "response").unbind(), TIMING) - .unwrap(); - } - let cancellation = CancelledError::new_err("cancelled"); - let cancelled = cancellation.value(py).clone(); - let error = logging.resume(py, Err(cancellation)).err().unwrap(); - assert!(error.value(py).is(&cancelled)); - let names: Vec = local(&locals, "logger") - .call_method0("names") - .unwrap() - .extract() - .unwrap(); - assert!(!names.iter().any(|name| name.contains("handler"))); - }); -} - -#[rstest] -#[case::hook_completed(false)] -#[case::hook_cancelled(true)] -fn failure_callbacks_run_after_the_failure_hook_however_it_ends(#[case] cancelled: bool) { - Python::initialize(); - Python::attach(|py| { - let locals = namespace( - py, - c"kwargs = {'logger': logger}\nfailure = ValueError('provider')", - ); - let (mut logging, _) = begin(py, &locals, true); - logging - .resume(py, Ok(local(&locals, "kwargs").unbind())) - .unwrap(); - let failure = PyErr::from_value(local(&locals, "failure")); - let failed = LifecycleEvent::Failed { - timing: TIMING, - origin: FailureOrigin::Call, - error: &failure, - }; - let step = logging.emit(py, failed).unwrap(); - assert!(awaits_deployment_hook(&step)); - let hook_result = if cancelled { - Err(CancelledError::new_err("cancelled")) - } else { - Ok(py.None()) - }; - assert!(matches!( - logging.resume(py, hook_result).unwrap(), - LifecycleStep::Await(_) - )); - run( - py, - &locals, - c" -assert logger.names()[-3:] == ['failure_hook', 'failure_handler', 'async_failure_handler'], logger.calls -assert all(value is failure for name, value in logger.calls if name.endswith('_handler')) -", - ); - }); -} - -#[rstest] -#[case::synchronous(false)] -#[case::asynchronous(true)] -fn a_limit_rejected_before_the_call_surfaces_as_the_callers_error(#[case] asynchronous: bool) { - Python::initialize(); - Python::attach(|py| { - let locals = namespace( - py, - c" -class BudgetExceeded(Exception): - pass - -rejection = BudgetExceeded('over budget') - -class LimitedLogger(StubLogger): - def check_limits(self, arguments): - raise rejection - -logger = LimitedLogger() -logger.hooks = {'pre': lambda kwargs: kwargs} -kwargs = {'logger': logger} -", - ); - let mut logging = legacy_call(py, &locals, asynchronous); - let kwargs = local(&locals, "kwargs") - .cast_into::() - .unwrap() - .unbind(); - let result = logging.begin(py, kwargs, 0.0).and_then(|step| match step { - LifecycleStep::Await(_) => logging.resume(py, Ok(local(&locals, "kwargs").unbind())), - step => Ok(step), - }); - let error = result.err().unwrap(); - assert!(error.value(py).is(local(&locals, "rejection"))); - }); -} diff --git a/litellm-rust/crates/callbacks-legacy-python/tests/payload.rs b/litellm-rust/crates/callbacks-legacy-python/tests/payload.rs deleted file mode 100644 index 5459b36af27..00000000000 --- a/litellm-rust/crates/callbacks-legacy-python/tests/payload.rs +++ /dev/null @@ -1,523 +0,0 @@ -use std::ffi::CStr; - -use litellm_auth::SecretValue; -use litellm_host::event::{MachineEvent, RawResponse, RequestContext, WireRequest}; -use litellm_host_python::{LifecycleEvent, LifecycleStep, PythonLifecycle, to_py}; -use proptest::prelude::*; -use pyo3::prelude::*; -use rstest::rstest; -use serde_json::{Map, Value, json}; - -use super::LegacyLogging; -use crate::PythonLogger; -use crate::test_support::{legacy_call, local, namespace, run}; - -/// The payload phases of `Logging` on top of `StubLogger`, with `pre_call` handing the -/// payload to the case's `on_pre_call`. -const PAYLOAD_LOGGER: &CStr = c" -class Request: - pass - -class PayloadLogger(StubLogger): - def update_from_kwargs(self, **update): - self.update = update - - def pre_call(self, input, api_key, additional_args): - self.record('pre_call', None) - self.pre = additional_args - self.pre_api_key = api_key - on_pre_call(additional_args) - - def post_call(self, original_response, api_key, additional_args): - self.record('post_call', None) - self.post = (original_response, api_key, additional_args) - -request = Request() -kwargs = {} -logger = PayloadLogger() -on_pre_call = lambda additional_args: None -check = lambda: None -"; - -const DOCUMENT: &str = "data:application/pdf;base64,YWJj"; -const EDITED: &str = "data:application/pdf;base64,ZWRpdGVk"; - -fn document(source: &str) -> Value { - json!({"type": "document_url", "document_url": source}) -} - -fn before_send(script: &CStr, body: Value) -> WireRequest { - before_send_with_secrets(script, json!({}), body, &[]) -} - -/// Runs `before_send` over `body` for a route whose parameters are `optional_params`, with -/// the Python objects `script` binds, then delivers the provider's raw response the way the -/// driver does and runs the script's `check()`. -fn before_send_with_secrets( - script: &CStr, - optional_params: Value, - body: Value, - secret_fields: &[&str], -) -> WireRequest { - before_send_bound(&[], script, optional_params, body, secret_fields) -} - -/// [`before_send_with_secrets`] with `bindings` placed in the namespace before `script` runs. -fn before_send_bound( - bindings: &[(&str, &Value)], - script: &CStr, - optional_params: Value, - body: Value, - secret_fields: &[&str], -) -> WireRequest { - Python::initialize(); - Python::attach(|py| { - let locals = namespace(py, PAYLOAD_LOGGER); - for &(name, value) in bindings { - locals.set_item(name, to_py(py, value).unwrap()).unwrap(); - } - run(py, &locals, script); - let mut logging = LegacyLogging { - logger: Some(PythonLogger::new(local(&locals, "logger").unbind())), - ..legacy_call(py, &locals, false) - }; - let context = RequestContext { - model: "model".into(), - custom_llm_provider: "provider".into(), - optional_params, - secret_fields: secret_fields.iter().map(|name| name.to_string()).collect(), - api_key: Some(SecretValue::new("route-key")), - }; - let wire = WireRequest { - url: "https://provider.invalid/ocr".into(), - headers: vec![("x-route".into(), "route".into())], - body, - }; - let step = logging.before_send(py, Box::new(wire), &context).unwrap(); - let raw = MachineEvent::ResponseReceived { - raw: RawResponse { - body: "raw response".into(), - }, - }; - assert!(matches!( - logging.emit(py, LifecycleEvent::Machine(&raw)).unwrap(), - LifecycleStep::Done - )); - run(py, &locals, c"check()"); - let LifecycleStep::Wire(wire) = step else { - panic!("before_send did not hand back the wire request"); - }; - *wire - }) -} - -#[rstest] -#[case::caller_keyword(c" -document = {'type': 'document_url', 'document_url': 'data:application/pdf;base64,YWJj'} -pages = [0] -kwargs = {'document': document, 'pages': pages} -observed = [] -on_pre_call = lambda args: observed.append( - (args['complete_input_dict']['document'] is document, args['complete_input_dict']['pages'] is pages) -) -def check(): - assert observed == [(True, True)], observed -")] -#[case::request_attribute_behind_an_omitted_keyword(c" -document = {'type': 'document_url', 'document_url': 'data:application/pdf;base64,YWJj'} -pages = [0] -request.document = document -kwargs = {'pages': pages} -observed = [] -on_pre_call = lambda args: observed.append( - (args['complete_input_dict']['document'] is document, args['complete_input_dict']['pages'] is pages) -) -def check(): - assert observed == [(True, True)], observed -")] -fn passthrough_keys_reach_pre_call_as_the_callers_own_objects(#[case] script: &CStr) { - let body = json!({"model": "model", "document": document(DOCUMENT), "pages": [0]}); - let wire = before_send(script, body.clone()); - assert_eq!(wire.body, body); -} - -#[test] -fn pre_call_edit_of_a_passthrough_object_reaches_the_caller_and_the_wire() { - let wire = before_send( - c" -document = {'type': 'document_url', 'document_url': 'data:application/pdf;base64,YWJj'} -kwargs = {'document': document} -def on_pre_call(args): - args['complete_input_dict']['document']['document_url'] = 'data:application/pdf;base64,ZWRpdGVk' -def check(): - assert document['document_url'] == 'data:application/pdf;base64,ZWRpdGVk' -", - json!({"document": document(DOCUMENT)}), - ); - assert_eq!(wire.body["document"], document(EDITED)); -} - -#[test] -fn a_body_key_the_route_rewrote_is_not_the_callers_object() { - let wire = before_send( - c" -document = {'type': 'document_url', 'document_url': 'https://example.invalid/scan.pdf'} -kwargs = {'document': document} -observed = [] -def on_pre_call(args): - observed.append(args['complete_input_dict']['document'] is document) - args['complete_input_dict']['document']['document_name'] = 'edited.pdf' -def check(): - assert observed == [False], observed - assert document == {'type': 'document_url', 'document_url': 'https://example.invalid/scan.pdf'} -", - json!({"document": document(DOCUMENT)}), - ); - assert_eq!( - wire.body["document"], - json!({"type": "document_url", "document_url": DOCUMENT, "document_name": "edited.pdf"}) - ); -} - -#[test] -fn a_caller_value_with_no_json_form_is_left_out_of_realiasing() { - let body = json!({"pages": [0]}); - let wire = before_send( - c" -opaque = object() -kwargs = {'pages': opaque} -observed = [] -on_pre_call = lambda args: observed.append(args['complete_input_dict']['pages']) -def check(): - assert observed == [[0]], observed -", - body.clone(), - ); - assert_eq!(wire.body, body); -} - -#[rstest] -#[case::body( - c" -def on_pre_call(args): - args['complete_input_dict'] = {'replacement': True} -" -)] -#[case::headers( - c" -def on_pre_call(args): - args['headers'] = {'x-replacement': 'yes'} -" -)] -fn rebinding_the_payload_envelope_does_not_reach_the_wire(#[case] script: &CStr) { - let body = json!({"document": document(DOCUMENT)}); - let wire = before_send(script, body.clone()); - assert_eq!(wire.body, body); - assert_eq!(wire.headers, [("x-route".to_string(), "route".to_string())]); -} - -#[test] -fn pre_call_header_edit_reaches_the_wire() { - let wire = before_send( - c" -def on_pre_call(args): - args['headers']['x-callback'] = 'edited' -", - json!({}), - ); - assert_eq!( - wire.headers, - [ - ("x-route".to_string(), "route".to_string()), - ("x-callback".to_string(), "edited".to_string()), - ] - ); -} - -#[test] -fn pre_call_receives_the_wire_request_and_the_logger_its_redacted_request() { - let body = json!({"model": "model", "document": document(DOCUMENT)}); - before_send_with_secrets( - c" -logger_fn = lambda *args: None -kwargs = { - 'litellm_call_id': 'call-1', - 'client_secret': 'shh', - 'proxy_server_request': {'body': {}}, - 'logger_fn': logger_fn, - 'litellm_request_debug': True, - 'ocr_cost_per_page': 0.05, -} -observed = [] -on_pre_call = observed.append -def check(): - [args] = observed - assert args['api_base'] == 'https://provider.invalid/ocr', args - assert args['complete_input_dict'] == { - 'model': 'model', - 'document': {'type': 'document_url', 'document_url': 'data:application/pdf;base64,YWJj'}, - }, args - update = logger.update - assert update['model'] == 'model' and update['custom_llm_provider'] == 'provider', update - assert update['litellm_params']['litellm_call_id'] == 'call-1', update - assert update['litellm_params']['api_base'] == 'https://provider.invalid/ocr', update - assert update['litellm_params']['logger_fn'] is logger_fn, update - assert update['litellm_params']['litellm_request_debug'] is True, update - assert update['litellm_params']['ocr_cost_per_page'] == 0.05, update - assert update['kwargs']['client_secret'] == '****', update - assert 'proxy_server_request' not in update['kwargs'], update - assert update['optional_params']['client_secret'] == '****', update -", - json!({"client_secret": "shh"}), - body, - &["client_secret"], - ); -} - -#[rstest] -#[case::added_key( - c" -def on_pre_call(args): - args['complete_input_dict']['include_image_base64'] = True -", - json!({"document": document(DOCUMENT), "include_image_base64": true}) -)] -#[case::replaced_document( - c" -document = {'type': 'document_url', 'document_url': 'data:application/pdf;base64,YWJj'} -kwargs = {'document': document} -def on_pre_call(args): - args['complete_input_dict']['document'] = { - 'type': 'document_url', 'document_url': 'data:application/pdf;base64,ZWRpdGVk' - } -def check(): - assert document['document_url'] == 'data:application/pdf;base64,YWJj', document -", - json!({"document": document(EDITED)}) -)] -#[case::retained_body_edited_after_rebinding( - c" -def on_pre_call(args): - retained = args['complete_input_dict'] - args['complete_input_dict'] = {'rebound': True} - retained['include_image_base64'] = True -", - json!({"document": document(DOCUMENT), "include_image_base64": true}) -)] -fn pre_call_body_edits_reach_the_wire(#[case] script: &CStr, #[case] expected: Value) { - let body = json!({"document": document(DOCUMENT)}); - let wire = before_send(script, body); - assert_eq!(wire.body, expected); -} - -#[test] -fn retained_headers_edited_after_rebinding_reach_the_wire() { - let wire = before_send( - c" -def on_pre_call(args): - retained = args['headers'] - args['headers'] = {'x-rebound': 'rebound'} - retained['x-retained'] = 'sent' -", - json!({}), - ); - assert_eq!( - wire.headers, - [ - ("x-route".to_string(), "route".to_string()), - ("x-retained".to_string(), "sent".to_string()), - ] - ); -} - -#[test] -fn post_call_receives_the_raw_response_the_route_key_and_the_body_and_headers_pre_call_saw() { - before_send( - c" -def check(): - original_response, api_key, additional_args = logger.post - assert original_response == 'raw response', original_response - assert api_key == logger.pre_api_key == 'route-key', (api_key, logger.pre_api_key) - assert additional_args == { - 'complete_input_dict': logger.pre['complete_input_dict'], - 'headers': logger.pre['headers'], - }, additional_args - assert additional_args['complete_input_dict'] is logger.pre['complete_input_dict'] - assert additional_args['headers'] is logger.pre['headers'] -", - json!({"document": document(DOCUMENT)}), - ); -} - -#[test] -fn every_request_runs_the_full_pre_call_and_post_call() { - let wire = before_send( - c" -def on_pre_call(args): - args['complete_input_dict']['include_image_base64'] = True -def check(): - assert logger.names() == ['pre_call', 'post_call'], logger.calls -", - json!({"document": document(DOCUMENT)}), - ); - assert_eq!( - wire.body, - json!({"document": document(DOCUMENT), "include_image_base64": true}) - ); -} - -/// What one pre-call callback does to the payload it is handed. -#[derive(Clone, Debug)] -enum Edit { - Nothing, - Set(String, Value), - Remove(String), - Rebind(Value), - RebindThenSetRetained(String, Value), -} - -impl Edit { - fn script(&self) -> Value { - match self { - Self::Nothing => json!({"kind": "nothing"}), - Self::Set(key, value) => json!({"kind": "set", "key": key, "value": value}), - Self::Remove(key) => json!({"kind": "remove", "key": key}), - Self::Rebind(value) => json!({"kind": "rebind", "value": value}), - Self::RebindThenSetRetained(key, value) => { - json!({"kind": "rebind_then_set_retained", "key": key, "value": value}) - } - } - } - - /// The legacy contract: the provider is sent the body object `pre_call` received, as - /// the callback left it. Rebinding the envelope's key points the envelope elsewhere and - /// leaves that object alone. - fn sent(&self, body: &Map) -> Value { - let mut sent = body.clone(); - match self { - Self::Nothing | Self::Rebind(_) => {} - Self::Set(key, value) | Self::RebindThenSetRetained(key, value) => { - sent.insert(key.clone(), value.clone()); - } - Self::Remove(key) => { - sent.remove(key); - } - } - Value::Object(sent) - } -} - -/// How the caller's keyword for a body key relates to what the route sends under it. -#[derive(Clone, Copy, Debug, PartialEq, Eq)] -enum Caller { - PassedUnchanged, - RewrittenByTheRoute, - NotPassed, -} - -const MODEL: &CStr = c" -aliased = {} -def on_pre_call(args): - body = args['complete_input_dict'] - aliased.update({name: body[name] is kwargs[name] for name in unchanged}) - kind = edit['kind'] - if kind == 'set': - body[edit['key']] = edit['value'] - elif kind == 'remove': - body.pop(edit['key'], None) - elif kind == 'rebind': - args['complete_input_dict'] = edit['value'] - elif kind == 'rebind_then_set_retained': - args['complete_input_dict'] = {} - body[edit['key']] = edit['value'] -def check(): - assert aliased == {name: True for name in unchanged}, aliased - assert logger.names() == ['pre_call', 'post_call'], logger.calls -"; - -fn json_value() -> impl Strategy { - let leaf = prop_oneof![ - Just(Value::Null), - any::().prop_map(Value::from), - any::().prop_map(Value::from), - any::() - .prop_filter("JSON has no NaN or infinity", |number| number.is_finite()) - .prop_map(Value::from), - ".{0,8}".prop_map(Value::from), - ]; - leaf.prop_recursive(3, 24, 4, |inner| { - prop_oneof![ - prop::collection::vec(inner.clone(), 0..4).prop_map(Value::from), - prop::collection::btree_map(key(), inner, 0..4) - .prop_map(|fields| Value::Object(fields.into_iter().collect())), - ] - }) -} - -fn key() -> impl Strategy { - "[a-z]{1,6}" -} - -fn caller() -> impl Strategy { - prop_oneof![ - Just(Caller::PassedUnchanged), - Just(Caller::RewrittenByTheRoute), - Just(Caller::NotPassed), - ] -} - -fn edit() -> impl Strategy { - prop_oneof![ - Just(Edit::Nothing), - (key(), json_value()).prop_map(|(key, value)| Edit::Set(key, value)), - key().prop_map(Edit::Remove), - json_value().prop_map(Edit::Rebind), - (key(), json_value()).prop_map(|(key, value)| Edit::RebindThenSetRetained(key, value)), - ] -} - -proptest! { - #![proptest_config(ProptestConfig::with_cases(128))] - - /// For any body, any caller keywords and any callback edit: every keyword the route - /// sends unchanged reaches `pre_call` as the caller's own object, and the provider is - /// sent exactly what the model says, so a callback that edits nothing changes nothing. - #[test] - fn the_wire_is_the_body_pre_call_received_as_the_callback_left_it( - fields in prop::collection::btree_map(key(), (json_value(), caller()), 0..5), - edit in edit(), - ) { - let body: Map = fields - .iter() - .map(|(name, (value, _))| (name.clone(), value.clone())) - .collect(); - let kwargs: Map = fields - .iter() - .filter_map(|(name, (value, caller))| match caller { - Caller::PassedUnchanged => Some((name.clone(), value.clone())), - Caller::RewrittenByTheRoute => Some((name.clone(), json!([value]))), - Caller::NotPassed => None, - }) - .collect(); - let unchanged: Value = fields - .iter() - .filter(|(_, (_, caller))| *caller == Caller::PassedUnchanged) - .map(|(name, _)| Value::from(name.clone())) - .collect(); - - let wire = before_send_bound( - &[ - ("kwargs", &Value::Object(kwargs)), - ("unchanged", &unchanged), - ("edit", &edit.script()), - ], - MODEL, - json!({}), - Value::Object(body.clone()), - &[], - ); - - prop_assert_eq!(wire.body, edit.sent(&body)); - prop_assert_eq!(wire.headers, [("x-route".to_string(), "route".to_string())]); - } -} diff --git a/litellm-rust/crates/callbacks-legacy-python/tests/support.rs b/litellm-rust/crates/callbacks-legacy-python/tests/support.rs deleted file mode 100644 index d0c02fa6da5..00000000000 --- a/litellm-rust/crates/callbacks-legacy-python/tests/support.rs +++ /dev/null @@ -1,205 +0,0 @@ -use std::ffi::CStr; - -use pyo3::prelude::*; -use pyo3::types::{PyDict, PyTuple}; - -use crate::{LegacyLogging, LegacySurface, PublicCall}; - -/// The parameters of every `callbacks_legacy_python` function, as the real module declares them. -/// `tests/test_litellm/rust_bridge/test_callbacks_legacy_python.py` pins this file to the Python -/// signatures, and [`namespace`] binds every fake call against it. -pub(crate) const PYTHON_CONTRACT: &str = include_str!("../python_contract.json"); - -/// Stand-ins for `callbacks_legacy_python`, the only Python module the crate calls. Tests -/// share one interpreter and run concurrently, so each fake is installed idempotently and -/// forwards to the per-test `StubLogger` it is handed (directly, or as `kwargs['logger']`). -/// Every fake is bound against the contract first, so a call the real module would reject -/// fails here too. -const STUBS: &CStr = c" -import contextvars -import inspect -import json -import sys -import traceback -import types - -for name in ('litellm', 'litellm.rust_bridge', 'litellm.rust_bridge.callbacks_legacy_python'): - sys.modules.setdefault(name, types.ModuleType(name)) - -legacy = sys.modules['litellm.rust_bridge.callbacks_legacy_python'] -CONTRACT = json.loads(python_contract) - - -def contracted(name, fake): - signature = inspect.Signature( - [inspect.Parameter(parameter, inspect.Parameter.POSITIONAL_OR_KEYWORD) for parameter in CONTRACT[name]] - ) - - def checked(*args, **kwargs): - signature.bind(*args, **kwargs) - return fake(*args, **kwargs) - - return checked - - -if not hasattr(legacy, 'is_internal'): - legacy.is_internal = contextvars.ContextVar('is_internal_call', default=False) - -FAKES = { - 'setup': lambda call_type, args, kwargs, start, asynchronous: types.SimpleNamespace( - logger=kwargs['logger_factory'](kwargs) if 'logger_factory' in kwargs else kwargs['logger'], - kwargs=kwargs, - ), - 'check_limits': lambda arguments: arguments['logger'].check_limits(arguments), - 'finalize': lambda response, logger, kwargs, start, end: logger.record('finalize', response), - 'update_logging': lambda logger, kwargs, model, optional_params, litellm_params, provider: logger.update_from_kwargs( - kwargs=kwargs, - model=model, - optional_params=optional_params, - litellm_params=litellm_params, - custom_llm_provider=provider, - ), - 'pre_call': lambda logger, input, api_key, additional_args: logger.pre_call(input, api_key, additional_args), - 'post_call': lambda logger, original_response, api_key, additional_args: logger.post_call( - original_response, api_key, additional_args - ), - 'defers_async_logging': lambda logger: bool(getattr(logger, '_defer_async_logging', False)), - 'defer_success': lambda logger, pending: setattr(logger, '_native_pending_logging', pending), - 'sync_success_for_async_call': lambda logger, response, start, end: logger.handle_sync_success_callbacks_for_async_calls( - response, start, end - ), - 'failure_handler': lambda logger, error, start, end, asynchronous: ( - logger.async_failure_handler if asynchronous else logger.failure_handler - )(error, ''.join(traceback.format_exception(error)), start, end), - 'submit_success': lambda logger, response, start, end: logger.record('submit', (response, start, end)), - 'async_success_handler': lambda logger, response, start, end: logger.async_success_handler(response, start, end), - 'enqueue_logging': lambda coroutine: coroutine.enqueue(), - 'restore_context': lambda logger: logger.record('restore', None), - 'custom_pricing_fields': lambda: ('ocr_cost_per_page',), - 'is_internal_call': lambda: legacy.is_internal.get(), - 'credential_list': lambda: [], - 'warn_unknown_credential': lambda name, loaded: None, - 'before_deployment_call': lambda kwargs, call_type: kwargs['logger'].hook('pre', kwargs, call_type), - 'after_deployment_success': lambda kwargs, response, call_type: kwargs['logger'].hook( - 'success', response, call_type - ), - 'after_deployment_failure': lambda kwargs, error, call_type: kwargs['logger'].hook('failure', error, call_type), - 'stream_opened': lambda logger: logger.record('stream_opened', None), - 'stream_success': lambda logger, request_body, chunks, start, end, first_chunk: logger.record( - 'stream_success', list(chunks) - ), - 'stream_failure': lambda logger, request_body, chunks, error: logger.record('stream_failure', error), -} -assert FAKES.keys() == CONTRACT.keys(), sorted(FAKES.keys() ^ CONTRACT.keys()) -for name, fake in FAKES.items(): - setattr(legacy, name, contracted(name, fake)) - - -unraisable = sys.modules.setdefault( - 'litellm_test_unraisable', types.ModuleType('litellm_test_unraisable') -) -if not hasattr(unraisable, 'events'): - unraisable.events = [] - sys.unraisablehook = lambda event: unraisable.events.append((event.object, event.exc_value)) - - -def unraisable_from(owner): - return [error for source, error in unraisable.events if source is owner] - - -class StubCoroutine: - def __init__(self, logger): - self.logger = logger - - def enqueue(self): - self.logger.record('enqueued', None) - self.logger.on_enqueue(self) - - def close(self): - self.logger.record('closed', None) - - -class StubLogger: - def __init__(self): - self.calls = [] - self.hooks = {} - self.on_enqueue = lambda coroutine: None - - def record(self, name, value): - self.calls.append((name, value)) - - def names(self): - return [name for name, _ in self.calls] - - def hook(self, phase, value, call_type): - self.record(phase + '_hook', call_type) - return self.hooks.get(phase, lambda value: 'awaitable')(value) - - def check_limits(self, arguments): - self.record('check_limits', arguments) - - def failure_handler(self, error, trace, start, end): - self.record('failure_handler', error) - - def async_failure_handler(self, error, trace, start, end): - self.record('async_failure_handler', error) - return 'awaitable' - - def success_handler(self, response, start, end): - self.record('success_handler', response) - - def async_success_handler(self, response, start, end): - self.record('async_success_handler', response) - return StubCoroutine(self) - - def handle_sync_success_callbacks_for_async_calls(self, response, start, end): - self.record('sync_success_for_async_call', response) - - -logger = StubLogger() -"; - -/// A namespace with the stubs, `StubLogger` and a fresh `logger`, after `script` ran in it. -pub(crate) fn namespace<'py>(py: Python<'py>, script: &CStr) -> Bound<'py, PyDict> { - let locals = PyDict::new(py); - locals.set_item("python_contract", PYTHON_CONTRACT).unwrap(); - py.run(STUBS, Some(&locals), Some(&locals)).unwrap(); - py.run(script, Some(&locals), Some(&locals)).unwrap(); - locals -} - -pub(crate) fn run(py: Python<'_>, locals: &Bound<'_, PyDict>, code: &CStr) { - py.run(code, Some(locals), Some(locals)).unwrap(); -} - -pub(crate) fn local<'py>(locals: &Bound<'py, PyDict>, name: &str) -> Bound<'py, PyAny> { - locals.get_item(name).unwrap().unwrap() -} - -/// A legacy call over the namespace's `kwargs` (or none) and `request` (or `None`). -pub(crate) fn legacy_call( - py: Python<'_>, - locals: &Bound<'_, PyDict>, - asynchronous: bool, -) -> LegacyLogging { - let request = locals - .get_item("request") - .unwrap() - .unwrap_or_else(|| py.None().into_bound(py)); - let kwargs = locals - .get_item("kwargs") - .unwrap() - .map(|kwargs| kwargs.cast_into::().unwrap()) - .unwrap_or_else(|| PyDict::new(py)); - let call = PublicCall::capture(&request, &PyTuple::empty(py), &kwargs).unwrap(); - LegacyLogging::new( - py, - LegacySurface { - call_type: "test", - input_description: "test input", - stream: None, - }, - call, - asynchronous, - ) -} diff --git a/litellm-rust/crates/callbacks-legacy-python/tests/terminal.rs b/litellm-rust/crates/callbacks-legacy-python/tests/terminal.rs deleted file mode 100644 index f68209233f2..00000000000 --- a/litellm-rust/crates/callbacks-legacy-python/tests/terminal.rs +++ /dev/null @@ -1,291 +0,0 @@ -use std::ffi::CStr; - -use litellm_host::event::{FailureOrigin, Timing}; -use litellm_host_python::{LifecycleEvent, LifecycleStep, PythonLifecycle}; -use pyo3::exceptions::PyRuntimeError; -use pyo3::exceptions::asyncio::CancelledError; -use pyo3::prelude::*; -use pyo3::types::PyDict; -use rstest::rstest; - -use super::LegacyLogging; -use crate::PythonLogger; -use crate::test_support::{legacy_call, local, namespace, run}; - -const TIMING: Timing = Timing { - start_time: 0.0, - end_time: 1.0, -}; - -fn logged(py: Python<'_>, locals: &Bound<'_, PyDict>, asynchronous: bool) -> LegacyLogging { - LegacyLogging { - logger: Some(PythonLogger::new(local(locals, "logger").unbind())), - ..legacy_call(py, locals, asynchronous) - } -} - -fn succeed( - py: Python<'_>, - locals: &Bound<'_, PyDict>, - logging: &mut LegacyLogging, -) -> LifecycleStep { - let response = local(locals, "response").unbind(); - logging - .emit( - py, - LifecycleEvent::Succeeded { - timing: TIMING, - response: &response, - }, - ) - .unwrap() -} - -fn fail(py: Python<'_>, locals: &Bound<'_, PyDict>, logging: &mut LegacyLogging) -> LifecycleStep { - let failure = PyErr::from_value(local(locals, "failure")); - logging - .emit( - py, - LifecycleEvent::Failed { - timing: TIMING, - origin: FailureOrigin::Host, - error: &failure, - }, - ) - .unwrap() -} - -#[rstest] -#[case::sync_listened(false, c"", &["submit"])] -#[case::async_listened( - true, - c"", - &["async_success_handler", "enqueued", "sync_success_for_async_call"] -)] -#[case::async_deferred(true, c"logger._defer_async_logging = True", &["sync_success_for_async_call"])] -#[case::async_with_fallbacks(true, c"kwargs = {'fallbacks': ['other']}", &["sync_success_for_async_call"])] -fn success_reaches_the_logging_handlers( - #[case] asynchronous: bool, - #[case] script: &CStr, - #[case] expected: &[&str], -) { - Python::initialize(); - Python::attach(|py| { - let locals = namespace(py, c"response = object()"); - run(py, &locals, script); - let mut logging = logged(py, &locals, asynchronous); - assert!(matches!( - succeed(py, &locals, &mut logging), - LifecycleStep::Done - )); - let names: Vec = local(&locals, "logger") - .call_method0("names") - .unwrap() - .extract() - .unwrap(); - assert_eq!(names, expected); - run( - py, - &locals, - c" -assert all(value is response for name, value in logger.calls if name.endswith('_handler')) -assert hasattr(logger, '_native_pending_logging') == getattr(logger, '_defer_async_logging', False) -", - ); - }); -} - -#[rstest] -#[case::synchronous(false, &["failure_handler"])] -#[case::asynchronous(true, &[])] -fn internal_calls_skip_failure_callbacks_only_when_asynchronous( - #[case] asynchronous: bool, - #[case] expected: &[&str], -) { - Python::initialize(); - Python::attach(|py| { - let locals = namespace(py, c"failure = ValueError('provider')"); - let mut logging = LegacyLogging { - internal: true, - ..logged(py, &locals, asynchronous) - }; - assert!(matches!( - fail(py, &locals, &mut logging), - LifecycleStep::Done - )); - let names: Vec = local(&locals, "logger") - .call_method0("names") - .unwrap() - .extract() - .unwrap(); - assert_eq!(names, expected); - }); -} - -#[test] -fn internal_async_calls_skip_the_async_success_fan_out() { - Python::initialize(); - Python::attach(|py| { - let locals = namespace(py, c"response = object()"); - let mut logging = LegacyLogging { - internal: true, - ..logged(py, &locals, true) - }; - succeed(py, &locals, &mut logging); - run( - py, - &locals, - c"assert logger.names() == ['sync_success_for_async_call'], logger.calls", - ); - }); -} - -#[test] -fn a_failing_success_callback_is_reported_without_replacing_the_response() { - Python::initialize(); - Python::attach(|py| { - let locals = namespace( - py, - c" -response = object() -failure = ValueError('terminal diagnostic') - -class FailingLogger(StubLogger): - def handle_sync_success_callbacks_for_async_calls(self, *args): - raise failure - -logger = FailingLogger() -", - ); - let mut logging = logged(py, &locals, true); - assert!(matches!( - succeed(py, &locals, &mut logging), - LifecycleStep::Done - )); - assert!( - logging - .response - .as_ref() - .unwrap() - .bind(py) - .is(local(&locals, "response")) - ); - run(py, &locals, c"assert unraisable_from(logger) == [failure]"); - }); -} - -#[rstest] -#[case::sync_listened(false, c"", &["failure_handler"])] -#[case::async_listened(true, c"", &["failure_handler", "async_failure_handler"])] -fn failure_reaches_the_logging_handlers( - #[case] asynchronous: bool, - #[case] script: &CStr, - #[case] expected: &[&str], -) { - Python::initialize(); - Python::attach(|py| { - let locals = namespace(py, c"failure = ValueError('provider')"); - run(py, &locals, script); - let mut logging = logged(py, &locals, asynchronous); - let step = fail(py, &locals, &mut logging); - let awaits_async_handler = expected.contains(&"async_failure_handler"); - assert_eq!( - matches!(step, LifecycleStep::Await(_)), - awaits_async_handler - ); - let names: Vec = local(&locals, "logger") - .call_method0("names") - .unwrap() - .extract() - .unwrap(); - assert_eq!(names, expected); - run( - py, - &locals, - c"assert all(value is failure for name, value in logger.calls if name.endswith('_handler'))", - ); - }); -} - -#[test] -fn a_failing_sync_failure_callback_keeps_the_error_and_still_runs_the_async_family() { - Python::initialize(); - Python::attach(|py| { - let locals = namespace( - py, - c" -failure = ValueError('selected') - -class FailingLogger(StubLogger): - def failure_handler(self, error, trace, start, end): - self.record('failure_handler', error) - raise RuntimeError('handler failed') - -logger = FailingLogger() -", - ); - let mut logging = logged(py, &locals, true); - assert!(matches!( - fail(py, &locals, &mut logging), - LifecycleStep::Await(_) - )); - assert!( - logging - .error - .as_ref() - .unwrap() - .bind(py) - .is(local(&locals, "failure")) - ); - run( - py, - &locals, - c"assert logger.names() == ['failure_handler', 'async_failure_handler'], logger.calls", - ); - }); -} - -#[rstest] -#[case::completed(None, true)] -#[case::handler_error(Some(false), true)] -#[case::cancelled(Some(true), false)] -fn the_async_failure_handler_ends_the_call_unless_it_was_cancelled( - #[case] error: Option, - #[case] done: bool, -) { - Python::initialize(); - Python::attach(|py| { - let locals = namespace(py, c"failure = ValueError('provider')"); - let mut logging = logged(py, &locals, true); - fail(py, &locals, &mut logging); - let result = match error { - None => Ok(py.None()), - Some(false) => Err(PyRuntimeError::new_err("handler failed")), - Some(true) => Err(CancelledError::new_err("cancelled")), - }; - let expected = result.as_ref().err().map(|error| error.value(py).clone()); - match logging.resume(py, result) { - Ok(step) => assert!(done && matches!(step, LifecycleStep::Done)), - Err(propagated) => { - assert!(!done); - assert!(propagated.value(py).is(expected.unwrap())); - } - } - }); -} - -#[test] -fn closing_restores_the_correlation_context_once() { - Python::initialize(); - Python::attach(|py| { - let locals = namespace(py, c""); - let mut logging = logged(py, &locals, true); - logging.close(py); - logging.close(py); - run( - py, - &locals, - c"assert logger.names() == ['restore'], logger.calls", - ); - }); -} diff --git a/litellm-rust/crates/core/Cargo.toml b/litellm-rust/crates/core/Cargo.toml index 4626781f3e3..d7096cdd774 100644 --- a/litellm-rust/crates/core/Cargo.toml +++ b/litellm-rust/crates/core/Cargo.toml @@ -4,7 +4,6 @@ version = "0.1.0" edition.workspace = true license.workspace = true repository.workspace = true -autotests = false [dependencies] litellm-secrets.workspace = true diff --git a/litellm-rust/crates/core/src/audio_transcription/mod.rs b/litellm-rust/crates/core/src/audio_transcription/mod.rs index af9c398c065..801fd5e9673 100644 --- a/litellm-rust/crates/core/src/audio_transcription/mod.rs +++ b/litellm-rust/crates/core/src/audio_transcription/mod.rs @@ -14,6 +14,3 @@ pub async fn audio_transcription(request: AudioTranscriptionRequest<'_>) -> Resu execute_audio_transcription_provider_call(prepare_audio_transcription_provider_call(request)?) .await } - -#[cfg(test)] -mod tests; diff --git a/litellm-rust/crates/core/src/chat_completions/mod.rs b/litellm-rust/crates/core/src/chat_completions/mod.rs index 81d35044d08..224c9d8cfed 100644 --- a/litellm-rust/crates/core/src/chat_completions/mod.rs +++ b/litellm-rust/crates/core/src/chat_completions/mod.rs @@ -51,6 +51,3 @@ pub fn chat_completions_decline_reason( .unsupported_reason(&messages, optional_params) .map(|reason| reason.0) } - -#[cfg(test)] -mod tests; diff --git a/litellm-rust/crates/core/src/chat_completions/prepare.rs b/litellm-rust/crates/core/src/chat_completions/prepare.rs index c8e6365121e..afea46221f5 100644 --- a/litellm-rust/crates/core/src/chat_completions/prepare.rs +++ b/litellm-rust/crates/core/src/chat_completions/prepare.rs @@ -143,3 +143,841 @@ pub(super) fn prepare_provider_request( timeout: request.timeout, }) } + +#[cfg(test)] +mod tests { + use litellm_llms::base_llm::chat::transformation::RequestAuth; + use serde_json::{Map, Value, json}; + + use super::{prepare_provider_request, resolve_request}; + use crate::chat_completions::{ + Error, + types::{ChatCompletionsRequest, ProviderChatCompletionsRequest}, + }; + + fn prepare_chat_completions_call( + request: ChatCompletionsRequest<'_>, + ) -> Result { + prepare_provider_request(resolve_request(request)?) + } + + fn request<'a>( + model: &'a str, + provider: Option<&'a str>, + messages: Value, + optional_params: Value, + ) -> ChatCompletionsRequest<'a> { + ChatCompletionsRequest { + model, + messages, + optional_params: match optional_params { + Value::Object(map) => map, + other => panic!("params must be an object, got {other}"), + }, + api_key: Some("sk-test"), + api_base: None, + custom_llm_provider: provider, + extra_headers: None, + timeout: None, + } + } + + /// `ProviderChatCompletionsRequest` deliberately has no `Debug` (its headers + /// carry resolved credentials), so unwrap the failure case by hand. + fn decline(request: ChatCompletionsRequest<'_>) -> Error { + match prepare_chat_completions_call(request) { + Err(error) => error, + Ok(prepared) => panic!("expected a decline, prepared a call to {}", prepared.url), + } + } + + #[test] + fn resolves_the_provider_from_the_model_prefix() { + let prepared = prepare_chat_completions_call(request( + "anthropic/claude-sonnet-4-5", + None, + json!([{"role": "user", "content": "hi"}]), + json!({"max_tokens": 16}), + )) + .expect("prepares"); + assert_eq!(prepared.model, "claude-sonnet-4-5"); + assert_eq!(prepared.url, "https://api.anthropic.com/v1/messages"); + assert_eq!(prepared.body["model"], json!("claude-sonnet-4-5")); + } + + #[test] + fn strips_an_explicit_provider_prefix_from_the_model() { + let prepared = prepare_chat_completions_call(request( + "anthropic/claude-sonnet-4-5", + Some("anthropic"), + json!([{"role": "user", "content": "hi"}]), + json!({}), + )) + .expect("prepares"); + assert_eq!(prepared.model, "claude-sonnet-4-5"); + } + + #[test] + fn adds_the_auth_and_default_headers() { + let prepared = prepare_chat_completions_call(request( + "claude-sonnet-4-5", + Some("anthropic"), + json!([{"role": "user", "content": "hi"}]), + json!({}), + )) + .expect("prepares"); + assert!( + prepared + .upstream_headers + .contains(&("x-api-key".to_string(), "sk-test".to_string())) + ); + assert!( + prepared + .upstream_headers + .contains(&("anthropic-version".to_string(), "2023-06-01".to_string())) + ); + assert!(matches!( + prepared.auth, + RequestAuth::Header { + name: "x-api-key", + .. + } + )); + } + + #[test] + fn the_deployment_credential_replaces_a_caller_supplied_auth_header() { + // Python builds `{**headers, **anthropic_headers}`, so the deployment's key + // overwrites a forwarded one. Honouring the caller's would let whoever sends + // the request choose the Anthropic principal it bills to. + let mut call = request( + "claude-sonnet-4-5", + Some("anthropic"), + json!([{"role": "user", "content": "hi"}]), + json!({}), + ); + call.extra_headers = Some(Map::from_iter([( + "X-Api-Key".to_string(), + json!("sk-caller"), + )])); + let prepared = prepare_chat_completions_call(call).expect("prepares"); + let keys: Vec<_> = prepared + .upstream_headers + .iter() + .filter(|(name, _)| name.eq_ignore_ascii_case("x-api-key")) + .collect(); + assert_eq!(keys.len(), 1, "got {:?}", prepared.upstream_headers); + assert_eq!(keys[0].1, "sk-test"); + } + + #[test] + fn a_forwarded_authorization_header_suppresses_the_resolved_api_key_header() { + // Anthropic's `validate_environment` pops `x-api-key` and sets `authorization` + // for an OAuth token, so re-adding the key here would put the credential into + // a header the host removed on purpose. + let mut call = request( + "claude-sonnet-4-5", + Some("anthropic"), + json!([{"role": "user", "content": "hi"}]), + json!({}), + ); + call.extra_headers = Some(Map::from_iter([ + ( + "Authorization".to_string(), + json!("Bearer sk-ant-oat01-token"), + ), + ("X-Api-Key".to_string(), json!("sk-caller")), + ])); + let prepared = prepare_chat_completions_call(call).expect("prepares"); + assert!( + !prepared + .upstream_headers + .iter() + .any(|(name, value)| name.eq_ignore_ascii_case("x-api-key") && value == "sk-test"), + "the resolved key must not be applied over an OAuth bearer, got {:?}", + prepared.upstream_headers + ); + assert!( + prepared + .upstream_headers + .iter() + .any(|(name, value)| name.eq_ignore_ascii_case("authorization") + && value == "Bearer sk-ant-oat01-token") + ); + } + + #[test] + fn an_unrelated_forwarded_authorization_does_not_defer_the_resolved_key() { + // Only an OAuth bearer replaces the credential. Python sends the deployment's + // `x-api-key` alongside any other forwarded `authorization`, so deferring on + // the mere presence of that header would drop the deployment's auth. + let mut call = request( + "claude-sonnet-4-5", + Some("anthropic"), + json!([{"role": "user", "content": "hi"}]), + json!({}), + ); + call.extra_headers = Some(Map::from_iter([ + ("Authorization".to_string(), json!("Bearer unrelated")), + ("X-Api-Key".to_string(), json!("sk-caller")), + ])); + let prepared = prepare_chat_completions_call(call).expect("prepares"); + let keys: Vec<_> = prepared + .upstream_headers + .iter() + .filter(|(name, _)| name.eq_ignore_ascii_case("x-api-key")) + .collect(); + assert_eq!(keys.len(), 1, "got {:?}", prepared.upstream_headers); + assert_eq!(keys[0].1, "sk-test"); + assert!( + prepared + .upstream_headers + .iter() + .any(|(name, value)| name.eq_ignore_ascii_case("authorization") + && value == "Bearer unrelated"), + "the unrelated authorization must survive, got {:?}", + prepared.upstream_headers + ); + } + + #[test] + fn declines_an_unsupported_request_before_resolving_credentials() { + let mut call = request( + "claude-sonnet-4-5", + Some("anthropic"), + json!([{"role": "user", "content": "hi"}]), + json!({"stream": true}), + ); + call.api_key = None; + // No api_key is set and no env is consulted: the gate must run first, so the + // error is the decline rather than a missing-credential error. + assert_eq!(decline(call), Error::Unsupported("streaming")); + } + + #[test] + fn rejects_an_unknown_provider() { + assert_eq!( + decline(request( + "openai/gpt-4o", + None, + json!([{"role": "user", "content": "hi"}]), + json!({}), + )), + Error::InvalidProvider("openai".to_string()) + ); + } + + #[test] + fn rejects_a_model_with_no_resolvable_provider() { + assert!(matches!( + decline(request( + "claude-sonnet-4-5", + None, + json!([{"role": "user", "content": "hi"}]), + json!({}), + )), + Error::InvalidProvider(_) + )); + } + + #[test] + fn rejects_an_empty_or_malformed_message_list() { + assert_eq!( + decline(request( + "anthropic/claude-sonnet-4-5", + None, + json!([]), + json!({}), + )), + Error::InvalidRequest("chat completions requires at least one message".to_string()) + ); + assert!(matches!( + decline(request( + "anthropic/claude-sonnet-4-5", + None, + json!("not a list"), + json!({}), + )), + Error::InvalidRequest(_) + )); + } + + #[test] + fn rejects_non_string_extra_headers() { + let mut call = request( + "anthropic/claude-sonnet-4-5", + None, + json!([{"role": "user", "content": "hi"}]), + json!({}), + ); + call.extra_headers = Some(Map::from_iter([("x-trace".to_string(), json!(7))])); + assert_eq!( + decline(call), + Error::Headers(litellm_http::request::HeaderError { + context: "chat completions", + name: "x-trace".to_string(), + actual: "number", + }) + ); + } + + #[test] + fn prepares_a_bedrock_call_without_resolving_credentials() { + let mut call = request( + "bedrock/us-east-1/anthropic.claude-v2", + None, + json!([{"role": "user", "content": "hi"}]), + json!({"maxTokens": 16}), + ); + call.api_key = None; + let prepared = prepare_chat_completions_call(call).expect("prepares"); + assert_eq!( + prepared.url, + "https://bedrock-runtime.us-east-1.amazonaws.com/model/anthropic.claude-v2/converse" + ); + assert_eq!( + prepared.auth, + RequestAuth::AwsSigV4 { + region: "us-east-1".to_string(), + service: "bedrock", + } + ); + // SigV4 signs the serialized body, so prepare must not have added an + // Authorization header; the handler does it. + assert!( + !prepared + .upstream_headers + .iter() + .any(|(name, _)| name.eq_ignore_ascii_case("authorization")) + ); + assert_eq!(prepared.body["inferenceConfig"], json!({"maxTokens": 16})); + } + + #[tokio::test] + async fn a_forwarded_client_header_does_not_enter_the_bedrock_signature() { + // Python signs only the AWS header set and reattaches the rest, so a header + // the caller forwarded rides along without joining the canonical request. + // Signing it makes Converse 403 on a deployment that works on Python. + let mut call = request( + "bedrock/us-east-1/anthropic.claude-v2", + None, + json!([{"role": "user", "content": "hi"}]), + json!({ + "maxTokens": 16, + "aws_access_key_id": "AKIDEXAMPLE", + "aws_secret_access_key": "wJalrXUtnFEMI/K7MDENG+bPxRfiCYEXAMPLEKEY" + }), + ); + // A key would resolve to a bearer token and never reach the signer. + call.api_key = None; + call.extra_headers = Some(Map::from_iter([( + "x-request-id".to_string(), + json!("abc-123"), + )])); + let prepared = prepare_chat_completions_call(call).expect("prepares"); + let signed = crate::chat_completions::handler::outbound_request(&prepared) + .await + .expect("signs"); + + let authorization = signed + .header("authorization") + .expect("carries an authorization header") + .to_string(); + assert!( + authorization.starts_with("AWS4-HMAC-SHA256"), + "expected a SigV4 signature, got {authorization}" + ); + assert!( + !authorization.contains("x-request-id"), + "forwarded header reached SignedHeaders: {authorization}" + ); + // It still goes on the wire, it is just not part of the signature. + assert!( + signed + .headers() + .iter() + .any(|(name, value)| name == "x-request-id" && value == "abc-123"), + "forwarded header was dropped instead of reattached" + ); + } + + #[tokio::test] + async fn a_forwarded_header_the_signer_computes_declines_to_python() { + // Reattaching the caller's copy next to the computed one puts the name on + // the wire twice and Bedrock rejects the pair, so a request carrying one + // has to go to Python instead of being signed here. + for forwarded in [ + "Authorization", + "x-amz-date", + "x-amz-security-token", + "Date", + ] { + let mut call = request( + "bedrock/us-east-1/anthropic.claude-v2", + None, + json!([{"role": "user", "content": "hi"}]), + json!({ + "maxTokens": 16, + "aws_access_key_id": "AKIDEXAMPLE", + "aws_secret_access_key": "wJalrXUtnFEMI/K7MDENG+bPxRfiCYEXAMPLEKEY" + }), + ); + call.api_key = None; + call.extra_headers = Some(Map::from_iter([(forwarded.to_string(), json!("forged"))])); + let prepared = prepare_chat_completions_call(call).expect("prepares"); + let error = crate::chat_completions::handler::outbound_request(&prepared) + .await + .expect_err("{forwarded} should decline instead of being signed"); + assert!( + matches!(error, Error::Unsupported(_)), + "{forwarded} declined as {error:?}, which the host would not fall back on" + ); + } + } + + #[test] + fn a_bedrock_deployment_bearer_outranks_a_forwarded_authorization() { + // `get_request_headers` assigns `headers["Authorization"]` unconditionally + // once a bearer token resolves, so the deployment's identity wins on + // Python. Keeping the caller's would authorize and bill the call as a + // different principal, and only when the deployment carries `rust: true`. + let mut call = request( + "bedrock/us-east-1/anthropic.claude-v2", + None, + json!([{"role": "user", "content": "hi"}]), + json!({"maxTokens": 16}), + ); + call.extra_headers = Some(Map::from_iter([( + "Authorization".to_string(), + json!("Bearer caller-supplied"), + )])); + let prepared = prepare_chat_completions_call(call).expect("prepares"); + let authorizations: Vec<_> = prepared + .upstream_headers + .iter() + .filter(|(name, _)| name.eq_ignore_ascii_case("authorization")) + .map(|(_, value)| value.as_str()) + .collect(); + assert_eq!( + authorizations, + vec!["Bearer sk-test"], + "the deployment token must be the only authorization on the wire" + ); + } + + #[test] + fn an_anthropic_forwarded_oauth_bearer_still_outranks_the_resolved_key() { + // The opposite precedence, and deliberate: Anthropic's own transform + // honours a forwarded OAuth bearer, so the Bedrock fix above must not be + // generalized into a rule that the configured key always wins. + // + // An OAuth bearer is the whole of that exception. This forwarded a plain + // `x-api-key` until round 17, which read as the same claim and was not: + // Python overwrites a forwarded `x-api-key` with the deployment's. + let mut call = request( + "claude-sonnet-4-5", + Some("anthropic"), + json!([{"role": "user", "content": "hi"}]), + json!({}), + ); + call.extra_headers = Some(Map::from_iter([( + "authorization".to_string(), + json!("Bearer sk-ant-oat01-forwarded"), + )])); + let prepared = prepare_chat_completions_call(call).expect("prepares"); + let keys: Vec<_> = prepared + .upstream_headers + .iter() + .filter(|(name, _)| name.eq_ignore_ascii_case("x-api-key")) + .map(|(_, value)| value.as_str()) + .collect(); + assert!(keys.is_empty(), "got {:?}", prepared.upstream_headers); + assert!( + prepared + .upstream_headers + .iter() + .any(|(name, value)| name.eq_ignore_ascii_case("authorization") + && value == "Bearer sk-ant-oat01-forwarded") + ); + } + + #[test] + fn a_bedrock_api_key_is_sent_as_a_bearer_token_instead_of_being_signed() { + // The configured bearer identity has its own account and quota boundary, + // so a request carrying one must not be signed as whatever principal the + // host's AWS credentials resolve to. + let prepared = prepare_chat_completions_call(request( + "bedrock/us-east-1/anthropic.claude-v2", + None, + json!([{"role": "user", "content": "hi"}]), + json!({"maxTokens": 16}), + )) + .expect("prepares"); + assert_eq!( + prepared.auth, + RequestAuth::Bearer { + token: "sk-test".to_string() + } + ); + assert!( + prepared + .upstream_headers + .iter() + .any(|(name, value)| name.eq_ignore_ascii_case("authorization") + && value == "Bearer sk-test"), + "prepare did not carry the bearer token" + ); + } + + fn decline_reason( + model: &str, + provider: Option<&str>, + messages: Value, + params: Value, + ) -> Option<&'static str> { + let params = match params { + Value::Object(map) => map, + other => panic!("params must be an object, got {other}"), + }; + crate::chat_completions::chat_completions_decline_reason(model, provider, messages, ¶ms) + } + + #[test] + fn the_gate_accepts_what_prepare_accepts() { + assert_eq!( + decline_reason( + "anthropic/claude-sonnet-4-5", + None, + json!([{"role": "user", "content": "hi"}]), + json!({"max_tokens": 16}), + ), + None + ); + } + + #[test] + fn the_gate_declines_without_resolving_credentials_or_calling_out() { + assert_eq!( + decline_reason( + "anthropic/claude-sonnet-4-5", + None, + json!([{"role": "user", "content": "hi"}]), + json!({"stream": true}), + ), + Some("streaming") + ); + assert_eq!( + decline_reason( + "openai/gpt-4o", + None, + json!([{"role": "user", "content": "hi"}]), + json!({}), + ), + Some("provider is not on the rust chat completions path") + ); + assert_eq!( + decline_reason( + "claude-sonnet-4-5", + None, + json!([{"role": "user", "content": "hi"}]), + json!({}), + ), + Some("provider is not on the rust chat completions path") + ); + assert_eq!( + decline_reason( + "anthropic/claude-sonnet-4-5", + None, + json!("nope"), + json!({}) + ), + Some("unreadable message list") + ); + assert_eq!( + decline_reason("anthropic/claude-sonnet-4-5", None, json!([]), json!({})), + Some("empty message list") + ); + } + + #[test] + fn the_gate_agrees_with_prepare_on_every_case_it_accepts() { + // A gate that accepts what prepare then declines would make the host emit + // its pre-call logging on a path that falls back, so pin the agreement. + for (messages, params) in [ + ( + json!([{"role": "user", "content": "hi"}]), + json!({"max_tokens": 8}), + ), + ( + json!([{"role": "system", "content": "s"}, {"role": "user", "content": "hi"}]), + json!({"temperature": 0.1}), + ), + ( + json!([{"role": "user", "content": "hi"}, {"role": "assistant", "content": "yo"}]), + json!({}), + ), + ] { + assert_eq!( + decline_reason( + "anthropic/claude-sonnet-4-5", + None, + messages.clone(), + params.clone() + ), + None, + "gate declined {messages}" + ); + prepare_chat_completions_call(request( + "anthropic/claude-sonnet-4-5", + None, + messages.clone(), + params, + )) + .unwrap_or_else(|error| panic!("prepare declined {messages}: {error}")); + } + } + + mod round_trip { + use tokio::{ + io::{AsyncReadExt, AsyncWriteExt}, + net::{TcpListener, TcpStream}, + }; + + use super::*; + use crate::chat_completions::chat_completions; + + async fn read_http_request(socket: &mut TcpStream) -> String { + let mut request = Vec::new(); + let mut buffer = [0_u8; 1024]; + let header_end = loop { + let n = socket.read(&mut buffer).await.expect("reads request"); + if n == 0 { + break request.len(); + } + request.extend_from_slice(&buffer[..n]); + if let Some(position) = request.windows(4).position(|window| window == b"\r\n\r\n") + { + break position + 4; + } + }; + let headers = String::from_utf8_lossy(&request[..header_end]); + let content_length = headers + .lines() + .find_map(|line| { + let (name, value) = line.split_once(':')?; + name.eq_ignore_ascii_case("content-length") + .then(|| value.trim().parse::().ok()) + .flatten() + }) + .unwrap_or(0); + while request.len().saturating_sub(header_end) < content_length { + let n = socket.read(&mut buffer).await.expect("reads body"); + if n == 0 { + break; + } + request.extend_from_slice(&buffer[..n]); + } + String::from_utf8(request).expect("request is utf8") + } + + fn http_response(status: &str, body: &str) -> String { + format!( + "HTTP/1.1 {status}\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{}", + body.len(), + body + ) + } + + /// Serve one request from a stub upstream and hand back what it received. + async fn serve_once( + status: &'static str, + body: &'static str, + ) -> (String, tokio::task::JoinHandle) { + let listener = TcpListener::bind("127.0.0.1:0").await.expect("binds"); + let port = listener.local_addr().expect("addr").port(); + let handle = tokio::spawn(async move { + let (mut socket, _) = listener.accept().await.expect("accepts"); + let received = read_http_request(&mut socket).await; + socket + .write_all(http_response(status, body).as_bytes()) + .await + .expect("writes response"); + socket.flush().await.expect("flushes"); + received + }); + (format!("http://127.0.0.1:{port}/v1/messages"), handle) + } + + fn call(api_base: &str, messages: Value, params: Value) -> ChatCompletionsRequest<'_> { + ChatCompletionsRequest { + model: "anthropic/claude-sonnet-4-5", + messages, + optional_params: match params { + Value::Object(map) => map, + other => panic!("params must be an object, got {other}"), + }, + api_key: Some("sk-test"), + api_base: Some(api_base), + custom_llm_provider: None, + extra_headers: None, + timeout: Some(std::time::Duration::from_secs(10)), + } + } + + const GOOD_BODY: &str = r#"{"id":"msg_1","type":"message","role":"assistant","model":"claude-sonnet-4-5-20260101","content":[{"type":"text","text":"hello"}],"stop_reason":"end_turn","stop_sequence":null,"usage":{"input_tokens":11,"output_tokens":4}}"#; + + #[tokio::test] + async fn round_trip_sends_the_translated_body_and_normalizes_the_response() { + let (api_base, handle) = serve_once("200 OK", GOOD_BODY).await; + let response = chat_completions(call( + &api_base, + json!([ + {"role": "system", "content": "be terse"}, + {"role": "user", "content": "hi"} + ]), + json!({"max_tokens": 16}), + )) + .await + .expect("call succeeds"); + + let received = handle.await.expect("server task"); + let sent: Value = serde_json::from_str( + received + .split_once("\r\n\r\n") + .expect("request has a body") + .1, + ) + .expect("body is json"); + assert_eq!( + sent["messages"], + json!([{"role": "user", "content": [{"type": "text", "text": "hi"}]}]) + ); + assert_eq!( + sent["system"], + json!([{"type": "text", "text": "be terse"}]) + ); + assert_eq!(sent["max_tokens"], json!(16)); + assert!(received.to_lowercase().contains("x-api-key: sk-test")); + + assert_eq!( + response.choices[0].message.content.as_deref(), + Some("hello") + ); + assert_eq!(response.usage.total_tokens, 15); + } + + #[tokio::test] + async fn a_response_it_cannot_normalize_is_reported_as_already_sent() { + // The provider was called and billed, so the host must not retry this + // on its own path. `MissingField` here would read as a pre-send + // decline and be retried; `InvalidResponse` cannot. + const NO_USAGE: &str = + r#"{"model":"m","content":[{"type":"text","text":"hi"}],"stop_reason":"end_turn"}"#; + let (api_base, handle) = serve_once("200 OK", NO_USAGE).await; + let err = chat_completions(call( + &api_base, + json!([{"role": "user", "content": "hi"}]), + json!({"max_tokens": 16}), + )) + .await + .expect_err("response cannot be normalized"); + handle.await.expect("server task"); + assert!( + matches!(err, Error::InvalidResponse(_)), + "expected a post-send error, got {err:?}" + ); + } + + #[tokio::test] + async fn a_tool_use_block_in_the_response_is_also_reported_as_already_sent() { + const TOOL_USE: &str = r#"{"model":"m","content":[{"type":"tool_use","id":"t","name":"f","input":{}}],"stop_reason":"tool_use","usage":{"input_tokens":1,"output_tokens":1}}"#; + let (api_base, handle) = serve_once("200 OK", TOOL_USE).await; + let err = chat_completions(call( + &api_base, + json!([{"role": "user", "content": "hi"}]), + json!({"max_tokens": 16}), + )) + .await + .expect_err("response cannot be normalized"); + handle.await.expect("server task"); + assert!( + matches!(err, Error::InvalidResponse(_)), + "expected a post-send error, got {err:?}" + ); + } + + #[tokio::test] + async fn an_upstream_error_status_keeps_its_code() { + let (api_base, handle) = + serve_once("429 Too Many Requests", r#"{"error":"slow down"}"#).await; + let err = chat_completions(call( + &api_base, + json!([{"role": "user", "content": "hi"}]), + json!({"max_tokens": 16}), + )) + .await + .expect_err("upstream rejects"); + handle.await.expect("server task"); + assert!( + matches!( + err, + Error::Transport(litellm_http::transport::Error::Http { status: 429, .. }) + ), + "expected a 429, got {err:?}" + ); + } + + #[tokio::test] + async fn a_connection_that_is_never_established_declines_instead_of_failing() { + // Nothing was sent, so nothing was billed and the host can still serve + // the request. Classing this with the post-send failures would turn a + // recoverable fallback into a user-facing error on exactly the + // deployments whose transport is configured only on the Python client. + let port = { + let listener = TcpListener::bind("127.0.0.1:0").await.expect("binds"); + listener.local_addr().expect("has an address").port() + // Dropped here, so the port is closed and the connect is refused. + }; + let err = chat_completions(call( + &format!("http://127.0.0.1:{port}/v1/messages"), + json!([{"role": "user", "content": "hi"}]), + json!({"max_tokens": 16}), + )) + .await + .expect_err("nothing is listening"); + assert!( + matches!( + err, + Error::Transport(litellm_http::transport::Error::Connect(_)) + ), + "expected a pre-send connect failure, got {err:?}" + ); + } + + #[test] + fn response_errors_collapse_to_one_variant_that_can_only_mean_already_sent() { + use crate::chat_completions::handler::as_response_error; + + for original in [ + Error::MissingField("usage"), + Error::Unsupported("non-text response content block"), + Error::InvalidRequest("whatever".to_string()), + Error::Auth(litellm_auth::Error::InvalidHeader), + ] { + let label = format!("{original:?}"); + assert!( + matches!(as_response_error(original), Error::InvalidResponse(_)), + "{label} must not stay retryable once the provider has answered" + ); + } + // An upstream status is already unambiguous, so it survives intact. + assert!(matches!( + as_response_error(Error::Transport(litellm_http::transport::Error::Http { + status: 500, + body: "boom".to_string() + })), + Error::Transport(litellm_http::transport::Error::Http { status: 500, .. }) + )); + } + } +} diff --git a/litellm-rust/crates/core/src/chat_completions/tests.rs b/litellm-rust/crates/core/src/chat_completions/tests.rs deleted file mode 100644 index dd5938cf168..00000000000 --- a/litellm-rust/crates/core/src/chat_completions/tests.rs +++ /dev/null @@ -1,833 +0,0 @@ -use litellm_llms::base_llm::chat::transformation::RequestAuth; -use serde_json::{Map, Value, json}; - -use super::{ - Error, - prepare::{prepare_provider_request, resolve_request}, -}; -use crate::chat_completions::types::{ChatCompletionsRequest, ProviderChatCompletionsRequest}; - -fn prepare_chat_completions_call( - request: ChatCompletionsRequest<'_>, -) -> Result { - prepare_provider_request(resolve_request(request)?) -} - -fn request<'a>( - model: &'a str, - provider: Option<&'a str>, - messages: Value, - optional_params: Value, -) -> ChatCompletionsRequest<'a> { - ChatCompletionsRequest { - model, - messages, - optional_params: match optional_params { - Value::Object(map) => map, - other => panic!("params must be an object, got {other}"), - }, - api_key: Some("sk-test"), - api_base: None, - custom_llm_provider: provider, - extra_headers: None, - timeout: None, - } -} - -/// `ProviderChatCompletionsRequest` deliberately has no `Debug` (its headers -/// carry resolved credentials), so unwrap the failure case by hand. -fn decline(request: ChatCompletionsRequest<'_>) -> Error { - match prepare_chat_completions_call(request) { - Err(error) => error, - Ok(prepared) => panic!("expected a decline, prepared a call to {}", prepared.url), - } -} - -#[test] -fn resolves_the_provider_from_the_model_prefix() { - let prepared = prepare_chat_completions_call(request( - "anthropic/claude-sonnet-4-5", - None, - json!([{"role": "user", "content": "hi"}]), - json!({"max_tokens": 16}), - )) - .expect("prepares"); - assert_eq!(prepared.model, "claude-sonnet-4-5"); - assert_eq!(prepared.url, "https://api.anthropic.com/v1/messages"); - assert_eq!(prepared.body["model"], json!("claude-sonnet-4-5")); -} - -#[test] -fn strips_an_explicit_provider_prefix_from_the_model() { - let prepared = prepare_chat_completions_call(request( - "anthropic/claude-sonnet-4-5", - Some("anthropic"), - json!([{"role": "user", "content": "hi"}]), - json!({}), - )) - .expect("prepares"); - assert_eq!(prepared.model, "claude-sonnet-4-5"); -} - -#[test] -fn adds_the_auth_and_default_headers() { - let prepared = prepare_chat_completions_call(request( - "claude-sonnet-4-5", - Some("anthropic"), - json!([{"role": "user", "content": "hi"}]), - json!({}), - )) - .expect("prepares"); - assert!( - prepared - .upstream_headers - .contains(&("x-api-key".to_string(), "sk-test".to_string())) - ); - assert!( - prepared - .upstream_headers - .contains(&("anthropic-version".to_string(), "2023-06-01".to_string())) - ); - assert!(matches!( - prepared.auth, - RequestAuth::Header { - name: "x-api-key", - .. - } - )); -} - -#[test] -fn the_deployment_credential_replaces_a_caller_supplied_auth_header() { - // Python builds `{**headers, **anthropic_headers}`, so the deployment's key - // overwrites a forwarded one. Honouring the caller's would let whoever sends - // the request choose the Anthropic principal it bills to. - let mut call = request( - "claude-sonnet-4-5", - Some("anthropic"), - json!([{"role": "user", "content": "hi"}]), - json!({}), - ); - call.extra_headers = Some(Map::from_iter([( - "X-Api-Key".to_string(), - json!("sk-caller"), - )])); - let prepared = prepare_chat_completions_call(call).expect("prepares"); - let keys: Vec<_> = prepared - .upstream_headers - .iter() - .filter(|(name, _)| name.eq_ignore_ascii_case("x-api-key")) - .collect(); - assert_eq!(keys.len(), 1, "got {:?}", prepared.upstream_headers); - assert_eq!(keys[0].1, "sk-test"); -} - -#[test] -fn a_forwarded_authorization_header_suppresses_the_resolved_api_key_header() { - // Anthropic's `validate_environment` pops `x-api-key` and sets `authorization` - // for an OAuth token, so re-adding the key here would put the credential into - // a header the host removed on purpose. - let mut call = request( - "claude-sonnet-4-5", - Some("anthropic"), - json!([{"role": "user", "content": "hi"}]), - json!({}), - ); - call.extra_headers = Some(Map::from_iter([ - ( - "Authorization".to_string(), - json!("Bearer sk-ant-oat01-token"), - ), - ("X-Api-Key".to_string(), json!("sk-caller")), - ])); - let prepared = prepare_chat_completions_call(call).expect("prepares"); - assert!( - !prepared - .upstream_headers - .iter() - .any(|(name, value)| name.eq_ignore_ascii_case("x-api-key") && value == "sk-test"), - "the resolved key must not be applied over an OAuth bearer, got {:?}", - prepared.upstream_headers - ); - assert!( - prepared - .upstream_headers - .iter() - .any(|(name, value)| name.eq_ignore_ascii_case("authorization") - && value == "Bearer sk-ant-oat01-token") - ); -} - -#[test] -fn an_unrelated_forwarded_authorization_does_not_defer_the_resolved_key() { - // Only an OAuth bearer replaces the credential. Python sends the deployment's - // `x-api-key` alongside any other forwarded `authorization`, so deferring on - // the mere presence of that header would drop the deployment's auth. - let mut call = request( - "claude-sonnet-4-5", - Some("anthropic"), - json!([{"role": "user", "content": "hi"}]), - json!({}), - ); - call.extra_headers = Some(Map::from_iter([ - ("Authorization".to_string(), json!("Bearer unrelated")), - ("X-Api-Key".to_string(), json!("sk-caller")), - ])); - let prepared = prepare_chat_completions_call(call).expect("prepares"); - let keys: Vec<_> = prepared - .upstream_headers - .iter() - .filter(|(name, _)| name.eq_ignore_ascii_case("x-api-key")) - .collect(); - assert_eq!(keys.len(), 1, "got {:?}", prepared.upstream_headers); - assert_eq!(keys[0].1, "sk-test"); - assert!( - prepared - .upstream_headers - .iter() - .any(|(name, value)| name.eq_ignore_ascii_case("authorization") - && value == "Bearer unrelated"), - "the unrelated authorization must survive, got {:?}", - prepared.upstream_headers - ); -} - -#[test] -fn declines_an_unsupported_request_before_resolving_credentials() { - let mut call = request( - "claude-sonnet-4-5", - Some("anthropic"), - json!([{"role": "user", "content": "hi"}]), - json!({"stream": true}), - ); - call.api_key = None; - // No api_key is set and no env is consulted: the gate must run first, so the - // error is the decline rather than a missing-credential error. - assert_eq!(decline(call), Error::Unsupported("streaming")); -} - -#[test] -fn rejects_an_unknown_provider() { - assert_eq!( - decline(request( - "openai/gpt-4o", - None, - json!([{"role": "user", "content": "hi"}]), - json!({}), - )), - Error::InvalidProvider("openai".to_string()) - ); -} - -#[test] -fn rejects_a_model_with_no_resolvable_provider() { - assert!(matches!( - decline(request( - "claude-sonnet-4-5", - None, - json!([{"role": "user", "content": "hi"}]), - json!({}), - )), - Error::InvalidProvider(_) - )); -} - -#[test] -fn rejects_an_empty_or_malformed_message_list() { - assert_eq!( - decline(request( - "anthropic/claude-sonnet-4-5", - None, - json!([]), - json!({}), - )), - Error::InvalidRequest("chat completions requires at least one message".to_string()) - ); - assert!(matches!( - decline(request( - "anthropic/claude-sonnet-4-5", - None, - json!("not a list"), - json!({}), - )), - Error::InvalidRequest(_) - )); -} - -#[test] -fn rejects_non_string_extra_headers() { - let mut call = request( - "anthropic/claude-sonnet-4-5", - None, - json!([{"role": "user", "content": "hi"}]), - json!({}), - ); - call.extra_headers = Some(Map::from_iter([("x-trace".to_string(), json!(7))])); - assert_eq!( - decline(call), - Error::Headers(litellm_http::request::HeaderError { - context: "chat completions", - name: "x-trace".to_string(), - actual: "number", - }) - ); -} - -#[test] -fn prepares_a_bedrock_call_without_resolving_credentials() { - let mut call = request( - "bedrock/us-east-1/anthropic.claude-v2", - None, - json!([{"role": "user", "content": "hi"}]), - json!({"maxTokens": 16}), - ); - call.api_key = None; - let prepared = prepare_chat_completions_call(call).expect("prepares"); - assert_eq!( - prepared.url, - "https://bedrock-runtime.us-east-1.amazonaws.com/model/anthropic.claude-v2/converse" - ); - assert_eq!( - prepared.auth, - RequestAuth::AwsSigV4 { - region: "us-east-1".to_string(), - service: "bedrock", - } - ); - // SigV4 signs the serialized body, so prepare must not have added an - // Authorization header; the handler does it. - assert!( - !prepared - .upstream_headers - .iter() - .any(|(name, _)| name.eq_ignore_ascii_case("authorization")) - ); - assert_eq!(prepared.body["inferenceConfig"], json!({"maxTokens": 16})); -} - -#[tokio::test] -async fn a_forwarded_client_header_does_not_enter_the_bedrock_signature() { - // Python signs only the AWS header set and reattaches the rest, so a header - // the caller forwarded rides along without joining the canonical request. - // Signing it makes Converse 403 on a deployment that works on Python. - let mut call = request( - "bedrock/us-east-1/anthropic.claude-v2", - None, - json!([{"role": "user", "content": "hi"}]), - json!({ - "maxTokens": 16, - "aws_access_key_id": "AKIDEXAMPLE", - "aws_secret_access_key": "wJalrXUtnFEMI/K7MDENG+bPxRfiCYEXAMPLEKEY" - }), - ); - // A key would resolve to a bearer token and never reach the signer. - call.api_key = None; - call.extra_headers = Some(Map::from_iter([( - "x-request-id".to_string(), - json!("abc-123"), - )])); - let prepared = prepare_chat_completions_call(call).expect("prepares"); - let signed = super::handler::outbound_request(&prepared) - .await - .expect("signs"); - - let authorization = signed - .header("authorization") - .expect("carries an authorization header") - .to_string(); - assert!( - authorization.starts_with("AWS4-HMAC-SHA256"), - "expected a SigV4 signature, got {authorization}" - ); - assert!( - !authorization.contains("x-request-id"), - "forwarded header reached SignedHeaders: {authorization}" - ); - // It still goes on the wire, it is just not part of the signature. - assert!( - signed - .headers() - .iter() - .any(|(name, value)| name == "x-request-id" && value == "abc-123"), - "forwarded header was dropped instead of reattached" - ); -} - -#[tokio::test] -async fn a_forwarded_header_the_signer_computes_declines_to_python() { - // Reattaching the caller's copy next to the computed one puts the name on - // the wire twice and Bedrock rejects the pair, so a request carrying one - // has to go to Python instead of being signed here. - for forwarded in [ - "Authorization", - "x-amz-date", - "x-amz-security-token", - "Date", - ] { - let mut call = request( - "bedrock/us-east-1/anthropic.claude-v2", - None, - json!([{"role": "user", "content": "hi"}]), - json!({ - "maxTokens": 16, - "aws_access_key_id": "AKIDEXAMPLE", - "aws_secret_access_key": "wJalrXUtnFEMI/K7MDENG+bPxRfiCYEXAMPLEKEY" - }), - ); - call.api_key = None; - call.extra_headers = Some(Map::from_iter([(forwarded.to_string(), json!("forged"))])); - let prepared = prepare_chat_completions_call(call).expect("prepares"); - let error = super::handler::outbound_request(&prepared) - .await - .expect_err("{forwarded} should decline instead of being signed"); - assert!( - matches!(error, Error::Unsupported(_)), - "{forwarded} declined as {error:?}, which the host would not fall back on" - ); - } -} - -#[test] -fn a_bedrock_deployment_bearer_outranks_a_forwarded_authorization() { - // `get_request_headers` assigns `headers["Authorization"]` unconditionally - // once a bearer token resolves, so the deployment's identity wins on - // Python. Keeping the caller's would authorize and bill the call as a - // different principal, and only when the deployment carries `rust: true`. - let mut call = request( - "bedrock/us-east-1/anthropic.claude-v2", - None, - json!([{"role": "user", "content": "hi"}]), - json!({"maxTokens": 16}), - ); - call.extra_headers = Some(Map::from_iter([( - "Authorization".to_string(), - json!("Bearer caller-supplied"), - )])); - let prepared = prepare_chat_completions_call(call).expect("prepares"); - let authorizations: Vec<_> = prepared - .upstream_headers - .iter() - .filter(|(name, _)| name.eq_ignore_ascii_case("authorization")) - .map(|(_, value)| value.as_str()) - .collect(); - assert_eq!( - authorizations, - vec!["Bearer sk-test"], - "the deployment token must be the only authorization on the wire" - ); -} - -#[test] -fn an_anthropic_forwarded_oauth_bearer_still_outranks_the_resolved_key() { - // The opposite precedence, and deliberate: Anthropic's own transform - // honours a forwarded OAuth bearer, so the Bedrock fix above must not be - // generalized into a rule that the configured key always wins. - // - // An OAuth bearer is the whole of that exception. This forwarded a plain - // `x-api-key` until round 17, which read as the same claim and was not: - // Python overwrites a forwarded `x-api-key` with the deployment's. - let mut call = request( - "claude-sonnet-4-5", - Some("anthropic"), - json!([{"role": "user", "content": "hi"}]), - json!({}), - ); - call.extra_headers = Some(Map::from_iter([( - "authorization".to_string(), - json!("Bearer sk-ant-oat01-forwarded"), - )])); - let prepared = prepare_chat_completions_call(call).expect("prepares"); - let keys: Vec<_> = prepared - .upstream_headers - .iter() - .filter(|(name, _)| name.eq_ignore_ascii_case("x-api-key")) - .map(|(_, value)| value.as_str()) - .collect(); - assert!(keys.is_empty(), "got {:?}", prepared.upstream_headers); - assert!( - prepared - .upstream_headers - .iter() - .any(|(name, value)| name.eq_ignore_ascii_case("authorization") - && value == "Bearer sk-ant-oat01-forwarded") - ); -} - -#[test] -fn a_bedrock_api_key_is_sent_as_a_bearer_token_instead_of_being_signed() { - // The configured bearer identity has its own account and quota boundary, - // so a request carrying one must not be signed as whatever principal the - // host's AWS credentials resolve to. - let prepared = prepare_chat_completions_call(request( - "bedrock/us-east-1/anthropic.claude-v2", - None, - json!([{"role": "user", "content": "hi"}]), - json!({"maxTokens": 16}), - )) - .expect("prepares"); - assert_eq!( - prepared.auth, - RequestAuth::Bearer { - token: "sk-test".to_string() - } - ); - assert!( - prepared - .upstream_headers - .iter() - .any(|(name, value)| name.eq_ignore_ascii_case("authorization") - && value == "Bearer sk-test"), - "prepare did not carry the bearer token" - ); -} - -fn decline_reason( - model: &str, - provider: Option<&str>, - messages: Value, - params: Value, -) -> Option<&'static str> { - let params = match params { - Value::Object(map) => map, - other => panic!("params must be an object, got {other}"), - }; - super::chat_completions_decline_reason(model, provider, messages, ¶ms) -} - -#[test] -fn the_gate_accepts_what_prepare_accepts() { - assert_eq!( - decline_reason( - "anthropic/claude-sonnet-4-5", - None, - json!([{"role": "user", "content": "hi"}]), - json!({"max_tokens": 16}), - ), - None - ); -} - -#[test] -fn the_gate_declines_without_resolving_credentials_or_calling_out() { - assert_eq!( - decline_reason( - "anthropic/claude-sonnet-4-5", - None, - json!([{"role": "user", "content": "hi"}]), - json!({"stream": true}), - ), - Some("streaming") - ); - assert_eq!( - decline_reason( - "openai/gpt-4o", - None, - json!([{"role": "user", "content": "hi"}]), - json!({}), - ), - Some("provider is not on the rust chat completions path") - ); - assert_eq!( - decline_reason( - "claude-sonnet-4-5", - None, - json!([{"role": "user", "content": "hi"}]), - json!({}), - ), - Some("provider is not on the rust chat completions path") - ); - assert_eq!( - decline_reason( - "anthropic/claude-sonnet-4-5", - None, - json!("nope"), - json!({}) - ), - Some("unreadable message list") - ); - assert_eq!( - decline_reason("anthropic/claude-sonnet-4-5", None, json!([]), json!({})), - Some("empty message list") - ); -} - -#[test] -fn the_gate_agrees_with_prepare_on_every_case_it_accepts() { - // A gate that accepts what prepare then declines would make the host emit - // its pre-call logging on a path that falls back, so pin the agreement. - for (messages, params) in [ - ( - json!([{"role": "user", "content": "hi"}]), - json!({"max_tokens": 8}), - ), - ( - json!([{"role": "system", "content": "s"}, {"role": "user", "content": "hi"}]), - json!({"temperature": 0.1}), - ), - ( - json!([{"role": "user", "content": "hi"}, {"role": "assistant", "content": "yo"}]), - json!({}), - ), - ] { - assert_eq!( - decline_reason( - "anthropic/claude-sonnet-4-5", - None, - messages.clone(), - params.clone() - ), - None, - "gate declined {messages}" - ); - prepare_chat_completions_call(request( - "anthropic/claude-sonnet-4-5", - None, - messages.clone(), - params, - )) - .unwrap_or_else(|error| panic!("prepare declined {messages}: {error}")); - } -} - -mod round_trip { - use tokio::{ - io::{AsyncReadExt, AsyncWriteExt}, - net::{TcpListener, TcpStream}, - }; - - use super::*; - use crate::chat_completions::chat_completions; - - async fn read_http_request(socket: &mut TcpStream) -> String { - let mut request = Vec::new(); - let mut buffer = [0_u8; 1024]; - let header_end = loop { - let n = socket.read(&mut buffer).await.expect("reads request"); - if n == 0 { - break request.len(); - } - request.extend_from_slice(&buffer[..n]); - if let Some(position) = request.windows(4).position(|window| window == b"\r\n\r\n") { - break position + 4; - } - }; - let headers = String::from_utf8_lossy(&request[..header_end]); - let content_length = headers - .lines() - .find_map(|line| { - let (name, value) = line.split_once(':')?; - name.eq_ignore_ascii_case("content-length") - .then(|| value.trim().parse::().ok()) - .flatten() - }) - .unwrap_or(0); - while request.len().saturating_sub(header_end) < content_length { - let n = socket.read(&mut buffer).await.expect("reads body"); - if n == 0 { - break; - } - request.extend_from_slice(&buffer[..n]); - } - String::from_utf8(request).expect("request is utf8") - } - - fn http_response(status: &str, body: &str) -> String { - format!( - "HTTP/1.1 {status}\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{}", - body.len(), - body - ) - } - - /// Serve one request from a stub upstream and hand back what it received. - async fn serve_once( - status: &'static str, - body: &'static str, - ) -> (String, tokio::task::JoinHandle) { - let listener = TcpListener::bind("127.0.0.1:0").await.expect("binds"); - let port = listener.local_addr().expect("addr").port(); - let handle = tokio::spawn(async move { - let (mut socket, _) = listener.accept().await.expect("accepts"); - let received = read_http_request(&mut socket).await; - socket - .write_all(http_response(status, body).as_bytes()) - .await - .expect("writes response"); - socket.flush().await.expect("flushes"); - received - }); - (format!("http://127.0.0.1:{port}/v1/messages"), handle) - } - - fn call(api_base: &str, messages: Value, params: Value) -> ChatCompletionsRequest<'_> { - ChatCompletionsRequest { - model: "anthropic/claude-sonnet-4-5", - messages, - optional_params: match params { - Value::Object(map) => map, - other => panic!("params must be an object, got {other}"), - }, - api_key: Some("sk-test"), - api_base: Some(api_base), - custom_llm_provider: None, - extra_headers: None, - timeout: Some(std::time::Duration::from_secs(10)), - } - } - - const GOOD_BODY: &str = r#"{"id":"msg_1","type":"message","role":"assistant","model":"claude-sonnet-4-5-20260101","content":[{"type":"text","text":"hello"}],"stop_reason":"end_turn","stop_sequence":null,"usage":{"input_tokens":11,"output_tokens":4}}"#; - - #[tokio::test] - async fn round_trip_sends_the_translated_body_and_normalizes_the_response() { - let (api_base, handle) = serve_once("200 OK", GOOD_BODY).await; - let response = chat_completions(call( - &api_base, - json!([ - {"role": "system", "content": "be terse"}, - {"role": "user", "content": "hi"} - ]), - json!({"max_tokens": 16}), - )) - .await - .expect("call succeeds"); - - let received = handle.await.expect("server task"); - let sent: Value = serde_json::from_str( - received - .split_once("\r\n\r\n") - .expect("request has a body") - .1, - ) - .expect("body is json"); - assert_eq!( - sent["messages"], - json!([{"role": "user", "content": [{"type": "text", "text": "hi"}]}]) - ); - assert_eq!( - sent["system"], - json!([{"type": "text", "text": "be terse"}]) - ); - assert_eq!(sent["max_tokens"], json!(16)); - assert!(received.to_lowercase().contains("x-api-key: sk-test")); - - assert_eq!( - response.choices[0].message.content.as_deref(), - Some("hello") - ); - assert_eq!(response.usage.total_tokens, 15); - } - - #[tokio::test] - async fn a_response_it_cannot_normalize_is_reported_as_already_sent() { - // The provider was called and billed, so the host must not retry this - // on its own path. `MissingField` here would read as a pre-send - // decline and be retried; `InvalidResponse` cannot. - const NO_USAGE: &str = - r#"{"model":"m","content":[{"type":"text","text":"hi"}],"stop_reason":"end_turn"}"#; - let (api_base, handle) = serve_once("200 OK", NO_USAGE).await; - let err = chat_completions(call( - &api_base, - json!([{"role": "user", "content": "hi"}]), - json!({"max_tokens": 16}), - )) - .await - .expect_err("response cannot be normalized"); - handle.await.expect("server task"); - assert!( - matches!(err, Error::InvalidResponse(_)), - "expected a post-send error, got {err:?}" - ); - } - - #[tokio::test] - async fn a_tool_use_block_in_the_response_is_also_reported_as_already_sent() { - const TOOL_USE: &str = r#"{"model":"m","content":[{"type":"tool_use","id":"t","name":"f","input":{}}],"stop_reason":"tool_use","usage":{"input_tokens":1,"output_tokens":1}}"#; - let (api_base, handle) = serve_once("200 OK", TOOL_USE).await; - let err = chat_completions(call( - &api_base, - json!([{"role": "user", "content": "hi"}]), - json!({"max_tokens": 16}), - )) - .await - .expect_err("response cannot be normalized"); - handle.await.expect("server task"); - assert!( - matches!(err, Error::InvalidResponse(_)), - "expected a post-send error, got {err:?}" - ); - } - - #[tokio::test] - async fn an_upstream_error_status_keeps_its_code() { - let (api_base, handle) = - serve_once("429 Too Many Requests", r#"{"error":"slow down"}"#).await; - let err = chat_completions(call( - &api_base, - json!([{"role": "user", "content": "hi"}]), - json!({"max_tokens": 16}), - )) - .await - .expect_err("upstream rejects"); - handle.await.expect("server task"); - assert!( - matches!( - err, - Error::Transport(litellm_http::transport::Error::Http { status: 429, .. }) - ), - "expected a 429, got {err:?}" - ); - } - - #[tokio::test] - async fn a_connection_that_is_never_established_declines_instead_of_failing() { - // Nothing was sent, so nothing was billed and the host can still serve - // the request. Classing this with the post-send failures would turn a - // recoverable fallback into a user-facing error on exactly the - // deployments whose transport is configured only on the Python client. - let port = { - let listener = TcpListener::bind("127.0.0.1:0").await.expect("binds"); - listener.local_addr().expect("has an address").port() - // Dropped here, so the port is closed and the connect is refused. - }; - let err = chat_completions(call( - &format!("http://127.0.0.1:{port}/v1/messages"), - json!([{"role": "user", "content": "hi"}]), - json!({"max_tokens": 16}), - )) - .await - .expect_err("nothing is listening"); - assert!( - matches!( - err, - Error::Transport(litellm_http::transport::Error::Connect(_)) - ), - "expected a pre-send connect failure, got {err:?}" - ); - } - - #[test] - fn response_errors_collapse_to_one_variant_that_can_only_mean_already_sent() { - use crate::chat_completions::handler::as_response_error; - - for original in [ - Error::MissingField("usage"), - Error::Unsupported("non-text response content block"), - Error::InvalidRequest("whatever".to_string()), - Error::Auth(litellm_auth::Error::InvalidHeader), - ] { - let label = format!("{original:?}"); - assert!( - matches!(as_response_error(original), Error::InvalidResponse(_)), - "{label} must not stay retryable once the provider has answered" - ); - } - // An upstream status is already unambiguous, so it survives intact. - assert!(matches!( - as_response_error(Error::Transport(litellm_http::transport::Error::Http { - status: 500, - body: "boom".to_string() - })), - Error::Transport(litellm_http::transport::Error::Http { status: 500, .. }) - )); - } -} diff --git a/litellm-rust/crates/core/src/messages/common_utils.rs b/litellm-rust/crates/core/src/messages/common_utils.rs index 015e026f6da..4327754ed05 100644 --- a/litellm-rust/crates/core/src/messages/common_utils.rs +++ b/litellm-rust/crates/core/src/messages/common_utils.rs @@ -26,3 +26,186 @@ pub(super) fn string_headers( ) -> Result, Error> { shared_string_headers(HEADER_CONTEXT, extra_headers).map_err(Error::from) } + +#[cfg(test)] +mod tests { + use std::{sync::Arc, time::Duration}; + + use futures_util::future::BoxFuture; + use litellm_secrets::{SecretValue, source::SecretSource}; + use serde_json::{Value, json}; + use tokio::{ + io::{AsyncReadExt, AsyncWriteExt}, + net::{TcpListener, TcpStream}, + }; + + use super::{messages_provider_config, string_headers, truncate_error_body}; + use crate::messages::{ + Error, + route::{LocalMessagesHost, MessagesCall, MessagesOutput, messages_machine}, + types::MessagesShaping, + }; + + struct RecordingSecrets { + values: Vec<(&'static str, String)>, + requested: std::sync::Mutex>, + } + + impl SecretSource for RecordingSecrets { + fn get_secret_str<'a>( + &'a self, + name: &'a str, + ) -> BoxFuture<'a, Result, litellm_secrets::Error>> { + Box::pin(async move { + self.requested.lock().unwrap().push(name.to_string()); + Ok(self + .values + .iter() + .find(|(key, _)| *key == name) + .map(|(_, value)| SecretValue::new(value.clone()))) + }) + } + } + + fn secrets_call() -> MessagesCall { + let Value::Object(body) = json!({ + "model": "claude-sonnet-4-5", + "max_tokens": 16, + "messages": [{"role": "user", "content": "hi"}] + }) else { + unreachable!("literal object") + }; + MessagesCall { + model: "claude-sonnet-4-5".into(), + body, + api_key: None, + api_base: None, + custom_llm_provider: Some("anthropic".into()), + extra_headers: None, + provider_specific_header: None, + timeout: Some(Duration::from_secs(5)), + shaping: MessagesShaping::default(), + } + } + + async fn read_http_request(socket: &mut TcpStream) -> String { + let mut request = Vec::new(); + let mut buffer = [0_u8; 1024]; + let header_end = loop { + let n = socket.read(&mut buffer).await.expect("reads request"); + if n == 0 { + break request.len(); + } + request.extend_from_slice(&buffer[..n]); + if let Some(position) = request.windows(4).position(|window| window == b"\r\n\r\n") { + break position + 4; + } + }; + let headers = String::from_utf8_lossy(&request[..header_end]); + let content_length = headers + .lines() + .find_map(|line| { + let (name, value) = line.split_once(':')?; + name.eq_ignore_ascii_case("content-length") + .then(|| value.trim().parse::().ok()) + .flatten() + }) + .unwrap_or(0); + while request.len().saturating_sub(header_end) < content_length { + let n = socket.read(&mut buffer).await.expect("reads body"); + if n == 0 { + break; + } + request.extend_from_slice(&buffer[..n]); + } + String::from_utf8(request).expect("request is utf8") + } + + #[tokio::test] + async fn route_reads_the_provider_credential_and_base_from_the_secret_source() { + let listener = TcpListener::bind("127.0.0.1:0").await.expect("binds"); + let addr = listener.local_addr().expect("addr"); + let server = tokio::spawn(async move { + let (mut socket, _) = listener.accept().await.expect("accepts request"); + let request = read_http_request(&mut socket).await; + let response_body = r#"{"id":"msg_1","type":"message","role":"assistant","content":[],"model":"claude-sonnet-4-5","stop_reason":"end_turn","usage":{"input_tokens":1,"output_tokens":1}}"#; + let response = format!( + "HTTP/1.1 200 OK\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{}", + response_body.len(), + response_body + ); + socket + .write_all(response.as_bytes()) + .await + .expect("writes response"); + request + }); + let secrets = Arc::new(RecordingSecrets { + values: vec![ + ("ANTHROPIC_API_KEY", "sk-from-manager".to_string()), + ("ANTHROPIC_BASE_URL", format!("http://{addr}")), + ], + requested: std::sync::Mutex::new(Vec::new()), + }); + + let output = litellm_host::run::run( + messages_machine(secrets.clone()), + &LocalMessagesHost::new(secrets_call()), + ) + .await + .expect("messages request succeeds"); + + assert!(matches!(output, MessagesOutput::Message(_))); + let request = server.await.expect("server task completes"); + assert!( + request + .to_ascii_lowercase() + .contains("x-api-key: sk-from-manager"), + "{request}" + ); + let requested = secrets.requested.lock().unwrap().clone(); + assert_eq!( + requested, + messages_provider_config("anthropic") + .unwrap() + .secret_names() + .iter() + .map(ToString::to_string) + .collect::>() + ); + } + + #[test] + fn provider_config_resolves_anthropic_and_azure_ai() { + assert!(messages_provider_config("anthropic").is_some()); + assert!(messages_provider_config("azure_ai").is_some()); + assert!(messages_provider_config("openai").is_none()); + } + + #[test] + fn truncate_error_body_caps_long_payloads() { + let body = "x".repeat(400); + let truncated = truncate_error_body(&body); + assert!(truncated.ends_with("... (truncated)")); + let prefix_chars = truncated + .strip_suffix("... (truncated)") + .expect("truncated marker present") + .chars() + .count(); + assert_eq!(prefix_chars, 256); + } + + #[test] + fn string_headers_rejects_non_string_values() { + let headers = json!({"x-count": 3}).as_object().unwrap().clone(); + let err = string_headers(Some(headers)).expect_err("non-string header rejected"); + assert_eq!( + err, + Error::Headers(litellm_http::request::HeaderError { + context: "messages", + name: "x-count".to_string(), + actual: "number", + }) + ); + } +} diff --git a/litellm-rust/crates/core/src/messages/mod.rs b/litellm-rust/crates/core/src/messages/mod.rs index 8795d4f8507..180eb08810e 100644 --- a/litellm-rust/crates/core/src/messages/mod.rs +++ b/litellm-rust/crates/core/src/messages/mod.rs @@ -46,6 +46,3 @@ pub async fn messages(request: MessagesRequest<'_>) -> Result &'static str { + match self { + Self::Mistral => "mistral/model", + Self::AzureAi => "azure_ai/model", + Self::VertexMistral => "vertex_ai/mistral-ocr-maas", + Self::AzureCohereParse => "azure_ai/cohere-parse", + Self::Cohere => "cohere/model", + } + } + + fn document_type(self) -> &'static str { + match self { + Self::Mistral | Self::AzureAi | Self::VertexMistral => "document_url", + Self::AzureCohereParse | Self::Cohere => "image_url", + } + } + + fn options(self) -> Value { + match self { + Self::Mistral | Self::AzureAi => json!({"pages": [0]}), + Self::VertexMistral => json!({"pages": [0], "vertex_project": "project-1"}), + Self::AzureCohereParse | Self::Cohere => json!({"output_format": "markdown"}), + } + } + } + + /// What the host does to the wire request in `before_send`. + #[derive(Clone, Copy, Debug)] + enum Host { + Detached, + ReplacesDocument, + } + + const REPLACED_DOCUMENT: &str = "data:image/png;base64,cmVwbGFjZWQ="; + + impl Host { + fn before_send(self, wire: WireRequest) -> WireRequest { + let Value::Object(fields) = wire.body else { + return wire; + }; + let body = fields + .into_iter() + .map(|(name, value)| match self { + Self::Detached => (name, value), + Self::ReplacesDocument if name == "document" => { + let document_type = value["type"].clone(); + let key = document_type.as_str().unwrap_or_default().to_string(); + (name, json!({"type": document_type, key: REPLACED_DOCUMENT})) + } + Self::ReplacesDocument => (name, value), + }) + .collect(); + WireRequest { + body: Value::Object(body), + ..wire + } + } + } + + struct Sent { + result: Result<(), Error>, + provider_body: Option, + } + + async fn send(route: Route, host: Host, document_base: &str) -> Sent { + let (base, seen, provider) = + mock_server(vec![MockResponse::json(json!({"pages": []}))]).await; + let document_type = route.document_type(); + let document = + json!({"type": document_type, document_type: format!("{document_base}/scan.png")}); + let request = wire_request_with_document(route.model(), &base, document, route.options()); + let local = + LocalOcrHost::new(request).with_before_send(move |wire, _| Ok(host.before_send(wire))); + let result = perform_ocr_with(local).await.map(|_| ()); + match result { + Ok(()) => provider.await.unwrap(), + Err(_) => provider.abort(), + } + let provider_body = seen + .lock() + .unwrap() + .first() + .map(|request| request_body(request)); + Sent { + result, + provider_body, + } + } + + fn served_document_uri() -> String { + use base64::Engine; + format!( + "data:image/png;base64,{}", + base64::engine::general_purpose::STANDARD.encode(SERVED_DOCUMENT) + ) + } + + #[rstest] + #[case::azure_ai(Route::AzureAi)] + #[case::vertex_mistral(Route::VertexMistral)] + #[case::azure_cohere_parse(Route::AzureCohereParse)] + #[tokio::test] + async fn inlining_routes_send_the_downloaded_document(#[case] route: Route) { + let (document_base, _documents) = document_server().await; + let sent = send(route, Host::Detached, &document_base).await; + sent.result.unwrap(); + assert_eq!( + sent.provider_body.unwrap()["document"][route.document_type()], + json!(served_document_uri()) + ); + } + + #[rstest] + #[tokio::test] + async fn document_replaced_by_the_host_reaches_the_provider( + #[values( + Route::Mistral, + Route::AzureAi, + Route::VertexMistral, + Route::AzureCohereParse, + Route::Cohere + )] + route: Route, + ) { + let (document_base, _documents) = document_server().await; + let sent = send(route, Host::ReplacesDocument, &document_base).await; + sent.result.unwrap(); + assert_eq!( + sent.provider_body.unwrap()["document"][route.document_type()], + json!(REPLACED_DOCUMENT) + ); + } +} diff --git a/litellm-rust/crates/core/src/ocr/mod.rs b/litellm-rust/crates/core/src/ocr/mod.rs index f298f106a5f..270a402c9fa 100644 --- a/litellm-rust/crates/core/src/ocr/mod.rs +++ b/litellm-rust/crates/core/src/ocr/mod.rs @@ -9,36 +9,210 @@ pub mod types; pub mod wire; #[cfg(test)] -#[path = "../../tests/aws_textract_ocr.rs"] -mod aws_textract_tests; +pub(crate) mod test_support { + use std::sync::{Arc, Mutex}; -#[cfg(test)] -#[path = "../../tests/azure_ai_ocr.rs"] -mod azure_ai_tests; -#[cfg(test)] -#[path = "../../tests/azure_document_intelligence_ocr.rs"] -mod azure_document_intelligence_tests; -#[cfg(test)] -#[path = "../../tests/cohere_ocr.rs"] -mod cohere_tests; -#[cfg(test)] -#[path = "../../tests/deepseek_ocr.rs"] -mod deepseek_tests; -#[cfg(test)] -#[path = "../../tests/ocr/document.rs"] -mod document_tests; -#[cfg(test)] -#[path = "../../tests/reducto_ocr.rs"] -mod reducto_tests; -#[cfg(test)] -#[path = "../../tests/ocr/support.rs"] -pub(crate) mod test_support; -#[cfg(test)] -#[path = "../../tests/ocr.rs"] -pub(crate) mod tests; -#[cfg(test)] -#[path = "../../tests/vertex_ai_deepseek_ocr.rs"] -mod vertex_ai_deepseek_tests; -#[cfg(test)] -#[path = "../../tests/vertex_ai_ocr.rs"] -mod vertex_ai_tests; + use futures_util::future::BoxFuture; + use litellm_host::event::WireRequest; + use litellm_llms::base_llm::ocr::{ + error::Error, + handler::{CallHooks, OcrClient}, + transformation::LiteLLMOcrResponse, + }; + use serde_json::{Value, json}; + use tokio::{ + io::{AsyncReadExt, AsyncWriteExt}, + net::TcpListener, + }; + + use crate::ocr::{ + route::{LocalOcrHost, ocr_machine}, + types::LiteLLMOcrRequest, + wire::{OcrWireRequest, decode_request}, + }; + + /// Stands in for a host with no hooks registered: the wire request goes out unchanged + /// and response events go nowhere. + pub(crate) struct NoHooks; + + impl CallHooks for NoHooks { + fn before_send(&self, wire: WireRequest) -> BoxFuture<'_, Result> { + Box::pin(async move { Ok(wire) }) + } + + fn response_received<'a>(&'a self, _body: &'a [u8]) -> BoxFuture<'a, Result<(), Error>> { + Box::pin(async { Ok(()) }) + } + } + + pub(crate) fn ocr_client() -> OcrClient { + let document_http = reqwest::Client::builder() + .redirect(reqwest::redirect::Policy::none()) + .build() + .expect("test document client builds"); + OcrClient::for_test(reqwest::Client::new(), document_http) + } + + pub(crate) async fn perform_ocr( + request: LiteLLMOcrRequest, + ) -> Result { + crate::ocr::client::perform(&ocr_client(), request).await + } + + pub(crate) async fn perform_ocr_with(host: LocalOcrHost) -> Result { + litellm_host::run::run(ocr_machine(ocr_client()), &host).await + } + + pub(crate) fn wire_request(model: &str, base: &str, options: Value) -> LiteLLMOcrRequest { + wire_request_with_document( + model, + base, + json!({"type":"document_url","document_url":"data:application/pdf;base64,YWJj"}), + options, + ) + } + + pub(crate) fn wire_request_with_document( + model: &str, + base: &str, + document: Value, + options: Value, + ) -> LiteLLMOcrRequest { + decode_request(OcrWireRequest { + model: model.into(), + document, + api_key: Some(litellm_auth::SecretValue::new("test-key")), + api_base: Some(base.into()), + custom_llm_provider: None, + extra_headers: None, + optional_params: options.as_object().unwrap().clone(), + input_sources: Default::default(), + timeout_seconds: Some(2.0), + }) + .unwrap() + } + + pub(crate) fn resolved_request( + request: LiteLLMOcrRequest, + ) -> crate::ocr::types::ResolvedOcrRequest { + request + .map_document(crate::ocr::document::prepare_document) + .unwrap() + } + + pub(crate) fn with_source(request: LiteLLMOcrRequest, source: &str) -> LiteLLMOcrRequest { + let request = resolved_request(request); + let document = request.document.clone().with_source(source.into()); + request.with_document(document.into()) + } + + pub(crate) fn request_body(request: &str) -> Value { + serde_json::from_str(request.split_once("\r\n\r\n").unwrap().1).unwrap() + } + + pub(crate) const SERVED_DOCUMENT: &[u8] = b"\x89PNG served document"; + + /// Serves [`SERVED_DOCUMENT`] as `image/png` to every connection until aborted. + pub(crate) async fn document_server() -> (String, tokio::task::JoinHandle<()>) { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let base = format!("http://{}", listener.local_addr().unwrap()); + let task = tokio::spawn(async move { + loop { + let (mut socket, _) = listener.accept().await.unwrap(); + let mut buffer = [0u8; 4096]; + let _ = socket.read(&mut buffer).await.unwrap(); + let head = format!( + "HTTP/1.1 200 OK\r\nContent-Type: image/png\r\nContent-Length: {}\r\nConnection: close\r\n\r\n", + SERVED_DOCUMENT.len() + ); + socket.write_all(head.as_bytes()).await.unwrap(); + socket.write_all(SERVED_DOCUMENT).await.unwrap(); + } + }); + (base, task) + } + + pub(crate) struct MockResponse { + pub status: u16, + pub headers: Vec<(&'static str, String)>, + pub body: Value, + } + + impl MockResponse { + pub fn json(body: Value) -> Self { + Self { + status: 200, + headers: vec![], + body, + } + } + } + + pub(crate) async fn mock_server( + responses: Vec, + ) -> (String, Arc>>, tokio::task::JoinHandle<()>) { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let base = format!("http://{}", listener.local_addr().unwrap()); + let requests = Arc::new(Mutex::new(Vec::new())); + let seen = requests.clone(); + let server_base = base.clone(); + let task = tokio::spawn(async move { + for response in responses { + let (mut socket, _) = listener.accept().await.unwrap(); + let mut bytes = Vec::new(); + let mut buffer = [0u8; 4096]; + let header_end = loop { + let n = socket.read(&mut buffer).await.unwrap(); + assert!(n > 0); + bytes.extend_from_slice(&buffer[..n]); + if let Some(index) = bytes.windows(4).position(|s| s == b"\r\n\r\n") { + break index + 4; + } + }; + let length = String::from_utf8_lossy(&bytes[..header_end]) + .lines() + .find_map(|line| { + let (name, value) = line.split_once(':')?; + name.eq_ignore_ascii_case("content-length") + .then(|| value.trim().parse::().unwrap()) + }) + .unwrap_or(0); + while bytes.len() < header_end + length { + let n = socket.read(&mut buffer).await.unwrap(); + assert!(n > 0); + bytes.extend_from_slice(&buffer[..n]); + } + seen.lock() + .unwrap() + .push(String::from_utf8_lossy(&bytes).into_owned()); + let body = serde_json::to_vec(&response.body).unwrap(); + let headers = response + .headers + .into_iter() + .map(|(name, value)| { + format!("{name}: {}\r\n", value.replace("{base}", &server_base)) + }) + .collect::(); + let head = format!( + "HTTP/1.1 {} OK\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n{}\r\n", + response.status, + body.len(), + headers + ); + socket.write_all(head.as_bytes()).await.unwrap(); + socket.write_all(&body).await.unwrap(); + } + }); + (base, requests, task) + } + + pub(crate) fn header<'a>(request: &'a str, name: &str) -> Option<&'a str> { + request + .lines() + .take_while(|line| !line.is_empty()) + .find_map(|line| { + let (key, value) = line.split_once(':')?; + key.eq_ignore_ascii_case(name).then(|| value.trim()) + }) + } +} diff --git a/litellm-rust/crates/core/src/ocr/route.rs b/litellm-rust/crates/core/src/ocr/route.rs index 26c9ac27102..7f83291bdab 100644 --- a/litellm-rust/crates/core/src/ocr/route.rs +++ b/litellm-rust/crates/core/src/ocr/route.rs @@ -207,3 +207,3589 @@ impl litellm_host::host::Host for LocalOcrHost { Ok(()) } } + +#[cfg(test)] +mod aws_textract_tests { + use std::{collections::BTreeMap, time::SystemTime}; + + use litellm_auth_aws::{Credentials, aws_signature_headers, sign_post}; + use litellm_llms::base_llm::ocr::error::Error; + use serde_json::{Value, json}; + use time::{PrimitiveDateTime, format_description}; + + use crate::ocr::{ + route::LocalOcrHost, + test_support::{ + MockResponse, header, mock_server, perform_ocr_with, request_body, + wire_request_with_document, + }, + types::LiteLLMOcrRequest, + }; + + const ACCESS_KEY_ID: &str = "AKIDEXAMPLE"; + const SECRET_ACCESS_KEY: &str = "wJalrXUtnFEMI/K7MDENG+bPxRfiCYEXAMPLEKEY"; + + fn textract_request(base: &str) -> LiteLLMOcrRequest { + textract_request_for("aws_textract/detect-document-text", base) + } + + fn textract_request_for(model: &str, base: &str) -> LiteLLMOcrRequest { + wire_request_with_document( + model, + &format!("{base}/"), + json!({"type": "image_url", "image_url": "data:image/png;base64,b3JpZ2luYWw="}), + json!({ + "aws_access_key_id": ACCESS_KEY_ID, + "aws_secret_access_key": SECRET_ACCESS_KEY, + "aws_region_name": "eu-west-1" + }), + ) + } + + fn textract_response() -> MockResponse { + MockResponse::json(json!({ + "DocumentMetadata": {"Pages": 1}, + "Blocks": [{"BlockType": "PAGE"}, {"BlockType": "LINE", "Text": "Invoice 12345"}] + })) + } + + /// Recomputes SigV4 over the bytes the server received, at the time the client claimed. + fn expected_authorization(url: &str, raw_request: &str) -> String { + let format = + format_description::parse_borrowed::<2>("[year][month][day]T[hour][minute][second]Z") + .unwrap(); + let signed_at: SystemTime = + PrimitiveDateTime::parse(header(raw_request, "x-amz-date").unwrap(), &format) + .unwrap() + .assume_utc() + .into(); + let headers: BTreeMap = ["content-type", "x-amz-target"] + .into_iter() + .map(|name| { + ( + name.to_string(), + header(raw_request, name).unwrap().to_string(), + ) + }) + .collect(); + let body = raw_request.split_once("\r\n\r\n").unwrap().1; + sign_post( + url, + body.as_bytes(), + &aws_signature_headers(&headers), + "eu-west-1", + "textract", + &Credentials::new(ACCESS_KEY_ID, SECRET_ACCESS_KEY, None, None, "test"), + signed_at, + ) + .unwrap()["Authorization"] + .clone() + } + + #[tokio::test] + async fn the_request_is_signed_for_textract_and_lines_become_the_page() { + let (base, seen, server) = mock_server(vec![textract_response()]).await; + + let response = perform_ocr_with(LocalOcrHost::new(textract_request(&base))) + .await + .unwrap(); + server.await.unwrap(); + + let raw = seen.lock().unwrap()[0].clone(); + assert_eq!( + header(&raw, "x-amz-target"), + Some("Textract.DetectDocumentText") + ); + assert_eq!( + header(&raw, "content-type"), + Some("application/x-amz-json-1.1") + ); + assert_eq!( + request_body(&raw), + json!({"Document": {"Bytes": "b3JpZ2luYWw="}}) + ); + assert_eq!( + header(&raw, "authorization"), + Some(expected_authorization(&format!("{base}/"), &raw).as_str()) + ); + assert_eq!(response.pages[0].markdown, "Invoice 12345"); + assert_eq!(response.usage_info.unwrap().pages_processed, Some(1)); + } + + #[tokio::test] + async fn a_body_rewritten_by_before_send_is_what_gets_signed_and_sent() { + let (base, seen, server) = mock_server(vec![textract_response()]).await; + let host = LocalOcrHost::new(textract_request(&base)).with_before_send(|mut wire, _| { + assert!( + !wire + .headers + .iter() + .any(|(name, _)| name.eq_ignore_ascii_case("authorization")), + "the hook ran after signing" + ); + wire.body["Document"]["Bytes"] = Value::from("cmVkYWN0ZWQ="); + Ok(wire) + }); + + perform_ocr_with(host).await.unwrap(); + server.await.unwrap(); + + let raw = seen.lock().unwrap()[0].clone(); + assert_eq!( + request_body(&raw), + json!({"Document": {"Bytes": "cmVkYWN0ZWQ="}}) + ); + assert_eq!( + header(&raw, "authorization"), + Some(expected_authorization(&format!("{base}/"), &raw).as_str()) + ); + } + + #[tokio::test] + async fn a_multi_page_rejection_reaches_the_caller_with_the_single_page_limit() { + let (base, _, server) = mock_server(vec![MockResponse { + status: 400, + headers: vec![], + body: json!({ + "__type": "UnsupportedDocumentException", + "Message": "Request has unsupported document format" + }), + }]) + .await; + + let error = perform_ocr_with(LocalOcrHost::new(textract_request(&base))) + .await + .unwrap_err(); + server.await.unwrap(); + + let Error::Provider { status, body, .. } = error else { + panic!("expected a provider error, got {error:?}"); + }; + assert_eq!(status, 400); + assert!( + body.contains("multi-page documents are not supported"), + "{body}" + ); + } + + #[tokio::test] + async fn analyze_document_asks_for_layout_and_tables_and_returns_markdown() { + let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ + "DocumentMetadata": {"Pages": 1}, + "Blocks": [ + {"Id": "l1", "BlockType": "LINE", "Text": "Quarterly Report"}, + {"Id": "t", "BlockType": "LAYOUT_TITLE", + "Relationships": [{"Type": "CHILD", "Ids": ["l1"]}]} + ] + }))]) + .await; + let request = textract_request_for("aws_textract/analyze-document", &base); + + let response = perform_ocr_with(LocalOcrHost::new(request)).await.unwrap(); + server.await.unwrap(); + + let raw = seen.lock().unwrap()[0].clone(); + assert_eq!( + header(&raw, "x-amz-target"), + Some("Textract.AnalyzeDocument") + ); + assert_eq!( + request_body(&raw)["FeatureTypes"], + json!(["LAYOUT", "TABLES"]) + ); + assert_eq!( + header(&raw, "authorization"), + Some(expected_authorization(&format!("{base}/"), &raw).as_str()) + ); + assert_eq!(response.pages[0].markdown, "# Quarterly Report"); + } +} + +#[cfg(test)] +mod azure_ai_tests { + use litellm_llms::base_llm::ocr::error::Error; + use serde_json::{Value, json}; + + use crate::ocr::route::LocalOcrHost; + use crate::ocr::test_support::{ + MockResponse, mock_server, perform_ocr, perform_ocr_with, wire_request, + }; + + #[tokio::test] + async fn facade_executes_azure_mistral_with_prepared_auth() { + let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ + "pages":[{"index":0,"markdown":"hello"}], + "usage_info":{"pages_processed":1} + }))]) + .await; + let mut request = wire_request( + "azure_ai/model", + &base, + json!({"include_image_base64":true}), + ); + request.credentials.api_key = None; + request.transport.extra_headers = vec![( + "Authorization".into(), + "Bearer python-prepared-token".into(), + )]; + + let result = perform_ocr(request).await.unwrap(); + server.await.unwrap(); + assert_eq!(result.pages[0].markdown, "hello"); + let requests = seen.lock().unwrap(); + assert_eq!(requests.len(), 1); + assert!(requests[0].starts_with("POST /providers/mistral/azure/ocr ")); + assert!( + requests[0] + .to_ascii_lowercase() + .contains("authorization: bearer python-prepared-token\r\n") + ); + let body: Value = + serde_json::from_str(requests[0].split_once("\r\n\r\n").unwrap().1).unwrap(); + assert_eq!( + body, + json!({ + "model":"model", + "document":{"type":"document_url","document_url":"data:application/pdf;base64,YWJj"}, + "include_image_base64":true + }) + ); + } + + #[tokio::test] + async fn facade_acquires_supplied_entra_token_for_final_request() { + let (base, seen, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await; + let mut request = wire_request( + "azure_ai/model", + &base, + json!({"azure_ad_token":"rust-owned-token"}), + ); + request.credentials.api_key = None; + + perform_ocr(request).await.unwrap(); + server.await.unwrap(); + + let requests = seen.lock().unwrap(); + assert_eq!(requests.len(), 1); + assert!( + requests[0] + .to_ascii_lowercase() + .contains("authorization: bearer rust-owned-token\r\n") + ); + } + + #[tokio::test] + async fn rejects_non_inline_body_after_guardrails() { + let request = wire_request("azure_ai/model", "http://127.0.0.1:1", json!({})); + let host = LocalOcrHost::new(request).with_before_send(|mut wire, _| { + wire.body["document"] = json!({ + "type":"document_url", + "document_url":"https://example.com/not-inline.pdf" + }); + Ok(wire) + }); + let error = perform_ocr_with(host).await.unwrap_err(); + assert!(error.to_string().contains("data URI")); + } + + mod transformation { + use std::sync::{ + Arc, + atomic::{AtomicUsize, Ordering}, + }; + + use litellm_auth::{ + ResolvedCredential, SecretValue, TokenFuture, TokenProvider, TokenProviderHandle, + }; + use rstest::rstest; + use serde_json::json; + + use super::*; + use crate::ocr::{ + test_support::{MockResponse, header, mock_server, perform_ocr}, + types::LiteLLMOcrRequest, + wire::decode_request, + }; + + #[derive(Debug)] + struct CountingToken { + token: fn(usize) -> String, + calls: AtomicUsize, + } + + impl CountingToken { + fn new(token: fn(usize) -> String) -> Arc { + Arc::new(Self { + token, + calls: AtomicUsize::new(0), + }) + } + + fn calls(&self) -> usize { + self.calls.load(Ordering::SeqCst) + } + } + + impl TokenProvider for CountingToken { + fn acquire(&self) -> TokenFuture<'_> { + let call = self.calls.fetch_add(1, Ordering::SeqCst) + 1; + let token = SecretValue::new((self.token)(call)); + Box::pin(async move { + Ok(ResolvedCredential::AccessToken { + token, + expires_on: None, + }) + }) + } + } + + fn numbered_token(call: usize) -> String { + format!("callback-{call}") + } + + fn azure_request( + provider: &Arc, + api_base: Option<&str>, + api_key: Option<&str>, + extra_headers: Value, + optional_params: Value, + ) -> LiteLLMOcrRequest { + let wire = serde_json::from_value(json!({ + "model": "azure_ai/mistral-ocr-latest", + "document": {"type":"document_url","document_url":"data:application/pdf;base64,YWJj"}, + "api_key": api_key, + "api_base": api_base, + "custom_llm_provider": null, + "extra_headers": extra_headers, + "optional_params": optional_params, + "timeout_seconds": 2.0 + })) + .unwrap(); + LiteLLMOcrRequest { + azure_ad_token_provider: Some(TokenProviderHandle::new(provider.clone())), + ..decode_request(wire).unwrap() + } + } + + fn ocr_page() -> MockResponse { + MockResponse::json(json!({"pages":[{"index":0,"markdown":"hello"}]})) + } + + #[tokio::test] + async fn token_provider_result_is_the_bearer_and_is_acquired_for_each_request() { + let provider = CountingToken::new(numbered_token); + let (base, seen, server) = mock_server(vec![ocr_page(), ocr_page()]).await; + + for _ in 0..2 { + perform_ocr(azure_request( + &provider, + Some(&base), + None, + Value::Null, + json!({}), + )) + .await + .unwrap(); + } + server.await.unwrap(); + + assert_eq!(provider.calls(), 2); + let requests = seen.lock().unwrap(); + assert_eq!( + requests + .iter() + .map(|request| header(request, "authorization")) + .collect::>(), + [Some("Bearer callback-1"), Some("Bearer callback-2")] + ); + } + + #[rstest] + #[case::api_key_skips_provider(Some("resource-key"), Value::Null, json!({}), "Bearer resource-key", 0)] + #[case::provider_beats_static_token( + None, + Value::Null, + json!({"azure_ad_token":"static-token"}), + "Bearer callback-1", + 1 + )] + #[case::header_wins_on_the_wire_but_provider_still_runs( + None, + json!({"Authorization":"Bearer override"}), + json!({}), + "Bearer override", + 1 + )] + #[tokio::test] + async fn credential_precedence( + #[case] api_key: Option<&str>, + #[case] extra_headers: Value, + #[case] optional_params: Value, + #[case] expected_authorization: &str, + #[case] expected_calls: usize, + ) { + let provider = CountingToken::new(numbered_token); + let (base, seen, server) = mock_server(vec![ocr_page()]).await; + + perform_ocr(azure_request( + &provider, + Some(&base), + api_key, + extra_headers, + optional_params, + )) + .await + .unwrap(); + server.await.unwrap(); + + assert_eq!(provider.calls(), expected_calls); + let requests = seen.lock().unwrap(); + assert_eq!(requests.len(), 1); + assert_eq!( + header(&requests[0], "authorization"), + Some(expected_authorization) + ); + } + + #[rstest] + #[case::missing_api_base( + false, + json!({}), + numbered_token, + |error: &Error| matches!(error, Error::Auth(litellm_auth::Error::MissingApiBase { + provider: "Azure AI", + environment_variable: "AZURE_AI_API_BASE", + })), + 0 + )] + #[case::unsupported_oidc_reference( + true, + json!({"azure_ad_token":"oidc/assertion","client_id":"client","tenant_id":"tenant"}), + numbered_token, + |error: &Error| matches!(error, Error::Auth(litellm_auth::Error::UnsupportedOidcReference)), + 0 + )] + #[case::empty_provider_token_ignores_static_token( + true, + json!({"azure_ad_token":"static-token"}), + |_| String::new(), + |error: &Error| matches!(error, Error::MissingAzureAiCredentials), + 1 + )] + #[tokio::test] + async fn credential_failures_send_no_provider_request( + #[case] with_api_base: bool, + #[case] optional_params: Value, + #[case] token: fn(usize) -> String, + #[case] expected: fn(&Error) -> bool, + #[case] expected_calls: usize, + ) { + let provider = CountingToken::new(token); + let (base, seen, server) = mock_server(vec![ocr_page()]).await; + + let error = perform_ocr(azure_request( + &provider, + with_api_base.then_some(base.as_str()), + None, + Value::Null, + optional_params, + )) + .await + .unwrap_err(); + server.abort(); + + assert!(expected(&error), "unexpected error: {error:?}"); + assert_eq!(provider.calls(), expected_calls); + assert!(seen.lock().unwrap().is_empty()); + } + } +} + +#[cfg(test)] +mod azure_document_intelligence_tests { + use litellm_host::event::{CallEvent, MachineEvent}; + use litellm_llms::base_llm::ocr::{error::Error, settings::OcrSettings}; + use rstest::rstest; + use serde_json::{Value, json}; + + use crate::ocr::route::LocalOcrHost; + use crate::ocr::{ + test_support::{ + MockResponse, mock_server, ocr_client, perform_ocr, perform_ocr_with, wire_request, + }, + wire::{OcrWireRequest, decode_request}, + }; + + fn query_value(url: &str, key: &str) -> Option { + url::Url::parse(url) + .unwrap() + .query_pairs() + .find_map(|(name, value)| (name == key).then(|| value.into_owned())) + } + + #[tokio::test] + async fn facade_maps_pages_features_and_url_document() { + let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ + "status":"succeeded", + "analyzeResult":{"pages":[]} + }))]) + .await; + let mut request = wire_request( + "azure_ai/doc-intelligence/prebuilt-read", + &base, + json!({"pages":[2,0,0,1],"features":["keyValuePairs","languages"]}), + ); + request.document = serde_json::from_value::< + litellm_llms::base_llm::ocr::transformation::OcrDocument, + >(json!({ + "type":"document_url", + "document_url":"https://example.com/document.pdf" + })) + .unwrap() + .into(); + + perform_ocr(request).await.unwrap(); + server.await.unwrap(); + let request = &seen.lock().unwrap()[0]; + let target = request.split_whitespace().nth(1).unwrap(); + let url = format!("{base}{target}"); + assert_eq!(query_value(&url, "pages").as_deref(), Some("1,2,3")); + assert_eq!( + query_value(&url, "features").as_deref(), + Some("keyValuePairs,languages") + ); + let body: Value = serde_json::from_str(request.split_once("\r\n\r\n").unwrap().1).unwrap(); + assert_eq!( + body, + json!({"urlSource":"https://example.com/document.pdf"}) + ); + } + + #[rstest] + #[case(json!({"pages":[true]}), Error::Pages("expected only integers or only strings".into()))] + #[case(json!({"pages":[1,"2"]}), Error::Pages("expected only integers or only strings".into()))] + #[case(json!({"pages":[-1]}), Error::Pages("negative page index".into()))] + #[case(json!({"pages":"1&&features=bad"}), Error::Pages("invalid native page range".into()))] + #[case(json!({"features":"languages&pages=1"}), Error::Features)] + #[case(json!({"req_format":"azure"}), Error::RequestFormat)] + #[tokio::test] + async fn rejects_invalid_pages_features_and_format( + #[case] options: Value, + #[case] expected: Error, + ) { + let (base, seen, server) = mock_server(vec![MockResponse::json(json!({}))]).await; + let result = decode_request(OcrWireRequest { + model: "azure_ai/doc-intelligence/prebuilt-read".into(), + document: json!({"type":"document_url","document_url":"https://example.com/a.pdf"}), + api_key: Some(litellm_auth::SecretValue::new("key")), + api_base: Some(base), + custom_llm_provider: None, + extra_headers: None, + optional_params: options.as_object().unwrap().clone(), + input_sources: Default::default(), + timeout_seconds: Some(2.0), + }); + let result = match result { + Ok(request) => perform_ocr(request).await, + Err(error) => Err(error), + }; + server.abort(); + let _ = server.await; + assert!( + seen.lock().unwrap().is_empty(), + "sent invalid options: {options}" + ); + let error = result.unwrap_err(); + assert_eq!( + std::mem::discriminant(&error), + std::mem::discriminant(&expected) + ); + assert_eq!(error.http_status_code(), Some(400)); + assert_eq!(error.to_string(), expected.to_string()); + } + + #[rstest] + #[case(json!({}))] + #[case(json!({"req_format":"litellm"}))] + #[tokio::test] + async fn missing_native_fields_keep_page_text_without_retaining_raw_response( + #[case] options: Value, + ) { + let operation = json!({ + "status":"succeeded", + "analyzeResult":{"pages":[{"pageNumber":1,"lines":[{"content":"hello"}]}]} + }); + let (base, seen, server) = mock_server(vec![MockResponse::json(operation)]).await; + let response = perform_ocr(wire_request( + "azure_ai/doc-intelligence/prebuilt-read", + &base, + options, + )) + .await + .unwrap(); + server.await.unwrap(); + + assert_eq!(response.pages.len(), 1); + assert_eq!(response.pages[0].index, 0); + assert_eq!(response.pages[0].markdown, "hello"); + assert_eq!(response.provider_native_response, None); + let serialized = response.into_json(); + assert_eq!(serialized.get("content"), Some(&Value::Null)); + assert_eq!(serialized.get("tables"), Some(&Value::Null)); + assert_eq!(serialized.get("keyValuePairs"), Some(&Value::Null)); + let requests = seen.lock().unwrap(); + assert_eq!(requests.len(), 1); + let target = requests[0].split_whitespace().nth(1).unwrap(); + let url = format!("{base}{target}"); + for field in ["pages", "features", "req_format"] { + assert_eq!(query_value(&url, field), None); + } + let body: Value = + serde_json::from_str(requests[0].split_once("\r\n\r\n").unwrap().1).unwrap(); + assert_eq!(body, json!({"base64Source":"YWJj"})); + } + + #[tokio::test] + async fn inline_document_decodes_to_base64_source() { + let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ + "status":"succeeded" + }))]) + .await; + let request = wire_request("azure_ai/doc-intelligence/prebuilt-read", &base, json!({})); + + perform_ocr(request).await.unwrap(); + server.await.unwrap(); + let request = &seen.lock().unwrap()[0]; + let body: Value = serde_json::from_str(request.split_once("\r\n\r\n").unwrap().1).unwrap(); + assert_eq!(body, json!({"base64Source":"YWJj"})); + } + + #[tokio::test] + async fn immediate_response_normalizes_pages_and_preserves_native() { + let operation = json!({ + "status":"succeeded", + "operationExtension":42, + "analyzeResult":{ + "content":"A\n\nB", + "tables":[{"cells":[]}], + "keyValuePairs":[{"key":{"content":"A"}}], + "pages":[{ + "pageNumber":"2", + "width":"8.5", + "height":11, + "unit":"inch", + "lines":[{"content":"A"},{"content":null},{"content":"B"}] + }] + } + }); + let (base, _, server) = mock_server(vec![MockResponse::json(operation.clone())]).await; + let result = perform_ocr(wire_request( + "azure_ai/doc-intelligence/prebuilt-read", + &base, + json!({"req_format":"native"}), + )) + .await + .unwrap(); + server.await.unwrap(); + + assert_eq!(result.pages[0].index, 1); + assert_eq!(result.pages[0].markdown, "A\n\nB"); + assert_eq!( + serde_json::to_value(&result.pages[0].dimensions).unwrap(), + json!({"width":816,"height":1056,"dpi":96}) + ); + assert_eq!(result.usage_info.as_ref().unwrap().pages_processed, Some(1)); + let serialized = result.clone().into_json(); + assert_eq!(serialized["content"], "A\n\nB"); + assert_eq!(serialized["tables"], json!([{"cells":[]}])); + assert_eq!( + serialized["keyValuePairs"], + json!([{"key":{"content":"A"}}]) + ); + assert!(serialized.get("key_value_pairs").is_none()); + assert_eq!( + result.provider_native_response.map(Value::Object), + Some(operation) + ); + } + + #[tokio::test] + async fn client_settings_choose_the_api_version_and_the_inch_to_pixel_dpi() { + let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ + "status":"succeeded", + "analyzeResult":{"pages":[{"pageNumber":1,"width":8.5,"height":11,"unit":"inch"}]} + }))]) + .await; + let client = ocr_client().with_settings(OcrSettings { + document_intelligence_api_version: "2099-01-01".into(), + document_intelligence_dpi: 72, + ..OcrSettings::default() + }); + + let result = crate::ocr::client::perform( + &client, + wire_request("azure_ai/doc-intelligence/prebuilt-read", &base, json!({})), + ) + .await + .unwrap(); + server.await.unwrap(); + + let target = seen.lock().unwrap()[0] + .split_whitespace() + .nth(1) + .unwrap() + .to_string(); + assert_eq!( + query_value(&format!("{base}{target}"), "api-version").as_deref(), + Some("2099-01-01") + ); + assert_eq!( + serde_json::to_value(&result.pages[0].dimensions).unwrap(), + json!({"width":612,"height":792,"dpi":72}) + ); + } + + #[tokio::test] + async fn accepted_response_polls_to_success_with_only_credentials() { + let operation = json!({"status":"succeeded","analyzeResult":{"pages":[]}}); + let (base, seen, server) = mock_server(vec![ + MockResponse { + status: 202, + headers: vec![("Operation-Location", "{base}/operation".into())], + body: json!({}), + }, + MockResponse { + status: 200, + headers: vec![("Retry-After", "0".into())], + body: json!({"status":"running"}), + }, + MockResponse::json(operation.clone()), + ]) + .await; + let mut request = wire_request( + "azure_ai/doc-intelligence/prebuilt-read", + &base, + json!({"req_format":"native"}), + ); + request + .transport + .extra_headers + .push(("X-Trace".into(), "initial-only".into())); + + let result = perform_ocr(request).await.unwrap(); + server.await.unwrap(); + assert_eq!( + result.provider_native_response.map(Value::Object), + Some(operation) + ); + let requests = seen.lock().unwrap(); + assert_eq!(requests.len(), 3); + assert!(requests[0].to_ascii_lowercase().contains("x-trace:")); + for poll in &requests[1..] { + assert!(!poll.to_ascii_lowercase().contains("x-trace:")); + assert!( + poll.to_ascii_lowercase() + .contains("ocp-apim-subscription-key: test-key") + ); + } + } + + #[tokio::test] + async fn accepted_response_emits_response_received_before_polling() { + let (base, seen, server) = mock_server(vec![ + MockResponse { + status: 202, + headers: vec![("Operation-Location", "{base}/operation".into())], + body: json!({"submitted": true}), + }, + MockResponse::json(json!({"status":"succeeded"})), + ]) + .await; + let request_count = seen.clone(); + let host = LocalOcrHost::new(wire_request( + "azure_ai/doc-intelligence/prebuilt-read", + &base, + json!({}), + )) + .with_observer(move |event| { + let CallEvent::Machine(MachineEvent::ResponseReceived { raw }) = event else { + return; + }; + match request_count.lock().unwrap().len() { + 1 => assert_eq!(raw.body, r#"{"submitted":true}"#), + 2 => assert!(raw.body.contains("succeeded")), + count => panic!("unexpected callback after {count} requests"), + } + }); + + perform_ocr_with(host).await.unwrap(); + server.await.unwrap(); + assert_eq!(seen.lock().unwrap().len(), 2); + } + + #[tokio::test] + async fn polling_forwards_bearer_credentials() { + let (base, seen, server) = mock_server(vec![ + MockResponse { + status: 202, + headers: vec![("Operation-Location", "{base}/operation".into())], + body: json!({}), + }, + MockResponse::json(json!({"status":"succeeded"})), + ]) + .await; + let mut request = wire_request("azure_ai/doc-intelligence/prebuilt-read", &base, json!({})); + request.credentials.api_key = None; + request.transport.extra_headers = vec![("Authorization".into(), "Bearer token".into())]; + + perform_ocr(request).await.unwrap(); + server.await.unwrap(); + let requests = seen.lock().unwrap(); + assert!( + requests[1] + .to_ascii_lowercase() + .contains("authorization: bearer token") + ); + } + + #[tokio::test] + async fn polling_does_not_follow_redirects() { + let (base, seen, server) = mock_server(vec![ + MockResponse { + status: 202, + headers: vec![("Operation-Location", "{base}/operation".into())], + body: json!({}), + }, + MockResponse { + status: 302, + headers: vec![("Location", "{base}/redirected".into())], + body: json!({}), + }, + MockResponse::json(json!({"status":"succeeded"})), + ]) + .await; + + let error = perform_ocr(wire_request( + "azure_ai/doc-intelligence/prebuilt-read", + &base, + json!({}), + )) + .await + .unwrap_err(); + + assert!(error.to_string().contains("status 302"), "{error}"); + assert_eq!(seen.lock().unwrap().len(), 2); + server.abort(); + } + + #[tokio::test] + async fn polling_rejects_terminal_failure() { + let (base, _, server) = mock_server(vec![ + MockResponse { + status: 202, + headers: vec![("Operation-Location", "{base}/operation".into())], + body: json!({}), + }, + MockResponse::json(json!({"status":"failed"})), + ]) + .await; + + let error = perform_ocr(wire_request( + "azure_ai/doc-intelligence/prebuilt-read", + &base, + json!({}), + )) + .await + .unwrap_err(); + server.await.unwrap(); + assert!(error.to_string().contains("status failed")); + } + + #[tokio::test] + async fn malformed_provider_pages_report_response_paths() { + for (analysis, path) in [ + (json!({"pages":null}), "pages"), + (json!({"pages":[null]}), "pages[0]"), + (json!({"pages":[{"lines":null}]}), "lines"), + (json!({"pages":[{"width":"bad"}]}), "width"), + ] { + let (base, _, server) = mock_server(vec![MockResponse::json(json!({ + "status":"succeeded", + "analyzeResult":analysis + }))]) + .await; + let error = perform_ocr(wire_request( + "azure_ai/doc-intelligence/prebuilt-read", + &base, + json!({}), + )) + .await + .unwrap_err(); + server.await.unwrap(); + assert!(error.to_string().contains(path), "{error}"); + } + } + + #[tokio::test] + async fn rejects_missing_invalid_and_cross_origin_operation_locations() { + for headers in [ + Vec::new(), + vec![("Operation-Location", "/relative".into())], + vec![("Operation-Location", "http://example.com/operation".into())], + vec![( + "Operation-Location", + "http://user:password@127.0.0.1/operation".into(), + )], + ] { + let (base, _, server) = mock_server(vec![MockResponse { + status: 202, + headers, + body: json!({}), + }]) + .await; + let error = perform_ocr(wire_request( + "azure_ai/doc-intelligence/prebuilt-read", + &base, + json!({}), + )) + .await + .unwrap_err(); + server.await.unwrap(); + assert!(error.to_string().contains("operation-location")); + } + } + + #[tokio::test] + async fn polling_deadline_bounds_retry_delay() { + let (base, _, server) = mock_server(vec![ + MockResponse { + status: 202, + headers: vec![("Operation-Location", "{base}/operation".into())], + body: json!({}), + }, + MockResponse { + status: 200, + headers: vec![("Retry-After", "9999".into())], + body: json!({"status":"notStarted"}), + }, + ]) + .await; + let request = wire_request("azure_ai/doc-intelligence/prebuilt-read", &base, json!({})); + let client = ocr_client().with_settings(OcrSettings { + poll_timeout: std::time::Duration::from_millis(100), + ..OcrSettings::default() + }); + + let error = tokio::time::timeout( + std::time::Duration::from_secs(1), + crate::ocr::client::perform(&client, request), + ) + .await + .unwrap() + .unwrap_err(); + server.await.unwrap(); + assert!(error.to_string().contains("timed out")); + } + + #[tokio::test] + async fn model_id_is_encoded_and_dot_segments_are_rejected() { + let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ + "status":"succeeded" + }))]) + .await; + perform_ocr(wire_request( + "azure_ai/doc-intelligence/a ?#é", + &base, + json!({}), + )) + .await + .unwrap(); + server.await.unwrap(); + assert!(seen.lock().unwrap()[0].contains("a%20%3F%23%C3%A9:analyze")); + + for model in [ + "azure_ai/doc-intelligence/.", + "azure_ai/doc-intelligence/..", + ] { + let error = perform_ocr(wire_request(model, "http://127.0.0.1:1", json!({}))) + .await + .unwrap_err(); + assert!(error.to_string().contains("dot segment")); + } + } + + mod transformation { + use std::sync::{Arc, Mutex}; + + use litellm_host::event::{CallEvent, MachineEvent}; + use litellm_llms::base_llm::ocr::transformation::OcrDocument; + use serde_json::{Value, json}; + + use super::*; + use crate::ocr::{ + route::LocalOcrHost, + test_support::{ + MockResponse, mock_server, perform_ocr, perform_ocr_with, wire_request, + }, + }; + + #[tokio::test] + async fn facade_maps_pages_features_and_url_document() { + let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ + "status":"succeeded", + "analyzeResult":{"pages":[]} + }))]) + .await; + let mut request = wire_request( + "azure_ai/doc-intelligence/prebuilt-read", + &base, + json!({"pages":[2,0,0,1],"features":["keyValuePairs","languages"], "future_option": {"nested":null}, "extra_body":{"provider_option":false}}), + ); + request.document = serde_json::from_value::(json!({ + "type":"document_url", + "document_url":"https://example.com/document.pdf" + })) + .unwrap() + .into(); + + perform_ocr(request).await.unwrap(); + server.await.unwrap(); + let request = &seen.lock().unwrap()[0]; + let target = request.split_whitespace().nth(1).unwrap(); + let url = format!("{base}{target}"); + assert_eq!(query_value(&url, "pages").as_deref(), Some("1,2,3")); + assert_eq!( + query_value(&url, "features").as_deref(), + Some("keyValuePairs,languages") + ); + let body: Value = + serde_json::from_str(request.split_once("\r\n\r\n").unwrap().1).unwrap(); + assert_eq!( + body, + json!({"urlSource":"https://example.com/document.pdf", "future_option":{"nested":null}, "provider_option":false}) + ); + } + + #[tokio::test] + async fn rejects_invalid_pages_features_and_format() { + for options in [ + json!({"pages":[true]}), + json!({"pages":[1,"2"]}), + json!({"pages":[-1]}), + json!({"pages":"1&&features=bad"}), + json!({"features":"languages&pages=1"}), + json!({"req_format":"azure"}), + ] { + let request = wire_request( + "azure_ai/doc-intelligence/prebuilt-read", + "http://127.0.0.1:1", + options.clone(), + ); + let rejected = perform_ocr(request).await.is_err(); + assert!(rejected, "accepted {options}"); + } + } + + #[tokio::test] + async fn immediate_response_normalizes_pages_and_preserves_native() { + let operation = json!({ + "status":"succeeded", + "operationExtension":42, + "analyzeResult":{ + "content":"A\n\nB", + "tables":[{"cells":[]}], + "keyValuePairs":[{"key":{"content":"A"}}], + "pages":[{ + "pageNumber":"2", + "width":"8.5", + "height":11, + "unit":"inch", + "lines":[{"content":"A"},{"content":null},{"content":"B"}] + }] + } + }); + let (base, _, server) = mock_server(vec![MockResponse::json(operation.clone())]).await; + let result = perform_ocr(wire_request( + "azure_ai/doc-intelligence/prebuilt-read", + &base, + json!({"req_format":"native"}), + )) + .await + .unwrap(); + server.await.unwrap(); + + assert_eq!(result.pages[0].index, 1); + assert_eq!(result.pages[0].markdown, "A\n\nB"); + assert_eq!( + serde_json::to_value(&result.pages[0].dimensions).unwrap(), + json!({"width":816,"height":1056,"dpi":96}) + ); + assert_eq!(result.usage_info.as_ref().unwrap().pages_processed, Some(1)); + let serialized = result.clone().into_json(); + assert_eq!(serialized["content"], "A\n\nB"); + assert_eq!(serialized["tables"], json!([{"cells":[]}])); + assert_eq!( + serialized["keyValuePairs"], + json!([{"key":{"content":"A"}}]) + ); + assert!(serialized.get("key_value_pairs").is_none()); + assert_eq!( + result.provider_native_response.as_ref(), + operation.as_object() + ); + } + + #[tokio::test] + async fn accepted_response_polls_to_success_with_only_credentials() { + let operation = json!({"status":"succeeded","analyzeResult":{"pages":[]}}); + let (base, seen, server) = mock_server(vec![ + MockResponse { + status: 202, + headers: vec![("Operation-Location", "{base}/operation".into())], + body: json!({}), + }, + MockResponse { + status: 200, + headers: vec![("Retry-After", "0".into())], + body: json!({"status":"running"}), + }, + MockResponse::json(operation.clone()), + ]) + .await; + let mut request = wire_request( + "azure_ai/doc-intelligence/prebuilt-read", + &base, + json!({"req_format":"native"}), + ); + request + .transport + .extra_headers + .push(("X-Trace".into(), "initial-only".into())); + + let result = perform_ocr(request).await.unwrap(); + server.await.unwrap(); + assert_eq!( + result.provider_native_response.as_ref(), + operation.as_object() + ); + let requests = seen.lock().unwrap(); + assert_eq!(requests.len(), 3); + assert!(requests[0].to_ascii_lowercase().contains("x-trace:")); + for poll in &requests[1..] { + assert!(!poll.to_ascii_lowercase().contains("x-trace:")); + assert!( + poll.to_ascii_lowercase() + .contains("ocp-apim-subscription-key: test-key") + ); + } + } + + #[tokio::test] + async fn accepted_response_emits_response_received_for_submission_and_completed_poll() { + let (base, seen, server) = mock_server(vec![ + MockResponse { + status: 202, + headers: vec![("Operation-Location", "{base}/operation".into())], + body: json!({"submitted": true}), + }, + MockResponse::json(json!({"status":"succeeded"})), + ]) + .await; + let responses_received = Arc::new(Mutex::new(Vec::new())); + let request_count = seen.clone(); + let observed = responses_received.clone(); + let host = LocalOcrHost::new(wire_request( + "azure_ai/doc-intelligence/prebuilt-read", + &base, + json!({}), + )) + .with_observer(move |event| { + if let CallEvent::Machine(MachineEvent::ResponseReceived { raw }) = event { + observed + .lock() + .unwrap() + .push((request_count.lock().unwrap().len(), raw.body.clone())); + } + }); + + perform_ocr_with(host).await.unwrap(); + server.await.unwrap(); + assert_eq!(seen.lock().unwrap().len(), 2); + assert_eq!( + *responses_received.lock().unwrap(), + [ + (1, r#"{"submitted":true}"#.to_string()), + (2, r#"{"status":"succeeded"}"#.to_string()), + ] + ); + } + } +} + +#[cfg(test)] +mod cohere_tests { + mod transformation { + use litellm_llms::{ + base_llm::ocr::{ + error::Error, + transformation::{BaseOcrConfig, OcrDocument, OcrResponseFormat}, + }, + cohere::ocr::transformation::*, + }; + use rstest::rstest; + use serde_json::{Value, json}; + + #[tokio::test] + async fn composed_body_preserves_native_document_fields_and_untyped_overrides() { + let request = crate::ocr::test_support::wire_request( + "cohere/parse", + "https://example.com", + json!({ + "output_format":"markdown", "timeout":30, + "extra_body":{ + "output_format": {"future":true}, + "document":{"type":"image_url","image_url":"https://example.com/a.png", + "provider_options":{"nested":[false,0,null]}} + } + }), + ); + let request = request.with_document( + serde_json::from_value(json!({ + "type":"image_url","image_url":"https://example.com/original.png" + })) + .unwrap(), + ); + let request = crate::ocr::prepare::prepare_request_for_test(request); + let http = CohereParseConfig + .prepare_request( + &request, + &crate::ocr::test_support::ocr_client(), + &crate::ocr::test_support::NoHooks, + ) + .await + .unwrap(); + let body: Value = serde_json::from_slice(http.body()).unwrap(); + assert_eq!( + body, + json!({ + "model":"parse", "output_format":{"future":true}, + "document":{"type":"image_url","image_url":"https://example.com/a.png", + "provider_options":{"nested":[false,0,null]}} + }) + ); + } + + #[tokio::test] + async fn explicit_null_options_use_defaults_before_http() { + let request = crate::ocr::test_support::wire_request( + "cohere/parse", + "https://example.com", + json!({"output_format":null,"req_format":null}), + ); + let request = request.with_document( + serde_json::from_value( + json!({"type":"image_url","image_url":"https://example.com/a.png"}), + ) + .unwrap(), + ); + assert_eq!( + request.response_format().unwrap(), + OcrResponseFormat::Litellm + ); + let request = crate::ocr::prepare::prepare_request_for_test(request); + let http = CohereParseConfig + .prepare_request( + &request, + &crate::ocr::test_support::ocr_client(), + &crate::ocr::test_support::NoHooks, + ) + .await + .unwrap(); + let body: Value = serde_json::from_slice(http.body()).unwrap(); + assert_eq!(body["output_format"], "markdown"); + assert!(body.get("req_format").is_none()); + } + + #[rstest] + #[case::cohere("cohere/parse-v5.0", "POST /v2/parse ")] + #[case::azure_ai("azure_ai/Cohere-parse-v5.0", "POST /providers/cohere/v2/parse ")] + #[tokio::test] + async fn route_sends_image_to_its_parse_endpoint_with_the_bearer_key( + #[case] model: &str, + #[case] request_line: &str, + ) { + use crate::ocr::test_support::{MockResponse, header, mock_server, perform_ocr}; + + let (base, seen, server) = + mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await; + let request = crate::ocr::test_support::wire_request(model, &base, json!({})) + .with_document( + serde_json::from_value::( + json!({"type":"image_url","image_url":"data:image/png;base64,YWJj"}), + ) + .unwrap() + .into(), + ); + + perform_ocr(request).await.unwrap(); + server.await.unwrap(); + + let requests = seen.lock().unwrap(); + assert_eq!(requests.len(), 1); + assert!(requests[0].starts_with(request_line), "{}", requests[0]); + assert_eq!( + header(&requests[0], "authorization"), + Some("Bearer test-key") + ); + } + + #[rstest] + #[tokio::test] + async fn route_rejects_non_image_document_without_a_request( + #[values("cohere/parse-v5.0", "azure_ai/Cohere-parse-v5.0")] model: &str, + ) { + use crate::ocr::test_support::{MockResponse, mock_server, perform_ocr}; + + let (base, seen, server) = + mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await; + + let error = perform_ocr(crate::ocr::test_support::wire_request( + model, + &base, + json!({}), + )) + .await + .unwrap_err(); + server.abort(); + + assert!(matches!(error, Error::CohereImageOnly), "{error:?}"); + assert!(seen.lock().unwrap().is_empty()); + } + } +} + +#[cfg(test)] +mod deepseek_tests { + use litellm_llms::{ + base_llm::ocr::transformation::{BaseOcrConfig, OcrDocument}, + vertex_ai::ocr::deepseek_transformation::{ + DeepSeekOcrParams, DeepSeekOcrResponse, VertexAIDeepSeekOCRConfig, + normalize_response as transform_ocr_response, + }, + }; + use rstest::rstest; + use serde_json::{Value, json}; + + fn document() -> OcrDocument { + serde_json::from_value(json!({"type":"image_url","image_url":"gs://bucket/a.png"})).unwrap() + } + + #[rstest] + #[case("stream", json!(true))] + #[case("temperature", json!(0.1))] + #[case("max_tokens", json!(1024))] + #[case("top_p", json!(0.9))] + #[case("n", json!(2))] + #[case("stop", json!("done"))] + #[case("stop", json!(["done", "stop"]))] + fn request_mapping_matches_python(#[case] name: &str, #[case] value: Value) { + let params: DeepSeekOcrParams = + serde_json::from_value(json!({name: value.clone(), "ignored": true})).unwrap(); + let result = serde_json::to_value( + VertexAIDeepSeekOCRConfig + .transform_ocr_request("deepseek-ai/deepseek-ocr-maas", document(), ¶ms, &[]) + .unwrap(), + ) + .unwrap(); + assert_eq!(result["model"], "deepseek-ai/deepseek-ocr-maas"); + assert_eq!( + result["messages"][0]["content"][0], + json!({"type":"image_url","image_url":"gs://bucket/a.png"}) + ); + assert_eq!(result[name], value); + assert!(result.get("ignored").is_none()); + } + + #[rstest] + #[case(json!({"type":"image_url","image_url":"data:image/png;base64,AA=="}))] + #[case(json!({"type":"document_url","document_url":"data:application/pdf;base64,AA=="}))] + fn request_maps_both_document_types_to_image_content(#[case] document: Value) { + let source = document + .get("image_url") + .or_else(|| document.get("document_url")) + .unwrap() + .clone(); + let request = VertexAIDeepSeekOCRConfig + .transform_ocr_request( + "deepseek-ai/deepseek-ocr-maas", + serde_json::from_value(document).unwrap(), + &DeepSeekOcrParams::default(), + &[], + ) + .unwrap(); + let result = serde_json::to_value(request).unwrap(); + assert_eq!( + result["messages"][0]["content"][0], + json!({"type":"image_url","image_url":source}) + ); + } + + #[rstest] + #[case(json!("# hello"), "# hello")] + #[case(json!("{broken"), "{broken")] + #[case(json!(" {\"pages\":[]} "), " {\"pages\":[]} ")] + #[case(json!({"pages":[]}), "")] + #[case(json!("[]"), "[]")] + #[case(json!("{\"pages\":[{\"markdown\":\"json text\"}]}"), "json text")] + #[case(json!({"pages":[{"markdown":"object"}]}), "object")] + fn response_codec_handles_text_json_and_objects( + #[case] content: Value, + #[case] expected: &str, + ) { + let structured = content + .as_object() + .is_some_and(|object| object.contains_key("pages")) + || content + .as_str() + .is_some_and(|text| text.contains("\"pages\"")); + let response: DeepSeekOcrResponse = serde_json::from_value( + json!({"choices":[{"message":{"content":content}}],"usage":{"prompt_tokens":1}}), + ) + .unwrap(); + let result = transform_ocr_response("model", response) + .unwrap() + .into_json(); + assert_eq!(result["pages"][0]["markdown"], expected); + assert_eq!(result["pages"][0]["index"], 0); + if structured { + assert!(result["usage_info"].is_null()); + } else { + assert_eq!(result["usage_info"]["prompt_tokens"], 1); + } + } + + #[test] + fn structured_result_maps_pages_usage_model_and_annotation() { + let response: DeepSeekOcrResponse = serde_json::from_value(json!({ + "choices":[{"message":{"content":{ + "pages":[{"index":2,"markdown":"page","images":[{"id":"one"}],"dimensions":{"width":10}}], + "model":"provider-model", + "usage_info":{"pages_processed":1}, + "document_annotation":{"language":"en"}, + "future":"kept" + }}}] + })) + .unwrap(); + let result = transform_ocr_response("requested", response) + .unwrap() + .into_json(); + assert_eq!(result["pages"][0]["index"], 2); + assert_eq!(result["pages"][0]["images"][0]["id"], "one"); + assert_eq!(result["model"], "provider-model"); + assert_eq!(result["usage_info"]["pages_processed"], 1); + assert_eq!(result["document_annotation"]["language"], "en"); + assert_eq!(result["future"], "kept"); + } + + #[test] + fn response_codec_rejects_missing_empty_and_malformed_content() { + for value in [ + json!({"choices":[{"message":{"content":{}}}]}), + json!({"choices":[]}), + json!({"choices":[{"message":{"content":""}}]}), + json!({"choices":[{"message":{"content":"{\"pages\":[{\"markdown\":42}]}"}}]}), + json!({"choices":[{"message":{"content":{"pages":[{"markdown":42}]}}}]}), + ] { + let result = serde_json::from_value::(value) + .map_err(|_| ()) + .and_then(|response| transform_ocr_response("model", response).map_err(|_| ())); + assert!(result.is_err()); + } + } +} + +#[cfg(test)] +mod reducto_tests { + use litellm_host::event::{CallEvent, MachineEvent, WireRequest}; + use litellm_llms::base_llm::ocr::{error::Error, transformation::OcrDocument}; + use rstest::rstest; + use serde_json::{Value, json}; + + use crate::ocr::route::LocalOcrHost; + use crate::ocr::test_support::{ + MockResponse, mock_server, perform_ocr, perform_ocr_with, wire_request, + }; + + fn request_body(request: &str) -> Value { + serde_json::from_str(request.split_once("\r\n\r\n").unwrap().1).unwrap() + } + + #[rstest] + #[case( + "reducto/parse-v3", + json!({ + "formatting":{"table_output_format":"html"}, + "retrieval":{"chunk_mode":"section"}, + "settings":{"ocr_system":"standard"}, + "future_ocr_option":true, + "extra_body":{"provider_option":"value"} + }), + "reducto://already.pdf", + json!({ + "input":"reducto://already.pdf", + "formatting":{"table_output_format":"html"}, + "retrieval":{"chunk_mode":"section"}, + "settings":{"ocr_system":"standard"}, + "future_ocr_option":true, + "provider_option":"value" + }) + )] + #[case( + "reducto/parse-legacy", + json!({ + "enhance":{"agentic":[{"type":"table"}]}, + "future_ocr_option":true, + "extra_body":{"provider_option":"value"} + }), + "reducto://legacy.pdf", + json!({ + "document_url":"reducto://legacy.pdf", + "options":{"enhance":{"agentic":[{"type":"table"}]}}, + "future_ocr_option":true, + "provider_option":"value" + }) + )] + #[tokio::test] + async fn request_mapping_matches_python( + #[case] model: &str, + #[case] options: Value, + #[case] source: &str, + #[case] expected: Value, + ) { + let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ + "result":{"chunks":[]} + }))]) + .await; + let request = + crate::ocr::test_support::with_source(wire_request(model, &base, options), source); + + perform_ocr(request).await.unwrap(); + server.await.unwrap(); + let requests = seen.lock().unwrap(); + assert_eq!(requests.len(), 1); + assert!(requests[0].starts_with("POST /parse ")); + assert_eq!(request_body(&requests[0]), expected); + } + + #[rstest] + #[case("parse-v3")] + #[case("parse-legacy")] + #[tokio::test] + async fn data_uri_upload_preserves_multipart_headers( + #[case] model: &str, + #[values("application/pdf", "image/png")] mime_type: &str, + ) { + let (base, seen, server) = mock_server(vec![ + MockResponse::json(json!({"file_id":"reducto://uploaded.pdf"})), + MockResponse::json(json!({"result":{"chunks":[{"content":"hello"}]}})), + ]) + .await; + let document = if mime_type.starts_with("image/") { + json!({"type":"image_url","image_url":format!("data:{mime_type};base64,YWJj")}) + } else { + json!({"type":"document_url","document_url":format!("data:{mime_type};base64,YWJj")}) + }; + let mut request = crate::ocr::types::LiteLLMOcrRequest { + document: serde_json::from_value::(document) + .unwrap() + .into(), + ..wire_request(&format!("reducto/{model}"), &base, json!({})) + }; + request.transport.extra_headers = vec![ + ("Content-Type".into(), "application/json".into()), + ("X-Trace".into(), "upload-test".into()), + ]; + + let response = perform_ocr(request).await.unwrap(); + server.await.unwrap(); + assert_eq!(response.pages[0].markdown, "hello"); + let requests = seen.lock().unwrap(); + assert_eq!(requests.len(), 2); + assert!(requests[0].starts_with("POST /upload ")); + assert!( + requests[0] + .to_ascii_lowercase() + .contains("content-type: multipart/form-data; boundary=") + ); + assert!(requests[0].contains("x-trace: upload-test")); + let multipart = requests[0].split_once("\r\n\r\n").unwrap().1; + assert!(multipart.contains(&format!("Content-Type: {mime_type}\r\n"))); + assert!(multipart.contains("\r\n\r\nabc\r\n--")); + assert!(requests[1].starts_with("POST /parse ")); + let source_field = if model == "parse-legacy" { + "document_url" + } else { + "input" + }; + assert_eq!( + request_body(&requests[1]), + json!({source_field:"reducto://uploaded.pdf"}) + ); + for request in requests.iter() { + assert!( + request + .to_ascii_lowercase() + .contains("authorization: bearer test-key\r\n") + ); + } + } + + #[tokio::test] + async fn response_received_stays_after_reducto_upload_and_parse() { + let (base, seen, server) = mock_server(vec![ + MockResponse::json(json!({"file_id":"reducto://uploaded.pdf"})), + MockResponse::json(json!({"result":{"chunks":[]}})), + ]) + .await; + let request_count = seen.clone(); + let host = LocalOcrHost::new(wire_request("reducto/parse-v3", &base, json!({}))) + .with_observer(move |event| { + if let CallEvent::Machine(MachineEvent::ResponseReceived { raw }) = event { + assert_eq!(request_count.lock().unwrap().len(), 2); + assert_eq!(raw.body, r#"{"result":{"chunks":[]}}"#); + } + }); + + perform_ocr_with(host).await.unwrap(); + server.await.unwrap(); + assert_eq!(seen.lock().unwrap().len(), 2); + } + + #[rstest] + #[case(json!({"file_id":""}))] + #[case(json!({}))] + #[case(json!({"file_id":null}))] + #[tokio::test] + async fn invalid_upload_ids_stop_before_parse(#[case] response: Value) { + let (base, seen, server) = mock_server(vec![MockResponse::json(response)]).await; + let error = perform_ocr(wire_request("reducto/parse-v3", &base, json!({}))) + .await + .unwrap_err(); + server.await.unwrap(); + assert!(error.to_string().contains("file_id")); + assert_eq!(seen.lock().unwrap().len(), 1); + } + + #[tokio::test] + async fn upload_failure_stops_before_parse() { + let (base, seen, server) = mock_server(vec![MockResponse { + status: 503, + headers: vec![], + body: json!({"error":"unavailable"}), + }]) + .await; + assert!( + perform_ocr(wire_request("reducto/parse-v3", &base, json!({}))) + .await + .is_err() + ); + server.await.unwrap(); + assert_eq!(seen.lock().unwrap().len(), 1); + } + + #[rstest] + #[case("https://example.com/a.pdf", Error::ReductoSource)] + #[case("reducto://", Error::RequestField { path: "document file id".into() })] + #[case("data:application/pdf;base64", Error::InvalidDataUri)] + #[case("data:application/pdf;base64,INVALID!", Error::InvalidDataUri)] + #[tokio::test] + async fn rejects_invalid_document_sources_before_network( + #[case] source: &str, + #[case] expected: Error, + ) { + let (base, seen, server) = mock_server(vec![MockResponse::json(json!({}))]).await; + let request = crate::ocr::test_support::with_source( + wire_request("reducto/parse-v3", &base, json!({})), + source, + ); + let result = perform_ocr(request).await; + server.abort(); + let _ = server.await; + assert!( + seen.lock().unwrap().is_empty(), + "sent invalid source: {source}" + ); + let error = result.unwrap_err(); + assert_eq!( + std::mem::discriminant(&error), + std::mem::discriminant(&expected) + ); + assert_eq!(error.http_status_code(), Some(400)); + assert_eq!(error.to_string(), expected.to_string()); + } + + #[test] + fn response_normalization_groups_blocks_and_distinguishes_null_result() { + use litellm_llms::reducto::ocr::transformation::{ + ReductoResponse, normalize_response as transform_ocr_response, + }; + + let raw = json!({"usage":{"num_pages":"2","credits":"3"},"result":{"type":"full","chunks":[ + {"blocks":[{ + "type":"Table", + "content":"B", + "bbox":{"left":0.1,"top":0.2,"width":0.8,"height":0.3,"page":2,"original_page":4}, + "confidence":"high", + "granular_confidence":{"parse_confidence":0.95,"extract_confidence":null}, + "image_url":null + }]}, + {"blocks":[{"content":"A","bbox":{"page":1},"type":"Text"},{"content":"C","bbox":{"page":1}}]} + ]}}); + let response: ReductoResponse = serde_json::from_value(raw).unwrap(); + let normalized = transform_ocr_response("parse-v3", response) + .unwrap() + .into_json(); + assert_eq!(normalized["pages"][0]["markdown"], "A\n\nC"); + assert_eq!(normalized["pages"][1]["markdown"], "B"); + assert_eq!(normalized["pages"][1]["blocks"][0]["type"], "Table"); + assert_eq!( + normalized["pages"][1]["blocks"][0]["bbox"], + json!({"left":0.1,"top":0.2,"width":0.8,"height":0.3,"page":2,"original_page":4}) + ); + assert_eq!(normalized["pages"][1]["blocks"][0]["confidence"], "high"); + assert_eq!( + normalized["pages"][1]["blocks"][0]["granular_confidence"]["parse_confidence"], + 0.95 + ); + assert!(normalized["pages"][1]["blocks"][0]["image_url"].is_null()); + assert_eq!(normalized["usage_info"]["pages_processed"], 2); + assert_eq!(normalized["usage_info"]["credits"], 3.0); + + let missing: ReductoResponse = + serde_json::from_value(json!({"chunks":[{"content":"text"}]})).unwrap(); + let missing = transform_ocr_response("parse-v3", missing).unwrap(); + assert_eq!(missing.pages[0].markdown, "text"); + let null: ReductoResponse = serde_json::from_value( + json!({"result":null,"chunks":[{"content":"ignored"}],"usage":null}), + ) + .unwrap(); + let null = transform_ocr_response("parse-v3", null).unwrap(); + assert!(null.pages.is_empty()); + } + + #[tokio::test] + async fn facade_omits_native_response_by_default_and_preserves_auth_priority() { + let raw = json!({"job_id":"job-1","result":{"chunks":[]}}); + let (base, seen, server) = mock_server(vec![MockResponse::json(raw)]).await; + let mut request = crate::ocr::test_support::with_source( + wire_request("reducto/parse-v3", &base, json!({})), + "reducto://ready.pdf", + ); + request.transport.extra_headers = vec![("authorization".into(), "Bearer existing".into())]; + + let response = perform_ocr(request).await.unwrap(); + server.await.unwrap(); + assert_eq!(response.provider_native_response, None); + assert!( + seen.lock().unwrap()[0] + .to_ascii_lowercase() + .contains("authorization: bearer existing") + ); + } + + #[tokio::test] + async fn native_format_retains_the_provider_response() { + let raw = json!({ + "result":{"chunks":[{"content":"native OCR response"}]}, + "usage":{"num_pages":1} + }); + let (base, _, server) = mock_server(vec![MockResponse::json(raw.clone())]).await; + let request = crate::ocr::test_support::with_source( + wire_request("reducto/parse-v3", &base, json!({"req_format":"native"})), + "reducto://ready.pdf", + ); + + let response = perform_ocr(request).await.unwrap(); + server.await.unwrap(); + + assert_eq!(response.pages[0].markdown, "native OCR response"); + assert_eq!(response.provider_native_response.as_ref(), raw.as_object()); + } + + #[tokio::test] + async fn unknown_model_reaches_parse_and_keeps_its_name() { + let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ + "result":{"chunks":[{"content":"future model response"}]} + }))]) + .await; + let request = crate::ocr::test_support::with_source( + wire_request("reducto/future-parse-model", &base, json!({})), + "reducto://ready.pdf", + ); + + let response = perform_ocr(request).await.unwrap(); + server.await.unwrap(); + + assert_eq!(response.model, "future-parse-model"); + assert_eq!(response.pages[0].markdown, "future model response"); + let requests = seen.lock().unwrap(); + assert!(requests[0].starts_with("POST /parse ")); + assert_eq!( + request_body(&requests[0]), + json!({"input":"reducto://ready.pdf"}) + ); + } + + #[tokio::test] + async fn guardrail_rewrites_document_before_upload() { + let (base, seen, server) = + mock_server(vec![MockResponse::json(json!({"result":{"chunks":[]}}))]).await; + let host = LocalOcrHost::new(wire_request("reducto/parse-v3", &base, json!({}))) + .with_before_send(|wire, _| { + assert_eq!( + wire.body["document_url"], + "data:application/pdf;base64,YWJj" + ); + Ok(WireRequest { + body: json!({"type":"document_url","document_url":"reducto://guarded.pdf"}), + ..wire + }) + }); + + perform_ocr_with(host).await.unwrap(); + server.await.unwrap(); + let requests = seen.lock().unwrap(); + assert_eq!(requests.len(), 1); + assert!(requests[0].starts_with("POST /parse ")); + assert!(requests[0].contains("reducto://guarded.pdf")); + } + + mod transformation { + use litellm_host::event::{CallEvent, MachineEvent, WireRequest}; + use litellm_llms::{ + base_llm::ocr::transformation::{BaseOcrConfig, OcrConnection, OcrRequestContext}, + reducto::ocr::transformation::*, + }; + use rstest::rstest; + + use super::*; + use crate::ocr::{ + route::LocalOcrHost, + test_support::{ + MockResponse, mock_server, perform_ocr, perform_ocr_with, wire_request, + }, + }; + + #[tokio::test] + async fn v3_options_preserve_explicit_null() { + let overrides = + serde_json::from_value(json!({"formatting":null,"settings":{},"unknown":true})) + .unwrap(); + let params = ReductoParseV3Config + .map_ocr_params(&overrides, "parse-v3") + .unwrap(); + let client = crate::ocr::test_support::ocr_client(); + let connection = OcrConnection::default(); + let document = serde_json::from_value( + json!({"type":"document_url","document_url":"reducto://ready.pdf"}), + ) + .unwrap(); + let body = ReductoParseV3Config + .async_transform_ocr_request( + "parse-v3", + document, + ¶ms, + &[], + OcrRequestContext { + client: &client, + connection: &connection, + }, + ) + .await + .unwrap(); + assert_eq!( + serde_json::to_value(body).unwrap(), + json!({ + "input":"reducto://ready.pdf", "formatting":null, "settings":{} + }) + ); + let absent = ReductoParseV3Config + .map_ocr_params( + &litellm_core_utils::call_arguments::CallArguments::default(), + "parse-v3", + ) + .unwrap(); + assert_eq!(serde_json::to_value(absent).unwrap(), json!({})); + } + + #[rstest] + #[case( + "reducto/parse-v3", + json!({ + "formatting":{"table_output_format":"html"}, + "retrieval":{"chunk_mode":"section"}, + "settings":{"ocr_system":"standard"}, + "future_ocr_option":true, + "extra_body":{"provider_option":"value"} + }), + "reducto://already.pdf", + json!({ + "input":"reducto://already.pdf", + "formatting":{"table_output_format":"html"}, + "retrieval":{"chunk_mode":"section"}, + "settings":{"ocr_system":"standard"}, + "future_ocr_option":true, + "provider_option":"value" + }) + )] + #[case( + "reducto/parse-legacy", + json!({ + "enhance":{"agentic":[{"type":"table"}]}, + "future_ocr_option":true, + "extra_body":{"provider_option":"value"} + }), + "reducto://legacy.pdf", + json!({ + "document_url":"reducto://legacy.pdf", + "options":{"enhance":{"agentic":[{"type":"table"}]}}, + "future_ocr_option":true, + "provider_option":"value" + }) + )] + #[tokio::test] + async fn request_mapping_matches_python( + #[case] model: &str, + #[case] options: Value, + #[case] source: &str, + #[case] expected: Value, + ) { + let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ + "result":{"chunks":[]} + }))]) + .await; + let request = + crate::ocr::test_support::with_source(wire_request(model, &base, options), source); + + perform_ocr(request).await.unwrap(); + server.await.unwrap(); + let requests = seen.lock().unwrap(); + assert_eq!(requests.len(), 1); + assert!(requests[0].starts_with("POST /parse ")); + assert_eq!(request_body(&requests[0]), expected); + } + + #[rstest] + #[case("parse-v3")] + #[case("parse-legacy")] + #[tokio::test] + async fn data_uri_upload_preserves_multipart_headers(#[case] model: &str) { + let (base, seen, server) = mock_server(vec![ + MockResponse::json(json!({"file_id":"reducto://uploaded.pdf"})), + MockResponse::json(json!({"result":{"chunks":[{"content":"hello"}]}})), + ]) + .await; + let mut request = wire_request(&format!("reducto/{model}"), &base, json!({})); + request.transport.extra_headers = vec![ + ("Content-Type".into(), "application/json".into()), + ("X-Trace".into(), "upload-test".into()), + ]; + + let response = perform_ocr(request).await.unwrap(); + server.await.unwrap(); + assert_eq!(response.pages[0].markdown, "hello"); + let requests = seen.lock().unwrap(); + assert_eq!(requests.len(), 2); + assert!(requests[0].starts_with("POST /upload ")); + assert!( + requests[0] + .to_ascii_lowercase() + .contains("content-type: multipart/form-data; boundary=") + ); + assert!(requests[0].contains("x-trace: upload-test")); + assert!(requests[0].contains("application/pdf")); + assert!(requests[0].contains("abc")); + assert!(requests[1].starts_with("POST /parse ")); + } + + #[tokio::test] + async fn response_received_stays_after_reducto_upload_and_parse() { + let (base, seen, server) = mock_server(vec![ + MockResponse::json(json!({"file_id":"reducto://uploaded.pdf"})), + MockResponse::json(json!({"result":{"chunks":[]}})), + ]) + .await; + let request_count = seen.clone(); + let host = LocalOcrHost::new(wire_request("reducto/parse-v3", &base, json!({}))) + .with_observer(move |event| { + if let CallEvent::Machine(MachineEvent::ResponseReceived { raw }) = event { + assert_eq!(request_count.lock().unwrap().len(), 2); + assert_eq!(raw.body, r#"{"result":{"chunks":[]}}"#); + } + }); + + perform_ocr_with(host).await.unwrap(); + server.await.unwrap(); + assert_eq!(seen.lock().unwrap().len(), 2); + } + + #[rstest] + #[case("https://example.com/a.pdf")] + #[case("reducto://")] + #[case("data:application/pdf;base64")] + #[case("data:application/pdf;base64,INVALID!")] + #[tokio::test] + async fn rejects_invalid_document_sources_before_network(#[case] source: &str) { + let request = crate::ocr::test_support::with_source( + wire_request("reducto/parse-v3", "http://127.0.0.1:1", json!({})), + source, + ); + assert!(perform_ocr(request).await.is_err()); + } + + #[tokio::test] + async fn facade_omits_native_response_by_default_and_preserves_auth_priority() { + let raw = json!({"job_id":"job-1","result":{"chunks":[]}}); + let (base, seen, server) = mock_server(vec![MockResponse::json(raw)]).await; + let mut request = crate::ocr::test_support::with_source( + wire_request("reducto/parse-v3", &base, json!({})), + "reducto://ready.pdf", + ); + request.transport.extra_headers = + vec![("authorization".into(), "Bearer existing".into())]; + + let response = perform_ocr(request).await.unwrap(); + server.await.unwrap(); + assert_eq!(response.provider_native_response, None); + assert!( + seen.lock().unwrap()[0] + .to_ascii_lowercase() + .contains("authorization: bearer existing") + ); + } + + #[rstest] + #[case("reducto/parse-v3")] + #[case("reducto/parse-legacy")] + #[tokio::test] + async fn guardrail_headers_reach_upload_and_parse(#[case] model: &str) { + let (base, seen, server) = mock_server(vec![ + MockResponse::json(json!({"file_id":"reducto://uploaded.pdf"})), + MockResponse::json(json!({"result":{"chunks":[]}})), + ]) + .await; + let mut request = wire_request(model, &base, json!({})); + request.transport.extra_headers = + vec![("authorization".into(), "Bearer original".into())]; + let host = LocalOcrHost::new(request).with_before_send(|wire, _| { + Ok(WireRequest { + headers: vec![("authorization".into(), "Bearer guarded".into())], + ..wire + }) + }); + + perform_ocr_with(host).await.unwrap(); + server.await.unwrap(); + let requests = seen.lock().unwrap(); + assert_eq!(requests.len(), 2); + assert!(requests[0].starts_with("POST /upload ")); + assert!(requests[1].starts_with("POST /parse ")); + for request in requests.iter() { + assert!(request.contains("authorization: Bearer guarded")); + assert!(!request.contains("Bearer original")); + } + } + } +} + +#[cfg(test)] +mod vertex_ai_tests { + use litellm_auth::InputSource; + use litellm_llms::base_llm::ocr::{settings::OcrSettings, transformation::OcrResponseFormat}; + use serde_json::{Value, json}; + + use crate::ocr::test_support::{ + MockResponse, mock_server, ocr_client, perform_ocr, wire_request, + }; + + fn request_body(request: &str) -> Value { + serde_json::from_str(request.split_once("\r\n\r\n").unwrap().1).unwrap() + } + + #[tokio::test] + async fn facade_executes_vertex_mistral_with_resolved_project_and_location() { + let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ + "pages":[{"index":0,"markdown":"hello"}], + "usage_info":{"pages_processed":1} + }))]) + .await; + let request = wire_request( + "vertex_ai/mistral-ocr-maas", + &base, + json!({ + "vertex_project":"project-1", + "vertex_location":"europe-west4", + "extract_footer":true + }), + ); + + let response = perform_ocr(request).await.unwrap(); + server.await.unwrap(); + assert_eq!(response.pages[0].markdown, "hello"); + let requests = seen.lock().unwrap(); + assert_eq!(requests.len(), 1); + assert!(requests[0].starts_with( + "POST /v1/projects/project-1/locations/europe-west4/publishers/mistralai/models/mistral-ocr-maas:rawPredict " + )); + assert!( + requests[0] + .to_ascii_lowercase() + .contains("authorization: bearer test-key") + ); + assert_eq!( + request_body(&requests[0]), + json!({ + "model":"mistral-ocr-maas", + "document":{"type":"document_url","document_url":"data:application/pdf;base64,YWJj"}, + "extract_footer":true + }) + ); + } + + #[tokio::test] + async fn configured_project_and_location_apply_when_the_call_sets_neither() { + let (base, seen, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await; + let client = ocr_client().with_settings(OcrSettings { + vertex_project: Some("configured-project".into()), + vertex_location: Some("europe-west4".into()), + ..OcrSettings::default() + }); + + crate::ocr::client::perform( + &client, + wire_request("vertex_ai/mistral-ocr-maas", &base, json!({})), + ) + .await + .unwrap(); + server.await.unwrap(); + assert!(seen.lock().unwrap()[0].starts_with( + "POST /v1/projects/configured-project/locations/europe-west4/publishers/mistralai/models/mistral-ocr-maas:rawPredict " + )); + } + + #[tokio::test] + async fn supplied_authorization_is_forwarded_without_a_static_token() { + let (base, seen, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await; + let mut request = wire_request( + "vertex_ai/model", + &base, + json!({"vertex_project":"project-1"}), + ); + request.credentials.api_key = None; + request.transport.extra_headers = vec![("authorization".into(), "Bearer supplied".into())]; + + perform_ocr(request).await.unwrap(); + server.await.unwrap(); + assert!( + seen.lock().unwrap()[0] + .to_ascii_lowercase() + .contains("authorization: bearer supplied") + ); + } + + #[tokio::test] + async fn invalid_credentials_fail_before_provider_http() { + let request = wire_request( + "vertex_ai/model", + "http://127.0.0.1:1", + json!({"vertex_credentials": true}), + ); + let error = perform_ocr(request).await.unwrap_err(); + assert!(error.to_string().contains("vertex_credentials")); + } + + #[tokio::test] + async fn request_controlled_api_base_is_rejected_before_vertex_auth() { + let mut request = wire_request( + "vertex_ai/mistral-ocr-maas", + "https://caller.example", + json!({"vertex_project":"project-1"}), + ); + request.credentials.api_base = Some(litellm_auth::Sourced::new( + "https://caller.example".into(), + InputSource::Request, + )); + + let error = perform_ocr(request).await.unwrap_err(); + assert!( + error + .to_string() + .contains("request-controlled Vertex AI endpoint") + ); + } + + #[tokio::test] + async fn adapters_build_complete_requests_and_share_mistral_normalization() { + use std::time::Duration; + + use litellm_llms::{ + base_llm::ocr::transformation::BaseOcrConfig, + mistral::ocr::transformation::MistralOcrConfig, + vertex_ai::ocr::transformation::VertexAiOcrConfig, + }; + + use crate::ocr::test_support::ocr_client; + + let client = ocr_client(); + let options = json!({ + "pages": [0, 2], + "include_image_base64": true, + "vertex_project": "project-1", + "vertex_location": "us-central1", + "unknown": "ignored" + }); + let direct = wire_request( + "mistral/mistral-ocr-maas", + "https://mistral.test", + options.clone(), + ); + let vertex = wire_request("vertex_ai/mistral-ocr-maas", "https://vertex.test", options); + let direct = crate::ocr::prepare::prepare_request_for_test( + crate::ocr::test_support::resolved_request(direct), + ); + let vertex = crate::ocr::prepare::prepare_request_for_test( + crate::ocr::test_support::resolved_request(vertex), + ); + let direct_http = MistralOcrConfig + .prepare_request(&direct, &client, &crate::ocr::test_support::NoHooks) + .await + .unwrap(); + let vertex_http = VertexAiOcrConfig + .prepare_request(&vertex, &client, &crate::ocr::test_support::NoHooks) + .await + .unwrap(); + assert_eq!(direct_http.url(), "https://mistral.test/v1/ocr"); + assert_eq!( + vertex_http.url(), + "https://vertex.test/v1/projects/project-1/locations/us-central1/publishers/mistralai/models/mistral-ocr-maas:rawPredict" + ); + for http in [&direct_http, &vertex_http] { + assert_eq!(http.header("authorization").unwrap(), "Bearer test-key"); + assert_eq!(http.header("content-type").unwrap(), "application/json"); + assert_eq!(http.timeout(), Some(Duration::from_secs(2))); + let body: Value = serde_json::from_slice(http.body()).unwrap(); + assert_eq!( + body, + json!({ + "model": "mistral-ocr-maas", + "document": {"type": "document_url", "document_url": "data:application/pdf;base64,YWJj"}, + "pages": [0, 2], + "include_image_base64": true, + "unknown": "ignored" + }) + ); + } + let payload = json!({"pages": [{"index": 0, "markdown": "hello"}], "extra": "preserved"}); + let raw = serde_json::to_vec(&payload).unwrap(); + let direct_response = MistralOcrConfig + .transform_ocr_response(&direct.model, &raw, OcrResponseFormat::Litellm) + .unwrap() + .into_json(); + let vertex_response = VertexAiOcrConfig + .transform_ocr_response(&vertex.model, &raw, OcrResponseFormat::Litellm) + .unwrap() + .into_json(); + assert_eq!(direct_response, vertex_response); + assert_eq!(direct_response["model"], "mistral-ocr-maas"); + assert_eq!(direct_response["object"], "ocr"); + assert_eq!(direct_response["extra"], "preserved"); + } + + mod transformation { + + use rstest::rstest; + use serde_json::{Value, json}; + + use crate::ocr::test_support::wire_request; + + #[rstest] + #[case::mistral(false)] + #[case::vertex(true)] + #[tokio::test] + async fn configs_build_complete_requests_and_share_mistral_normalization( + #[case] use_vertex: bool, + ) { + use std::time::Duration; + + use litellm_llms::{ + base_llm::ocr::transformation::BaseOcrConfig, + mistral::ocr::transformation::MistralOcrConfig, + vertex_ai::ocr::transformation::VertexAiOcrConfig, + }; + + use crate::ocr::test_support::ocr_client; + + let client = ocr_client(); + let options = json!({ + "pages": [0, 2], + "include_image_base64": true, + "vertex_project": "project-1", + "vertex_location": "us-central1", + "unknown": "preserved" + }); + let direct = wire_request( + "mistral/mistral-ocr-maas", + "https://mistral.test", + options.clone(), + ); + let vertex = wire_request("vertex_ai/mistral-ocr-maas", "https://vertex.test", options); + let direct = crate::ocr::prepare::prepare_request_for_test( + crate::ocr::test_support::resolved_request(direct), + ); + let vertex = crate::ocr::prepare::prepare_request_for_test( + crate::ocr::test_support::resolved_request(vertex), + ); + let direct_http = MistralOcrConfig + .prepare_request(&direct, &client, &crate::ocr::test_support::NoHooks) + .await + .unwrap(); + let vertex_http = VertexAiOcrConfig + .prepare_request(&vertex, &client, &crate::ocr::test_support::NoHooks) + .await + .unwrap(); + assert_eq!(direct_http.url(), "https://mistral.test/v1/ocr"); + assert_eq!( + vertex_http.url(), + "https://vertex.test/v1/projects/project-1/locations/us-central1/publishers/mistralai/models/mistral-ocr-maas:rawPredict" + ); + let http = if use_vertex { + &vertex_http + } else { + &direct_http + }; + assert_eq!(http.header("authorization").unwrap(), "Bearer test-key"); + assert_eq!(http.header("content-type").unwrap(), "application/json"); + assert_eq!(http.timeout(), Some(Duration::from_secs(2))); + let body: Value = serde_json::from_slice(http.body()).unwrap(); + assert_eq!( + body, + json!({ + "model": "mistral-ocr-maas", + "document": {"type": "document_url", "document_url": "data:application/pdf;base64,YWJj"}, + "pages": [0, 2], + "include_image_base64": true, + "unknown": "preserved" + }) + ); + let payload = serde_json::to_vec( + &json!({"pages": [{"index": 0, "markdown": "hello"}], "extra": "preserved"}), + ) + .unwrap(); + let direct_response = MistralOcrConfig + .transform_ocr_response(&direct.model, &payload, Default::default()) + .unwrap() + .into_json(); + let vertex_response = VertexAiOcrConfig + .transform_ocr_response(&vertex.model, &payload, Default::default()) + .unwrap() + .into_json(); + assert_eq!(direct_response, vertex_response); + assert_eq!(direct_response["model"], "mistral-ocr-maas"); + assert_eq!(direct_response["object"], "ocr"); + assert_eq!(direct_response["extra"], "preserved"); + } + } +} + +#[cfg(test)] +mod vertex_ai_deepseek_tests { + use litellm_auth::InputSource; + use serde_json::{Value, json}; + + use crate::ocr::test_support::{MockResponse, mock_server, perform_ocr, wire_request}; + + fn request_body(request: &str) -> Value { + serde_json::from_str(request.split_once("\r\n\r\n").unwrap().1).unwrap() + } + + #[tokio::test] + async fn facade_executes_vertex_deepseek_at_the_openai_endpoint() { + let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ + "choices":[{"message":{"content":"recognized"}}], + "usage":{"prompt_tokens":1} + }))]) + .await; + let request = wire_request( + "vertex_ai/deepseek-ocr-maas", + &base, + json!({ + "vertex_project":"project-1", + "vertex_location":"europe-west4", + "temperature":0.1, + "future_ocr_option":true, + "extra_body":{"provider_option":"value"} + }), + ); + let request = crate::ocr::test_support::with_source(request, "gs://bucket/document.pdf"); + + let response = perform_ocr(request).await.unwrap(); + server.await.unwrap(); + assert_eq!(response.pages[0].markdown, "recognized"); + assert_eq!( + response.usage_info.unwrap().extra_fields["prompt_tokens"], + 1 + ); + let requests = seen.lock().unwrap(); + assert!(requests[0].starts_with( + "POST /v1/projects/project-1/locations/europe-west4/endpoints/openapi/chat/completions " + )); + assert!( + requests[0] + .to_ascii_lowercase() + .contains("authorization: bearer test-key") + ); + let body = request_body(&requests[0]); + assert_eq!(body["model"], "deepseek-ai/deepseek-ocr-maas"); + assert_eq!(body["temperature"], 0.1); + assert_eq!(body["future_ocr_option"], true); + assert!(body.get("extra_body").is_none()); + assert_eq!( + body["messages"][0]["content"][0], + json!({"type":"image_url","image_url":"gs://bucket/document.pdf"}) + ); + } + + #[test] + fn host_registration_selects_deepseek_without_affecting_mistral() { + assert!(crate::ocr::arguments::is_supported_request( + "deepseek-ocr-maas", + Some("vertex_ai") + )); + assert!(crate::ocr::arguments::is_supported_request( + "mistral-ocr-maas", + Some("vertex_ai") + )); + } + + #[tokio::test] + async fn request_controlled_api_base_is_rejected_before_vertex_auth() { + let mut request = wire_request( + "vertex_ai/deepseek-ocr-maas", + "https://caller.example", + json!({"vertex_project":"project-1"}), + ); + request.credentials.api_base = Some(litellm_auth::Sourced::new( + "https://caller.example".into(), + InputSource::Request, + )); + + let error = perform_ocr(request).await.unwrap_err(); + assert!( + error + .to_string() + .contains("request-controlled Vertex AI endpoint") + ); + } + + mod deepseek_transformation { + use serde_json::json; + + use super::*; + use crate::ocr::test_support::{MockResponse, mock_server, perform_ocr, wire_request}; + + #[tokio::test] + async fn facade_executes_vertex_deepseek_at_the_openai_endpoint() { + let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ + "choices":[{"message":{"content":"recognized"}}], + "usage":{"prompt_tokens":1} + }))]) + .await; + let request = wire_request( + "vertex_ai/deepseek-ocr-maas", + &base, + json!({ + "vertex_project":"project-1", + "vertex_location":"europe-west4", + "temperature":0.1, + "future_ocr_option":true, + "extra_body":{"provider_option":"value"} + }), + ); + let request = + crate::ocr::test_support::with_source(request, "gs://bucket/document.pdf"); + + let response = perform_ocr(request).await.unwrap(); + server.await.unwrap(); + assert_eq!(response.pages[0].markdown, "recognized"); + assert_eq!( + response.usage_info.unwrap().extra_fields["prompt_tokens"], + 1 + ); + let requests = seen.lock().unwrap(); + assert!(requests[0].starts_with( + "POST /v1/projects/project-1/locations/europe-west4/endpoints/openapi/chat/completions " + )); + assert!( + requests[0] + .to_ascii_lowercase() + .contains("authorization: bearer test-key") + ); + let body = request_body(&requests[0]); + assert_eq!(body["model"], "deepseek-ai/deepseek-ocr-maas"); + assert_eq!(body["temperature"], 0.1); + assert_eq!(body["future_ocr_option"], true); + assert_eq!(body["provider_option"], "value"); + assert!(body.get("vertex_project").is_none()); + assert!(body.get("extra_body").is_none()); + assert_eq!( + body["messages"][0]["content"][0], + json!({"type":"image_url","image_url":"gs://bucket/document.pdf"}) + ); + } + } +} + +#[cfg(test)] +pub(crate) mod tests { + use std::sync::{Arc, Mutex}; + + use futures_util::future::BoxFuture; + use litellm_auth_gcp::VertexAuth; + use litellm_host::{ + event::{CallEvent, MachineEvent, WireRequest}, + host::{Host, HostOp, HostResult}, + machine::{HostFailure, Machine, MachineStep}, + }; + use litellm_http::{ + HttpClientPool, HttpSettings, Resolution, + media::{PublicDnsResolver, UrlPolicy}, + }; + use litellm_llms::base_llm::ocr::{ + error::Error as OcrError, + handler::OcrClient, + settings::OcrSettings, + transformation::{ + BaseOcrConfig, LiteLLMOcrResponse, OCR_RESPONSE_MAX_BYTES, OcrTransportConfig, + }, + }; + use litellm_secrets::source::SecretSource; + use rstest::rstest; + use serde_json::{Value, json}; + + use crate::ocr::route::{LocalOcrHost, OcrOp, OcrOpResult, ocr_machine}; + use crate::ocr::{ + test_support::{ + MockResponse, mock_server, ocr_client, perform_ocr, perform_ocr_with, wire_request, + }, + wire::{OcrWireRequest, decode_request}, + }; + + struct RecordingSecretSource { + names: Arc>>, + values: &'static [(&'static str, &'static str)], + api_base: String, + } + + impl SecretSource for RecordingSecretSource { + fn get_secret_str<'a>( + &'a self, + name: &'a str, + ) -> BoxFuture<'a, Result, litellm_secrets::Error>> + { + self.names.lock().unwrap().push(name.to_owned()); + Box::pin(async move { + Ok(match name { + "MISTRAL_AZURE_API_BASE" => Some(self.api_base.clone()), + _ => self + .values + .iter() + .find(|(key, _)| *key == name) + .map(|(_, value)| value.to_string()), + } + .map(litellm_secrets::SecretValue::new)) + }) + } + } + + #[rstest] + #[case::mistral("mistral/model", json!({}))] + #[case::vertex("vertex_ai/mistral-ocr-latest", json!({"vertex_project":"test-project", "vertex_location":"us-central1"}))] + #[tokio::test] + async fn ocr_contract_upstream_error_preserves_status_body_and_headers( + #[case] model: &str, + #[case] options: Value, + ) { + let payload = json!({"message": format!("{} END-OF-PROVIDER-BODY", "x".repeat(4096))}); + let expected_body = serde_json::to_string(&payload).unwrap(); + let (base, seen, server) = mock_server(vec![MockResponse { + status: 422, + headers: vec![ + ("Retry-After", "17".into()), + ("X-Request-ID", "request-123".into()), + ("X-Future-Header", "retained".into()), + ], + body: payload, + }]) + .await; + let error = perform_ocr(wire_request(model, &base, options)) + .await + .unwrap_err(); + server.await.unwrap(); + assert_eq!(seen.lock().unwrap().len(), 1); + let OcrError::Provider { + status, + body, + headers, + } = error + else { + panic!("expected provider error, got {error:?}"); + }; + assert_eq!(status, 422); + for (name, value) in [ + ("retry-after", "17"), + ("x-request-id", "request-123"), + ("x-future-header", "retained"), + ] { + assert!( + headers + .iter() + .any(|(key, actual)| key.eq_ignore_ascii_case(name) && actual == value) + ); + } + assert_eq!( + body.len(), + expected_body.len(), + "provider error body was truncated" + ); + assert_eq!(body, expected_body); + } + + #[test] + fn request_boundary_selects_mistral_and_rejects_unknown_providers() { + let request = OcrWireRequest { + model: "mistral/model".into(), + document: json!({"type":"document_url","document_url":"https://example.com/doc.pdf"}), + api_key: Some(litellm_auth::SecretValue::new("key")), + api_base: None, + custom_llm_provider: None, + extra_headers: None, + optional_params: json!({"extract_header":true,"unknown":42}) + .as_object() + .unwrap() + .clone(), + input_sources: Default::default(), + timeout_seconds: None, + }; + assert!(decode_request(request).is_ok()); + assert!( + decode_request(OcrWireRequest { + model: "model".into(), + document: json!({"type":"document_url","document_url":"https://example.com/doc.pdf"}), + api_key: Some(litellm_auth::SecretValue::new("key")), + api_base: None, + custom_llm_provider: Some("unknown".into()), + extra_headers: None, + optional_params: serde_json::Map::new(), + input_sources: Default::default(), + timeout_seconds: None, + }) + .is_err() + ); + } + + #[tokio::test] + async fn facade_executes_direct_mistral_once() { + let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ + "pages":[{"index":0,"markdown":"hello","custom":"preserved"}], + "usage_info":{"pages_processed":1} + }))]) + .await; + let result = perform_ocr(wire_request( + "mistral/model", + &base, + json!({"pages":"0,2-4","extract_header":true,"unknown":"ignored"}), + )) + .await + .unwrap(); + server.await.unwrap(); + assert_eq!(result.pages[0].markdown, "hello"); + assert_eq!(result.pages[0].extra_fields["custom"], "preserved"); + let requests = seen.lock().unwrap(); + assert_eq!(requests.len(), 1); + assert!(requests[0].starts_with("POST /v1/ocr ")); + assert!( + requests[0] + .to_ascii_lowercase() + .contains("authorization: bearer test-key\r\n") + ); + let body: Value = + serde_json::from_str(requests[0].split_once("\r\n\r\n").unwrap().1).unwrap(); + assert_eq!( + body, + json!({ + "model":"model", + "document":{"type":"document_url","document_url":"data:application/pdf;base64,YWJj"}, + "pages":"0,2-4", + "extract_header":true, + "unknown":"ignored" + }) + ); + } + + #[tokio::test] + async fn facade_retains_native_response_when_requested() { + let provider_response = json!({ + "pages":[{"index":0,"markdown":"hello"}], + "usage_info":{"pages_processed":1}, + "provider_only":"preserved" + }); + let (base, _, server) = + mock_server(vec![MockResponse::json(provider_response.clone())]).await; + let response = perform_ocr(wire_request( + "mistral/model", + &base, + json!({"req_format":"native"}), + )) + .await + .unwrap(); + + server.await.unwrap(); + assert_eq!( + response.provider_native_response.map(Value::Object), + Some(provider_response) + ); + } + + #[rstest] + #[case::plain_key(&[("MISTRAL_API_KEY", "plain")], "plain")] + #[case::azure_key_wins(&[("MISTRAL_AZURE_API_KEY", "azure"), ("MISTRAL_API_KEY", "plain")], "azure")] + #[case::empty_azure_key_falls_through(&[("MISTRAL_AZURE_API_KEY", ""), ("MISTRAL_API_KEY", "plain")], "plain")] + #[tokio::test] + async fn mistral_env_fallbacks_follow_python_through_the_injected_secret_source( + #[case] secrets: &'static [(&'static str, &'static str)], + #[case] expected_key: &str, + ) { + let (base, seen, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await; + let names = Arc::new(Mutex::new(Vec::new())); + let client = ocr_client().with_secrets(Arc::new(RecordingSecretSource { + names: names.clone(), + values: secrets, + api_base: base.clone(), + })); + let request = decode_request(OcrWireRequest { + model: "mistral/model".into(), + document: json!({"type":"document_url","document_url":"data:application/pdf;base64,YWJj"}), + api_key: None, + api_base: None, + custom_llm_provider: None, + extra_headers: None, + optional_params: Default::default(), + input_sources: Default::default(), + timeout_seconds: Some(2.0), + }) + .unwrap(); + + crate::ocr::client::perform(&client, request).await.unwrap(); + server.await.unwrap(); + assert_eq!( + *names.lock().unwrap(), + litellm_llms::mistral::ocr::transformation::MistralOcrConfig.secret_names() + ); + assert!(seen.lock().unwrap()[0].contains(&format!("authorization: Bearer {expected_key}"))); + } + + #[tokio::test] + async fn mistral_ocr_resolves_provider_secrets_before_transformation() { + let (base, seen, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await; + let names = Arc::new(Mutex::new(Vec::new())); + let client = ocr_client().with_secrets(Arc::new(RecordingSecretSource { + names: names.clone(), + values: &[("MISTRAL_API_KEY", "source-key")], + api_base: base.clone(), + })); + let request = decode_request(OcrWireRequest { + model: "mistral/mistral-ocr-latest".into(), + document: json!({ + "type":"document_url", + "document_url":"data:application/pdf;base64,YWJj" + }), + api_key: None, + api_base: None, + custom_llm_provider: None, + extra_headers: None, + optional_params: Default::default(), + input_sources: Default::default(), + timeout_seconds: Some(2.0), + }) + .unwrap(); + + crate::ocr::client::perform(&client, request).await.unwrap(); + server.await.unwrap(); + assert_eq!( + *names.lock().unwrap(), + litellm_llms::mistral::ocr::transformation::MistralOcrConfig.secret_names() + ); + assert!(seen.lock().unwrap()[0].contains("authorization: Bearer source-key")); + } + + #[tokio::test] + async fn ocr_client_uses_the_injected_http_pool_configuration() { + let (base, seen, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await; + let settings = HttpSettings { + user_agent: Some("host-owned/1".into()), + ..HttpSettings::default() + }; + let client = OcrClient::new( + &HttpClientPool::new(Arc::new(PublicDnsResolver)), + &Resolution::from(&settings).config, + UrlPolicy::default(), + VertexAuth::default(), + OcrSettings::default(), + Arc::new(litellm_secrets::source::EnvironmentSecrets::default()), + ) + .unwrap(); + crate::ocr::client::perform(&client, wire_request("mistral/model", &base, json!({}))) + .await + .unwrap(); + server.await.unwrap(); + assert!(seen.lock().unwrap()[0].contains("user-agent: host-owned/1")); + } + + fn event_name(event: &CallEvent) -> &'static str { + match event { + CallEvent::Started { .. } => "started", + CallEvent::Machine(MachineEvent::ResponseReceived { .. }) => "response", + CallEvent::Succeeded { .. } => "success", + CallEvent::Failed { .. } => "failure", + } + } + + fn recording_host( + request: crate::ocr::types::LiteLLMOcrRequest, + events: Arc>>, + block: bool, + ) -> LocalOcrHost { + let before_send_events = events.clone(); + LocalOcrHost::new(request) + .with_before_send(move |wire, _| { + before_send_events.lock().unwrap().push("before_send"); + if block { + return Err(OcrError::InvalidRequest("blocked".into())); + } + Ok(wire) + }) + .with_observer(move |event| events.lock().unwrap().push(event_name(event))) + } + + #[tokio::test] + async fn lifecycle_sends_headers_returned_by_the_before_send_operation() { + let (base, seen, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await; + let host = LocalOcrHost::new(wire_request("mistral/model", &base, json!({}))) + .with_before_send(|mut wire, _| { + wire.headers + .push(("x-core-callback".into(), "edited".into())); + Ok(wire) + }); + + perform_ocr_with(host).await.unwrap(); + server.await.unwrap(); + + assert!(seen.lock().unwrap()[0].contains("x-core-callback: edited")); + } + + #[tokio::test] + async fn before_send_context_names_the_route_and_its_secrets() { + let (base, _, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await; + let observed = Arc::new(Mutex::new(None)); + let captured = observed.clone(); + let host = LocalOcrHost::new(wire_request( + "mistral/model", + &base, + json!({"pages": [0], "req_format": "native"}), + )) + .with_before_send(move |wire, context| { + *captured.lock().unwrap() = Some((wire.clone(), context.clone())); + Ok(wire) + }); + perform_ocr_with(host).await.unwrap(); + server.await.unwrap(); + let (wire, context) = observed.lock().unwrap().take().unwrap(); + assert_eq!(context.custom_llm_provider, "mistral"); + assert_eq!(context.model, "model"); + assert_eq!(wire.body["pages"], json!([0])); + assert!(context.secret_fields.is_empty()); + assert_eq!(context.optional_params["req_format"], "native"); + + let (base, _, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await; + let observed = Arc::new(Mutex::new(None)); + let captured = observed.clone(); + let request = wire_request( + "azure_ai/model", + &base, + json!({"client_secret": "shh", "tenant_id": "t"}), + ); + let request = request.with_document(crate::ocr::types::OcrDocumentInput::Bytes { + bytes: b"abc".as_slice().into(), + file_name: None, + mime_type: Some("application/pdf".into()), + }); + let host = LocalOcrHost::new(request).with_before_send(move |wire, context| { + *captured.lock().unwrap() = Some(context.clone()); + Ok(wire) + }); + perform_ocr_with(host).await.unwrap(); + server.await.unwrap(); + let context = observed.lock().unwrap().take().unwrap(); + assert_eq!(context.secret_fields, ["client_secret"]); + } + + #[tokio::test] + async fn lifecycle_orders_hooks_and_emits_one_success() { + let (base, seen, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await; + let events = Arc::new(Mutex::new(Vec::new())); + let host = recording_host( + wire_request("mistral/model", &base, json!({})), + events.clone(), + false, + ); + perform_ocr_with(host).await.unwrap(); + server.await.unwrap(); + assert_eq!( + *events.lock().unwrap(), + ["started", "before_send", "response", "success"] + ); + assert_eq!(seen.lock().unwrap().len(), 1); + } + + #[tokio::test] + async fn lifecycle_blocking_prevents_execution_and_emits_one_failure() { + let events = Arc::new(Mutex::new(Vec::new())); + let host = recording_host( + wire_request("mistral/model", "http://127.0.0.1:1", json!({})), + events.clone(), + true, + ); + let error = perform_ocr_with(host).await.unwrap_err(); + assert!(matches!(error, OcrError::InvalidRequest(message) if message == "blocked")); + assert_eq!( + *events.lock().unwrap(), + ["started", "before_send", "failure"] + ); + } + + #[tokio::test] + async fn upstream_failure_emits_one_terminal_failure() { + let (base, seen, server) = mock_server(vec![MockResponse { + status: 500, + headers: vec![], + body: json!({"error":"failed"}), + }]) + .await; + let events = Arc::new(Mutex::new(Vec::new())); + let host = recording_host( + wire_request("mistral/model", &base, json!({})), + events.clone(), + false, + ); + assert!(perform_ocr_with(host).await.is_err()); + server.await.unwrap(); + assert_eq!( + *events.lock().unwrap(), + ["started", "before_send", "failure"] + ); + assert_eq!(seen.lock().unwrap().len(), 1); + } + + /// Drives the machine by hand, answering every op through `host` except `before_send`, + /// which `intercept` answers so a test can fail or cancel exactly there. + async fn drive_until( + client: OcrClient, + host: &LocalOcrHost, + mut intercept: impl FnMut(WireRequest) -> Result>, + ) -> ( + Result, + Vec<&'static str>, + crate::ocr::route::OcrMachine, + ) { + let mut machine = ocr_machine(client); + let mut result = None; + let mut ops = Vec::new(); + let outcome = loop { + let op = match machine.resume(result.take()).await { + Ok(MachineStep::Host(op)) => op, + Ok(MachineStep::Complete(response)) => break Ok(response), + Err(error) => break Err(error), + }; + let answer = match op { + HostOp::Route(op) => { + ops.push(match op { + OcrOp::ProjectRequest => "ProjectRequest", + OcrOp::ReadDocument => "ReadDocument", + OcrOp::AcquireAzureAdToken => "AcquireAzureAdToken", + }); + host.route(op) + .await + .map(HostResult::Route) + .map_err(HostFailure::Error) + } + HostOp::BeforeSend { wire, .. } => { + ops.push("BeforeSend"); + intercept(*wire).map(|wire| HostResult::BeforeSend(Box::new(wire))) + } + HostOp::Emit(event) => { + let event = CallEvent::Machine(event); + ops.push(event_name(&event)); + host.emit(&event) + .await + .map(|()| HostResult::Emitted) + .map_err(HostFailure::Error) + } + }; + match answer { + Ok(answer) => result = Some(answer), + Err(failure) => break machine.interrupt(failure).await, + } + }; + (outcome, ops, machine) + } + + #[tokio::test] + async fn failed_before_send_does_not_replay_or_reach_transport() { + let host = LocalOcrHost::new(wire_request( + "mistral/model", + "http://127.0.0.1:1", + json!({}), + )); + let (outcome, ops, mut machine) = drive_until(ocr_client(), &host, |_| { + Err(HostFailure::Error(OcrError::InvalidRequest( + "before_send failed".into(), + ))) + }) + .await; + assert!( + matches!(outcome, Err(OcrError::InvalidRequest(message)) if message == "before_send failed") + ); + assert_eq!(ops, ["ProjectRequest", "BeforeSend"]); + assert!(machine.resume(None).await.is_err()); + } + + #[tokio::test] + async fn invalid_provider_response_emits_response_received_before_normalization_failure() { + let (base, seen, server) = + mock_server(vec![MockResponse::json(json!({"pages":"invalid"}))]).await; + let responses_received = Arc::new(Mutex::new(Vec::new())); + let observed = responses_received.clone(); + let host = LocalOcrHost::new(wire_request("mistral/model", &base, json!({}))) + .with_observer(move |event| { + if let CallEvent::Machine(MachineEvent::ResponseReceived { raw }) = event { + observed.lock().unwrap().push(raw.body.clone()); + } + }); + let error = perform_ocr_with(host).await.unwrap_err(); + server.await.unwrap(); + assert!(matches!(error, OcrError::ResponseField { .. })); + assert_eq!(seen.lock().unwrap().len(), 1); + assert_eq!( + *responses_received.lock().unwrap(), + [r#"{"pages":"invalid"}"#] + ); + } + + #[tokio::test] + async fn direct_native_host_drives_the_same_state_machine() { + let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ + "pages":[{"index":0,"markdown":"native"}] + }))]) + .await; + let host = LocalOcrHost::new(wire_request("mistral/model", &base, json!({}))); + let (outcome, ops, mut machine) = drive_until(ocr_client(), &host, Ok).await; + server.await.unwrap(); + assert_eq!(outcome.unwrap().pages[0].markdown, "native"); + assert_eq!(seen.lock().unwrap().len(), 1); + assert_eq!(ops, ["ProjectRequest", "BeforeSend", "response"]); + assert!(matches!( + machine.resume(None).await, + Err(OcrError::InvalidRequest(_)) + )); + } + + async fn drive_native_file_call( + request: crate::ocr::types::LiteLLMOcrRequest, + content: Result, + ) -> (Result, usize) { + let reads = Arc::new(Mutex::new(0)); + let counted = reads.clone(); + let content = Mutex::new(Some(content)); + let host = LocalOcrHost::new(request).with_reader(move || { + *counted.lock().unwrap() += 1; + content.lock().unwrap().take().unwrap() + }); + let outcome = perform_ocr_with(host).await; + let reads = *reads.lock().unwrap(); + (outcome, reads) + } + + #[tokio::test] + async fn host_reader_documents_are_read_once_at_the_core_selected_point_and_encoded() { + let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ + "pages":[{"index":0,"markdown":"file"}] + }))]) + .await; + let request = wire_request("mistral/model", &base, json!({})).with_document( + crate::ocr::types::OcrDocumentInput::HostReader { + mime_type: Some("application/pdf".into()), + }, + ); + let (response, reads) = drive_native_file_call( + request, + Ok(crate::ocr::types::OcrFileContent { + bytes: b"abc".as_slice().into(), + file_name: Some("scan.png".into()), + }), + ) + .await; + server.await.unwrap(); + assert_eq!(response.unwrap().pages[0].markdown, "file"); + assert_eq!(reads, 1); + assert!(seen.lock().unwrap()[0].contains("data:application/pdf;base64,YWJj")); + } + + #[tokio::test] + async fn host_reader_failures_and_empty_files_fail_before_the_provider_is_called() { + let (base, seen, _server) = mock_server(vec![]).await; + let request = wire_request("mistral/model", &base, json!({})); + let failure = OcrError::InvalidRequest("reader exploded".into()); + let (response, reads) = drive_native_file_call( + request + .with_document(crate::ocr::types::OcrDocumentInput::HostReader { mime_type: None }), + Err(failure.clone()), + ) + .await; + assert!( + matches!(response.unwrap_err(), OcrError::InvalidRequest(message) if message == "reader exploded") + ); + assert_eq!(reads, 1); + + let request = wire_request("mistral/model", &base, json!({})); + let (response, _) = drive_native_file_call( + request + .with_document(crate::ocr::types::OcrDocumentInput::HostReader { mime_type: None }), + Ok(crate::ocr::types::OcrFileContent { + bytes: Default::default(), + file_name: None, + }), + ) + .await; + assert!(matches!(response.unwrap_err(), OcrError::EmptyFile)); + assert!(seen.lock().unwrap().is_empty()); + } + + #[tokio::test] + async fn path_documents_are_read_by_core_without_a_host_operation() { + let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ + "pages":[{"index":0,"markdown":"path"}] + }))]) + .await; + let dir = std::env::temp_dir().join(format!("litellm-ocr-{}", rand::random::())); + std::fs::create_dir_all(&dir).unwrap(); + let path = dir.join("scan.png"); + std::fs::write(&path, b"abc").unwrap(); + let request = wire_request("mistral/model", &base, json!({})).with_document( + crate::ocr::types::OcrDocumentInput::Path { + path: path.clone(), + mime_type: None, + }, + ); + let (response, reads) = + drive_native_file_call(request, Err(OcrError::InvalidRequest("unused".into()))).await; + server.await.unwrap(); + std::fs::remove_dir_all(&dir).unwrap(); + assert_eq!(response.unwrap().pages[0].markdown, "path"); + assert_eq!(reads, 0); + assert!(seen.lock().unwrap()[0].contains("data:image/png;base64,YWJj")); + + let (base, seen, _server) = mock_server(vec![]).await; + let request = wire_request("mistral/model", &base, json!({})); + let (response, _) = drive_native_file_call( + request.with_document(crate::ocr::types::OcrDocumentInput::Path { + path: path.clone(), + mime_type: None, + }), + Err(OcrError::InvalidRequest("unused".into())), + ) + .await; + assert!(matches!( + response.unwrap_err(), + OcrError::FileRead { path: failed, source } if failed == path && source.kind() == std::io::ErrorKind::NotFound + )); + assert!(seen.lock().unwrap().is_empty()); + } + + #[tokio::test] + async fn cancellation_at_before_send_prevents_execution_and_further_resumption() { + let host = LocalOcrHost::new(wire_request( + "mistral/model", + "http://127.0.0.1:1", + json!({}), + )); + let (outcome, ops, mut machine) = drive_until(ocr_client(), &host, |_| { + Err(HostFailure::Cancelled(OcrError::InvalidRequest( + "cancelled".into(), + ))) + }) + .await; + assert!( + matches!(outcome, Err(OcrError::InvalidRequest(message)) if message == "cancelled") + ); + assert_eq!(ops, ["ProjectRequest", "BeforeSend"]); + assert!(machine.resume(Some(HostResult::Emitted)).await.is_err()); + } + + #[tokio::test] + async fn missing_host_result_preserves_pending_operation() { + let request = wire_request("mistral/model", "http://127.0.0.1:1", json!({})); + let mut machine = ocr_machine(ocr_client()); + assert!(matches!( + machine.resume(None).await.unwrap(), + MachineStep::Host(HostOp::Route(OcrOp::ProjectRequest)) + )); + assert!(machine.resume(None).await.is_err()); + assert!(matches!( + machine + .resume(Some(HostResult::Route(OcrOpResult::Request { + request: Box::new(request), + caller_token: false, + }))) + .await + .unwrap(), + MachineStep::Host(HostOp::BeforeSend { .. }) + )); + } + + async fn read_bounded_response( + response: Vec, + limit: usize, + ) -> Result { + use tokio::io::{AsyncReadExt, AsyncWriteExt}; + + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + let server = tokio::spawn(async move { + let (mut socket, _) = listener.accept().await.unwrap(); + let mut request = [0; 4096]; + assert!(socket.read(&mut request).await.unwrap() > 0); + socket.write_all(&response).await.unwrap(); + std::future::pending::<()>().await; + }); + let response = reqwest::Client::new() + .get(format!("http://{address}")) + .send() + .await + .unwrap(); + let result = tokio::time::timeout( + std::time::Duration::from_secs(2), + litellm_llms::base_llm::ocr::handler::read_response_bytes(response, limit), + ) + .await; + server.abort(); + let _ = server.await; + result.expect("bounded reads must finish without waiting for the rest of an oversized body") + } + + #[tokio::test] + async fn response_limit_accepts_exact_size_and_rejects_declared_and_chunked_overflow() { + use litellm_llms::base_llm::ocr::error::Error; + + for response in [ + "HTTP/1.1 200 OK\r\nContent-Length: 8\r\n\r\nabcdefgh", + "HTTP/1.1 200 OK\r\nTransfer-Encoding: chunked\r\n\r\n4\r\nabcd\r\n4\r\nefgh\r\n0\r\n\r\n", + ] { + assert_eq!( + read_bounded_response(response.as_bytes().to_vec(), 8) + .await + .unwrap(), + "abcdefgh" + ); + } + for response in [ + "HTTP/1.1 200 OK\r\nContent-Length: 9\r\n\r\n", + "HTTP/1.1 200 OK\r\nTransfer-Encoding: chunked\r\n\r\n4\r\nabcd\r\n5\r\nefghi\r\n", + ] { + assert!(matches!( + read_bounded_response(response.as_bytes().to_vec(), 8).await, + Err(Error::TooLarge { limit: 8 }) + )); + } + } + + #[rstest] + #[case::declared("Content-Length: 1000000")] + #[case::chunked("Transfer-Encoding: chunked")] + #[tokio::test] + async fn oversized_error_retains_http_status_and_bounded_diagnostics_without_draining( + #[case] headers: &str, + ) { + let prefix = "x".repeat(4096); + let body = if headers.starts_with("Transfer") { + format!("{:x}\r\n{prefix}\r\n", prefix.len()) + } else { + prefix.clone() + }; + let response = format!("HTTP/1.1 429 Too Many Requests\r\n{headers}\r\n\r\n{body}"); + let error = read_bounded_response(response.into_bytes(), prefix.len()) + .await + .unwrap_err(); + match error { + OcrError::Transport(litellm_http::transport::Error::Http { status, body }) => { + assert_eq!(status, 429); + assert_eq!(body, prefix); + } + error => panic!("unexpected error: {error}"), + } + } + + #[test] + fn response_limit_is_validated_and_not_forwarded_to_the_provider() { + let request = wire_request( + "mistral/model", + "http://localhost", + json!({"max_response_bytes": 123}), + ); + assert_eq!(request.transport.max_response_bytes, 123); + assert!(!request.optional_params.contains_key("max_response_bytes")); + for value in [ + json!(0), + json!(-1), + json!(true), + json!("123"), + json!(1.5), + json!(OCR_RESPONSE_MAX_BYTES + 1), + Value::Null, + ] { + let wire = serde_json::from_value(json!({ + "model": "mistral/model", "document": {"type": "document_url", "document_url": "data:application/pdf;base64,YWJj"}, + "optional_params": {"max_response_bytes": value} + })).unwrap(); + let Err(error) = decode_request(wire) else { + panic!("invalid response limit accepted") + }; + assert!(error.to_string().contains("max_response_bytes")); + } + } + + #[derive(Debug)] + struct PendingToken { + entered: Arc, + dropped: Arc, + } + + struct TokenFutureDrop(Arc); + + impl Drop for TokenFutureDrop { + fn drop(&mut self) { + self.0.store(true, std::sync::atomic::Ordering::SeqCst); + } + } + + impl litellm_auth::TokenProvider for PendingToken { + fn acquire(&self) -> litellm_auth::TokenFuture<'_> { + Box::pin(async move { + let _guard = TokenFutureDrop(self.dropped.clone()); + self.entered.notify_one(); + std::future::pending().await + }) + } + } + + #[tokio::test] + async fn interrupt_drops_provider_captures_before_returning() { + use std::sync::atomic::{AtomicBool, Ordering}; + + let entered = Arc::new(tokio::sync::Notify::new()); + let dropped = Arc::new(AtomicBool::new(false)); + let request = wire_request("azure_ai/mistral-ocr", "https://example.invalid", json!({})); + let request = crate::ocr::types::LiteLLMOcrRequest { + transport: OcrTransportConfig { + extra_headers: vec![("authorization".into(), "Bearer test-key".into())], + ..request.transport + }, + azure_ad_token_provider: Some(litellm_auth::TokenProviderHandle::new(Arc::new( + PendingToken { + entered: entered.clone(), + dropped: dropped.clone(), + }, + ))), + ..request + }; + let host = LocalOcrHost::new(request); + let mut machine = ocr_machine(ocr_client()); + let mut result = None; + tokio::time::timeout(std::time::Duration::from_secs(2), async { + loop { + tokio::select! { + _ = entered.notified() => break, + step = machine.resume(result.take()) => { + result = Some(match step.unwrap() { + MachineStep::Host(HostOp::Route(op)) => HostResult::Route(host.route(op).await.unwrap()), + MachineStep::Host(HostOp::BeforeSend { wire, .. }) => { + HostResult::BeforeSend(wire) + } + MachineStep::Host(HostOp::Emit(_)) => HostResult::Emitted, + MachineStep::Complete(_) => panic!("pending provider completed"), + }); + } + } + } + }) + .await + .unwrap(); + assert!(!dropped.load(Ordering::SeqCst)); + let selected = OcrError::InvalidRequest("cancelled".into()); + let acknowledgement = machine.interrupt(HostFailure::Cancelled(selected.clone())); + assert!( + dropped.load(Ordering::SeqCst), + "interrupt returned while provider captures were still alive" + ); + assert!( + matches!(acknowledgement.await, Err(OcrError::InvalidRequest(message)) if message == "cancelled") + ); + } + + struct CallerTokenHost { + request: Mutex>, + trace: Mutex>, + } + + impl Host for CallerTokenHost { + async fn route(&self, op: OcrOp) -> Result { + match op { + OcrOp::ProjectRequest => { + self.trace.lock().unwrap().push("project".into()); + Ok(OcrOpResult::Request { + request: Box::new(self.request.lock().unwrap().take().unwrap()), + caller_token: true, + }) + } + OcrOp::AcquireAzureAdToken => { + self.trace.lock().unwrap().push("token".into()); + Ok(OcrOpResult::AzureAdToken( + litellm_auth::ResolvedCredential::Static(litellm_auth::SecretValue::new( + "caller-token", + )), + )) + } + OcrOp::ReadDocument => Err(OcrError::InvalidRequest("no reader".into())), + } + } + + async fn before_send( + &self, + wire: WireRequest, + _: &litellm_host::event::RequestContext, + ) -> Result { + let is_authorization = |name: &str| name.eq_ignore_ascii_case("authorization"); + let authorization = wire + .headers + .iter() + .find(|(name, _)| is_authorization(name)) + .map(|(_, value)| value.clone()) + .unwrap_or_default(); + self.trace + .lock() + .unwrap() + .push(format!("before_send:{authorization}")); + let headers = wire + .headers + .into_iter() + .map(|(name, value)| match is_authorization(&name) { + true => (name, "Bearer edited".to_string()), + false => (name, value), + }) + .collect(); + Ok(WireRequest { headers, ..wire }) + } + } + + #[tokio::test] + async fn the_callers_azure_token_is_acquired_before_before_send_which_can_still_replace_it() { + let (base, seen, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await; + let mut request = wire_request("azure_ai/model", &base, json!({})); + request.credentials.api_key = None; + let host = CallerTokenHost { + request: Mutex::new(Some(request)), + trace: Mutex::new(Vec::new()), + }; + + litellm_host::run::run(ocr_machine(ocr_client()), &host) + .await + .unwrap(); + server.await.unwrap(); + + assert_eq!( + *host.trace.lock().unwrap(), + ["project", "token", "before_send:Bearer caller-token"] + ); + assert!( + seen.lock().unwrap()[0] + .to_ascii_lowercase() + .contains("authorization: bearer edited\r\n") + ); + } + + #[tokio::test] + async fn interrupting_an_in_flight_provider_request_closes_its_connection() { + use tokio::io::AsyncReadExt; + + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let base = format!("http://{}", listener.local_addr().unwrap()); + let received = Arc::new(tokio::sync::Notify::new()); + let server_received = received.clone(); + let server = tokio::spawn(async move { + let (mut socket, _) = listener.accept().await.unwrap(); + let mut request = Vec::new(); + let mut buffer = [0u8; 4096]; + while !request.windows(4).any(|window| window == b"\r\n\r\n") { + let read = socket.read(&mut buffer).await.unwrap(); + request.extend_from_slice(&buffer[..read]); + } + server_received.notify_one(); + loop { + if socket.read(&mut buffer).await.unwrap() == 0 { + break; + } + } + }); + let host = LocalOcrHost::new(wire_request("mistral/model", &base, json!({}))); + let mut machine = ocr_machine(ocr_client()); + let mut result = None; + tokio::time::timeout(std::time::Duration::from_secs(2), async { + loop { + tokio::select! { + _ = received.notified() => break, + step = machine.resume(result.take()) => { + result = Some(match step.unwrap() { + MachineStep::Host(HostOp::Route(op)) => HostResult::Route(host.route(op).await.unwrap()), + MachineStep::Host(HostOp::BeforeSend { wire, .. }) => HostResult::BeforeSend(wire), + MachineStep::Host(HostOp::Emit(_)) => HostResult::Emitted, + MachineStep::Complete(_) => panic!("the stalled provider completed"), + }); + } + } + } + }) + .await + .unwrap(); + + let cancelled = OcrError::InvalidRequest("cancelled".into()); + assert!( + machine + .interrupt(HostFailure::Cancelled(cancelled)) + .await + .is_err() + ); + tokio::time::timeout(std::time::Duration::from_secs(1), server) + .await + .expect("the provider connection stayed open after the interrupt") + .unwrap(); + } +} diff --git a/litellm-rust/crates/core/src/audio_transcription/tests.rs b/litellm-rust/crates/core/tests/audio_transcription.rs similarity index 95% rename from litellm-rust/crates/core/src/audio_transcription/tests.rs rename to litellm-rust/crates/core/tests/audio_transcription.rs index 8ccf7a07a0f..aa5aef0149b 100644 --- a/litellm-rust/crates/core/src/audio_transcription/tests.rs +++ b/litellm-rust/crates/core/tests/audio_transcription.rs @@ -4,11 +4,9 @@ use std::{ thread, }; +use litellm_core::audio_transcription::{audio_transcription, types::AudioTranscriptionRequest}; use serde_json::{Map, json}; -use super::audio_transcription; -use crate::audio_transcription::types::AudioTranscriptionRequest; - #[tokio::test] async fn bedrock_request_is_signed_and_contains_audio() { let listener = TcpListener::bind("127.0.0.1:0").expect("listener"); diff --git a/litellm-rust/crates/core/tests/aws_textract_ocr.rs b/litellm-rust/crates/core/tests/aws_textract_ocr.rs deleted file mode 100644 index c536317ad5c..00000000000 --- a/litellm-rust/crates/core/tests/aws_textract_ocr.rs +++ /dev/null @@ -1,193 +0,0 @@ -use std::{collections::BTreeMap, time::SystemTime}; - -use litellm_auth_aws::{Credentials, aws_signature_headers, sign_post}; -use litellm_llms::base_llm::ocr::error::Error; -use serde_json::{Value, json}; -use time::{PrimitiveDateTime, format_description}; - -use crate::ocr::{ - route::LocalOcrHost, - test_support::{ - MockResponse, header, mock_server, perform_ocr_with, request_body, - wire_request_with_document, - }, - types::LiteLLMOcrRequest, -}; - -const ACCESS_KEY_ID: &str = "AKIDEXAMPLE"; -const SECRET_ACCESS_KEY: &str = "wJalrXUtnFEMI/K7MDENG+bPxRfiCYEXAMPLEKEY"; - -fn textract_request(base: &str) -> LiteLLMOcrRequest { - textract_request_for("aws_textract/detect-document-text", base) -} - -fn textract_request_for(model: &str, base: &str) -> LiteLLMOcrRequest { - wire_request_with_document( - model, - &format!("{base}/"), - json!({"type": "image_url", "image_url": "data:image/png;base64,b3JpZ2luYWw="}), - json!({ - "aws_access_key_id": ACCESS_KEY_ID, - "aws_secret_access_key": SECRET_ACCESS_KEY, - "aws_region_name": "eu-west-1" - }), - ) -} - -fn textract_response() -> MockResponse { - MockResponse::json(json!({ - "DocumentMetadata": {"Pages": 1}, - "Blocks": [{"BlockType": "PAGE"}, {"BlockType": "LINE", "Text": "Invoice 12345"}] - })) -} - -/// Recomputes SigV4 over the bytes the server received, at the time the client claimed. -fn expected_authorization(url: &str, raw_request: &str) -> String { - let format = - format_description::parse_borrowed::<2>("[year][month][day]T[hour][minute][second]Z") - .unwrap(); - let signed_at: SystemTime = - PrimitiveDateTime::parse(header(raw_request, "x-amz-date").unwrap(), &format) - .unwrap() - .assume_utc() - .into(); - let headers: BTreeMap = ["content-type", "x-amz-target"] - .into_iter() - .map(|name| { - ( - name.to_string(), - header(raw_request, name).unwrap().to_string(), - ) - }) - .collect(); - let body = raw_request.split_once("\r\n\r\n").unwrap().1; - sign_post( - url, - body.as_bytes(), - &aws_signature_headers(&headers), - "eu-west-1", - "textract", - &Credentials::new(ACCESS_KEY_ID, SECRET_ACCESS_KEY, None, None, "test"), - signed_at, - ) - .unwrap()["Authorization"] - .clone() -} - -#[tokio::test] -async fn the_request_is_signed_for_textract_and_lines_become_the_page() { - let (base, seen, server) = mock_server(vec![textract_response()]).await; - - let response = perform_ocr_with(LocalOcrHost::new(textract_request(&base))) - .await - .unwrap(); - server.await.unwrap(); - - let raw = seen.lock().unwrap()[0].clone(); - assert_eq!( - header(&raw, "x-amz-target"), - Some("Textract.DetectDocumentText") - ); - assert_eq!( - header(&raw, "content-type"), - Some("application/x-amz-json-1.1") - ); - assert_eq!( - request_body(&raw), - json!({"Document": {"Bytes": "b3JpZ2luYWw="}}) - ); - assert_eq!( - header(&raw, "authorization"), - Some(expected_authorization(&format!("{base}/"), &raw).as_str()) - ); - assert_eq!(response.pages[0].markdown, "Invoice 12345"); - assert_eq!(response.usage_info.unwrap().pages_processed, Some(1)); -} - -#[tokio::test] -async fn a_body_rewritten_by_before_send_is_what_gets_signed_and_sent() { - let (base, seen, server) = mock_server(vec![textract_response()]).await; - let host = LocalOcrHost::new(textract_request(&base)).with_before_send(|mut wire, _| { - assert!( - !wire - .headers - .iter() - .any(|(name, _)| name.eq_ignore_ascii_case("authorization")), - "the hook ran after signing" - ); - wire.body["Document"]["Bytes"] = Value::from("cmVkYWN0ZWQ="); - Ok(wire) - }); - - perform_ocr_with(host).await.unwrap(); - server.await.unwrap(); - - let raw = seen.lock().unwrap()[0].clone(); - assert_eq!( - request_body(&raw), - json!({"Document": {"Bytes": "cmVkYWN0ZWQ="}}) - ); - assert_eq!( - header(&raw, "authorization"), - Some(expected_authorization(&format!("{base}/"), &raw).as_str()) - ); -} - -#[tokio::test] -async fn a_multi_page_rejection_reaches_the_caller_with_the_single_page_limit() { - let (base, _, server) = mock_server(vec![MockResponse { - status: 400, - headers: vec![], - body: json!({ - "__type": "UnsupportedDocumentException", - "Message": "Request has unsupported document format" - }), - }]) - .await; - - let error = perform_ocr_with(LocalOcrHost::new(textract_request(&base))) - .await - .unwrap_err(); - server.await.unwrap(); - - let Error::Provider { status, body, .. } = error else { - panic!("expected a provider error, got {error:?}"); - }; - assert_eq!(status, 400); - assert!( - body.contains("multi-page documents are not supported"), - "{body}" - ); -} - -#[tokio::test] -async fn analyze_document_asks_for_layout_and_tables_and_returns_markdown() { - let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ - "DocumentMetadata": {"Pages": 1}, - "Blocks": [ - {"Id": "l1", "BlockType": "LINE", "Text": "Quarterly Report"}, - {"Id": "t", "BlockType": "LAYOUT_TITLE", - "Relationships": [{"Type": "CHILD", "Ids": ["l1"]}]} - ] - }))]) - .await; - let request = textract_request_for("aws_textract/analyze-document", &base); - - let response = perform_ocr_with(LocalOcrHost::new(request)).await.unwrap(); - server.await.unwrap(); - - let raw = seen.lock().unwrap()[0].clone(); - assert_eq!( - header(&raw, "x-amz-target"), - Some("Textract.AnalyzeDocument") - ); - assert_eq!( - request_body(&raw)["FeatureTypes"], - json!(["LAYOUT", "TABLES"]) - ); - assert_eq!( - header(&raw, "authorization"), - Some(expected_authorization(&format!("{base}/"), &raw).as_str()) - ); - assert_eq!(response.pages[0].markdown, "# Quarterly Report"); -} diff --git a/litellm-rust/crates/core/tests/azure_ai_ocr.rs b/litellm-rust/crates/core/tests/azure_ai_ocr.rs deleted file mode 100644 index 1492aaaeb11..00000000000 --- a/litellm-rust/crates/core/tests/azure_ai_ocr.rs +++ /dev/null @@ -1,293 +0,0 @@ -use litellm_llms::base_llm::ocr::error::Error; -use serde_json::{Value, json}; - -use super::test_support::{MockResponse, mock_server, perform_ocr, perform_ocr_with, wire_request}; -use crate::ocr::route::LocalOcrHost; - -#[tokio::test] -async fn facade_executes_azure_mistral_with_prepared_auth() { - let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ - "pages":[{"index":0,"markdown":"hello"}], - "usage_info":{"pages_processed":1} - }))]) - .await; - let mut request = wire_request( - "azure_ai/model", - &base, - json!({"include_image_base64":true}), - ); - request.credentials.api_key = None; - request.transport.extra_headers = vec![( - "Authorization".into(), - "Bearer python-prepared-token".into(), - )]; - - let result = perform_ocr(request).await.unwrap(); - server.await.unwrap(); - assert_eq!(result.pages[0].markdown, "hello"); - let requests = seen.lock().unwrap(); - assert_eq!(requests.len(), 1); - assert!(requests[0].starts_with("POST /providers/mistral/azure/ocr ")); - assert!( - requests[0] - .to_ascii_lowercase() - .contains("authorization: bearer python-prepared-token\r\n") - ); - let body: Value = serde_json::from_str(requests[0].split_once("\r\n\r\n").unwrap().1).unwrap(); - assert_eq!( - body, - json!({ - "model":"model", - "document":{"type":"document_url","document_url":"data:application/pdf;base64,YWJj"}, - "include_image_base64":true - }) - ); -} - -#[tokio::test] -async fn facade_acquires_supplied_entra_token_for_final_request() { - let (base, seen, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await; - let mut request = wire_request( - "azure_ai/model", - &base, - json!({"azure_ad_token":"rust-owned-token"}), - ); - request.credentials.api_key = None; - - perform_ocr(request).await.unwrap(); - server.await.unwrap(); - - let requests = seen.lock().unwrap(); - assert_eq!(requests.len(), 1); - assert!( - requests[0] - .to_ascii_lowercase() - .contains("authorization: bearer rust-owned-token\r\n") - ); -} - -#[tokio::test] -async fn rejects_non_inline_body_after_guardrails() { - let request = wire_request("azure_ai/model", "http://127.0.0.1:1", json!({})); - let host = LocalOcrHost::new(request).with_before_send(|mut wire, _| { - wire.body["document"] = json!({ - "type":"document_url", - "document_url":"https://example.com/not-inline.pdf" - }); - Ok(wire) - }); - let error = perform_ocr_with(host).await.unwrap_err(); - assert!(error.to_string().contains("data URI")); -} - -mod transformation { - use std::sync::{ - Arc, - atomic::{AtomicUsize, Ordering}, - }; - - use litellm_auth::{ - ResolvedCredential, SecretValue, TokenFuture, TokenProvider, TokenProviderHandle, - }; - use rstest::rstest; - use serde_json::json; - - use super::*; - use crate::ocr::{ - test_support::{MockResponse, header, mock_server, perform_ocr}, - types::LiteLLMOcrRequest, - wire::decode_request, - }; - - #[derive(Debug)] - struct CountingToken { - token: fn(usize) -> String, - calls: AtomicUsize, - } - - impl CountingToken { - fn new(token: fn(usize) -> String) -> Arc { - Arc::new(Self { - token, - calls: AtomicUsize::new(0), - }) - } - - fn calls(&self) -> usize { - self.calls.load(Ordering::SeqCst) - } - } - - impl TokenProvider for CountingToken { - fn acquire(&self) -> TokenFuture<'_> { - let call = self.calls.fetch_add(1, Ordering::SeqCst) + 1; - let token = SecretValue::new((self.token)(call)); - Box::pin(async move { - Ok(ResolvedCredential::AccessToken { - token, - expires_on: None, - }) - }) - } - } - - fn numbered_token(call: usize) -> String { - format!("callback-{call}") - } - - fn azure_request( - provider: &Arc, - api_base: Option<&str>, - api_key: Option<&str>, - extra_headers: Value, - optional_params: Value, - ) -> LiteLLMOcrRequest { - let wire = serde_json::from_value(json!({ - "model": "azure_ai/mistral-ocr-latest", - "document": {"type":"document_url","document_url":"data:application/pdf;base64,YWJj"}, - "api_key": api_key, - "api_base": api_base, - "custom_llm_provider": null, - "extra_headers": extra_headers, - "optional_params": optional_params, - "timeout_seconds": 2.0 - })) - .unwrap(); - LiteLLMOcrRequest { - azure_ad_token_provider: Some(TokenProviderHandle::new(provider.clone())), - ..decode_request(wire).unwrap() - } - } - - fn ocr_page() -> MockResponse { - MockResponse::json(json!({"pages":[{"index":0,"markdown":"hello"}]})) - } - - #[tokio::test] - async fn token_provider_result_is_the_bearer_and_is_acquired_for_each_request() { - let provider = CountingToken::new(numbered_token); - let (base, seen, server) = mock_server(vec![ocr_page(), ocr_page()]).await; - - for _ in 0..2 { - perform_ocr(azure_request( - &provider, - Some(&base), - None, - Value::Null, - json!({}), - )) - .await - .unwrap(); - } - server.await.unwrap(); - - assert_eq!(provider.calls(), 2); - let requests = seen.lock().unwrap(); - assert_eq!( - requests - .iter() - .map(|request| header(request, "authorization")) - .collect::>(), - [Some("Bearer callback-1"), Some("Bearer callback-2")] - ); - } - - #[rstest] - #[case::api_key_skips_provider(Some("resource-key"), Value::Null, json!({}), "Bearer resource-key", 0)] - #[case::provider_beats_static_token( - None, - Value::Null, - json!({"azure_ad_token":"static-token"}), - "Bearer callback-1", - 1 - )] - #[case::header_wins_on_the_wire_but_provider_still_runs( - None, - json!({"Authorization":"Bearer override"}), - json!({}), - "Bearer override", - 1 - )] - #[tokio::test] - async fn credential_precedence( - #[case] api_key: Option<&str>, - #[case] extra_headers: Value, - #[case] optional_params: Value, - #[case] expected_authorization: &str, - #[case] expected_calls: usize, - ) { - let provider = CountingToken::new(numbered_token); - let (base, seen, server) = mock_server(vec![ocr_page()]).await; - - perform_ocr(azure_request( - &provider, - Some(&base), - api_key, - extra_headers, - optional_params, - )) - .await - .unwrap(); - server.await.unwrap(); - - assert_eq!(provider.calls(), expected_calls); - let requests = seen.lock().unwrap(); - assert_eq!(requests.len(), 1); - assert_eq!( - header(&requests[0], "authorization"), - Some(expected_authorization) - ); - } - - #[rstest] - #[case::missing_api_base( - false, - json!({}), - numbered_token, - |error: &Error| matches!(error, Error::Auth(litellm_auth::Error::MissingApiBase { - provider: "Azure AI", - environment_variable: "AZURE_AI_API_BASE", - })), - 0 - )] - #[case::unsupported_oidc_reference( - true, - json!({"azure_ad_token":"oidc/assertion","client_id":"client","tenant_id":"tenant"}), - numbered_token, - |error: &Error| matches!(error, Error::Auth(litellm_auth::Error::UnsupportedOidcReference)), - 0 - )] - #[case::empty_provider_token_ignores_static_token( - true, - json!({"azure_ad_token":"static-token"}), - |_| String::new(), - |error: &Error| matches!(error, Error::MissingAzureAiCredentials), - 1 - )] - #[tokio::test] - async fn credential_failures_send_no_provider_request( - #[case] with_api_base: bool, - #[case] optional_params: Value, - #[case] token: fn(usize) -> String, - #[case] expected: fn(&Error) -> bool, - #[case] expected_calls: usize, - ) { - let provider = CountingToken::new(token); - let (base, seen, server) = mock_server(vec![ocr_page()]).await; - - let error = perform_ocr(azure_request( - &provider, - with_api_base.then_some(base.as_str()), - None, - Value::Null, - optional_params, - )) - .await - .unwrap_err(); - server.abort(); - - assert!(expected(&error), "unexpected error: {error:?}"); - assert_eq!(provider.calls(), expected_calls); - assert!(seen.lock().unwrap().is_empty()); - } -} diff --git a/litellm-rust/crates/core/tests/azure_document_intelligence_ocr.rs b/litellm-rust/crates/core/tests/azure_document_intelligence_ocr.rs deleted file mode 100644 index 6dc9bfa5e7e..00000000000 --- a/litellm-rust/crates/core/tests/azure_document_intelligence_ocr.rs +++ /dev/null @@ -1,712 +0,0 @@ -use litellm_host::event::{CallEvent, MachineEvent}; -use litellm_llms::base_llm::ocr::{error::Error, settings::OcrSettings}; -use rstest::rstest; -use serde_json::{Value, json}; - -use super::{ - test_support::{ - MockResponse, mock_server, ocr_client, perform_ocr, perform_ocr_with, wire_request, - }, - wire::{OcrWireRequest, decode_request}, -}; -use crate::ocr::route::LocalOcrHost; - -fn query_value(url: &str, key: &str) -> Option { - url::Url::parse(url) - .unwrap() - .query_pairs() - .find_map(|(name, value)| (name == key).then(|| value.into_owned())) -} - -#[tokio::test] -async fn facade_maps_pages_features_and_url_document() { - let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ - "status":"succeeded", - "analyzeResult":{"pages":[]} - }))]) - .await; - let mut request = wire_request( - "azure_ai/doc-intelligence/prebuilt-read", - &base, - json!({"pages":[2,0,0,1],"features":["keyValuePairs","languages"]}), - ); - request.document = - serde_json::from_value::(json!({ - "type":"document_url", - "document_url":"https://example.com/document.pdf" - })) - .unwrap() - .into(); - - perform_ocr(request).await.unwrap(); - server.await.unwrap(); - let request = &seen.lock().unwrap()[0]; - let target = request.split_whitespace().nth(1).unwrap(); - let url = format!("{base}{target}"); - assert_eq!(query_value(&url, "pages").as_deref(), Some("1,2,3")); - assert_eq!( - query_value(&url, "features").as_deref(), - Some("keyValuePairs,languages") - ); - let body: Value = serde_json::from_str(request.split_once("\r\n\r\n").unwrap().1).unwrap(); - assert_eq!( - body, - json!({"urlSource":"https://example.com/document.pdf"}) - ); -} - -#[rstest] -#[case(json!({"pages":[true]}), Error::Pages("expected only integers or only strings".into()))] -#[case(json!({"pages":[1,"2"]}), Error::Pages("expected only integers or only strings".into()))] -#[case(json!({"pages":[-1]}), Error::Pages("negative page index".into()))] -#[case(json!({"pages":"1&&features=bad"}), Error::Pages("invalid native page range".into()))] -#[case(json!({"features":"languages&pages=1"}), Error::Features)] -#[case(json!({"req_format":"azure"}), Error::RequestFormat)] -#[tokio::test] -async fn rejects_invalid_pages_features_and_format( - #[case] options: Value, - #[case] expected: Error, -) { - let (base, seen, server) = mock_server(vec![MockResponse::json(json!({}))]).await; - let result = decode_request(OcrWireRequest { - model: "azure_ai/doc-intelligence/prebuilt-read".into(), - document: json!({"type":"document_url","document_url":"https://example.com/a.pdf"}), - api_key: Some(litellm_auth::SecretValue::new("key")), - api_base: Some(base), - custom_llm_provider: None, - extra_headers: None, - optional_params: options.as_object().unwrap().clone(), - input_sources: Default::default(), - timeout_seconds: Some(2.0), - }); - let result = match result { - Ok(request) => perform_ocr(request).await, - Err(error) => Err(error), - }; - server.abort(); - let _ = server.await; - assert!( - seen.lock().unwrap().is_empty(), - "sent invalid options: {options}" - ); - let error = result.unwrap_err(); - assert_eq!( - std::mem::discriminant(&error), - std::mem::discriminant(&expected) - ); - assert_eq!(error.http_status_code(), Some(400)); - assert_eq!(error.to_string(), expected.to_string()); -} - -#[rstest] -#[case(json!({}))] -#[case(json!({"req_format":"litellm"}))] -#[tokio::test] -async fn missing_native_fields_keep_page_text_without_retaining_raw_response( - #[case] options: Value, -) { - let operation = json!({ - "status":"succeeded", - "analyzeResult":{"pages":[{"pageNumber":1,"lines":[{"content":"hello"}]}]} - }); - let (base, seen, server) = mock_server(vec![MockResponse::json(operation)]).await; - let response = perform_ocr(wire_request( - "azure_ai/doc-intelligence/prebuilt-read", - &base, - options, - )) - .await - .unwrap(); - server.await.unwrap(); - - assert_eq!(response.pages.len(), 1); - assert_eq!(response.pages[0].index, 0); - assert_eq!(response.pages[0].markdown, "hello"); - assert_eq!(response.provider_native_response, None); - let serialized = response.into_json(); - assert_eq!(serialized.get("content"), Some(&Value::Null)); - assert_eq!(serialized.get("tables"), Some(&Value::Null)); - assert_eq!(serialized.get("keyValuePairs"), Some(&Value::Null)); - let requests = seen.lock().unwrap(); - assert_eq!(requests.len(), 1); - let target = requests[0].split_whitespace().nth(1).unwrap(); - let url = format!("{base}{target}"); - for field in ["pages", "features", "req_format"] { - assert_eq!(query_value(&url, field), None); - } - let body: Value = serde_json::from_str(requests[0].split_once("\r\n\r\n").unwrap().1).unwrap(); - assert_eq!(body, json!({"base64Source":"YWJj"})); -} - -#[tokio::test] -async fn inline_document_decodes_to_base64_source() { - let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ - "status":"succeeded" - }))]) - .await; - let request = wire_request("azure_ai/doc-intelligence/prebuilt-read", &base, json!({})); - - perform_ocr(request).await.unwrap(); - server.await.unwrap(); - let request = &seen.lock().unwrap()[0]; - let body: Value = serde_json::from_str(request.split_once("\r\n\r\n").unwrap().1).unwrap(); - assert_eq!(body, json!({"base64Source":"YWJj"})); -} - -#[tokio::test] -async fn immediate_response_normalizes_pages_and_preserves_native() { - let operation = json!({ - "status":"succeeded", - "operationExtension":42, - "analyzeResult":{ - "content":"A\n\nB", - "tables":[{"cells":[]}], - "keyValuePairs":[{"key":{"content":"A"}}], - "pages":[{ - "pageNumber":"2", - "width":"8.5", - "height":11, - "unit":"inch", - "lines":[{"content":"A"},{"content":null},{"content":"B"}] - }] - } - }); - let (base, _, server) = mock_server(vec![MockResponse::json(operation.clone())]).await; - let result = perform_ocr(wire_request( - "azure_ai/doc-intelligence/prebuilt-read", - &base, - json!({"req_format":"native"}), - )) - .await - .unwrap(); - server.await.unwrap(); - - assert_eq!(result.pages[0].index, 1); - assert_eq!(result.pages[0].markdown, "A\n\nB"); - assert_eq!( - serde_json::to_value(&result.pages[0].dimensions).unwrap(), - json!({"width":816,"height":1056,"dpi":96}) - ); - assert_eq!(result.usage_info.as_ref().unwrap().pages_processed, Some(1)); - let serialized = result.clone().into_json(); - assert_eq!(serialized["content"], "A\n\nB"); - assert_eq!(serialized["tables"], json!([{"cells":[]}])); - assert_eq!( - serialized["keyValuePairs"], - json!([{"key":{"content":"A"}}]) - ); - assert!(serialized.get("key_value_pairs").is_none()); - assert_eq!( - result.provider_native_response.map(Value::Object), - Some(operation) - ); -} - -#[tokio::test] -async fn client_settings_choose_the_api_version_and_the_inch_to_pixel_dpi() { - let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ - "status":"succeeded", - "analyzeResult":{"pages":[{"pageNumber":1,"width":8.5,"height":11,"unit":"inch"}]} - }))]) - .await; - let client = ocr_client().with_settings(OcrSettings { - document_intelligence_api_version: "2099-01-01".into(), - document_intelligence_dpi: 72, - ..OcrSettings::default() - }); - - let result = crate::ocr::client::perform( - &client, - wire_request("azure_ai/doc-intelligence/prebuilt-read", &base, json!({})), - ) - .await - .unwrap(); - server.await.unwrap(); - - let target = seen.lock().unwrap()[0] - .split_whitespace() - .nth(1) - .unwrap() - .to_string(); - assert_eq!( - query_value(&format!("{base}{target}"), "api-version").as_deref(), - Some("2099-01-01") - ); - assert_eq!( - serde_json::to_value(&result.pages[0].dimensions).unwrap(), - json!({"width":612,"height":792,"dpi":72}) - ); -} - -#[tokio::test] -async fn accepted_response_polls_to_success_with_only_credentials() { - let operation = json!({"status":"succeeded","analyzeResult":{"pages":[]}}); - let (base, seen, server) = mock_server(vec![ - MockResponse { - status: 202, - headers: vec![("Operation-Location", "{base}/operation".into())], - body: json!({}), - }, - MockResponse { - status: 200, - headers: vec![("Retry-After", "0".into())], - body: json!({"status":"running"}), - }, - MockResponse::json(operation.clone()), - ]) - .await; - let mut request = wire_request( - "azure_ai/doc-intelligence/prebuilt-read", - &base, - json!({"req_format":"native"}), - ); - request - .transport - .extra_headers - .push(("X-Trace".into(), "initial-only".into())); - - let result = perform_ocr(request).await.unwrap(); - server.await.unwrap(); - assert_eq!( - result.provider_native_response.map(Value::Object), - Some(operation) - ); - let requests = seen.lock().unwrap(); - assert_eq!(requests.len(), 3); - assert!(requests[0].to_ascii_lowercase().contains("x-trace:")); - for poll in &requests[1..] { - assert!(!poll.to_ascii_lowercase().contains("x-trace:")); - assert!( - poll.to_ascii_lowercase() - .contains("ocp-apim-subscription-key: test-key") - ); - } -} - -#[tokio::test] -async fn accepted_response_emits_response_received_before_polling() { - let (base, seen, server) = mock_server(vec![ - MockResponse { - status: 202, - headers: vec![("Operation-Location", "{base}/operation".into())], - body: json!({"submitted": true}), - }, - MockResponse::json(json!({"status":"succeeded"})), - ]) - .await; - let request_count = seen.clone(); - let host = LocalOcrHost::new(wire_request( - "azure_ai/doc-intelligence/prebuilt-read", - &base, - json!({}), - )) - .with_observer(move |event| { - let CallEvent::Machine(MachineEvent::ResponseReceived { raw }) = event else { - return; - }; - match request_count.lock().unwrap().len() { - 1 => assert_eq!(raw.body, r#"{"submitted":true}"#), - 2 => assert!(raw.body.contains("succeeded")), - count => panic!("unexpected callback after {count} requests"), - } - }); - - perform_ocr_with(host).await.unwrap(); - server.await.unwrap(); - assert_eq!(seen.lock().unwrap().len(), 2); -} - -#[tokio::test] -async fn polling_forwards_bearer_credentials() { - let (base, seen, server) = mock_server(vec![ - MockResponse { - status: 202, - headers: vec![("Operation-Location", "{base}/operation".into())], - body: json!({}), - }, - MockResponse::json(json!({"status":"succeeded"})), - ]) - .await; - let mut request = wire_request("azure_ai/doc-intelligence/prebuilt-read", &base, json!({})); - request.credentials.api_key = None; - request.transport.extra_headers = vec![("Authorization".into(), "Bearer token".into())]; - - perform_ocr(request).await.unwrap(); - server.await.unwrap(); - let requests = seen.lock().unwrap(); - assert!( - requests[1] - .to_ascii_lowercase() - .contains("authorization: bearer token") - ); -} - -#[tokio::test] -async fn polling_does_not_follow_redirects() { - let (base, seen, server) = mock_server(vec![ - MockResponse { - status: 202, - headers: vec![("Operation-Location", "{base}/operation".into())], - body: json!({}), - }, - MockResponse { - status: 302, - headers: vec![("Location", "{base}/redirected".into())], - body: json!({}), - }, - MockResponse::json(json!({"status":"succeeded"})), - ]) - .await; - - let error = perform_ocr(wire_request( - "azure_ai/doc-intelligence/prebuilt-read", - &base, - json!({}), - )) - .await - .unwrap_err(); - - assert!(error.to_string().contains("status 302"), "{error}"); - assert_eq!(seen.lock().unwrap().len(), 2); - server.abort(); -} - -#[tokio::test] -async fn polling_rejects_terminal_failure() { - let (base, _, server) = mock_server(vec![ - MockResponse { - status: 202, - headers: vec![("Operation-Location", "{base}/operation".into())], - body: json!({}), - }, - MockResponse::json(json!({"status":"failed"})), - ]) - .await; - - let error = perform_ocr(wire_request( - "azure_ai/doc-intelligence/prebuilt-read", - &base, - json!({}), - )) - .await - .unwrap_err(); - server.await.unwrap(); - assert!(error.to_string().contains("status failed")); -} - -#[tokio::test] -async fn malformed_provider_pages_report_response_paths() { - for (analysis, path) in [ - (json!({"pages":null}), "pages"), - (json!({"pages":[null]}), "pages[0]"), - (json!({"pages":[{"lines":null}]}), "lines"), - (json!({"pages":[{"width":"bad"}]}), "width"), - ] { - let (base, _, server) = mock_server(vec![MockResponse::json(json!({ - "status":"succeeded", - "analyzeResult":analysis - }))]) - .await; - let error = perform_ocr(wire_request( - "azure_ai/doc-intelligence/prebuilt-read", - &base, - json!({}), - )) - .await - .unwrap_err(); - server.await.unwrap(); - assert!(error.to_string().contains(path), "{error}"); - } -} - -#[tokio::test] -async fn rejects_missing_invalid_and_cross_origin_operation_locations() { - for headers in [ - Vec::new(), - vec![("Operation-Location", "/relative".into())], - vec![("Operation-Location", "http://example.com/operation".into())], - vec![( - "Operation-Location", - "http://user:password@127.0.0.1/operation".into(), - )], - ] { - let (base, _, server) = mock_server(vec![MockResponse { - status: 202, - headers, - body: json!({}), - }]) - .await; - let error = perform_ocr(wire_request( - "azure_ai/doc-intelligence/prebuilt-read", - &base, - json!({}), - )) - .await - .unwrap_err(); - server.await.unwrap(); - assert!(error.to_string().contains("operation-location")); - } -} - -#[tokio::test] -async fn polling_deadline_bounds_retry_delay() { - let (base, _, server) = mock_server(vec![ - MockResponse { - status: 202, - headers: vec![("Operation-Location", "{base}/operation".into())], - body: json!({}), - }, - MockResponse { - status: 200, - headers: vec![("Retry-After", "9999".into())], - body: json!({"status":"notStarted"}), - }, - ]) - .await; - let request = wire_request("azure_ai/doc-intelligence/prebuilt-read", &base, json!({})); - let client = ocr_client().with_settings(OcrSettings { - poll_timeout: std::time::Duration::from_millis(100), - ..OcrSettings::default() - }); - - let error = tokio::time::timeout( - std::time::Duration::from_secs(1), - crate::ocr::client::perform(&client, request), - ) - .await - .unwrap() - .unwrap_err(); - server.await.unwrap(); - assert!(error.to_string().contains("timed out")); -} - -#[tokio::test] -async fn model_id_is_encoded_and_dot_segments_are_rejected() { - let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ - "status":"succeeded" - }))]) - .await; - perform_ocr(wire_request( - "azure_ai/doc-intelligence/a ?#é", - &base, - json!({}), - )) - .await - .unwrap(); - server.await.unwrap(); - assert!(seen.lock().unwrap()[0].contains("a%20%3F%23%C3%A9:analyze")); - - for model in [ - "azure_ai/doc-intelligence/.", - "azure_ai/doc-intelligence/..", - ] { - let error = perform_ocr(wire_request(model, "http://127.0.0.1:1", json!({}))) - .await - .unwrap_err(); - assert!(error.to_string().contains("dot segment")); - } -} - -mod transformation { - use std::sync::{Arc, Mutex}; - - use litellm_host::event::{CallEvent, MachineEvent}; - use litellm_llms::base_llm::ocr::transformation::OcrDocument; - use serde_json::{Value, json}; - - use super::*; - use crate::ocr::{ - route::LocalOcrHost, - test_support::{MockResponse, mock_server, perform_ocr, perform_ocr_with, wire_request}, - }; - - #[tokio::test] - async fn facade_maps_pages_features_and_url_document() { - let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ - "status":"succeeded", - "analyzeResult":{"pages":[]} - }))]) - .await; - let mut request = wire_request( - "azure_ai/doc-intelligence/prebuilt-read", - &base, - json!({"pages":[2,0,0,1],"features":["keyValuePairs","languages"], "future_option": {"nested":null}, "extra_body":{"provider_option":false}}), - ); - request.document = serde_json::from_value::(json!({ - "type":"document_url", - "document_url":"https://example.com/document.pdf" - })) - .unwrap() - .into(); - - perform_ocr(request).await.unwrap(); - server.await.unwrap(); - let request = &seen.lock().unwrap()[0]; - let target = request.split_whitespace().nth(1).unwrap(); - let url = format!("{base}{target}"); - assert_eq!(query_value(&url, "pages").as_deref(), Some("1,2,3")); - assert_eq!( - query_value(&url, "features").as_deref(), - Some("keyValuePairs,languages") - ); - let body: Value = serde_json::from_str(request.split_once("\r\n\r\n").unwrap().1).unwrap(); - assert_eq!( - body, - json!({"urlSource":"https://example.com/document.pdf", "future_option":{"nested":null}, "provider_option":false}) - ); - } - - #[tokio::test] - async fn rejects_invalid_pages_features_and_format() { - for options in [ - json!({"pages":[true]}), - json!({"pages":[1,"2"]}), - json!({"pages":[-1]}), - json!({"pages":"1&&features=bad"}), - json!({"features":"languages&pages=1"}), - json!({"req_format":"azure"}), - ] { - let request = wire_request( - "azure_ai/doc-intelligence/prebuilt-read", - "http://127.0.0.1:1", - options.clone(), - ); - let rejected = perform_ocr(request).await.is_err(); - assert!(rejected, "accepted {options}"); - } - } - - #[tokio::test] - async fn immediate_response_normalizes_pages_and_preserves_native() { - let operation = json!({ - "status":"succeeded", - "operationExtension":42, - "analyzeResult":{ - "content":"A\n\nB", - "tables":[{"cells":[]}], - "keyValuePairs":[{"key":{"content":"A"}}], - "pages":[{ - "pageNumber":"2", - "width":"8.5", - "height":11, - "unit":"inch", - "lines":[{"content":"A"},{"content":null},{"content":"B"}] - }] - } - }); - let (base, _, server) = mock_server(vec![MockResponse::json(operation.clone())]).await; - let result = perform_ocr(wire_request( - "azure_ai/doc-intelligence/prebuilt-read", - &base, - json!({"req_format":"native"}), - )) - .await - .unwrap(); - server.await.unwrap(); - - assert_eq!(result.pages[0].index, 1); - assert_eq!(result.pages[0].markdown, "A\n\nB"); - assert_eq!( - serde_json::to_value(&result.pages[0].dimensions).unwrap(), - json!({"width":816,"height":1056,"dpi":96}) - ); - assert_eq!(result.usage_info.as_ref().unwrap().pages_processed, Some(1)); - let serialized = result.clone().into_json(); - assert_eq!(serialized["content"], "A\n\nB"); - assert_eq!(serialized["tables"], json!([{"cells":[]}])); - assert_eq!( - serialized["keyValuePairs"], - json!([{"key":{"content":"A"}}]) - ); - assert!(serialized.get("key_value_pairs").is_none()); - assert_eq!( - result.provider_native_response.as_ref(), - operation.as_object() - ); - } - - #[tokio::test] - async fn accepted_response_polls_to_success_with_only_credentials() { - let operation = json!({"status":"succeeded","analyzeResult":{"pages":[]}}); - let (base, seen, server) = mock_server(vec![ - MockResponse { - status: 202, - headers: vec![("Operation-Location", "{base}/operation".into())], - body: json!({}), - }, - MockResponse { - status: 200, - headers: vec![("Retry-After", "0".into())], - body: json!({"status":"running"}), - }, - MockResponse::json(operation.clone()), - ]) - .await; - let mut request = wire_request( - "azure_ai/doc-intelligence/prebuilt-read", - &base, - json!({"req_format":"native"}), - ); - request - .transport - .extra_headers - .push(("X-Trace".into(), "initial-only".into())); - - let result = perform_ocr(request).await.unwrap(); - server.await.unwrap(); - assert_eq!( - result.provider_native_response.as_ref(), - operation.as_object() - ); - let requests = seen.lock().unwrap(); - assert_eq!(requests.len(), 3); - assert!(requests[0].to_ascii_lowercase().contains("x-trace:")); - for poll in &requests[1..] { - assert!(!poll.to_ascii_lowercase().contains("x-trace:")); - assert!( - poll.to_ascii_lowercase() - .contains("ocp-apim-subscription-key: test-key") - ); - } - } - - #[tokio::test] - async fn accepted_response_emits_response_received_for_submission_and_completed_poll() { - let (base, seen, server) = mock_server(vec![ - MockResponse { - status: 202, - headers: vec![("Operation-Location", "{base}/operation".into())], - body: json!({"submitted": true}), - }, - MockResponse::json(json!({"status":"succeeded"})), - ]) - .await; - let responses_received = Arc::new(Mutex::new(Vec::new())); - let request_count = seen.clone(); - let observed = responses_received.clone(); - let host = LocalOcrHost::new(wire_request( - "azure_ai/doc-intelligence/prebuilt-read", - &base, - json!({}), - )) - .with_observer(move |event| { - if let CallEvent::Machine(MachineEvent::ResponseReceived { raw }) = event { - observed - .lock() - .unwrap() - .push((request_count.lock().unwrap().len(), raw.body.clone())); - } - }); - - perform_ocr_with(host).await.unwrap(); - server.await.unwrap(); - assert_eq!(seen.lock().unwrap().len(), 2); - assert_eq!( - *responses_received.lock().unwrap(), - [ - (1, r#"{"submitted":true}"#.to_string()), - (2, r#"{"status":"succeeded"}"#.to_string()), - ] - ); - } -} diff --git a/litellm-rust/crates/core/tests/cohere_ocr.rs b/litellm-rust/crates/core/tests/cohere_ocr.rs deleted file mode 100644 index 12824f58b1d..00000000000 --- a/litellm-rust/crates/core/tests/cohere_ocr.rs +++ /dev/null @@ -1,136 +0,0 @@ -mod transformation { - use litellm_llms::{ - base_llm::ocr::{ - error::Error, - transformation::{BaseOcrConfig, OcrDocument, OcrResponseFormat}, - }, - cohere::ocr::transformation::*, - }; - use rstest::rstest; - use serde_json::{Value, json}; - - #[tokio::test] - async fn composed_body_preserves_native_document_fields_and_untyped_overrides() { - let request = crate::ocr::test_support::wire_request( - "cohere/parse", - "https://example.com", - json!({ - "output_format":"markdown", "timeout":30, - "extra_body":{ - "output_format": {"future":true}, - "document":{"type":"image_url","image_url":"https://example.com/a.png", - "provider_options":{"nested":[false,0,null]}} - } - }), - ); - let request = request.with_document( - serde_json::from_value(json!({ - "type":"image_url","image_url":"https://example.com/original.png" - })) - .unwrap(), - ); - let request = crate::ocr::prepare::prepare_request_for_test(request); - let http = CohereParseConfig - .prepare_request( - &request, - &crate::ocr::test_support::ocr_client(), - &crate::ocr::test_support::NoHooks, - ) - .await - .unwrap(); - let body: Value = serde_json::from_slice(http.body()).unwrap(); - assert_eq!( - body, - json!({ - "model":"parse", "output_format":{"future":true}, - "document":{"type":"image_url","image_url":"https://example.com/a.png", - "provider_options":{"nested":[false,0,null]}} - }) - ); - } - - #[tokio::test] - async fn explicit_null_options_use_defaults_before_http() { - let request = crate::ocr::test_support::wire_request( - "cohere/parse", - "https://example.com", - json!({"output_format":null,"req_format":null}), - ); - let request = request.with_document( - serde_json::from_value( - json!({"type":"image_url","image_url":"https://example.com/a.png"}), - ) - .unwrap(), - ); - assert_eq!( - request.response_format().unwrap(), - OcrResponseFormat::Litellm - ); - let request = crate::ocr::prepare::prepare_request_for_test(request); - let http = CohereParseConfig - .prepare_request( - &request, - &crate::ocr::test_support::ocr_client(), - &crate::ocr::test_support::NoHooks, - ) - .await - .unwrap(); - let body: Value = serde_json::from_slice(http.body()).unwrap(); - assert_eq!(body["output_format"], "markdown"); - assert!(body.get("req_format").is_none()); - } - - #[rstest] - #[case::cohere("cohere/parse-v5.0", "POST /v2/parse ")] - #[case::azure_ai("azure_ai/Cohere-parse-v5.0", "POST /providers/cohere/v2/parse ")] - #[tokio::test] - async fn route_sends_image_to_its_parse_endpoint_with_the_bearer_key( - #[case] model: &str, - #[case] request_line: &str, - ) { - use crate::ocr::test_support::{MockResponse, header, mock_server, perform_ocr}; - - let (base, seen, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await; - let request = crate::ocr::test_support::wire_request(model, &base, json!({})) - .with_document( - serde_json::from_value::( - json!({"type":"image_url","image_url":"data:image/png;base64,YWJj"}), - ) - .unwrap() - .into(), - ); - - perform_ocr(request).await.unwrap(); - server.await.unwrap(); - - let requests = seen.lock().unwrap(); - assert_eq!(requests.len(), 1); - assert!(requests[0].starts_with(request_line), "{}", requests[0]); - assert_eq!( - header(&requests[0], "authorization"), - Some("Bearer test-key") - ); - } - - #[rstest] - #[tokio::test] - async fn route_rejects_non_image_document_without_a_request( - #[values("cohere/parse-v5.0", "azure_ai/Cohere-parse-v5.0")] model: &str, - ) { - use crate::ocr::test_support::{MockResponse, mock_server, perform_ocr}; - - let (base, seen, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await; - - let error = perform_ocr(crate::ocr::test_support::wire_request( - model, - &base, - json!({}), - )) - .await - .unwrap_err(); - server.abort(); - - assert!(matches!(error, Error::CohereImageOnly), "{error:?}"); - assert!(seen.lock().unwrap().is_empty()); - } -} diff --git a/litellm-rust/crates/core/tests/deepseek_ocr.rs b/litellm-rust/crates/core/tests/deepseek_ocr.rs deleted file mode 100644 index 96e7451769d..00000000000 --- a/litellm-rust/crates/core/tests/deepseek_ocr.rs +++ /dev/null @@ -1,133 +0,0 @@ -use litellm_llms::{ - base_llm::ocr::transformation::{BaseOcrConfig, OcrDocument}, - vertex_ai::ocr::deepseek_transformation::{ - DeepSeekOcrParams, DeepSeekOcrResponse, VertexAIDeepSeekOCRConfig, - normalize_response as transform_ocr_response, - }, -}; -use rstest::rstest; -use serde_json::{Value, json}; - -fn document() -> OcrDocument { - serde_json::from_value(json!({"type":"image_url","image_url":"gs://bucket/a.png"})).unwrap() -} - -#[rstest] -#[case("stream", json!(true))] -#[case("temperature", json!(0.1))] -#[case("max_tokens", json!(1024))] -#[case("top_p", json!(0.9))] -#[case("n", json!(2))] -#[case("stop", json!("done"))] -#[case("stop", json!(["done", "stop"]))] -fn request_mapping_matches_python(#[case] name: &str, #[case] value: Value) { - let params: DeepSeekOcrParams = - serde_json::from_value(json!({name: value.clone(), "ignored": true})).unwrap(); - let result = serde_json::to_value( - VertexAIDeepSeekOCRConfig - .transform_ocr_request("deepseek-ai/deepseek-ocr-maas", document(), ¶ms, &[]) - .unwrap(), - ) - .unwrap(); - assert_eq!(result["model"], "deepseek-ai/deepseek-ocr-maas"); - assert_eq!( - result["messages"][0]["content"][0], - json!({"type":"image_url","image_url":"gs://bucket/a.png"}) - ); - assert_eq!(result[name], value); - assert!(result.get("ignored").is_none()); -} - -#[rstest] -#[case(json!({"type":"image_url","image_url":"data:image/png;base64,AA=="}))] -#[case(json!({"type":"document_url","document_url":"data:application/pdf;base64,AA=="}))] -fn request_maps_both_document_types_to_image_content(#[case] document: Value) { - let source = document - .get("image_url") - .or_else(|| document.get("document_url")) - .unwrap() - .clone(); - let request = VertexAIDeepSeekOCRConfig - .transform_ocr_request( - "deepseek-ai/deepseek-ocr-maas", - serde_json::from_value(document).unwrap(), - &DeepSeekOcrParams::default(), - &[], - ) - .unwrap(); - let result = serde_json::to_value(request).unwrap(); - assert_eq!( - result["messages"][0]["content"][0], - json!({"type":"image_url","image_url":source}) - ); -} - -#[rstest] -#[case(json!("# hello"), "# hello")] -#[case(json!("{broken"), "{broken")] -#[case(json!(" {\"pages\":[]} "), " {\"pages\":[]} ")] -#[case(json!({"pages":[]}), "")] -#[case(json!("[]"), "[]")] -#[case(json!("{\"pages\":[{\"markdown\":\"json text\"}]}"), "json text")] -#[case(json!({"pages":[{"markdown":"object"}]}), "object")] -fn response_codec_handles_text_json_and_objects(#[case] content: Value, #[case] expected: &str) { - let structured = content - .as_object() - .is_some_and(|object| object.contains_key("pages")) - || content - .as_str() - .is_some_and(|text| text.contains("\"pages\"")); - let response: DeepSeekOcrResponse = serde_json::from_value( - json!({"choices":[{"message":{"content":content}}],"usage":{"prompt_tokens":1}}), - ) - .unwrap(); - let result = transform_ocr_response("model", response) - .unwrap() - .into_json(); - assert_eq!(result["pages"][0]["markdown"], expected); - assert_eq!(result["pages"][0]["index"], 0); - if structured { - assert!(result["usage_info"].is_null()); - } else { - assert_eq!(result["usage_info"]["prompt_tokens"], 1); - } -} - -#[test] -fn structured_result_maps_pages_usage_model_and_annotation() { - let response: DeepSeekOcrResponse = serde_json::from_value(json!({ - "choices":[{"message":{"content":{ - "pages":[{"index":2,"markdown":"page","images":[{"id":"one"}],"dimensions":{"width":10}}], - "model":"provider-model", - "usage_info":{"pages_processed":1}, - "document_annotation":{"language":"en"}, - "future":"kept" - }}}] - })) - .unwrap(); - let result = transform_ocr_response("requested", response) - .unwrap() - .into_json(); - assert_eq!(result["pages"][0]["index"], 2); - assert_eq!(result["pages"][0]["images"][0]["id"], "one"); - assert_eq!(result["model"], "provider-model"); - assert_eq!(result["usage_info"]["pages_processed"], 1); - assert_eq!(result["document_annotation"]["language"], "en"); - assert_eq!(result["future"], "kept"); -} - -#[test] -fn response_codec_rejects_missing_empty_and_malformed_content() { - for value in [ - json!({"choices":[{"message":{"content":{}}}]}), - json!({"choices":[]}), - json!({"choices":[{"message":{"content":""}}]}), - json!({"choices":[{"message":{"content":"{\"pages\":[{\"markdown\":42}]}"}}]}), - json!({"choices":[{"message":{"content":{"pages":[{"markdown":42}]}}}]}), - ] { - let result = serde_json::from_value::(value) - .map_err(|_| ()) - .and_then(|response| transform_ocr_response("model", response).map_err(|_| ())); - assert!(result.is_err()); - } -} diff --git a/litellm-rust/crates/core/src/messages/tests.rs b/litellm-rust/crates/core/tests/messages.rs similarity index 79% rename from litellm-rust/crates/core/src/messages/tests.rs rename to litellm-rust/crates/core/tests/messages.rs index ce48752864a..18af8a7d619 100644 --- a/litellm-rust/crates/core/src/messages/tests.rs +++ b/litellm-rust/crates/core/tests/messages.rs @@ -1,7 +1,11 @@ use std::{sync::Arc, time::Duration}; use futures_util::future::BoxFuture; -use litellm_http::request::{has_bearer_auth, has_header}; +use litellm_core::messages::{ + Error, messages, + route::{LocalMessagesHost, MessagesCall, messages_machine}, + types::{MessagesRequest, MessagesShaping}, +}; use litellm_secrets::{SecretValue, source::SecretSource}; use serde_json::{Map, Value, json}; use tokio::{ @@ -9,14 +13,6 @@ use tokio::{ net::{TcpListener, TcpStream}, }; -use super::{ - Error, - common_utils::{messages_provider_config, string_headers, truncate_error_body}, - messages, - route::{LocalMessagesHost, MessagesCall, MessagesOutput, messages_machine}, -}; -use crate::messages::types::{MessagesRequest, MessagesShaping}; - struct RecordingSecrets { values: Vec<(&'static str, String)>, fails: bool, @@ -73,55 +69,6 @@ fn secrets_call() -> MessagesCall { } } -#[tokio::test] -async fn route_reads_the_provider_credential_and_base_from_the_secret_source() { - let listener = TcpListener::bind("127.0.0.1:0").await.expect("binds"); - let addr = listener.local_addr().expect("addr"); - let server = tokio::spawn(async move { - let (mut socket, _) = listener.accept().await.expect("accepts request"); - let request = read_http_request(&mut socket).await; - let response_body = r#"{"id":"msg_1","type":"message","role":"assistant","content":[],"model":"claude-sonnet-4-5","stop_reason":"end_turn","usage":{"input_tokens":1,"output_tokens":1}}"#; - socket - .write_all(write_response(response_body).as_bytes()) - .await - .expect("writes response"); - request - }); - let secrets = Arc::new(RecordingSecrets::new( - vec![ - ("ANTHROPIC_API_KEY", "sk-from-manager".to_string()), - ("ANTHROPIC_BASE_URL", format!("http://{addr}")), - ], - false, - )); - - let output = litellm_host::run::run( - messages_machine(secrets.clone()), - &LocalMessagesHost::new(secrets_call()), - ) - .await - .expect("messages request succeeds"); - - assert!(matches!(output, MessagesOutput::Message(_))); - let request = server.await.expect("server task completes"); - assert!( - request - .to_ascii_lowercase() - .contains("x-api-key: sk-from-manager"), - "{request}" - ); - let requested = secrets.requested.lock().unwrap().clone(); - assert_eq!( - requested, - messages_provider_config("anthropic") - .unwrap() - .secret_names() - .iter() - .map(ToString::to_string) - .collect::>() - ); -} - #[tokio::test] async fn route_surfaces_a_secret_manager_failure_before_the_call() { let Err(error) = litellm_host::run::run( @@ -179,75 +126,6 @@ fn write_response(body: &str) -> String { ) } -#[test] -fn provider_config_resolves_anthropic_and_azure_ai() { - assert!(messages_provider_config("anthropic").is_some()); - assert!(messages_provider_config("azure_ai").is_some()); - assert!(messages_provider_config("openai").is_none()); -} - -#[test] -fn truncate_error_body_caps_long_payloads() { - let body = "x".repeat(400); - let truncated = truncate_error_body(&body); - assert!(truncated.ends_with("... (truncated)")); - let prefix_chars = truncated - .strip_suffix("... (truncated)") - .expect("truncated marker present") - .chars() - .count(); - assert_eq!(prefix_chars, 256); -} - -#[test] -fn string_headers_rejects_non_string_values() { - let headers = json!({"x-count": 3}).as_object().unwrap().clone(); - let err = string_headers(Some(headers)).expect_err("non-string header rejected"); - assert_eq!( - err, - Error::Headers(litellm_http::request::HeaderError { - context: "messages", - name: "x-count".to_string(), - actual: "number", - }) - ); -} - -#[test] -fn has_header_is_case_insensitive() { - let headers = vec![("X-Api-Key".to_string(), "secret".to_string())]; - assert!(has_header(&headers, "x-api-key")); - assert!(!has_header(&headers, "authorization")); -} - -#[test] -fn has_bearer_auth_requires_a_nonempty_bearer_token() { - assert!(has_bearer_auth(&[( - "Authorization".to_string(), - "Bearer tok".to_string() - )])); - assert!(has_bearer_auth(&[( - "authorization".to_string(), - "bearer tok".to_string() - )])); - assert!(!has_bearer_auth(&[( - "authorization".to_string(), - "Bearer ".to_string() - )])); - assert!(!has_bearer_auth(&[( - "authorization".to_string(), - String::new() - )])); - assert!(!has_bearer_auth(&[( - "authorization".to_string(), - "Basic abc".to_string() - )])); - assert!(!has_bearer_auth(&[( - "x-api-key".to_string(), - "sk".to_string() - )])); -} - #[tokio::test] async fn messages_round_trip_builds_azure_request_and_passes_response_through() { let listener = TcpListener::bind("127.0.0.1:0").await.expect("binds"); diff --git a/litellm-rust/crates/core/tests/ocr.rs b/litellm-rust/crates/core/tests/ocr.rs deleted file mode 100644 index 26cd2153e8d..00000000000 --- a/litellm-rust/crates/core/tests/ocr.rs +++ /dev/null @@ -1,1033 +0,0 @@ -use std::sync::{Arc, Mutex}; - -use futures_util::future::BoxFuture; -use litellm_auth_gcp::VertexAuth; -use litellm_host::{ - event::{CallEvent, MachineEvent, WireRequest}, - host::{Host, HostOp, HostResult}, - machine::{HostFailure, Machine, MachineStep}, -}; -use litellm_http::{ - HttpClientPool, HttpSettings, Resolution, - media::{PublicDnsResolver, UrlPolicy}, -}; -use litellm_llms::base_llm::ocr::{ - error::Error as OcrError, - handler::OcrClient, - settings::OcrSettings, - transformation::{ - BaseOcrConfig, LiteLLMOcrResponse, OCR_RESPONSE_MAX_BYTES, OcrTransportConfig, - }, -}; -use litellm_secrets::source::SecretSource; -use rstest::rstest; -use serde_json::{Value, json}; - -use super::{ - test_support::{ - MockResponse, mock_server, ocr_client, perform_ocr, perform_ocr_with, wire_request, - }, - wire::{OcrWireRequest, decode_request}, -}; -use crate::ocr::route::{LocalOcrHost, OcrOp, OcrOpResult, ocr_machine}; - -struct RecordingSecretSource { - names: Arc>>, - values: &'static [(&'static str, &'static str)], - api_base: String, -} - -impl SecretSource for RecordingSecretSource { - fn get_secret_str<'a>( - &'a self, - name: &'a str, - ) -> BoxFuture<'a, Result, litellm_secrets::Error>> { - self.names.lock().unwrap().push(name.to_owned()); - Box::pin(async move { - Ok(match name { - "MISTRAL_AZURE_API_BASE" => Some(self.api_base.clone()), - _ => self - .values - .iter() - .find(|(key, _)| *key == name) - .map(|(_, value)| value.to_string()), - } - .map(litellm_secrets::SecretValue::new)) - }) - } -} - -#[rstest] -#[case::mistral("mistral/model", json!({}))] -#[case::vertex("vertex_ai/mistral-ocr-latest", json!({"vertex_project":"test-project", "vertex_location":"us-central1"}))] -#[tokio::test] -async fn ocr_contract_upstream_error_preserves_status_body_and_headers( - #[case] model: &str, - #[case] options: Value, -) { - let payload = json!({"message": format!("{} END-OF-PROVIDER-BODY", "x".repeat(4096))}); - let expected_body = serde_json::to_string(&payload).unwrap(); - let (base, seen, server) = mock_server(vec![MockResponse { - status: 422, - headers: vec![ - ("Retry-After", "17".into()), - ("X-Request-ID", "request-123".into()), - ("X-Future-Header", "retained".into()), - ], - body: payload, - }]) - .await; - let error = perform_ocr(wire_request(model, &base, options)) - .await - .unwrap_err(); - server.await.unwrap(); - assert_eq!(seen.lock().unwrap().len(), 1); - let OcrError::Provider { - status, - body, - headers, - } = error - else { - panic!("expected provider error, got {error:?}"); - }; - assert_eq!(status, 422); - for (name, value) in [ - ("retry-after", "17"), - ("x-request-id", "request-123"), - ("x-future-header", "retained"), - ] { - assert!( - headers - .iter() - .any(|(key, actual)| key.eq_ignore_ascii_case(name) && actual == value) - ); - } - assert_eq!( - body.len(), - expected_body.len(), - "provider error body was truncated" - ); - assert_eq!(body, expected_body); -} - -#[test] -fn request_boundary_selects_mistral_and_rejects_unknown_providers() { - let request = OcrWireRequest { - model: "mistral/model".into(), - document: json!({"type":"document_url","document_url":"https://example.com/doc.pdf"}), - api_key: Some(litellm_auth::SecretValue::new("key")), - api_base: None, - custom_llm_provider: None, - extra_headers: None, - optional_params: json!({"extract_header":true,"unknown":42}) - .as_object() - .unwrap() - .clone(), - input_sources: Default::default(), - timeout_seconds: None, - }; - assert!(decode_request(request).is_ok()); - assert!( - decode_request(OcrWireRequest { - model: "model".into(), - document: json!({"type":"document_url","document_url":"https://example.com/doc.pdf"}), - api_key: Some(litellm_auth::SecretValue::new("key")), - api_base: None, - custom_llm_provider: Some("unknown".into()), - extra_headers: None, - optional_params: serde_json::Map::new(), - input_sources: Default::default(), - timeout_seconds: None, - }) - .is_err() - ); -} - -#[tokio::test] -async fn facade_executes_direct_mistral_once() { - let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ - "pages":[{"index":0,"markdown":"hello","custom":"preserved"}], - "usage_info":{"pages_processed":1} - }))]) - .await; - let result = perform_ocr(wire_request( - "mistral/model", - &base, - json!({"pages":"0,2-4","extract_header":true,"unknown":"ignored"}), - )) - .await - .unwrap(); - server.await.unwrap(); - assert_eq!(result.pages[0].markdown, "hello"); - assert_eq!(result.pages[0].extra_fields["custom"], "preserved"); - let requests = seen.lock().unwrap(); - assert_eq!(requests.len(), 1); - assert!(requests[0].starts_with("POST /v1/ocr ")); - assert!( - requests[0] - .to_ascii_lowercase() - .contains("authorization: bearer test-key\r\n") - ); - let body: Value = serde_json::from_str(requests[0].split_once("\r\n\r\n").unwrap().1).unwrap(); - assert_eq!( - body, - json!({ - "model":"model", - "document":{"type":"document_url","document_url":"data:application/pdf;base64,YWJj"}, - "pages":"0,2-4", - "extract_header":true, - "unknown":"ignored" - }) - ); -} - -#[tokio::test] -async fn facade_retains_native_response_when_requested() { - let provider_response = json!({ - "pages":[{"index":0,"markdown":"hello"}], - "usage_info":{"pages_processed":1}, - "provider_only":"preserved" - }); - let (base, _, server) = mock_server(vec![MockResponse::json(provider_response.clone())]).await; - let response = perform_ocr(wire_request( - "mistral/model", - &base, - json!({"req_format":"native"}), - )) - .await - .unwrap(); - - server.await.unwrap(); - assert_eq!( - response.provider_native_response.map(Value::Object), - Some(provider_response) - ); -} - -#[rstest] -#[case::plain_key(&[("MISTRAL_API_KEY", "plain")], "plain")] -#[case::azure_key_wins(&[("MISTRAL_AZURE_API_KEY", "azure"), ("MISTRAL_API_KEY", "plain")], "azure")] -#[case::empty_azure_key_falls_through(&[("MISTRAL_AZURE_API_KEY", ""), ("MISTRAL_API_KEY", "plain")], "plain")] -#[tokio::test] -async fn mistral_env_fallbacks_follow_python_through_the_injected_secret_source( - #[case] secrets: &'static [(&'static str, &'static str)], - #[case] expected_key: &str, -) { - let (base, seen, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await; - let names = Arc::new(Mutex::new(Vec::new())); - let client = ocr_client().with_secrets(Arc::new(RecordingSecretSource { - names: names.clone(), - values: secrets, - api_base: base.clone(), - })); - let request = decode_request(OcrWireRequest { - model: "mistral/model".into(), - document: json!({"type":"document_url","document_url":"data:application/pdf;base64,YWJj"}), - api_key: None, - api_base: None, - custom_llm_provider: None, - extra_headers: None, - optional_params: Default::default(), - input_sources: Default::default(), - timeout_seconds: Some(2.0), - }) - .unwrap(); - - crate::ocr::client::perform(&client, request).await.unwrap(); - server.await.unwrap(); - assert_eq!( - *names.lock().unwrap(), - litellm_llms::mistral::ocr::transformation::MistralOcrConfig.secret_names() - ); - assert!(seen.lock().unwrap()[0].contains(&format!("authorization: Bearer {expected_key}"))); -} - -#[tokio::test] -async fn mistral_ocr_resolves_provider_secrets_before_transformation() { - let (base, seen, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await; - let names = Arc::new(Mutex::new(Vec::new())); - let client = ocr_client().with_secrets(Arc::new(RecordingSecretSource { - names: names.clone(), - values: &[("MISTRAL_API_KEY", "source-key")], - api_base: base.clone(), - })); - let request = decode_request(OcrWireRequest { - model: "mistral/mistral-ocr-latest".into(), - document: json!({ - "type":"document_url", - "document_url":"data:application/pdf;base64,YWJj" - }), - api_key: None, - api_base: None, - custom_llm_provider: None, - extra_headers: None, - optional_params: Default::default(), - input_sources: Default::default(), - timeout_seconds: Some(2.0), - }) - .unwrap(); - - crate::ocr::client::perform(&client, request).await.unwrap(); - server.await.unwrap(); - assert_eq!( - *names.lock().unwrap(), - litellm_llms::mistral::ocr::transformation::MistralOcrConfig.secret_names() - ); - assert!(seen.lock().unwrap()[0].contains("authorization: Bearer source-key")); -} - -#[tokio::test] -async fn ocr_client_uses_the_injected_http_pool_configuration() { - let (base, seen, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await; - let settings = HttpSettings { - user_agent: Some("host-owned/1".into()), - ..HttpSettings::default() - }; - let client = OcrClient::new( - &HttpClientPool::new(Arc::new(PublicDnsResolver)), - &Resolution::from(&settings).config, - UrlPolicy::default(), - VertexAuth::default(), - OcrSettings::default(), - Arc::new(litellm_secrets::source::EnvironmentSecrets::default()), - ) - .unwrap(); - crate::ocr::client::perform(&client, wire_request("mistral/model", &base, json!({}))) - .await - .unwrap(); - server.await.unwrap(); - assert!(seen.lock().unwrap()[0].contains("user-agent: host-owned/1")); -} - -fn event_name(event: &CallEvent) -> &'static str { - match event { - CallEvent::Started { .. } => "started", - CallEvent::Machine(MachineEvent::ResponseReceived { .. }) => "response", - CallEvent::Succeeded { .. } => "success", - CallEvent::Failed { .. } => "failure", - } -} - -fn recording_host( - request: crate::ocr::types::LiteLLMOcrRequest, - events: Arc>>, - block: bool, -) -> LocalOcrHost { - let before_send_events = events.clone(); - LocalOcrHost::new(request) - .with_before_send(move |wire, _| { - before_send_events.lock().unwrap().push("before_send"); - if block { - return Err(OcrError::InvalidRequest("blocked".into())); - } - Ok(wire) - }) - .with_observer(move |event| events.lock().unwrap().push(event_name(event))) -} - -#[tokio::test] -async fn lifecycle_sends_headers_returned_by_the_before_send_operation() { - let (base, seen, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await; - let host = LocalOcrHost::new(wire_request("mistral/model", &base, json!({}))).with_before_send( - |mut wire, _| { - wire.headers - .push(("x-core-callback".into(), "edited".into())); - Ok(wire) - }, - ); - - perform_ocr_with(host).await.unwrap(); - server.await.unwrap(); - - assert!(seen.lock().unwrap()[0].contains("x-core-callback: edited")); -} - -#[tokio::test] -async fn before_send_context_names_the_route_and_its_secrets() { - let (base, _, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await; - let observed = Arc::new(Mutex::new(None)); - let captured = observed.clone(); - let host = LocalOcrHost::new(wire_request( - "mistral/model", - &base, - json!({"pages": [0], "req_format": "native"}), - )) - .with_before_send(move |wire, context| { - *captured.lock().unwrap() = Some((wire.clone(), context.clone())); - Ok(wire) - }); - perform_ocr_with(host).await.unwrap(); - server.await.unwrap(); - let (wire, context) = observed.lock().unwrap().take().unwrap(); - assert_eq!(context.custom_llm_provider, "mistral"); - assert_eq!(context.model, "model"); - assert_eq!(wire.body["pages"], json!([0])); - assert!(context.secret_fields.is_empty()); - assert_eq!(context.optional_params["req_format"], "native"); - - let (base, _, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await; - let observed = Arc::new(Mutex::new(None)); - let captured = observed.clone(); - let request = wire_request( - "azure_ai/model", - &base, - json!({"client_secret": "shh", "tenant_id": "t"}), - ); - let request = request.with_document(crate::ocr::types::OcrDocumentInput::Bytes { - bytes: b"abc".as_slice().into(), - file_name: None, - mime_type: Some("application/pdf".into()), - }); - let host = LocalOcrHost::new(request).with_before_send(move |wire, context| { - *captured.lock().unwrap() = Some(context.clone()); - Ok(wire) - }); - perform_ocr_with(host).await.unwrap(); - server.await.unwrap(); - let context = observed.lock().unwrap().take().unwrap(); - assert_eq!(context.secret_fields, ["client_secret"]); -} - -#[tokio::test] -async fn lifecycle_orders_hooks_and_emits_one_success() { - let (base, seen, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await; - let events = Arc::new(Mutex::new(Vec::new())); - let host = recording_host( - wire_request("mistral/model", &base, json!({})), - events.clone(), - false, - ); - perform_ocr_with(host).await.unwrap(); - server.await.unwrap(); - assert_eq!( - *events.lock().unwrap(), - ["started", "before_send", "response", "success"] - ); - assert_eq!(seen.lock().unwrap().len(), 1); -} - -#[tokio::test] -async fn lifecycle_blocking_prevents_execution_and_emits_one_failure() { - let events = Arc::new(Mutex::new(Vec::new())); - let host = recording_host( - wire_request("mistral/model", "http://127.0.0.1:1", json!({})), - events.clone(), - true, - ); - let error = perform_ocr_with(host).await.unwrap_err(); - assert!(matches!(error, OcrError::InvalidRequest(message) if message == "blocked")); - assert_eq!( - *events.lock().unwrap(), - ["started", "before_send", "failure"] - ); -} - -#[tokio::test] -async fn upstream_failure_emits_one_terminal_failure() { - let (base, seen, server) = mock_server(vec![MockResponse { - status: 500, - headers: vec![], - body: json!({"error":"failed"}), - }]) - .await; - let events = Arc::new(Mutex::new(Vec::new())); - let host = recording_host( - wire_request("mistral/model", &base, json!({})), - events.clone(), - false, - ); - assert!(perform_ocr_with(host).await.is_err()); - server.await.unwrap(); - assert_eq!( - *events.lock().unwrap(), - ["started", "before_send", "failure"] - ); - assert_eq!(seen.lock().unwrap().len(), 1); -} - -/// Drives the machine by hand, answering every op through `host` except `before_send`, -/// which `intercept` answers so a test can fail or cancel exactly there. -async fn drive_until( - client: OcrClient, - host: &LocalOcrHost, - mut intercept: impl FnMut(WireRequest) -> Result>, -) -> ( - Result, - Vec<&'static str>, - crate::ocr::route::OcrMachine, -) { - let mut machine = ocr_machine(client); - let mut result = None; - let mut ops = Vec::new(); - let outcome = loop { - let op = match machine.resume(result.take()).await { - Ok(MachineStep::Host(op)) => op, - Ok(MachineStep::Complete(response)) => break Ok(response), - Err(error) => break Err(error), - }; - let answer = match op { - HostOp::Route(op) => { - ops.push(match op { - OcrOp::ProjectRequest => "ProjectRequest", - OcrOp::ReadDocument => "ReadDocument", - OcrOp::AcquireAzureAdToken => "AcquireAzureAdToken", - }); - host.route(op) - .await - .map(HostResult::Route) - .map_err(HostFailure::Error) - } - HostOp::BeforeSend { wire, .. } => { - ops.push("BeforeSend"); - intercept(*wire).map(|wire| HostResult::BeforeSend(Box::new(wire))) - } - HostOp::Emit(event) => { - let event = CallEvent::Machine(event); - ops.push(event_name(&event)); - host.emit(&event) - .await - .map(|()| HostResult::Emitted) - .map_err(HostFailure::Error) - } - }; - match answer { - Ok(answer) => result = Some(answer), - Err(failure) => break machine.interrupt(failure).await, - } - }; - (outcome, ops, machine) -} - -#[tokio::test] -async fn failed_before_send_does_not_replay_or_reach_transport() { - let host = LocalOcrHost::new(wire_request( - "mistral/model", - "http://127.0.0.1:1", - json!({}), - )); - let (outcome, ops, mut machine) = drive_until(ocr_client(), &host, |_| { - Err(HostFailure::Error(OcrError::InvalidRequest( - "before_send failed".into(), - ))) - }) - .await; - assert!( - matches!(outcome, Err(OcrError::InvalidRequest(message)) if message == "before_send failed") - ); - assert_eq!(ops, ["ProjectRequest", "BeforeSend"]); - assert!(machine.resume(None).await.is_err()); -} - -#[tokio::test] -async fn invalid_provider_response_emits_response_received_before_normalization_failure() { - let (base, seen, server) = - mock_server(vec![MockResponse::json(json!({"pages":"invalid"}))]).await; - let responses_received = Arc::new(Mutex::new(Vec::new())); - let observed = responses_received.clone(); - let host = LocalOcrHost::new(wire_request("mistral/model", &base, json!({}))).with_observer( - move |event| { - if let CallEvent::Machine(MachineEvent::ResponseReceived { raw }) = event { - observed.lock().unwrap().push(raw.body.clone()); - } - }, - ); - let error = perform_ocr_with(host).await.unwrap_err(); - server.await.unwrap(); - assert!(matches!(error, OcrError::ResponseField { .. })); - assert_eq!(seen.lock().unwrap().len(), 1); - assert_eq!( - *responses_received.lock().unwrap(), - [r#"{"pages":"invalid"}"#] - ); -} - -#[tokio::test] -async fn direct_native_host_drives_the_same_state_machine() { - let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ - "pages":[{"index":0,"markdown":"native"}] - }))]) - .await; - let host = LocalOcrHost::new(wire_request("mistral/model", &base, json!({}))); - let (outcome, ops, mut machine) = drive_until(ocr_client(), &host, Ok).await; - server.await.unwrap(); - assert_eq!(outcome.unwrap().pages[0].markdown, "native"); - assert_eq!(seen.lock().unwrap().len(), 1); - assert_eq!(ops, ["ProjectRequest", "BeforeSend", "response"]); - assert!(matches!( - machine.resume(None).await, - Err(OcrError::InvalidRequest(_)) - )); -} - -async fn drive_native_file_call( - request: crate::ocr::types::LiteLLMOcrRequest, - content: Result, -) -> (Result, usize) { - let reads = Arc::new(Mutex::new(0)); - let counted = reads.clone(); - let content = Mutex::new(Some(content)); - let host = LocalOcrHost::new(request).with_reader(move || { - *counted.lock().unwrap() += 1; - content.lock().unwrap().take().unwrap() - }); - let outcome = perform_ocr_with(host).await; - let reads = *reads.lock().unwrap(); - (outcome, reads) -} - -#[tokio::test] -async fn host_reader_documents_are_read_once_at_the_core_selected_point_and_encoded() { - let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ - "pages":[{"index":0,"markdown":"file"}] - }))]) - .await; - let request = wire_request("mistral/model", &base, json!({})).with_document( - crate::ocr::types::OcrDocumentInput::HostReader { - mime_type: Some("application/pdf".into()), - }, - ); - let (response, reads) = drive_native_file_call( - request, - Ok(crate::ocr::types::OcrFileContent { - bytes: b"abc".as_slice().into(), - file_name: Some("scan.png".into()), - }), - ) - .await; - server.await.unwrap(); - assert_eq!(response.unwrap().pages[0].markdown, "file"); - assert_eq!(reads, 1); - assert!(seen.lock().unwrap()[0].contains("data:application/pdf;base64,YWJj")); -} - -#[tokio::test] -async fn host_reader_failures_and_empty_files_fail_before_the_provider_is_called() { - let (base, seen, _server) = mock_server(vec![]).await; - let request = wire_request("mistral/model", &base, json!({})); - let failure = OcrError::InvalidRequest("reader exploded".into()); - let (response, reads) = drive_native_file_call( - request.with_document(crate::ocr::types::OcrDocumentInput::HostReader { mime_type: None }), - Err(failure.clone()), - ) - .await; - assert!( - matches!(response.unwrap_err(), OcrError::InvalidRequest(message) if message == "reader exploded") - ); - assert_eq!(reads, 1); - - let request = wire_request("mistral/model", &base, json!({})); - let (response, _) = drive_native_file_call( - request.with_document(crate::ocr::types::OcrDocumentInput::HostReader { mime_type: None }), - Ok(crate::ocr::types::OcrFileContent { - bytes: Default::default(), - file_name: None, - }), - ) - .await; - assert!(matches!(response.unwrap_err(), OcrError::EmptyFile)); - assert!(seen.lock().unwrap().is_empty()); -} - -#[tokio::test] -async fn path_documents_are_read_by_core_without_a_host_operation() { - let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ - "pages":[{"index":0,"markdown":"path"}] - }))]) - .await; - let dir = std::env::temp_dir().join(format!("litellm-ocr-{}", rand::random::())); - std::fs::create_dir_all(&dir).unwrap(); - let path = dir.join("scan.png"); - std::fs::write(&path, b"abc").unwrap(); - let request = wire_request("mistral/model", &base, json!({})).with_document( - crate::ocr::types::OcrDocumentInput::Path { - path: path.clone(), - mime_type: None, - }, - ); - let (response, reads) = - drive_native_file_call(request, Err(OcrError::InvalidRequest("unused".into()))).await; - server.await.unwrap(); - std::fs::remove_dir_all(&dir).unwrap(); - assert_eq!(response.unwrap().pages[0].markdown, "path"); - assert_eq!(reads, 0); - assert!(seen.lock().unwrap()[0].contains("data:image/png;base64,YWJj")); - - let (base, seen, _server) = mock_server(vec![]).await; - let request = wire_request("mistral/model", &base, json!({})); - let (response, _) = drive_native_file_call( - request.with_document(crate::ocr::types::OcrDocumentInput::Path { - path: path.clone(), - mime_type: None, - }), - Err(OcrError::InvalidRequest("unused".into())), - ) - .await; - assert!(matches!( - response.unwrap_err(), - OcrError::FileRead { path: failed, source } if failed == path && source.kind() == std::io::ErrorKind::NotFound - )); - assert!(seen.lock().unwrap().is_empty()); -} - -#[tokio::test] -async fn cancellation_at_before_send_prevents_execution_and_further_resumption() { - let host = LocalOcrHost::new(wire_request( - "mistral/model", - "http://127.0.0.1:1", - json!({}), - )); - let (outcome, ops, mut machine) = drive_until(ocr_client(), &host, |_| { - Err(HostFailure::Cancelled(OcrError::InvalidRequest( - "cancelled".into(), - ))) - }) - .await; - assert!(matches!(outcome, Err(OcrError::InvalidRequest(message)) if message == "cancelled")); - assert_eq!(ops, ["ProjectRequest", "BeforeSend"]); - assert!(machine.resume(Some(HostResult::Emitted)).await.is_err()); -} - -#[tokio::test] -async fn missing_host_result_preserves_pending_operation() { - let request = wire_request("mistral/model", "http://127.0.0.1:1", json!({})); - let mut machine = ocr_machine(ocr_client()); - assert!(matches!( - machine.resume(None).await.unwrap(), - MachineStep::Host(HostOp::Route(OcrOp::ProjectRequest)) - )); - assert!(machine.resume(None).await.is_err()); - assert!(matches!( - machine - .resume(Some(HostResult::Route(OcrOpResult::Request { - request: Box::new(request), - caller_token: false, - }))) - .await - .unwrap(), - MachineStep::Host(HostOp::BeforeSend { .. }) - )); -} - -async fn read_bounded_response(response: Vec, limit: usize) -> Result { - use tokio::io::{AsyncReadExt, AsyncWriteExt}; - - let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); - let address = listener.local_addr().unwrap(); - let server = tokio::spawn(async move { - let (mut socket, _) = listener.accept().await.unwrap(); - let mut request = [0; 4096]; - assert!(socket.read(&mut request).await.unwrap() > 0); - socket.write_all(&response).await.unwrap(); - std::future::pending::<()>().await; - }); - let response = reqwest::Client::new() - .get(format!("http://{address}")) - .send() - .await - .unwrap(); - let result = tokio::time::timeout( - std::time::Duration::from_secs(2), - litellm_llms::base_llm::ocr::handler::read_response_bytes(response, limit), - ) - .await; - server.abort(); - let _ = server.await; - result.expect("bounded reads must finish without waiting for the rest of an oversized body") -} - -#[tokio::test] -async fn response_limit_accepts_exact_size_and_rejects_declared_and_chunked_overflow() { - use litellm_llms::base_llm::ocr::error::Error; - - for response in [ - "HTTP/1.1 200 OK\r\nContent-Length: 8\r\n\r\nabcdefgh", - "HTTP/1.1 200 OK\r\nTransfer-Encoding: chunked\r\n\r\n4\r\nabcd\r\n4\r\nefgh\r\n0\r\n\r\n", - ] { - assert_eq!( - read_bounded_response(response.as_bytes().to_vec(), 8) - .await - .unwrap(), - "abcdefgh" - ); - } - for response in [ - "HTTP/1.1 200 OK\r\nContent-Length: 9\r\n\r\n", - "HTTP/1.1 200 OK\r\nTransfer-Encoding: chunked\r\n\r\n4\r\nabcd\r\n5\r\nefghi\r\n", - ] { - assert!(matches!( - read_bounded_response(response.as_bytes().to_vec(), 8).await, - Err(Error::TooLarge { limit: 8 }) - )); - } -} - -#[rstest] -#[case::declared("Content-Length: 1000000")] -#[case::chunked("Transfer-Encoding: chunked")] -#[tokio::test] -async fn oversized_error_retains_http_status_and_bounded_diagnostics_without_draining( - #[case] headers: &str, -) { - let prefix = "x".repeat(4096); - let body = if headers.starts_with("Transfer") { - format!("{:x}\r\n{prefix}\r\n", prefix.len()) - } else { - prefix.clone() - }; - let response = format!("HTTP/1.1 429 Too Many Requests\r\n{headers}\r\n\r\n{body}"); - let error = read_bounded_response(response.into_bytes(), prefix.len()) - .await - .unwrap_err(); - match error { - OcrError::Transport(litellm_http::transport::Error::Http { status, body }) => { - assert_eq!(status, 429); - assert_eq!(body, prefix); - } - error => panic!("unexpected error: {error}"), - } -} - -#[test] -fn response_limit_is_validated_and_not_forwarded_to_the_provider() { - let request = wire_request( - "mistral/model", - "http://localhost", - json!({"max_response_bytes": 123}), - ); - assert_eq!(request.transport.max_response_bytes, 123); - assert!(!request.optional_params.contains_key("max_response_bytes")); - for value in [ - json!(0), - json!(-1), - json!(true), - json!("123"), - json!(1.5), - json!(OCR_RESPONSE_MAX_BYTES + 1), - Value::Null, - ] { - let wire = serde_json::from_value(json!({ - "model": "mistral/model", "document": {"type": "document_url", "document_url": "data:application/pdf;base64,YWJj"}, - "optional_params": {"max_response_bytes": value} - })).unwrap(); - let Err(error) = decode_request(wire) else { - panic!("invalid response limit accepted") - }; - assert!(error.to_string().contains("max_response_bytes")); - } -} - -#[derive(Debug)] -struct PendingToken { - entered: Arc, - dropped: Arc, -} - -struct TokenFutureDrop(Arc); - -impl Drop for TokenFutureDrop { - fn drop(&mut self) { - self.0.store(true, std::sync::atomic::Ordering::SeqCst); - } -} - -impl litellm_auth::TokenProvider for PendingToken { - fn acquire(&self) -> litellm_auth::TokenFuture<'_> { - Box::pin(async move { - let _guard = TokenFutureDrop(self.dropped.clone()); - self.entered.notify_one(); - std::future::pending().await - }) - } -} - -#[tokio::test] -async fn interrupt_drops_provider_captures_before_returning() { - use std::sync::atomic::{AtomicBool, Ordering}; - - let entered = Arc::new(tokio::sync::Notify::new()); - let dropped = Arc::new(AtomicBool::new(false)); - let request = wire_request("azure_ai/mistral-ocr", "https://example.invalid", json!({})); - let request = crate::ocr::types::LiteLLMOcrRequest { - transport: OcrTransportConfig { - extra_headers: vec![("authorization".into(), "Bearer test-key".into())], - ..request.transport - }, - azure_ad_token_provider: Some(litellm_auth::TokenProviderHandle::new(Arc::new( - PendingToken { - entered: entered.clone(), - dropped: dropped.clone(), - }, - ))), - ..request - }; - let host = LocalOcrHost::new(request); - let mut machine = ocr_machine(ocr_client()); - let mut result = None; - tokio::time::timeout(std::time::Duration::from_secs(2), async { - loop { - tokio::select! { - _ = entered.notified() => break, - step = machine.resume(result.take()) => { - result = Some(match step.unwrap() { - MachineStep::Host(HostOp::Route(op)) => HostResult::Route(host.route(op).await.unwrap()), - MachineStep::Host(HostOp::BeforeSend { wire, .. }) => { - HostResult::BeforeSend(wire) - } - MachineStep::Host(HostOp::Emit(_)) => HostResult::Emitted, - MachineStep::Complete(_) => panic!("pending provider completed"), - }); - } - } - } - }) - .await - .unwrap(); - assert!(!dropped.load(Ordering::SeqCst)); - let selected = OcrError::InvalidRequest("cancelled".into()); - let acknowledgement = machine.interrupt(HostFailure::Cancelled(selected.clone())); - assert!( - dropped.load(Ordering::SeqCst), - "interrupt returned while provider captures were still alive" - ); - assert!( - matches!(acknowledgement.await, Err(OcrError::InvalidRequest(message)) if message == "cancelled") - ); -} - -struct CallerTokenHost { - request: Mutex>, - trace: Mutex>, -} - -impl Host for CallerTokenHost { - async fn route(&self, op: OcrOp) -> Result { - match op { - OcrOp::ProjectRequest => { - self.trace.lock().unwrap().push("project".into()); - Ok(OcrOpResult::Request { - request: Box::new(self.request.lock().unwrap().take().unwrap()), - caller_token: true, - }) - } - OcrOp::AcquireAzureAdToken => { - self.trace.lock().unwrap().push("token".into()); - Ok(OcrOpResult::AzureAdToken( - litellm_auth::ResolvedCredential::Static(litellm_auth::SecretValue::new( - "caller-token", - )), - )) - } - OcrOp::ReadDocument => Err(OcrError::InvalidRequest("no reader".into())), - } - } - - async fn before_send( - &self, - wire: WireRequest, - _: &litellm_host::event::RequestContext, - ) -> Result { - let is_authorization = |name: &str| name.eq_ignore_ascii_case("authorization"); - let authorization = wire - .headers - .iter() - .find(|(name, _)| is_authorization(name)) - .map(|(_, value)| value.clone()) - .unwrap_or_default(); - self.trace - .lock() - .unwrap() - .push(format!("before_send:{authorization}")); - let headers = wire - .headers - .into_iter() - .map(|(name, value)| match is_authorization(&name) { - true => (name, "Bearer edited".to_string()), - false => (name, value), - }) - .collect(); - Ok(WireRequest { headers, ..wire }) - } -} - -#[tokio::test] -async fn the_callers_azure_token_is_acquired_before_before_send_which_can_still_replace_it() { - let (base, seen, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await; - let mut request = wire_request("azure_ai/model", &base, json!({})); - request.credentials.api_key = None; - let host = CallerTokenHost { - request: Mutex::new(Some(request)), - trace: Mutex::new(Vec::new()), - }; - - litellm_host::run::run(ocr_machine(ocr_client()), &host) - .await - .unwrap(); - server.await.unwrap(); - - assert_eq!( - *host.trace.lock().unwrap(), - ["project", "token", "before_send:Bearer caller-token"] - ); - assert!( - seen.lock().unwrap()[0] - .to_ascii_lowercase() - .contains("authorization: bearer edited\r\n") - ); -} - -#[tokio::test] -async fn interrupting_an_in_flight_provider_request_closes_its_connection() { - use tokio::io::AsyncReadExt; - - let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); - let base = format!("http://{}", listener.local_addr().unwrap()); - let received = Arc::new(tokio::sync::Notify::new()); - let server_received = received.clone(); - let server = tokio::spawn(async move { - let (mut socket, _) = listener.accept().await.unwrap(); - let mut request = Vec::new(); - let mut buffer = [0u8; 4096]; - while !request.windows(4).any(|window| window == b"\r\n\r\n") { - let read = socket.read(&mut buffer).await.unwrap(); - request.extend_from_slice(&buffer[..read]); - } - server_received.notify_one(); - loop { - if socket.read(&mut buffer).await.unwrap() == 0 { - break; - } - } - }); - let host = LocalOcrHost::new(wire_request("mistral/model", &base, json!({}))); - let mut machine = ocr_machine(ocr_client()); - let mut result = None; - tokio::time::timeout(std::time::Duration::from_secs(2), async { - loop { - tokio::select! { - _ = received.notified() => break, - step = machine.resume(result.take()) => { - result = Some(match step.unwrap() { - MachineStep::Host(HostOp::Route(op)) => HostResult::Route(host.route(op).await.unwrap()), - MachineStep::Host(HostOp::BeforeSend { wire, .. }) => HostResult::BeforeSend(wire), - MachineStep::Host(HostOp::Emit(_)) => HostResult::Emitted, - MachineStep::Complete(_) => panic!("the stalled provider completed"), - }); - } - } - } - }) - .await - .unwrap(); - - let cancelled = OcrError::InvalidRequest("cancelled".into()); - assert!( - machine - .interrupt(HostFailure::Cancelled(cancelled)) - .await - .is_err() - ); - tokio::time::timeout(std::time::Duration::from_secs(1), server) - .await - .expect("the provider connection stayed open after the interrupt") - .unwrap(); -} diff --git a/litellm-rust/crates/core/tests/ocr/document.rs b/litellm-rust/crates/core/tests/ocr/document.rs deleted file mode 100644 index 855548dc6bf..00000000000 --- a/litellm-rust/crates/core/tests/ocr/document.rs +++ /dev/null @@ -1,152 +0,0 @@ -use litellm_host::event::WireRequest; -use litellm_llms::base_llm::ocr::error::Error; -use rstest::rstest; -use serde_json::{Value, json}; - -use super::test_support::{ - MockResponse, SERVED_DOCUMENT, document_server, mock_server, perform_ocr_with, request_body, - wire_request_with_document, -}; -use crate::ocr::route::LocalOcrHost; - -#[derive(Clone, Copy, Debug)] -enum Route { - Mistral, - AzureAi, - VertexMistral, - AzureCohereParse, - Cohere, -} - -impl Route { - fn model(self) -> &'static str { - match self { - Self::Mistral => "mistral/model", - Self::AzureAi => "azure_ai/model", - Self::VertexMistral => "vertex_ai/mistral-ocr-maas", - Self::AzureCohereParse => "azure_ai/cohere-parse", - Self::Cohere => "cohere/model", - } - } - - fn document_type(self) -> &'static str { - match self { - Self::Mistral | Self::AzureAi | Self::VertexMistral => "document_url", - Self::AzureCohereParse | Self::Cohere => "image_url", - } - } - - fn options(self) -> Value { - match self { - Self::Mistral | Self::AzureAi => json!({"pages": [0]}), - Self::VertexMistral => json!({"pages": [0], "vertex_project": "project-1"}), - Self::AzureCohereParse | Self::Cohere => json!({"output_format": "markdown"}), - } - } -} - -/// What the host does to the wire request in `before_send`. -#[derive(Clone, Copy, Debug)] -enum Host { - Detached, - ReplacesDocument, -} - -const REPLACED_DOCUMENT: &str = "data:image/png;base64,cmVwbGFjZWQ="; - -impl Host { - fn before_send(self, wire: WireRequest) -> WireRequest { - let Value::Object(fields) = wire.body else { - return wire; - }; - let body = fields - .into_iter() - .map(|(name, value)| match self { - Self::Detached => (name, value), - Self::ReplacesDocument if name == "document" => { - let document_type = value["type"].clone(); - let key = document_type.as_str().unwrap_or_default().to_string(); - (name, json!({"type": document_type, key: REPLACED_DOCUMENT})) - } - Self::ReplacesDocument => (name, value), - }) - .collect(); - WireRequest { - body: Value::Object(body), - ..wire - } - } -} - -struct Sent { - result: Result<(), Error>, - provider_body: Option, -} - -async fn send(route: Route, host: Host, document_base: &str) -> Sent { - let (base, seen, provider) = mock_server(vec![MockResponse::json(json!({"pages": []}))]).await; - let document_type = route.document_type(); - let document = - json!({"type": document_type, document_type: format!("{document_base}/scan.png")}); - let request = wire_request_with_document(route.model(), &base, document, route.options()); - let local = - LocalOcrHost::new(request).with_before_send(move |wire, _| Ok(host.before_send(wire))); - let result = perform_ocr_with(local).await.map(|_| ()); - match result { - Ok(()) => provider.await.unwrap(), - Err(_) => provider.abort(), - } - let provider_body = seen - .lock() - .unwrap() - .first() - .map(|request| request_body(request)); - Sent { - result, - provider_body, - } -} - -fn served_document_uri() -> String { - use base64::Engine; - format!( - "data:image/png;base64,{}", - base64::engine::general_purpose::STANDARD.encode(SERVED_DOCUMENT) - ) -} - -#[rstest] -#[case::azure_ai(Route::AzureAi)] -#[case::vertex_mistral(Route::VertexMistral)] -#[case::azure_cohere_parse(Route::AzureCohereParse)] -#[tokio::test] -async fn inlining_routes_send_the_downloaded_document(#[case] route: Route) { - let (document_base, _documents) = document_server().await; - let sent = send(route, Host::Detached, &document_base).await; - sent.result.unwrap(); - assert_eq!( - sent.provider_body.unwrap()["document"][route.document_type()], - json!(served_document_uri()) - ); -} - -#[rstest] -#[tokio::test] -async fn document_replaced_by_the_host_reaches_the_provider( - #[values( - Route::Mistral, - Route::AzureAi, - Route::VertexMistral, - Route::AzureCohereParse, - Route::Cohere - )] - route: Route, -) { - let (document_base, _documents) = document_server().await; - let sent = send(route, Host::ReplacesDocument, &document_base).await; - sent.result.unwrap(); - assert_eq!( - sent.provider_body.unwrap()["document"][route.document_type()], - json!(REPLACED_DOCUMENT) - ); -} diff --git a/litellm-rust/crates/core/tests/ocr/support.rs b/litellm-rust/crates/core/tests/ocr/support.rs deleted file mode 100644 index 974fa3d6655..00000000000 --- a/litellm-rust/crates/core/tests/ocr/support.rs +++ /dev/null @@ -1,203 +0,0 @@ -use std::sync::{Arc, Mutex}; - -use futures_util::future::BoxFuture; -use litellm_host::event::WireRequest; -use litellm_llms::base_llm::ocr::{ - error::Error, - handler::{CallHooks, OcrClient}, - transformation::LiteLLMOcrResponse, -}; -use serde_json::{Value, json}; -use tokio::{ - io::{AsyncReadExt, AsyncWriteExt}, - net::TcpListener, -}; - -use crate::ocr::{ - route::{LocalOcrHost, ocr_machine}, - types::LiteLLMOcrRequest, - wire::{OcrWireRequest, decode_request}, -}; - -/// Stands in for a host with no hooks registered: the wire request goes out unchanged -/// and response events go nowhere. -pub(crate) struct NoHooks; - -impl CallHooks for NoHooks { - fn before_send(&self, wire: WireRequest) -> BoxFuture<'_, Result> { - Box::pin(async move { Ok(wire) }) - } - - fn response_received<'a>(&'a self, _body: &'a [u8]) -> BoxFuture<'a, Result<(), Error>> { - Box::pin(async { Ok(()) }) - } -} - -pub(crate) fn ocr_client() -> OcrClient { - let document_http = reqwest::Client::builder() - .redirect(reqwest::redirect::Policy::none()) - .build() - .expect("test document client builds"); - OcrClient::for_test(reqwest::Client::new(), document_http) -} - -pub(crate) async fn perform_ocr(request: LiteLLMOcrRequest) -> Result { - crate::ocr::client::perform(&ocr_client(), request).await -} - -pub(crate) async fn perform_ocr_with(host: LocalOcrHost) -> Result { - litellm_host::run::run(ocr_machine(ocr_client()), &host).await -} - -pub(crate) fn wire_request(model: &str, base: &str, options: Value) -> LiteLLMOcrRequest { - wire_request_with_document( - model, - base, - json!({"type":"document_url","document_url":"data:application/pdf;base64,YWJj"}), - options, - ) -} - -pub(crate) fn wire_request_with_document( - model: &str, - base: &str, - document: Value, - options: Value, -) -> LiteLLMOcrRequest { - decode_request(OcrWireRequest { - model: model.into(), - document, - api_key: Some(litellm_auth::SecretValue::new("test-key")), - api_base: Some(base.into()), - custom_llm_provider: None, - extra_headers: None, - optional_params: options.as_object().unwrap().clone(), - input_sources: Default::default(), - timeout_seconds: Some(2.0), - }) - .unwrap() -} - -pub(crate) fn resolved_request( - request: LiteLLMOcrRequest, -) -> crate::ocr::types::ResolvedOcrRequest { - request - .map_document(crate::ocr::document::prepare_document) - .unwrap() -} - -pub(crate) fn with_source(request: LiteLLMOcrRequest, source: &str) -> LiteLLMOcrRequest { - let request = resolved_request(request); - let document = request.document.clone().with_source(source.into()); - request.with_document(document.into()) -} - -pub(crate) fn request_body(request: &str) -> Value { - serde_json::from_str(request.split_once("\r\n\r\n").unwrap().1).unwrap() -} - -pub(crate) const SERVED_DOCUMENT: &[u8] = b"\x89PNG served document"; - -/// Serves [`SERVED_DOCUMENT`] as `image/png` to every connection until aborted. -pub(crate) async fn document_server() -> (String, tokio::task::JoinHandle<()>) { - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let base = format!("http://{}", listener.local_addr().unwrap()); - let task = tokio::spawn(async move { - loop { - let (mut socket, _) = listener.accept().await.unwrap(); - let mut buffer = [0u8; 4096]; - let _ = socket.read(&mut buffer).await.unwrap(); - let head = format!( - "HTTP/1.1 200 OK\r\nContent-Type: image/png\r\nContent-Length: {}\r\nConnection: close\r\n\r\n", - SERVED_DOCUMENT.len() - ); - socket.write_all(head.as_bytes()).await.unwrap(); - socket.write_all(SERVED_DOCUMENT).await.unwrap(); - } - }); - (base, task) -} - -pub(crate) struct MockResponse { - pub status: u16, - pub headers: Vec<(&'static str, String)>, - pub body: Value, -} - -impl MockResponse { - pub fn json(body: Value) -> Self { - Self { - status: 200, - headers: vec![], - body, - } - } -} - -pub(crate) async fn mock_server( - responses: Vec, -) -> (String, Arc>>, tokio::task::JoinHandle<()>) { - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let base = format!("http://{}", listener.local_addr().unwrap()); - let requests = Arc::new(Mutex::new(Vec::new())); - let seen = requests.clone(); - let server_base = base.clone(); - let task = tokio::spawn(async move { - for response in responses { - let (mut socket, _) = listener.accept().await.unwrap(); - let mut bytes = Vec::new(); - let mut buffer = [0u8; 4096]; - let header_end = loop { - let n = socket.read(&mut buffer).await.unwrap(); - assert!(n > 0); - bytes.extend_from_slice(&buffer[..n]); - if let Some(index) = bytes.windows(4).position(|s| s == b"\r\n\r\n") { - break index + 4; - } - }; - let length = String::from_utf8_lossy(&bytes[..header_end]) - .lines() - .find_map(|line| { - let (name, value) = line.split_once(':')?; - name.eq_ignore_ascii_case("content-length") - .then(|| value.trim().parse::().unwrap()) - }) - .unwrap_or(0); - while bytes.len() < header_end + length { - let n = socket.read(&mut buffer).await.unwrap(); - assert!(n > 0); - bytes.extend_from_slice(&buffer[..n]); - } - seen.lock() - .unwrap() - .push(String::from_utf8_lossy(&bytes).into_owned()); - let body = serde_json::to_vec(&response.body).unwrap(); - let headers = response - .headers - .into_iter() - .map(|(name, value)| { - format!("{name}: {}\r\n", value.replace("{base}", &server_base)) - }) - .collect::(); - let head = format!( - "HTTP/1.1 {} OK\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n{}\r\n", - response.status, - body.len(), - headers - ); - socket.write_all(head.as_bytes()).await.unwrap(); - socket.write_all(&body).await.unwrap(); - } - }); - (base, requests, task) -} - -pub(crate) fn header<'a>(request: &'a str, name: &str) -> Option<&'a str> { - request - .lines() - .take_while(|line| !line.is_empty()) - .find_map(|line| { - let (key, value) = line.split_once(':')?; - key.eq_ignore_ascii_case(name).then(|| value.trim()) - }) -} diff --git a/litellm-rust/crates/core/tests/reducto_ocr.rs b/litellm-rust/crates/core/tests/reducto_ocr.rs deleted file mode 100644 index 83e7754122b..00000000000 --- a/litellm-rust/crates/core/tests/reducto_ocr.rs +++ /dev/null @@ -1,584 +0,0 @@ -use litellm_host::event::{CallEvent, MachineEvent, WireRequest}; -use litellm_llms::base_llm::ocr::{error::Error, transformation::OcrDocument}; -use rstest::rstest; -use serde_json::{Value, json}; - -use super::test_support::{MockResponse, mock_server, perform_ocr, perform_ocr_with, wire_request}; -use crate::ocr::route::LocalOcrHost; - -fn request_body(request: &str) -> Value { - serde_json::from_str(request.split_once("\r\n\r\n").unwrap().1).unwrap() -} - -#[rstest] -#[case( - "reducto/parse-v3", - json!({ - "formatting":{"table_output_format":"html"}, - "retrieval":{"chunk_mode":"section"}, - "settings":{"ocr_system":"standard"}, - "future_ocr_option":true, - "extra_body":{"provider_option":"value"} - }), - "reducto://already.pdf", - json!({ - "input":"reducto://already.pdf", - "formatting":{"table_output_format":"html"}, - "retrieval":{"chunk_mode":"section"}, - "settings":{"ocr_system":"standard"}, - "future_ocr_option":true, - "provider_option":"value" - }) -)] -#[case( - "reducto/parse-legacy", - json!({ - "enhance":{"agentic":[{"type":"table"}]}, - "future_ocr_option":true, - "extra_body":{"provider_option":"value"} - }), - "reducto://legacy.pdf", - json!({ - "document_url":"reducto://legacy.pdf", - "options":{"enhance":{"agentic":[{"type":"table"}]}}, - "future_ocr_option":true, - "provider_option":"value" - }) -)] -#[tokio::test] -async fn request_mapping_matches_python( - #[case] model: &str, - #[case] options: Value, - #[case] source: &str, - #[case] expected: Value, -) { - let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ - "result":{"chunks":[]} - }))]) - .await; - let request = super::test_support::with_source(wire_request(model, &base, options), source); - - perform_ocr(request).await.unwrap(); - server.await.unwrap(); - let requests = seen.lock().unwrap(); - assert_eq!(requests.len(), 1); - assert!(requests[0].starts_with("POST /parse ")); - assert_eq!(request_body(&requests[0]), expected); -} - -#[rstest] -#[case("parse-v3")] -#[case("parse-legacy")] -#[tokio::test] -async fn data_uri_upload_preserves_multipart_headers( - #[case] model: &str, - #[values("application/pdf", "image/png")] mime_type: &str, -) { - let (base, seen, server) = mock_server(vec![ - MockResponse::json(json!({"file_id":"reducto://uploaded.pdf"})), - MockResponse::json(json!({"result":{"chunks":[{"content":"hello"}]}})), - ]) - .await; - let document = if mime_type.starts_with("image/") { - json!({"type":"image_url","image_url":format!("data:{mime_type};base64,YWJj")}) - } else { - json!({"type":"document_url","document_url":format!("data:{mime_type};base64,YWJj")}) - }; - let mut request = crate::ocr::types::LiteLLMOcrRequest { - document: serde_json::from_value::(document) - .unwrap() - .into(), - ..wire_request(&format!("reducto/{model}"), &base, json!({})) - }; - request.transport.extra_headers = vec![ - ("Content-Type".into(), "application/json".into()), - ("X-Trace".into(), "upload-test".into()), - ]; - - let response = perform_ocr(request).await.unwrap(); - server.await.unwrap(); - assert_eq!(response.pages[0].markdown, "hello"); - let requests = seen.lock().unwrap(); - assert_eq!(requests.len(), 2); - assert!(requests[0].starts_with("POST /upload ")); - assert!( - requests[0] - .to_ascii_lowercase() - .contains("content-type: multipart/form-data; boundary=") - ); - assert!(requests[0].contains("x-trace: upload-test")); - let multipart = requests[0].split_once("\r\n\r\n").unwrap().1; - assert!(multipart.contains(&format!("Content-Type: {mime_type}\r\n"))); - assert!(multipart.contains("\r\n\r\nabc\r\n--")); - assert!(requests[1].starts_with("POST /parse ")); - let source_field = if model == "parse-legacy" { - "document_url" - } else { - "input" - }; - assert_eq!( - request_body(&requests[1]), - json!({source_field:"reducto://uploaded.pdf"}) - ); - for request in requests.iter() { - assert!( - request - .to_ascii_lowercase() - .contains("authorization: bearer test-key\r\n") - ); - } -} - -#[tokio::test] -async fn response_received_stays_after_reducto_upload_and_parse() { - let (base, seen, server) = mock_server(vec![ - MockResponse::json(json!({"file_id":"reducto://uploaded.pdf"})), - MockResponse::json(json!({"result":{"chunks":[]}})), - ]) - .await; - let request_count = seen.clone(); - let host = LocalOcrHost::new(wire_request("reducto/parse-v3", &base, json!({}))).with_observer( - move |event| { - if let CallEvent::Machine(MachineEvent::ResponseReceived { raw }) = event { - assert_eq!(request_count.lock().unwrap().len(), 2); - assert_eq!(raw.body, r#"{"result":{"chunks":[]}}"#); - } - }, - ); - - perform_ocr_with(host).await.unwrap(); - server.await.unwrap(); - assert_eq!(seen.lock().unwrap().len(), 2); -} - -#[rstest] -#[case(json!({"file_id":""}))] -#[case(json!({}))] -#[case(json!({"file_id":null}))] -#[tokio::test] -async fn invalid_upload_ids_stop_before_parse(#[case] response: Value) { - let (base, seen, server) = mock_server(vec![MockResponse::json(response)]).await; - let error = perform_ocr(wire_request("reducto/parse-v3", &base, json!({}))) - .await - .unwrap_err(); - server.await.unwrap(); - assert!(error.to_string().contains("file_id")); - assert_eq!(seen.lock().unwrap().len(), 1); -} - -#[tokio::test] -async fn upload_failure_stops_before_parse() { - let (base, seen, server) = mock_server(vec![MockResponse { - status: 503, - headers: vec![], - body: json!({"error":"unavailable"}), - }]) - .await; - assert!( - perform_ocr(wire_request("reducto/parse-v3", &base, json!({}))) - .await - .is_err() - ); - server.await.unwrap(); - assert_eq!(seen.lock().unwrap().len(), 1); -} - -#[rstest] -#[case("https://example.com/a.pdf", Error::ReductoSource)] -#[case("reducto://", Error::RequestField { path: "document file id".into() })] -#[case("data:application/pdf;base64", Error::InvalidDataUri)] -#[case("data:application/pdf;base64,INVALID!", Error::InvalidDataUri)] -#[tokio::test] -async fn rejects_invalid_document_sources_before_network( - #[case] source: &str, - #[case] expected: Error, -) { - let (base, seen, server) = mock_server(vec![MockResponse::json(json!({}))]).await; - let request = super::test_support::with_source( - wire_request("reducto/parse-v3", &base, json!({})), - source, - ); - let result = perform_ocr(request).await; - server.abort(); - let _ = server.await; - assert!( - seen.lock().unwrap().is_empty(), - "sent invalid source: {source}" - ); - let error = result.unwrap_err(); - assert_eq!( - std::mem::discriminant(&error), - std::mem::discriminant(&expected) - ); - assert_eq!(error.http_status_code(), Some(400)); - assert_eq!(error.to_string(), expected.to_string()); -} - -#[test] -fn response_normalization_groups_blocks_and_distinguishes_null_result() { - use litellm_llms::reducto::ocr::transformation::{ - ReductoResponse, normalize_response as transform_ocr_response, - }; - - let raw = json!({"usage":{"num_pages":"2","credits":"3"},"result":{"type":"full","chunks":[ - {"blocks":[{ - "type":"Table", - "content":"B", - "bbox":{"left":0.1,"top":0.2,"width":0.8,"height":0.3,"page":2,"original_page":4}, - "confidence":"high", - "granular_confidence":{"parse_confidence":0.95,"extract_confidence":null}, - "image_url":null - }]}, - {"blocks":[{"content":"A","bbox":{"page":1},"type":"Text"},{"content":"C","bbox":{"page":1}}]} - ]}}); - let response: ReductoResponse = serde_json::from_value(raw).unwrap(); - let normalized = transform_ocr_response("parse-v3", response) - .unwrap() - .into_json(); - assert_eq!(normalized["pages"][0]["markdown"], "A\n\nC"); - assert_eq!(normalized["pages"][1]["markdown"], "B"); - assert_eq!(normalized["pages"][1]["blocks"][0]["type"], "Table"); - assert_eq!( - normalized["pages"][1]["blocks"][0]["bbox"], - json!({"left":0.1,"top":0.2,"width":0.8,"height":0.3,"page":2,"original_page":4}) - ); - assert_eq!(normalized["pages"][1]["blocks"][0]["confidence"], "high"); - assert_eq!( - normalized["pages"][1]["blocks"][0]["granular_confidence"]["parse_confidence"], - 0.95 - ); - assert!(normalized["pages"][1]["blocks"][0]["image_url"].is_null()); - assert_eq!(normalized["usage_info"]["pages_processed"], 2); - assert_eq!(normalized["usage_info"]["credits"], 3.0); - - let missing: ReductoResponse = - serde_json::from_value(json!({"chunks":[{"content":"text"}]})).unwrap(); - let missing = transform_ocr_response("parse-v3", missing).unwrap(); - assert_eq!(missing.pages[0].markdown, "text"); - let null: ReductoResponse = serde_json::from_value( - json!({"result":null,"chunks":[{"content":"ignored"}],"usage":null}), - ) - .unwrap(); - let null = transform_ocr_response("parse-v3", null).unwrap(); - assert!(null.pages.is_empty()); -} - -#[tokio::test] -async fn facade_omits_native_response_by_default_and_preserves_auth_priority() { - let raw = json!({"job_id":"job-1","result":{"chunks":[]}}); - let (base, seen, server) = mock_server(vec![MockResponse::json(raw)]).await; - let mut request = super::test_support::with_source( - wire_request("reducto/parse-v3", &base, json!({})), - "reducto://ready.pdf", - ); - request.transport.extra_headers = vec![("authorization".into(), "Bearer existing".into())]; - - let response = perform_ocr(request).await.unwrap(); - server.await.unwrap(); - assert_eq!(response.provider_native_response, None); - assert!( - seen.lock().unwrap()[0] - .to_ascii_lowercase() - .contains("authorization: bearer existing") - ); -} - -#[tokio::test] -async fn native_format_retains_the_provider_response() { - let raw = json!({ - "result":{"chunks":[{"content":"native OCR response"}]}, - "usage":{"num_pages":1} - }); - let (base, _, server) = mock_server(vec![MockResponse::json(raw.clone())]).await; - let request = super::test_support::with_source( - wire_request("reducto/parse-v3", &base, json!({"req_format":"native"})), - "reducto://ready.pdf", - ); - - let response = perform_ocr(request).await.unwrap(); - server.await.unwrap(); - - assert_eq!(response.pages[0].markdown, "native OCR response"); - assert_eq!(response.provider_native_response.as_ref(), raw.as_object()); -} - -#[tokio::test] -async fn unknown_model_reaches_parse_and_keeps_its_name() { - let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ - "result":{"chunks":[{"content":"future model response"}]} - }))]) - .await; - let request = super::test_support::with_source( - wire_request("reducto/future-parse-model", &base, json!({})), - "reducto://ready.pdf", - ); - - let response = perform_ocr(request).await.unwrap(); - server.await.unwrap(); - - assert_eq!(response.model, "future-parse-model"); - assert_eq!(response.pages[0].markdown, "future model response"); - let requests = seen.lock().unwrap(); - assert!(requests[0].starts_with("POST /parse ")); - assert_eq!( - request_body(&requests[0]), - json!({"input":"reducto://ready.pdf"}) - ); -} - -#[tokio::test] -async fn guardrail_rewrites_document_before_upload() { - let (base, seen, server) = - mock_server(vec![MockResponse::json(json!({"result":{"chunks":[]}}))]).await; - let host = LocalOcrHost::new(wire_request("reducto/parse-v3", &base, json!({}))) - .with_before_send(|wire, _| { - assert_eq!( - wire.body["document_url"], - "data:application/pdf;base64,YWJj" - ); - Ok(WireRequest { - body: json!({"type":"document_url","document_url":"reducto://guarded.pdf"}), - ..wire - }) - }); - - perform_ocr_with(host).await.unwrap(); - server.await.unwrap(); - let requests = seen.lock().unwrap(); - assert_eq!(requests.len(), 1); - assert!(requests[0].starts_with("POST /parse ")); - assert!(requests[0].contains("reducto://guarded.pdf")); -} - -mod transformation { - use litellm_host::event::{CallEvent, MachineEvent, WireRequest}; - use litellm_llms::{ - base_llm::ocr::transformation::{BaseOcrConfig, OcrConnection, OcrRequestContext}, - reducto::ocr::transformation::*, - }; - use rstest::rstest; - - use super::*; - use crate::ocr::{ - route::LocalOcrHost, - test_support::{MockResponse, mock_server, perform_ocr, perform_ocr_with, wire_request}, - }; - - #[tokio::test] - async fn v3_options_preserve_explicit_null() { - let overrides = - serde_json::from_value(json!({"formatting":null,"settings":{},"unknown":true})) - .unwrap(); - let params = ReductoParseV3Config - .map_ocr_params(&overrides, "parse-v3") - .unwrap(); - let client = crate::ocr::test_support::ocr_client(); - let connection = OcrConnection::default(); - let document = serde_json::from_value( - json!({"type":"document_url","document_url":"reducto://ready.pdf"}), - ) - .unwrap(); - let body = ReductoParseV3Config - .async_transform_ocr_request( - "parse-v3", - document, - ¶ms, - &[], - OcrRequestContext { - client: &client, - connection: &connection, - }, - ) - .await - .unwrap(); - assert_eq!( - serde_json::to_value(body).unwrap(), - json!({ - "input":"reducto://ready.pdf", "formatting":null, "settings":{} - }) - ); - let absent = ReductoParseV3Config - .map_ocr_params( - &litellm_core_utils::call_arguments::CallArguments::default(), - "parse-v3", - ) - .unwrap(); - assert_eq!(serde_json::to_value(absent).unwrap(), json!({})); - } - - #[rstest] - #[case( - "reducto/parse-v3", - json!({ - "formatting":{"table_output_format":"html"}, - "retrieval":{"chunk_mode":"section"}, - "settings":{"ocr_system":"standard"}, - "future_ocr_option":true, - "extra_body":{"provider_option":"value"} - }), - "reducto://already.pdf", - json!({ - "input":"reducto://already.pdf", - "formatting":{"table_output_format":"html"}, - "retrieval":{"chunk_mode":"section"}, - "settings":{"ocr_system":"standard"}, - "future_ocr_option":true, - "provider_option":"value" - }) - )] - #[case( - "reducto/parse-legacy", - json!({ - "enhance":{"agentic":[{"type":"table"}]}, - "future_ocr_option":true, - "extra_body":{"provider_option":"value"} - }), - "reducto://legacy.pdf", - json!({ - "document_url":"reducto://legacy.pdf", - "options":{"enhance":{"agentic":[{"type":"table"}]}}, - "future_ocr_option":true, - "provider_option":"value" - }) - )] - #[tokio::test] - async fn request_mapping_matches_python( - #[case] model: &str, - #[case] options: Value, - #[case] source: &str, - #[case] expected: Value, - ) { - let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ - "result":{"chunks":[]} - }))]) - .await; - let request = - crate::ocr::test_support::with_source(wire_request(model, &base, options), source); - - perform_ocr(request).await.unwrap(); - server.await.unwrap(); - let requests = seen.lock().unwrap(); - assert_eq!(requests.len(), 1); - assert!(requests[0].starts_with("POST /parse ")); - assert_eq!(request_body(&requests[0]), expected); - } - - #[rstest] - #[case("parse-v3")] - #[case("parse-legacy")] - #[tokio::test] - async fn data_uri_upload_preserves_multipart_headers(#[case] model: &str) { - let (base, seen, server) = mock_server(vec![ - MockResponse::json(json!({"file_id":"reducto://uploaded.pdf"})), - MockResponse::json(json!({"result":{"chunks":[{"content":"hello"}]}})), - ]) - .await; - let mut request = wire_request(&format!("reducto/{model}"), &base, json!({})); - request.transport.extra_headers = vec![ - ("Content-Type".into(), "application/json".into()), - ("X-Trace".into(), "upload-test".into()), - ]; - - let response = perform_ocr(request).await.unwrap(); - server.await.unwrap(); - assert_eq!(response.pages[0].markdown, "hello"); - let requests = seen.lock().unwrap(); - assert_eq!(requests.len(), 2); - assert!(requests[0].starts_with("POST /upload ")); - assert!( - requests[0] - .to_ascii_lowercase() - .contains("content-type: multipart/form-data; boundary=") - ); - assert!(requests[0].contains("x-trace: upload-test")); - assert!(requests[0].contains("application/pdf")); - assert!(requests[0].contains("abc")); - assert!(requests[1].starts_with("POST /parse ")); - } - - #[tokio::test] - async fn response_received_stays_after_reducto_upload_and_parse() { - let (base, seen, server) = mock_server(vec![ - MockResponse::json(json!({"file_id":"reducto://uploaded.pdf"})), - MockResponse::json(json!({"result":{"chunks":[]}})), - ]) - .await; - let request_count = seen.clone(); - let host = LocalOcrHost::new(wire_request("reducto/parse-v3", &base, json!({}))) - .with_observer(move |event| { - if let CallEvent::Machine(MachineEvent::ResponseReceived { raw }) = event { - assert_eq!(request_count.lock().unwrap().len(), 2); - assert_eq!(raw.body, r#"{"result":{"chunks":[]}}"#); - } - }); - - perform_ocr_with(host).await.unwrap(); - server.await.unwrap(); - assert_eq!(seen.lock().unwrap().len(), 2); - } - - #[rstest] - #[case("https://example.com/a.pdf")] - #[case("reducto://")] - #[case("data:application/pdf;base64")] - #[case("data:application/pdf;base64,INVALID!")] - #[tokio::test] - async fn rejects_invalid_document_sources_before_network(#[case] source: &str) { - let request = crate::ocr::test_support::with_source( - wire_request("reducto/parse-v3", "http://127.0.0.1:1", json!({})), - source, - ); - assert!(perform_ocr(request).await.is_err()); - } - - #[tokio::test] - async fn facade_omits_native_response_by_default_and_preserves_auth_priority() { - let raw = json!({"job_id":"job-1","result":{"chunks":[]}}); - let (base, seen, server) = mock_server(vec![MockResponse::json(raw)]).await; - let mut request = crate::ocr::test_support::with_source( - wire_request("reducto/parse-v3", &base, json!({})), - "reducto://ready.pdf", - ); - request.transport.extra_headers = vec![("authorization".into(), "Bearer existing".into())]; - - let response = perform_ocr(request).await.unwrap(); - server.await.unwrap(); - assert_eq!(response.provider_native_response, None); - assert!( - seen.lock().unwrap()[0] - .to_ascii_lowercase() - .contains("authorization: bearer existing") - ); - } - - #[rstest] - #[case("reducto/parse-v3")] - #[case("reducto/parse-legacy")] - #[tokio::test] - async fn guardrail_headers_reach_upload_and_parse(#[case] model: &str) { - let (base, seen, server) = mock_server(vec![ - MockResponse::json(json!({"file_id":"reducto://uploaded.pdf"})), - MockResponse::json(json!({"result":{"chunks":[]}})), - ]) - .await; - let mut request = wire_request(model, &base, json!({})); - request.transport.extra_headers = vec![("authorization".into(), "Bearer original".into())]; - let host = LocalOcrHost::new(request).with_before_send(|wire, _| { - Ok(WireRequest { - headers: vec![("authorization".into(), "Bearer guarded".into())], - ..wire - }) - }); - - perform_ocr_with(host).await.unwrap(); - server.await.unwrap(); - let requests = seen.lock().unwrap(); - assert_eq!(requests.len(), 2); - assert!(requests[0].starts_with("POST /upload ")); - assert!(requests[1].starts_with("POST /parse ")); - for request in requests.iter() { - assert!(request.contains("authorization: Bearer guarded")); - assert!(!request.contains("Bearer original")); - } - } -} diff --git a/litellm-rust/crates/core/tests/vertex_ai_deepseek_ocr.rs b/litellm-rust/crates/core/tests/vertex_ai_deepseek_ocr.rs deleted file mode 100644 index 2e8d69f5f64..00000000000 --- a/litellm-rust/crates/core/tests/vertex_ai_deepseek_ocr.rs +++ /dev/null @@ -1,143 +0,0 @@ -use litellm_auth::InputSource; -use serde_json::{Value, json}; - -use super::test_support::{MockResponse, mock_server, perform_ocr, wire_request}; - -fn request_body(request: &str) -> Value { - serde_json::from_str(request.split_once("\r\n\r\n").unwrap().1).unwrap() -} - -#[tokio::test] -async fn facade_executes_vertex_deepseek_at_the_openai_endpoint() { - let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ - "choices":[{"message":{"content":"recognized"}}], - "usage":{"prompt_tokens":1} - }))]) - .await; - let request = wire_request( - "vertex_ai/deepseek-ocr-maas", - &base, - json!({ - "vertex_project":"project-1", - "vertex_location":"europe-west4", - "temperature":0.1, - "future_ocr_option":true, - "extra_body":{"provider_option":"value"} - }), - ); - let request = super::test_support::with_source(request, "gs://bucket/document.pdf"); - - let response = perform_ocr(request).await.unwrap(); - server.await.unwrap(); - assert_eq!(response.pages[0].markdown, "recognized"); - assert_eq!( - response.usage_info.unwrap().extra_fields["prompt_tokens"], - 1 - ); - let requests = seen.lock().unwrap(); - assert!(requests[0].starts_with( - "POST /v1/projects/project-1/locations/europe-west4/endpoints/openapi/chat/completions " - )); - assert!( - requests[0] - .to_ascii_lowercase() - .contains("authorization: bearer test-key") - ); - let body = request_body(&requests[0]); - assert_eq!(body["model"], "deepseek-ai/deepseek-ocr-maas"); - assert_eq!(body["temperature"], 0.1); - assert_eq!(body["future_ocr_option"], true); - assert!(body.get("extra_body").is_none()); - assert_eq!( - body["messages"][0]["content"][0], - json!({"type":"image_url","image_url":"gs://bucket/document.pdf"}) - ); -} - -#[test] -fn host_registration_selects_deepseek_without_affecting_mistral() { - assert!(crate::ocr::arguments::is_supported_request( - "deepseek-ocr-maas", - Some("vertex_ai") - )); - assert!(crate::ocr::arguments::is_supported_request( - "mistral-ocr-maas", - Some("vertex_ai") - )); -} - -#[tokio::test] -async fn request_controlled_api_base_is_rejected_before_vertex_auth() { - let mut request = wire_request( - "vertex_ai/deepseek-ocr-maas", - "https://caller.example", - json!({"vertex_project":"project-1"}), - ); - request.credentials.api_base = Some(litellm_auth::Sourced::new( - "https://caller.example".into(), - InputSource::Request, - )); - - let error = perform_ocr(request).await.unwrap_err(); - assert!( - error - .to_string() - .contains("request-controlled Vertex AI endpoint") - ); -} - -mod deepseek_transformation { - use serde_json::json; - - use super::*; - use crate::ocr::test_support::{MockResponse, mock_server, perform_ocr, wire_request}; - - #[tokio::test] - async fn facade_executes_vertex_deepseek_at_the_openai_endpoint() { - let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ - "choices":[{"message":{"content":"recognized"}}], - "usage":{"prompt_tokens":1} - }))]) - .await; - let request = wire_request( - "vertex_ai/deepseek-ocr-maas", - &base, - json!({ - "vertex_project":"project-1", - "vertex_location":"europe-west4", - "temperature":0.1, - "future_ocr_option":true, - "extra_body":{"provider_option":"value"} - }), - ); - let request = crate::ocr::test_support::with_source(request, "gs://bucket/document.pdf"); - - let response = perform_ocr(request).await.unwrap(); - server.await.unwrap(); - assert_eq!(response.pages[0].markdown, "recognized"); - assert_eq!( - response.usage_info.unwrap().extra_fields["prompt_tokens"], - 1 - ); - let requests = seen.lock().unwrap(); - assert!(requests[0].starts_with( - "POST /v1/projects/project-1/locations/europe-west4/endpoints/openapi/chat/completions " - )); - assert!( - requests[0] - .to_ascii_lowercase() - .contains("authorization: bearer test-key") - ); - let body = request_body(&requests[0]); - assert_eq!(body["model"], "deepseek-ai/deepseek-ocr-maas"); - assert_eq!(body["temperature"], 0.1); - assert_eq!(body["future_ocr_option"], true); - assert_eq!(body["provider_option"], "value"); - assert!(body.get("vertex_project").is_none()); - assert!(body.get("extra_body").is_none()); - assert_eq!( - body["messages"][0]["content"][0], - json!({"type":"image_url","image_url":"gs://bucket/document.pdf"}) - ); - } -} diff --git a/litellm-rust/crates/core/tests/vertex_ai_ocr.rs b/litellm-rust/crates/core/tests/vertex_ai_ocr.rs deleted file mode 100644 index 035f3fe944d..00000000000 --- a/litellm-rust/crates/core/tests/vertex_ai_ocr.rs +++ /dev/null @@ -1,293 +0,0 @@ -use litellm_auth::InputSource; -use litellm_llms::base_llm::ocr::{settings::OcrSettings, transformation::OcrResponseFormat}; -use serde_json::{Value, json}; - -use super::test_support::{MockResponse, mock_server, ocr_client, perform_ocr, wire_request}; - -fn request_body(request: &str) -> Value { - serde_json::from_str(request.split_once("\r\n\r\n").unwrap().1).unwrap() -} - -#[tokio::test] -async fn facade_executes_vertex_mistral_with_resolved_project_and_location() { - let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ - "pages":[{"index":0,"markdown":"hello"}], - "usage_info":{"pages_processed":1} - }))]) - .await; - let request = wire_request( - "vertex_ai/mistral-ocr-maas", - &base, - json!({ - "vertex_project":"project-1", - "vertex_location":"europe-west4", - "extract_footer":true - }), - ); - - let response = perform_ocr(request).await.unwrap(); - server.await.unwrap(); - assert_eq!(response.pages[0].markdown, "hello"); - let requests = seen.lock().unwrap(); - assert_eq!(requests.len(), 1); - assert!(requests[0].starts_with( - "POST /v1/projects/project-1/locations/europe-west4/publishers/mistralai/models/mistral-ocr-maas:rawPredict " - )); - assert!( - requests[0] - .to_ascii_lowercase() - .contains("authorization: bearer test-key") - ); - assert_eq!( - request_body(&requests[0]), - json!({ - "model":"mistral-ocr-maas", - "document":{"type":"document_url","document_url":"data:application/pdf;base64,YWJj"}, - "extract_footer":true - }) - ); -} - -#[tokio::test] -async fn configured_project_and_location_apply_when_the_call_sets_neither() { - let (base, seen, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await; - let client = ocr_client().with_settings(OcrSettings { - vertex_project: Some("configured-project".into()), - vertex_location: Some("europe-west4".into()), - ..OcrSettings::default() - }); - - crate::ocr::client::perform( - &client, - wire_request("vertex_ai/mistral-ocr-maas", &base, json!({})), - ) - .await - .unwrap(); - server.await.unwrap(); - assert!(seen.lock().unwrap()[0].starts_with( - "POST /v1/projects/configured-project/locations/europe-west4/publishers/mistralai/models/mistral-ocr-maas:rawPredict " - )); -} - -#[tokio::test] -async fn supplied_authorization_is_forwarded_without_a_static_token() { - let (base, seen, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await; - let mut request = wire_request( - "vertex_ai/model", - &base, - json!({"vertex_project":"project-1"}), - ); - request.credentials.api_key = None; - request.transport.extra_headers = vec![("authorization".into(), "Bearer supplied".into())]; - - perform_ocr(request).await.unwrap(); - server.await.unwrap(); - assert!( - seen.lock().unwrap()[0] - .to_ascii_lowercase() - .contains("authorization: bearer supplied") - ); -} - -#[tokio::test] -async fn invalid_credentials_fail_before_provider_http() { - let request = wire_request( - "vertex_ai/model", - "http://127.0.0.1:1", - json!({"vertex_credentials": true}), - ); - let error = perform_ocr(request).await.unwrap_err(); - assert!(error.to_string().contains("vertex_credentials")); -} - -#[tokio::test] -async fn request_controlled_api_base_is_rejected_before_vertex_auth() { - let mut request = wire_request( - "vertex_ai/mistral-ocr-maas", - "https://caller.example", - json!({"vertex_project":"project-1"}), - ); - request.credentials.api_base = Some(litellm_auth::Sourced::new( - "https://caller.example".into(), - InputSource::Request, - )); - - let error = perform_ocr(request).await.unwrap_err(); - assert!( - error - .to_string() - .contains("request-controlled Vertex AI endpoint") - ); -} - -#[tokio::test] -async fn adapters_build_complete_requests_and_share_mistral_normalization() { - use std::time::Duration; - - use litellm_llms::{ - base_llm::ocr::transformation::BaseOcrConfig, - mistral::ocr::transformation::MistralOcrConfig, - vertex_ai::ocr::transformation::VertexAiOcrConfig, - }; - - use crate::ocr::test_support::ocr_client; - - let client = ocr_client(); - let options = json!({ - "pages": [0, 2], - "include_image_base64": true, - "vertex_project": "project-1", - "vertex_location": "us-central1", - "unknown": "ignored" - }); - let direct = wire_request( - "mistral/mistral-ocr-maas", - "https://mistral.test", - options.clone(), - ); - let vertex = wire_request("vertex_ai/mistral-ocr-maas", "https://vertex.test", options); - let direct = crate::ocr::prepare::prepare_request_for_test( - super::test_support::resolved_request(direct), - ); - let vertex = crate::ocr::prepare::prepare_request_for_test( - super::test_support::resolved_request(vertex), - ); - let direct_http = MistralOcrConfig - .prepare_request(&direct, &client, &crate::ocr::test_support::NoHooks) - .await - .unwrap(); - let vertex_http = VertexAiOcrConfig - .prepare_request(&vertex, &client, &crate::ocr::test_support::NoHooks) - .await - .unwrap(); - assert_eq!(direct_http.url(), "https://mistral.test/v1/ocr"); - assert_eq!( - vertex_http.url(), - "https://vertex.test/v1/projects/project-1/locations/us-central1/publishers/mistralai/models/mistral-ocr-maas:rawPredict" - ); - for http in [&direct_http, &vertex_http] { - assert_eq!(http.header("authorization").unwrap(), "Bearer test-key"); - assert_eq!(http.header("content-type").unwrap(), "application/json"); - assert_eq!(http.timeout(), Some(Duration::from_secs(2))); - let body: Value = serde_json::from_slice(http.body()).unwrap(); - assert_eq!( - body, - json!({ - "model": "mistral-ocr-maas", - "document": {"type": "document_url", "document_url": "data:application/pdf;base64,YWJj"}, - "pages": [0, 2], - "include_image_base64": true, - "unknown": "ignored" - }) - ); - } - let payload = json!({"pages": [{"index": 0, "markdown": "hello"}], "extra": "preserved"}); - let raw = serde_json::to_vec(&payload).unwrap(); - let direct_response = MistralOcrConfig - .transform_ocr_response(&direct.model, &raw, OcrResponseFormat::Litellm) - .unwrap() - .into_json(); - let vertex_response = VertexAiOcrConfig - .transform_ocr_response(&vertex.model, &raw, OcrResponseFormat::Litellm) - .unwrap() - .into_json(); - assert_eq!(direct_response, vertex_response); - assert_eq!(direct_response["model"], "mistral-ocr-maas"); - assert_eq!(direct_response["object"], "ocr"); - assert_eq!(direct_response["extra"], "preserved"); -} - -mod transformation { - - use rstest::rstest; - use serde_json::{Value, json}; - - use crate::ocr::test_support::wire_request; - - #[rstest] - #[case::mistral(false)] - #[case::vertex(true)] - #[tokio::test] - async fn configs_build_complete_requests_and_share_mistral_normalization( - #[case] use_vertex: bool, - ) { - use std::time::Duration; - - use litellm_llms::{ - base_llm::ocr::transformation::BaseOcrConfig, - mistral::ocr::transformation::MistralOcrConfig, - vertex_ai::ocr::transformation::VertexAiOcrConfig, - }; - - use crate::ocr::test_support::ocr_client; - - let client = ocr_client(); - let options = json!({ - "pages": [0, 2], - "include_image_base64": true, - "vertex_project": "project-1", - "vertex_location": "us-central1", - "unknown": "preserved" - }); - let direct = wire_request( - "mistral/mistral-ocr-maas", - "https://mistral.test", - options.clone(), - ); - let vertex = wire_request("vertex_ai/mistral-ocr-maas", "https://vertex.test", options); - let direct = crate::ocr::prepare::prepare_request_for_test( - crate::ocr::test_support::resolved_request(direct), - ); - let vertex = crate::ocr::prepare::prepare_request_for_test( - crate::ocr::test_support::resolved_request(vertex), - ); - let direct_http = MistralOcrConfig - .prepare_request(&direct, &client, &crate::ocr::test_support::NoHooks) - .await - .unwrap(); - let vertex_http = VertexAiOcrConfig - .prepare_request(&vertex, &client, &crate::ocr::test_support::NoHooks) - .await - .unwrap(); - assert_eq!(direct_http.url(), "https://mistral.test/v1/ocr"); - assert_eq!( - vertex_http.url(), - "https://vertex.test/v1/projects/project-1/locations/us-central1/publishers/mistralai/models/mistral-ocr-maas:rawPredict" - ); - let http = if use_vertex { - &vertex_http - } else { - &direct_http - }; - assert_eq!(http.header("authorization").unwrap(), "Bearer test-key"); - assert_eq!(http.header("content-type").unwrap(), "application/json"); - assert_eq!(http.timeout(), Some(Duration::from_secs(2))); - let body: Value = serde_json::from_slice(http.body()).unwrap(); - assert_eq!( - body, - json!({ - "model": "mistral-ocr-maas", - "document": {"type": "document_url", "document_url": "data:application/pdf;base64,YWJj"}, - "pages": [0, 2], - "include_image_base64": true, - "unknown": "preserved" - }) - ); - let payload = serde_json::to_vec( - &json!({"pages": [{"index": 0, "markdown": "hello"}], "extra": "preserved"}), - ) - .unwrap(); - let direct_response = MistralOcrConfig - .transform_ocr_response(&direct.model, &payload, Default::default()) - .unwrap() - .into_json(); - let vertex_response = VertexAiOcrConfig - .transform_ocr_response(&vertex.model, &payload, Default::default()) - .unwrap() - .into_json(); - assert_eq!(direct_response, vertex_response); - assert_eq!(direct_response["model"], "mistral-ocr-maas"); - assert_eq!(direct_response["object"], "ocr"); - assert_eq!(direct_response["extra"], "preserved"); - } -} diff --git a/litellm-rust/crates/http/src/request.rs b/litellm-rust/crates/http/src/request.rs index 874a0f3abf9..fcf296793a5 100644 --- a/litellm-rust/crates/http/src/request.rs +++ b/litellm-rust/crates/http/src/request.rs @@ -217,13 +217,25 @@ mod tests { "Authorization".to_string(), "Bearer abc".to_string() )])); + assert!(has_bearer_auth(&[( + "authorization".to_string(), + "bearer abc".to_string() + )])); assert!(!has_bearer_auth(&[( "Authorization".to_string(), "Bearer ".to_string() )])); + assert!(!has_bearer_auth(&[( + "authorization".to_string(), + String::new() + )])); assert!(!has_bearer_auth(&[( "Authorization".to_string(), "Basic abc".to_string() )])); + assert!(!has_bearer_auth(&[( + "x-api-key".to_string(), + "abc".to_string() + )])); } } diff --git a/litellm-rust/crates/llms/src/anthropic/chat/transformation.rs b/litellm-rust/crates/llms/src/anthropic/chat/transformation.rs index 6fc4f00b981..fd86c5ca25a 100644 --- a/litellm-rust/crates/llms/src/anthropic/chat/transformation.rs +++ b/litellm-rust/crates/llms/src/anthropic/chat/transformation.rs @@ -218,7 +218,3 @@ fn anthropic_body( ); Value::Object(body) } - -#[cfg(test)] -#[path = "tests.rs"] -mod tests; diff --git a/litellm-rust/crates/llms/src/bedrock/chat/converse_transformation.rs b/litellm-rust/crates/llms/src/bedrock/chat/converse_transformation.rs index 09c456f1d0a..b5db88d7dc4 100644 --- a/litellm-rust/crates/llms/src/bedrock/chat/converse_transformation.rs +++ b/litellm-rust/crates/llms/src/bedrock/chat/converse_transformation.rs @@ -302,7 +302,3 @@ fn has_blank_text(message: &ChatMessage) -> bool { }), } } - -#[cfg(test)] -#[path = "tests.rs"] -mod tests; diff --git a/litellm-rust/crates/llms/src/anthropic/chat/tests.rs b/litellm-rust/crates/llms/tests/anthropic_chat_transformation.rs similarity index 96% rename from litellm-rust/crates/llms/src/anthropic/chat/tests.rs rename to litellm-rust/crates/llms/tests/anthropic_chat_transformation.rs index 3777347d240..ed22a1d141d 100644 --- a/litellm-rust/crates/llms/src/anthropic/chat/tests.rs +++ b/litellm-rust/crates/llms/tests/anthropic_chat_transformation.rs @@ -1,7 +1,11 @@ -use serde_json::json; - -use super::*; -use crate::base_llm::chat::transformation::Error; +use litellm_llms::{ + anthropic::chat::transformation::ANTHROPIC_CHAT_COMPLETIONS_CONFIG, + base_llm::chat::transformation::{ + BaseConfig, Error, ProviderChatResponseData, RequestAuth, Unsupported, + }, +}; +use litellm_types::{llms::openai::ChatMessage, utils::ChatCompletionsResponse}; +use serde_json::{Map, Value, json}; fn messages(value: Value) -> Vec { serde_json::from_value(value).expect("valid messages") @@ -205,7 +209,8 @@ fn declines_tool_calls_tool_results_and_multimodal_content() { ); assert_eq!( reason( - json!([{"role": "user", "content": [ + json!([ + {"role": "user", "content": [ {"type": "image_url", "image_url": {"url": "https://x/y.png"}} ]}]), json!({}) @@ -214,7 +219,8 @@ fn declines_tool_calls_tool_results_and_multimodal_content() { ); assert_eq!( reason( - json!([{"role": "user", "content": [ + json!([ + {"role": "user", "content": [ {"type": "text", "text": "hi", "cache_control": {"type": "ephemeral"}} ]}]), json!({}) diff --git a/litellm-rust/crates/llms/src/bedrock/chat/tests.rs b/litellm-rust/crates/llms/tests/bedrock_converse_transformation.rs similarity index 98% rename from litellm-rust/crates/llms/src/bedrock/chat/tests.rs rename to litellm-rust/crates/llms/tests/bedrock_converse_transformation.rs index d7ecde47c6b..4127bcfa19d 100644 --- a/litellm-rust/crates/llms/src/bedrock/chat/tests.rs +++ b/litellm-rust/crates/llms/tests/bedrock_converse_transformation.rs @@ -1,7 +1,11 @@ -use serde_json::json; - -use super::*; -use crate::base_llm::chat::transformation::Error; +use litellm_llms::{ + base_llm::chat::transformation::{ + BaseConfig, Error, ProviderChatResponseData, RequestAuth, Unsupported, + }, + bedrock::chat::converse_transformation::BEDROCK_CHAT_COMPLETIONS_CONFIG, +}; +use litellm_types::{llms::openai::ChatMessage, utils::ChatCompletionsResponse}; +use serde_json::{Map, Value, json}; fn messages(value: Value) -> Vec { serde_json::from_value(value).expect("valid messages") diff --git a/litellm-rust/crates/types/src/utils.rs b/litellm-rust/crates/types/src/utils.rs index 5ca56ec9e49..af0ba2c01c9 100644 --- a/litellm-rust/crates/types/src/utils.rs +++ b/litellm-rust/crates/types/src/utils.rs @@ -56,7 +56,7 @@ pub struct ChatCompletionsChoice { /// /// There is deliberately no `id`: Python mints the `chatcmpl-…` id on the /// `ModelResponse` it already created, and echoing the provider's own id here -/// would change it. Pinned by `response_carries_no_id` in `tests.rs`. +/// would change it. Pinned by `response_carries_no_id` in the Anthropic chat transformation tests. #[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] pub struct ChatCompletionsResponse { pub created: u64, From e64e635185b25cb9ee649dba70a372a869270264 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 24 Sep 2026 13:35:06 -0700 Subject: [PATCH 076/120] fix(cost-map): add video and reasoning output prices to vertex gemini-omni-1.1-flash (#43036) Co-authored-by: kerry Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/model_prices_and_context_window_backup.json | 2 ++ model_prices_and_context_window.json | 2 ++ 2 files changed, 4 insertions(+) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index dad89fc58a8..90a268aa3e2 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -68612,7 +68612,9 @@ "input_cost_per_token": 1.5e-06, "litellm_provider": "vertex_ai", "mode": "chat", + "output_cost_per_reasoning_token": 9e-06, "output_cost_per_token": 9e-06, + "output_cost_per_video_token": 1.75e-05, "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing" }, "vertex_ai/gemini-omni-1.1-flash-preview": { diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index dad89fc58a8..90a268aa3e2 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -68612,7 +68612,9 @@ "input_cost_per_token": 1.5e-06, "litellm_provider": "vertex_ai", "mode": "chat", + "output_cost_per_reasoning_token": 9e-06, "output_cost_per_token": 9e-06, + "output_cost_per_video_token": 1.75e-05, "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing" }, "vertex_ai/gemini-omni-1.1-flash-preview": { From 77eccaca78d25820c54c6b2703711efe3f95224f Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 24 Sep 2026 15:40:36 -0500 Subject: [PATCH 077/120] feat(proxy): server-side Team Usage export beyond the top-N key cap (#42996) * feat(proxy): add uncapped server-side team usage export route GET /team/daily/activity/export answers the same scoping as /team/daily/activity/aggregated with one unbounded rollup query, so keys past USAGE_TOP_API_KEYS_LIMIT are included. Supports daily, daily_with_keys, daily_with_users and daily_with_models export types as CSV (default) or JSON. The PTU flat-cost sentinel stays in the plain daily rollup and is excluded from the keyed and per-model exports Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * feat(ui): export team usage server-side when the key list was truncated When the aggregated spend response reports api_key truncation, EntityUsage passes a serverExport into the export modal that downloads CSV or JSON from GET /team/daily/activity/export instead of building the file from the truncated on-screen data. apiClient gains a responseType option so the download can arrive as a Blob, and truncation no longer blocks the export button Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(proxy): cover team usage export types, sentinel handling and scope Unit tests pin the uncapped key rollup past USAGE_TOP_API_KEYS_LIMIT, PTU sentinel inclusion in the daily rollup and exclusion elsewhere, the per-user fold, and the CSV column layout. Integration tests exercise the route against a live proxy, including member scope denial Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * refactor(proxy): tidy team usage export route Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): use membership test for export type branch (PLR1714) Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * style(ui): format exportBlockedReason test with prettier Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): satisfy type-discipline gate in team usage export Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): pass export rows as a sequence to the response model Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(proxy-behavior): cover team usage export in the daily activity scope matrix Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * feat(proxy): carry PTU flat cost and escape formulas in team usage export Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(ui): keep the truncation export block on surfaces without a server export Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * chore(ui): drop redundant comments in team export call and modal test Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(integration): audit cells for team usage export Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(integration): tighten team usage export audit cells Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): type the export params tuple and fold user keys in one pass Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * style(ui): bring entity usage export helpers under eslint budgets Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * style(ui): prettier-format UsagePageView after merge Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: yassin Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/_types.py | 4 + .../common_daily_activity.py | 293 +++++++++- .../management_endpoints/team_endpoints.py | 183 +++++- .../management_endpoints/team_endpoints.py | 45 ++ .../endpointaudit/coverage_allowlist.txt | 1 + .../spend/test_team_daily_activity_export.py | 522 ++++++++++++++++++ .../management/test_team_daily_activity.py | 12 +- .../test_common_daily_activity.py | 305 ++++++++++ .../test_team_endpoints.py | 127 +++++ .../components/EntityUsage/EntityUsage.tsx | 19 +- .../_components/components/UsagePageView.tsx | 11 +- .../EntityUsageExportModal.test.tsx | 31 +- .../EntityUsageExportModal.tsx | 8 +- .../EntityUsageExport/UsageExportHeader.tsx | 5 +- .../exportBlockedReason.test.ts | 4 + .../EntityUsageExport/exportBlockedReason.ts | 4 +- .../src/components/EntityUsageExport/types.ts | 3 + .../EntityUsageExport/utils.test.ts | 37 ++ .../src/components/EntityUsageExport/utils.ts | 71 ++- .../src/components/networking.tsx | 31 ++ ui/litellm-dashboard/src/lib/http/client.ts | 10 +- ui/litellm-dashboard/src/lib/http/schema.d.ts | 147 ++++- 22 files changed, 1812 insertions(+), 61 deletions(-) create mode 100644 tests/integration/spend/test_team_daily_activity_export.py diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 73e3e0e6ee0..b6de36f8423 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -307,6 +307,7 @@ class KeyManagementRoutes(str, enum.Enum): # team usage routes TEAM_DAILY_ACTIVITY = "/team/daily/activity" TEAM_DAILY_ACTIVITY_AGGREGATED = "/team/daily/activity/aggregated" + TEAM_DAILY_ACTIVITY_EXPORT = "/team/daily/activity/export" TEAM_DAILY_ACTIVITY_AGGREGATED_SEARCH = "/team/daily/activity/aggregated/search" # team spend-log viewing @@ -719,6 +720,7 @@ class LiteLLMRoutes(enum.Enum): "/team/permissions_bulk_update", "/team/daily/activity", "/team/daily/activity/aggregated", + "/team/daily/activity/export", "/team/daily/activity/aggregated/search", "/team/spend/by_user", # gateway request counts (SGR); deployment-wide, admin-only @@ -890,6 +892,7 @@ class LiteLLMRoutes(enum.Enum): "/team/permissions_update", "/team/daily/activity", "/team/daily/activity/aggregated", + "/team/daily/activity/export", "/team/daily/activity/aggregated/search", "/team/spend/by_user", "/team/{team_id}/members/me", @@ -990,6 +993,7 @@ class LiteLLMRoutes(enum.Enum): "/user/daily/activity", "/team/daily/activity", "/team/daily/activity/aggregated", + "/team/daily/activity/export", "/team/daily/activity/aggregated/search", "/tag/daily/activity", "/tag/list", diff --git a/litellm/proxy/management_endpoints/common_daily_activity.py b/litellm/proxy/management_endpoints/common_daily_activity.py index bf8e7bc15fb..1d797206ead 100644 --- a/litellm/proxy/management_endpoints/common_daily_activity.py +++ b/litellm/proxy/management_endpoints/common_daily_activity.py @@ -1,4 +1,6 @@ import asyncio +import dataclasses +import itertools from collections.abc import Awaitable, Callable, Mapping, Sequence from collections.abc import Set as AbstractSet from datetime import datetime, timedelta, timezone @@ -35,6 +37,10 @@ from litellm.types.proxy.management_endpoints.common_daily_activity import ( SpendAnalyticsPaginatedResponse, SpendMetrics, ) +from litellm.types.proxy.management_endpoints.team_endpoints import ( + TeamDailyActivityExportRow, + TeamDailyActivityExportType, +) if TYPE_CHECKING: from prisma.models import ( @@ -198,7 +204,7 @@ class _AggregatedQueryKwargs(TypedDict): include_current_utc_day: ReadOnly[bool] -_SqlQuery = tuple[str, list[str]] +_SqlQuery = tuple[str, Sequence[str]] async def _query_raw_optional( @@ -974,6 +980,291 @@ def _build_entity_rollup_sql_query( return sql_query, sql_params +def _build_export_sql_query( + *, + table_name: str, + entity_id_field: str, + entity_id: str | list[str] | None, # mutable-ok: filter union shared with the paginated path + start_date: str, + end_date: str, + api_key: str | list[str] | None, # mutable-ok: filter union shared with the paginated path + exclude_entity_ids: list[str] | None, # mutable-ok: filter union shared with the paginated path + timezone_offset_minutes: int | None, + export_type: TeamDailyActivityExportType, +) -> tuple[str, tuple[str, ...]]: + """One unbounded rollup for the export route, on the aggregated path's WHERE clause. + + No LIMIT anywhere: the export exists so a caller can reach keys past + USAGE_TOP_API_KEYS_LIMIT. PTU sentinel rows stay in `daily` so per-team + totals match breakdown.entities, and are excluded from the key, user and + model exports where the flat-cost row has no meaning. + """ + pg_table: Final = _PRISMA_TO_PG_TABLE.get(table_name) + if pg_table is None: + raise ValueError(f"Unknown table name: {table_name}") + + adjusted_start, adjusted_end = _adjust_dates_for_timezone(start_date, end_date, timezone_offset_minutes) + where_clause, where_params = _build_aggregated_where_clause( + entity_id_field=entity_id_field, + entity_id=entity_id, + adjusted_start=adjusted_start, + adjusted_end=adjusted_end, + model=None, + api_key=api_key, + exclude_entity_ids=exclude_entity_ids, + ) + + keyed: Final = export_type in ("daily_with_keys", "daily_with_users") + by_model: Final = export_type == "daily_with_models" + group_extras: Final = tuple(field for field in ("api_key" if keyed else "", "model" if by_model else "") if field) + group_by: Final = f'date, "{entity_id_field}"' + "".join(f", {field}" for field in group_extras) + sentinel_clause: Final = f" AND api_key <> ${len(where_params) + 1}" if (keyed or by_model) else "" + sentinel_params: Final = (PTU_SENTINEL_API_KEY,) if (keyed or by_model) else () + + sql_query: Final = f""" + SELECT + date, + "{entity_id_field}" AS entity_id, + {"api_key" if keyed else "NULL::text AS api_key"}, + {"model" if by_model else "NULL::text AS model"},{_rollup_metric_select(table_name)} + FROM "{pg_table}" + WHERE {where_clause}{sentinel_clause} + GROUP BY {group_by} + ORDER BY {group_by} + """ + + return sql_query, (*where_params, *sentinel_params) + + +class _ExportRow(_RollupMetricsRow): + entity_id: str | None + model: str | None + + +def _export_team_alias(entity_metadata_field: Mapping[str, dict[str, object]] | None, entity_id: str) -> str | None: + alias: Final = _entity_metadata(entity_metadata_field, entity_id).get("team_alias") + return alias if isinstance(alias, str) else None + + +@dataclasses.dataclass(frozen=True, slots=True) +class _ExportMetrics: + spend: float + api_requests: int + successful_requests: int + failed_requests: int + total_tokens: int + prompt_tokens: int + completion_tokens: int + cache_read_input_tokens: int + cache_creation_input_tokens: int + + @classmethod + def from_record(cls, record: _RollupMetricsRow) -> "_ExportMetrics": + prompt_tokens: Final = record.prompt_tokens or 0 + completion_tokens: Final = record.completion_tokens or 0 + return cls( + spend=record.spend or 0.0, + api_requests=record.api_requests or 0, + successful_requests=record.successful_requests or 0, + failed_requests=record.failed_requests or 0, + total_tokens=prompt_tokens + completion_tokens, + prompt_tokens=prompt_tokens, + completion_tokens=completion_tokens, + cache_read_input_tokens=record.cache_read_input_tokens or 0, + cache_creation_input_tokens=record.cache_creation_input_tokens or 0, + ) + + @classmethod + def zero(cls) -> "_ExportMetrics": + return cls( + spend=0.0, + api_requests=0, + successful_requests=0, + failed_requests=0, + total_tokens=0, + prompt_tokens=0, + completion_tokens=0, + cache_read_input_tokens=0, + cache_creation_input_tokens=0, + ) + + def __add__(self, other: "_ExportMetrics") -> "_ExportMetrics": + return _ExportMetrics( + spend=self.spend + other.spend, + api_requests=self.api_requests + other.api_requests, + successful_requests=self.successful_requests + other.successful_requests, + failed_requests=self.failed_requests + other.failed_requests, + total_tokens=self.total_tokens + other.total_tokens, + prompt_tokens=self.prompt_tokens + other.prompt_tokens, + completion_tokens=self.completion_tokens + other.completion_tokens, + cache_read_input_tokens=self.cache_read_input_tokens + other.cache_read_input_tokens, + cache_creation_input_tokens=self.cache_creation_input_tokens + other.cache_creation_input_tokens, + ) + + +def _export_base_row( + record: _ExportRow, + entity_metadata_field: Mapping[str, dict[str, object]] | None, +) -> TeamDailyActivityExportRow: + entity_id: Final = record.entity_id or "Unassigned" + metrics: Final = _ExportMetrics.from_record(record) + return TeamDailyActivityExportRow( + date=record.date, + team_id=entity_id, + team_alias=_export_team_alias(entity_metadata_field, entity_id), + model=record.model, + spend=metrics.spend, + flat_cost=_reported_flat_cost(record), + api_requests=metrics.api_requests, + successful_requests=metrics.successful_requests, + failed_requests=metrics.failed_requests, + total_tokens=metrics.total_tokens, + prompt_tokens=metrics.prompt_tokens, + completion_tokens=metrics.completion_tokens, + cache_read_input_tokens=metrics.cache_read_input_tokens, + cache_creation_input_tokens=metrics.cache_creation_input_tokens, + ) + + +def _export_key_row( + record: _ExportRow, + entity_metadata_field: Mapping[str, dict[str, object]] | None, + api_key_metadata: Mapping[str, _KeyMetadataDict], +) -> TeamDailyActivityExportRow: + entity_id: Final = record.entity_id or "Unassigned" + metadata: Final = _key_metadata(api_key_metadata, record.api_key or "") + metrics: Final = _ExportMetrics.from_record(record) + return TeamDailyActivityExportRow( + date=record.date, + team_id=entity_id, + team_alias=_export_team_alias(entity_metadata_field, entity_id), + api_key=record.api_key, + key_alias=metadata.key_alias, + user_id=metadata.user_id, + user_email=metadata.user_email, + spend=metrics.spend, + api_requests=metrics.api_requests, + successful_requests=metrics.successful_requests, + failed_requests=metrics.failed_requests, + total_tokens=metrics.total_tokens, + prompt_tokens=metrics.prompt_tokens, + completion_tokens=metrics.completion_tokens, + cache_read_input_tokens=metrics.cache_read_input_tokens, + cache_creation_input_tokens=metrics.cache_creation_input_tokens, + ) + + +def _fold_export_users( + records: Sequence[_ExportRow], + entity_metadata_field: Mapping[str, dict[str, object]] | None, + api_key_metadata: Mapping[str, _KeyMetadataDict], +) -> tuple[TeamDailyActivityExportRow, ...]: + """Fold (date, team, api_key) rows into (date, team, user) rows.""" + + def bucket_of(record: _ExportRow) -> tuple[str, str, str]: + return ( + record.date, + record.entity_id or "Unassigned", + _key_metadata(api_key_metadata, record.api_key or "").user_id or "Unassigned", + ) + + key_sets: Final = MappingProxyType( + { + bucket: frozenset(record.api_key or "" for record in group) + for bucket, group in itertools.groupby(sorted(records, key=bucket_of), key=bucket_of) + } + ) + sums: Final[dict[tuple[str, str, str], _ExportMetrics]] = {} # mutable-ok: local fold accumulator + emails: Final[dict[tuple[str, str, str], str | None]] = {} # mutable-ok: local fold accumulator + for record in records: + metadata = _key_metadata(api_key_metadata, record.api_key or "") + bucket_key = bucket_of(record) + sums[bucket_key] = sums.get(bucket_key, _ExportMetrics.zero()) + _ExportMetrics.from_record(record) + emails.setdefault(bucket_key, metadata.user_email) + if emails[bucket_key] is None and metadata.user_email is not None: + emails[bucket_key] = metadata.user_email + return tuple( + _export_folded_user_row( + bucket_key, sums[bucket_key], emails[bucket_key], len(key_sets[bucket_key]), entity_metadata_field + ) + for bucket_key in sorted(sums) + ) + + +def _export_folded_user_row( + bucket_key: tuple[str, str, str], + metrics: _ExportMetrics, + user_email: str | None, + keys: int, + entity_metadata_field: Mapping[str, dict[str, object]] | None, +) -> TeamDailyActivityExportRow: + date, entity_id, user_id = bucket_key + return TeamDailyActivityExportRow( + date=date, + team_id=entity_id, + team_alias=_export_team_alias(entity_metadata_field, entity_id), + user_id=user_id if user_id != "Unassigned" else None, + user_email=user_email, + keys=keys, + spend=metrics.spend, + api_requests=metrics.api_requests, + successful_requests=metrics.successful_requests, + failed_requests=metrics.failed_requests, + total_tokens=metrics.total_tokens, + prompt_tokens=metrics.prompt_tokens, + completion_tokens=metrics.completion_tokens, + cache_read_input_tokens=metrics.cache_read_input_tokens, + cache_creation_input_tokens=metrics.cache_creation_input_tokens, + ) + + +async def get_daily_activity_export_rows( + *, + prisma_client: PrismaClient, + table_name: str, + entity_id_field: str, + entity_id: str | list[str] | None, # mutable-ok: filter union shared with the paginated path + entity_metadata_field: Mapping[str, dict[str, object]] | None, + start_date: str, + end_date: str, + api_key: str | list[str] | None, # mutable-ok: filter union shared with the paginated path + exclude_entity_ids: list[str] | None, # mutable-ok: filter union shared with the paginated path + timezone_offset_minutes: int | None, + export_type: TeamDailyActivityExportType, +) -> tuple[TeamDailyActivityExportRow, ...]: + """Every (date, entity[, api_key|model]) rollup row in the range, uncapped.""" + sql_query, sql_params = _build_export_sql_query( + table_name=table_name, + entity_id_field=entity_id_field, + entity_id=entity_id, + start_date=start_date, + end_date=end_date, + api_key=api_key, + exclude_entity_ids=exclude_entity_ids, + timezone_offset_minutes=timezone_offset_minutes, + export_type=export_type, + ) + raw_rows: Final = await _query_raw_optional(prisma_client, (sql_query, sql_params)) + records: Final = tuple(_ExportRow(**row) for row in (raw_rows or ())) + + if export_type in ("daily", "daily_with_models"): + return await asyncio.to_thread( + lambda: tuple(_export_base_row(record, entity_metadata_field) for record in records) + ) + + api_keys: Final = frozenset(record.api_key for record in records if record.api_key) + api_key_metadata: Final = ( + await get_api_key_metadata(prisma_client, api_keys, _spend_logs_window(frozenset(r.date for r in records))) + if api_keys + else _EMPTY_KEY_METADATA + ) + if export_type == "daily_with_keys": + return await asyncio.to_thread( + lambda: tuple(_export_key_row(record, entity_metadata_field, api_key_metadata) for record in records) + ) + return await asyncio.to_thread(_fold_export_users, records, entity_metadata_field, api_key_metadata) + + def _aggregate_spend_records_sync( *, records: Sequence[DailySpendRecord], diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 493c83c730a..8a6cd1218ee 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -11,6 +11,8 @@ All /team management endpoints import asyncio import copy +import csv +import io import json import math import traceback @@ -33,7 +35,8 @@ from typing import ( ) import fastapi -from fastapi import APIRouter, Depends, Header, HTTPException, Request, status +from fastapi import APIRouter, Depends, Header, HTTPException, Query, Request, Response, status +from fastapi.responses import JSONResponse from pydantic import BaseModel, JsonValue, TypeAdapter, ValidationError from typing_extensions import ReadOnly, TypedDict, assert_never @@ -126,6 +129,7 @@ from litellm.proxy.hooks.model_max_budget_limiter import ( ) from litellm.proxy.management_endpoints.common_daily_activity import ( get_daily_activity_aggregated, + get_daily_activity_export_rows, ) from litellm.proxy.management_endpoints.common_utils import ( _check_disable_global_guardrails_caller_permission, @@ -206,6 +210,11 @@ from litellm.types.proxy.management_endpoints.team_endpoints import ( BulkUpdateTeamMemberPermissionsRequest, BulkUpdateTeamMemberPermissionsResponse, GetTeamMemberPermissionsResponse, + TeamDailyActivityExportFormat, + TeamDailyActivityExportMetadata, + TeamDailyActivityExportResponse, + TeamDailyActivityExportRow, + TeamDailyActivityExportType, TeamIdSearchFilter, TeamIdSearchMatch, TeamKeyActivitySearchWhere, @@ -6809,6 +6818,178 @@ async def get_team_daily_activity_aggregated( ) +_EXPORT_CSV_METRIC_HEADERS: Final = ( + "Spend ($)", + "Requests", + "Successful Requests", + "Failed Requests", + "Total Tokens", + "Prompt Tokens", + "Completion Tokens", + "Cache Read Input Tokens", + "Cache Creation Input Tokens", +) + + +def _export_csv_headers(export_type: TeamDailyActivityExportType) -> tuple[str, ...]: + base: Final = ("Date", "Team", "Team ID") + if export_type == "daily_with_keys": + return (*base, "Key Alias", "Key ID", "User ID", "User Email", *_EXPORT_CSV_METRIC_HEADERS) + if export_type == "daily_with_users": + return (*base, "User ID", "User Email", "Keys", *_EXPORT_CSV_METRIC_HEADERS) + if export_type == "daily_with_models": + return ( + *base, + "Model", + "Spend ($)", + "Requests", + "Successful", + "Failed", + "Total Tokens", + "Prompt Tokens", + "Completion Tokens", + "Cache Read Input Tokens", + "Cache Creation Input Tokens", + ) + return (*base, *_EXPORT_CSV_METRIC_HEADERS) + + +def _csv_safe(value: str) -> str: + return "'" + value if value[:1] in ("=", "+", "-", "@", "\t", "\r") else value + + +def _export_csv_record(row: TeamDailyActivityExportRow) -> dict[str, object]: + return { # mutable-ok: csv.DictWriter consumes a plain mapping per row + "Date": row.date, + "Team": _csv_safe(row.team_alias) if row.team_alias else "-", + "Team ID": row.team_id, + "Key Alias": _csv_safe(row.key_alias) if row.key_alias else "-", + "Key ID": row.api_key or "-", + "User ID": _csv_safe(row.user_id) if row.user_id else "-", + "User Email": _csv_safe(row.user_email) if row.user_email else "-", + "Keys": row.keys, + "Model": _csv_safe(row.model) if row.model else "-", + "Spend ($)": f"{row.spend:.4f}", + "Flat Cost ($)": f"{row.flat_cost:.4f}", + "Total Cost ($)": f"{row.spend + row.flat_cost:.4f}", + "Requests": row.api_requests, + "Successful Requests": row.successful_requests, + "Failed Requests": row.failed_requests, + "Successful": row.successful_requests, + "Failed": row.failed_requests, + "Total Tokens": row.total_tokens, + "Prompt Tokens": row.prompt_tokens, + "Completion Tokens": row.completion_tokens, + "Cache Read Input Tokens": row.cache_read_input_tokens, + "Cache Creation Input Tokens": row.cache_creation_input_tokens, + } + + +def _team_export_csv(export_type: TeamDailyActivityExportType, rows: Sequence[TeamDailyActivityExportRow]) -> str: + base_headers: Final = _export_csv_headers(export_type) + spend_index: Final = base_headers.index("Spend ($)") + 1 + headers: Final = ( + (*base_headers[:spend_index], "Flat Cost ($)", "Total Cost ($)", *base_headers[spend_index:]) + if sum(row.flat_cost for row in rows) > 0 + else base_headers + ) + buffer: Final = io.StringIO() + writer: Final = csv.DictWriter(buffer, fieldnames=headers, extrasaction="ignore") + writer.writeheader() + writer.writerows(_export_csv_record(row) for row in rows) + return buffer.getvalue() + + +@router.get( + "/team/daily/activity/export", + response_model=TeamDailyActivityExportResponse, + responses={200: {"content": {"text/csv": {}, "application/json": {}}}}, # mutable-ok: OpenAPI content map + tags=["team management"], # mutable-ok: fastapi's decorator signature types tags as a list +) +async def get_team_daily_activity_export( + user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)], + start_date: str | None = None, + end_date: str | None = None, + export_type: TeamDailyActivityExportType = "daily", + format: TeamDailyActivityExportFormat = "csv", + team_id: str | None = None, + exclude_team_ids: str | None = None, + timezone_offset: Annotated[int | None, Query(alias="timezone")] = None, +) -> Response: + """ + Server-side Team Usage export, not subject to USAGE_TOP_API_KEYS_LIMIT. + + Same scoping as /team/daily/activity/aggregated, answered by one unbounded + rollup query, returned as CSV or JSON. For daily_with_keys, + daily_with_users and daily_with_models the PTU sentinel flat-cost rows are + excluded, so metadata totals under those export types cover request spend + only; the plain daily export includes them. + """ + from litellm.proxy.proxy_server import ( + prisma_client, + proxy_logging_obj, + user_api_key_cache, + ) + + if prisma_client is None: + raise _daily_activity_error(status_code=500, message=CommonProxyErrors.db_not_connected_error.value) + + range_error: Final = _aggregated_date_range_error(start_date, end_date) + if range_error is not None or start_date is None or end_date is None: + raise _daily_activity_error(status_code=400, message=range_error or "Please provide start_date and end_date") + + scope: Final = await _resolve_team_daily_activity_scope( + team_ids=team_id, + exclude_team_ids=exclude_team_ids, + api_key=None, + user_api_key_dict=user_api_key_dict, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + + rows: Final = await get_daily_activity_export_rows( + prisma_client=prisma_client, + table_name="litellm_dailyteamspend", + entity_id_field="team_id", + entity_id=scope.team_ids, + entity_metadata_field=scope.team_alias_metadata, + start_date=start_date, + end_date=end_date, + api_key=scope.api_key_filter, + exclude_entity_ids=scope.exclude_team_ids, + timezone_offset_minutes=timezone_offset, + export_type=export_type, + ) + + now: Final = datetime.now(timezone.utc) + metadata: Final = TeamDailyActivityExportMetadata( + export_date=now.isoformat(), + export_type=export_type, + start_date=start_date, + end_date=end_date, + team_ids=list(scope.team_ids) if scope.team_ids else None, # mutable-ok: response model field type + total_spend=sum(row.spend for row in rows), + total_flat_cost=sum(row.flat_cost for row in rows), + total_api_requests=sum(row.api_requests for row in rows), + total_successful_requests=sum(row.successful_requests for row in rows), + total_failed_requests=sum(row.failed_requests for row in rows), + total_tokens=sum(row.total_tokens for row in rows), + ) + + if format == "json": + return JSONResponse( + content=TeamDailyActivityExportResponse(metadata=metadata, data=rows).model_dump(mode="json") + ) + return Response( + content=_team_export_csv(export_type, rows), + media_type="text/csv; charset=utf-8", + headers={ # mutable-ok: starlette Response headers is a dict + "Content-Disposition": f'attachment; filename="team_usage_{export_type}_{now.date().isoformat()}.csv"' + }, + ) + + def _team_key_search_where(*, search: str, scope: _TeamDailyActivityScope) -> TeamKeyActivitySearchWhere: """Caller scoping lives inside the same Prisma where as the search term so `take` never trims visible matches in favour of keys the caller is not allowed to see.""" diff --git a/litellm/types/proxy/management_endpoints/team_endpoints.py b/litellm/types/proxy/management_endpoints/team_endpoints.py index aac2703e918..9b87d114a81 100644 --- a/litellm/types/proxy/management_endpoints/team_endpoints.py +++ b/litellm/types/proxy/management_endpoints/team_endpoints.py @@ -277,3 +277,48 @@ class TeamUserSpendResponse(BaseModel): start_date: str end_date: str results: tuple[TeamUserSpendRow, ...] + + +TeamDailyActivityExportType = Literal["daily", "daily_with_keys", "daily_with_users", "daily_with_models"] +TeamDailyActivityExportFormat = Literal["csv", "json"] + + +class TeamDailyActivityExportRow(BaseModel): + date: str + team_id: str + team_alias: str | None = None + api_key: str | None = None + key_alias: str | None = None + user_id: str | None = None + user_email: str | None = None + keys: int | None = None + model: str | None = None + spend: float + flat_cost: float = 0.0 + api_requests: int + successful_requests: int + failed_requests: int + total_tokens: int + prompt_tokens: int + completion_tokens: int + cache_read_input_tokens: int + cache_creation_input_tokens: int + + +class TeamDailyActivityExportMetadata(BaseModel): + export_date: str + export_type: TeamDailyActivityExportType + start_date: str + end_date: str + team_ids: list[str] | None + total_spend: float + total_flat_cost: float = 0.0 + total_api_requests: int + total_successful_requests: int + total_failed_requests: int + total_tokens: int + + +class TeamDailyActivityExportResponse(BaseModel): + metadata: TeamDailyActivityExportMetadata + data: list[TeamDailyActivityExportRow] diff --git a/terraform/provider/tools/endpointaudit/coverage_allowlist.txt b/terraform/provider/tools/endpointaudit/coverage_allowlist.txt index b0f28a9c740..7aacf7ceab9 100644 --- a/terraform/provider/tools/endpointaudit/coverage_allowlist.txt +++ b/terraform/provider/tools/endpointaudit/coverage_allowlist.txt @@ -28,6 +28,7 @@ GET /tag/user-agent/per-user-analytics GET /tag/wau GET /team/daily/activity GET /team/daily/activity/aggregated +GET /team/daily/activity/export GET /team/daily/activity/aggregated/search GET /team/spend/by_user GET /team/spend/report diff --git a/tests/integration/spend/test_team_daily_activity_export.py b/tests/integration/spend/test_team_daily_activity_export.py new file mode 100644 index 00000000000..b35d3fe0c8a --- /dev/null +++ b/tests/integration/spend/test_team_daily_activity_export.py @@ -0,0 +1,522 @@ +import csv +import io +import os +import signal +import uuid +from concurrent.futures import ThreadPoolExecutor +from datetime import datetime, timedelta, timezone +from hashlib import sha256 +from pathlib import Path +from typing import Final + +import httpx +import openai +import pytest +from integration._support.client import Gateway, Scenario, eventually, object_value, string_value +from integration._support.database import read_rows +from integration._support.process import group_members, owned_proxy, owned_proxy_process + + +def _export_range() -> dict[str, str]: + today: Final = datetime.now(timezone.utc) + return { + "start_date": (today - timedelta(days=1)).strftime("%Y-%m-%d"), + "end_date": (today + timedelta(days=1)).strftime("%Y-%m-%d"), + "timezone": "0", + } + + +def _team_with_three_keys( + gateway: Gateway, scenario: Scenario, model: str +) -> tuple[str, tuple[str, ...], tuple[str, ...], dict[str, float]]: + team: Final = scenario.team(models=[model]) + keys: Final = tuple(scenario.key(team_id=team, models=[model]) for _ in range(3)) + digests: Final = tuple(sha256(key.encode()).hexdigest() for key in keys) + for key in keys: + reply: Final = gateway.chat(model, key=key, text=f"team export {uuid.uuid4().hex}") + assert reply["usage"]["total_tokens"] == 40, reply + daily: Final = eventually( + lambda: read_rows('SELECT api_key, spend FROM "LiteLLM_DailyTeamSpend" WHERE team_id=%s', (team,)), + lambda values: len({row["api_key"] for row in values}) == 3, + seconds=70, + ) + spend_by_key: Final = {row["api_key"]: float(row["spend"]) for row in daily} + return team, keys, digests, spend_by_key + + +def _export_json(gateway: Gateway, **params: str) -> httpx.Response: + return gateway.request("GET", "/team/daily/activity/export", params={**_export_range(), **params}) + + +def test_team_activity_export_returns_every_key_beyond_the_top_n_cap(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + model: Final = scenario.model(input_cost_per_token=0.001, output_cost_per_token=0.002) + team: Final = scenario.team(models=[model]) + keys: Final = tuple(scenario.key(team_id=team, models=[model]) for _ in range(3)) + digests: Final = tuple(sha256(key.encode()).hexdigest() for key in keys) + for key in keys: + reply: Final = gateway.chat(model, key=key, text=f"team export {uuid.uuid4().hex}") + assert reply["usage"]["total_tokens"] == 40, reply + daily: Final = eventually( + lambda: read_rows('SELECT api_key, spend FROM "LiteLLM_DailyTeamSpend" WHERE team_id=%s', (team,)), + lambda values: len({row["api_key"] for row in values}) == 3, + seconds=70, + ) + spend_by_key: Final = {row["api_key"]: float(row["spend"]) for row in daily} + response: Final = gateway.request( + "GET", + "/team/daily/activity/export", + params={ + **_export_range(), + "team_id": team, + "export_type": "daily_with_keys", + "format": "json", + }, + ) + assert response.status_code == 200, response.text + body: Final = object_value(response.json()) + rows: Final = tuple(object_value(row) for row in body["data"]) + assert sorted(string_value(row["api_key"]) for row in rows) == sorted(digests), response.text + for row in rows: + assert row["team_id"] == team, response.text + assert float(row["spend"]) == pytest.approx(spend_by_key[string_value(row["api_key"])]), response.text + metadata: Final = object_value(body["metadata"]) + assert ( + metadata["export_type"], + metadata["team_ids"], + metadata["total_api_requests"], + metadata["total_successful_requests"], + metadata["total_failed_requests"], + ) == ("daily_with_keys", [team], 3, 3, 0), response.text + assert float(metadata["total_spend"]) == pytest.approx(sum(spend_by_key.values())), response.text + + +def test_team_activity_export_csv_downloads_every_key(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + model: Final = scenario.model(input_cost_per_token=0.001, output_cost_per_token=0.002) + team: Final = scenario.team(models=[model]) + keys: Final = tuple(scenario.key(team_id=team, models=[model]) for _ in range(3)) + digests: Final = tuple(sha256(key.encode()).hexdigest() for key in keys) + for key in keys: + reply: Final = gateway.chat(model, key=key, text=f"team export {uuid.uuid4().hex}") + assert reply["usage"]["total_tokens"] == 40, reply + daily: Final = eventually( + lambda: read_rows('SELECT api_key, spend FROM "LiteLLM_DailyTeamSpend" WHERE team_id=%s', (team,)), + lambda values: len({row["api_key"] for row in values}) == 3, + seconds=70, + ) + spend_by_key: Final = {row["api_key"]: float(row["spend"]) for row in daily} + response: Final = gateway.request( + "GET", + "/team/daily/activity/export", + params={ + **_export_range(), + "team_id": team, + "export_type": "daily_with_keys", + "format": "csv", + }, + ) + assert response.status_code == 200, response.text + assert response.headers["content-type"].startswith("text/csv"), response.headers + assert "attachment" in response.headers["content-disposition"], response.headers + records: Final = tuple(csv.DictReader(io.StringIO(response.text))) + assert len(records) == 3, response.text + assert sorted(record["Key ID"] for record in records) == sorted(digests), response.text + assert sorted(record["Team ID"] for record in records) == [team, team, team], response.text + for record in records: + assert record["Spend ($)"] == f"{spend_by_key[record['Key ID']]:.4f}", response.text + + +def test_team_activity_export_denies_a_member_another_team(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + model: Final = scenario.model(input_cost_per_token=0.001, output_cost_per_token=0.002) + team_a: Final = scenario.team(models=[model]) + team_b: Final = scenario.team(models=[model]) + member: Final = scenario.user(user_role="internal_user", teams=[team_a]) + member_key: Final = scenario.key(user_id=member, team_id=team_a, models=[model]) + reply: Final = gateway.chat(model, key=member_key, text=f"team export {uuid.uuid4().hex}") + assert reply["usage"]["total_tokens"] == 40, reply + daily: Final = eventually( + lambda: read_rows('SELECT api_key, spend FROM "LiteLLM_DailyTeamSpend" WHERE team_id=%s', (team_a,)), + lambda values: len(values) == 1, + seconds=70, + ) + denied: Final = gateway.request( + "GET", + "/team/daily/activity/export", + params={**_export_range(), "team_id": team_b, "export_type": "daily", "format": "json"}, + key=member_key, + ) + assert denied.status_code == 404, denied.text + assert f"User does not belong to Team= {team_b}" in denied.text, denied.text + allowed: Final = gateway.request( + "GET", + "/team/daily/activity/export", + params={**_export_range(), "team_id": team_a, "export_type": "daily", "format": "json"}, + key=member_key, + ) + assert allowed.status_code == 200, allowed.text + rows: Final = tuple(object_value(row) for row in object_value(allowed.json())["data"]) + assert len(rows) == 1, allowed.text + assert rows[0]["team_id"] == team_a, allowed.text + assert float(rows[0]["spend"]) == pytest.approx(float(daily[0]["spend"])), allowed.text + + +def test_export_daily_total_matches_the_capped_aggregated_team_spend(gateway: Gateway, tmp_path: Path) -> None: + with owned_proxy(gateway, tmp_path, {"USAGE_TOP_API_KEYS_LIMIT": "2"}, workers=2) as candidate: + with candidate.scenario() as scenario: + model: Final = scenario.model(input_cost_per_token=0.001, output_cost_per_token=0.002) + team, keys, digests, spend_by_key = _team_with_three_keys(candidate, scenario, model) + aggregated: Final = candidate.request( + "GET", + "/team/daily/activity/aggregated", + params={**_export_range(), "team_ids": team}, + ) + assert aggregated.status_code == 200, aggregated.text + body: Final = object_value(aggregated.json()) + metadata: Final = object_value(body["metadata"]) + assert metadata["api_key_limit"] == 2, aggregated.text + assert metadata["total_api_keys"] == 3, aggregated.text + day: Final = object_value(body["results"][0]) + breakdown: Final = object_value(day["breakdown"]) + assert len(object_value(breakdown["api_keys"])) == 2, aggregated.text + team_spend: Final = float( + object_value(object_value(object_value(breakdown["entities"])[team])["metrics"])["spend"] + ) + + response: Final = _export_json(candidate, team_id=team, export_type="daily", format="json") + assert response.status_code == 200, response.text + rows: Final = tuple(object_value(row) for row in object_value(response.json())["data"]) + assert len(rows) == 1, response.text + assert rows[0]["team_id"] == team, response.text + assert float(rows[0]["spend"]) == pytest.approx(team_spend), response.text + assert float(rows[0]["spend"]) == pytest.approx(sum(spend_by_key.values())), response.text + + +def test_export_users_folds_spend_per_user_and_leaves_keyless_keys_unassigned(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + model: Final = scenario.model(input_cost_per_token=0.001, output_cost_per_token=0.002) + team: Final = scenario.team(models=[model]) + user_a: Final = scenario.user(user_role="internal_user", teams=[team]) + user_b: Final = scenario.user(user_role="internal_user", teams=[team]) + key_a: Final = scenario.key(team_id=team, user_id=user_a, models=[model]) + key_b: Final = scenario.key(team_id=team, user_id=user_b, models=[model]) + key_none: Final = scenario.key(team_id=team, models=[model]) + for key in (key_a, key_b, key_none): + reply: Final = gateway.chat(model, key=key, text=f"team export {uuid.uuid4().hex}") + assert reply["usage"]["total_tokens"] == 40, reply + daily: Final = eventually( + lambda: read_rows('SELECT api_key, spend FROM "LiteLLM_DailyTeamSpend" WHERE team_id=%s', (team,)), + lambda values: len({row["api_key"] for row in values}) == 3, + seconds=70, + ) + spend_by_key: Final = {row["api_key"]: float(row["spend"]) for row in daily} + response: Final = _export_json(gateway, team_id=team, export_type="daily_with_users", format="json") + assert response.status_code == 200, response.text + rows: Final = tuple(object_value(row) for row in object_value(response.json())["data"]) + by_user: Final = {row["user_id"]: row for row in rows} + assert by_user[user_a]["spend"] == pytest.approx(spend_by_key[sha256(key_a.encode()).hexdigest()]), ( + response.text + ) + assert by_user[user_b]["spend"] == pytest.approx(spend_by_key[sha256(key_b.encode()).hexdigest()]), ( + response.text + ) + assert None in by_user, response.text + assert by_user[None]["spend"] == pytest.approx(spend_by_key[sha256(key_none.encode()).hexdigest()]), ( + response.text + ) + metadata: Final = object_value(object_value(response.json())["metadata"]) + assert float(metadata["total_spend"]) == pytest.approx(sum(spend_by_key.values())), response.text + + +def test_export_models_reports_one_row_per_model_with_matching_spend(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + upstream_a: Final = f"openai/export-{uuid.uuid4().hex}" + upstream_b: Final = f"openai/export-{uuid.uuid4().hex}" + model_a: Final = scenario.model(model=upstream_a, input_cost_per_token=0.001, output_cost_per_token=0.002) + model_b: Final = scenario.model(model=upstream_b, input_cost_per_token=0.0005, output_cost_per_token=0.001) + upstream_models: Final = (upstream_a, upstream_b) + team: Final = scenario.team(models=[model_a, model_b]) + key: Final = scenario.key(team_id=team, models=[model_a, model_b]) + for model in (model_a, model_b): + reply: Final = gateway.chat(model, key=key, text=f"team export {uuid.uuid4().hex}") + assert reply["usage"]["total_tokens"] == 40, reply + daily: Final = eventually( + lambda: read_rows('SELECT model, spend FROM "LiteLLM_DailyTeamSpend" WHERE team_id=%s', (team,)), + lambda values: len({row["model"] for row in values}) == 2, + seconds=70, + ) + spend_by_model: Final = {row["model"]: float(row["spend"]) for row in daily} + + response: Final = _export_json(gateway, team_id=team, export_type="daily_with_models", format="json") + assert response.status_code == 200, response.text + rows: Final = tuple(object_value(row) for row in object_value(response.json())["data"]) + assert {row["model"] for row in rows} == set(upstream_models), response.text + for row in rows: + assert float(row["spend"]) == pytest.approx(spend_by_model[row["model"]]), response.text + + csv_response: Final = _export_json(gateway, team_id=team, export_type="daily_with_models", format="csv") + assert csv_response.status_code == 200, csv_response.text + records: Final = tuple(csv.DictReader(io.StringIO(csv_response.text))) + assert sorted(record["Model"] for record in records) == sorted(upstream_models), csv_response.text + + +def test_export_without_team_id_returns_only_the_callers_teams(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + model: Final = scenario.model(input_cost_per_token=0.001, output_cost_per_token=0.002) + team_a: Final = scenario.team(models=[model]) + team_b: Final = scenario.team(models=[model]) + member: Final = scenario.user(user_role="internal_user", teams=[team_a]) + member_key: Final = scenario.key(user_id=member, team_id=team_a, models=[model]) + other_key: Final = scenario.key(team_id=team_b, models=[model]) + reply: Final = gateway.chat(model, key=member_key, text=f"team export {uuid.uuid4().hex}") + assert reply["usage"]["total_tokens"] == 40, reply + reply_b: Final = gateway.chat(model, key=other_key, text=f"team export {uuid.uuid4().hex}") + assert reply_b["usage"]["total_tokens"] == 40, reply_b + eventually( + lambda: read_rows('SELECT team_id FROM "LiteLLM_DailyTeamSpend" WHERE team_id=%s', (team_b,)), + lambda values: len(values) == 1, + seconds=70, + ) + eventually( + lambda: read_rows('SELECT team_id FROM "LiteLLM_DailyTeamSpend" WHERE team_id=%s', (team_a,)), + lambda values: len(values) == 1, + seconds=70, + ) + + response: Final = gateway.request( + "GET", + "/team/daily/activity/export", + params={**_export_range(), "export_type": "daily", "format": "json"}, + key=member_key, + ) + assert response.status_code == 200, response.text + rows: Final = tuple(object_value(row) for row in object_value(response.json())["data"]) + assert len(rows) == 1, response.text + assert rows[0]["team_id"] == team_a, response.text + + +def test_export_rejects_requests_without_a_valid_key(gateway: Gateway) -> None: + params: Final = {**_export_range(), "export_type": "daily", "format": "json"} + anonymous: Final = gateway.client.get("/team/daily/activity/export", params=params) + assert anonymous.status_code == 401, anonymous.text + garbage: Final = gateway.request("GET", "/team/daily/activity/export", params=params, key="sk-nope") + assert garbage.status_code == 401, garbage.text + + +def test_export_rejects_bad_parameters(gateway: Gateway) -> None: + weekly: Final = _export_json(gateway, export_type="weekly", format="json") + assert weekly.status_code == 422, weekly.text + xml: Final = _export_json(gateway, export_type="daily", format="xml") + assert xml.status_code == 422, xml.text + no_dates: Final = gateway.request( + "GET", "/team/daily/activity/export", params={"export_type": "daily", "format": "json"} + ) + assert no_dates.status_code == 400, no_dates.text + assert "start_date and end_date" in no_dates.text, no_dates.text + reversed_range: Final = gateway.request( + "GET", + "/team/daily/activity/export", + params={"start_date": "2026-09-25", "end_date": "2026-09-23", "export_type": "daily", "format": "json"}, + ) + assert reversed_range.status_code == 400, reversed_range.text + assert "end_date must be on or after start_date" in reversed_range.text, reversed_range.text + bad_date: Final = gateway.request( + "GET", + "/team/daily/activity/export", + params={"start_date": "2026-13-40", "end_date": "2026-12-31", "export_type": "daily", "format": "json"}, + ) + assert bad_date.status_code == 400, bad_date.text + assert "valid YYYY-MM-DD" in bad_date.text, bad_date.text + + +def test_export_of_a_team_without_spend_returns_empty(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + model: Final = scenario.model(input_cost_per_token=0.001, output_cost_per_token=0.002) + team: Final = scenario.team(models=[model]) + fresh: Final = _export_json(gateway, team_id=team, export_type="daily", format="json") + assert fresh.status_code == 200, fresh.text + body: Final = object_value(fresh.json()) + assert body["data"] == [], fresh.text + assert float(object_value(body["metadata"])["total_spend"]) == 0, fresh.text + unknown: Final = _export_json(gateway, team_id=str(uuid.uuid4()), export_type="daily", format="json") + assert unknown.status_code == 200, unknown.text + unknown_body: Final = object_value(unknown.json()) + assert unknown_body["data"] == [], unknown.text + assert float(object_value(unknown_body["metadata"])["total_spend"]) == 0, unknown.text + + +def test_export_csv_is_deterministic_and_omits_flat_cost_without_ptu(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + model: Final = scenario.model(input_cost_per_token=0.001, output_cost_per_token=0.002) + _team_with_three_keys(gateway, scenario, model) + params: Final = {**_export_range(), "export_type": "daily_with_keys", "format": "csv"} + first: Final = gateway.request("GET", "/team/daily/activity/export", params=params) + second: Final = gateway.request("GET", "/team/daily/activity/export", params=params) + assert first.status_code == 200 and second.status_code == 200, first.text + assert first.text == second.text, "daily_with_keys csv is not byte-identical across calls" + header: Final = first.text.splitlines()[0] + assert "Flat Cost" not in header and "Total Cost" not in header, header + + +def test_export_csv_escapes_formula_like_key_aliases(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + model: Final = scenario.model(input_cost_per_token=0.001, output_cost_per_token=0.002) + team: Final = scenario.team(models=[model]) + alias: Final = f'=HYPERLINK("http://x.{uuid.uuid4().hex}","x")' + keys: Final = ( + scenario.key(team_id=team, models=[model], key_alias=alias), + scenario.key(team_id=team, models=[model]), + ) + for key in keys: + reply: Final = gateway.chat(model, key=key, text=f"team export {uuid.uuid4().hex}") + assert reply["usage"]["total_tokens"] == 40, reply + digests: Final = tuple(sha256(key.encode()).hexdigest() for key in keys) + eventually( + lambda: read_rows('SELECT api_key FROM "LiteLLM_DailyTeamSpend" WHERE team_id=%s', (team,)), + lambda values: len({row["api_key"] for row in values}) == 2, + seconds=70, + ) + response: Final = _export_json(gateway, team_id=team, export_type="daily_with_keys", format="csv") + assert response.status_code == 200, response.text + records: Final = {record["Key ID"]: record for record in csv.DictReader(io.StringIO(response.text))} + assert records[digests[0]]["Key Alias"] == "'" + alias, response.text + assert records[digests[1]]["Key Alias"] == "-", response.text + + +def test_aggregated_route_keeps_the_top_n_key_cap(gateway: Gateway, tmp_path: Path) -> None: + with owned_proxy(gateway, tmp_path, {"USAGE_TOP_API_KEYS_LIMIT": "2"}, workers=2) as candidate: + with candidate.scenario() as scenario: + model: Final = scenario.model(input_cost_per_token=0.001, output_cost_per_token=0.002) + team, keys, digests, spend_by_key = _team_with_three_keys(candidate, scenario, model) + response: Final = candidate.request( + "GET", + "/team/daily/activity/aggregated", + params={**_export_range(), "team_ids": team}, + ) + assert response.status_code == 200, response.text + body: Final = object_value(response.json()) + metadata: Final = object_value(body["metadata"]) + assert metadata["api_key_limit"] == 2, response.text + assert metadata["total_api_keys"] == 3, response.text + breakdown: Final = object_value(object_value(body["results"][0])["breakdown"]) + assert len(object_value(breakdown["api_keys"])) == 2, response.text + + +def test_paginated_team_daily_activity_still_lists_the_team(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + model: Final = scenario.model(input_cost_per_token=0.001, output_cost_per_token=0.002) + team, keys, digests, spend_by_key = _team_with_three_keys(gateway, scenario, model) + response: Final = gateway.request( + "GET", + "/team/daily/activity", + params={ + "team_ids": team, + "start_date": _export_range()["start_date"], + "end_date": _export_range()["end_date"], + }, + ) + assert response.status_code == 200, response.text + results: Final = object_value(response.json())["results"] + assert isinstance(results, list), response.text + days: Final = tuple( + object_value(day) + for day in results + if team in object_value(object_value(object_value(day)["breakdown"])["entities"]) + ) + assert len(days) == 1, response.text + entity: Final = object_value(object_value(object_value(days[0]["breakdown"])["entities"])[team]) + assert float(object_value(entity["metrics"])["spend"]) == pytest.approx(sum(spend_by_key.values())), ( + response.text + ) + + +def test_openai_sdk_chat_still_lands_one_spend_log(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + model: Final = scenario.model(input_cost_per_token=0.001, output_cost_per_token=0.002) + team: Final = scenario.team(models=[model]) + key: Final = scenario.key(team_id=team, models=[model]) + client: Final = openai.OpenAI(base_url=f"{gateway.client.base_url}/v1", api_key=key, max_retries=0) + reply: Final = client.chat.completions.create( + model=model, messages=[{"role": "user", "content": f"sdk {uuid.uuid4().hex}"}], stream=False + ) + rows: Final = eventually( + lambda: read_rows('SELECT request_id FROM "LiteLLM_SpendLogs" WHERE request_id=%s', (reply.id,)), + lambda values: len(values) == 1, + seconds=70, + ) + assert len(rows) == 1 and rows[0]["request_id"] == reply.id, rows + + +def test_export_and_chat_burst_survives_worker_kill(gateway: Gateway, tmp_path: Path) -> None: + with owned_proxy_process(gateway, tmp_path, {"USAGE_TOP_API_KEYS_LIMIT": "2"}, workers=2) as owned: + candidate: Final = owned.gateway + with candidate.scenario() as scenario: + model: Final = scenario.model(input_cost_per_token=0.001, output_cost_per_token=0.002) + team, keys, digests, spend_by_key = _team_with_three_keys(candidate, scenario, model) + + workers: Final = eventually( + lambda: tuple(member for member in group_members(owned.process.pid) if member.pid != owned.process.pid), + lambda members: len(members) >= 2, + seconds=30, + ) + assert len(workers) >= 2, workers + + params: Final = { + **_export_range(), + "team_id": team, + "export_type": "daily_with_keys", + "format": "json", + } + + def burst(tag: str) -> tuple[tuple[httpx.Response, ...], tuple[httpx.Response, ...]]: + with ThreadPoolExecutor(max_workers=30) as pool: + futures: Final = tuple( + ( + pool.submit( + candidate.request, + "POST", + "/v1/chat/completions", + { + "model": model, + "messages": [{"role": "user", "content": f"{tag}-{index}-{uuid.uuid4().hex}"}], + }, + key=keys[index % 3], + ) + if index % 2 == 0 + else pool.submit(candidate.request, "GET", "/team/daily/activity/export", params=params) + ) + for index in range(30) + ) + results: Final = tuple(future.result() for future in futures) + return results[0::2], results[1::2] + + chat_a, export_a = burst("bursta") + assert all(response.status_code == 200 for response in chat_a), [r.text for r in chat_a] + assert all(response.status_code == 200 for response in export_a), [r.text for r in export_a] + + victim: Final = workers[0] + os.kill(victim.pid, signal.SIGKILL) + + chat_b, export_b = burst("burstb") + all_chats: Final = chat_a + chat_b + all_exports: Final = export_a + export_b + assert all(response.status_code == 200 for response in all_chats), [ + (r.status_code, r.text) for r in all_chats + ] + for response in all_exports: + assert response.status_code == 200, response.text + returned: Final = {string_value(row["api_key"]) for row in object_value(response.json())["data"]} + assert returned == set(digests), response.text + chat_ids: Final = tuple(string_value(object_value(r.json())["id"]) for r in all_chats) + assert len(set(chat_ids)) == 30 + id_slots: Final = ", ".join("%s" for _ in chat_ids) + rows: Final = eventually( + lambda: read_rows( + f'SELECT request_id, COUNT(*)::int AS n FROM "LiteLLM_SpendLogs" WHERE request_id IN ({id_slots}) GROUP BY request_id', + chat_ids, + ), + lambda values: len(values) == 30, + seconds=70, + ) + assert all(row["n"] == 1 for row in rows), rows diff --git a/tests/proxy_behavior/management/test_team_daily_activity.py b/tests/proxy_behavior/management/test_team_daily_activity.py index 9bbc8fdde29..f85eb12c403 100644 --- a/tests/proxy_behavior/management/test_team_daily_activity.py +++ b/tests/proxy_behavior/management/test_team_daily_activity.py @@ -48,8 +48,9 @@ _DATES = "start_date=2024-01-01&end_date=2024-12-31" "/team/daily/activity", "/team/daily/activity/aggregated", "/team/daily/activity/aggregated/search", + "/team/daily/activity/export", ), - ids=("paginated", "aggregated", "search"), + ids=("paginated", "aggregated", "search", "export"), ) @pytest.mark.parametrize( "actor,team,expected_status", @@ -59,16 +60,15 @@ _DATES = "start_date=2024-01-01&end_date=2024-12-31" async def test_team_daily_activity_matrix( actor: Actor, team: str, expected_status: int, endpoint: str, proxy_client, world ): + filter_param = "team_id" if endpoint.endswith("/export") else "team_ids" query = _DATES + ("&search=x" if endpoint.endswith("/search") else "") if team == "alpha": - query += f"&team_ids={world.team_alpha_id}" + query += f"&{filter_param}={world.team_alpha_id}" elif team == "beta": - query += f"&team_ids={world.team_beta_id}" + query += f"&{filter_param}={world.team_beta_id}" resp = await proxy_client.get( f"{endpoint}?{query}", headers={"Authorization": f"Bearer {world.keys[actor].cleartext}"}, ) - assert ( - resp.status_code == expected_status - ), f"{actor.value} -> {team}: {resp.status_code} {resp.text}" + assert resp.status_code == expected_status, f"{actor.value} -> {team}: {resp.status_code} {resp.text}" diff --git a/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py b/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py index baaf3f4ba2f..cecdb937c40 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py +++ b/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py @@ -25,6 +25,7 @@ from litellm.proxy.management_endpoints.common_daily_activity import ( get_api_key_metadata, get_daily_activity, get_daily_activity_aggregated, + get_daily_activity_export_rows, global_rollup_reconciled_through, update_metrics, ) @@ -2868,3 +2869,307 @@ def test_spend_logs_window_is_none_when_no_date_parses(): from litellm.proxy.management_endpoints.common_daily_activity import _spend_logs_window assert _spend_logs_window({"garbage", ""}) is None + + +_DAILY_TEAM_SPEND_DDL: Final = """ + CREATE TABLE "LiteLLM_DailyTeamSpend" ( + id TEXT PRIMARY KEY, + team_id TEXT, + date TEXT NOT NULL, + api_key TEXT NOT NULL, + model TEXT, + model_group TEXT, + custom_llm_provider TEXT, + mcp_namespaced_tool_name TEXT, + endpoint TEXT, + prompt_tokens BIGINT DEFAULT 0, + completion_tokens BIGINT DEFAULT 0, + cache_read_input_tokens BIGINT DEFAULT 0, + cache_creation_input_tokens BIGINT DEFAULT 0, + compression_saved_tokens BIGINT DEFAULT 0, + compression_savings_spend DOUBLE PRECISION DEFAULT 0, + prompt_caching_savings_spend DOUBLE PRECISION DEFAULT 0, + gateway_injected_caching_savings_spend DOUBLE PRECISION DEFAULT 0, + autorouter_savings_spend DOUBLE PRECISION DEFAULT 0, + spend DOUBLE PRECISION DEFAULT 0, + ptu_flat_cost DOUBLE PRECISION DEFAULT 0, + api_requests BIGINT DEFAULT 0, + successful_requests BIGINT DEFAULT 0, + failed_requests BIGINT DEFAULT 0, + total_response_time_ms BIGINT DEFAULT 0, + timed_requests BIGINT DEFAULT 0 + ) +""" + + +def _seed_daily_team_spend(conn: psycopg.Connection, rows: Sequence[tuple[object, ...]]) -> None: + with conn.cursor() as cur: + cur.execute(_DAILY_TEAM_SPEND_DDL) + cur.executemany( + """ + INSERT INTO "LiteLLM_DailyTeamSpend" + (id, team_id, date, api_key, model, model_group, custom_llm_provider, + endpoint, prompt_tokens, spend, ptu_flat_cost, api_requests, successful_requests) + VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s) + """, + rows, + ) + conn.commit() + + +def _team_spend_row( + row_id: str, + team_id: str, + api_key: str, + spend: float, + *, + date: str = "2026-06-01", + model: str = "gpt-5", + ptu_flat_cost: float = 0.0, +) -> tuple[object, ...]: + return ( + row_id, + team_id, + date, + api_key, + model, + "", + "openai", + "/v1/chat/completions", + 10, + spend, + ptu_flat_cost, + 1, + 1, + ) + + +def _export_prisma(conn: psycopg.Connection, token_rows: Sequence[SimpleNamespace] = ()) -> MagicMock: + mock_prisma = MagicMock() + mock_prisma.db = MagicMock() + mock_prisma.db.query_raw = _psycopg_query_raw(conn, []) + mock_prisma.db.litellm_verificationtoken.find_many = AsyncMock(return_value=list(token_rows)) + mock_prisma.db.litellm_deletedverificationtoken.find_many = AsyncMock(return_value=[]) + mock_prisma.db.litellm_usertable.find_many = AsyncMock(return_value=[]) + return mock_prisma + + +@pytest.mark.asyncio +async def test_export_keys_returns_every_key_beyond_the_top_n_cap( + _aggregated_postgresql: psycopg.Connection, +): + """The export route exists because the aggregated route caps the per-key arm at + USAGE_TOP_API_KEYS_LIMIT. With more keys than the cap every one of them must + land in the export, while the PTU sentinel stays out of the key view.""" + n_keys: Final = USAGE_TOP_API_KEYS_LIMIT + 7 + _seed_daily_team_spend( + _aggregated_postgresql, + [ + *[_team_spend_row(f"row-{i:03d}", "team-1", f"key-{i:03d}", float(i + 1)) for i in range(n_keys)], + _team_spend_row("row-ptu", "team-1", PTU_SENTINEL_API_KEY, 0.0, ptu_flat_cost=1000.0), + ], + ) + + rows = await get_daily_activity_export_rows( + prisma_client=_export_prisma(_aggregated_postgresql), + table_name="litellm_dailyteamspend", + entity_id_field="team_id", + entity_id="team-1", + entity_metadata_field=None, + start_date="2026-06-01", + end_date="2026-06-01", + api_key=None, + exclude_entity_ids=None, + timezone_offset_minutes=None, + export_type="daily_with_keys", + ) + + assert {row.api_key for row in rows} == {f"key-{i:03d}" for i in range(n_keys)} + assert len(rows) == n_keys + assert all(row.team_id == "team-1" for row in rows) + by_key: Final = {row.api_key: row for row in rows} + assert by_key["key-000"].spend == pytest.approx(1.0) + assert sum(row.spend for row in rows) == pytest.approx(n_keys * (n_keys + 1) / 2) + assert all(row.total_tokens == 10 and row.api_requests == 1 for row in rows) + + +@pytest.mark.asyncio +async def test_export_daily_keeps_ptu_sentinel_in_the_team_rollup( + _aggregated_postgresql: psycopg.Connection, +): + """The plain daily export groups by (date, team), so the sentinel's flat cost + must land in the team row exactly like breakdown.entities on the aggregated + route; dropping it would silently under-report team spend.""" + _seed_daily_team_spend( + _aggregated_postgresql, + [ + _team_spend_row("row-1", "team-1", "key-1", 2.0), + _team_spend_row("row-ptu", "team-1", PTU_SENTINEL_API_KEY, 0.0, ptu_flat_cost=0.0), + ], + ) + with _aggregated_postgresql.cursor() as cur: + cur.execute("UPDATE \"LiteLLM_DailyTeamSpend\" SET spend = 1000.0 WHERE id = 'row-ptu'") + _aggregated_postgresql.commit() + + rows = await get_daily_activity_export_rows( + prisma_client=_export_prisma(_aggregated_postgresql), + table_name="litellm_dailyteamspend", + entity_id_field="team_id", + entity_id="team-1", + entity_metadata_field={"team-1": {"team_alias": "Alpha"}}, + start_date="2026-06-01", + end_date="2026-06-01", + api_key=None, + exclude_entity_ids=None, + timezone_offset_minutes=None, + export_type="daily", + ) + + assert len(rows) == 1 + assert rows[0].team_id == "team-1" + assert rows[0].team_alias == "Alpha" + assert rows[0].api_key is None + assert rows[0].spend == pytest.approx(1002.0) + + +@pytest.mark.asyncio +async def test_export_users_folds_keys_into_one_row_per_user( + _aggregated_postgresql: psycopg.Connection, +): + """daily_with_users runs the per-key rollup then folds in Python: two keys of + user-1 merge into one row with keys=2 and summed metrics, and the distinct + user keeps its own row.""" + _seed_daily_team_spend( + _aggregated_postgresql, + [ + _team_spend_row("row-1", "team-1", "key-1", 2.0), + _team_spend_row("row-2", "team-1", "key-2", 3.0), + _team_spend_row("row-3", "team-1", "key-3", 5.0), + ], + ) + tokens: Final = tuple( + SimpleNamespace(token=token, key_alias=None, team_id="team-1", user_id=user_id) + for token, user_id in (("key-1", "user-1"), ("key-2", "user-1"), ("key-3", "user-2")) + ) + + rows = await get_daily_activity_export_rows( + prisma_client=_export_prisma(_aggregated_postgresql, tokens), + table_name="litellm_dailyteamspend", + entity_id_field="team_id", + entity_id="team-1", + entity_metadata_field=None, + start_date="2026-06-01", + end_date="2026-06-01", + api_key=None, + exclude_entity_ids=None, + timezone_offset_minutes=None, + export_type="daily_with_users", + ) + + assert [(row.user_id, row.keys, row.spend, row.api_requests, row.total_tokens) for row in rows] == [ + ("user-1", 2, 5.0, 2, 20), + ("user-2", 1, 5.0, 1, 10), + ] + + +@pytest.mark.asyncio +async def test_export_models_rolls_up_per_team_and_model( + _aggregated_postgresql: psycopg.Connection, +): + _seed_daily_team_spend( + _aggregated_postgresql, + [ + _team_spend_row("row-1", "team-1", "key-1", 2.0, model="gpt-5"), + _team_spend_row("row-2", "team-1", "key-2", 3.0, model="gpt-5"), + _team_spend_row("row-3", "team-1", "key-1", 5.0, model="claude"), + ], + ) + + rows = await get_daily_activity_export_rows( + prisma_client=_export_prisma(_aggregated_postgresql), + table_name="litellm_dailyteamspend", + entity_id_field="team_id", + entity_id="team-1", + entity_metadata_field=None, + start_date="2026-06-01", + end_date="2026-06-01", + api_key=None, + exclude_entity_ids=None, + timezone_offset_minutes=None, + export_type="daily_with_models", + ) + + assert [(row.model, row.spend, row.api_requests) for row in rows] == [ + ("claude", 5.0, 1), + ("gpt-5", 5.0, 2), + ] + + +@pytest.mark.asyncio +async def test_export_daily_reports_ptu_flat_cost_on_the_team_row( + _aggregated_postgresql: psycopg.Connection, ptu_cost_attribution_enabled +): + """The CSV the dashboard hands to finance must match the client-side export, + which shows flat cost columns once any PTU spend exists for the day.""" + from litellm.proxy.management_endpoints.team_endpoints import _team_export_csv + + _seed_daily_team_spend( + _aggregated_postgresql, + [ + _team_spend_row("row-1", "team-1", "key-1", 2.0), + _team_spend_row("row-ptu", "team-1", PTU_SENTINEL_API_KEY, 0.0, ptu_flat_cost=240.0), + ], + ) + + rows = await get_daily_activity_export_rows( + prisma_client=_export_prisma(_aggregated_postgresql), + table_name="litellm_dailyteamspend", + entity_id_field="team_id", + entity_id="team-1", + entity_metadata_field=None, + start_date="2026-06-01", + end_date="2026-06-01", + api_key=None, + exclude_entity_ids=None, + timezone_offset_minutes=None, + export_type="daily", + ) + + assert len(rows) == 1 + assert rows[0].flat_cost == pytest.approx(240.0) + header: Final = _team_export_csv("daily", rows).splitlines()[0] + assert "Spend ($),Flat Cost ($),Total Cost ($)" in header + record: Final = _team_export_csv("daily", rows).splitlines()[1].split(",") + spend_index: Final = header.split(",").index("Spend ($)") + assert record[spend_index : spend_index + 3] == ["2.0000", "240.0000", "242.0000"] + + +@pytest.mark.asyncio +async def test_export_csv_omits_flat_cost_columns_when_no_ptu_spend_exists( + _aggregated_postgresql: psycopg.Connection, +): + from litellm.proxy.management_endpoints.team_endpoints import _team_export_csv + + _seed_daily_team_spend( + _aggregated_postgresql, + [_team_spend_row("row-1", "team-1", "key-1", 2.0)], + ) + + rows = await get_daily_activity_export_rows( + prisma_client=_export_prisma(_aggregated_postgresql), + table_name="litellm_dailyteamspend", + entity_id_field="team_id", + entity_id="team-1", + entity_metadata_field=None, + start_date="2026-06-01", + end_date="2026-06-01", + api_key=None, + exclude_entity_ids=None, + timezone_offset_minutes=None, + export_type="daily", + ) + + assert rows[0].flat_cost == 0.0 + header: Final = _team_export_csv("daily", rows).splitlines()[0] + assert "Flat Cost" not in header + assert "Total Cost" not in header diff --git a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py index 38241926f8e..6d902cb7fec 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -17101,3 +17101,130 @@ def test_list_team_v2_answers_503_no_db_connection_when_the_callers_user_read_hi assert response.status_code == 503, response.text assert response.json() == _DB_OUTAGE_503_BODY + + +def test_team_export_csv_columns_match_the_dashboard_client_layout(): + import csv + import io + + from litellm.proxy.management_endpoints.team_endpoints import _team_export_csv + from litellm.types.proxy.management_endpoints.team_endpoints import TeamDailyActivityExportRow + + row: Final = TeamDailyActivityExportRow( + date="2026-06-01", + team_id="team-1", + team_alias=None, + api_key="key-1", + key_alias="key-alias-1", + user_id="user-1", + user_email="u@example.com", + spend=1.5, + api_requests=2, + successful_requests=2, + failed_requests=0, + total_tokens=30, + prompt_tokens=20, + completion_tokens=10, + cache_read_input_tokens=5, + cache_creation_input_tokens=4, + ) + + records: Final = list(csv.DictReader(io.StringIO(_team_export_csv("daily_with_keys", (row,))))) + + assert records == [ + { + "Date": "2026-06-01", + "Team": "-", + "Team ID": "team-1", + "Key Alias": "key-alias-1", + "Key ID": "key-1", + "User ID": "user-1", + "User Email": "u@example.com", + "Spend ($)": "1.5000", + "Requests": "2", + "Successful Requests": "2", + "Failed Requests": "0", + "Total Tokens": "30", + "Prompt Tokens": "20", + "Completion Tokens": "10", + "Cache Read Input Tokens": "5", + "Cache Creation Input Tokens": "4", + } + ] + + +def test_team_export_csv_omits_key_columns_for_the_plain_daily_scope(): + import csv + import io + + from litellm.proxy.management_endpoints.team_endpoints import _team_export_csv + from litellm.types.proxy.management_endpoints.team_endpoints import TeamDailyActivityExportRow + + row: Final = TeamDailyActivityExportRow( + date="2026-06-01", + team_id="team-1", + team_alias="Alpha", + spend=1.5, + api_requests=2, + successful_requests=2, + failed_requests=0, + total_tokens=30, + prompt_tokens=20, + completion_tokens=10, + cache_read_input_tokens=5, + cache_creation_input_tokens=4, + ) + + text: Final = _team_export_csv("daily", (row,)) + + assert text.splitlines()[0] == ( + "Date,Team,Team ID,Spend ($),Requests,Successful Requests,Failed Requests," + "Total Tokens,Prompt Tokens,Completion Tokens,Cache Read Input Tokens,Cache Creation Input Tokens" + ) + assert list(csv.reader(io.StringIO(text)))[1] == [ + "2026-06-01", + "Alpha", + "team-1", + "1.5000", + "2", + "2", + "0", + "30", + "20", + "10", + "5", + "4", + ] + + +def test_team_export_csv_escapes_formula_aliases_and_keeps_dash_placeholder(): + import csv + import io + + from litellm.proxy.management_endpoints.team_endpoints import _team_export_csv + from litellm.types.proxy.management_endpoints.team_endpoints import TeamDailyActivityExportRow + + row: Final = TeamDailyActivityExportRow( + date="2026-06-01", + team_id="team-1", + team_alias='=HYPERLINK("http://evil.example","x")', + key_alias="@cmd", + user_id=None, + user_email=None, + spend=1.5, + api_requests=2, + successful_requests=2, + failed_requests=0, + total_tokens=30, + prompt_tokens=20, + completion_tokens=10, + cache_read_input_tokens=5, + cache_creation_input_tokens=4, + ) + + record: Final = next(csv.DictReader(io.StringIO(_team_export_csv("daily_with_keys", (row,))))) + + assert record["Team"] == "'=HYPERLINK(\"http://evil.example\",\"x\")" + assert record["Key Alias"] == "'@cmd" + assert record["User ID"] == "-" + assert record["User Email"] == "-" diff --git a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/EntityUsage.tsx b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/EntityUsage.tsx index 16dc41c3ba8..d1ebd8b98a0 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/EntityUsage.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/EntityUsage.tsx @@ -26,7 +26,7 @@ import UserDropdown from "@/components/common_components/UserDropdown"; import { ActivityMetrics, processActivityData } from "@/components/activity_metrics"; import { UsageExportHeader } from "@/components/EntityUsageExport"; import { getApiKeyTruncation, getExportBlockedReason } from "@/components/EntityUsageExport/exportBlockedReason"; -import type { EntityType } from "@/components/EntityUsageExport/types"; +import type { EntityType, ServerExport } from "@/components/EntityUsageExport/types"; import { agentDailyActivityCall, customerDailyActivityCall, @@ -34,6 +34,7 @@ import { tagDailyActivityCall, teamDailyActivityAggregatedCall, teamDailyActivityCall, + teamDailyActivityExportCall, teamDailyActivityKeySearchCall, userDailyActivityCall, } from "@/components/networking"; @@ -685,7 +686,20 @@ const EntityUsage: React.FC = ({ { key: "endpoints", label: "Endpoint Activity", content: }, ]; - const spendFetchState = { coversRange, cancelled, failed, apiKeyTruncation }; + const serverExport: ServerExport | undefined = + entityType === "team" && apiKeyTruncation !== undefined && accessToken && startTime && endTime + ? (scope, format) => + teamDailyActivityExportCall({ + accessToken, + startTime, + endTime, + teamIds: entityFilterArg as string[] | null, + exportType: scope, + format, + }) + : undefined; + + const spendFetchState = { coversRange, cancelled, failed, apiKeyTruncation: serverExport ? null : apiKeyTruncation }; return (
@@ -719,6 +733,7 @@ const EntityUsage: React.FC = ({ filterOptions={getAllTags() || undefined} teams={teams || []} exportBlockedReason={getExportBlockedReason(spendFetchState)} + serverExport={serverExport} /> diff --git a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsagePageView.tsx b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsagePageView.tsx index e0171d97423..4c5db0def62 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsagePageView.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsagePageView.tsx @@ -253,14 +253,15 @@ const UsagePage: React.FC = ({ teams, organizations }) => { // Read through the same range stamp as the tiles, so the export is blocked from the first // render of a new range rather than from whenever the fetch effect gets around to running. + const apiKeyTruncation = getApiKeyTruncation( + userSpendData.metadata?.api_key_limit, + userSpendData.metadata?.total_api_keys, + ); const spendFetchState = { coversRange: activeAggregated !== null || paginatedResult.coversRange, cancelled: paginatedResult.cancelled, failed: paginatedResult.failed, - apiKeyTruncation: getApiKeyTruncation( - userSpendData.metadata?.api_key_limit, - userSpendData.metadata?.total_api_keys, - ), + apiKeyTruncation, }; const exportBlockedReason = getExportBlockedReason(spendFetchState); @@ -877,7 +878,7 @@ const UsagePage: React.FC = ({ teams, organizations }) => { diff --git a/ui/litellm-dashboard/src/components/EntityUsageExport/EntityUsageExportModal.test.tsx b/ui/litellm-dashboard/src/components/EntityUsageExport/EntityUsageExportModal.test.tsx index cb04323e4c9..6f74d241166 100644 --- a/ui/litellm-dashboard/src/components/EntityUsageExport/EntityUsageExportModal.test.tsx +++ b/ui/litellm-dashboard/src/components/EntityUsageExport/EntityUsageExportModal.test.tsx @@ -1,14 +1,3 @@ -/** - * Tests for EntityUsageExportModal component - * - * Validates core export functionality: - * - Renders modal with correct default state (CSV format, daily scope) - * - User can select export type (daily vs daily_with_models) - * - User can switch format (CSV vs JSON) - * - Export button triggers data generation with correct parameters - * - Modal closes after successful export - */ - import { describe, it, expect, vi, beforeEach } from "vitest"; import { screen } from "@testing-library/react"; import { renderWithProviders } from "../../../tests/test-utils"; @@ -20,6 +9,7 @@ vi.mock("./utils", () => { return { handleExportCSV: vi.fn(), handleExportJSON: vi.fn(), + handleServerExport: vi.fn(async () => undefined), generateExportData: vi.fn(() => [{ Date: "2025-10-01" }]), generateMetadata: vi.fn(() => ({ meta: true })), }; @@ -114,4 +104,23 @@ describe("EntityUsageExportModal", () => { // Modal closes after export expect(baseProps.onClose).toHaveBeenCalled(); }); + + it("routes the export through the server export when one is provided, so truncated key lists still export", async () => { + /** + * When the spend fetch was capped at the top-N keys, the caller supplies a + * serverExport that hits the uncapped export route. The modal must defer to + * it instead of generating a CSV from the truncated on-screen data. + */ + const user = userEvent.setup(); + const { handleExportCSV, handleServerExport } = await import("./utils"); + const serverExport = vi.fn(async () => new Blob(["csv"])); + + renderWithProviders(); + + await user.click(screen.getByRole("button", { name: /Export CSV/i })); + + expect(handleServerExport).toHaveBeenCalledWith(serverExport, "daily", "team", "csv"); + expect(handleExportCSV).not.toHaveBeenCalled(); + expect(baseProps.onClose).toHaveBeenCalled(); + }); }); diff --git a/ui/litellm-dashboard/src/components/EntityUsageExport/EntityUsageExportModal.tsx b/ui/litellm-dashboard/src/components/EntityUsageExport/EntityUsageExportModal.tsx index ec688e9fc3d..5e177a3efe6 100644 --- a/ui/litellm-dashboard/src/components/EntityUsageExport/EntityUsageExportModal.tsx +++ b/ui/litellm-dashboard/src/components/EntityUsageExport/EntityUsageExportModal.tsx @@ -10,7 +10,7 @@ import ExportFormatSelector from "./ExportFormatSelector"; import ExportSummary from "./ExportSummary"; import ExportTypeSelector from "./ExportTypeSelector"; import type { EntityUsageExportModalProps, ExportFormat, ExportScope } from "./types"; -import { handleExportCSV, handleExportJSON } from "./utils"; +import { handleExportCSV, handleExportJSON, handleServerExport } from "./utils"; const EntityUsageExportModal: React.FC = ({ isOpen, @@ -20,6 +20,7 @@ const EntityUsageExportModal: React.FC = ({ dateRange, selectedFilters, customTitle, + serverExport, }) => { const [exportFormat, setExportFormat] = useState("csv"); const [exportScope, setExportScope] = useState("daily"); @@ -35,7 +36,10 @@ const EntityUsageExportModal: React.FC = ({ const formatToUse = format || exportFormat; setIsExporting(true); try { - if (formatToUse === "csv") { + if (serverExport) { + await handleServerExport(serverExport, exportScope, entityType, formatToUse); + toast.success(`${entityLabel} usage data exported successfully as ${formatToUse.toUpperCase()}`); + } else if (formatToUse === "csv") { handleExportCSV(spendData, exportScope, entityLabel, entityType, teamAliasMap); toast.success(`${entityLabel} usage data exported successfully as CSV`); } else { diff --git a/ui/litellm-dashboard/src/components/EntityUsageExport/UsageExportHeader.tsx b/ui/litellm-dashboard/src/components/EntityUsageExport/UsageExportHeader.tsx index 388a211d9bb..10825e18f3d 100644 --- a/ui/litellm-dashboard/src/components/EntityUsageExport/UsageExportHeader.tsx +++ b/ui/litellm-dashboard/src/components/EntityUsageExport/UsageExportHeader.tsx @@ -16,7 +16,7 @@ import { useComboboxAnchor, } from "@/components/ui/combobox"; import EntityUsageExportModal from "./EntityUsageExportModal"; -import type { EntitySpendData, EntityType } from "./types"; +import type { EntitySpendData, EntityType, ServerExport } from "./types"; import type { Team } from "@/components/key_team_helpers/key_list"; interface UsageExportHeaderProps { @@ -35,6 +35,7 @@ interface UsageExportHeaderProps { compactLayout?: boolean; teams?: Team[]; exportBlockedReason?: string; + serverExport?: ServerExport; } const UsageExportHeader: React.FC = ({ @@ -52,6 +53,7 @@ const UsageExportHeader: React.FC = ({ compactLayout = false, teams = [], exportBlockedReason, + serverExport, }) => { const anchor = useComboboxAnchor(); const [isExportModalOpen, setIsExportModalOpen] = useState(false); @@ -142,6 +144,7 @@ const UsageExportHeader: React.FC = ({ selectedFilters={selectedFilters} customTitle={customTitle} teams={teams} + serverExport={serverExport} /> ); diff --git a/ui/litellm-dashboard/src/components/EntityUsageExport/exportBlockedReason.test.ts b/ui/litellm-dashboard/src/components/EntityUsageExport/exportBlockedReason.test.ts index 8491b31f5f9..3d84ef4703a 100644 --- a/ui/litellm-dashboard/src/components/EntityUsageExport/exportBlockedReason.test.ts +++ b/ui/litellm-dashboard/src/components/EntityUsageExport/exportBlockedReason.test.ts @@ -39,6 +39,10 @@ describe("getExportBlockedReason", () => { expect(reason).toMatch(/100 highest-spend keys of 3000/); expect(reason).toMatch(/USAGE_TOP_API_KEYS_LIMIT/); }); + + it("does not block on truncation when a server export will cover every key", () => { + expect(getExportBlockedReason(state({ apiKeyTruncation: null }))).toBeUndefined(); + }); }); describe("getApiKeyTruncation", () => { diff --git a/ui/litellm-dashboard/src/components/EntityUsageExport/exportBlockedReason.ts b/ui/litellm-dashboard/src/components/EntityUsageExport/exportBlockedReason.ts index 6c5a5f83231..32601680bad 100644 --- a/ui/litellm-dashboard/src/components/EntityUsageExport/exportBlockedReason.ts +++ b/ui/litellm-dashboard/src/components/EntityUsageExport/exportBlockedReason.ts @@ -7,7 +7,7 @@ export interface UsageFetchState { coversRange: boolean; cancelled: boolean; failed: boolean; - apiKeyTruncation: ApiKeyTruncation | undefined; + apiKeyTruncation?: ApiKeyTruncation | null; } export const getApiKeyTruncation = (apiKeyLimit: unknown, totalApiKeys: unknown): ApiKeyTruncation | undefined => { @@ -25,7 +25,7 @@ export const getExportBlockedReason = ({ if (cancelled) return "Loading was stopped before the whole range arrived, so an export would under-report. Reload the page to load it all."; if (!coversRange) return "Spend data is still loading, so an export would under-report. Wait for it to finish."; - if (apiKeyTruncation !== undefined) + if (apiKeyTruncation) return `Only the ${apiKeyTruncation.limit} highest-spend keys of ${apiKeyTruncation.total} were loaded, so a per-team export would under-report. Raise USAGE_TOP_API_KEYS_LIMIT on the proxy to load more keys.`; return undefined; }; diff --git a/ui/litellm-dashboard/src/components/EntityUsageExport/types.ts b/ui/litellm-dashboard/src/components/EntityUsageExport/types.ts index 15f193ecc3f..9cca4a2ab2b 100644 --- a/ui/litellm-dashboard/src/components/EntityUsageExport/types.ts +++ b/ui/litellm-dashboard/src/components/EntityUsageExport/types.ts @@ -17,6 +17,8 @@ export interface EntitySpendData { }; } +export type ServerExport = (exportScope: ExportScope, format: ExportFormat) => Promise; + export interface EntityUsageExportModalProps { isOpen: boolean; onClose: () => void; @@ -26,6 +28,7 @@ export interface EntityUsageExportModalProps { selectedFilters: string[]; customTitle?: string; teams?: Team[]; + serverExport?: ServerExport; } export interface ExportMetadata { diff --git a/ui/litellm-dashboard/src/components/EntityUsageExport/utils.test.ts b/ui/litellm-dashboard/src/components/EntityUsageExport/utils.test.ts index 3f9cf58ec20..4273a54e07d 100644 --- a/ui/litellm-dashboard/src/components/EntityUsageExport/utils.test.ts +++ b/ui/litellm-dashboard/src/components/EntityUsageExport/utils.test.ts @@ -14,6 +14,7 @@ import { getEntityBreakdown, handleExportCSV, handleExportJSON, + handleServerExport, resolveEntities, } from "./utils"; @@ -3005,4 +3006,40 @@ describe("EntityUsageExport utils", () => { ]); }); }); + + describe("handleServerExport", () => { + beforeEach(() => { + document.body.innerHTML = ""; + window.URL.createObjectURL = vi.fn(() => "blob:mock-url"); + window.URL.revokeObjectURL = vi.fn(); + }); + + afterEach(() => { + vi.restoreAllMocks(); + }); + + it("passes the chosen scope and format to the server export and downloads the returned blob", async () => { + const serverBlob = new Blob(["payload"], { type: "text/csv" }); + const serverExport = vi.fn(async () => serverBlob); + const createObjectURLSpy = vi.spyOn(window.URL, "createObjectURL"); + const appendChildSpy = vi.spyOn(document.body, "appendChild"); + + await handleServerExport(serverExport, "daily_with_keys", "team", "csv"); + + expect(serverExport).toHaveBeenCalledWith("daily_with_keys", "csv"); + expect(createObjectURLSpy).toHaveBeenCalledWith(serverBlob); + const attached = appendChildSpy.mock.calls[0][0] as HTMLAnchorElement; + const today = new Date().toISOString().split("T")[0]; + expect(attached.download).toBe(`team_usage_daily_with_keys_${today}.csv`); + }); + + it("lets a server failure propagate so the modal can toast it instead of downloading nothing", async () => { + const serverExport = vi.fn(async () => { + throw new Error("upstream 500"); + }); + + await expect(handleServerExport(serverExport, "daily", "team", "json")).rejects.toThrow("upstream 500"); + expect(document.body.querySelector("a")).toBeNull(); + }); + }); }); diff --git a/ui/litellm-dashboard/src/components/EntityUsageExport/utils.ts b/ui/litellm-dashboard/src/components/EntityUsageExport/utils.ts index 95ce584cc89..861ea585142 100644 --- a/ui/litellm-dashboard/src/components/EntityUsageExport/utils.ts +++ b/ui/litellm-dashboard/src/components/EntityUsageExport/utils.ts @@ -2,21 +2,27 @@ import { formatNumberWithCommas } from "@/utils/dataUtils"; import type { DateRangePickerValue } from "@/components/shared/date_picker_types"; import Papa from "papaparse"; import { keyActivityLabel } from "@/components/UsagePage/keyActivityLabel"; -import type { EntityBreakdown, EntitySpendData, EntityType, ExportMetadata, ExportScope } from "./types"; +import type { + EntityBreakdown, + EntitySpendData, + EntityType, + ExportFormat, + ExportMetadata, + ExportScope, + ServerExport, +} from "./types"; const resolveEntityDisplay = ( entity: string, teamAliasMap: Record, entityMetadata?: Record, -): { id: string; alias: string } => ({ - id: entity, - alias: - teamAliasMap[entity] || - entityMetadata?.team_alias || - entityMetadata?.user_email || - entityMetadata?.user_alias || - entity, -}); +): { id: string; alias: string } => { + const alias = + [teamAliasMap[entity], entityMetadata?.team_alias, entityMetadata?.user_email, entityMetadata?.user_alias].find( + Boolean, + ) ?? entity; + return { id: entity, alias }; +}; // Mirrors backend SpendMetrics fields (litellm/types/activity_tracking.py). // If the backend adds a field, add it here too. @@ -375,7 +381,7 @@ export const generateDailyWithModelsData = ( const { id, alias } = resolveEntityDisplay(entity, teamAliasMap, dailyEntityMetadata[entity]); Object.entries(models).forEach(([model, metrics]: [string, any]) => { - dailyModelBreakdown.push({ + const row = { Date: day.date, [entityLabel]: alias, [`${entityLabel} ID`]: id, @@ -389,7 +395,8 @@ export const generateDailyWithModelsData = ( "Completion Tokens": metrics.completionTokens, "Cache Read Input Tokens": metrics.cacheReadInputTokens, "Cache Creation Input Tokens": metrics.cacheCreationInputTokens, - }); + }; + dailyModelBreakdown.push(row); }); }); }); @@ -449,6 +456,28 @@ export const generateMetadata = ( }; }; +export const downloadBlob = (blob: Blob, fileName: string): void => { + const url = window.URL.createObjectURL(blob); + const a = document.createElement("a"); + a.href = url; + a.download = fileName; + document.body.appendChild(a); + a.click(); + document.body.removeChild(a); + window.URL.revokeObjectURL(url); +}; + +export const handleServerExport = async ( + serverExport: ServerExport, + exportScope: ExportScope, + entityType: EntityType, + format: ExportFormat, +): Promise => { + const blob = await serverExport(exportScope, format); + const fileName = `${entityType}_usage_${exportScope}_${new Date().toISOString().split("T")[0]}.${format}`; + downloadBlob(blob, fileName); +}; + export const handleExportCSV = ( spendData: EntitySpendData, exportScope: ExportScope, @@ -459,15 +488,8 @@ export const handleExportCSV = ( const data = generateExportData(spendData, exportScope, entityLabel, teamAliasMap); const csv = Papa.unparse(data); const blob = new Blob([csv], { type: "text/csv;charset=utf-8;" }); - const url = window.URL.createObjectURL(blob); - const a = document.createElement("a"); - a.href = url; const fileName = `${entityType}_usage_${exportScope}_${new Date().toISOString().split("T")[0]}.csv`; - a.download = fileName; - document.body.appendChild(a); - a.click(); - document.body.removeChild(a); - window.URL.revokeObjectURL(url); + downloadBlob(blob, fileName); }; export const handleExportJSON = ( @@ -487,13 +509,6 @@ export const handleExportJSON = ( }; const jsonString = JSON.stringify(exportObject, null, 2); const blob = new Blob([jsonString], { type: "application/json" }); - const url = window.URL.createObjectURL(blob); - const a = document.createElement("a"); - a.href = url; const fileName = `${entityType}_usage_${exportScope}_${new Date().toISOString().split("T")[0]}.json`; - a.download = fileName; - document.body.appendChild(a); - a.click(); - document.body.removeChild(a); - window.URL.revokeObjectURL(url); + downloadBlob(blob, fileName); }; diff --git a/ui/litellm-dashboard/src/components/networking.tsx b/ui/litellm-dashboard/src/components/networking.tsx index edaf56a8a16..19af2e79bdd 100644 --- a/ui/litellm-dashboard/src/components/networking.tsx +++ b/ui/litellm-dashboard/src/components/networking.tsx @@ -111,6 +111,7 @@ import type { CoordinationRedisTestResponse, } from "@/app/(dashboard)/caching/_components/coordination_redis_settings/types"; import { MCP_TOOLS_PREVIEW_FORBIDDEN_MESSAGE } from "./mcp_tools/constants"; +import type { ExportFormat, ExportScope } from "./EntityUsageExport/types"; import type { ComplexityRouterConfigPayload } from "./add_model/build_complexity_router_config"; import type { AutoRouterPresetsResponse } from "@/lib/autorouter_presets"; import type { VectorStoreIndex } from "@/app/(dashboard)/vector-stores/_components/IndexesTab"; @@ -1467,6 +1468,36 @@ export const teamDailyActivityAggregatedCall = async ( } }; +export const teamDailyActivityExportCall = async ({ + accessToken, + startTime, + endTime, + teamIds, + exportType, + format, +}: { + accessToken: string; + startTime: Date; + endTime: Date; + teamIds: string[] | null; + exportType: ExportScope; + format: ExportFormat; +}): Promise => { + return apiClient.get(`/team/daily/activity/export`, { + accessToken, + responseType: "blob", + query: { + start_date: formatDate(startTime), + end_date: formatDate(endTime), + timezone: new Date().getTimezoneOffset().toString(), + export_type: exportType, + format, + team_id: teamIds && teamIds.length > 0 ? teamIds.join(",") : undefined, + exclude_team_ids: "litellm-dashboard", + }, + }); +}; + export const teamDailyActivityKeySearchCall = async ( accessToken: string, startTime: Date, diff --git a/ui/litellm-dashboard/src/lib/http/client.ts b/ui/litellm-dashboard/src/lib/http/client.ts index c78b2b10011..974c07736e9 100644 --- a/ui/litellm-dashboard/src/lib/http/client.ts +++ b/ui/litellm-dashboard/src/lib/http/client.ts @@ -22,6 +22,8 @@ export interface RequestOptions { body?: unknown; /** Sent verbatim (FormData, Blob, pre-stringified text); disables JSON handling. */ rawBody?: BodyInit; + /** Response body handling. Defaults to JSON parsing; use this for downloads. */ + responseType?: "json" | "blob" | "text"; query?: QueryParams; headers?: Record; signal?: AbortSignal; @@ -138,7 +140,7 @@ export function createApiClient(config: ApiClientConfig): ApiClient { const doFetch: typeof fetch = (input, init) => (fetchImpl ?? fetch)(input, init); async function request(method: HttpMethod, path: string, options: RequestOptions = {}): Promise { - const { accessToken, body, rawBody, query, headers: extraHeaders, signal, credentials } = options; + const { accessToken, body, rawBody, query, headers: extraHeaders, signal, credentials, responseType } = options; const url = appendQuery(`${getBaseUrl()}${path}`, query); @@ -177,6 +179,12 @@ export function createApiClient(config: ApiClientConfig): ApiClient { throw new ApiError(message, response.status, errorBody); } + if (responseType === "blob") { + return (await response.blob()) as T; + } + if (responseType === "text") { + return (await response.text()) as T; + } const text = await response.text(); return (text ? JSON.parse(text) : undefined) as T; } diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index a398b2db93c..5d0bb56936a 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -15666,6 +15666,32 @@ export interface paths { patch?: never; trace?: never; }; + "/team/daily/activity/export": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + /** + * Get Team Daily Activity Export + * @description Server-side Team Usage export, not subject to USAGE_TOP_API_KEYS_LIMIT. + * + * Same scoping as /team/daily/activity/aggregated, answered by one unbounded + * rollup query, returned as CSV or JSON. For daily_with_keys, + * daily_with_users and daily_with_models the PTU sentinel flat-cost rows are + * excluded, so metadata totals under those export types cover request spend + * only; the plain daily export includes them. + */ + get: operations["get_team_daily_activity_export_team_daily_activity_export_get"]; + put?: never; + post?: never; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; "/team/delete": { parameters: { query?: never; @@ -31399,7 +31425,7 @@ export interface components { * @description Enum for key management routes * @enum {string} */ - KeyManagementRoutes: "/key/generate" | "/key/update" | "/key/delete" | "/key/regenerate" | "/key/service-account/generate" | "/key/{key_id}/regenerate" | "/key/block" | "/key/unblock" | "/key/bulk_update" | "/team/key/bulk_update" | "/key/{key_id}/reset_spend" | "/key/access_group_assignment" | "/auto_router/manage" | "/key/info" | "/key/health" | "/key/list" | "/key/aliases" | "/team/daily/activity" | "/team/daily/activity/aggregated" | "/team/daily/activity/aggregated/search" | "/spend/logs" | "/spend/logs/v2"; + KeyManagementRoutes: "/key/generate" | "/key/update" | "/key/delete" | "/key/regenerate" | "/key/service-account/generate" | "/key/{key_id}/regenerate" | "/key/block" | "/key/unblock" | "/key/bulk_update" | "/team/key/bulk_update" | "/key/{key_id}/reset_spend" | "/key/access_group_assignment" | "/auto_router/manage" | "/key/info" | "/key/health" | "/key/list" | "/key/aliases" | "/team/daily/activity" | "/team/daily/activity/aggregated" | "/team/daily/activity/export" | "/team/daily/activity/aggregated/search" | "/spend/logs" | "/spend/logs/v2"; /** * KeyManagementSystem * @enum {string} @@ -42812,6 +42838,87 @@ export interface components { /** Team Id */ team_id: string; }; + /** TeamDailyActivityExportMetadata */ + TeamDailyActivityExportMetadata: { + /** End Date */ + end_date: string; + /** Export Date */ + export_date: string; + /** + * Export Type + * @enum {string} + */ + export_type: "daily" | "daily_with_keys" | "daily_with_users" | "daily_with_models"; + /** Start Date */ + start_date: string; + /** Team Ids */ + team_ids: string[] | null; + /** Total Api Requests */ + total_api_requests: number; + /** Total Failed Requests */ + total_failed_requests: number; + /** + * Total Flat Cost + * @default 0 + */ + total_flat_cost: number; + /** Total Spend */ + total_spend: number; + /** Total Successful Requests */ + total_successful_requests: number; + /** Total Tokens */ + total_tokens: number; + }; + /** TeamDailyActivityExportResponse */ + TeamDailyActivityExportResponse: { + /** Data */ + data: components["schemas"]["TeamDailyActivityExportRow"][]; + metadata: components["schemas"]["TeamDailyActivityExportMetadata"]; + }; + /** TeamDailyActivityExportRow */ + TeamDailyActivityExportRow: { + /** Api Key */ + api_key?: string | null; + /** Api Requests */ + api_requests: number; + /** Cache Creation Input Tokens */ + cache_creation_input_tokens: number; + /** Cache Read Input Tokens */ + cache_read_input_tokens: number; + /** Completion Tokens */ + completion_tokens: number; + /** Date */ + date: string; + /** Failed Requests */ + failed_requests: number; + /** + * Flat Cost + * @default 0 + */ + flat_cost: number; + /** Key Alias */ + key_alias?: string | null; + /** Keys */ + keys?: number | null; + /** Model */ + model?: string | null; + /** Prompt Tokens */ + prompt_tokens: number; + /** Spend */ + spend: number; + /** Successful Requests */ + successful_requests: number; + /** Team Alias */ + team_alias?: string | null; + /** Team Id */ + team_id: string; + /** Total Tokens */ + total_tokens: number; + /** User Email */ + user_email?: string | null; + /** User Id */ + user_id?: string | null; + }; /** * TeamListItem * @description A team item in the paginated list response, enriched with computed fields. @@ -66682,6 +66789,44 @@ export interface operations { }; }; }; + get_team_daily_activity_export_team_daily_activity_export_get: { + parameters: { + query?: { + start_date?: string | null; + end_date?: string | null; + export_type?: "daily" | "daily_with_keys" | "daily_with_users" | "daily_with_models"; + format?: "csv" | "json"; + team_id?: string | null; + exclude_team_ids?: string | null; + timezone?: number | null; + }; + header?: never; + path?: never; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["TeamDailyActivityExportResponse"]; + "text/csv": unknown; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; delete_team_team_delete_post: { parameters: { query?: never; From fc87a06f009ea0926005b29d072f6e4d0bacf843 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 24 Sep 2026 15:42:36 -0500 Subject: [PATCH 078/120] fix(proxy): stop leaking periodic tasks on every DB config reload (#42784) Co-authored-by: yassin Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../SlackAlerting/slack_alerting.py | 19 ++-- litellm/router.py | 6 +- .../router_strategy/base_routing_strategy.py | 16 ++- .../SlackAlerting/test_slack_alerting.py | 26 +++++ .../test_router_routing_groups.py | 105 ++++++++++++++++++ 5 files changed, 162 insertions(+), 10 deletions(-) diff --git a/litellm/integrations/SlackAlerting/slack_alerting.py b/litellm/integrations/SlackAlerting/slack_alerting.py index 8d0d044ff93..17ec3ed787d 100644 --- a/litellm/integrations/SlackAlerting/slack_alerting.py +++ b/litellm/integrations/SlackAlerting/slack_alerting.py @@ -118,6 +118,7 @@ class SlackAlerting(CustomBatchLogger): self.default_webhook_url = default_webhook_url self.flush_lock = asyncio.Lock() self.periodic_started = False + self._periodic_flush_task: asyncio.Task[None] | None = None self.hanging_request_check = AlertingHangingRequestCheck( slack_alerting_object=self, ) @@ -129,6 +130,12 @@ class SlackAlerting(CustomBatchLogger): self.digest_lock = asyncio.Lock() super().__init__(**kwargs, flush_lock=self.flush_lock) + def _ensure_periodic_flush_task(self) -> None: + if self.periodic_started and (self._periodic_flush_task is None or not self._periodic_flush_task.done()): + return + self._periodic_flush_task = asyncio.create_task(self.periodic_flush()) + self.periodic_started = True + def update_values( self, alerting: list | None = None, @@ -141,17 +148,14 @@ class SlackAlerting(CustomBatchLogger): ): if alerting is not None: self.alerting = alerting - asyncio.create_task(self.periodic_flush()) - self.periodic_started = True + self._ensure_periodic_flush_task() if alerting_threshold is not None: self.alerting_threshold = alerting_threshold if alert_types is not None: self.alert_types = alert_types if alerting_args is not None: self.alerting_args = SlackAlertingArgs(**alerting_args) - if not self.periodic_started: - asyncio.create_task(self.periodic_flush()) - self.periodic_started = True + self._ensure_periodic_flush_task() if alert_type_config is not None: for key, val in alert_type_config.items(): self.alert_type_config[key] = AlertTypeConfig(**val) if isinstance(val, dict) else val @@ -1446,9 +1450,8 @@ Model Info: return # Start periodic flush if not already started - if not self.periodic_started and self.alerting is not None and len(self.alerting) > 0: - asyncio.create_task(self.periodic_flush()) - self.periodic_started = True + if self.alerting is not None and len(self.alerting) > 0: + self._ensure_periodic_flush_task() if "webhook" in self.alerting and alert_type == "budget_alerts" and user_info is not None: await self.send_webhook_alert(webhook_event=user_info) diff --git a/litellm/router.py b/litellm/router.py index 8960cd92cd8..da92eef9102 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -118,6 +118,7 @@ from litellm.llms.openai_like.model_info import ( MODEL_INFO_REFRESH_SECONDS, get_openai_compatible_model_info, ) +from litellm.router_strategy.base_routing_strategy import BaseRoutingStrategy from litellm.router_strategy.budget_limiter import RouterBudgetLimiting from litellm.router_strategy.complexity_router.context_compaction import ( arm_compaction, @@ -1377,6 +1378,9 @@ class Router: `_init_routing_groups`) so repeated `update_settings` calls don't accumulate dead selectors that keep receiving callback events. """ + for selector in selectors: + if isinstance(selector, BaseRoutingStrategy): + selector.retire() selector_ids: Final = {id(s) for s in selectors if s is not None} if not selector_ids: return @@ -12117,7 +12121,7 @@ class Router: ) rebuild_routing_groups = True elif var == "routing_strategy_args": - routing_args_updated = True + routing_args_updated = value != self.routing_strategy_args setattr(self, var, value) else: verbose_router_logger.debug("Setting %s is not allowed", var) diff --git a/litellm/router_strategy/base_routing_strategy.py b/litellm/router_strategy/base_routing_strategy.py index 686d57e2b77..79d457f2836 100644 --- a/litellm/router_strategy/base_routing_strategy.py +++ b/litellm/router_strategy/base_routing_strategy.py @@ -40,10 +40,24 @@ class BaseRoutingStrategy(ABC): self.periodic_sync_in_memory_spend_with_redis(default_sync_interval=default_sync_interval) ) + def cancel_sync_task(self) -> None: + if self._sync_task is not None: + self._sync_task.cancel() + + def retire(self) -> None: + self.cancel_sync_task() + if not self.redis_increment_operation_queue: + return + try: + loop: Final = asyncio.get_running_loop() + except RuntimeError: + return + loop.create_task(self._push_in_memory_increments_to_redis()) + async def cleanup(self): """Cleanup method to be called when shutting down""" if self._sync_task is not None: - self._sync_task.cancel() + self.cancel_sync_task() try: await self._sync_task except asyncio.CancelledError: diff --git a/tests/test_litellm/integrations/SlackAlerting/test_slack_alerting.py b/tests/test_litellm/integrations/SlackAlerting/test_slack_alerting.py index 2d5eb78950c..b9e5ff2eeb7 100644 --- a/tests/test_litellm/integrations/SlackAlerting/test_slack_alerting.py +++ b/tests/test_litellm/integrations/SlackAlerting/test_slack_alerting.py @@ -526,3 +526,29 @@ async def test_async_send_batch_collapses_only_identical_alerts() -> None: {"text": f"[Num Alerts: 2]\n\n{THRESHOLD_ALERT}"}, {"text": CROSSED_ALERT}, ) + + +def _periodic_flush_tasks() -> list[asyncio.Task[object]]: + return [ + t + for t in asyncio.all_tasks() + if t.get_coro() is not None and t.get_coro().__qualname__ == "SlackAlerting.periodic_flush" + ] + + +@pytest.mark.asyncio +async def test_update_values_repeated_alerting_reload_keeps_single_periodic_flush_task() -> None: + slack_alerting: Final = SlackAlerting(alerting=["slack"]) + try: + for _ in range(5): + slack_alerting.update_values(alerting=["slack"]) + await asyncio.sleep(0) + flush_tasks: Final = _periodic_flush_tasks() + assert len(flush_tasks) == 1, f"expected 1 periodic_flush task, found {len(flush_tasks)}" + finally: + for t in _periodic_flush_tasks(): + t.cancel() + try: + await t + except asyncio.CancelledError: + pass diff --git a/tests/test_litellm/router_strategy/test_router_routing_groups.py b/tests/test_litellm/router_strategy/test_router_routing_groups.py index 425f68dda18..534aea47885 100644 --- a/tests/test_litellm/router_strategy/test_router_routing_groups.py +++ b/tests/test_litellm/router_strategy/test_router_routing_groups.py @@ -18,6 +18,7 @@ from pydantic import ValidationError import litellm from litellm import Router +from litellm.caching.redis_cache import RedisPipelineIncrementOperation from litellm.integrations.custom_logger import CustomLogger from litellm.types.router import DeploymentTypedDict, FallbackAccessCheck, RoutingGroup, RoutingStrategy from litellm.utils import Rules, function_setup @@ -2149,3 +2150,107 @@ async def test_caller_cannot_spoof_a_priority_group_to_bypass_fallback_gates( **{metadata_bucket: {"pre_routing_selected_model": "priority-group"}}, ) assert checked == ["priority-group"] + + +def _sync_task_count() -> int: + return sum( + 1 + for t in asyncio.all_tasks() + if t.get_coro() is not None + and t.get_coro().__qualname__ == "BaseRoutingStrategy.periodic_sync_in_memory_spend_with_redis" + ) + + +@pytest.mark.asyncio +async def test_update_settings_same_routing_strategy_args_does_not_leak_sync_tasks(monkeypatch) -> None: + monkeypatch.setattr(litellm, "callbacks", []) + monkeypatch.setattr(litellm, "input_callback", []) + router: Final = Router( + model_list=_model_list(), + routing_strategy="usage-based-routing-v2", + routing_strategy_args={"ttl": 60}, + ) + try: + assert _sync_task_count() == 1 + selector_before: Final = router.lowesttpm_logger_v2 + + for _ in range(5): + router.update_settings(routing_strategy_args={"ttl": 60}) + await asyncio.sleep(0) + assert _sync_task_count() == 1 + assert router.lowesttpm_logger_v2 is selector_before, "same routing_strategy_args must not rebuild the selector" + + router.update_settings(routing_strategy_args={"ttl": 120}) + await asyncio.sleep(0) + assert _sync_task_count() == 1 + assert router.lowesttpm_logger_v2.routing_args.ttl == 120 + finally: + for t in [ + t + for t in asyncio.all_tasks() + if t.get_coro() is not None + and t.get_coro().__qualname__ == "BaseRoutingStrategy.periodic_sync_in_memory_spend_with_redis" + ]: + t.cancel() + try: + await t + except asyncio.CancelledError: + pass + + +class _RecordingRedisCache: + def __init__(self) -> None: + self.increment_lists: list[list[RedisPipelineIncrementOperation]] = [] + + async def async_increment_pipeline(self, increment_list: list[RedisPipelineIncrementOperation]) -> list[float]: + self.increment_lists.append(list(increment_list)) + return [float(op["increment_value"]) for op in increment_list] + + +@pytest.mark.asyncio +async def test_update_settings_changed_routing_strategy_args_flushes_replaced_selector_queue( + monkeypatch, +) -> None: + monkeypatch.setattr(litellm, "callbacks", []) + monkeypatch.setattr(litellm, "input_callback", []) + router: Final = Router( + model_list=_model_list(), + routing_strategy="usage-based-routing-v2", + routing_strategy_args={"ttl": 60}, + ) + try: + redis_cache: Final = _RecordingRedisCache() + replaced: Final = router.lowesttpm_logger_v2 + replaced.dual_cache.redis_cache = redis_cache + replaced.redis_increment_operation_queue.append( + RedisPipelineIncrementOperation(key="rpm-key", increment_value=3, ttl=60) + ) + + router.update_settings(routing_strategy_args={"ttl": 120}) + await asyncio.sleep(0) + await asyncio.gather( + *( + t + for t in asyncio.all_tasks() + if t.get_coro() is not None + and t.get_coro().__qualname__ == "BaseRoutingStrategy._push_in_memory_increments_to_redis" + ) + ) + + assert router.lowesttpm_logger_v2 is not replaced + assert redis_cache.increment_lists == [ + [RedisPipelineIncrementOperation(key="rpm-key", increment_value=3, ttl=60)] + ] + assert replaced.redis_increment_operation_queue == [] + finally: + for t in [ + t + for t in asyncio.all_tasks() + if t.get_coro() is not None + and t.get_coro().__qualname__ == "BaseRoutingStrategy.periodic_sync_in_memory_spend_with_redis" + ]: + t.cancel() + try: + await t + except asyncio.CancelledError: + pass From 82146bff43d86052675e3d492661c21c17b58c01 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 24 Sep 2026 13:43:47 -0700 Subject: [PATCH 079/120] chore(cost-map): add azure deprecation dates from the Models API for five realtime and transcribe rows (#43037) Co-authored-by: kerry Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/model_prices_and_context_window_backup.json | 5 +++++ model_prices_and_context_window.json | 5 +++++ 2 files changed, 10 insertions(+) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 90a268aa3e2..5a48b2d293b 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -69776,6 +69776,7 @@ "cache_creation_input_audio_token_cost": 4e-07, "cache_read_input_audio_token_cost": 4e-07, "cache_read_input_token_cost": 4e-07, + "deprecation_date": "2026-10-31", "input_cost_per_audio_token": 3.2e-05, "input_cost_per_image_token": 5e-06, "input_cost_per_token": 4e-06, @@ -69807,6 +69808,7 @@ "supports_tool_choice": true }, "azure/gpt-live-1": { + "deprecation_date": "2027-09-10", "input_cost_per_second": 0.000833333333333, "litellm_provider": "azure", "mode": "realtime", @@ -69824,6 +69826,7 @@ "supports_function_calling": true }, "azure/gpt-live-transcribe": { + "deprecation_date": "2028-02-01", "input_cost_per_second": 0.000283333333333, "litellm_provider": "azure", "max_input_tokens": 32000, @@ -69845,6 +69848,7 @@ "supports_audio_input": true }, "azure/gpt-transcribe": { + "deprecation_date": "2028-02-01", "input_cost_per_second": 7.5e-05, "litellm_provider": "azure", "mode": "audio_transcription", @@ -69863,6 +69867,7 @@ "supports_audio_input": true }, "azure/gpt-realtime-translate": { + "deprecation_date": "2027-05-06", "input_cost_per_second": 0.000566666666667, "litellm_provider": "azure", "max_input_tokens": 32000, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 90a268aa3e2..5a48b2d293b 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -69776,6 +69776,7 @@ "cache_creation_input_audio_token_cost": 4e-07, "cache_read_input_audio_token_cost": 4e-07, "cache_read_input_token_cost": 4e-07, + "deprecation_date": "2026-10-31", "input_cost_per_audio_token": 3.2e-05, "input_cost_per_image_token": 5e-06, "input_cost_per_token": 4e-06, @@ -69807,6 +69808,7 @@ "supports_tool_choice": true }, "azure/gpt-live-1": { + "deprecation_date": "2027-09-10", "input_cost_per_second": 0.000833333333333, "litellm_provider": "azure", "mode": "realtime", @@ -69824,6 +69826,7 @@ "supports_function_calling": true }, "azure/gpt-live-transcribe": { + "deprecation_date": "2028-02-01", "input_cost_per_second": 0.000283333333333, "litellm_provider": "azure", "max_input_tokens": 32000, @@ -69845,6 +69848,7 @@ "supports_audio_input": true }, "azure/gpt-transcribe": { + "deprecation_date": "2028-02-01", "input_cost_per_second": 7.5e-05, "litellm_provider": "azure", "mode": "audio_transcription", @@ -69863,6 +69867,7 @@ "supports_audio_input": true }, "azure/gpt-realtime-translate": { + "deprecation_date": "2027-05-06", "input_cost_per_second": 0.000566666666667, "litellm_provider": "azure", "max_input_tokens": 32000, From 5e4b1b9df0ac766fc0640c186d994977dcf023d0 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Thu, 24 Sep 2026 13:44:56 -0700 Subject: [PATCH 080/120] fix(proxy): pass team member spend rows as jsonb so a $0 flush cannot poison the pool connection (#43029) Prisma types a raw array parameter from the first batch a connection sees. After a flush in which every member cost was a whole number (a free model), the connection's cached statement expected int8[] and every later fractional batch on it failed with "improper binary format in array element", so member spend silently stopped landing while team spend kept rising. The rows now travel as one JSON document unpacked by jsonb_to_recordset with the column types declared in SQL, so Postgres types the numbers and the batch shape no longer matters. --- litellm/proxy/db/db_spend_update_writer.py | 34 +++--- .../spend/test_team_member_spend_flush.py | 106 ++++++++++++++++++ .../proxy/db/test_db_spend_update_writer.py | 12 +- 3 files changed, 133 insertions(+), 19 deletions(-) create mode 100644 tests/integration/spend/test_team_member_spend_flush.py diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index d3ac37a3e0f..17e6152bef6 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -124,6 +124,12 @@ class _SpendIncrement(TypedDict): increment: ReadOnly[float] +class _MemberSpendRow(TypedDict): + user_id: ReadOnly[str] + team_id: ReadOnly[str] + cost: ReadOnly[float] + + class _SpendBatch(Protocol): litellm_usertable: BatchTable litellm_verificationtoken: BatchTable @@ -351,17 +357,22 @@ _TEAM_ADVISORY_LOCK_SQL: Final = "SELECT pg_advisory_xact_lock(hashtext($1)) IS # One statement adds every member's cost to their membership row. A missing row is created only # while the user is still on the team's roster, so a spend flush landing after a removal never -# recreates the member. +# recreates the member. The rows travel as one JSON document, not as a numeric array: Prisma +# types a raw array parameter from the first batch a connection sees, so after an all-$0 batch +# (integers) every later fractional batch on that connection failed with "improper binary format". _TEAM_MEMBER_SPEND_SQL: Final = """ INSERT INTO "LiteLLM_TeamMembership" (user_id, team_id, spend, total_spend) -SELECT p.user_id, p.team_id, p.cost, p.cost -FROM unnest($1::text[], $2::text[], $3::float8[]) AS p(user_id, team_id, cost) +SELECT member.user_id, member.team_id, member.cost, member.cost +FROM jsonb_to_recordset($1::jsonb) AS member(user_id text, team_id text, cost float8) WHERE EXISTS ( SELECT 1 FROM "LiteLLM_TeamTable" t - WHERE t.team_id = p.team_id - AND t.members_with_roles @> jsonb_build_array(jsonb_build_object('user_id', p.user_id)) + WHERE t.team_id = member.team_id + AND t.members_with_roles @> jsonb_build_array(jsonb_build_object('user_id', member.user_id)) +) + OR EXISTS ( + SELECT 1 FROM "LiteLLM_TeamMembership" m + WHERE m.user_id = member.user_id AND m.team_id = member.team_id ) - OR EXISTS (SELECT 1 FROM "LiteLLM_TeamMembership" m WHERE m.user_id = p.user_id AND m.team_id = p.team_id) ON CONFLICT (user_id, team_id) DO UPDATE SET spend = "LiteLLM_TeamMembership".spend + EXCLUDED.spend, total_spend = "LiteLLM_TeamMembership".total_spend + EXCLUDED.total_spend @@ -371,15 +382,12 @@ SET spend = "LiteLLM_TeamMembership".spend + EXCLUDED.spend, async def _write_team_member_spend(transaction: _SpendTransaction, spend_by_member_key: Mapping[str, float]) -> None: # key is "team_id::::user_id::"; locks are taken in sorted team_id order like the team endpoints rows: Final = sorted((key.split("::")[1], key.split("::")[3], cost) for key, cost in spend_by_member_key.items()) - team_ids: Final = tuple(team_id for team_id, _user_id, _cost in rows) - for team_id in dict.fromkeys(team_ids): + for team_id in dict.fromkeys(team_id for team_id, _user_id, _cost in rows): _ = await transaction.execute_raw(_TEAM_ADVISORY_LOCK_SQL, team_id) - _ = await transaction.execute_raw( - _TEAM_MEMBER_SPEND_SQL, - tuple(user_id for _team_id, user_id, _cost in rows), - team_ids, - tuple(cost for _team_id, _user_id, cost in rows), + members: Final = tuple( + _MemberSpendRow(user_id=user_id, team_id=team_id, cost=cost) for team_id, user_id, cost in rows ) + _ = await transaction.execute_raw(_TEAM_MEMBER_SPEND_SQL, json.dumps(members)) def get_llm_router(): diff --git a/tests/integration/spend/test_team_member_spend_flush.py b/tests/integration/spend/test_team_member_spend_flush.py new file mode 100644 index 00000000000..2433731f733 --- /dev/null +++ b/tests/integration/spend/test_team_member_spend_flush.py @@ -0,0 +1,106 @@ +"""Team member spend keeps landing after a flush in which every cost was a whole number. + +The proxy runs on a one-connection pool so every spend flush reuses the same database +connection. A batch of $0 requests (a free model here) is the whole-number batch, and the +fractional batches that follow it must still land on that connection. + +The $0 batch has to be flushed on its own before the paid request is sent. The spend log +row cannot prove that, since a separate monitor writes spend logs whenever they queue up, +but the daily user spend row is written by the flush cycle right after the member spend +statement, so its arrival means the whole-number batch has already been sent. +""" + +from pathlib import Path +from typing import Final + +import pytest +from integration._support.client import Gateway, eventually, string_value +from integration._support.database import read_rows +from integration._support.process import owned_proxy +from pydantic import JsonValue + +SINGLE_CONNECTION_CONFIG: Final = """ +model_list: [] +general_settings: + master_key: os.environ/LITELLM_MASTER_KEY + database_url: os.environ/DATABASE_URL + store_model_in_db: true + proxy_batch_write_at: 1 + proxy_batch_polling_interval: 1 + database_connection_pool_limit: 1 +router_settings: + disable_cooldowns: true +""" + + +def _member_row(team_id: str, user_id: str) -> list[dict[str, JsonValue]]: + return read_rows( + 'SELECT spend, total_spend FROM "LiteLLM_TeamMembership" WHERE team_id=%s AND user_id=%s', + (team_id, user_id), + ) + + +def _member_spend_is(rows: list[dict[str, JsonValue]], amount: float) -> bool: + return len(rows) == 1 and all( + float(str(rows[0][column])) == pytest.approx(amount) for column in ("spend", "total_spend") + ) + + +def _daily_user_spend_rows(user_id: str) -> list[dict[str, JsonValue]]: + return read_rows('SELECT spend FROM "LiteLLM_DailyUserSpend" WHERE user_id=%s', (user_id,)) + + +def _logged_spend(request_id: str) -> float: + rows: Final = eventually( + lambda: read_rows('SELECT spend FROM "LiteLLM_SpendLogs" WHERE request_id=%s', (request_id,)), + lambda found: len(found) == 1, + seconds=30, + ) + return float(str(rows[0]["spend"])) + + +def _chat(gateway: Gateway, key: str, model: str) -> str: + response: Final = gateway.request( + "POST", + "/v1/chat/completions", + {"model": model, "messages": [{"role": "user", "content": "member spend control"}]}, + key=key, + ) + assert response.status_code == 200, response.text + return string_value(response.json()["id"]) + + +def test_fractional_member_spend_lands_after_a_whole_number_flush_on_the_same_connection( + gateway: Gateway, tmp_path: Path +) -> None: + config: Final = tmp_path / "single_connection_proxy.yaml" + config.write_text(SINGLE_CONNECTION_CONFIG) + with owned_proxy(gateway, tmp_path, {}, config=config) as candidate, candidate.scenario() as scenario: + free: Final = scenario.model(input_cost_per_token=0, output_cost_per_token=0, num_retries=0) + paid: Final = scenario.model(input_cost_per_token=0.001, output_cost_per_token=0.002, num_retries=0) + team: Final = scenario.team(models=[free, paid]) + first: Final = scenario.user() + second: Final = scenario.user() + candidate.post( + "/team/member_add", + {"team_id": team, "member": [{"role": "user", "user_id": first}, {"role": "user", "user_id": second}]}, + ) + first_key: Final = scenario.key(team_id=team, user_id=first) + second_key: Final = scenario.key(team_id=team, user_id=second) + + _chat(candidate, first_key, free) + flushed: Final = eventually(lambda: _daily_user_spend_rows(first), lambda rows: len(rows) == 1, seconds=30) + assert float(str(flushed[0]["spend"])) == 0 + assert _member_spend_is(_member_row(team, first), 0) + + paid_spend: Final = _logged_spend(_chat(candidate, second_key, paid)) + assert paid_spend > 0 + eventually(lambda: _member_row(team, second), lambda rows: _member_spend_is(rows, paid_spend), seconds=30) + + repeat_spend: Final = _logged_spend(_chat(candidate, first_key, paid)) + eventually(lambda: _member_row(team, first), lambda rows: _member_spend_is(rows, repeat_spend), seconds=30) + eventually( + lambda: read_rows('SELECT spend FROM "LiteLLM_TeamTable" WHERE team_id=%s', (team,)), + lambda rows: float(str(rows[0]["spend"])) == pytest.approx(paid_spend + repeat_spend), + seconds=30, + ) diff --git a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py index bf3b9aed234..7abb6e1ef92 100644 --- a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py +++ b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py @@ -994,11 +994,11 @@ async def test_commit_spend_updates_to_db_writes_team_member_spend_in_one_roster assert lock_statement is _TEAM_ADVISORY_LOCK_SQL assert locked_team_id == team_id assert "pg_advisory_xact_lock(hashtext($1))" in lock_statement - statement, user_ids, team_ids, costs = spend_call.args + statement, members = spend_call.args assert statement is _TEAM_MEMBER_SPEND_SQL - assert (list(user_ids), list(team_ids), list(costs)) == ([user_id], [team_id], [response_cost]) + assert json.loads(members) == [{"user_id": user_id, "team_id": team_id, "cost": response_cost}] assert 'INSERT INTO "LiteLLM_TeamMembership"' in statement - assert "members_with_roles @> jsonb_build_array(jsonb_build_object('user_id', p.user_id))" in statement + assert "members_with_roles @> jsonb_build_array(jsonb_build_object('user_id', member.user_id))" in statement assert "ON CONFLICT (user_id, team_id) DO UPDATE" in statement assert 'spend = "LiteLLM_TeamMembership".spend + EXCLUDED.spend' in statement assert 'total_spend = "LiteLLM_TeamMembership".total_spend + EXCLUDED.total_spend' in statement @@ -1007,7 +1007,7 @@ async def test_commit_spend_updates_to_db_writes_team_member_spend_in_one_roster @pytest.mark.asyncio async def test_commit_spend_updates_to_db_orders_team_member_rows_by_team_then_user(): """ - The member spend statement touches rows in the order of its input arrays, so the batch + The member spend statement touches rows in the order of its input rows, so the batch is handed over sorted by (team_id, user_id), with each cost kept next to its member, and each distinct team is locked once, in `sorted(team_ids)` order, the order /team/delete locks in, so a concurrent flush and delete cannot deadlock. `eng` and `eng2` pin that: @@ -1034,13 +1034,13 @@ async def test_commit_spend_updates_to_db_orders_team_member_rows_by_team_then_u ) *lock_calls, spend_call = mock_transaction.execute_raw.await_args_list - _statement, user_ids, team_ids, costs = spend_call.args + _statement, members = spend_call.args assert [lock_call.args for lock_call in lock_calls] == [ (_TEAM_ADVISORY_LOCK_SQL, "eng"), (_TEAM_ADVISORY_LOCK_SQL, "eng-b"), (_TEAM_ADVISORY_LOCK_SQL, "eng2"), ] - assert list(zip(team_ids, user_ids, costs)) == [ + assert [(row["team_id"], row["user_id"], row["cost"]) for row in json.loads(members)] == [ ("eng", "user_x", 0.3), ("eng", "user_y", 0.2), ("eng-b", "user_x", 0.4), From fc29fb513cf7b7cffb67018ceb968b4c20001414 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 24 Sep 2026 14:05:31 -0700 Subject: [PATCH 081/120] fix(vertex_ai): keep batch output_file_id null until Vertex reports outputInfo (#43030) * fix(vertex_ai): keep batch output_file_id null until Vertex reports outputInfo Vertex only sets outputInfo.gcsOutputDirectory once a batch job has written output. Falling back to outputConfig's outputUriPrefix named the per-model directory shared by every batch of the deployment, an object that never exists, so the proxy minted a managed file for it under the first key and every other key's file calls on that id were 403s * fix(vertex_ai): treat a null gcsOutputDirectory as no output file yet --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- .../llms/vertex_ai/batches/transformation.py | 27 +--- .../test_vertex_batch_output_info_wire.py | 2 +- .../vertex_ai/batches/test_transformation.py | 120 +++++++++++------- .../test_vertex_ai_batch_transformation.py | 9 +- 4 files changed, 84 insertions(+), 74 deletions(-) diff --git a/litellm/llms/vertex_ai/batches/transformation.py b/litellm/llms/vertex_ai/batches/transformation.py index a7dbb058465..11c99d130be 100644 --- a/litellm/llms/vertex_ai/batches/transformation.py +++ b/litellm/llms/vertex_ai/batches/transformation.py @@ -253,30 +253,15 @@ class VertexAIBatchTransformation: return uris[0] @classmethod - def _get_output_file_id_from_vertex_ai_batch_response(cls, response: VertexBatchPredictionResponse) -> str: + def _get_output_file_id_from_vertex_ai_batch_response(cls, response: VertexBatchPredictionResponse) -> str | None: """ - Gets the output file id from the Vertex AI Batch response + Gets the output file id from the Vertex AI Batch response, None until Vertex reports outputInfo """ - output_info: Final = response.get("outputInfo") or OutputInfo() - output_file_id: str = output_info.get("gcsOutputDirectory", "") - if output_file_id: - output_file_id = output_file_id.rstrip("/") + "/predictions.jsonl" - if output_file_id and output_file_id != "/predictions.jsonl": - return output_file_id - - output_config: Final = response.get("outputConfig") - if output_config is None: - return output_file_id - - gcs_destination: Final = output_config.get("gcsDestination") - if gcs_destination is None: - return output_file_id - - output_uri_prefix: Final = gcs_destination.get("outputUriPrefix", "") - if output_uri_prefix.endswith("/predictions.jsonl"): - return output_uri_prefix - return output_uri_prefix.rstrip("/") + "/predictions.jsonl" + gcs_output_directory: Final = (output_info.get("gcsOutputDirectory") or "").rstrip("/") + if not gcs_output_directory: + return None + return f"{gcs_output_directory}/predictions.jsonl" @classmethod def _get_batch_job_status_from_vertex_ai_batch_response( diff --git a/tests/integration/providers/test_vertex_batch_output_info_wire.py b/tests/integration/providers/test_vertex_batch_output_info_wire.py index a7ac896076c..de41e0790d1 100644 --- a/tests/integration/providers/test_vertex_batch_output_info_wire.py +++ b/tests/integration/providers/test_vertex_batch_output_info_wire.py @@ -116,7 +116,7 @@ def test_vertex_batch_create_survives_explicit_null_output_info(gateway: Gateway "batch", "validating", _encoded(INPUT_FILE_ID, model, "file-"), - _encoded(f"{OUTPUT_PREFIX}/predictions.jsonl", model, "file-"), + None, None, "24h", ), response.text diff --git a/tests/test_litellm/llms/vertex_ai/batches/test_transformation.py b/tests/test_litellm/llms/vertex_ai/batches/test_transformation.py index ae045edbec5..5f969620a80 100644 --- a/tests/test_litellm/llms/vertex_ai/batches/test_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/batches/test_transformation.py @@ -13,6 +13,8 @@ There are no real I/O seams here; ``uuid.uuid4`` is the only nondeterministic dependency and is patched where the displayName is asserted. """ +from collections.abc import Mapping +from typing import Final from unittest.mock import patch import pytest @@ -35,8 +37,7 @@ INPUT_FILE = ( ENDPOINT_ID = "7768560373388541952" ENDPOINT_INPUT_FILE = ( - f"gs://litellm-testing-bucket/litellm-vertex-files/endpoints/{ENDPOINT_ID}/" - "e9412502-2c91-42a6-8e61-f5c294cc0fc8" + f"gs://litellm-testing-bucket/litellm-vertex-files/endpoints/{ENDPOINT_ID}/e9412502-2c91-42a6-8e61-f5c294cc0fc8" ) @@ -248,9 +249,78 @@ def test_get_input_file_id_empty_uris(): # =========================================================================== # -# _get_output_file_id_from_vertex_ai_batch_response +# _get_output_file_id_from_vertex_ai_batch_response: None until Vertex reports outputInfo # =========================================================================== # +SHARED_OUTPUT_PREFIX: Final = "gs://bucket/litellm-vertex-files/publishers/google/models/gemini-2.5-flash" +SUCCEEDED_OUTPUT_DIRECTORY: Final = f"{SHARED_OUTPUT_PREFIX}/prediction-model-2026-09-24T19:41:00.000000Z" + + +def _vertex_job(state: str) -> dict[str, object]: + return { + "name": "projects/510528649030/locations/us-central1/batchPredictionJobs/3814889423749775360", + "state": state, + "createTime": "2026-09-24T19:37:25.775603Z", + "inputConfig": { + "instancesFormat": "jsonl", + "gcsSource": {"uris": [f"{SHARED_OUTPUT_PREFIX}/0586ba52-4f8b-4988-aa8d-3573550a4b0f"]}, + }, + "outputConfig": { + "predictionsFormat": "jsonl", + "gcsDestination": {"outputUriPrefix": SHARED_OUTPUT_PREFIX}, + }, + } + + +@pytest.mark.parametrize( + "vertex_state,output_info_field,expected_status,expected_output_file_id", + [ + ("JOB_STATE_PENDING", {}, "validating", None), + ("JOB_STATE_RUNNING", {"outputInfo": {}}, "in_progress", None), + ("JOB_STATE_CANCELLED", {"outputInfo": None}, "cancelled", None), + ( + "JOB_STATE_SUCCEEDED", + {"outputInfo": {"gcsOutputDirectory": SUCCEEDED_OUTPUT_DIRECTORY}}, + "completed", + f"{SUCCEEDED_OUTPUT_DIRECTORY}/predictions.jsonl", + ), + ], + ids=["create_or_pending", "running", "cancelled", "succeeded"], +) +def test_transform_vertex_response_output_file_id_is_none_until_output_info( + vertex_state: str, + output_info_field: Mapping[str, object], + expected_status: str, + expected_output_file_id: str | None, +) -> None: + batch: Final = T.transform_vertex_ai_batch_response_to_openai_batch_response( + {**_vertex_job(vertex_state), **output_info_field} + ) + + assert batch.status == expected_status + assert batch.output_file_id == expected_output_file_id + + +@pytest.mark.parametrize( + "response", + [ + {}, + {"outputConfig": {}}, + {"outputInfo": None}, + {"outputInfo": {"gcsOutputDirectory": ""}}, + {"outputInfo": {"gcsOutputDirectory": None}}, + ], + ids=[ + "no_fields", + "output_config_without_destination", + "null_output_info", + "empty_output_directory", + "null_output_directory", + ], +) +def test_get_output_file_id_is_none_without_output_directory(response: Mapping[str, object]) -> None: + assert T._get_output_file_id_from_vertex_ai_batch_response(response) is None + def test_get_output_file_id_from_output_info(): # outputInfo branch: rstrip trailing slash, append predictions.jsonl @@ -267,49 +337,7 @@ def test_get_output_file_id_output_info_no_trailing_slash(): ) -def test_get_output_file_id_empty_output_info_falls_through_to_output_config(): - # gcsOutputDirectory missing -> "" -> the "/predictions.jsonl" guard skips - # the outputInfo branch, falls through to outputConfig - resp = { - "outputInfo": {}, - "outputConfig": {"gcsDestination": {"outputUriPrefix": "gs://b/cfg"}}, - } - assert T._get_output_file_id_from_vertex_ai_batch_response(resp) == "gs://b/cfg/predictions.jsonl" - - -def test_get_output_file_id_output_info_explicit_none_falls_through_to_output_config(): - resp = { - "outputInfo": None, - "outputConfig": {"gcsDestination": {"outputUriPrefix": "gs://b/cfg"}}, - } - assert T._get_output_file_id_from_vertex_ai_batch_response(resp) == "gs://b/cfg/predictions.jsonl" - - -def test_get_output_file_id_output_info_explicit_none_and_no_output_config(): - assert T._get_output_file_id_from_vertex_ai_batch_response({"outputInfo": None}) == "" - - -def test_get_output_file_id_no_output_info_and_no_output_config(): - assert T._get_output_file_id_from_vertex_ai_batch_response({}) == "" - - -def test_get_output_file_id_output_config_missing_gcs_destination(): - # outputConfig present but no gcsDestination -> returns the running "" value - assert T._get_output_file_id_from_vertex_ai_batch_response({"outputConfig": {}}) == "" - - -def test_get_output_file_id_output_config_already_has_suffix(): - # outputUriPrefix already ends in /predictions.jsonl -> returned as-is (no double append) - resp = {"outputConfig": {"gcsDestination": {"outputUriPrefix": "gs://b/cfg/predictions.jsonl"}}} - assert T._get_output_file_id_from_vertex_ai_batch_response(resp) == "gs://b/cfg/predictions.jsonl" - - -def test_get_output_file_id_output_config_strips_trailing_slash(): - resp = {"outputConfig": {"gcsDestination": {"outputUriPrefix": "gs://b/cfg/"}}} - assert T._get_output_file_id_from_vertex_ai_batch_response(resp) == "gs://b/cfg/predictions.jsonl" - - -def test_get_output_file_id_output_info_takes_precedence_over_output_config(): +def test_get_output_file_id_output_info_ignores_output_uri_prefix(): resp = { "outputInfo": {"gcsOutputDirectory": "gs://from-info"}, "outputConfig": {"gcsDestination": {"outputUriPrefix": "gs://from-config"}}, diff --git a/tests/test_litellm/llms/vertex_ai/test_vertex_ai_batch_transformation.py b/tests/test_litellm/llms/vertex_ai/test_vertex_ai_batch_transformation.py index 8a135dac3bb..3bc66d017a0 100644 --- a/tests/test_litellm/llms/vertex_ai/test_vertex_ai_batch_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/test_vertex_ai_batch_transformation.py @@ -26,12 +26,12 @@ def test_output_file_id_uses_predictions_jsonl_with_output_info(): ) -def test_output_file_id_falls_back_to_output_uri_prefix_with_predictions_jsonl(): +def test_output_file_id_is_none_until_output_info(): response = { "outputInfo": {}, "outputConfig": { "gcsDestination": { - "outputUriPrefix": "gs://test-bucket/litellm-vertex-files/publishers/google/models/gemini-2.5-pro/prediction-model-456" + "outputUriPrefix": "gs://test-bucket/litellm-vertex-files/publishers/google/models/gemini-2.5-pro" } }, } @@ -42,10 +42,7 @@ def test_output_file_id_falls_back_to_output_uri_prefix_with_predictions_jsonl() ) ) - assert ( - output_file_id - == "gs://test-bucket/litellm-vertex-files/publishers/google/models/gemini-2.5-pro/prediction-model-456/predictions.jsonl" - ) + assert output_file_id is None def test_vertex_ai_cancel_batch(): From 0f0ac4ad59cd72d8968b13e79d4e7ecf62479306 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 24 Sep 2026 14:22:36 -0700 Subject: [PATCH 082/120] fix(playground): stop following streamed tokens, add jump to bottom button (#42968) * fix(playground): only auto-scroll the chat while pinned to the bottom Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(playground): keep scroll pin through programmatic scrolls Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(playground): stop forcing the chat to scroll to the bottom on every update Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(playground): drop scroll pinning integration tests Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(playground): scroll only the chat pane, not the page, while streaming Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * style(playground): format ChatUI with prettier Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * feat(playground): stop following streamed tokens, add jump to bottom button Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(playground): jump to bottom lands on the last message, not the spacer Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: ryan Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../chat_ui/ChatUI.integration.test.tsx | 54 ++++++++ .../playground/components/chat_ui/ChatUI.tsx | 124 ++++++++++-------- 2 files changed, 126 insertions(+), 52 deletions(-) diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatUI.integration.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatUI.integration.test.tsx index e79d382ae39..a713c581e47 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatUI.integration.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatUI.integration.test.tsx @@ -791,4 +791,58 @@ describe("ChatUI", () => { expect(screen.getByPlaceholderText("Select a Model")).toBeEnabled(); }); }); + + it("sends scroll the chat pane to the new message, tokens do not, jump button scrolls to bottom", async () => { + const scrollTopSetter = vi.spyOn(HTMLElement.prototype, "scrollTop", "set"); + let streamChunk: ((chunk: string, model?: string) => void) | undefined; + vi.mocked(makeOpenAIChatCompletionRequest).mockImplementation(async (...args) => { + streamChunk = args[1] as (chunk: string, model?: string) => void; + }); + + render( + , + ); + + await waitFor(() => { + expect(screen.getByText("Test Key")).toBeInTheDocument(); + }); + + await selectComboboxOption("Select an endpoint", "/v1/chat/completions"); + await selectComboboxOption("Select a Model", "Model 1"); + const messageInput = screen.getByPlaceholderText("Type your message... (Shift+Enter for new line)"); + await act(async () => { + fireEvent.change(messageInput, { target: { value: "hello" } }); + }); + await act(async () => { + fireEvent.keyDown(messageInput, { key: "Enter", code: "Enter" }); + }); + + await waitFor(() => { + expect(makeOpenAIChatCompletionRequest).toHaveBeenCalledTimes(1); + }); + + expect(scrollTopSetter).toHaveBeenCalled(); + scrollTopSetter.mockClear(); + + const scrollIntoViewMock = vi.mocked(Element.prototype.scrollIntoView); + const scrollIntoViewCallsBeforeTokens = scrollIntoViewMock.mock.calls.length; + + await act(async () => { + streamChunk?.("Hello world", "Model 1"); + }); + + expect(scrollTopSetter).not.toHaveBeenCalled(); + expect(scrollIntoViewMock.mock.calls.length).toBe(scrollIntoViewCallsBeforeTokens); + + const user = userEvent.setup(); + await user.click(screen.getByRole("button", { name: "Jump to bottom" })); + + expect(scrollTopSetter).toHaveBeenCalled(); + }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatUI.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatUI.tsx index ae0fabe5ef2..a87214720ef 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatUI.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatUI.tsx @@ -1,6 +1,7 @@ "use client"; import { + ArrowDown, Bot, Code2, Database, @@ -287,7 +288,7 @@ const ChatUI: React.FC = ({ // Code Interpreter state (using custom hook) const codeInterpreter = useCodeInterpreter(); - const chatEndRef = useRef(null); + const chatScrollRef = useRef(null); // Fetch MCP servers and toolsets const loadMCPServers = async () => { @@ -516,18 +517,20 @@ const ChatUI: React.FC = ({ }, [accessToken, apiKeySource, apiKey, endpointType, customProxyBaseUrl, selectedAgent]); useEffect(() => { - // Scroll to the bottom of the chat whenever chatHistory updates - if (chatEndRef.current) { - // Add a small delay to ensure content is rendered - setTimeout(() => { - chatEndRef.current?.scrollIntoView({ - behavior: "smooth", - block: "end", // Keep the scroll position at the end - }); - }, 100); - } + const el = chatScrollRef.current; + if (!el || chatHistory.at(-1)?.role !== "user") return; + const userMessages = el.querySelectorAll('[data-role="user"]'); + const last = userMessages[userMessages.length - 1]; + if (last) el.scrollTop = last.offsetTop; }, [chatHistory]); + const scrollToLastMessage = () => { + const el = chatScrollRef.current; + const messages = el?.querySelectorAll("[data-role]"); + const last = messages?.[messages.length - 1]; + if (el && last) el.scrollTop = last.offsetTop + last.offsetHeight - el.clientHeight; + }; + const handleCancelRequest = () => { if (abortControllerRef.current) { abortControllerRef.current.abort(); @@ -1801,51 +1804,68 @@ const ChatUI: React.FC = ({ )}
-
- {chatHistory.length === 0 && ( -
-
- )} - - {chatHistory.map((message, index) => ( -
- -
- ))} - - {isLoading && - mcpEvents.length > 0 && - (endpointType === EndpointType.RESPONSES || endpointType === EndpointType.CHAT) && - chatHistory.length > 0 && - chatHistory[chatHistory.length - 1].role === "user" && ( -
-
-
-
-
- Assistant -
- -
+
+
+ {chatHistory.length === 0 && ( +
+
)} - {isLoading && ( -
- -
+ {chatHistory.map((message, index) => ( +
+ +
+ ))} + + {isLoading && + mcpEvents.length > 0 && + (endpointType === EndpointType.RESPONSES || endpointType === EndpointType.CHAT) && + chatHistory.length > 0 && + chatHistory[chatHistory.length - 1].role === "user" && ( +
+
+
+
+
+ Assistant +
+ +
+
+ )} + + {isLoading && ( +
+ +
+ )} + {chatHistory.length > 0 &&
} +
+ {chatHistory.length > 0 && ( + )} -
From 72eb2ef651ca46f90bbbf665e99ac1c58c231d89 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 24 Sep 2026 14:33:02 -0700 Subject: [PATCH 083/120] fix(cost-map): drop the priority input price from vertex gemini-2.5-flash-image (#43050) Co-authored-by: kerry Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/model_prices_and_context_window_backup.json | 2 -- model_prices_and_context_window.json | 2 -- 2 files changed, 4 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 5a48b2d293b..680b6cceab6 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -26212,7 +26212,6 @@ "input_cost_per_token": 3e-07, "input_cost_per_token_batches": 1.5e-07, "input_cost_per_token_flex": 1.5e-07, - "input_cost_per_token_priority": 5.4e-07, "litellm_provider": "vertex_ai-language-models", "max_input_tokens": 32768, "max_output_tokens": 32768, @@ -49416,7 +49415,6 @@ "input_cost_per_token": 3e-07, "input_cost_per_token_batches": 1.5e-07, "input_cost_per_token_flex": 1.5e-07, - "input_cost_per_token_priority": 5.4e-07, "litellm_provider": "vertex_ai-language-models", "max_input_tokens": 32768, "max_output_tokens": 32768, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 5a48b2d293b..680b6cceab6 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -26212,7 +26212,6 @@ "input_cost_per_token": 3e-07, "input_cost_per_token_batches": 1.5e-07, "input_cost_per_token_flex": 1.5e-07, - "input_cost_per_token_priority": 5.4e-07, "litellm_provider": "vertex_ai-language-models", "max_input_tokens": 32768, "max_output_tokens": 32768, @@ -49416,7 +49415,6 @@ "input_cost_per_token": 3e-07, "input_cost_per_token_batches": 1.5e-07, "input_cost_per_token_flex": 1.5e-07, - "input_cost_per_token_priority": 5.4e-07, "litellm_provider": "vertex_ai-language-models", "max_input_tokens": 32768, "max_output_tokens": 32768, From 86ba4fc16f86731ba8736a62b1ea6dadeb2b8a48 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 24 Sep 2026 14:35:03 -0700 Subject: [PATCH 084/120] feat(compat-matrix): resolve and install the Claude Code CLI per run (#43038) * feat(compat-matrix): resolve and install the Claude Code CLI per run * chore(compat-matrix): drop the installer header comment --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- .github/workflows/compat-matrix-image.yml | 13 +++- tests/e2e/claude_code/cli_driver.py | 1 + tests/e2e/claude_code/cron_vm/Dockerfile | 7 -- tests/e2e/claude_code/cron_vm/README.md | 71 +++++++++++-------- .../cron_vm/install_claude_code.sh | 38 ++++++++++ tests/e2e/claude_code/cron_vm/run_daily.sh | 62 +++++++++++----- 6 files changed, 137 insertions(+), 55 deletions(-) create mode 100755 tests/e2e/claude_code/cron_vm/install_claude_code.sh diff --git a/.github/workflows/compat-matrix-image.yml b/.github/workflows/compat-matrix-image.yml index c554096cd7e..f08792c904c 100644 --- a/.github/workflows/compat-matrix-image.yml +++ b/.github/workflows/compat-matrix-image.yml @@ -4,6 +4,7 @@ on: pull_request: paths: - tests/e2e/claude_code/cron_vm/** + - tests/e2e/claude_code/pr_gate_version_resolver.py - .github/workflows/compat-matrix-image.yml workflow_dispatch: @@ -28,6 +29,14 @@ jobs: - name: Build the Render cron image run: docker build -f tests/e2e/claude_code/cron_vm/Dockerfile -t compat-matrix:${{ github.sha }} tests/e2e - - name: Run the pinned binaries as the cron user + - name: Resolve and install the Claude Code CLI as the cron user run: | - docker run --rm compat-matrix:${{ github.sha }} bash -c 'set -e; whoami; claude --version; gh --version; uv --version' + docker run --rm compat-matrix:${{ github.sha }} bash -c ' + set -euo pipefail + whoami + gh --version + uv --version + version="$(uv run --no-project --python 3.12 python /opt/litellm/tests/e2e/claude_code/pr_gate_version_resolver.py)" + /opt/litellm/tests/e2e/claude_code/cron_vm/install_claude_code.sh "${version}" /tmp/claude-cli + /tmp/claude-cli/claude --version + ' diff --git a/tests/e2e/claude_code/cli_driver.py b/tests/e2e/claude_code/cli_driver.py index a01d8ab3e7c..6996849de80 100644 --- a/tests/e2e/claude_code/cli_driver.py +++ b/tests/e2e/claude_code/cli_driver.py @@ -297,6 +297,7 @@ def run_claude( } env["ANTHROPIC_BASE_URL"] = base_url env["ANTHROPIC_AUTH_TOKEN"] = api_key + env["DISABLE_AUTOUPDATER"] = "1" # Hand the CLI a fresh empty HOME so a compromised claude package # or a model-directed Read tool call can't see the runtime user's # real dotfiles. Created here, removed in the `finally` below diff --git a/tests/e2e/claude_code/cron_vm/Dockerfile b/tests/e2e/claude_code/cron_vm/Dockerfile index 623d6b1840a..05c1387c5a2 100644 --- a/tests/e2e/claude_code/cron_vm/Dockerfile +++ b/tests/e2e/claude_code/cron_vm/Dockerfile @@ -4,8 +4,6 @@ ARG GH_VERSION=2.101.0 ARG GH_SHA256=9bca2d1c16825f109907a23307628a2f0698fbf99662b73a5cf0b020293072b8 ARG UV_VERSION=0.10.9 ARG UV_SHA256=20d79708222611fa540b5c9ed84f352bcd3937740e51aacc0f8b15b271c57594 -ARG CLAUDE_CODE_VERSION=2.1.228 -ARG CLAUDE_CODE_SHA256=d535985e6941a3eb00179ccd7f52ceb0c6623a0305a518ebc4e6514f84a94c99 SHELL ["/bin/bash", "-o", "pipefail", "-c"] @@ -23,11 +21,6 @@ RUN curl -fsSLo /tmp/uv.tar.gz "https://github.com/astral-sh/uv/releases/downloa && tar -xzf /tmp/uv.tar.gz -C /usr/local/bin --strip-components=1 uv-x86_64-unknown-linux-gnu/uv \ && rm /tmp/uv.tar.gz -RUN curl -fsSLo /tmp/claude "https://downloads.claude.ai/claude-code-releases/${CLAUDE_CODE_VERSION}/linux-x64/claude" \ - && echo "${CLAUDE_CODE_SHA256} /tmp/claude" | sha256sum -c - \ - && install -m 0755 /tmp/claude /usr/local/bin/claude \ - && rm /tmp/claude - RUN groupadd --gid 1000 populator && useradd --uid 1000 --gid 1000 --create-home populator ENV HOME=/home/populator \ diff --git a/tests/e2e/claude_code/cron_vm/README.md b/tests/e2e/claude_code/cron_vm/README.md index ed2bf4ab436..4c16b6e237a 100644 --- a/tests/e2e/claude_code/cron_vm/README.md +++ b/tests/e2e/claude_code/cron_vm/README.md @@ -16,16 +16,18 @@ than as a GitHub Action or on a dedicated VM. Trade-offs: clone of litellm plus a cold `uv sync`. That adds a few minutes on top of the ~10 minute test run; the job's 12 hour ceiling is nowhere near. -- ⚠️ The Claude Code CLI version under test is pinned in the - `Dockerfile` (`CLAUDE_CODE_VERSION` + its checksum). Bumping it is a - PR, see the gotchas below. +- ✅ The Claude Code CLI under test is chosen on every run (the newest + npm release published at least 3 days ago) and downloaded + checksum-verified, so the matrix follows CLI releases without a PR; + see the gotchas for pinning a run. ## Layout | File | Purpose | | --- | --- | -| `Dockerfile` | The image Render builds: Debian bookworm-slim plus pinned, checksum-verified `gh`, `uv`, and the Claude Code CLI, with this `tests/e2e/` tree copied to `/opt/litellm/tests/e2e/`. Runs as the non-root user `populator` (uid/gid 1000, which is what Render's secret files are readable by). | -| `run_daily.sh` | The actual cron job. Resolves versions, clones the worktree, boots the proxy, runs pytest, builds the JSON, opens (or updates) a docs PR, sweeps stale compat-matrix PRs. | +| `Dockerfile` | The image Render builds: Debian bookworm-slim plus pinned, checksum-verified `gh` and `uv`, with this `tests/e2e/` tree copied to `/opt/litellm/tests/e2e/`. Runs as the non-root user `populator` (uid/gid 1000, which is what Render's secret files are readable by). | +| `run_daily.sh` | The actual cron job. Resolves versions, clones the worktree, installs the Claude Code CLI under test, boots the proxy, runs pytest, builds the JSON, opens (or updates) a docs PR, sweeps stale compat-matrix PRs. | +| `install_claude_code.sh` | Downloads one Claude Code release (` `) from the vendor's native release channel, verifies it against the sha256 in that release's `manifest.json`, and refuses a binary whose `--version` disagrees. Run by the cron and by the `compat-matrix-image` GitHub workflow. | | `build_matrix.py` | Tiny Python CLI that wraps `claude_code.matrix_builder.build_from_paths`. Exists only because the bash script needs *some* way to render the per-cell aggregation, and the builder is already Python. | | `check_regressions.py` | Tiny Python CLI that wraps `claude_code.matrix_builder.find_regressions`. Diffs the freshly built matrix against the currently-published one and exits `3` if any cell flipped green→red, which gates auto-merge. | | `litellm-compat-matrix.env.example` | The service's env vars, one per line, with what each is for. | @@ -35,10 +37,7 @@ than as a GitHub Action or on a dedicated VM. Trade-offs: 1. **Resolves the latest LiteLLM final release tag** (newest bare `vX.Y.Z`, skipping `-rc.N`/`-dev.N` pre-releases) by paging the GitHub Releases API (`curl | jq`). -2. **Reads the Claude Code CLI version** via `claude --version`. That - is whatever the `Dockerfile` pins; the job never upgrades it on its - own. -3. **Clones the worktree** at `~/litellm-cron-worktree/` (a +2. **Clones the worktree** at `~/litellm-cron-worktree/` (a `--filter=blob:none` clone, so only the checked-out tag's blobs are fetched), `git checkout --force `, then `uv sync --frozen --no-install-project` against a uv-managed CPython 3.12 followed by @@ -52,12 +51,20 @@ than as a GitHub Action or on a dedicated VM. Trade-offs: `claude_code/` so the tree's EKS-harness `conftest.py` (whose imports the stable venv doesn't install) is never loaded. The tag's own `tests/e2e/` is deliberately not used. +3. **Resolves and installs the Claude Code CLI under test**: + `pr_gate_version_resolver.py` (run on the venv, the image has no + Python of its own) picks the newest `@anthropic-ai/claude-code` npm + release published at least 3 days ago, the same buffer the PR gate + uses, unless `CLAUDE_CODE_VERSION` pins one, and + `install_claude_code.sh` downloads that release's `linux-x64` binary + into the run's scratch dir, verified against the release manifest. 4. **Boots the proxy** as a `setsid` background process on port `4100` bound to loopback, then polls `/health/liveliness` until it's up. -5. **Runs pytest** on `tests/e2e/claude_code/` with `LITELLM_PROXY_URL` - pointed at the proxy and `COMPAT_RESULTS_PATH` set so the conftest - hook writes the per-test results artifact. Test failures become - `fail` cells in the JSON, not script errors. +5. **Runs pytest** on `tests/e2e/claude_code/` with that CLI first on + `PATH`, `LITELLM_PROXY_URL` pointed at the proxy, and + `COMPAT_RESULTS_PATH` set so the conftest hook writes the per-test + results artifact. Test failures become `fail` cells in the JSON, not + script errors. 6. **Builds `compatibility-matrix.json`** by handing the artifact + manifest to `build_matrix.py`. 7. **Opens or updates a docs PR**: `gh repo clone` of `litellm-docs` @@ -70,7 +77,8 @@ than as a GitHub Action or on a dedicated VM. Trade-offs: branch ... already exists" is treated as success). If the JSON is byte-identical to what `main` already publishes, the push is skipped entirely. These PRs are not gated on a second human review. -8. **Gates auto-merge on a regression check**: before enabling + + **Auto-merge is gated on a regression check**: before enabling auto-merge, `check_regressions.py` diffs the new matrix against the one currently on `main`. Auto-merge (`gh pr merge --auto --squash`) is only enabled when **no cell flipped green→red** — i.e. every @@ -82,7 +90,7 @@ than as a GitHub Action or on a dedicated VM. Trade-offs: auto-merge a prior same-day run enabled is explicitly disabled — so a human reviews before it lands on the public table. The check fails *closed*: if it errors, auto-merge is withheld. -9. **Sweeps stale compat-matrix PRs**: once today's PR exists, every +8. **Sweeps stale compat-matrix PRs**: once today's PR exists, every other open `compat-matrix/*` PR that the publishing account opened from a branch on the docs repo itself is closed (and its bot-owned branch deleted), so at most one compat-matrix PR is ever open — the @@ -165,7 +173,8 @@ curl -fsS -X POST "https://api.render.com/v1/services/${CRON_ID}/deploys" \ curl -fsS "https://api.render.com/v1/services/${CRON_ID}/deploys?limit=1" \ -H "Authorization: Bearer ${RENDER_API_KEY}" -# A run that does NOT open a PR (first-time validation, CLI bumps): +# A run that does NOT open a PR (first-time validation, a CLI pinned +# with CLAUDE_CODE_VERSION): # set SKIP_PUBLISH=1 on the service, trigger a run, then remove it. # The matrix JSON is printed at the end of the run's log (nothing on # the container's disk outlives the run) and saved to @@ -217,21 +226,25 @@ docker run --rm --platform linux/amd64 \ or fine-grained Contents:RW + Pull requests:RW). It is delivered as a file, not an env var, so pytest, the proxy, and the claude CLI never inherit it; manual runs export `GITHUB_TOKEN` instead. -- **Bumping the Claude Code CLI is a PR.** Change `CLAUDE_CODE_VERSION` - in the `Dockerfile` and set `CLAUDE_CODE_SHA256` to the `linux-x64` - checksum from - `https://downloads.claude.ai/claude-code-releases//manifest.json`. - The first run on a new CLI is the riskiest one: if the new CLI - changes its wire format the matrix run can produce systematic - failures, so trigger a `SKIP_PUBLISH=1` run before the next scheduled - fire. `gh` and `uv` bump the same way, with the checksum from the - release's `gh__checksums.txt` and the tarball's `.sha256` - sidecar respectively. +- **The Claude Code CLI is chosen per run, not pinned.** Each run + tests the newest `@anthropic-ai/claude-code` npm release published + at least 3 days ago, downloaded from + `https://downloads.claude.ai/claude-code-releases//linux-x64/claude` + and verified against the sha256 in that release's `manifest.json`. + A CLI release that breaks a cell shows up as a green→red flip, which + withholds auto-merge on that day's docs PR for review. To rerun the + matrix on one specific CLI, set `CLAUDE_CODE_VERSION` on the run. + `gh` and `uv` stay pinned in the `Dockerfile`; bump them in a PR with + the checksum from the release's `gh__checksums.txt` and the + tarball's `.sha256` sidecar respectively. - **A local build on Apple silicon only proves the image assembles.** Under QEMU the Claude Code binary (a Bun executable) dies with - `CPU lacks AVX support` and `gh` panics in the Go runtime, so - `claude --version` and a full run are verified with a - `SKIP_PUBLISH=1` run on Render, not locally. + `CPU lacks AVX support` and `gh` panics in the Go runtime, so the CLI + download and `claude --version` are verified by the + `compat-matrix-image` GitHub workflow (an x86 runner that builds the + image and runs `install_claude_code.sh` in it on every PR touching + this directory) and a full run with a `SKIP_PUBLISH=1` run on Render, + not locally. - **Nothing persists between runs.** A failed run leaves no half-installed venv behind, but also no cache: don't expect a rerun to be faster than the first one. diff --git a/tests/e2e/claude_code/cron_vm/install_claude_code.sh b/tests/e2e/claude_code/cron_vm/install_claude_code.sh new file mode 100755 index 00000000000..123401cc78a --- /dev/null +++ b/tests/e2e/claude_code/cron_vm/install_claude_code.sh @@ -0,0 +1,38 @@ +#!/usr/bin/env bash + +set -Eeuo pipefail + +RELEASES_URL="https://downloads.claude.ai/claude-code-releases" + +log() { printf '==> %s\n' "$*" >&2; } +die() { printf 'ERROR: %s\n' "$*" >&2; exit 1; } + +[[ $# -eq 2 ]] || die "usage: $(basename "$0") " +VERSION="$1" +DEST_DIR="$2" +[[ "${VERSION}" =~ ^[0-9]+\.[0-9]+\.[0-9]+$ ]] || die "not a Claude Code release version: '${VERSION}'" + +mkdir -p "${DEST_DIR}" +MANIFEST="${DEST_DIR}/manifest.json" +curl -fsSL --retry 3 --retry-all-errors --output "${MANIFEST}" "${RELEASES_URL}/${VERSION}/manifest.json" \ + || die "no release manifest for claude code ${VERSION} at ${RELEASES_URL}" +CHECKSUM="$(jq -r '.platforms["linux-x64"].checksum // empty' "${MANIFEST}")" +[[ "${CHECKSUM}" =~ ^[0-9a-f]{64}$ ]] || die "manifest for claude code ${VERSION} carries no linux-x64 sha256" + +log "downloading claude code ${VERSION} (linux-x64)" +DOWNLOAD="${DEST_DIR}/claude.download" +curl -fsSL --retry 3 --retry-all-errors --output "${DOWNLOAD}" "${RELEASES_URL}/${VERSION}/linux-x64/claude" +echo "${CHECKSUM} ${DOWNLOAD}" | sha256sum -c - >/dev/null \ + || die "claude code ${VERSION} sha256 mismatch; refusing to install" +chmod 0755 "${DOWNLOAD}" +mv "${DOWNLOAD}" "${DEST_DIR}/claude" + +PROBE_HOME="$(mktemp -d -t claude-probe-home.XXXXXX)" +trap 'rm -rf "${PROBE_HOME}"' EXIT +REPORTED="$( + env -i HOME="${PROBE_HOME}" PATH="${PATH}" DISABLE_AUTOUPDATER=1 \ + "${DEST_DIR}/claude" --version | awk '{print $1}' +)" || die "claude code ${VERSION} could not run --version" +[[ "${REPORTED}" == "${VERSION}" ]] \ + || die "installed claude code reports '${REPORTED}', expected ${VERSION}" +log "installed claude code ${VERSION} at ${DEST_DIR}/claude" diff --git a/tests/e2e/claude_code/cron_vm/run_daily.sh b/tests/e2e/claude_code/cron_vm/run_daily.sh index 172dd03b614..af01057eff7 100755 --- a/tests/e2e/claude_code/cron_vm/run_daily.sh +++ b/tests/e2e/claude_code/cron_vm/run_daily.sh @@ -7,15 +7,21 @@ # 1. Resolve the latest LiteLLM final release tag from the GitHub # Releases API. # 2. Update a long-lived worktree at $WORKTREE to that tag and `uv sync` it. -# 3. Boot the proxy as a background subprocess on $PROXY_PORT (default +# 3. Resolve the Claude Code CLI version under test (the newest npm +# release at least 3 days old, via pr_gate_version_resolver.py, or +# $CLAUDE_CODE_VERSION when set) and download that release's +# linux-x64 binary into the run's scratch dir, checksum-verified +# against the vendor's release manifest (install_claude_code.sh). +# 4. Boot the proxy as a background subprocess on $PROXY_PORT (default # 4100; a separate port from the human-tended :4000 proxy). -# 4. Run `pytest tests/e2e/claude_code/` against the proxy. Test -# failures become `fail` cells in the JSON, not script errors. -# 5. Hand the per-test results artifact + manifest to a small Python +# 5. Run `pytest tests/e2e/claude_code/` against the proxy with that +# CLI first on PATH. Test failures become `fail` cells in the JSON, +# not script errors. +# 6. Hand the per-test results artifact + manifest to a small Python # CLI (`build_matrix.py`) that wraps the existing # `matrix_builder.build_from_paths` to produce the published # compatibility-matrix.json. -# 6. `gh repo clone` litellm-docs, write the JSON to a deterministic +# 7. `gh repo clone` litellm-docs, write the JSON to a deterministic # branch (`compat-matrix/--`), commit, # push the branch straight to BerriAI/litellm-docs (mateo-berri has # write access), `gh pr create`, then — *only if no cell regressed @@ -23,7 +29,7 @@ # auto-merge so the PR merges itself once required checks pass. A # green→red regression leaves auto-merge off for human review; an # already-red cell (red→red) does not block. -# 7. Sweep stale compat-matrix PRs: once today's PR exists, close any +# 8. Sweep stale compat-matrix PRs: once today's PR exists, close any # other open `compat-matrix/*` PR (and delete its bot-owned branch) # so at most ONE compat-matrix PR is ever open — the newest. A # gate-withheld PR that nobody triages is superseded by the next @@ -33,7 +39,7 @@ # rather than spawning a new one. If the JSON is byte-identical to the # docs branch, we skip the push entirely. # -# Required commands on $PATH: git, uv, gh, jq, curl, claude. +# Required commands on $PATH: git, uv, gh, jq, curl. # Required state: a litellm checkout at $LITELLM_REPO (this file lives in # it); $WORKTREE is created on first run. # @@ -51,6 +57,9 @@ DOCS_BRANCH="${DOCS_BRANCH:-main}" DOCS_TARGET_PATH="${DOCS_TARGET_PATH:-src/data/compatibility-matrix.json}" SKIP_PUBLISH="${SKIP_PUBLISH:-0}" PYTEST_K="${PYTEST_K:-}" +# Empty means "resolve it": the newest @anthropic-ai/claude-code npm +# release published at least 3 days ago. Set it to pin a manual run. +CLAUDE_CODE_VERSION="${CLAUDE_CODE_VERSION:-}" # The e2e suite uses PEP 695 `type` aliases, so the venv needs Python # >= 3.12 (also what repo CI runs) even when the host's system python is # older. uv fetches a managed CPython of this version on first use -- @@ -108,7 +117,7 @@ trap cleanup EXIT INT TERM log() { printf '==> %s\n' "$*" >&2; } die() { printf 'ERROR: %s\n' "$*" >&2; exit 1; } -for cmd in git uv gh jq curl claude; do +for cmd in git uv gh jq curl; do command -v "${cmd}" >/dev/null 2>&1 || die "missing required command: ${cmd}" done @@ -199,10 +208,6 @@ LITELLM_VERSION="$( [[ -n "${LITELLM_VERSION}" ]] || die "could not resolve latest PEP 440 final release (vX.Y.Z) in 5 pages of releases" log "resolved litellm: ${LITELLM_VERSION}" -CLAUDE_CODE_VERSION="$(claude --version 2>/dev/null | awk '{print $1}')" -[[ -n "${CLAUDE_CODE_VERSION}" ]] || die "could not read 'claude --version'" -log "local claude code: ${CLAUDE_CODE_VERSION}" - # --------------------------------------------------------------------------- # 2. Update the worktree to that tag # --------------------------------------------------------------------------- @@ -314,7 +319,27 @@ PROXY_CONFIG="${WORKTREE}/tests/e2e/claude_code/test_config.yaml" [[ -f "${PROXY_CONFIG}" ]] || die "proxy config not found at ${PROXY_CONFIG} (shim incomplete?)" # --------------------------------------------------------------------------- -# 3. Boot the proxy +# 3. Resolve and install the Claude Code CLI under test +# --------------------------------------------------------------------------- + +# The resolver is stdlib-only, but the image ships no python of its +# own, so it runs on the venv the sync above just built. Its 3-day +# publish-age buffer (PRD #26476) keeps a release that gets pulled or +# patched within days from ever driving the published matrix. +if [[ -z "${CLAUDE_CODE_VERSION}" ]]; then + CLAUDE_CODE_VERSION="$( + cd "${WORKTREE}" \ + && "${WORKTREE_UV}" run --no-sync python "${POPULATOR_DIR}/../pr_gate_version_resolver.py" + )" || die "could not resolve the Claude Code version to test" + log "resolved claude code: ${CLAUDE_CODE_VERSION}" +else + log "CLAUDE_CODE_VERSION set; testing claude code ${CLAUDE_CODE_VERSION}" +fi +CLAUDE_CLI_DIR="${WORKDIR}/claude-cli" +"${POPULATOR_DIR}/install_claude_code.sh" "${CLAUDE_CODE_VERSION}" "${CLAUDE_CLI_DIR}" + +# --------------------------------------------------------------------------- +# 4. Boot the proxy # --------------------------------------------------------------------------- log "starting proxy on 127.0.0.1:${PROXY_PORT}" @@ -350,7 +375,7 @@ curl -fsS "${HEALTH_URL}" >/dev/null \ || { tail -50 "${WORKDIR}/proxy.log" >&2; die "proxy did not become healthy"; } # --------------------------------------------------------------------------- -# 4. Run pytest +# 5. Run pytest # --------------------------------------------------------------------------- RESULTS_JSON="${WORKDIR}/compat-results.json" @@ -374,6 +399,7 @@ set +e && LITELLM_PROXY_URL="http://127.0.0.1:${PROXY_PORT}" \ LITELLM_MASTER_KEY="${PROXY_API_KEY}" \ COMPAT_RESULTS_PATH="${RESULTS_JSON}" \ + PATH="${CLAUDE_CLI_DIR}:${PATH}" \ "${WORKTREE_UV}" run --no-sync pytest "${PYTEST_ARGS[@]}" ) PYTEST_EXIT=$? @@ -386,7 +412,7 @@ log "pytest exit code: ${PYTEST_EXIT} (failures become 'fail' cells, not script [[ -f "${RESULTS_JSON}" ]] || die "pytest did not produce ${RESULTS_JSON}" # --------------------------------------------------------------------------- -# 5. Build the matrix JSON +# 6. Build the matrix JSON # --------------------------------------------------------------------------- MATRIX_JSON="${WORKDIR}/compatibility-matrix.json" @@ -402,7 +428,7 @@ log "building ${MATRIX_JSON}" ) # --------------------------------------------------------------------------- -# 6. Open a docs-repo PR +# 7. Open a docs-repo PR # --------------------------------------------------------------------------- if [[ "${SKIP_PUBLISH}" == "1" ]]; then @@ -634,7 +660,9 @@ else || die "auto-merge still armed on ${BRANCH_NAME} (enabled ${AUTOMERGE_ARMED}) after --disable-auto" fi -# --- Stale-PR sweep ---------------------------------------------------------- +# --------------------------------------------------------------------------- +# 8. Sweep stale compat-matrix PRs +# --------------------------------------------------------------------------- # Keep at most ONE compat-matrix PR open: today's. Any other open # `compat-matrix/*` PR is a leftover from a day whose regression gate # withheld auto-merge and nobody triaged it; the PR we just opened or From 2be2d68ac367a84bc7035948202d59c5fa7402ea Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Thu, 24 Sep 2026 14:48:35 -0700 Subject: [PATCH 085/120] test(e2e/ui): hide the LiteAdmin button in the shared admin session (#43033) #42443 pins a floating LiteAdmin button to the bottom-right corner, where it covers the logs page's next-page control. Flip the per-user Hide LiteAdmin switch during global setup so every spec reusing the admin storage state loads with the button hidden. --- tests/e2e/ui/globalSetup.ts | 4 ++++ tests/e2e/ui/helpers/navigation.ts | 14 ++++++++++++++ 2 files changed, 18 insertions(+) diff --git a/tests/e2e/ui/globalSetup.ts b/tests/e2e/ui/globalSetup.ts index 9447c93a72e..8b3ad204767 100644 --- a/tests/e2e/ui/globalSetup.ts +++ b/tests/e2e/ui/globalSetup.ts @@ -2,6 +2,7 @@ import { chromium, expect, request } from "@playwright/test"; import { users, Role, STORAGE_PATHS } from "./fixtures/users"; import { ARTIFACT_DIR, UI_BASE_URL } from "./constants"; import { expectUnrestrictedDashboard, setInvitedUserPassword } from "./helpers/userOnboarding"; +import { hideLiteAdmin } from "./helpers/navigation"; import * as fs from "fs"; import * as path from "path"; @@ -75,6 +76,9 @@ async function globalSetup() { if (await dismiss.isVisible({ timeout: 1_500 }).catch(() => false)) { await dismiss.click(); } + if (role === Role.ProxyAdmin) { + await hideLiteAdmin(page); + } // The login flow stores a post-login return URL in the litellm_return_url // cookie. If the snapshot captures it before the app consumes it, every // test inheriting this storageState gets yanked to that stale URL the diff --git a/tests/e2e/ui/helpers/navigation.ts b/tests/e2e/ui/helpers/navigation.ts index e0e7b4da396..78288fab743 100644 --- a/tests/e2e/ui/helpers/navigation.ts +++ b/tests/e2e/ui/helpers/navigation.ts @@ -62,6 +62,20 @@ export async function dismissFeedbackPopup(page: PlaywrightPage): Promise } } +export async function hideLiteAdmin(page: PlaywrightPage): Promise { + await page.getByRole("button", { name: /Account menu/i }).click(); + const panel = page.getByTestId("sidebar-account-menu-panel"); + await expect(panel).toBeVisible({ timeout: 5_000 }); + const toggle = panel.getByRole("switch", { name: "Toggle hide LiteAdmin" }); + if ((await toggle.getAttribute("aria-checked")) !== "true") { + await toggle.click(); + } + await expect(toggle).toHaveAttribute("aria-checked", "true"); + await page.keyboard.press("Escape"); + await expect(panel).toBeHidden(); + await expect(page.getByRole("button", { name: "LiteAdmin", exact: true })).toBeHidden(); +} + /** * Click on a team ID in the table. Team IDs are rendered differently depending * on the component version — try button first (Tremor Button), fall back to From 1628978db707dfff9d17844d345e0587e893c307 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 24 Sep 2026 14:49:21 -0700 Subject: [PATCH 086/120] ci: fix the litellm-tests unit job (sysmon, codecov on failure, env -i allowlist, selection errors, reruns param) (#42900) * ci: fix the litellm-tests unit job with sysmon coverage, an env allowlist and coverage upload on failure * ci: fail the unit shard when circleci tests split errors * ci: exit the unit shard cleanly when circleci tests split assigns it no files --------- Co-authored-by: yuneng --- .circleci/tests.yml | 33 ++++++++++++++++++++++----------- 1 file changed, 22 insertions(+), 11 deletions(-) diff --git a/.circleci/tests.yml b/.circleci/tests.yml index 6ee1eb662e6..1afb935453f 100644 --- a/.circleci/tests.yml +++ b/.circleci/tests.yml @@ -74,6 +74,7 @@ commands: steps: - run: name: Install Codecov CLI (pinned v11.3.1) + when: always command: | curl -sSLf -o /tmp/codecov https://cli.codecov.io/v11.3.1/linux/codecov curl -sSLf -o /tmp/codecov.SHA256SUM https://cli.codecov.io/v11.3.1/linux/codecov.SHA256SUM @@ -90,7 +91,6 @@ commands: uv run --no-sync python -c "import litellm_enterprise; print('litellm-enterprise OK:', litellm_enterprise.__file__)" setup_test_deps: steps: - - checkout - install_uv - install_rust - restore_cache: @@ -165,9 +165,6 @@ commands: jobs: unit: parameters: - tests_path: - type: string - default: tests/unit flag: type: string default: unit @@ -180,27 +177,39 @@ jobs: pull_request_url: type: string default: "" + reruns: + type: integer + default: 0 machine: image: ubuntu-2204:2024.04.1 resource_class: large working_directory: ~/project parallelism: << parameters.shards >> environment: + COVERAGE_CORE: sysmon LITELLM_LOCAL_MODEL_COST_MAP: "True" steps: - - setup_test_deps + - checkout - skip_unless_relevant: base_ref: << parameters.base_ref >> pull_request_url: << parameters.pull_request_url >> + - setup_test_deps - run: - name: "Run << parameters.tests_path >> shard" + name: "Run << parameters.flag >> shard" no_output_timeout: 20m command: | mkdir -p test-results/<< parameters.flag >> - mapfile -t files < <(find << parameters.tests_path >> -name 'test_*.py' | sort | circleci tests split --split-by=timings --timings-type=filename) - if [ "${#files[@]}" -eq 0 ]; then echo "shard ${CIRCLE_NODE_INDEX} received no << parameters.tests_path >> files; nothing to run"; exit 0; fi + selection="$(find tests/unit -name 'test_*.py' | sort)" || { echo "test selection failed for << parameters.flag >>"; exit 1; } + [ -n "${selection}" ] || { echo "test selection produced no files for << parameters.flag >>"; exit 1; } + shard="$(printf '%s\n' "${selection}" | circleci tests split --split-by=timings --timings-type=filename)" || { echo "circleci tests split failed for << parameters.flag >>"; exit 1; } + [ -n "${shard}" ] || { echo "shard ${CIRCLE_NODE_INDEX} received no << parameters.flag >> files; nothing to run"; exit 0; } + mapfile -t files < <(printf '%s\n' "${shard}") + rerun_args=(-p no:rerunfailures) + if [ "<< parameters.reruns >>" -gt 0 ]; then rerun_args=(--reruns << parameters.reruns >> --reruns-delay 1 --rerun-except "from pytest-timeout"); fi + test_env=(PATH="$PATH" HOME="$HOME" CI=true COVERAGE_CORE="$COVERAGE_CORE" LITELLM_LOCAL_MODEL_COST_MAP="$LITELLM_LOCAL_MODEL_COST_MAP") set +e - uv run --no-sync pytest "${files[@]}" -p no:rerunfailures -p no:pytest-retry --timeout=90 -n 4 --dist=loadscope --tb=short --durations=20 -o junit_family=xunit1 --junitxml=test-results/<< parameters.flag >>/junit.xml --cov=./litellm --cov=./enterprise/litellm_enterprise --cov-report=xml:coverage.xml --cov-config=pyproject.toml + env -i "${test_env[@]}" \ + uv run --no-sync pytest "${files[@]}" "${rerun_args[@]}" -p no:pytest-retry --timeout=90 -n 4 --dist=loadscope --tb=short --durations=20 -o junit_family=xunit1 --junitxml=test-results/<< parameters.flag >>/junit.xml --cov=./litellm --cov=./enterprise/litellm_enterprise --cov-report=xml:coverage.xml --cov-config=pyproject.toml status=$? set -e if [ "$status" -eq 5 ]; then echo "pytest collected no tests from the shard; passing"; exit 0; fi @@ -224,6 +233,7 @@ jobs: resource_class: large working_directory: ~/project steps: + - checkout - setup_test_deps - run: name: Checkout litellm-docs @@ -250,16 +260,17 @@ jobs: resource_class: large working_directory: ~/project steps: - - setup_test_deps + - checkout - skip_unless_relevant: base_ref: << parameters.base_ref >> pull_request_url: << parameters.pull_request_url >> + - setup_test_deps - start_postgres: image: postgres:16@sha256:e17e86066e5ef83e0952a9347f5c792b7ece00972e2aa787a6986f471b3dd3d5 - start_redis - run: name: Run owned integration contracts - command: bash .circleci/scripts/run_integration.sh << parameters.suite >> + command: env -i PATH="$PATH" HOME="$HOME" CIRCLE_SHA1="$CIRCLE_SHA1" CIRCLE_WORKFLOW_ID="$CIRCLE_WORKFLOW_ID" bash .circleci/scripts/run_integration.sh << parameters.suite >> no_output_timeout: 15m - run: name: Stop owned database and Redis From b72a03050140e8514ccefdab15ea9b277569ba60 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 24 Sep 2026 14:56:52 -0700 Subject: [PATCH 087/120] test: take keys out of the legacy proxy, enterprise and mcp unit tests before moving them (#42901) * ci: fix the litellm-tests unit job with sysmon coverage, an env allowlist and coverage upload on failure * test: replace key-dependent proxy, enterprise and mcp unit tests with synthetic values and integration and e2e coverage * test: drop key reads at the legacy proxy, enterprise and mcp paths and wire the gemini pass-through split * ci: fail the unit shard when circleci tests split errors * test: drop restating comments from the gemini pass-through split * ci: exit the unit shard cleanly when circleci tests split assigns it no files --------- Co-authored-by: yuneng --- .github/workflows/test-unit-proxy-db.yml | 2 +- .../test_token_counter_gemini_contents_e2e.py | 78 ++++ .../test_prometheus_unit_tests.py | 10 +- .../routing/test_user_config_routing.py | 103 +++++ .../mcp_tests/test_aresponses_api_with_mcp.py | 377 +---------------- .../test_aresponses_api_with_mcp_providers.py | 389 ++++++++++++++++++ .../test_proxy_custom_auth.py | 5 +- .../test_proxy_pass_user_config.py | 114 ----- tests/proxy_unit_tests/test_proxy_server.py | 53 --- .../test_proxy_server_gemini_pass_through.py | 51 +++ .../test_proxy_token_counter.py | 159 +------ tests/proxy_unit_tests/test_proxy_utils.py | 4 +- 12 files changed, 652 insertions(+), 693 deletions(-) create mode 100644 tests/e2e/llm_translation/test_token_counter_gemini_contents_e2e.py create mode 100644 tests/integration/routing/test_user_config_routing.py create mode 100644 tests/mcp_tests/test_aresponses_api_with_mcp_providers.py delete mode 100644 tests/proxy_unit_tests/test_proxy_pass_user_config.py create mode 100644 tests/proxy_unit_tests/test_proxy_server_gemini_pass_through.py diff --git a/.github/workflows/test-unit-proxy-db.yml b/.github/workflows/test-unit-proxy-db.yml index 4af7a161984..73015ac6e02 100644 --- a/.github/workflows/test-unit-proxy-db.yml +++ b/.github/workflows/test-unit-proxy-db.yml @@ -106,6 +106,7 @@ jobs: - test-group: proxy-server-core test-path: >- tests/proxy_unit_tests/test_proxy_server.py + tests/proxy_unit_tests/test_proxy_server_gemini_pass_through.py tests/proxy_unit_tests/test_aproxy_startup.py workers: 4 dist: loadscope @@ -115,7 +116,6 @@ jobs: tests/proxy_unit_tests/test_proxy_config_unit_test.py tests/proxy_unit_tests/test_proxy_routes.py tests/proxy_unit_tests/test_server_root_path.py - tests/proxy_unit_tests/test_proxy_pass_user_config.py tests/proxy_unit_tests/test_proxy_token_counter.py tests/proxy_unit_tests/test_request_size_limit_middleware.py tests/proxy_unit_tests/test_multipart_bypass_repro.py diff --git a/tests/e2e/llm_translation/test_token_counter_gemini_contents_e2e.py b/tests/e2e/llm_translation/test_token_counter_gemini_contents_e2e.py new file mode 100644 index 00000000000..b2b8f46ab0a --- /dev/null +++ b/tests/e2e/llm_translation/test_token_counter_gemini_contents_e2e.py @@ -0,0 +1,78 @@ +"""Live e2e: `/utils/token_counter?call_endpoint=true` counts Gemini `contents` upstream. + +Google's countTokens API is the only tokenizer that knows Gemini's real token +boundaries, so the proxy must forward `contents` to it for both the AI Studio and +Vertex deployments and hand back the provider's `promptTokensDetails`. Claude on +Vertex is covered by `/v1/messages/count_tokens`; this is the Gemini `contents` +route the claude_code rows never reach +""" + +from __future__ import annotations + +import pytest +from e2e_config import unique_marker +from e2e_http import require_successful_call +from proxy_client import ProxyClient +from pydantic import BaseModel + +pytestmark = pytest.mark.e2e + +GEMINI_DEPLOYMENTS = ("gemini-2.5-flash", "gemini-2.5-flash-vertex") + + +class _Part(BaseModel): + text: str + + +class _Content(BaseModel): + parts: tuple[_Part, ...] + + +class _TokenCountBody(BaseModel): + model: str + contents: tuple[_Content, ...] + + +class _CallEndpoint(BaseModel): + call_endpoint: bool = True + + +class _ModalityTokens(BaseModel): + modality: str + tokenCount: int + + +class _CountTokensUpstream(BaseModel): + totalTokens: int + promptTokensDetails: tuple[_ModalityTokens, ...] + + +class _TokenCountResponse(BaseModel): + total_tokens: int + request_model: str + model_used: str + tokenizer_type: str + original_response: _CountTokensUpstream + + +class TestGeminiContentsTokenCounting: + @pytest.mark.parametrize("model", GEMINI_DEPLOYMENTS) + def test_contents_are_counted_by_the_provider_endpoint( + self, proxy: ProxyClient, scoped_key: str, model: str + ) -> None: + text = f"Hello world, how are you doing today? {unique_marker()}" + body = _TokenCountBody(model=model, contents=(_Content(parts=(_Part(text=text),)),)) + + result = proxy.transport.send( + "/utils/token_counter", + headers=proxy.transport.bearer(scoped_key), + json=body, + params=_CallEndpoint(), + ) + + require_successful_call(result) + counted = _TokenCountResponse.model_validate_json(result.body) + assert counted.request_model == model, counted + assert counted.original_response.totalTokens == counted.total_tokens > 0, counted + assert counted.original_response.promptTokensDetails, counted + assert all(detail.tokenCount > 0 for detail in counted.original_response.promptTokensDetails), counted diff --git a/tests/enterprise/litellm_enterprise/integrations/test_prometheus_unit_tests.py b/tests/enterprise/litellm_enterprise/integrations/test_prometheus_unit_tests.py index 28fd03daf37..5b26dad269e 100644 --- a/tests/enterprise/litellm_enterprise/integrations/test_prometheus_unit_tests.py +++ b/tests/enterprise/litellm_enterprise/integrations/test_prometheus_unit_tests.py @@ -13,7 +13,6 @@ import asyncio from dotenv import load_dotenv load_dotenv() -import os from unittest.mock import MagicMock @@ -165,9 +164,9 @@ async def test_prometheus_metric_tracking(): "model_name": "gpt-5-mini", # openai model name "litellm_params": { # params for litellm completion/embedding call "model": "azure/gpt-4.1-mini", - "api_key": os.getenv("AZURE_AI_API_KEY"), - "api_version": os.getenv("AZURE_AI_API_VERSION"), - "api_base": os.getenv("AZURE_AI_API_BASE"), + "api_key": "sk-azure-unit-test", + "api_version": "2025-01-01-preview", + "api_base": "https://unit-test.openai.azure.com", }, "model_info": {"id": "azure-model-id"}, }, @@ -180,9 +179,6 @@ async def test_prometheus_metric_tracking(): }, ], provider_budget_config=provider_budget_config, - redis_host=os.getenv("REDIS_HOST"), - redis_port=int(os.getenv("REDIS_PORT", 6379)), - redis_password=os.getenv("REDIS_PASSWORD"), ) try: diff --git a/tests/integration/routing/test_user_config_routing.py b/tests/integration/routing/test_user_config_routing.py new file mode 100644 index 00000000000..1c50a78201b --- /dev/null +++ b/tests/integration/routing/test_user_config_routing.py @@ -0,0 +1,103 @@ +import json +import uuid +from pathlib import Path +from typing import Final + +from integration._support.client import Gateway +from integration._support.process import owned_proxy +from integration._support.wire import Reply, Request, wire_server + +USER_KEY: Final = "sk-user-supplied-" + uuid.uuid4().hex + + +def _completion(request: Request) -> Reply: + if request.target != "/v1/chat/completions": + return Reply(status=404, body=b"{}") + body: Final = json.loads(request.body) + return Reply( + body=json.dumps( + { + "id": "chatcmpl-user-config", + "object": "chat.completion", + "created": 0, + "model": body["model"], + "choices": [ + {"index": 0, "message": {"role": "assistant", "content": "routed"}, "finish_reason": "stop"} + ], + "usage": {"prompt_tokens": 3, "completion_tokens": 1, "total_tokens": 4}, + } + ).encode() + ) + + +def _user_config(upstream_url: str) -> dict[str, object]: + return { + "model_list": [ + { + "model_name": "user-config-deployment", + "litellm_params": { + "model": "openai/gpt-4.1-mini", + "api_base": upstream_url + "/v1", + "api_key": USER_KEY, + }, + } + ], + "num_retries": 0, + } + + +def _opt_in_config(directory: Path, upstream_url: str) -> Path: + config: Final = directory / "allow_client_side_credentials_config.yaml" + config.write_text( + json.dumps( + { + "model_list": [ + { + "model_name": "admin-deployment", + "litellm_params": { + "model": "openai/gpt-4.1-mini", + "api_base": upstream_url + "/v1", + "api_key": "sk-admin-configured", + }, + } + ], + "general_settings": { + "master_key": "os.environ/LITELLM_MASTER_KEY", + "database_url": "os.environ/DATABASE_URL", + "store_model_in_db": True, + "allow_client_side_credentials": True, + }, + } + ) + ) + return config + + +def _request_body(upstream_url: str) -> dict[str, object]: + return { + "model": "user-config-deployment", + "messages": [{"role": "user", "content": "user config control"}], + "user_config": _user_config(upstream_url), + } + + +def test_user_config_routes_to_the_user_supplied_deployment_when_opted_in(gateway: Gateway, tmp_path: Path) -> None: + with wire_server(_completion) as upstream: + config: Final = _opt_in_config(tmp_path, upstream.url) + with owned_proxy(gateway, tmp_path, {}, config=config) as candidate: + response: Final = candidate.request("POST", "/v1/chat/completions", _request_body(upstream.url)) + assert response.status_code == 200, response.text + assert response.json()["choices"][0]["message"]["content"] == "routed" + outbound: Final = tuple(upstream.received.get_nowait() for _ in range(upstream.received.qsize())) + completions: Final = tuple(request for request in outbound if request.target == "/v1/chat/completions") + assert len(completions) == 1, outbound + assert completions[0].headers["authorization"] == f"Bearer {USER_KEY}" + assert json.loads(completions[0].body)["model"] == "gpt-4.1-mini" + + +def test_user_config_is_rejected_without_the_opt_in(gateway: Gateway) -> None: + with wire_server(_completion) as upstream: + response: Final = gateway.request("POST", "/v1/chat/completions", _request_body(upstream.url)) + assert response.status_code == 401, response.text + assert "user_config is not allowed in request body" in response.text + assert upstream.received.empty() diff --git a/tests/mcp_tests/test_aresponses_api_with_mcp.py b/tests/mcp_tests/test_aresponses_api_with_mcp.py index eb6f78b57a1..de0dc78af43 100644 --- a/tests/mcp_tests/test_aresponses_api_with_mcp.py +++ b/tests/mcp_tests/test_aresponses_api_with_mcp.py @@ -1,5 +1,3 @@ -import logging -import os import pytest from mcp.types import Tool as MCPTool from typing import List, Any, cast @@ -846,161 +844,6 @@ async def test_streaming_mcp_events_validation(): assert mock_get_tools.called, "MCP tools should have been fetched" -@pytest.mark.asyncio -@pytest.mark.parametrize( - "model", - [ - pytest.param("gpt-4o-mini", id="openai"), - pytest.param("claude-haiku-4-5", id="anthropic"), - ], -) -async def test_streaming_responses_api_with_mcp_tools( - model: str, caplog: pytest.LogCaptureFixture -): - """ - Test the streaming responses API with MCP tools when using server_url="litellm_proxy" - - Under the hood the follow occurs - - - MCP: responses called litellm MCP manager.list_tools (MOCKED) - - Request 1: Made to model under test with fetched tools (REAL LLM CALL) - - MCP: Execute tool call from request 1 and returns result (MOCKED) - - Request 2: Made to model under test with fetched tools and tool results (REAL LLM CALL) - - Return the user the result of request 2 - """ - # Skip test if API keys are not set for the respective models - if ("claude" in model.lower() or "anthropic" in model.lower()) and not os.getenv( - "ANTHROPIC_API_KEY" - ): - pytest.skip("ANTHROPIC_API_KEY not set, skipping anthropic model test") - if ("gpt" in model.lower() or "openai" in model.lower()) and not os.getenv( - "OPENAI_API_KEY" - ): - pytest.skip("OPENAI_API_KEY not set, skipping openai model test") - - from unittest.mock import AsyncMock, patch - - print("🧪 Testing basic streaming with MCP tools...") - - # Mock MCP tools that would be returned from the manager - mock_mcp_tools = [ - MCPTool.model_validate({ - "name": "search_repo", - "description": "Search BerriAI/litellm repository for information", - "inputSchema": { - "type": "object", - "properties": { - "query": {"type": "string", "description": "Search query"} - }, - "required": ["query"], - }, - }, by_name=False) - ] - - # Only mock the MCP-specific operations, let LLM responses be real - with caplog.at_level(logging.ERROR): - with ( - patch.object( - LiteLLM_Proxy_MCP_Handler, - "_get_mcp_tools_from_manager", - new_callable=AsyncMock, - ) as mock_get_tools, - patch.object( - LiteLLM_Proxy_MCP_Handler, - "_execute_tool_calls", - new_callable=AsyncMock, - ) as mock_execute_tools, - ): - # Setup MCP mocks only - mock_get_tools.return_value = (mock_mcp_tools, ["litellm_proxy"]) - - # Create a dynamic mock that will match the actual tool call ID from the LLM response - def mock_execute_tool_calls_side_effect( - tool_calls, user_api_key_auth, **kwargs - ): - """Mock function that returns results matching the actual tool call IDs from the LLM""" - results = [] - for tool_call in tool_calls: - # Extract call_id from the tool call - call_id = None - if isinstance(tool_call, dict): - call_id = tool_call.get("call_id") or tool_call.get("id") - elif hasattr(tool_call, "call_id"): - call_id = tool_call.call_id - elif hasattr(tool_call, "id"): - call_id = tool_call.id - - if call_id: - results.append( - { - "tool_call_id": call_id, - "result": "LiteLLM is a unified interface for 100+ LLMs that translates inputs to provider-specific completion endpoints and provides consistent OpenAI-format output.", - } - ) - return results - - mock_execute_tools.side_effect = mock_execute_tool_calls_side_effect - - # Make the actual call - LLM responses will be real - mcp_tool_config = cast( - Any, - { - "type": "mcp", - "server_url": "litellm_proxy", - "require_approval": "never", - }, - ) - response = await litellm.aresponses( - model=model, - tools=[mcp_tool_config], - tool_choice="required", - input=[ - { - "role": "user", - "type": "message", - "content": "give me a TLDR of what BerriAI/litellm is about", - } - ], - stream=True, - ) - - print(f"📋 Response type: {type(response)}") - assert hasattr( - response, "__aiter__" - ), "Response should be an async streaming response" - - # Collect streaming chunks - chunks = [] - async for chunk in response: - chunks.append(chunk) - print(f"📦 Chunk type: {getattr(chunk, 'type', 'unknown')}") - - print(f"📊 Total chunks received: {len(chunks)}") - - # Verify MCP mocks were called (may be called multiple times in streaming) - assert ( - mock_get_tools.call_count >= 1 - ), f"Expected MCP tools to be fetched at least once, got {mock_get_tools.call_count}" - print(f"MCP tools fetched: {len(mock_mcp_tools)}") - - # Verify we got a response - assert response is not None - assert len(chunks) > 0, "Should have received streaming chunks" - - print("Basic streaming responses API with MCP tools test passed!") - - lite_errors = [ - record - for record in caplog.records - if record.levelno >= logging.ERROR - and ("LiteLLM" in record.name or "LiteLLM" in record.getMessage()) - ] - assert not lite_errors, "Unexpected LiteLLM errors: " + ", ".join( - record.getMessage() for record in lite_errors - ) - - @pytest.mark.asyncio async def test_mcp_parameter_preparation_helpers(): """ @@ -1215,7 +1058,7 @@ async def test_no_duplicate_mcp_tools_in_streaming_e2e(): The test mocks the MCP manager response but validates the actual tools sent to the LLM to ensure no duplication occurs. """ - from unittest.mock import AsyncMock, patch, call + from unittest.mock import AsyncMock, patch from litellm.responses.mcp.litellm_proxy_mcp_handler import ( LiteLLM_Proxy_MCP_Handler, ) @@ -1432,221 +1275,3 @@ async def test_no_duplicate_mcp_tools_in_streaming_e2e(): "tools_per_call": [len(tools) for tools in llm_call_tools], "duplicate_tools_found": False, } - - -@pytest.mark.asyncio -@pytest.mark.parametrize("model", ["gpt-4o-mini"]) -async def test_streaming_mcp_event_order_and_response_id_consistency( - model: str, caplog: pytest.LogCaptureFixture -): - """ - Test that: - 1. Streaming events are emitted in correct order (response.created, response.in_progress, response.output_item.added before MCP events) - 2. All response lifecycle events share the same response ID within a cycle - """ - if ("gpt" in model.lower() or "openai" in model.lower()) and not os.getenv( - "OPENAI_API_KEY" - ): - pytest.skip("OPENAI_API_KEY not set, skipping openai model test") - - from unittest.mock import AsyncMock, patch - - mock_mcp_tools = [ - MCPTool.model_validate({ - "name": "get_weather", - "description": "Get weather for a city", - "inputSchema": { - "type": "object", - "properties": { - "city": {"type": "string", "description": "City name"} - }, - "required": ["city"], - }, - }, by_name=False) - ] - - with caplog.at_level(logging.ERROR): - with ( - patch.object( - LiteLLM_Proxy_MCP_Handler, - "_get_mcp_tools_from_manager", - new_callable=AsyncMock, - ) as mock_get_tools, - patch.object( - LiteLLM_Proxy_MCP_Handler, - "_execute_tool_calls", - new_callable=AsyncMock, - ) as mock_execute_tools, - ): - mock_get_tools.return_value = (mock_mcp_tools, ["litellm_proxy"]) - - def mock_execute_side_effect(tool_calls, user_api_key_auth, **kwargs): - results = [] - for tool_call in tool_calls: - call_id = None - if isinstance(tool_call, dict): - call_id = tool_call.get("call_id") or tool_call.get("id") - elif hasattr(tool_call, "call_id"): - call_id = tool_call.call_id - elif hasattr(tool_call, "id"): - call_id = tool_call.id - if call_id: - results.append( - { - "tool_call_id": call_id, - "result": "Sunny, 72°F", - } - ) - return results - - mock_execute_tools.side_effect = mock_execute_side_effect - - mcp_tool_config = cast( - Any, - { - "type": "mcp", - "server_url": "litellm_proxy", - "require_approval": "never", - }, - ) - - response = await litellm.aresponses( - model=model, - tools=[mcp_tool_config], - input=[ - { - "role": "user", - "type": "message", - "content": "What's the weather in San Francisco?", - } - ], - stream=True, - ) - - events = [] - async for chunk in response: - events.append(chunk) - - assert len(events) > 0, "Should receive streaming events" - - created_idx = next( - ( - i - for i, e in enumerate(events) - if getattr(e, "type", None) == "response.created" - ), - None, - ) - in_progress_idx = next( - ( - i - for i, e in enumerate(events) - if getattr(e, "type", None) == "response.in_progress" - ), - None, - ) - output_item_added_idx = next( - ( - i - for i, e in enumerate(events) - if getattr(e, "type", None) == "response.output_item.added" - ), - None, - ) - mcp_in_progress_idx = next( - ( - i - for i, e in enumerate(events) - if "mcp_list_tools.in_progress" in str(getattr(e, "type", "")) - ), - None, - ) - completed_idx = next( - ( - i - for i, e in enumerate(events) - if getattr(e, "type", None) == "response.completed" - ), - None, - ) - - assert created_idx is not None, "response.created event should be present" - assert ( - in_progress_idx is not None - ), "response.in_progress event should be present" - assert ( - output_item_added_idx is not None - ), "response.output_item.added event should be present" - - assert ( - created_idx < in_progress_idx - ), "response.created should come before response.in_progress" - assert ( - in_progress_idx < output_item_added_idx - ), "response.in_progress should come before response.output_item.added" - - if mcp_in_progress_idx is not None: - assert ( - output_item_added_idx < mcp_in_progress_idx - ), "response.output_item.added should come before response.mcp_list_tools.in_progress" - - response_ids = [] - for i, event in enumerate(events): - event_type = getattr(event, "type", None) - if hasattr(event, "response"): - response_obj = getattr(event, "response", None) - if response_obj and hasattr(response_obj, "id"): - event_type_value = ( - event_type.value - if hasattr(event_type, "value") - else str(event_type) - ) - if any( - x in event_type_value - for x in [ - "response.created", - "response.in_progress", - "response.completed", - ] - ): - response_ids.append((i, event_type_value, response_obj.id)) - - assert ( - len(response_ids) >= 2 - ), f"Should have at least 2 response lifecycle events. Found {len(response_ids)}" - - cycles = [] - current_cycle = [] - current_id = None - - for idx, event_type, resp_id in response_ids: - if current_id is None or resp_id == current_id: - current_cycle.append((idx, event_type, resp_id)) - current_id = resp_id - else: - if current_cycle: - cycles.append(current_cycle) - current_cycle = [(idx, event_type, resp_id)] - current_id = resp_id - if current_cycle: - cycles.append(current_cycle) - - for cycle_num, cycle in enumerate(cycles): - cycle_ids = set(resp_id for _, _, resp_id in cycle) - assert ( - len(cycle_ids) == 1 - ), f"Cycle {cycle_num + 1} should have consistent response ID. Found {len(cycle_ids)} unique IDs" - - assert ( - completed_idx is not None - ), "response.completed event should be present" - - lite_errors = [ - record - for record in caplog.records - if record.levelno >= logging.ERROR - and ("LiteLLM" in record.name or "LiteLLM" in record.getMessage()) - ] - assert not lite_errors, "Unexpected LiteLLM errors: " + ", ".join( - record.getMessage() for record in lite_errors - ) diff --git a/tests/mcp_tests/test_aresponses_api_with_mcp_providers.py b/tests/mcp_tests/test_aresponses_api_with_mcp_providers.py new file mode 100644 index 00000000000..72a0415cf88 --- /dev/null +++ b/tests/mcp_tests/test_aresponses_api_with_mcp_providers.py @@ -0,0 +1,389 @@ +import logging +import os +import pytest +from mcp.types import Tool as MCPTool +from typing import Any, cast + +import litellm +from litellm.responses.mcp.litellm_proxy_mcp_handler import LiteLLM_Proxy_MCP_Handler + + +class MockUserAPIKeyAuth: + """Mock UserAPIKeyAuth for testing""" + + def __init__(self): + self.api_key = "test_key" + self.user_id = "test_user" + self.team_id = "test_team" + self.user_email = "test@example.com" + self.max_budget = 100.0 + self.spend = 0.0 + self.models = [] + self.aliases = {} + self.config = {} + self.permissions = {} + self.metadata = {} + self.object_permission_id = "test_permission_id" + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "model", + [ + pytest.param("gpt-4o-mini", id="openai"), + pytest.param("claude-haiku-4-5", id="anthropic"), + ], +) +async def test_streaming_responses_api_with_mcp_tools( + model: str, caplog: pytest.LogCaptureFixture +): + """ + Test the streaming responses API with MCP tools when using server_url="litellm_proxy" + + Under the hood the follow occurs + + - MCP: responses called litellm MCP manager.list_tools (MOCKED) + - Request 1: Made to model under test with fetched tools (REAL LLM CALL) + - MCP: Execute tool call from request 1 and returns result (MOCKED) + - Request 2: Made to model under test with fetched tools and tool results (REAL LLM CALL) + + Return the user the result of request 2 + """ + if ("claude" in model.lower() or "anthropic" in model.lower()) and not os.getenv( + "ANTHROPIC_API_KEY" + ): + pytest.skip("ANTHROPIC_API_KEY not set, skipping anthropic model test") + if ("gpt" in model.lower() or "openai" in model.lower()) and not os.getenv( + "OPENAI_API_KEY" + ): + pytest.skip("OPENAI_API_KEY not set, skipping openai model test") + + from unittest.mock import AsyncMock, patch + + print("🧪 Testing basic streaming with MCP tools...") + + mock_mcp_tools = [ + MCPTool.model_validate({ + "name": "search_repo", + "description": "Search BerriAI/litellm repository for information", + "inputSchema": { + "type": "object", + "properties": { + "query": {"type": "string", "description": "Search query"} + }, + "required": ["query"], + }, + }, by_name=False) + ] + + with caplog.at_level(logging.ERROR): + with ( + patch.object( + LiteLLM_Proxy_MCP_Handler, + "_get_mcp_tools_from_manager", + new_callable=AsyncMock, + ) as mock_get_tools, + patch.object( + LiteLLM_Proxy_MCP_Handler, + "_execute_tool_calls", + new_callable=AsyncMock, + ) as mock_execute_tools, + ): + mock_get_tools.return_value = (mock_mcp_tools, ["litellm_proxy"]) + + def mock_execute_tool_calls_side_effect( + tool_calls, user_api_key_auth, **kwargs + ): + """Mock function that returns results matching the actual tool call IDs from the LLM""" + results = [] + for tool_call in tool_calls: + call_id = None + if isinstance(tool_call, dict): + call_id = tool_call.get("call_id") or tool_call.get("id") + elif hasattr(tool_call, "call_id"): + call_id = tool_call.call_id + elif hasattr(tool_call, "id"): + call_id = tool_call.id + + if call_id: + results.append( + { + "tool_call_id": call_id, + "result": "LiteLLM is a unified interface for 100+ LLMs that translates inputs to provider-specific completion endpoints and provides consistent OpenAI-format output.", + } + ) + return results + + mock_execute_tools.side_effect = mock_execute_tool_calls_side_effect + + mcp_tool_config = cast( + Any, + { + "type": "mcp", + "server_url": "litellm_proxy", + "require_approval": "never", + }, + ) + response = await litellm.aresponses( + model=model, + tools=[mcp_tool_config], + tool_choice="required", + input=[ + { + "role": "user", + "type": "message", + "content": "give me a TLDR of what BerriAI/litellm is about", + } + ], + stream=True, + ) + + print(f"📋 Response type: {type(response)}") + assert hasattr( + response, "__aiter__" + ), "Response should be an async streaming response" + + chunks = [] + async for chunk in response: + chunks.append(chunk) + print(f"📦 Chunk type: {getattr(chunk, 'type', 'unknown')}") + + print(f"📊 Total chunks received: {len(chunks)}") + + assert ( + mock_get_tools.call_count >= 1 + ), f"Expected MCP tools to be fetched at least once, got {mock_get_tools.call_count}" + print(f"MCP tools fetched: {len(mock_mcp_tools)}") + + assert response is not None + assert len(chunks) > 0, "Should have received streaming chunks" + + print("Basic streaming responses API with MCP tools test passed!") + + lite_errors = [ + record + for record in caplog.records + if record.levelno >= logging.ERROR + and ("LiteLLM" in record.name or "LiteLLM" in record.getMessage()) + ] + assert not lite_errors, "Unexpected LiteLLM errors: " + ", ".join( + record.getMessage() for record in lite_errors + ) + + + +@pytest.mark.asyncio +@pytest.mark.parametrize("model", ["gpt-4o-mini"]) +async def test_streaming_mcp_event_order_and_response_id_consistency( + model: str, caplog: pytest.LogCaptureFixture +): + """ + Test that: + 1. Streaming events are emitted in correct order (response.created, response.in_progress, response.output_item.added before MCP events) + 2. All response lifecycle events share the same response ID within a cycle + """ + if ("gpt" in model.lower() or "openai" in model.lower()) and not os.getenv( + "OPENAI_API_KEY" + ): + pytest.skip("OPENAI_API_KEY not set, skipping openai model test") + + from unittest.mock import AsyncMock, patch + + mock_mcp_tools = [ + MCPTool.model_validate({ + "name": "get_weather", + "description": "Get weather for a city", + "inputSchema": { + "type": "object", + "properties": { + "city": {"type": "string", "description": "City name"} + }, + "required": ["city"], + }, + }, by_name=False) + ] + + with caplog.at_level(logging.ERROR): + with ( + patch.object( + LiteLLM_Proxy_MCP_Handler, + "_get_mcp_tools_from_manager", + new_callable=AsyncMock, + ) as mock_get_tools, + patch.object( + LiteLLM_Proxy_MCP_Handler, + "_execute_tool_calls", + new_callable=AsyncMock, + ) as mock_execute_tools, + ): + mock_get_tools.return_value = (mock_mcp_tools, ["litellm_proxy"]) + + def mock_execute_side_effect(tool_calls, user_api_key_auth, **kwargs): + results = [] + for tool_call in tool_calls: + call_id = None + if isinstance(tool_call, dict): + call_id = tool_call.get("call_id") or tool_call.get("id") + elif hasattr(tool_call, "call_id"): + call_id = tool_call.call_id + elif hasattr(tool_call, "id"): + call_id = tool_call.id + if call_id: + results.append( + { + "tool_call_id": call_id, + "result": "Sunny, 72°F", + } + ) + return results + + mock_execute_tools.side_effect = mock_execute_side_effect + + mcp_tool_config = cast( + Any, + { + "type": "mcp", + "server_url": "litellm_proxy", + "require_approval": "never", + }, + ) + + response = await litellm.aresponses( + model=model, + tools=[mcp_tool_config], + input=[ + { + "role": "user", + "type": "message", + "content": "What's the weather in San Francisco?", + } + ], + stream=True, + ) + + events = [] + async for chunk in response: + events.append(chunk) + + assert len(events) > 0, "Should receive streaming events" + + created_idx = next( + ( + i + for i, e in enumerate(events) + if getattr(e, "type", None) == "response.created" + ), + None, + ) + in_progress_idx = next( + ( + i + for i, e in enumerate(events) + if getattr(e, "type", None) == "response.in_progress" + ), + None, + ) + output_item_added_idx = next( + ( + i + for i, e in enumerate(events) + if getattr(e, "type", None) == "response.output_item.added" + ), + None, + ) + mcp_in_progress_idx = next( + ( + i + for i, e in enumerate(events) + if "mcp_list_tools.in_progress" in str(getattr(e, "type", "")) + ), + None, + ) + completed_idx = next( + ( + i + for i, e in enumerate(events) + if getattr(e, "type", None) == "response.completed" + ), + None, + ) + + assert created_idx is not None, "response.created event should be present" + assert ( + in_progress_idx is not None + ), "response.in_progress event should be present" + assert ( + output_item_added_idx is not None + ), "response.output_item.added event should be present" + + assert ( + created_idx < in_progress_idx + ), "response.created should come before response.in_progress" + assert ( + in_progress_idx < output_item_added_idx + ), "response.in_progress should come before response.output_item.added" + + if mcp_in_progress_idx is not None: + assert ( + output_item_added_idx < mcp_in_progress_idx + ), "response.output_item.added should come before response.mcp_list_tools.in_progress" + + response_ids = [] + for i, event in enumerate(events): + event_type = getattr(event, "type", None) + if hasattr(event, "response"): + response_obj = getattr(event, "response", None) + if response_obj and hasattr(response_obj, "id"): + event_type_value = ( + event_type.value + if hasattr(event_type, "value") + else str(event_type) + ) + if any( + x in event_type_value + for x in [ + "response.created", + "response.in_progress", + "response.completed", + ] + ): + response_ids.append((i, event_type_value, response_obj.id)) + + assert ( + len(response_ids) >= 2 + ), f"Should have at least 2 response lifecycle events. Found {len(response_ids)}" + + cycles = [] + current_cycle = [] + current_id = None + + for idx, event_type, resp_id in response_ids: + if current_id is None or resp_id == current_id: + current_cycle.append((idx, event_type, resp_id)) + current_id = resp_id + else: + if current_cycle: + cycles.append(current_cycle) + current_cycle = [(idx, event_type, resp_id)] + current_id = resp_id + if current_cycle: + cycles.append(current_cycle) + + for cycle_num, cycle in enumerate(cycles): + cycle_ids = set(resp_id for _, _, resp_id in cycle) + assert ( + len(cycle_ids) == 1 + ), f"Cycle {cycle_num + 1} should have consistent response ID. Found {len(cycle_ids)} unique IDs" + + assert ( + completed_idx is not None + ), "response.completed event should be present" + + lite_errors = [ + record + for record in caplog.records + if record.levelno >= logging.ERROR + and ("LiteLLM" in record.name or "LiteLLM" in record.getMessage()) + ] + assert not lite_errors, "Unexpected LiteLLM errors: " + ", ".join( + record.getMessage() for record in lite_errors + ) diff --git a/tests/proxy_unit_tests/test_proxy_custom_auth.py b/tests/proxy_unit_tests/test_proxy_custom_auth.py index b575e4c85c6..dbbad0dab1e 100644 --- a/tests/proxy_unit_tests/test_proxy_custom_auth.py +++ b/tests/proxy_unit_tests/test_proxy_custom_auth.py @@ -53,8 +53,7 @@ def test_custom_auth(client): "max_tokens": 10, } # Your bearer token - token = os.getenv("PROXY_MASTER_KEY") - print(f"token: {token}") + token = "sk-unit-test-master" headers = {"Authorization": f"Bearer {token}"} with pytest.raises(Exception, match="Authentication Error, Failed custom auth") as exc_info: client.post("/chat/completions", json=test_data, headers=headers) @@ -71,7 +70,7 @@ def test_custom_auth_bearer(client): "max_tokens": 10, } # Your bearer token - token = os.getenv("PROXY_MASTER_KEY") + token = "sk-unit-test-master" headers = {"Authorization": f"WITHOUT BEAR Er {token}"} with pytest.raises(Exception, match="CustomAuth - Malformed API Key passed in") as exc_info: diff --git a/tests/proxy_unit_tests/test_proxy_pass_user_config.py b/tests/proxy_unit_tests/test_proxy_pass_user_config.py deleted file mode 100644 index 91911c142ea..00000000000 --- a/tests/proxy_unit_tests/test_proxy_pass_user_config.py +++ /dev/null @@ -1,114 +0,0 @@ -import sys, os -import traceback -from dotenv import load_dotenv - -load_dotenv() -import io - -# this file is to test litellm/proxy - -import pytest, logging, asyncio -import litellm -from litellm import embedding, completion, completion_cost, Timeout -from litellm import RateLimitError - -# Configure logging -logging.basicConfig( - level=logging.DEBUG, # Set the desired logging level - format="%(asctime)s - %(levelname)s - %(message)s", -) - -# test /chat/completion request to the proxy -from fastapi.testclient import TestClient -from fastapi import FastAPI -from litellm.proxy.proxy_server import ( - router, - save_worker_config, - initialize, -) # Replace with the actual module where your FastAPI router is defined - -# Your bearer token -token = "sk-1234" - -headers = {"Authorization": f"Bearer {token}"} - - -@pytest.fixture(scope="function") -def client_no_auth(): - # Assuming litellm.proxy.proxy_server is an object - from litellm.proxy.proxy_server import cleanup_router_config_variables - - cleanup_router_config_variables() - filepath = os.path.dirname(os.path.abspath(__file__)) - config_fp = f"{filepath}/test_configs/test_config_no_auth.yaml" - # initialize can get run in parallel, it sets specific variables for the fast api app, sinc eit gets run in parallel different tests use the wrong variables - asyncio.run(initialize(config=config_fp, debug=True)) - app = FastAPI() - app.include_router(router) # Include your router in the test app - - return TestClient(app) - - -@pytest.mark.skipif( - os.environ.get("AZURE_AI_API_KEY") is None - or os.environ.get("OPENAI_API_KEY") is None, - reason="AZURE_AI_API_KEY or OPENAI_API_KEY not set - skipping integration test", -) -def test_chat_completion(client_no_auth): - global headers - - from litellm.types.router import RouterConfig, ModelConfig - from litellm.types.completion import CompletionRequest - - user_config = RouterConfig( - model_list=[ - ModelConfig( - model_name="user-azure-instance", - litellm_params=CompletionRequest( - model="azure/gpt-4.1-mini", - api_key=os.getenv("AZURE_AI_API_KEY"), - api_version=os.getenv("AZURE_API_VERSION"), - api_base=os.getenv("AZURE_AI_API_BASE"), - timeout=10, - ), - tpm=240000, - rpm=1800, - ), - ModelConfig( - model_name="user-openai-instance", - litellm_params=CompletionRequest( - model="gpt-3.5-turbo", - api_key=os.getenv("OPENAI_API_KEY"), - timeout=10, - ), - tpm=240000, - rpm=1800, - ), - ], - num_retries=2, - allowed_fails=3, - fallbacks=[{"user-azure-instance": ["user-openai-instance"]}], - ).dict() - - try: - # Your test data - test_data = { - "model": "user-azure-instance", - "messages": [ - {"role": "user", "content": "hi"}, - ], - "max_tokens": 10, - "user_config": user_config, - } - - print("testing proxy server with chat completions") - response = client_no_auth.post("/v1/chat/completions", json=test_data) - print(f"response - {response.text}") - assert response.status_code == 200 - result = response.json() - print(f"Received response: {result}") - except Exception as e: - pytest.fail(f"LiteLLM Proxy test failed. Exception - {str(e)}") - - -# Run the test diff --git a/tests/proxy_unit_tests/test_proxy_server.py b/tests/proxy_unit_tests/test_proxy_server.py index ed0380058a5..5be27b3ad72 100644 --- a/tests/proxy_unit_tests/test_proxy_server.py +++ b/tests/proxy_unit_tests/test_proxy_server.py @@ -2065,59 +2065,6 @@ async def test_add_callback_via_key_litellm_pre_call_utils_langsmith( assert new_data["failure_callback"] == expected_failure_callbacks -@pytest.mark.skipif( - not os.getenv("GEMINI_API_KEY") and not os.getenv("GOOGLE_API_KEY"), - reason="Requires GEMINI_API_KEY or GOOGLE_API_KEY.", -) -@pytest.mark.asyncio -async def test_gemini_pass_through_endpoint(): - from starlette.datastructures import URL - - from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( - Request, - Response, - gemini_proxy_route, - ) - - body = b""" - { - "contents": [{ - "parts":[{ - "text": "The quick brown fox jumps over the lazy dog." - }] - }] - } - """ - - # Construct the scope dictionary - scope = { - "type": "http", - "method": "POST", - "path": "/gemini/v1beta/models/gemini-2.5-flash:countTokens", - "query_string": b"key=sk-1234", - "headers": [ - (b"content-type", b"application/json"), - ], - } - - # Create a new Request object - async def async_receive(): - return {"type": "http.request", "body": body, "more_body": False} - - request = Request( - scope=scope, - receive=async_receive, - ) - - resp = await gemini_proxy_route( - endpoint="v1beta/models/gemini-2.5-flash:countTokens?key=sk-1234", - request=request, - fastapi_response=Response(), - ) - - print(resp.body) - - @pytest.mark.parametrize("hidden", [True, False]) @pytest.mark.asyncio async def test_model_info_alias_without_prisma(hidden): diff --git a/tests/proxy_unit_tests/test_proxy_server_gemini_pass_through.py b/tests/proxy_unit_tests/test_proxy_server_gemini_pass_through.py new file mode 100644 index 00000000000..2453ec3bfe3 --- /dev/null +++ b/tests/proxy_unit_tests/test_proxy_server_gemini_pass_through.py @@ -0,0 +1,51 @@ +import os + +import pytest + + +@pytest.mark.skipif( + not os.getenv("GEMINI_API_KEY") and not os.getenv("GOOGLE_API_KEY"), + reason="Requires GEMINI_API_KEY or GOOGLE_API_KEY.", +) +@pytest.mark.asyncio +async def test_gemini_pass_through_endpoint(): + from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( + Request, + Response, + gemini_proxy_route, + ) + + body = b""" + { + "contents": [{ + "parts":[{ + "text": "The quick brown fox jumps over the lazy dog." + }] + }] + } + """ + + scope = { + "type": "http", + "method": "POST", + "path": "/gemini/v1beta/models/gemini-2.5-flash:countTokens", + "query_string": b"key=sk-1234", + "headers": [ + (b"content-type", b"application/json"), + ], + } + + async def async_receive(): + return {"type": "http.request", "body": body, "more_body": False} + + request = Request( + scope=scope, + receive=async_receive, + ) + + await gemini_proxy_route( + endpoint="v1beta/models/gemini-2.5-flash:countTokens?key=sk-1234", + request=request, + fastapi_response=Response(), + ) + diff --git a/tests/proxy_unit_tests/test_proxy_token_counter.py b/tests/proxy_unit_tests/test_proxy_token_counter.py index 39ec4bb1887..8590e959961 100644 --- a/tests/proxy_unit_tests/test_proxy_token_counter.py +++ b/tests/proxy_unit_tests/test_proxy_token_counter.py @@ -2,10 +2,7 @@ # 1. Generate a Key, and use it to make a call -import json import logging -import os -import tempfile from unittest.mock import AsyncMock, MagicMock, patch import httpx @@ -35,79 +32,6 @@ from litellm.types.utils import TokenCountResponse verbose_proxy_logger.setLevel(level=logging.DEBUG) -def get_vertex_ai_creds_json() -> dict: - # Define the path to the vertex_key.json file - print("loading vertex ai credentials") - filepath = os.path.dirname(os.path.abspath(__file__)) - vertex_key_path = filepath + "/vertex_key.json" - # Read the existing content of the file or create an empty dictionary - try: - with open(vertex_key_path, "r") as file: - # Read the file content - print("Read vertexai file path") - content = file.read() - - # If the file is empty or not valid JSON, create an empty dictionary - if not content or not content.strip(): - service_account_key_data = {} - else: - # Attempt to load the existing JSON content - file.seek(0) - service_account_key_data = json.load(file) - except FileNotFoundError: - # If the file doesn't exist, create an empty dictionary - service_account_key_data = {} - - # Update the service_account_key_data with environment variables - private_key_id = os.environ.get("VERTEX_AI_PRIVATE_KEY_ID", "") - private_key = os.environ.get("VERTEX_AI_PRIVATE_KEY", "") - private_key = private_key.replace("\\n", "\n") - service_account_key_data["private_key_id"] = private_key_id - service_account_key_data["private_key"] = private_key - - return service_account_key_data - - -def load_vertex_ai_credentials(): - # Define the path to the vertex_key.json file - print("loading vertex ai credentials") - filepath = os.path.dirname(os.path.abspath(__file__)) - vertex_key_path = filepath + "/vertex_key.json" - - # Read the existing content of the file or create an empty dictionary - try: - with open(vertex_key_path, "r") as file: - # Read the file content - print("Read vertexai file path") - content = file.read() - - # If the file is empty or not valid JSON, create an empty dictionary - if not content or not content.strip(): - service_account_key_data = {} - else: - # Attempt to load the existing JSON content - file.seek(0) - service_account_key_data = json.load(file) - except FileNotFoundError: - # If the file doesn't exist, create an empty dictionary - service_account_key_data = {} - - # Update the service_account_key_data with environment variables - private_key_id = os.environ.get("VERTEX_AI_PRIVATE_KEY_ID", "") - private_key = os.environ.get("VERTEX_AI_PRIVATE_KEY", "") - private_key = private_key.replace("\\n", "\n") - service_account_key_data["private_key_id"] = private_key_id - service_account_key_data["private_key"] = private_key - - # Create a temporary file - with tempfile.NamedTemporaryFile(mode="w+", delete=False) as temp_file: - # Write the updated content to the temporary files - json.dump(service_account_key_data, temp_file, indent=2) - - # Export the temporary file as GOOGLE_APPLICATION_CREDENTIALS - os.environ["GOOGLE_APPLICATION_CREDENTIALS"] = os.path.abspath(temp_file.name) - - @pytest.mark.asyncio async def test_vLLM_token_counting(): """ @@ -223,10 +147,12 @@ async def test_anthropic_messages_count_tokens_endpoint(): - Should return response in Anthropic format: {"input_tokens": } - Should work as wrapper around internal token_counter function """ - from litellm.proxy.anthropic_endpoints.endpoints import count_tokens - from fastapi import Request from unittest.mock import MagicMock + from fastapi import Request + + from litellm.proxy.anthropic_endpoints.endpoints import count_tokens + # Mock request object mock_request = MagicMock(spec=Request) mock_request_data = { @@ -295,10 +221,12 @@ async def test_anthropic_messages_count_tokens_with_non_anthropic_model(): - Should still work and return Anthropic format - Should call internal token_counter with from_anthropic_endpoint=True """ - from litellm.proxy.anthropic_endpoints.endpoints import count_tokens - from fastapi import Request from unittest.mock import MagicMock + from fastapi import Request + + from litellm.proxy.anthropic_endpoints.endpoints import count_tokens + # Mock request object mock_request = MagicMock(spec=Request) mock_request_data = { @@ -435,10 +363,12 @@ async def test_anthropic_endpoint_error_handling(): """ Test error handling in the /v1/messages/count_tokens endpoint """ - from litellm.proxy.anthropic_endpoints.endpoints import count_tokens - from fastapi import Request, HTTPException from unittest.mock import MagicMock + from fastapi import HTTPException, Request + + from litellm.proxy.anthropic_endpoints.endpoints import count_tokens + # Mock request object mock_request = MagicMock(spec=Request) mock_user_api_key_dict = MagicMock() @@ -474,8 +404,10 @@ async def test_anthropic_endpoint_error_handling(): @pytest.mark.asyncio async def test_factory_anthropic_endpoint_calls_anthropic_counter(): """Test that /v1/messages/count_tokens with Anthropic model uses Anthropic counter.""" - from unittest.mock import patch, AsyncMock, MagicMock + from unittest.mock import AsyncMock, MagicMock, patch + from fastapi.testclient import TestClient + from litellm.proxy.proxy_server import app # Mock the global handler instance in token_counter module @@ -531,8 +463,10 @@ async def test_factory_anthropic_endpoint_calls_anthropic_counter(): @pytest.mark.asyncio async def test_factory_gpt4_endpoint_does_not_call_anthropic_counter(): """Test that /v1/messages/count_tokens with GPT-4 does NOT use Anthropic counter.""" - from unittest.mock import patch, AsyncMock, MagicMock + from unittest.mock import AsyncMock, MagicMock, patch + from fastapi.testclient import TestClient + from litellm.proxy.proxy_server import app # Mock the global handler instance in token_counter module @@ -590,8 +524,10 @@ async def test_factory_gpt4_endpoint_does_not_call_anthropic_counter(): @pytest.mark.asyncio async def test_factory_normal_token_counter_endpoint_does_not_call_anthropic(): """Test that /utils/token_counter does NOT use Anthropic counter even with Anthropic model.""" - from unittest.mock import patch, AsyncMock, MagicMock + from unittest.mock import AsyncMock, MagicMock, patch + from fastapi.testclient import TestClient + from litellm.proxy.proxy_server import app # Mock the global handler instance in token_counter module @@ -678,57 +614,6 @@ async def test_factory_registration(): assert not counter.should_use_token_counting_api(custom_llm_provider=None) -@pytest.mark.skip( - reason="Requires Google/Vertex AI credentials (GEMINI_API_KEY or VERTEX_AI_PRIVATE_KEY)." -) -@pytest.mark.asyncio -@pytest.mark.parametrize("model_name", ["gemini-2.5-pro", "vertex-ai-gemini-2.5-pro"]) -async def test_vertex_ai_gemini_token_counting_with_contents(model_name): - """ - Test token counting for Vertex AI Gemini model using contents format with call_endpoint=True - """ - load_vertex_ai_credentials() - llm_router = Router( - model_list=[ - { - "model_name": "gemini-2.5-pro", - "litellm_params": { - "model": "gemini/gemini-2.5-pro", - }, - }, - { - "model_name": "vertex-ai-gemini-2.5-pro", - "litellm_params": { - "model": "vertex_ai/gemini-2.5-pro", - }, - }, - ] - ) - - setattr(litellm.proxy.proxy_server, "llm_router", llm_router) - - # Test with contents format and call_endpoint=True - response = await token_counter( - request=TokenCountRequest( - model=model_name, - contents=[ - {"parts": [{"text": "Hello world, how are you doing today? i am ij"}]} - ], - ), - call_endpoint=True, - ) - - print("Vertex AI Gemini token counting response:", response) - - # validate we have original response - assert response.original_response is not None - assert response.original_response.get("totalTokens") is not None - assert response.original_response.get("promptTokensDetails") is not None - - prompt_tokens_details = response.original_response.get("promptTokensDetails") - assert prompt_tokens_details is not None - - @pytest.mark.asyncio async def test_bedrock_count_tokens_endpoint(): """ @@ -779,7 +664,7 @@ async def test_vertex_ai_anthropic_token_counting(): This tests the token counting implementation for Vertex AI partner models without making actual API calls. Mocks at the handler level to test the full flow. """ - from unittest.mock import AsyncMock, patch, MagicMock + from unittest.mock import patch # Mock the Vertex AI partner models token counter response mock_token_response = { diff --git a/tests/proxy_unit_tests/test_proxy_utils.py b/tests/proxy_unit_tests/test_proxy_utils.py index 1134f41a940..ab787bbbe27 100644 --- a/tests/proxy_unit_tests/test_proxy_utils.py +++ b/tests/proxy_unit_tests/test_proxy_utils.py @@ -438,7 +438,7 @@ def test_is_request_body_safe_global_enabled( "model_name": "gpt-3.5-turbo", "litellm_params": { "model": "gpt-3.5-turbo", - "api_key": os.getenv("OPENAI_API_KEY"), + "api_key": "sk-openai-unit-test", }, } ] @@ -475,7 +475,7 @@ def test_is_request_body_safe_model_enabled( "model_name": "fireworks_ai/*", "litellm_params": { "model": "fireworks_ai/*", - "api_key": os.getenv("FIREWORKS_API_KEY"), + "api_key": "sk-fireworks-unit-test", "configurable_clientside_auth_params": ( ["api_base"] if allow_client_side_credentials else [] ), From e3f087315de3eac8ec5c78b31fca218b1f846892 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 24 Sep 2026 17:06:30 -0500 Subject: [PATCH 088/120] feat(terraform): add display_name to litellm_model resource and model data sources (#42987) * feat(terraform): add display_name to litellm_model resource and model data sources Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(terraform): persist display_name on update and read /model/info data envelope Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * chore(terraform): drop PATCH /model/{model_id}/update from endpoint audit allowlist Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(terraform): surface external display_name removal as drift on refresh Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * ci(terraform): rerun after uv download timeout Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: yassin Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- terraform/provider/CHANGELOG.md | 2 + terraform/provider/docs/data-sources/model.md | 1 + .../provider/docs/data-sources/models.md | 1 + terraform/provider/docs/resources/model.md | 2 + .../provider/litellm/data_source_model.go | 24 +- .../litellm/data_source_model_test.go | 12 +- terraform/provider/litellm/resource_model.go | 5 + .../provider/litellm/resource_model_crud.go | 35 ++- .../provider/litellm/resource_model_test.go | 244 ++++++++++++++++++ terraform/provider/litellm/types.go | 23 +- terraform/provider/litellm/utils.go | 7 + .../endpointaudit/coverage_allowlist.txt | 1 - 12 files changed, 333 insertions(+), 24 deletions(-) create mode 100644 terraform/provider/litellm/resource_model_test.go diff --git a/terraform/provider/CHANGELOG.md b/terraform/provider/CHANGELOG.md index 079bb7d8667..4123b7e6b51 100644 --- a/terraform/provider/CHANGELOG.md +++ b/terraform/provider/CHANGELOG.md @@ -16,6 +16,7 @@ longer signal it. ### Added +- **model**: Optional `display_name` argument on `litellm_model`, sent as `model_info.display_name` and returned as `display_name` by `/v1/models`, so client model pickers show a readable name; changes are persisted through `/model/{id}/update` since `/model/update` ignores `model_info`; also exported by the `litellm_model` and `litellm_models` data sources - **key**: Computed `server_metadata` attribute on `litellm_key` exposing every metadata entry the proxy stores, so metadata created outside Terraform is visible in state and drift on it shows on refresh, while `metadata` keeps tracking only the declared entries and updates keep preserving undeclared ones - **team_member_add**: `tpm_limit`, `rpm_limit`, `budget_duration`, and `allowed_models` attributes on `litellm_team_member_add`, applied to every member of the resource; `budget_duration` and `allowed_models` ride on `/team/member_add`, while the limits are sent through `/team/member_update`, which is where the proxy accepts them - **team**: Optional `team_id` argument on `litellm_team`, so teams can be created with a stable, human-readable ID instead of a provider-generated UUID; changing it forces replacement @@ -48,6 +49,7 @@ longer signal it. ### Fixed +- **model**: `litellm_model` refresh now reads the `{"data": [...]}` envelope `/model/info` returns, so `model_info` fields changed outside Terraform show up as drift instead of silently keeping the previous state - **key**: An update that changes `team_id` and fails because the key was already cascade-deleted along with its previous team now recovers by recreating the key under the new team, instead of aborting the apply. The key's absence is confirmed against the proxy first, so an unrelated failure still errors out, and a `team_id` change between two teams that both still exist stays a plain in-place update - **credential**: create now reports a `credential_name` collision as a clear error naming the `terraform import` command that adopts the existing credential, instead of surfacing the proxy's raw 500 with a Prisma `Unique constraint failed` message. New `adopt_existing` argument (default `false`) opts into taking the existing credential over during create, which makes `apply` idempotent again once state loses track of a credential that still exists on the proxy. Requires a proxy that answers 409 on the collision; older proxies are still detected by their 500 message - **credential**: credential names and `model_id` are now percent-encoded in request URLs, so a name containing `/`, `?`, `#` or spaces reaches the proxy intact instead of being cut at the first reserved character and read, updated or deleted as a different credential diff --git a/terraform/provider/docs/data-sources/model.md b/terraform/provider/docs/data-sources/model.md index 6976ff1523a..587652bcbd8 100644 --- a/terraform/provider/docs/data-sources/model.md +++ b/terraform/provider/docs/data-sources/model.md @@ -43,6 +43,7 @@ In addition to all arguments above, the following attributes are exported: * `tier` - Model tier (`free` or `paid`). * `mode` - Model mode, e.g. `chat` or `embedding`. * `team_id` - Team the deployment is scoped to, if any. +* `display_name` - Human-readable name returned by `/v1/models`, if configured. * `db_model` - Whether the deployment is stored in the database (as opposed to config). ## Security Note diff --git a/terraform/provider/docs/data-sources/models.md b/terraform/provider/docs/data-sources/models.md index 7862dc30ab7..1cf0fd36ab3 100644 --- a/terraform/provider/docs/data-sources/models.md +++ b/terraform/provider/docs/data-sources/models.md @@ -41,4 +41,5 @@ In addition to all arguments above, the following attributes are exported: * `tier` - Model tier (`free` or `paid`). * `mode` - Model mode, e.g. `chat` or `embedding`. * `team_id` - Team the deployment is scoped to, if any. + * `display_name` - Human-readable name returned by `/v1/models`, if configured. * `db_model` - Whether the deployment is stored in the database. diff --git a/terraform/provider/docs/resources/model.md b/terraform/provider/docs/resources/model.md index 0409b48b391..5bb68bfe918 100644 --- a/terraform/provider/docs/resources/model.md +++ b/terraform/provider/docs/resources/model.md @@ -126,6 +126,8 @@ The following arguments are supported: * `team_id` - (Optional) string. Associate the model with a specific team. +* `display_name` - (Optional) string. Human-readable name stored in `model_info.display_name` and returned as `display_name` by `/v1/models`, so clients such as Claude Code and Claude Desktop show it in their model picker instead of `model_name`. When unset, clients fall back to `model_name`. + * `mode` - (Optional) string. The intended use of the model. Valid values are: * `completion` * `embedding` diff --git a/terraform/provider/litellm/data_source_model.go b/terraform/provider/litellm/data_source_model.go index 78af04ac160..6b9993d267e 100644 --- a/terraform/provider/litellm/data_source_model.go +++ b/terraform/provider/litellm/data_source_model.go @@ -23,14 +23,15 @@ type modelInfoParams struct { } type modelInfoMeta struct { - ID string `json:"id"` - DBModel bool `json:"db_model"` - BaseModel string `json:"base_model"` - Tier string `json:"tier"` - Mode string `json:"mode"` - TeamID string `json:"team_id"` - CreatedAt string `json:"created_at"` - UpdatedAt string `json:"updated_at"` + ID string `json:"id"` + DBModel bool `json:"db_model"` + BaseModel string `json:"base_model"` + Tier string `json:"tier"` + Mode string `json:"mode"` + TeamID string `json:"team_id"` + DisplayName string `json:"display_name"` + CreatedAt string `json:"created_at"` + UpdatedAt string `json:"updated_at"` } type modelInfoEntry struct { @@ -112,6 +113,10 @@ func dataSourceLiteLLMModel() *schema.Resource { Type: schema.TypeString, Computed: true, }, + "display_name": { + Type: schema.TypeString, + Computed: true, + }, "db_model": { Type: schema.TypeBool, Computed: true, @@ -161,6 +166,7 @@ func dataSourceLiteLLMModelRead(d *schema.ResourceData, m interface{}) error { d.Set("tier", entry.ModelInfo.Tier) d.Set("mode", entry.ModelInfo.Mode) d.Set("team_id", entry.ModelInfo.TeamID) + d.Set("display_name", entry.ModelInfo.DisplayName) d.Set("db_model", entry.ModelInfo.DBModel) log.Printf("[INFO] Successfully read model with ID: %s", modelID) @@ -197,6 +203,7 @@ func dataSourceLiteLLMModels() *schema.Resource { "tier": {Type: schema.TypeString, Computed: true}, "mode": {Type: schema.TypeString, Computed: true}, "team_id": {Type: schema.TypeString, Computed: true}, + "display_name": {Type: schema.TypeString, Computed: true}, "db_model": {Type: schema.TypeBool, Computed: true}, }, }, @@ -247,6 +254,7 @@ func dataSourceLiteLLMModelsRead(d *schema.ResourceData, m interface{}) error { "tier": entry.ModelInfo.Tier, "mode": entry.ModelInfo.Mode, "team_id": entry.ModelInfo.TeamID, + "display_name": entry.ModelInfo.DisplayName, "db_model": entry.ModelInfo.DBModel, }) } diff --git a/terraform/provider/litellm/data_source_model_test.go b/terraform/provider/litellm/data_source_model_test.go index 97d7f07dcd8..b46a355853a 100644 --- a/terraform/provider/litellm/data_source_model_test.go +++ b/terraform/provider/litellm/data_source_model_test.go @@ -35,7 +35,8 @@ func TestDataSourceModelReadSingleObject(t *testing.T) { "base_model": "gpt-4o", "tier": "paid", "mode": "chat", - "team_id": "team-1" + "team_id": "team-1", + "display_name": "GPT-4o" } } }`)) @@ -66,6 +67,7 @@ func TestDataSourceModelReadSingleObject(t *testing.T) { "tier": "paid", "mode": "chat", "team_id": "team-1", + "display_name": "GPT-4o", "db_model": true, } for attr, want := range checks { @@ -115,7 +117,7 @@ func TestDataSourceModelsRead(t *testing.T) { w.Header().Set("Content-Type", "application/json") w.Write([]byte(`{ "data": [ - {"model_name": "a", "litellm_params": {"model": "openai/a", "custom_llm_provider": "openai"}, "model_info": {"id": "id-1", "db_model": true}}, + {"model_name": "a", "litellm_params": {"model": "openai/a", "custom_llm_provider": "openai"}, "model_info": {"id": "id-1", "db_model": true, "display_name": "Model A"}}, {"model_name": "b", "litellm_params": {"model": "anthropic/b", "custom_llm_provider": "anthropic"}, "model_info": {"id": "id-2"}} ] }`)) @@ -143,7 +145,11 @@ func TestDataSourceModelsRead(t *testing.T) { t.Fatalf("expected 2 models, got %d", len(models)) } first := models[0].(map[string]interface{}) - if first["model_name"] != "a" || first["custom_llm_provider"] != "openai" || first["db_model"] != true { + if first["model_name"] != "a" || first["custom_llm_provider"] != "openai" || first["db_model"] != true || first["display_name"] != "Model A" { t.Errorf("unexpected first model: %v", first) } + second := models[1].(map[string]interface{}) + if second["display_name"] != "" { + t.Errorf("expected empty display_name for model without one, got %v", second["display_name"]) + } } diff --git a/terraform/provider/litellm/resource_model.go b/terraform/provider/litellm/resource_model.go index b0a7304718b..85cb1d038bf 100644 --- a/terraform/provider/litellm/resource_model.go +++ b/terraform/provider/litellm/resource_model.go @@ -93,6 +93,11 @@ func resourceLiteLLMModel() *schema.Resource { Type: schema.TypeString, Optional: true, }, + "display_name": { + Type: schema.TypeString, + Optional: true, + Description: "Human-readable name returned as display_name by /v1/models, shown in client model pickers instead of model_name", + }, "mode": { Type: schema.TypeString, Optional: true, diff --git a/terraform/provider/litellm/resource_model_crud.go b/terraform/provider/litellm/resource_model_crud.go index fc5d5b09dd5..c7db2a673e5 100644 --- a/terraform/provider/litellm/resource_model_crud.go +++ b/terraform/provider/litellm/resource_model_crud.go @@ -4,6 +4,7 @@ import ( "encoding/json" "fmt" "log" + "net/url" "strconv" "strings" "time" @@ -53,6 +54,7 @@ func retryModelRead(d *schema.ResourceData, m interface{}, maxRetries int) error const ( endpointModelNew = "/model/new" endpointModelUpdate = "/model/update" + endpointModelPatch = "/model/%s/update" endpointModelInfo = "/model/info" endpointModelDelete = "/model/delete" ) @@ -246,12 +248,13 @@ func createOrUpdateModel(d *schema.ResourceData, m interface{}, isUpdate bool) e ModelName: d.Get("model_name").(string), LiteLLMParams: litellmParams, ModelInfo: ModelInfo{ - ID: modelID, - DBModel: true, - BaseModel: pricingBaseModel, - Tier: d.Get("tier").(string), - Mode: d.Get("mode").(string), - TeamID: d.Get("team_id").(string), + ID: modelID, + DBModel: true, + BaseModel: pricingBaseModel, + Tier: d.Get("tier").(string), + Mode: d.Get("mode").(string), + TeamID: d.Get("team_id").(string), + DisplayName: d.Get("display_name").(string), }, Additional: make(map[string]interface{}), } @@ -275,6 +278,12 @@ func createOrUpdateModel(d *schema.ResourceData, m interface{}, isUpdate bool) e return fmt.Errorf("failed to %s model: %w", map[bool]string{true: "update", false: "create"}[isUpdate], err) } + if isUpdate && d.HasChange("display_name") { + if err := patchModelDisplayName(client, modelID, d.Get("display_name").(string)); err != nil { + return fmt.Errorf("failed to update model display_name: %w", err) + } + } + d.SetId(modelID) log.Printf("[INFO] Model created with ID %s. Starting retry mechanism to read the model...", modelID) @@ -282,6 +291,19 @@ func createOrUpdateModel(d *schema.ResourceData, m interface{}, isUpdate bool) e return retryModelRead(d, m, 5) } +// /model/update only merges litellm_params, so model_info changes go through the PATCH endpoint. +func patchModelDisplayName(client *Client, modelID, displayName string) error { + resp, err := MakeRequest(client, "PATCH", fmt.Sprintf(endpointModelPatch, url.PathEscape(modelID)), ModelInfoPatch{ + ModelInfo: ModelInfoPatchFields{ID: modelID, DisplayName: displayName}, + }) + if err != nil { + return err + } + defer resp.Body.Close() + _, err = handleAPIResponse(resp, nil, client) + return err +} + func resourceLiteLLMModelCreate(d *schema.ResourceData, m interface{}) error { return createOrUpdateModel(d, m, false) } @@ -327,6 +349,7 @@ func resourceLiteLLMModelRead(d *schema.ResourceData, m interface{}) error { d.Set("tier", GetStringValue(modelResp.ModelInfo.Tier, d.Get("tier").(string))) d.Set("mode", GetStringValue(modelResp.ModelInfo.Mode, d.Get("mode").(string))) d.Set("team_id", GetStringValue(modelResp.ModelInfo.TeamID, d.Get("team_id").(string))) + d.Set("display_name", modelResp.ModelInfo.DisplayName) // Preserve credential name from state since it might not be returned by API d.Set("litellm_credential_name", d.Get("litellm_credential_name").(string)) diff --git a/terraform/provider/litellm/resource_model_test.go b/terraform/provider/litellm/resource_model_test.go new file mode 100644 index 00000000000..0be3c39a3c5 --- /dev/null +++ b/terraform/provider/litellm/resource_model_test.go @@ -0,0 +1,244 @@ +package litellm + +import ( + "context" + "encoding/json" + "net/http" + "net/http/httptest" + "testing" + + "github.com/hashicorp/terraform-plugin-sdk/v2/helper/schema" + "github.com/hashicorp/terraform-plugin-sdk/v2/terraform" +) + +func modelInfoBody(displayName string) string { + modelInfo := map[string]interface{}{ + "id": "model-123", + "db_model": true, + "base_model": "claude-sonnet-4-5", + "tier": "free", + "mode": "chat", + } + if displayName != "" { + modelInfo["display_name"] = displayName + } + body, _ := json.Marshal(map[string]interface{}{ + "model_name": "sonnet-4-5-anthropic", + "litellm_params": map[string]interface{}{"model": "anthropic/claude-sonnet-4-5", "custom_llm_provider": "anthropic"}, + "model_info": modelInfo, + }) + return string(body) +} + +func modelInfoDataEnvelope(displayName string) string { + return `{"data": [` + modelInfoBody(displayName) + `]}` +} + +func TestResourceLiteLLMModelCreateSendsDisplayName(t *testing.T) { + var createPayload map[string]interface{} + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch r.URL.Path { + case "/model/new": + if err := json.NewDecoder(r.Body).Decode(&createPayload); err != nil { + t.Errorf("failed to decode create payload: %v", err) + } + w.Write([]byte(modelInfoBody("Claude Sonnet 4.5"))) + case "/model/info": + w.Write([]byte(modelInfoBody("Claude Sonnet 4.5"))) + default: + t.Errorf("unexpected request: %s %s", r.Method, r.URL.Path) + w.WriteHeader(http.StatusNotFound) + } + })) + defer srv.Close() + + d := schema.TestResourceDataRaw(t, resourceLiteLLMModel().Schema, map[string]interface{}{ + "model_name": "sonnet-4-5-anthropic", + "custom_llm_provider": "anthropic", + "base_model": "claude-sonnet-4-5", + "model_api_key": "sk-ant-test", + "mode": "chat", + "display_name": "Claude Sonnet 4.5", + }) + + if err := resourceLiteLLMModelCreate(d, NewClient(srv.URL, "test-key", true)); err != nil { + t.Fatalf("create failed: %v", err) + } + + modelInfo, ok := createPayload["model_info"].(map[string]interface{}) + if !ok { + t.Fatalf("expected model_info object in create payload, got %v", createPayload["model_info"]) + } + if modelInfo["display_name"] != "Claude Sonnet 4.5" { + t.Errorf("expected model_info.display_name 'Claude Sonnet 4.5', got %v", modelInfo["display_name"]) + } + if got := d.Get("display_name").(string); got != "Claude Sonnet 4.5" { + t.Errorf("expected state display_name 'Claude Sonnet 4.5', got %q", got) + } +} + +func TestResourceLiteLLMModelCreateOmitsUnsetDisplayName(t *testing.T) { + var createPayload map[string]interface{} + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch r.URL.Path { + case "/model/new": + if err := json.NewDecoder(r.Body).Decode(&createPayload); err != nil { + t.Errorf("failed to decode create payload: %v", err) + } + w.Write([]byte(modelInfoBody(""))) + case "/model/info": + w.Write([]byte(modelInfoBody(""))) + default: + t.Errorf("unexpected request: %s %s", r.Method, r.URL.Path) + w.WriteHeader(http.StatusNotFound) + } + })) + defer srv.Close() + + d := schema.TestResourceDataRaw(t, resourceLiteLLMModel().Schema, map[string]interface{}{ + "model_name": "sonnet-4-5-anthropic", + "custom_llm_provider": "anthropic", + "base_model": "claude-sonnet-4-5", + "model_api_key": "sk-ant-test", + }) + + if err := resourceLiteLLMModelCreate(d, NewClient(srv.URL, "test-key", true)); err != nil { + t.Fatalf("create failed: %v", err) + } + + modelInfo := createPayload["model_info"].(map[string]interface{}) + if _, present := modelInfo["display_name"]; present { + t.Errorf("expected display_name to be omitted from model_info when unset, got %v", modelInfo["display_name"]) + } + if got := d.Get("display_name").(string); got != "" { + t.Errorf("expected empty state display_name, got %q", got) + } +} + +func TestResourceLiteLLMModelReadDisplayName(t *testing.T) { + cases := map[string]struct { + serverBody string + want string + }{ + "server value wins inside data envelope": {serverBody: modelInfoDataEnvelope("Renamed In Admin UI"), want: "Renamed In Admin UI"}, + "server value wins unwrapped": {serverBody: modelInfoBody("Renamed In Admin UI"), want: "Renamed In Admin UI"}, + "external removal clears state": {serverBody: modelInfoDataEnvelope(""), want: ""}, + } + for name, tc := range cases { + t.Run(name, func(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path != "/model/info" { + t.Errorf("unexpected request: %s %s", r.Method, r.URL.Path) + } + w.Write([]byte(tc.serverBody)) + })) + defer srv.Close() + + d := schema.TestResourceDataRaw(t, resourceLiteLLMModel().Schema, map[string]interface{}{ + "model_name": "sonnet-4-5-anthropic", + "custom_llm_provider": "anthropic", + "base_model": "claude-sonnet-4-5", + "display_name": "Claude Sonnet 4.5", + }) + d.SetId("model-123") + + if err := resourceLiteLLMModelRead(d, NewClient(srv.URL, "test-key", true)); err != nil { + t.Fatalf("read failed: %v", err) + } + if got := d.Get("display_name").(string); got != tc.want { + t.Errorf("expected display_name %q, got %q", tc.want, got) + } + }) + } +} + +func updateResourceData(t *testing.T, oldDisplayName, newDisplayName string) *schema.ResourceData { + t.Helper() + res := resourceLiteLLMModel() + attrs := map[string]string{ + "model_name": "sonnet-4-5-anthropic", + "custom_llm_provider": "anthropic", + "base_model": "claude-sonnet-4-5", + } + if oldDisplayName != "" { + attrs["display_name"] = oldDisplayName + } + state := &terraform.InstanceState{ID: "model-123", Attributes: attrs} + diff, err := res.Diff(context.Background(), state, &terraform.ResourceConfig{Config: map[string]interface{}{ + "model_name": "sonnet-4-5-anthropic", + "custom_llm_provider": "anthropic", + "base_model": "claude-sonnet-4-5", + "display_name": newDisplayName, + }}, nil) + if err != nil { + t.Fatalf("diff failed: %v", err) + } + d, err := schema.InternalMap(res.Schema).Data(state, diff) + if err != nil { + t.Fatalf("data failed: %v", err) + } + return d +} + +func TestResourceLiteLLMModelUpdatePatchesDisplayName(t *testing.T) { + cases := map[string]struct { + newName string + }{ + "changed name is patched": {newName: "Claude Sonnet 4.5 v2"}, + "cleared name is patched": {newName: ""}, + } + for name, tc := range cases { + t.Run(name, func(t *testing.T) { + var patchPayload map[string]interface{} + var patchPath string + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch { + case r.Method == http.MethodPost && r.URL.Path == "/model/update": + w.Write([]byte(modelInfoBody("Claude Sonnet 4.5"))) + case r.Method == http.MethodPatch: + patchPath = r.URL.Path + if err := json.NewDecoder(r.Body).Decode(&patchPayload); err != nil { + t.Errorf("failed to decode patch payload: %v", err) + } + w.Write([]byte(modelInfoBody(tc.newName))) + case r.URL.Path == "/model/info": + w.Write([]byte(modelInfoDataEnvelope(tc.newName))) + default: + t.Errorf("unexpected request: %s %s", r.Method, r.URL.Path) + w.WriteHeader(http.StatusNotFound) + } + })) + defer srv.Close() + + d := updateResourceData(t, "Claude Sonnet 4.5", tc.newName) + if err := resourceLiteLLMModelUpdate(d, NewClient(srv.URL, "test-key", true)); err != nil { + t.Fatalf("update failed: %v", err) + } + if patchPath != "/model/model-123/update" { + t.Fatalf("expected PATCH /model/model-123/update, got %q", patchPath) + } + modelInfo := patchPayload["model_info"].(map[string]interface{}) + if modelInfo["display_name"] != tc.newName { + t.Errorf("expected patched display_name %q, got %v", tc.newName, modelInfo["display_name"]) + } + if got := d.Get("display_name").(string); got != tc.newName { + t.Errorf("expected state display_name %q, got %q", tc.newName, got) + } + }) + } +} + +func TestResourceLiteLLMModelUpdateSkipsPatchWhenDisplayNameUnchanged(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Method == http.MethodPatch { + t.Errorf("unexpected PATCH %s", r.URL.Path) + } + w.Write([]byte(modelInfoDataEnvelope("Claude Sonnet 4.5"))) + })) + defer srv.Close() + + d := updateResourceData(t, "Claude Sonnet 4.5", "Claude Sonnet 4.5") + if err := resourceLiteLLMModelUpdate(d, NewClient(srv.URL, "test-key", true)); err != nil { + t.Fatalf("update failed: %v", err) + } +} diff --git a/terraform/provider/litellm/types.go b/terraform/provider/litellm/types.go index 8bcf7dc4fe3..a8784b8a6a9 100644 --- a/terraform/provider/litellm/types.go +++ b/terraform/provider/litellm/types.go @@ -25,6 +25,16 @@ type ModelResponse struct { Additional map[string]interface{} `json:"additional"` } +// ModelInfoPatch is the body for PATCH /model/{id}/update; display_name is sent even when empty so it can be cleared. +type ModelInfoPatch struct { + ModelInfo ModelInfoPatchFields `json:"model_info"` +} + +type ModelInfoPatchFields struct { + ID string `json:"id"` + DisplayName string `json:"display_name"` +} + // ModelRequest represents a request to create or update a model. type ModelRequest struct { ModelName string `json:"model_name"` @@ -108,12 +118,13 @@ type LiteLLMParams struct { // ModelInfo represents information about a model. type ModelInfo struct { - ID string `json:"id"` - DBModel bool `json:"db_model"` - BaseModel string `json:"base_model"` - Tier string `json:"tier"` - Mode string `json:"mode"` - TeamID string `json:"team_id,omitempty"` + ID string `json:"id"` + DBModel bool `json:"db_model"` + BaseModel string `json:"base_model"` + Tier string `json:"tier"` + Mode string `json:"mode"` + TeamID string `json:"team_id,omitempty"` + DisplayName string `json:"display_name,omitempty"` } // Key represents a LiteLLM API key. diff --git a/terraform/provider/litellm/utils.go b/terraform/provider/litellm/utils.go index f8f66afba3c..ce1dae55f59 100644 --- a/terraform/provider/litellm/utils.go +++ b/terraform/provider/litellm/utils.go @@ -55,6 +55,13 @@ func handleAPIResponse(resp *http.Response, reqBody interface{}, client *Client) resp.Status, client.redactSensitiveData(string(bodyBytes)), client.redactSensitiveData(string(reqBodyBytes))) } + var envelope struct { + Data []json.RawMessage `json:"data"` + } + if err := json.Unmarshal(bodyBytes, &envelope); err == nil && len(envelope.Data) > 0 { + bodyBytes = envelope.Data[0] + } + var modelResp ModelResponse if err := json.Unmarshal(bodyBytes, &modelResp); err != nil { return nil, fmt.Errorf("failed to parse response: %v", err) diff --git a/terraform/provider/tools/endpointaudit/coverage_allowlist.txt b/terraform/provider/tools/endpointaudit/coverage_allowlist.txt index 7aacf7ceab9..e4574031d86 100644 --- a/terraform/provider/tools/endpointaudit/coverage_allowlist.txt +++ b/terraform/provider/tools/endpointaudit/coverage_allowlist.txt @@ -97,7 +97,6 @@ GET /guardrails/{guardrail_id} GET /prompts/{prompt_id} GET /prompts/{prompt_id}/versions PATCH /guardrails/{guardrail_id} -PATCH /model/{model_id}/update PATCH /prompts/{prompt_id} PATCH /team/{team_id} POST /team/model/add From de8aeff6c6d2f360b4ade529ecb90159c896bf1e Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 24 Sep 2026 17:07:10 -0500 Subject: [PATCH 089/120] feat(proxy_cli): add --validate_config dry-run flag (#41705) * feat(proxy_cli): add --validate_config dry-run flag Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * style(tests): format --validate_config CliRunner calls Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy_cli): run --validate_config before the ollama auto-start Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(proxy_cli): restore file and add ollama validate_config regression test Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Co-authored-by: yassin --- litellm/proxy/proxy_cli.py | 29 ++++++ tests/test_litellm/proxy/test_proxy_cli.py | 103 +++++++++++++++++++++ 2 files changed, 132 insertions(+) diff --git a/litellm/proxy/proxy_cli.py b/litellm/proxy/proxy_cli.py index ac69f2e3894..27c03d2d5d7 100644 --- a/litellm/proxy/proxy_cli.py +++ b/litellm/proxy/proxy_cli.py @@ -210,6 +210,25 @@ class ProxyInitializationHelpers: response: Final = httpx.get(url=f"http://{host}:{port}/health") print(json.dumps(response.json(), indent=4)) + @staticmethod + def _run_config_validation(config: str | None) -> None: + if config is None: + raise click.UsageError("--validate_config requires --config ") + import asyncio + + from litellm.proxy.proxy_server import ProxyConfig + + async def _load() -> int: + _, model_list, _ = await ProxyConfig().load_config(router=None, config_file_path=config) + return len(model_list) + + try: + model_count: Final = asyncio.run(_load()) + except Exception as error: + click.echo(f"LiteLLM: config validation failed: {error}", err=True) + raise click.exceptions.Exit(1) from error + click.echo(f"LiteLLM: config OK ({model_count} models)") + @staticmethod def _run_test_chat_completion( host: str, @@ -887,6 +906,12 @@ class ProxyInitializationHelpers: default=False, help="Skip starting the server after setup (useful for migrations only)", ) +@click.option( + "--validate_config", + is_flag=True, + default=False, + help="Load and validate the config file (including mcp_servers) without starting the server, then exit. Exit code 1 on any config error.", +) @click.option( "--keepalive_timeout", default=None, @@ -1027,6 +1052,7 @@ def run_server( log_config, use_prisma_db_push: bool, skip_server_startup, + validate_config: bool, keepalive_timeout, timeout_worker_healthcheck, max_requests_before_restart, @@ -1069,6 +1095,9 @@ def run_server( if version is True: ProxyInitializationHelpers._echo_litellm_version() return + if validate_config is True: + ProxyInitializationHelpers._run_config_validation(config) + return if model and "ollama" in model and api_base is None: ProxyInitializationHelpers._run_ollama_serve() if health is True: diff --git a/tests/test_litellm/proxy/test_proxy_cli.py b/tests/test_litellm/proxy/test_proxy_cli.py index b84efda3308..a275dd62400 100644 --- a/tests/test_litellm/proxy/test_proxy_cli.py +++ b/tests/test_litellm/proxy/test_proxy_cli.py @@ -3224,3 +3224,106 @@ class TestLibpqSslParamTranslation: assert query["sslmode"] == ["require"] assert query["sslcert"] == ["/certs/rds-bundle.pem"] assert query["sslaccept"] == ["strict"] + + +@pytest.mark.xdist_group("proxy_cli") +class TestValidateConfigFlag: + def test_validate_config_valid_config_exits_zero(self, tmp_path, monkeypatch): + from click.testing import CliRunner + + from litellm.proxy.proxy_cli import run_server + + monkeypatch.delenv("DATABASE_URL", raising=False) + monkeypatch.delenv("DIRECT_URL", raising=False) + config_path = tmp_path / "config.yaml" + config_path.write_text( + yaml.safe_dump( + { + "model_list": [ + { + "model_name": "gpt-4o", + "litellm_params": { + "model": "openai/gpt-4o", + "api_key": "sk-fake", + }, + } + ] + } + ) + ) + + result = CliRunner().invoke(run_server, ["--config", str(config_path), "--validate_config"]) + + assert result.exit_code == 0, f"exit_code={result.exit_code}, output={result.output}" + assert "config OK" in result.output + + def test_validate_config_invalid_mcp_server_exits_one(self, tmp_path, monkeypatch): + from click.testing import CliRunner + + from litellm.proxy.proxy_cli import run_server + + monkeypatch.delenv("DATABASE_URL", raising=False) + monkeypatch.delenv("DIRECT_URL", raising=False) + config_path = tmp_path / "config.yaml" + config_path.write_text( + yaml.safe_dump( + { + "mcp_servers": { + "zapier": { + "url": "https://example.com/mcp", + "transport": "http", + "per_server_oauth_discovery": "yes", + } + } + } + ) + ) + + result = CliRunner().invoke(run_server, ["--config", str(config_path), "--validate_config"]) + + assert result.exit_code == 1, f"exit_code={result.exit_code}, output={result.output}" + assert "per_server_oauth_discovery must be a boolean" in result.output + + def test_validate_config_without_config_is_usage_error(self, monkeypatch): + from click.testing import CliRunner + + from litellm.proxy.proxy_cli import run_server + + result = CliRunner().invoke(run_server, ["--validate_config"]) + + assert result.exit_code != 0 + assert "--validate_config requires --config" in result.output + + @patch("subprocess.Popen") + def test_validate_config_with_ollama_model_does_not_start_ollama(self, mock_popen, tmp_path, monkeypatch): + from click.testing import CliRunner + + from litellm.proxy.proxy_cli import run_server + + monkeypatch.delenv("DATABASE_URL", raising=False) + monkeypatch.delenv("DIRECT_URL", raising=False) + config_path = tmp_path / "config.yaml" + config_path.write_text( + yaml.safe_dump( + { + "model_list": [ + { + "model_name": "gpt-4o", + "litellm_params": { + "model": "openai/gpt-4o", + "api_key": "sk-fake", + }, + } + ] + } + ) + ) + + result = CliRunner().invoke( + run_server, + ["--config", str(config_path), "--model", "ollama/llama3", "--validate_config"], + ) + + assert result.exit_code == 0, f"exit_code={result.exit_code}, output={result.output}" + assert "config OK" in result.output + mock_popen.assert_not_called() From 7faeb15ff3dfa5e227e32fb69b3a2b6945b3ad72 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 24 Sep 2026 15:17:31 -0700 Subject: [PATCH 090/120] fix(e2e): skip unpublished npm versions in the Claude Code PR-gate resolver (#43053) npm keeps an unpublished version's timestamp in the packument's time map but drops it from versions, so the resolver could hand npm install a version it refuses with ETARGET. Only versions still present in versions are candidates now. Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- .../_pr_gate_unit_tests/__init__.py | 0 .../test_pr_gate_version_resolver.py | 60 +++++++++++++++++++ .../claude_code/pr_gate_version_resolver.py | 7 ++- 3 files changed, 66 insertions(+), 1 deletion(-) create mode 100644 tests/e2e/claude_code/_pr_gate_unit_tests/__init__.py create mode 100644 tests/e2e/claude_code/_pr_gate_unit_tests/test_pr_gate_version_resolver.py diff --git a/tests/e2e/claude_code/_pr_gate_unit_tests/__init__.py b/tests/e2e/claude_code/_pr_gate_unit_tests/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/e2e/claude_code/_pr_gate_unit_tests/test_pr_gate_version_resolver.py b/tests/e2e/claude_code/_pr_gate_unit_tests/test_pr_gate_version_resolver.py new file mode 100644 index 00000000000..0ed7e2bf083 --- /dev/null +++ b/tests/e2e/claude_code/_pr_gate_unit_tests/test_pr_gate_version_resolver.py @@ -0,0 +1,60 @@ +"""Unit tests for the Claude Code PR-gate version resolver. + +Markerless harness tests: they feed the resolver a hand-built packument and a +fixed clock, so they run without a proxy, never reach the npm registry, and +carry no `e2e` marker. +""" + +from __future__ import annotations + +from datetime import datetime, timezone +from typing import Final, Mapping + +import pytest + +from claude_code.pr_gate_version_resolver import NoEligibleVersionError, resolve_pr_gate_version + +NOW: Final = datetime(2026, 4, 25, 12, 0, tzinfo=timezone.utc) +INSIDE_THE_2_1_88_WINDOW: Final = datetime(2026, 4, 3, 12, 0, tzinfo=timezone.utc) + + +def _packument(times: Mapping[str, str], unpublished: frozenset[str] = frozenset()) -> dict[str, object]: + return { + "name": "@anthropic-ai/claude-code", + "time": {"created": "2024-01-01T00:00:00.000Z", "modified": "2026-04-25T00:00:00.000Z", **times}, + "versions": {version: {"version": version} for version in times if version not in unpublished}, + } + + +def test_skips_a_version_npm_has_unpublished() -> None: + metadata: Final = _packument( + { + "2.1.87": "2026-03-28T20:00:00.000Z", + "2.1.88": "2026-03-30T22:36:48.424Z", + "2.1.89": "2026-03-31T23:32:40.000Z", + }, + unpublished=frozenset({"2.1.88"}), + ) + assert resolve_pr_gate_version(metadata=metadata, as_of=INSIDE_THE_2_1_88_WINDOW) == "2.1.87" + + +def test_raises_when_the_only_old_enough_version_is_unpublished() -> None: + metadata: Final = _packument( + {"2.1.88": "2026-03-30T22:36:48.424Z", "2.1.89": "2026-03-31T23:32:40.000Z"}, + unpublished=frozenset({"2.1.88"}), + ) + with pytest.raises(NoEligibleVersionError): + resolve_pr_gate_version(metadata=metadata, as_of=INSIDE_THE_2_1_88_WINDOW) + + +def test_picks_the_newest_published_version_at_least_min_age_old() -> None: + metadata: Final = _packument( + { + "2.1.118": "2026-04-15T10:00:00.000Z", + "2.1.119": "2026-04-21T10:00:00.000Z", + "2.2.0-rc.1": "2026-04-22T10:00:00.000Z", + "2.1.120": "2026-04-23T10:00:00.000Z", + "2.1.121": "2026-04-25T11:00:00.000Z", + } + ) + assert resolve_pr_gate_version(metadata=metadata, as_of=NOW) == "2.1.119" diff --git a/tests/e2e/claude_code/pr_gate_version_resolver.py b/tests/e2e/claude_code/pr_gate_version_resolver.py index 82e12a2bf15..756dafb7d73 100644 --- a/tests/e2e/claude_code/pr_gate_version_resolver.py +++ b/tests/e2e/claude_code/pr_gate_version_resolver.py @@ -80,7 +80,9 @@ def resolve_pr_gate_version( "Newest" means newest by **publish time**, not semver string order — if a patch lands on an older major after a newer release, the - patched line is the eligible one. + patched line is the eligible one. A version npm has unpublished keeps + its ``time`` entry but drops out of ``versions``, so only versions + still present in ``versions`` are candidates. Args: metadata: Pre-fetched npm packument (skips the HTTP call). Useful @@ -101,6 +103,7 @@ def resolve_pr_gate_version( metadata = fetch(package_name) times = metadata.get("time") or {} + versions = metadata.get("versions") or {} if as_of is None: as_of = datetime.now(timezone.utc) cutoff = as_of - min_age @@ -109,6 +112,8 @@ def resolve_pr_gate_version( for version, raw_ts in times.items(): if version in _TIME_META_KEYS: continue + if version not in versions: + continue if not isinstance(raw_ts, str): continue if "-" in version: From 040b37fa49b99be78b997f76ad46bd6f10c07f27 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 24 Sep 2026 15:18:55 -0700 Subject: [PATCH 091/120] chore(cost-map): move azure gpt-realtime-2.1-mini deprecation date to the later Models API date (#43058) Co-authored-by: kerry Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/model_prices_and_context_window_backup.json | 2 +- model_prices_and_context_window.json | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 680b6cceab6..9bec83f6b08 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -6092,7 +6092,7 @@ "cache_creation_input_audio_token_cost": 3e-07, "cache_read_input_audio_token_cost": 3e-07, "cache_read_input_token_cost": 6e-08, - "deprecation_date": "2027-06-25", + "deprecation_date": "2027-07-31", "input_cost_per_audio_token": 1e-05, "input_cost_per_image_token": 8e-07, "input_cost_per_token": 6e-07, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 680b6cceab6..9bec83f6b08 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -6092,7 +6092,7 @@ "cache_creation_input_audio_token_cost": 3e-07, "cache_read_input_audio_token_cost": 3e-07, "cache_read_input_token_cost": 6e-08, - "deprecation_date": "2027-06-25", + "deprecation_date": "2027-07-31", "input_cost_per_audio_token": 1e-05, "input_cost_per_image_token": 8e-07, "input_cost_per_token": 6e-07, From 4aa3ff47fe1fd525c54edd9adc4e206cf439baea Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 24 Sep 2026 15:23:05 -0700 Subject: [PATCH 092/120] docs(github): require UI before/after screenshots and intentional UX change note in PR template (#43021) * docs(github): require UI before/after screenshots and intentional UX change note in PR template Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * docs(github): move intentional change note into TLDR rules and dedupe screenshots Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * docs(github): refresh user flow screenshots with new commits Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: kerry Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .github/pull_request_template.md | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/.github/pull_request_template.md b/.github/pull_request_template.md index 4b3878bed11..1fe0c602036 100644 --- a/.github/pull_request_template.md +++ b/.github/pull_request_template.md @@ -6,7 +6,8 @@ ## TLDR - + Problem this solves: @@ -28,7 +29,8 @@ How it solves it: No LiteLLM internals: never name functions, files, DB tables, config classes, hooks, callbacks, or code paths. "The upload hands back an ID that looks like OpenAI's own `file-abc123` instead of the scrambled one the gateway returned" is right, "no managed-file row was registered" is wrong Keep the two lists step-for-step identical until they diverge, so the changed step is obvious If the bug had a security or authorization consequence, end each list with what another user could or could no longer do - Regenerate this section whenever new commits change the PR's behavior, so it never describes an older revision + Regenerate this section, screenshots included, whenever new commits change the PR's behavior, so it never describes an older revision + If the PR changes what an Admin UI page shows, embed a before and an after screenshot of that page right after its list, taken at the same URL on the same data, with the rows, fields, or controls that changed boxed in red so a reader spots the difference without reading the steps. These are the UI screenshots for Screenshots / Proof of Fix too: embed them once here and have that section's Before and After steps point back to them instead of repeating the images Example: From bf0187072bb360153400a17895e65dc10f4a4110 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 24 Sep 2026 15:49:59 -0700 Subject: [PATCH 093/120] ci: move caching, proxy-extras, gateway and enterprise tests into tests/unit and run them from litellm-tests (#42902) * ci: fix the litellm-tests unit job with sysmon coverage, an env allowlist and coverage upload on failure * test: replace key-dependent proxy, enterprise and mcp unit tests with synthetic values and integration and e2e coverage * test: drop key reads at the legacy proxy, enterprise and mcp paths and wire the gemini pass-through split * ci: move caching, proxy-extras, gateway and enterprise tests into tests/unit and run them from litellm-tests under their legacy flags * ci: move caching, proxy-extras, gateway and enterprise tests into tests/unit and run them from litellm-tests under their legacy flags * ci: fail the unit shard when circleci tests split errors * test: drop restating comments from the gemini pass-through split * ci: exit the unit shard cleanly when circleci tests split assigns it no files --------- Co-authored-by: yuneng --- .circleci/scripts/unit_selection.sh | 63 +++++++++++++++++++ .circleci/tests.yml | 33 +++++++++- .github/scripts/assert_ci_coverage.py | 12 +++- .github/workflows/_test-unit-base.yml | 33 ++++++++-- .github/workflows/test-unit.yml | 22 ++++--- Makefile | 2 +- tests/test_litellm/test_assert_ci_coverage.py | 2 +- .../proxy => unit/caching}/__init__.py | 0 .../caching}/test_cache_preset_key.py | 0 .../caching}/test_caching_handler.py | 0 .../test_responses_stream_cache_keys.py | 0 .../caching}/test_unit_test_caching.py | 0 tests/unit/conftest.py | 2 +- tests/{ => unit}/enterprise/conftest.py | 0 .../send_emails/__init__.py | 0 .../send_emails/test_base_email.py | 0 .../send_emails/test_endpoints.py | 0 .../send_emails/test_resend_email.py | 0 .../send_emails/test_sendgrid_email.py | 0 .../test_prometheus_logging_callbacks.py | 0 .../unit/enterprise/integrations/__init__.py | 0 .../integrations/test_custom_guardrail.py | 0 .../integrations/test_prometheus.py | 0 .../test_prometheus_unit_tests.py | 0 tests/unit/enterprise/proxy/__init__.py | 0 tests/unit/enterprise/proxy/auth/__init__.py | 0 .../proxy/auth/test_route_checks.py | 0 .../proxy/auth/test_user_api_key_auth.py | 0 .../enterprise/proxy/guardrails/__init__.py | 0 .../enterprise}/proxy/guardrails/conftest.py | 0 .../test_apply_guardrail_endpoint.py | 0 .../test_bedrock_apply_guardrail.py | 0 tests/unit/enterprise/proxy/hooks/__init__.py | 0 .../proxy/hooks/test_managed_files.py | 0 .../proxy/management_endpoints/__init__.py | 0 .../test_internal_user_endpoints.py | 0 .../test_project_endpoints_prisma.py | 0 .../test_afile_retrieve_returns_unified_id.py | 0 .../proxy/test_audit_logging_endpoints.py | 0 .../test_batch_retrieve_input_file_id.py | 0 ...trieve_registers_missing_output_file_id.py | 0 ..._retrieve_returns_unified_input_file_id.py | 0 ..._batch_update_db_managed_output_file_id.py | 0 .../test_deleted_file_returns_403_not_404.py | 0 .../proxy/test_enterprise_routes.py | 0 .../proxy/test_file_deletion_blocking.py | 0 .../proxy/test_managed_files_access_check.py | 0 .../proxy/test_managed_files_hook.py | 0 tests/unit/gateway/__init__.py | 0 .../gateway}/test_launch.py | 0 tests/unit/litellm_proxy_extras/__init__.py | 0 .../test_litellm_proxy_extras_logging.py | 0 .../test_litellm_proxy_extras_utils.py | 6 +- 53 files changed, 152 insertions(+), 23 deletions(-) create mode 100755 .circleci/scripts/unit_selection.sh rename tests/{test_litellm/enterprise/proxy => unit/caching}/__init__.py (100%) rename tests/{local_testing => unit/caching}/test_cache_preset_key.py (100%) rename tests/{local_testing => unit/caching}/test_caching_handler.py (100%) rename tests/{local_testing => unit/caching}/test_responses_stream_cache_keys.py (100%) rename tests/{local_testing => unit/caching}/test_unit_test_caching.py (100%) rename tests/{ => unit}/enterprise/conftest.py (100%) create mode 100644 tests/unit/enterprise/enterprise_callbacks/send_emails/__init__.py rename tests/{test_litellm => unit}/enterprise/enterprise_callbacks/send_emails/test_base_email.py (100%) rename tests/{test_litellm => unit}/enterprise/enterprise_callbacks/send_emails/test_endpoints.py (100%) rename tests/{test_litellm => unit}/enterprise/enterprise_callbacks/send_emails/test_resend_email.py (100%) rename tests/{test_litellm => unit}/enterprise/enterprise_callbacks/send_emails/test_sendgrid_email.py (100%) rename tests/{enterprise/litellm_enterprise => unit/enterprise}/enterprise_callbacks/test_prometheus_logging_callbacks.py (100%) create mode 100644 tests/unit/enterprise/integrations/__init__.py rename tests/{enterprise/litellm_enterprise => unit/enterprise}/integrations/test_custom_guardrail.py (100%) rename tests/{enterprise/litellm_enterprise => unit/enterprise}/integrations/test_prometheus.py (100%) rename tests/{enterprise/litellm_enterprise => unit/enterprise}/integrations/test_prometheus_unit_tests.py (100%) create mode 100644 tests/unit/enterprise/proxy/__init__.py create mode 100644 tests/unit/enterprise/proxy/auth/__init__.py rename tests/{enterprise/litellm_enterprise => unit/enterprise}/proxy/auth/test_route_checks.py (100%) rename tests/{enterprise/litellm_enterprise => unit/enterprise}/proxy/auth/test_user_api_key_auth.py (100%) create mode 100644 tests/unit/enterprise/proxy/guardrails/__init__.py rename tests/{enterprise/litellm_enterprise => unit/enterprise}/proxy/guardrails/conftest.py (100%) rename tests/{enterprise/litellm_enterprise => unit/enterprise}/proxy/guardrails/test_apply_guardrail_endpoint.py (100%) rename tests/{enterprise/litellm_enterprise => unit/enterprise}/proxy/guardrails/test_bedrock_apply_guardrail.py (100%) create mode 100644 tests/unit/enterprise/proxy/hooks/__init__.py rename tests/{enterprise/litellm_enterprise => unit/enterprise}/proxy/hooks/test_managed_files.py (100%) create mode 100644 tests/unit/enterprise/proxy/management_endpoints/__init__.py rename tests/{enterprise/litellm_enterprise => unit/enterprise}/proxy/management_endpoints/test_internal_user_endpoints.py (100%) rename tests/{enterprise/litellm_enterprise => unit/enterprise}/proxy/management_endpoints/test_project_endpoints_prisma.py (100%) rename tests/{test_litellm => unit}/enterprise/proxy/test_afile_retrieve_returns_unified_id.py (100%) rename tests/{enterprise/litellm_enterprise => unit/enterprise}/proxy/test_audit_logging_endpoints.py (100%) rename tests/{test_litellm => unit}/enterprise/proxy/test_batch_retrieve_input_file_id.py (100%) rename tests/{test_litellm => unit}/enterprise/proxy/test_batch_retrieve_registers_missing_output_file_id.py (100%) rename tests/{test_litellm => unit}/enterprise/proxy/test_batch_retrieve_returns_unified_input_file_id.py (100%) rename tests/{test_litellm => unit}/enterprise/proxy/test_batch_update_db_managed_output_file_id.py (100%) rename tests/{test_litellm => unit}/enterprise/proxy/test_deleted_file_returns_403_not_404.py (100%) rename tests/{test_litellm => unit}/enterprise/proxy/test_enterprise_routes.py (100%) rename tests/{test_litellm => unit}/enterprise/proxy/test_file_deletion_blocking.py (100%) rename tests/{test_litellm => unit}/enterprise/proxy/test_managed_files_access_check.py (100%) rename tests/{test_litellm => unit}/enterprise/proxy/test_managed_files_hook.py (100%) create mode 100644 tests/unit/gateway/__init__.py rename tests/{test_gateway => unit/gateway}/test_launch.py (100%) create mode 100644 tests/unit/litellm_proxy_extras/__init__.py rename tests/{litellm-proxy-extras => unit/litellm_proxy_extras}/test_litellm_proxy_extras_logging.py (100%) rename tests/{litellm-proxy-extras => unit/litellm_proxy_extras}/test_litellm_proxy_extras_utils.py (99%) diff --git a/.circleci/scripts/unit_selection.sh b/.circleci/scripts/unit_selection.sh new file mode 100755 index 00000000000..8d2b8a42691 --- /dev/null +++ b/.circleci/scripts/unit_selection.sh @@ -0,0 +1,63 @@ +#!/usr/bin/env bash +set -euo pipefail + +flag="${1:?usage: unit_selection.sh }" + +legacy_flags=( + caching-local + enterprise-package + enterprise-routing + proxy-extras + proxy-infra +) + +legacy_paths() { + case "$1" in + caching-local) echo tests/unit/caching ;; + enterprise-package) + echo tests/unit/enterprise/integrations + echo tests/unit/enterprise/proxy/auth + echo tests/unit/enterprise/proxy/guardrails + echo tests/unit/enterprise/proxy/hooks + echo tests/unit/enterprise/proxy/management_endpoints + echo tests/unit/enterprise/proxy/test_audit_logging_endpoints.py + echo tests/unit/enterprise/enterprise_callbacks/test_prometheus_logging_callbacks.py ;; + enterprise-routing) + echo tests/unit/enterprise/enterprise_callbacks/send_emails + echo tests/unit/enterprise/proxy/test_afile_retrieve_returns_unified_id.py + echo tests/unit/enterprise/proxy/test_batch_retrieve_input_file_id.py + echo tests/unit/enterprise/proxy/test_batch_retrieve_registers_missing_output_file_id.py + echo tests/unit/enterprise/proxy/test_batch_retrieve_returns_unified_input_file_id.py + echo tests/unit/enterprise/proxy/test_batch_update_db_managed_output_file_id.py + echo tests/unit/enterprise/proxy/test_deleted_file_returns_403_not_404.py + echo tests/unit/enterprise/proxy/test_enterprise_routes.py + echo tests/unit/enterprise/proxy/test_file_deletion_blocking.py + echo tests/unit/enterprise/proxy/test_managed_files_access_check.py + echo tests/unit/enterprise/proxy/test_managed_files_hook.py ;; + proxy-extras) echo tests/unit/litellm_proxy_extras ;; + proxy-infra) echo tests/unit/gateway ;; + *) echo "unit_selection.sh: unknown flag $1" >&2; exit 1 ;; + esac +} + +expand() { + while read -r path; do + if [ -d "$path" ]; then + find "$path" -name 'test_*.py' + elif [ -f "$path" ]; then + echo "$path" + else + echo "unit_selection.sh: $path does not exist" >&2 + exit 1 + fi + done +} + +if [ "$flag" = unit ]; then + comm -23 \ + <(find tests/unit -name 'test_*.py' | sort) \ + <(for legacy in "${legacy_flags[@]}"; do legacy_paths "$legacy"; done | expand | sort) + exit 0 +fi + +legacy_paths "$flag" | expand | sort diff --git a/.circleci/tests.yml b/.circleci/tests.yml index 1afb935453f..38fe44bb625 100644 --- a/.circleci/tests.yml +++ b/.circleci/tests.yml @@ -171,6 +171,12 @@ jobs: shards: type: integer default: 6 + workers: + type: integer + default: 4 + dist: + type: string + default: loadscope base_ref: type: string default: "" @@ -199,17 +205,19 @@ jobs: no_output_timeout: 20m command: | mkdir -p test-results/<< parameters.flag >> - selection="$(find tests/unit -name 'test_*.py' | sort)" || { echo "test selection failed for << parameters.flag >>"; exit 1; } - [ -n "${selection}" ] || { echo "test selection produced no files for << parameters.flag >>"; exit 1; } + selection="$(bash .circleci/scripts/unit_selection.sh << parameters.flag >>)" || { echo "unit_selection.sh failed for << parameters.flag >>"; exit 1; } + [ -n "${selection}" ] || { echo "unit_selection.sh produced no files for << parameters.flag >>"; exit 1; } shard="$(printf '%s\n' "${selection}" | circleci tests split --split-by=timings --timings-type=filename)" || { echo "circleci tests split failed for << parameters.flag >>"; exit 1; } [ -n "${shard}" ] || { echo "shard ${CIRCLE_NODE_INDEX} received no << parameters.flag >> files; nothing to run"; exit 0; } mapfile -t files < <(printf '%s\n' "${shard}") + xdist_args=() + if [ "<< parameters.workers >>" -gt 0 ]; then xdist_args=(-n << parameters.workers >> --dist=<< parameters.dist >>); fi rerun_args=(-p no:rerunfailures) if [ "<< parameters.reruns >>" -gt 0 ]; then rerun_args=(--reruns << parameters.reruns >> --reruns-delay 1 --rerun-except "from pytest-timeout"); fi test_env=(PATH="$PATH" HOME="$HOME" CI=true COVERAGE_CORE="$COVERAGE_CORE" LITELLM_LOCAL_MODEL_COST_MAP="$LITELLM_LOCAL_MODEL_COST_MAP") set +e env -i "${test_env[@]}" \ - uv run --no-sync pytest "${files[@]}" "${rerun_args[@]}" -p no:pytest-retry --timeout=90 -n 4 --dist=loadscope --tb=short --durations=20 -o junit_family=xunit1 --junitxml=test-results/<< parameters.flag >>/junit.xml --cov=./litellm --cov=./enterprise/litellm_enterprise --cov-report=xml:coverage.xml --cov-config=pyproject.toml + uv run --no-sync pytest "${files[@]}" "${rerun_args[@]}" -p no:pytest-retry --timeout=90 "${xdist_args[@]}" --tb=short --durations=20 -o junit_family=xunit1 --junitxml=test-results/<< parameters.flag >>/junit.xml --cov=./litellm --cov=./enterprise/litellm_enterprise --cov-report=xml:coverage.xml --cov-config=pyproject.toml status=$? set -e if [ "$status" -eq 5 ]; then echo "pytest collected no tests from the shard; passing"; exit 0; fi @@ -293,6 +301,25 @@ workflows: - unit: base_ref: << pipeline.event.name == "pull_request" and pipeline.event.github.pull_request.base.ref or "" >> pull_request_url: << pipeline.event.name == "pull_request" and pipeline.event.github.pull_request.url or "" >> + - unit: + name: unit-<< matrix.flag >> + shards: 1 + workers: 2 + reruns: 2 + matrix: + parameters: + flag: [caching-local, proxy-extras, enterprise-routing] + base_ref: << pipeline.event.name == "pull_request" and pipeline.event.github.pull_request.base.ref or "" >> + pull_request_url: << pipeline.event.name == "pull_request" and pipeline.event.github.pull_request.url or "" >> + - unit: + name: unit-<< matrix.flag >> + shards: 1 + reruns: 2 + matrix: + parameters: + flag: [enterprise-package, proxy-infra] + base_ref: << pipeline.event.name == "pull_request" and pipeline.event.github.pull_request.base.ref or "" >> + pull_request_url: << pipeline.event.name == "pull_request" and pipeline.event.github.pull_request.url or "" >> - documentation - integration: name: integration-<< matrix.suite >> diff --git a/.github/scripts/assert_ci_coverage.py b/.github/scripts/assert_ci_coverage.py index 2e008fe7ade..d8246225a3b 100644 --- a/.github/scripts/assert_ci_coverage.py +++ b/.github/scripts/assert_ci_coverage.py @@ -120,6 +120,13 @@ def _invoked_test_tokens(scalars: Iterable[Scalar]) -> frozenset[str]: ) +def _unit_selection_tokens(repo_root: pathlib.Path = REPO_ROOT) -> frozenset[str]: + script: Final = repo_root / ".circleci/scripts/unit_selection.sh" + if not script.is_file(): + return frozenset() + return frozenset(match.group(0).rstrip("/") for match in TEST_TOKEN_RE.finditer(_uncommented(script.read_text()))) + + def _built_dockerfile_tokens(scalars: Iterable[Scalar]) -> frozenset[str]: return frozenset( match.group(0) @@ -611,7 +618,10 @@ def main() -> int: scalars = _all_scalars() integration_paths, ownership_findings = _integration_ownership() - test_findings = _uncovered_tests(allowlist, _invoked_test_tokens(scalars) | integration_paths) + ownership_findings + test_findings = ( + _uncovered_tests(allowlist, _invoked_test_tokens(scalars) | _unit_selection_tokens() | integration_paths) + + ownership_findings + ) dockerfile_findings = _uncovered_dockerfiles(allowlist, _built_dockerfile_tokens(scalars)) stale_findings = _stale_allowlist_paths(allowlist, test_files=_test_files(), dockerfiles=_dockerfiles()) diff --git a/.github/workflows/_test-unit-base.yml b/.github/workflows/_test-unit-base.yml index 8faddd11df8..ef1dc53b4a6 100644 --- a/.github/workflows/_test-unit-base.yml +++ b/.github/workflows/_test-unit-base.yml @@ -13,6 +13,15 @@ on: have its path existence-checked like any other token. required: true type: string + fork-flag: + description: >- + Codecov flag of the `.circleci/tests.yml` job that now owns part of + this shard. CircleCI does not run on pull requests from forks, so on + those events this shard also runs the files + `.circleci/scripts/unit_selection.sh` lists for the flag. + required: false + type: string + default: "" workers: description: "Number of pytest-xdist workers" required: false @@ -92,6 +101,7 @@ jobs: pull-requests: read outputs: decision: ${{ steps.changes.outputs.decision }} + has-coverage: ${{ steps.tests.outputs.has-coverage }} steps: - uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0 @@ -160,10 +170,13 @@ jobs: uv run --no-sync prisma generate --schema litellm/proxy/schema.prisma - name: Run tests + id: tests if: steps.changes.outputs.decision != 'skip' timeout-minutes: ${{ inputs.timeout-minutes }} env: TEST_PATH: ${{ inputs.test-path }} + FORK_FLAG: ${{ inputs.fork-flag }} + IS_FORK: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.repo.full_name != github.repository }} MAX_FAILURES: ${{ inputs.max-failures }} WORKERS: ${{ inputs.workers }} RERUNS: ${{ inputs.reruns }} @@ -171,9 +184,18 @@ jobs: DIST: ${{ inputs.dist }} COVERAGE_CORE: sysmon run: | + echo "has-coverage=false" >> "$GITHUB_OUTPUT" + selection="${TEST_PATH}" + if [ "${IS_FORK}" = "true" ] && [ -n "${FORK_FLAG}" ]; then + selection="${TEST_PATH} $(bash .circleci/scripts/unit_selection.sh "${FORK_FLAG}" | tr '\n' ' ')" + fi + if [ -z "${selection// /}" ]; then + echo "shard selection is empty on this event (CircleCI flag ${FORK_FLAG:-none} owns it); nothing to run" + exit 0 + fi pytest_args=() existing_paths=0 - for token in ${TEST_PATH:?}; do + for token in ${selection}; do case "${token}" in -*) pytest_args+=("${token}") ;; *) @@ -187,7 +209,7 @@ jobs: esac done if [ "${existing_paths}" -eq 0 ]; then - echo "No path in TEST_PATH exists (${TEST_PATH}); nothing to run" + echo "No path in the selection exists (${selection}); nothing to run" exit 0 fi xdist_args=() @@ -209,8 +231,11 @@ jobs: --cov-config=pyproject.toml status=$? set -e + if [ -f coverage.xml ]; then + echo "has-coverage=true" >> "$GITHUB_OUTPUT" + fi if [ "$status" -eq 5 ]; then - echo "pytest collected no tests from ${TEST_PATH}; passing" + echo "pytest collected no tests from ${selection}; passing" exit 0 fi exit "$status" @@ -226,7 +251,7 @@ jobs: upload-coverage: name: Upload coverage to Codecov needs: run - if: always() && needs.run.outputs.decision != 'skip' + if: always() && needs.run.outputs.decision != 'skip' && needs.run.outputs.has-coverage == 'true' runs-on: ubuntu-latest permissions: contents: read diff --git a/.github/workflows/test-unit.yml b/.github/workflows/test-unit.yml index 686bbc89467..94de6040038 100644 --- a/.github/workflows/test-unit.yml +++ b/.github/workflows/test-unit.yml @@ -35,6 +35,10 @@ concurrency: # already a matrix and carries a shard-coverage guard that reads that file by # name. Folding it in here is a follow-up, together with generalising that guard # into assert_ci_coverage.py. +# +# `fork-flag` names the `.circleci/tests.yml` job that now runs part of the +# shard under the same Codecov flag. CircleCI does not build pull requests from +# forks, so the shard still runs those files there and skips them elsewhere. jobs: unit: name: ${{ matrix.shard }} @@ -65,10 +69,10 @@ jobs: - shard: enterprise-routing artifact-name: enterprise-routing test-path: >- - tests/test_litellm/enterprise tests/test_litellm/google_genai tests/test_litellm/router_utils tests/test_litellm/router_strategy + fork-flag: enterprise-routing workers: 2 reruns: 2 timeout-minutes: 20 @@ -200,7 +204,7 @@ jobs: tests/test_litellm/proxy/types_utils tests/test_litellm/proxy/logging_endpoints tests/test_litellm/proxy/test_*.py - tests/test_gateway + fork-flag: proxy-infra workers: 4 reruns: 2 timeout-minutes: 20 @@ -208,11 +212,8 @@ jobs: - shard: caching-local artifact-name: caching-local - test-path: >- - tests/local_testing/test_cache_preset_key.py - tests/local_testing/test_caching_handler.py - tests/local_testing/test_responses_stream_cache_keys.py - tests/local_testing/test_unit_test_caching.py + test-path: "" + fork-flag: caching-local workers: 2 reruns: 2 timeout-minutes: 20 @@ -220,7 +221,8 @@ jobs: - shard: proxy-extras artifact-name: proxy-extras - test-path: "tests/litellm-proxy-extras" + test-path: "" + fork-flag: proxy-extras workers: 2 reruns: 2 timeout-minutes: 20 @@ -228,7 +230,8 @@ jobs: - shard: enterprise-package artifact-name: enterprise-package - test-path: "tests/enterprise" + test-path: "" + fork-flag: enterprise-package workers: 4 reruns: 2 timeout-minutes: 20 @@ -247,6 +250,7 @@ jobs: uses: ./.github/workflows/_test-unit-base.yml with: test-path: ${{ matrix.test-path }} + fork-flag: ${{ matrix.fork-flag || '' }} workers: ${{ matrix.workers }} reruns: ${{ matrix.reruns }} timeout-minutes: ${{ matrix.timeout-minutes }} diff --git a/Makefile b/Makefile index ab7fab6aa99..6263b646c17 100644 --- a/Makefile +++ b/Makefile @@ -332,7 +332,7 @@ test-unit-core-utils: install-test-deps $(UV_RUN) pytest tests/test_litellm/litellm_core_utils --tb=short -vv -n 2 --durations=20 test-unit-other: install-test-deps - $(UV_RUN) pytest tests/test_litellm/caching tests/test_litellm/responses tests/test_litellm/secret_managers tests/test_litellm/vector_stores tests/test_litellm/a2a_protocol tests/test_litellm/anthropic_interface tests/test_litellm/completion_extras tests/test_litellm/containers tests/test_litellm/enterprise tests/test_litellm/experimental_mcp_client tests/test_litellm/google_genai tests/test_litellm/images tests/test_litellm/interactions tests/test_litellm/passthrough tests/test_litellm/router_strategy tests/test_litellm/router_utils tests/test_litellm/types --tb=short -vv -n 4 --durations=20 + $(UV_RUN) pytest tests/test_litellm/caching tests/test_litellm/responses tests/test_litellm/secret_managers tests/test_litellm/vector_stores tests/test_litellm/a2a_protocol tests/test_litellm/anthropic_interface tests/test_litellm/completion_extras tests/test_litellm/containers tests/unit/enterprise tests/test_litellm/experimental_mcp_client tests/test_litellm/google_genai tests/test_litellm/images tests/test_litellm/interactions tests/test_litellm/passthrough tests/test_litellm/router_strategy tests/test_litellm/router_utils tests/test_litellm/types --tb=short -vv -n 4 --durations=20 test-unit-root: install-test-deps $(UV_RUN) pytest tests/test_litellm/test_*.py --tb=short -vv -n 4 --durations=20 diff --git a/tests/test_litellm/test_assert_ci_coverage.py b/tests/test_litellm/test_assert_ci_coverage.py index 983707db606..8524a905745 100644 --- a/tests/test_litellm/test_assert_ci_coverage.py +++ b/tests/test_litellm/test_assert_ci_coverage.py @@ -340,7 +340,7 @@ def test_a_dockerfile_directory_entry_is_stale_because_only_an_exact_path_exempt def test_a_workflow_that_names_a_file_clears_it_from_the_slice_check(): named = coverage._workflow_named_tokens() assert named, "the workflows must name some test paths or the check proves nothing" - assert any(coverage._token_covers(token, "tests/local_testing/test_caching_handler.py") for token in named) + assert any(coverage._token_covers(token, "tests/proxy_unit_tests/test_proxy_custom_logger.py") for token in named) def test_the_slice_check_credits_only_workflows_never_the_circleci_config(): diff --git a/tests/test_litellm/enterprise/proxy/__init__.py b/tests/unit/caching/__init__.py similarity index 100% rename from tests/test_litellm/enterprise/proxy/__init__.py rename to tests/unit/caching/__init__.py diff --git a/tests/local_testing/test_cache_preset_key.py b/tests/unit/caching/test_cache_preset_key.py similarity index 100% rename from tests/local_testing/test_cache_preset_key.py rename to tests/unit/caching/test_cache_preset_key.py diff --git a/tests/local_testing/test_caching_handler.py b/tests/unit/caching/test_caching_handler.py similarity index 100% rename from tests/local_testing/test_caching_handler.py rename to tests/unit/caching/test_caching_handler.py diff --git a/tests/local_testing/test_responses_stream_cache_keys.py b/tests/unit/caching/test_responses_stream_cache_keys.py similarity index 100% rename from tests/local_testing/test_responses_stream_cache_keys.py rename to tests/unit/caching/test_responses_stream_cache_keys.py diff --git a/tests/local_testing/test_unit_test_caching.py b/tests/unit/caching/test_unit_test_caching.py similarity index 100% rename from tests/local_testing/test_unit_test_caching.py rename to tests/unit/caching/test_unit_test_caching.py diff --git a/tests/unit/conftest.py b/tests/unit/conftest.py index b3bb19a8b8a..202ecb80d7b 100644 --- a/tests/unit/conftest.py +++ b/tests/unit/conftest.py @@ -11,7 +11,7 @@ import litellm # noqa: E402 # litellm reads LITELLM_LOCAL_MODEL_COST_MAP at im import litellm.router as litellm_router_module # noqa: E402 # same import-time dependency import litellm.utils as litellm_utils_module # noqa: E402 # same import-time dependency -LOOPBACK_HOSTS: Final = ["127.0.0.1", "::1"] +LOOPBACK_HOSTS: Final = ["127.0.0.1", "::1", "localhost"] AMBIENT_AZURE_CREDENTIAL_ENV_VARS: Final = ( "AZURE_AD_TOKEN", "AZURE_TENANT_ID", diff --git a/tests/enterprise/conftest.py b/tests/unit/enterprise/conftest.py similarity index 100% rename from tests/enterprise/conftest.py rename to tests/unit/enterprise/conftest.py diff --git a/tests/unit/enterprise/enterprise_callbacks/send_emails/__init__.py b/tests/unit/enterprise/enterprise_callbacks/send_emails/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/enterprise/enterprise_callbacks/send_emails/test_base_email.py b/tests/unit/enterprise/enterprise_callbacks/send_emails/test_base_email.py similarity index 100% rename from tests/test_litellm/enterprise/enterprise_callbacks/send_emails/test_base_email.py rename to tests/unit/enterprise/enterprise_callbacks/send_emails/test_base_email.py diff --git a/tests/test_litellm/enterprise/enterprise_callbacks/send_emails/test_endpoints.py b/tests/unit/enterprise/enterprise_callbacks/send_emails/test_endpoints.py similarity index 100% rename from tests/test_litellm/enterprise/enterprise_callbacks/send_emails/test_endpoints.py rename to tests/unit/enterprise/enterprise_callbacks/send_emails/test_endpoints.py diff --git a/tests/test_litellm/enterprise/enterprise_callbacks/send_emails/test_resend_email.py b/tests/unit/enterprise/enterprise_callbacks/send_emails/test_resend_email.py similarity index 100% rename from tests/test_litellm/enterprise/enterprise_callbacks/send_emails/test_resend_email.py rename to tests/unit/enterprise/enterprise_callbacks/send_emails/test_resend_email.py diff --git a/tests/test_litellm/enterprise/enterprise_callbacks/send_emails/test_sendgrid_email.py b/tests/unit/enterprise/enterprise_callbacks/send_emails/test_sendgrid_email.py similarity index 100% rename from tests/test_litellm/enterprise/enterprise_callbacks/send_emails/test_sendgrid_email.py rename to tests/unit/enterprise/enterprise_callbacks/send_emails/test_sendgrid_email.py diff --git a/tests/enterprise/litellm_enterprise/enterprise_callbacks/test_prometheus_logging_callbacks.py b/tests/unit/enterprise/enterprise_callbacks/test_prometheus_logging_callbacks.py similarity index 100% rename from tests/enterprise/litellm_enterprise/enterprise_callbacks/test_prometheus_logging_callbacks.py rename to tests/unit/enterprise/enterprise_callbacks/test_prometheus_logging_callbacks.py diff --git a/tests/unit/enterprise/integrations/__init__.py b/tests/unit/enterprise/integrations/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/enterprise/litellm_enterprise/integrations/test_custom_guardrail.py b/tests/unit/enterprise/integrations/test_custom_guardrail.py similarity index 100% rename from tests/enterprise/litellm_enterprise/integrations/test_custom_guardrail.py rename to tests/unit/enterprise/integrations/test_custom_guardrail.py diff --git a/tests/enterprise/litellm_enterprise/integrations/test_prometheus.py b/tests/unit/enterprise/integrations/test_prometheus.py similarity index 100% rename from tests/enterprise/litellm_enterprise/integrations/test_prometheus.py rename to tests/unit/enterprise/integrations/test_prometheus.py diff --git a/tests/enterprise/litellm_enterprise/integrations/test_prometheus_unit_tests.py b/tests/unit/enterprise/integrations/test_prometheus_unit_tests.py similarity index 100% rename from tests/enterprise/litellm_enterprise/integrations/test_prometheus_unit_tests.py rename to tests/unit/enterprise/integrations/test_prometheus_unit_tests.py diff --git a/tests/unit/enterprise/proxy/__init__.py b/tests/unit/enterprise/proxy/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/enterprise/proxy/auth/__init__.py b/tests/unit/enterprise/proxy/auth/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/enterprise/litellm_enterprise/proxy/auth/test_route_checks.py b/tests/unit/enterprise/proxy/auth/test_route_checks.py similarity index 100% rename from tests/enterprise/litellm_enterprise/proxy/auth/test_route_checks.py rename to tests/unit/enterprise/proxy/auth/test_route_checks.py diff --git a/tests/enterprise/litellm_enterprise/proxy/auth/test_user_api_key_auth.py b/tests/unit/enterprise/proxy/auth/test_user_api_key_auth.py similarity index 100% rename from tests/enterprise/litellm_enterprise/proxy/auth/test_user_api_key_auth.py rename to tests/unit/enterprise/proxy/auth/test_user_api_key_auth.py diff --git a/tests/unit/enterprise/proxy/guardrails/__init__.py b/tests/unit/enterprise/proxy/guardrails/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/enterprise/litellm_enterprise/proxy/guardrails/conftest.py b/tests/unit/enterprise/proxy/guardrails/conftest.py similarity index 100% rename from tests/enterprise/litellm_enterprise/proxy/guardrails/conftest.py rename to tests/unit/enterprise/proxy/guardrails/conftest.py diff --git a/tests/enterprise/litellm_enterprise/proxy/guardrails/test_apply_guardrail_endpoint.py b/tests/unit/enterprise/proxy/guardrails/test_apply_guardrail_endpoint.py similarity index 100% rename from tests/enterprise/litellm_enterprise/proxy/guardrails/test_apply_guardrail_endpoint.py rename to tests/unit/enterprise/proxy/guardrails/test_apply_guardrail_endpoint.py diff --git a/tests/enterprise/litellm_enterprise/proxy/guardrails/test_bedrock_apply_guardrail.py b/tests/unit/enterprise/proxy/guardrails/test_bedrock_apply_guardrail.py similarity index 100% rename from tests/enterprise/litellm_enterprise/proxy/guardrails/test_bedrock_apply_guardrail.py rename to tests/unit/enterprise/proxy/guardrails/test_bedrock_apply_guardrail.py diff --git a/tests/unit/enterprise/proxy/hooks/__init__.py b/tests/unit/enterprise/proxy/hooks/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/enterprise/litellm_enterprise/proxy/hooks/test_managed_files.py b/tests/unit/enterprise/proxy/hooks/test_managed_files.py similarity index 100% rename from tests/enterprise/litellm_enterprise/proxy/hooks/test_managed_files.py rename to tests/unit/enterprise/proxy/hooks/test_managed_files.py diff --git a/tests/unit/enterprise/proxy/management_endpoints/__init__.py b/tests/unit/enterprise/proxy/management_endpoints/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/enterprise/litellm_enterprise/proxy/management_endpoints/test_internal_user_endpoints.py b/tests/unit/enterprise/proxy/management_endpoints/test_internal_user_endpoints.py similarity index 100% rename from tests/enterprise/litellm_enterprise/proxy/management_endpoints/test_internal_user_endpoints.py rename to tests/unit/enterprise/proxy/management_endpoints/test_internal_user_endpoints.py diff --git a/tests/enterprise/litellm_enterprise/proxy/management_endpoints/test_project_endpoints_prisma.py b/tests/unit/enterprise/proxy/management_endpoints/test_project_endpoints_prisma.py similarity index 100% rename from tests/enterprise/litellm_enterprise/proxy/management_endpoints/test_project_endpoints_prisma.py rename to tests/unit/enterprise/proxy/management_endpoints/test_project_endpoints_prisma.py diff --git a/tests/test_litellm/enterprise/proxy/test_afile_retrieve_returns_unified_id.py b/tests/unit/enterprise/proxy/test_afile_retrieve_returns_unified_id.py similarity index 100% rename from tests/test_litellm/enterprise/proxy/test_afile_retrieve_returns_unified_id.py rename to tests/unit/enterprise/proxy/test_afile_retrieve_returns_unified_id.py diff --git a/tests/enterprise/litellm_enterprise/proxy/test_audit_logging_endpoints.py b/tests/unit/enterprise/proxy/test_audit_logging_endpoints.py similarity index 100% rename from tests/enterprise/litellm_enterprise/proxy/test_audit_logging_endpoints.py rename to tests/unit/enterprise/proxy/test_audit_logging_endpoints.py diff --git a/tests/test_litellm/enterprise/proxy/test_batch_retrieve_input_file_id.py b/tests/unit/enterprise/proxy/test_batch_retrieve_input_file_id.py similarity index 100% rename from tests/test_litellm/enterprise/proxy/test_batch_retrieve_input_file_id.py rename to tests/unit/enterprise/proxy/test_batch_retrieve_input_file_id.py diff --git a/tests/test_litellm/enterprise/proxy/test_batch_retrieve_registers_missing_output_file_id.py b/tests/unit/enterprise/proxy/test_batch_retrieve_registers_missing_output_file_id.py similarity index 100% rename from tests/test_litellm/enterprise/proxy/test_batch_retrieve_registers_missing_output_file_id.py rename to tests/unit/enterprise/proxy/test_batch_retrieve_registers_missing_output_file_id.py diff --git a/tests/test_litellm/enterprise/proxy/test_batch_retrieve_returns_unified_input_file_id.py b/tests/unit/enterprise/proxy/test_batch_retrieve_returns_unified_input_file_id.py similarity index 100% rename from tests/test_litellm/enterprise/proxy/test_batch_retrieve_returns_unified_input_file_id.py rename to tests/unit/enterprise/proxy/test_batch_retrieve_returns_unified_input_file_id.py diff --git a/tests/test_litellm/enterprise/proxy/test_batch_update_db_managed_output_file_id.py b/tests/unit/enterprise/proxy/test_batch_update_db_managed_output_file_id.py similarity index 100% rename from tests/test_litellm/enterprise/proxy/test_batch_update_db_managed_output_file_id.py rename to tests/unit/enterprise/proxy/test_batch_update_db_managed_output_file_id.py diff --git a/tests/test_litellm/enterprise/proxy/test_deleted_file_returns_403_not_404.py b/tests/unit/enterprise/proxy/test_deleted_file_returns_403_not_404.py similarity index 100% rename from tests/test_litellm/enterprise/proxy/test_deleted_file_returns_403_not_404.py rename to tests/unit/enterprise/proxy/test_deleted_file_returns_403_not_404.py diff --git a/tests/test_litellm/enterprise/proxy/test_enterprise_routes.py b/tests/unit/enterprise/proxy/test_enterprise_routes.py similarity index 100% rename from tests/test_litellm/enterprise/proxy/test_enterprise_routes.py rename to tests/unit/enterprise/proxy/test_enterprise_routes.py diff --git a/tests/test_litellm/enterprise/proxy/test_file_deletion_blocking.py b/tests/unit/enterprise/proxy/test_file_deletion_blocking.py similarity index 100% rename from tests/test_litellm/enterprise/proxy/test_file_deletion_blocking.py rename to tests/unit/enterprise/proxy/test_file_deletion_blocking.py diff --git a/tests/test_litellm/enterprise/proxy/test_managed_files_access_check.py b/tests/unit/enterprise/proxy/test_managed_files_access_check.py similarity index 100% rename from tests/test_litellm/enterprise/proxy/test_managed_files_access_check.py rename to tests/unit/enterprise/proxy/test_managed_files_access_check.py diff --git a/tests/test_litellm/enterprise/proxy/test_managed_files_hook.py b/tests/unit/enterprise/proxy/test_managed_files_hook.py similarity index 100% rename from tests/test_litellm/enterprise/proxy/test_managed_files_hook.py rename to tests/unit/enterprise/proxy/test_managed_files_hook.py diff --git a/tests/unit/gateway/__init__.py b/tests/unit/gateway/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_gateway/test_launch.py b/tests/unit/gateway/test_launch.py similarity index 100% rename from tests/test_gateway/test_launch.py rename to tests/unit/gateway/test_launch.py diff --git a/tests/unit/litellm_proxy_extras/__init__.py b/tests/unit/litellm_proxy_extras/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/litellm-proxy-extras/test_litellm_proxy_extras_logging.py b/tests/unit/litellm_proxy_extras/test_litellm_proxy_extras_logging.py similarity index 100% rename from tests/litellm-proxy-extras/test_litellm_proxy_extras_logging.py rename to tests/unit/litellm_proxy_extras/test_litellm_proxy_extras_logging.py diff --git a/tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py b/tests/unit/litellm_proxy_extras/test_litellm_proxy_extras_utils.py similarity index 99% rename from tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py rename to tests/unit/litellm_proxy_extras/test_litellm_proxy_extras_utils.py index bb329264a11..755c7617701 100644 --- a/tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py +++ b/tests/unit/litellm_proxy_extras/test_litellm_proxy_extras_utils.py @@ -9,7 +9,7 @@ import pytest sys.path.insert( 0, os.path.abspath( - os.path.join(os.path.dirname(__file__), "../../litellm-proxy-extras") + os.path.join(os.path.dirname(__file__), "../../../litellm-proxy-extras") ), ) @@ -23,7 +23,7 @@ from litellm_proxy_extras.utils import ( _MIGRATIONS_DIR = os.path.abspath( os.path.join( os.path.dirname(__file__), - "../../litellm-proxy-extras/litellm_proxy_extras/migrations", + "../../../litellm-proxy-extras/litellm_proxy_extras/migrations", ) ) @@ -999,7 +999,7 @@ class TestJWTKeyMappingCascade: schema_paths = glob.glob( os.path.abspath( os.path.join( - os.path.dirname(__file__), "../../**/schema.prisma" + os.path.dirname(__file__), "../../../**/schema.prisma" ) ), recursive=True, From 7b25a151bd72dbed46137a89799b298d6da4d87e Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 24 Sep 2026 15:52:44 -0700 Subject: [PATCH 094/120] feat(proxy): let callbacks filter the model listing routes per caller (#43027) * feat(proxy): let callbacks filter the model listing routes per caller * fix(proxy): offer every listed name to the listing callback, agent groups and deployment lookups included * fix(proxy): hide aliases of a team model by its public name and offer /model/info lookups the listed name * fix(proxy): map a malformed model listing filter return to the proxy error contract and document legacy team names --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- litellm/integrations/custom_logger.py | 18 + litellm/proxy/proxy_server.py | 101 ++++- litellm/proxy/utils.py | 62 ++- .../proxy_server/test_routes_model_info.py | 13 +- .../proxy/test_model_list_callback_filter.py | 425 ++++++++++++++++++ 5 files changed, 596 insertions(+), 23 deletions(-) create mode 100644 tests/test_litellm/proxy/test_model_list_callback_filter.py diff --git a/litellm/integrations/custom_logger.py b/litellm/integrations/custom_logger.py index 326abd5c6a3..d4162369a35 100644 --- a/litellm/integrations/custom_logger.py +++ b/litellm/integrations/custom_logger.py @@ -421,6 +421,24 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac ): # raise exception if invalid, return a str for the user to receive - if rejected, or return a modified dictionary for passing into litellm pass + async def async_filter_listed_models( + self, + user_api_key_dict: UserAPIKeyAuth, + model_names: Sequence[str], + ) -> Sequence[str]: + """Runs on the model listing routes (`/v1/models`, `/v1/models/{id}`, `/model/info`, + `/model_group/info`) with the public model names the route would otherwise return, so a + lookup of one model may offer just that name: decide per name, never by position in the + sequence. Return the names to keep as a sequence of strings; a name left out disappears + from every listing, any alias of it offered in the same call goes with it, and + `/v1/models/{id}` answers 404 for it, exactly as for a model that does not exist. Names + outside `model_names` are ignored, so a callback can only narrow the listing, never widen + it. Under `use_team_public_model_name: false`, `/v1/models` and `/model_group/info` list a + team model by its internal routing name while `/model/info` keeps its public name, so hide + both names to hide it on every route. + """ + return model_names + async def async_post_call_response_headers_hook( self, data: dict, diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index c0071aa7c81..a869150f7e8 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -152,6 +152,7 @@ from litellm.router_utils.auto_router_tuning_baseline import ( snapshot_tuning_baselines, tuning_limit_violation, ) +from litellm.router_utils.common_utils import resolve_model_group_alias from litellm.router_utils.routing_groups import parse_routing_groups from litellm.types.caching import RedisPipelineIncrementOperation from litellm.types.utils import ( @@ -11104,6 +11105,40 @@ class ProxyStartupEvent: #### API ENDPOINTS #### +async def _names_hidden_by_listing_callbacks( + user_api_key_dict: UserAPIKeyAuth, model_names: Sequence[str] +) -> frozenset[str]: + hidden: Final = await proxy_logging_obj.hidden_by_listing_callbacks(user_api_key_dict, model_names) + if not hidden or llm_router is None: + return hidden + aliases: Final = llm_router.model_group_alias + internal_to_public: Final = TeamModelNameTranslator.build_internal_to_public_map(llm_router, general_settings) + return hidden | frozenset( + alias + for alias in aliases + if (target := resolve_model_group_alias(aliases, alias)) is not None + and internal_to_public.get(target, target) in hidden + ) + + +async def _entries_kept_by_listing_callbacks( + entries: Sequence[tuple[str, str]], user_api_key_dict: UserAPIKeyAuth +) -> tuple[tuple[str, str], ...]: + hidden: Final = await _names_hidden_by_listing_callbacks( + user_api_key_dict, tuple(response_id for response_id, _ in entries) + ) + if not hidden: + return tuple(entries) + return tuple(entry for entry in entries if entry[0] not in hidden) + + +async def _deployment_hidden_by_listing_callbacks(deployment: Deployment, user_api_key_dict: UserAPIKeyAuth) -> bool: + listed_name: Final = _translate_model_name_for_response(deployment.model_dump(exclude_none=True)).get("model_name") + if not isinstance(listed_name, str): + return False + return listed_name in await _names_hidden_by_listing_callbacks(user_api_key_dict, (listed_name,)) + + @router.get("/v1/models", dependencies=[Depends(user_api_key_auth)], tags=["model management"]) @router.get( "/models", dependencies=[Depends(user_api_key_auth)], tags=["model management"] @@ -11254,7 +11289,9 @@ async def model_list( # The internal routing key drives the metadata/fallback lookup, while the # public name is what the client sees as the model id. model_data = [] - admin_entries: Final = TeamModelNameTranslator.listing_entries(all_models, llm_router, settings) + admin_entries: Final = await _entries_kept_by_listing_callbacks( + TeamModelNameTranslator.listing_entries(all_models, llm_router, settings), user_api_key_dict + ) for response_id, lookup_id in admin_entries: model_info = create_model_info_response( model_id=lookup_id, @@ -11310,7 +11347,10 @@ async def model_list( # public name is what the client sees as the model id. model_data = [] entries: Final = alias_listing_entries( - TeamModelNameTranslator.listing_entries(all_models, llm_router, settings), caller_aliases + await _entries_kept_by_listing_callbacks( + TeamModelNameTranslator.listing_entries(all_models, llm_router, settings), user_api_key_dict + ), + caller_aliases, ) for response_id, lookup_id in entries: model_info = create_model_info_response( @@ -11404,13 +11444,24 @@ async def model_info( llm_router=llm_router, ) hidden_names: Final = blocked_names | unhealthy_names - if hidden_names: - all_models = [m for m in all_models if m not in hidden_names] + internal_to_public: Final = TeamModelNameTranslator.build_internal_to_public_map(llm_router, settings) + callback_hidden_names: Final = await _names_hidden_by_listing_callbacks( + user_api_key_dict, + tuple( + response_id + for response_id, _ in TeamModelNameTranslator.listing_entries( + tuple(m for m in all_models if m not in hidden_names), llm_router, settings + ) + ), + ) + if hidden_names or callback_hidden_names: + all_models = [ + m for m in all_models if m not in hidden_names and internal_to_public.get(m, m) not in callback_hidden_names + ] undiscoverable_names: Final = undiscoverable_model_names( all_models, llm_router, user_api_key_dict, team_id or user_api_key_dict.team_id ) - internal_to_public: Final = TeamModelNameTranslator.build_internal_to_public_map(llm_router, settings) aliased_model_id: Final = alias_target( model_id, caller_alias_maps( @@ -15730,7 +15781,7 @@ async def model_info_v1( if litellm_model_id is not None: # user is trying to get specific model from litellm router deployment_info: Final = llm_router.get_deployment(model_id=litellm_model_id) - if deployment_info is None: + if deployment_info is None or await _deployment_hidden_by_listing_callbacks(deployment_info, user_api_key_dict): raise HTTPException( status_code=400, detail={"error": f"Model id = {litellm_model_id} not found on litellm proxy"}, @@ -15819,10 +15870,17 @@ async def model_info_v1( general_settings=general_settings, llm_router=llm_router, ) - visible_models: Final = discoverable_rows( + servable_rows: Final = discoverable_rows( (model for model in all_models if model.get("model_name") not in hidden_names), user_api_key_dict, ) + listed_names: Final = tuple( + dict.fromkeys(name for model in servable_rows if isinstance(name := model.get("model_name"), str)) + ) + callback_hidden_names: Final = await _names_hidden_by_listing_callbacks(user_api_key_dict, listed_names) + visible_models: Final = tuple( + model for model in servable_rows if model.get("model_name") not in callback_hidden_names + ) verbose_proxy_logger.debug("all_models: %s", visible_models) return _model_info_json_response(visible_models) @@ -15871,7 +15929,7 @@ async def model_deprecations( def _get_model_group_info( - llm_router: Router, all_models_str: list[str], model_group: str | None + llm_router: Router, all_models_str: Sequence[str], model_group: str | None ) -> list[ModelGroupInfoProxy]: model_groups: Final[list[ModelGroupInfoProxy]] = [] @@ -16104,23 +16162,34 @@ async def model_group_info( undiscoverable_group_names: Final = undiscoverable_model_names( all_models_str, llm_router, user_api_key_dict, user_api_key_dict.team_id ) - model_groups: list[ModelGroupInfoProxy] = _get_model_group_info( - llm_router=llm_router, - all_models_str=[name for name in all_models_str if name not in undiscoverable_group_names], - model_group=model_group, - ) + listed_group_names: Final = tuple(name for name in all_models_str if name not in undiscoverable_group_names) # Append A2A agents to model groups from litellm.proxy.agent_endpoints.model_list_helpers import ( append_agents_to_model_group, ) - model_groups = await append_agents_to_model_group( - model_groups=model_groups, + model_groups: Final = await append_agents_to_model_group( + model_groups=_get_model_group_info( + llm_router=llm_router, all_models_str=listed_group_names, model_group=model_group + ), user_api_key_dict=user_api_key_dict, ) + internal_to_public: Final = TeamModelNameTranslator.build_internal_to_public_map(llm_router, general_settings) + public_group_names: Final = tuple( + internal_to_public.get(group.model_group, group.model_group) for group in model_groups + ) + callback_hidden_names: Final = await _names_hidden_by_listing_callbacks( + user_api_key_dict, tuple(dict.fromkeys(public_group_names)) + ) - return {"data": model_groups} + return { + "data": [ + group + for group, public_name in zip(model_groups, public_group_names, strict=True) + if public_name not in callback_hidden_names + ] + } @router.get( diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index c617047fad9..b8cc30ad8a7 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -36,6 +36,7 @@ from typing import ( Final, Generic, Literal, + NoReturn, Optional, Protocol, TypeAlias, @@ -106,7 +107,7 @@ except ImportError: raise ImportError("backoff is not installed. Please install it via 'pip install backoff'") from fastapi import HTTPException, status -from pydantic import TypeAdapter +from pydantic import TypeAdapter, ValidationError import litellm import litellm.litellm_core_utils @@ -1136,11 +1137,51 @@ class _CallbackCapabilities: # avoids the per-request ``get_custom_logger_compatible_class`` walk for # every string entry in ``litellm.callbacks``. resolved_callbacks: tuple[object, ...] = field(default_factory=tuple) + listed_models_filters: tuple[CustomLogger, ...] = field(default_factory=tuple) + + +def _overrides_hook(callback: CustomLogger, hook_name: str) -> bool: + leaf_to_base: Final = takewhile(lambda klass: klass is not CustomLogger, type(callback).__mro__) + return any(hook_name in klass.__dict__ for klass in leaf_to_base) def _overrides_moderation_hook(callback: CustomLogger) -> bool: - leaf_to_base: Final = takewhile(lambda klass: klass is not CustomLogger, type(callback).__mro__) - return any("async_moderation_hook" in klass.__dict__ for klass in leaf_to_base) + return _overrides_hook(callback, "async_moderation_hook") + + +_LISTED_MODEL_NAMES: Final = TypeAdapter(tuple[str, ...]) + + +@dataclass(frozen=True, slots=True) +class MalformedListingFilterReturn: + callback: str + tag: Literal["malformed_listing_filter_return"] = "malformed_listing_filter_return" + + +async def _names_kept_by_listing_callbacks( + callbacks: Sequence[CustomLogger], + user_api_key_dict: UserAPIKeyAuth, + model_names: tuple[str, ...], +) -> tuple[str, ...] | MalformedListingFilterReturn: + if not callbacks or not model_names: + return model_names + returned: Final = await callbacks[0].async_filter_listed_models(user_api_key_dict, model_names) + try: + kept: Final = frozenset(_LISTED_MODEL_NAMES.validate_python(returned)) + except ValidationError: + return MalformedListingFilterReturn(callback=type(callbacks[0]).__name__) + return await _names_kept_by_listing_callbacks( + callbacks[1:], user_api_key_dict, tuple(name for name in model_names if name in kept) + ) + + +def _raise_malformed_listing_filter_return(error: MalformedListingFilterReturn) -> NoReturn: + raise ProxyException( + message=f"{error.callback}.async_filter_listed_models must return a sequence of model names", + type=ProxyErrorTypes.internal_server_error, + param=None, + code=500, + ) class ProxyLogging: @@ -2808,6 +2849,9 @@ class ProxyLogging: has_moderation_override=has_moderation_override, iterator_overrides=tuple(iterator_overrides), resolved_callbacks=tuple(resolved_callbacks), + listed_models_filters=tuple( + callback for callback in resolved_callbacks if _overrides_hook(callback, "async_filter_listed_models") + ), ) # Limit cache to handle test churn without leaking; production # callback lists are stable so this rarely grows past 1 entry. @@ -3715,6 +3759,18 @@ class ProxyLogging: verbose_proxy_logger.exception("Error in post_call_response_headers_hook: %s", str(e)) return merged_headers + async def hidden_by_listing_callbacks( + self, user_api_key_dict: UserAPIKeyAuth, model_names: Sequence[str] + ) -> frozenset[str]: + filters: Final = ProxyLogging._callback_capabilities().listed_models_filters + if not filters: + return frozenset() + candidates: Final = tuple(model_names) + kept: Final = await _names_kept_by_listing_callbacks(filters, user_api_key_dict, candidates) + if isinstance(kept, MalformedListingFilterReturn): + _raise_malformed_listing_filter_return(kept) + return frozenset(candidates).difference(kept) + @staticmethod def _build_litellm_call_info(data: dict, response: object) -> dict[str, object]: """ diff --git a/tests/test_litellm/proxy/proxy_server/test_routes_model_info.py b/tests/test_litellm/proxy/proxy_server/test_routes_model_info.py index 636dc0f4d77..5175d92084c 100644 --- a/tests/test_litellm/proxy/proxy_server/test_routes_model_info.py +++ b/tests/test_litellm/proxy/proxy_server/test_routes_model_info.py @@ -322,7 +322,9 @@ def test_get_proxy_model_info_shows_litellm_params_pricing_and_names_it_as_an_ov def test_get_proxy_model_info_names_config_model_info_pricing_as_an_override(monkeypatch, local_model_cost_map): """Pricing declared under ``model_info`` in config.yaml overrides the cost map too.""" info = _enriched_model_info( - monkeypatch, {"model": "openai/gpt-5.6"}, {"id": "dep-config", "db_model": False, "output_cost_per_token": 7e-06} + monkeypatch, + {"model": "openai/gpt-5.6"}, + {"id": "dep-config", "db_model": False, "output_cost_per_token": 7e-06}, ) assert info["pricing_overrides"] == ("output_cost_per_token",) assert info["output_cost_per_token"] == 7e-06 @@ -399,7 +401,9 @@ def test_model_info_reports_null_cost_for_unpriced_deployment_and_zero_for_decla def enriched_cost(model_name: str) -> tuple: deployment = router.get_model_list(model_name=model_name)[0] - info = proxy_server._enrich_model_info_with_litellm_data({**deployment, "model_info": dict(deployment["model_info"])})["model_info"] + info = proxy_server._enrich_model_info_with_litellm_data( + {**deployment, "model_info": dict(deployment["model_info"])} + )["model_info"] return info.get("input_cost_per_token"), info.get("output_cost_per_token") assert enriched_cost("vllm-unpriced") == (None, None) @@ -643,7 +647,6 @@ def model_group_info_router(monkeypatch): monkeypatch.setattr(proxy_server, "user_model", None) monkeypatch.setattr(proxy_server, "general_settings", {}) monkeypatch.setattr(proxy_server, "prisma_client", None) - monkeypatch.setattr(proxy_server, "proxy_logging_obj", None) monkeypatch.setattr(proxy_server, "user_api_key_cache", None) monkeypatch.setattr(proxy_server, "_get_model_group_info", model_group_info) @@ -671,7 +674,9 @@ def test_model_group_info_proxy_admin_ignores_key_model_restriction( @pytest.mark.parametrize("admin_role", ["proxy_admin", "proxy_admin_viewer"]) -def test_model_group_info_proxy_admin_expands_wildcard_deployments(client, auth_as, model_group_info_router, admin_role): +def test_model_group_info_proxy_admin_expands_wildcard_deployments( + client, auth_as, model_group_info_router, admin_role +): from litellm.proxy._types import LitellmUserRoles from litellm.proxy.auth.model_checks import get_known_models_from_wildcard diff --git a/tests/test_litellm/proxy/test_model_list_callback_filter.py b/tests/test_litellm/proxy/test_model_list_callback_filter.py new file mode 100644 index 00000000000..00fbfee24ed --- /dev/null +++ b/tests/test_litellm/proxy/test_model_list_callback_filter.py @@ -0,0 +1,425 @@ +""" +Tests for `CustomLogger.async_filter_listed_models` on the model listing endpoints: +GET /v1/models (`model_list`, OpenAI and Anthropic shapes), GET /v1/models/{id} +(`model_info`), GET /v1/model/info (`model_info_v1`) and GET /model_group/info +(`model_group_info`). A registered callback that overrides the hook decides per +caller which of the names the route would list are kept; the rest disappear and +`/v1/models/{id}` answers 404 for them. +""" + +import json +from collections.abc import Sequence + +import pytest +from fastapi import HTTPException +from starlette.requests import Request + +import litellm +from litellm import Router +from litellm.integrations.custom_logger import CustomLogger +from litellm.proxy import proxy_server +from litellm.proxy._types import LitellmUserRoles, ProxyException, UserAPIKeyAuth +from litellm.proxy.utils import ProxyLogging + + +class _Gate(CustomLogger): + def __init__(self, hidden: frozenset[str] = frozenset(), extra: tuple[str, ...] = ()) -> None: + super().__init__() + self.hidden = hidden + self.extra = extra + self.seen: list[tuple[str, ...]] = [] + + async def async_filter_listed_models( + self, user_api_key_dict: UserAPIKeyAuth, model_names: Sequence[str] + ) -> Sequence[str]: + self.seen.append(tuple(model_names)) + return [*(name for name in model_names if name not in self.hidden), *self.extra] + + +class _InferenceOnlyGate(CustomLogger): + async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type): + if data.get("model") == "restricted-model": + raise HTTPException(status_code=403, detail="not entitled to this model") + return data + + +class _RaisingGate(CustomLogger): + async def async_filter_listed_models( + self, user_api_key_dict: UserAPIKeyAuth, model_names: Sequence[str] + ) -> Sequence[str]: + raise HTTPException(status_code=503, detail="entitlement service down") + + +class _ReversingGate(CustomLogger): + async def async_filter_listed_models( + self, user_api_key_dict: UserAPIKeyAuth, model_names: Sequence[str] + ) -> Sequence[str]: + return list(reversed(model_names)) + + +class _StringReturningGate(CustomLogger): + async def async_filter_listed_models(self, user_api_key_dict: UserAPIKeyAuth, model_names: Sequence[str]) -> str: + return "open-model" + + +def _deployment(model_name: str, model: str = "openai/gpt-4o", **model_info): + return { + "model_name": model_name, + "litellm_params": {"model": model, "api_key": "sk-fake"}, + "model_info": {"id": f"{model_name}-id", **model_info}, + } + + +def _install_router(monkeypatch, *deployments, **router_kwargs) -> Router: + router = Router(model_list=list(deployments), **router_kwargs) + monkeypatch.setattr(proxy_server, "llm_router", router) + monkeypatch.setattr(proxy_server, "llm_model_list", router.model_list) + monkeypatch.setattr(proxy_server, "prisma_client", None) + monkeypatch.setattr(proxy_server, "general_settings", {}) + monkeypatch.setattr(proxy_server, "user_model", None) + return router + + +def _register(monkeypatch, *callbacks: CustomLogger) -> None: + monkeypatch.setattr(litellm, "callbacks", list(callbacks)) + ProxyLogging._callback_capabilities_cache.clear() + + +@pytest.fixture +def two_model_router(monkeypatch) -> Router: + return _install_router(monkeypatch, _deployment("open-model"), _deployment("restricted-model")) + + +@pytest.fixture +def team_router(monkeypatch) -> Router: + return _install_router( + monkeypatch, + _deployment("gpt-4"), + _deployment("model_name_team1_abc", team_id="team1", team_public_model_name="team-gpt"), + _deployment("model_name_team1_def", team_id="team1", team_public_model_name="team-chat"), + ) + + +@pytest.fixture +def team_admin_privileges(monkeypatch) -> None: + from litellm.proxy.management_endpoints import common_utils + + async def _is_team_admin(**kwargs) -> bool: + return True + + monkeypatch.setattr(common_utils, "_user_has_admin_privileges", _is_team_admin) + + +def _non_admin(**kwargs) -> UserAPIKeyAuth: + return UserAPIKeyAuth(api_key="sk-test", user_role=LitellmUserRoles.INTERNAL_USER, **kwargs) + + +def _admin() -> UserAPIKeyAuth: + return UserAPIKeyAuth(api_key="sk-test", user_id="u", user_role=LitellmUserRoles.PROXY_ADMIN, team_models=[]) + + +def _team_member() -> UserAPIKeyAuth: + return UserAPIKeyAuth( + api_key="sk-test", + user_id="u", + user_role=LitellmUserRoles.INTERNAL_USER, + team_id="team1", + team_models=["model_name_team1_abc", "model_name_team1_def"], + models=["model_name_team1_abc", "model_name_team1_def"], + ) + + +def _anthropic_request() -> Request: + return Request( + scope={ + "type": "http", + "method": "GET", + "path": "/v1/models", + "query_string": b"", + "headers": [(b"anthropic-version", b"2023-06-01")], + } + ) + + +async def _v1_models(user_api_key_dict: UserAPIKeyAuth, **kwargs) -> list[str]: + response = await proxy_server.model_list(user_api_key_dict=user_api_key_dict, **kwargs) + return [m["id"] for m in response["data"]] + + +async def _v1_model_info_names(user_api_key_dict: UserAPIKeyAuth) -> list[str]: + response = await proxy_server.model_info_v1(user_api_key_dict=user_api_key_dict) + return [row["model_name"] for row in json.loads(response.body)["data"]] + + +async def _model_groups(user_api_key_dict: UserAPIKeyAuth) -> list[str]: + response = await proxy_server.model_group_info(user_api_key_dict=user_api_key_dict) + return [group.model_group for group in response["data"]] + + +async def _model_by_id_status(model_id: str, user_api_key_dict: UserAPIKeyAuth) -> int: + try: + response = await proxy_server.model_info(model_id=model_id, user_api_key_dict=user_api_key_dict) + except HTTPException as error: + return error.status_code + assert response["id"] == model_id + return 200 + + +@pytest.mark.asyncio +async def test_v1_models_lists_only_the_names_the_callback_keeps(two_model_router, monkeypatch): + _register(monkeypatch, _Gate(hidden=frozenset({"restricted-model"}))) + + assert await _v1_models(_non_admin()) == ["open-model"] + assert await _v1_models(_admin()) == ["open-model"] + assert await _v1_models(_non_admin(), request=_anthropic_request()) == ["open-model"] + + +@pytest.mark.asyncio +async def test_v1_models_scope_expand_applies_the_callback(two_model_router, team_admin_privileges, monkeypatch): + _register(monkeypatch, _Gate(hidden=frozenset({"restricted-model"}))) + + assert await _v1_models(_non_admin(), scope="expand") == ["open-model"] + assert await _v1_models(_admin(), scope="expand") == ["open-model"] + + +@pytest.mark.asyncio +async def test_v1_models_by_id_answers_404_for_a_name_the_callback_leaves_out(two_model_router, monkeypatch): + _register(monkeypatch, _Gate(hidden=frozenset({"restricted-model"}))) + + assert await _model_by_id_status("restricted-model", _non_admin()) == 404 + assert await _model_by_id_status("open-model", _non_admin()) == 200 + + +@pytest.mark.asyncio +async def test_v1_model_info_lists_only_the_rows_the_callback_keeps(two_model_router, monkeypatch): + _register(monkeypatch, _Gate(hidden=frozenset({"restricted-model"}))) + + assert await _v1_model_info_names(_non_admin()) == ["open-model"] + assert await _v1_model_info_names(_admin()) == ["open-model"] + + +@pytest.mark.asyncio +async def test_model_group_info_lists_only_the_groups_the_callback_keeps(two_model_router, monkeypatch): + _register(monkeypatch, _Gate(hidden=frozenset({"restricted-model"}))) + + assert await _model_groups(_non_admin()) == ["open-model"] + assert await _model_groups(_admin()) == ["open-model"] + + +@pytest.mark.asyncio +async def test_a_callback_without_the_hook_changes_no_listing(two_model_router, monkeypatch): + _register(monkeypatch, _InferenceOnlyGate()) + + assert await _v1_models(_non_admin()) == ["open-model", "restricted-model"] + assert await _model_by_id_status("restricted-model", _non_admin()) == 200 + assert await _v1_model_info_names(_non_admin()) == ["open-model", "restricted-model"] + assert await _model_groups(_non_admin()) == ["open-model", "restricted-model"] + + +@pytest.mark.asyncio +async def test_callback_cannot_add_a_name_it_was_not_offered(two_model_router, monkeypatch): + _register(monkeypatch, _Gate(extra=("ghost-model",))) + + assert await _v1_models(_non_admin()) == ["open-model", "restricted-model"] + assert await _model_by_id_status("ghost-model", _non_admin()) == 404 + + +@pytest.mark.asyncio +async def test_callbacks_narrow_in_registration_order(monkeypatch): + _install_router(monkeypatch, _deployment("a"), _deployment("b"), _deployment("c")) + first: _Gate = _Gate(hidden=frozenset({"a"})) + second: _Gate = _Gate(hidden=frozenset({"b"})) + _register(monkeypatch, first, second) + + assert await _v1_models(_non_admin()) == ["c"] + assert first.seen == [("a", "b", "c")] + assert second.seen == [("b", "c")] + + +@pytest.mark.asyncio +async def test_callback_sees_and_filters_team_models_by_their_public_name(team_router, monkeypatch): + gate: _Gate = _Gate(hidden=frozenset({"team-gpt"})) + _register(monkeypatch, gate) + + assert await _v1_models(_team_member()) == ["team-chat"] + assert await _model_by_id_status("team-gpt", _team_member()) == 404 + assert await _model_by_id_status("team-chat", _team_member()) == 200 + assert all("team-gpt" in seen and "model_name_team1_abc" not in seen for seen in gate.seen) + + +@pytest.mark.asyncio +async def test_callback_sees_public_team_names_on_every_listing_route(team_router, monkeypatch): + gate: _Gate = _Gate(hidden=frozenset({"team-gpt"})) + _register(monkeypatch, gate) + + assert await _v1_models(_team_member()) == ["team-chat"] + assert await _v1_model_info_names(_team_member()) == ["team-chat"] + assert await _model_groups(_team_member()) == ["model_name_team1_def"] + assert await _model_by_id_status("team-gpt", _team_member()) == 404 + assert len(gate.seen) == 4 + assert all(sorted(seen) == ["team-chat", "team-gpt"] for seen in gate.seen) + + +@pytest.mark.asyncio +async def test_router_alias_follows_its_hidden_target(monkeypatch): + _install_router( + monkeypatch, + _deployment("open-model"), + _deployment("restricted-model"), + model_group_alias={"mini": "restricted-model", "wide": "open-model"}, + ) + _register(monkeypatch, _Gate(hidden=frozenset({"restricted-model"}))) + + assert sorted(await _v1_models(_non_admin())) == ["open-model", "wide"] + assert sorted(await _v1_model_info_names(_non_admin())) == ["open-model", "wide"] + assert sorted(await _model_groups(_non_admin())) == ["open-model", "wide"] + + _register(monkeypatch, _Gate(hidden=frozenset({"mini"}))) + + assert sorted(await _v1_models(_non_admin())) == ["open-model", "restricted-model", "wide"] + + +@pytest.mark.asyncio +async def test_router_alias_of_a_team_model_follows_its_hidden_public_name(monkeypatch): + _install_router( + monkeypatch, + _deployment("gpt-4"), + _deployment("model_name_team1_abc", team_id="team1", team_public_model_name="team-gpt"), + model_group_alias={"team-alias": "model_name_team1_abc"}, + ) + caller: UserAPIKeyAuth = _non_admin( + user_id="u", + team_id="team1", + team_models=["model_name_team1_abc", "team-alias"], + models=["model_name_team1_abc", "team-alias"], + ) + _register(monkeypatch, _Gate(hidden=frozenset())) + assert sorted(await _v1_models(caller)) == ["team-alias", "team-gpt"] + + _register(monkeypatch, _Gate(hidden=frozenset({"team-gpt"}))) + assert await _v1_models(caller) == [] + assert await _model_groups(caller) == [] + + +@pytest.mark.asyncio +async def test_v1_model_info_offers_only_the_rows_the_caller_would_see(monkeypatch): + _install_router(monkeypatch, _deployment("open-model"), _deployment("hidden-model", discoverable=False)) + gate: _Gate = _Gate() + _register(monkeypatch, gate) + + assert await _v1_model_info_names(_non_admin()) == ["open-model"] + assert await _v1_model_info_names(_admin()) == ["open-model", "hidden-model"] + assert gate.seen == [("open-model",), ("open-model", "hidden-model")] + + +@pytest.mark.asyncio +async def test_listing_keeps_its_order_whatever_order_the_callback_returns(monkeypatch): + _install_router(monkeypatch, _deployment("a"), _deployment("b"), _deployment("c")) + _register(monkeypatch, _ReversingGate()) + + assert await _v1_models(_non_admin()) == ["a", "b", "c"] + assert await _v1_model_info_names(_non_admin()) == ["a", "b", "c"] + + +@pytest.mark.asyncio +async def test_a_callback_returning_a_string_is_an_error_not_an_empty_listing(two_model_router, monkeypatch): + _register(monkeypatch, _StringReturningGate()) + + with pytest.raises(ProxyException, match=r"_StringReturningGate\.async_filter_listed_models") as raised: + await _v1_models(_non_admin()) + assert raised.value.code == "500" + + +@pytest.mark.asyncio +async def test_alias_of_a_hidden_model_is_not_listed(two_model_router, monkeypatch): + _register(monkeypatch, _Gate(hidden=frozenset({"restricted-model"}))) + caller = _non_admin(aliases={"mini": "restricted-model", "wide": "open-model"}) + + assert await _v1_models(caller) == ["open-model", "wide"] + assert await _model_by_id_status("mini", caller) == 404 + assert await _model_by_id_status("wide", caller) == 200 + + +@pytest.mark.asyncio +async def test_callback_error_reaches_the_caller(two_model_router, monkeypatch): + _register(monkeypatch, _RaisingGate()) + + with pytest.raises(HTTPException) as raised: + await _v1_models(_non_admin()) + assert raised.value.status_code == 503 + + +@pytest.mark.asyncio +async def test_hidden_model_still_routes_for_direct_requests(two_model_router, monkeypatch): + _register(monkeypatch, _Gate(hidden=frozenset({"restricted-model"}))) + assert "restricted-model" not in await _v1_models(_non_admin()) + + deployment = two_model_router.get_available_deployment( + model="restricted-model", messages=[{"role": "user", "content": "hi"}] + ) + assert deployment["model_name"] == "restricted-model" + + +@pytest.mark.asyncio +async def test_model_group_info_offers_a2a_agent_groups_to_the_callback(two_model_router, monkeypatch): + from litellm.proxy._types import LiteLLM_ObjectPermissionTable + from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry + from litellm.types.agents import AgentResponse + + monkeypatch.setattr( + global_agent_registry, + "agent_list", + [AgentResponse(agent_id="agent-1", agent_name="helper", agent_card_params={})], + ) + caller = _non_admin(object_permission=LiteLLM_ObjectPermissionTable(object_permission_id="p1", agents=["agent-1"])) + gate: _Gate = _Gate() + _register(monkeypatch, gate) + + assert await _model_groups(caller) == ["open-model", "restricted-model", "a2a/helper"] + assert gate.seen == [("open-model", "restricted-model", "a2a/helper")] + + _register(monkeypatch, _Gate(hidden=frozenset({"a2a/helper", "restricted-model"}))) + + assert await _model_groups(caller) == ["open-model"] + + +async def _v1_model_info_by_deployment_id(deployment_id: str, user_api_key_dict: UserAPIKeyAuth) -> int | list[str]: + try: + response = await proxy_server.model_info_v1(user_api_key_dict=user_api_key_dict, litellm_model_id=deployment_id) + except HTTPException as error: + return error.status_code + return [row["model_name"] for row in json.loads(response.body)["data"]] + + +@pytest.mark.asyncio +async def test_v1_model_info_by_deployment_id_answers_like_an_unknown_id_for_a_hidden_model( + two_model_router, monkeypatch +): + _register(monkeypatch, _Gate(hidden=frozenset({"restricted-model"}))) + + assert await _v1_model_info_by_deployment_id("restricted-model-id", _non_admin()) == 400 + assert await _v1_model_info_by_deployment_id("no-such-id", _non_admin()) == 400 + assert await _v1_model_info_by_deployment_id("open-model-id", _non_admin()) == ["open-model"] + + +@pytest.mark.asyncio +async def test_v1_model_info_by_deployment_id_offers_the_public_team_name(team_router, monkeypatch): + gate: _Gate = _Gate(hidden=frozenset({"team-gpt"})) + _register(monkeypatch, gate) + + assert await _v1_model_info_by_deployment_id("model_name_team1_abc-id", _team_member()) == 400 + assert await _v1_model_info_by_deployment_id("model_name_team1_def-id", _team_member()) == ["team-chat"] + assert gate.seen == [("team-gpt",), ("team-chat",)] + + +@pytest.mark.asyncio +async def test_v1_model_info_by_deployment_id_offers_the_name_its_listing_shows_in_legacy_mode( + team_router, monkeypatch +): + monkeypatch.setattr(proxy_server, "general_settings", {"use_team_public_model_name": False}) + gate: _Gate = _Gate(hidden=frozenset({"team-gpt"})) + _register(monkeypatch, gate) + + assert await _v1_model_info_names(_team_member()) == ["team-chat"] + assert await _v1_model_info_by_deployment_id("model_name_team1_abc-id", _team_member()) == 400 + assert gate.seen[-1] == ("team-gpt",) From 248f0eb159c6c2788a28d4671104ffc4ce24904d Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 24 Sep 2026 22:59:11 +0000 Subject: [PATCH 095/120] ci: move tests/proxy_unit_tests to tests/unit/proxy and run the proxy-db shards from litellm-tests (#42903) * ci: fix the litellm-tests unit job with sysmon coverage, an env allowlist and coverage upload on failure * test: replace key-dependent proxy, enterprise and mcp unit tests with synthetic values and integration and e2e coverage * test: drop key reads at the legacy proxy, enterprise and mcp paths and wire the gemini pass-through split * ci: move caching, proxy-extras, gateway and enterprise tests into tests/unit and run them from litellm-tests under their legacy flags * ci: move caching, proxy-extras, gateway and enterprise tests into tests/unit and run them from litellm-tests under their legacy flags * ci: move tests/proxy_unit_tests to tests/unit/proxy and run the proxy-db shards from litellm-tests * ci: fail the unit shard when circleci tests split errors * test: drop restating comments from the gemini pass-through split * build: point the local proxy unit targets at the nested tests/unit/proxy tree * ci: exit the unit shard cleanly when circleci tests split assigns it no files --------- Co-authored-by: yuneng --- .circleci/scripts/classify_changes.sh | 2 +- .circleci/scripts/unit_selection.sh | 72 +++++++++ .circleci/tests.yml | 30 +++- .github/scripts/assert_ci_coverage.py | 1 - .github/workflows/test-unit-proxy-db.yml | 97 ++++------- .github/workflows/test-unit.yml | 8 +- Makefile | 10 +- litellm/llms/litellm_proxy/skills/README.md | 2 +- .../user_api_key_auth_code_coverage.py | 4 +- .../image_endpoints/test_azure_routes.py | 3 +- .../test_litellm/test_circleci_path_filter.py | 2 +- .../proxy/__init__.py} | 0 tests/unit/proxy/auth/__init__.py | 0 .../proxy/auth}/test_auth_checks.py | 0 .../test_default_end_user_budget_simple.py | 0 .../proxy/auth}/test_jwt.py | 0 .../auth}/test_models_fallback_endpoint.py | 0 .../auth}/test_multipart_bypass_repro.py | 0 .../proxy/auth}/test_proxy_routes.py | 0 .../proxy/auth}/test_user_api_key_auth.py | 0 tests/unit/proxy/common_utils/__init__.py | 0 .../common_utils}/test_check_batch_cost.py | 0 .../test_check_responses_cost.py | 0 .../test_proxy_encrypt_decrypt.py | 0 .../common_utils}/test_realtime_cache.py | 0 tests/unit/proxy/conftest.py | 150 ++++++++++++++++++ tests/unit/proxy/db/__init__.py | 0 .../proxy/db/db_transaction_queue/__init__.py | 0 .../test_e2e_pod_lock_manager.py | 0 .../proxy/db}/test_update_daily_tag_spend.py | 0 .../proxy/example_config_yaml/__init__.py | 0 .../example_config_yaml/aliases_config.yaml | 0 .../example_config_yaml/azure_config.yaml | 0 .../example_config_yaml/cache_no_params.yaml | 0 .../cache_with_params.yaml | 0 .../config_with_env_vars.yaml | 0 .../config_with_include.yaml | 0 .../config_with_missing_include.yaml | 0 .../config_with_multiple_includes.yaml | 0 .../example_config_yaml/included_models.yaml | 0 .../example_config_yaml/langfuse_config.yaml | 0 .../example_config_yaml/load_balancer.yaml | 0 .../example_config_yaml/models_file_1.yaml | 0 .../example_config_yaml/models_file_2.yaml | 0 .../opentelemetry_config.yaml | 0 .../example_config_yaml/simple_config.yaml | 0 tests/unit/proxy/google_endpoints/__init__.py | 0 .../test_gemini_agents_endpoints.py | 0 .../test_google_endpoint_routing.py | 0 .../test_google_gemini_proxy_request.py | 0 tests/unit/proxy/hooks/__init__.py | 0 .../proxy/hooks}/test_banned_keyword_list.py | 0 ...test_unit_test_max_model_budget_limiter.py | 0 .../proxy/management_endpoints/__init__.py | 0 .../test_jwt_key_mapping.py | 2 +- .../test_key_generate_prisma.py | 0 .../unit/proxy/management_helpers/__init__.py | 0 .../test_audit_logs_proxy.py | 0 tests/unit/proxy/middleware/__init__.py | 0 .../test_request_size_limit_middleware.py | 0 tests/unit/proxy/public_endpoints/__init__.py | 0 .../test_blog_posts_endpoint.py | 0 tests/unit/proxy/response_polling/__init__.py | 0 .../test_response_polling_handler.py | 0 tests/unit/proxy/spend_tracking/__init__.py | 0 .../test_search_api_logging.py | 0 .../proxy}/test_aproxy_startup.py | 0 tests/unit/proxy/test_configs/__init__.py | 0 .../proxy}/test_configs/custom_auth.py | 0 ...st_cloudflare_azure_with_cache_config.yaml | 0 .../proxy}/test_configs/test_config.yaml | 0 .../test_configs/test_config_custom_auth.yaml | 0 .../test_configs/test_config_no_auth.yaml | 0 .../test_configs/test_guardrails_config.yaml | 0 .../proxy}/test_custom_callback_input.py | 0 .../proxy}/test_custom_logger_s3_gcs.py | 0 .../proxy}/test_custom_tokenizer_bug.py | 0 .../proxy}/test_db_schema_changes.py | 0 .../test_deprecated_key_grace_period.py | 0 .../proxy}/test_get_favicon.py | 0 .../proxy}/test_get_image.py | 0 .../test_prisma_client_backoff_retry.py | 0 .../proxy}/test_prompt_test_endpoint.py | 0 .../proxy}/test_proxy_config_unit_test.py | 2 +- .../proxy}/test_proxy_custom_auth.py | 0 .../proxy}/test_proxy_reject_logging.py | 0 .../proxy}/test_proxy_server.py | 4 +- .../proxy}/test_proxy_setting_guardrails.py | 0 .../proxy}/test_proxy_token_counter.py | 0 .../proxy}/test_proxy_utils.py | 0 .../proxy}/test_reducto_ocr_route.py | 0 .../test_response_polling_pre_call_checks.py | 0 .../proxy}/test_server_root_path.py | 0 .../proxy}/test_ui_path_detection.py | 0 .../proxy}/test_unit_test_proxy_hooks.py | 0 .../proxy}/test_update_spend.py | 0 .../test_zero_cost_model_budget_bypass.py | 0 .../proxy}/vertex_key.json | 0 .../skills}/test_skills_db.py | 2 +- tests/unit/skills/test_skills_main.py | 2 +- 100 files changed, 305 insertions(+), 88 deletions(-) rename tests/{proxy_unit_tests/test_key_generate_dynamodb.py => unit/proxy/__init__.py} (100%) create mode 100644 tests/unit/proxy/auth/__init__.py rename tests/{proxy_unit_tests => unit/proxy/auth}/test_auth_checks.py (100%) rename tests/{proxy_unit_tests => unit/proxy/auth}/test_default_end_user_budget_simple.py (100%) rename tests/{proxy_unit_tests => unit/proxy/auth}/test_jwt.py (100%) rename tests/{proxy_unit_tests => unit/proxy/auth}/test_models_fallback_endpoint.py (100%) rename tests/{proxy_unit_tests => unit/proxy/auth}/test_multipart_bypass_repro.py (100%) rename tests/{proxy_unit_tests => unit/proxy/auth}/test_proxy_routes.py (100%) rename tests/{proxy_unit_tests => unit/proxy/auth}/test_user_api_key_auth.py (100%) create mode 100644 tests/unit/proxy/common_utils/__init__.py rename tests/{proxy_unit_tests => unit/proxy/common_utils}/test_check_batch_cost.py (100%) rename tests/{proxy_unit_tests => unit/proxy/common_utils}/test_check_responses_cost.py (100%) rename tests/{proxy_unit_tests => unit/proxy/common_utils}/test_proxy_encrypt_decrypt.py (100%) rename tests/{proxy_unit_tests => unit/proxy/common_utils}/test_realtime_cache.py (100%) create mode 100644 tests/unit/proxy/conftest.py create mode 100644 tests/unit/proxy/db/__init__.py create mode 100644 tests/unit/proxy/db/db_transaction_queue/__init__.py rename tests/{proxy_unit_tests => unit/proxy/db/db_transaction_queue}/test_e2e_pod_lock_manager.py (100%) rename tests/{proxy_unit_tests => unit/proxy/db}/test_update_daily_tag_spend.py (100%) create mode 100644 tests/unit/proxy/example_config_yaml/__init__.py rename tests/{proxy_unit_tests => unit/proxy}/example_config_yaml/aliases_config.yaml (100%) rename tests/{proxy_unit_tests => unit/proxy}/example_config_yaml/azure_config.yaml (100%) rename tests/{proxy_unit_tests => unit/proxy}/example_config_yaml/cache_no_params.yaml (100%) rename tests/{proxy_unit_tests => unit/proxy}/example_config_yaml/cache_with_params.yaml (100%) rename tests/{proxy_unit_tests => unit/proxy}/example_config_yaml/config_with_env_vars.yaml (100%) rename tests/{proxy_unit_tests => unit/proxy}/example_config_yaml/config_with_include.yaml (100%) rename tests/{proxy_unit_tests => unit/proxy}/example_config_yaml/config_with_missing_include.yaml (100%) rename tests/{proxy_unit_tests => unit/proxy}/example_config_yaml/config_with_multiple_includes.yaml (100%) rename tests/{proxy_unit_tests => unit/proxy}/example_config_yaml/included_models.yaml (100%) rename tests/{proxy_unit_tests => unit/proxy}/example_config_yaml/langfuse_config.yaml (100%) rename tests/{proxy_unit_tests => unit/proxy}/example_config_yaml/load_balancer.yaml (100%) rename tests/{proxy_unit_tests => unit/proxy}/example_config_yaml/models_file_1.yaml (100%) rename tests/{proxy_unit_tests => unit/proxy}/example_config_yaml/models_file_2.yaml (100%) rename tests/{proxy_unit_tests => unit/proxy}/example_config_yaml/opentelemetry_config.yaml (100%) rename tests/{proxy_unit_tests => unit/proxy}/example_config_yaml/simple_config.yaml (100%) create mode 100644 tests/unit/proxy/google_endpoints/__init__.py rename tests/{proxy_unit_tests => unit/proxy/google_endpoints}/test_gemini_agents_endpoints.py (100%) rename tests/{proxy_unit_tests => unit/proxy/google_endpoints}/test_google_endpoint_routing.py (100%) rename tests/{proxy_unit_tests => unit/proxy/google_endpoints}/test_google_gemini_proxy_request.py (100%) create mode 100644 tests/unit/proxy/hooks/__init__.py rename tests/{proxy_unit_tests => unit/proxy/hooks}/test_banned_keyword_list.py (100%) rename tests/{proxy_unit_tests => unit/proxy/hooks}/test_unit_test_max_model_budget_limiter.py (100%) create mode 100644 tests/unit/proxy/management_endpoints/__init__.py rename tests/{proxy_unit_tests => unit/proxy/management_endpoints}/test_jwt_key_mapping.py (99%) rename tests/{proxy_unit_tests => unit/proxy/management_endpoints}/test_key_generate_prisma.py (100%) create mode 100644 tests/unit/proxy/management_helpers/__init__.py rename tests/{proxy_unit_tests => unit/proxy/management_helpers}/test_audit_logs_proxy.py (100%) create mode 100644 tests/unit/proxy/middleware/__init__.py rename tests/{proxy_unit_tests => unit/proxy/middleware}/test_request_size_limit_middleware.py (100%) create mode 100644 tests/unit/proxy/public_endpoints/__init__.py rename tests/{proxy_unit_tests => unit/proxy/public_endpoints}/test_blog_posts_endpoint.py (100%) create mode 100644 tests/unit/proxy/response_polling/__init__.py rename tests/{proxy_unit_tests => unit/proxy/response_polling}/test_response_polling_handler.py (100%) create mode 100644 tests/unit/proxy/spend_tracking/__init__.py rename tests/{proxy_unit_tests => unit/proxy/spend_tracking}/test_search_api_logging.py (100%) rename tests/{proxy_unit_tests => unit/proxy}/test_aproxy_startup.py (100%) create mode 100644 tests/unit/proxy/test_configs/__init__.py rename tests/{proxy_unit_tests => unit/proxy}/test_configs/custom_auth.py (100%) rename tests/{proxy_unit_tests => unit/proxy}/test_configs/test_cloudflare_azure_with_cache_config.yaml (100%) rename tests/{proxy_unit_tests => unit/proxy}/test_configs/test_config.yaml (100%) rename tests/{proxy_unit_tests => unit/proxy}/test_configs/test_config_custom_auth.yaml (100%) rename tests/{proxy_unit_tests => unit/proxy}/test_configs/test_config_no_auth.yaml (100%) rename tests/{proxy_unit_tests => unit/proxy}/test_configs/test_guardrails_config.yaml (100%) rename tests/{proxy_unit_tests => unit/proxy}/test_custom_callback_input.py (100%) rename tests/{proxy_unit_tests => unit/proxy}/test_custom_logger_s3_gcs.py (100%) rename tests/{proxy_unit_tests => unit/proxy}/test_custom_tokenizer_bug.py (100%) rename tests/{proxy_unit_tests => unit/proxy}/test_db_schema_changes.py (100%) rename tests/{proxy_unit_tests => unit/proxy}/test_deprecated_key_grace_period.py (100%) rename tests/{proxy_unit_tests => unit/proxy}/test_get_favicon.py (100%) rename tests/{proxy_unit_tests => unit/proxy}/test_get_image.py (100%) rename tests/{proxy_unit_tests => unit/proxy}/test_prisma_client_backoff_retry.py (100%) rename tests/{proxy_unit_tests => unit/proxy}/test_prompt_test_endpoint.py (100%) rename tests/{proxy_unit_tests => unit/proxy}/test_proxy_config_unit_test.py (99%) rename tests/{proxy_unit_tests => unit/proxy}/test_proxy_custom_auth.py (100%) rename tests/{proxy_unit_tests => unit/proxy}/test_proxy_reject_logging.py (100%) rename tests/{proxy_unit_tests => unit/proxy}/test_proxy_server.py (99%) rename tests/{proxy_unit_tests => unit/proxy}/test_proxy_setting_guardrails.py (100%) rename tests/{proxy_unit_tests => unit/proxy}/test_proxy_token_counter.py (100%) rename tests/{proxy_unit_tests => unit/proxy}/test_proxy_utils.py (100%) rename tests/{proxy_unit_tests => unit/proxy}/test_reducto_ocr_route.py (100%) rename tests/{proxy_unit_tests => unit/proxy}/test_response_polling_pre_call_checks.py (100%) rename tests/{proxy_unit_tests => unit/proxy}/test_server_root_path.py (100%) rename tests/{proxy_unit_tests => unit/proxy}/test_ui_path_detection.py (100%) rename tests/{proxy_unit_tests => unit/proxy}/test_unit_test_proxy_hooks.py (100%) rename tests/{proxy_unit_tests => unit/proxy}/test_update_spend.py (100%) rename tests/{proxy_unit_tests => unit/proxy}/test_zero_cost_model_budget_bypass.py (100%) rename tests/{proxy_unit_tests => unit/proxy}/vertex_key.json (100%) rename tests/{proxy_unit_tests => unit/skills}/test_skills_db.py (98%) diff --git a/.circleci/scripts/classify_changes.sh b/.circleci/scripts/classify_changes.sh index 01bc8290199..ad265a5e39f 100755 --- a/.circleci/scripts/classify_changes.sh +++ b/.circleci/scripts/classify_changes.sh @@ -31,7 +31,7 @@ while IFS= read -r file || [ -n "$file" ]; do case "$file" in model_prices_and_context_window.json | litellm/model_prices_and_context_window_backup.json | model_prices_and_context_window.schema.json) has_cost_map=true ;; - tests/test_litellm/* | tests/proxy_unit_tests/*) : ;; + tests/test_litellm/* | tests/proxy_unit_tests/* | tests/unit/proxy/*) : ;; *) outside_cost_map_set=true ;; esac done diff --git a/.circleci/scripts/unit_selection.sh b/.circleci/scripts/unit_selection.sh index 8d2b8a42691..2c60c5b1334 100755 --- a/.circleci/scripts/unit_selection.sh +++ b/.circleci/scripts/unit_selection.sh @@ -7,6 +7,18 @@ legacy_flags=( caching-local enterprise-package enterprise-routing + proxy-db-auth-checks + proxy-db-budgets + proxy-db-custom-logging + proxy-db-db-and-spend + proxy-db-endpoints-and-responses + proxy-db-guardrails-hooks + proxy-db-jwt-and-keys + proxy-db-key-generation + proxy-db-logging-misc + proxy-db-proxy-runtime + proxy-db-proxy-server-core + proxy-db-proxy-utils proxy-extras proxy-infra ) @@ -34,6 +46,66 @@ legacy_paths() { echo tests/unit/enterprise/proxy/test_file_deletion_blocking.py echo tests/unit/enterprise/proxy/test_managed_files_access_check.py echo tests/unit/enterprise/proxy/test_managed_files_hook.py ;; + proxy-db-auth-checks) + echo tests/unit/proxy/auth/test_auth_checks.py + echo tests/unit/proxy/auth/test_user_api_key_auth.py + echo tests/unit/proxy/test_deprecated_key_grace_period.py ;; + proxy-db-budgets) + echo tests/unit/proxy/auth/test_default_end_user_budget_simple.py + echo tests/unit/proxy/hooks/test_unit_test_max_model_budget_limiter.py + echo tests/unit/proxy/test_zero_cost_model_budget_bypass.py ;; + proxy-db-custom-logging) + echo tests/unit/proxy/test_custom_callback_input.py + echo tests/unit/proxy/test_custom_logger_s3_gcs.py ;; + proxy-db-db-and-spend) + echo tests/unit/proxy/common_utils/test_proxy_encrypt_decrypt.py + echo tests/unit/proxy/db/db_transaction_queue/test_e2e_pod_lock_manager.py + echo tests/unit/proxy/db/test_update_daily_tag_spend.py + echo tests/unit/proxy/test_db_schema_changes.py + echo tests/unit/proxy/test_prisma_client_backoff_retry.py + echo tests/unit/proxy/test_update_spend.py + echo tests/unit/skills/test_skills_db.py ;; + proxy-db-endpoints-and-responses) + echo tests/unit/proxy/auth/test_models_fallback_endpoint.py + echo tests/unit/proxy/common_utils/test_check_batch_cost.py + echo tests/unit/proxy/common_utils/test_check_responses_cost.py + echo tests/unit/proxy/common_utils/test_realtime_cache.py + echo tests/unit/proxy/google_endpoints/test_gemini_agents_endpoints.py + echo tests/unit/proxy/google_endpoints/test_google_endpoint_routing.py + echo tests/unit/proxy/google_endpoints/test_google_gemini_proxy_request.py + echo tests/unit/proxy/public_endpoints/test_blog_posts_endpoint.py + echo tests/unit/proxy/response_polling/test_response_polling_handler.py + echo tests/unit/proxy/test_custom_tokenizer_bug.py + echo tests/unit/proxy/test_get_favicon.py + echo tests/unit/proxy/test_get_image.py + echo tests/unit/proxy/test_prompt_test_endpoint.py + echo tests/unit/proxy/test_reducto_ocr_route.py + echo tests/unit/proxy/test_response_polling_pre_call_checks.py + echo tests/unit/proxy/test_ui_path_detection.py ;; + proxy-db-guardrails-hooks) + echo tests/unit/proxy/hooks/test_banned_keyword_list.py + echo tests/unit/proxy/test_proxy_setting_guardrails.py + echo tests/unit/proxy/test_unit_test_proxy_hooks.py ;; + proxy-db-jwt-and-keys) + echo tests/unit/proxy/auth/test_jwt.py + echo tests/unit/proxy/management_endpoints/test_jwt_key_mapping.py + echo tests/unit/proxy/test_proxy_custom_auth.py ;; + proxy-db-key-generation) echo tests/unit/proxy/management_endpoints/test_key_generate_prisma.py ;; + proxy-db-logging-misc) + echo tests/unit/proxy/management_helpers/test_audit_logs_proxy.py + echo tests/unit/proxy/spend_tracking/test_search_api_logging.py + echo tests/unit/proxy/test_proxy_reject_logging.py ;; + proxy-db-proxy-runtime) + echo tests/unit/proxy/auth/test_multipart_bypass_repro.py + echo tests/unit/proxy/auth/test_proxy_routes.py + echo tests/unit/proxy/middleware/test_request_size_limit_middleware.py + echo tests/unit/proxy/test_proxy_config_unit_test.py + echo tests/unit/proxy/test_proxy_token_counter.py + echo tests/unit/proxy/test_server_root_path.py ;; + proxy-db-proxy-server-core) + echo tests/unit/proxy/test_aproxy_startup.py + echo tests/unit/proxy/test_proxy_server.py ;; + proxy-db-proxy-utils) echo tests/unit/proxy/test_proxy_utils.py ;; proxy-extras) echo tests/unit/litellm_proxy_extras ;; proxy-infra) echo tests/unit/gateway ;; *) echo "unit_selection.sh: unknown flag $1" >&2; exit 1 ;; diff --git a/.circleci/tests.yml b/.circleci/tests.yml index 38fe44bb625..08e735637b7 100644 --- a/.circleci/tests.yml +++ b/.circleci/tests.yml @@ -317,7 +317,35 @@ workflows: reruns: 2 matrix: parameters: - flag: [enterprise-package, proxy-infra] + flag: + - enterprise-package + - proxy-infra + - proxy-db-auth-checks + - proxy-db-jwt-and-keys + - proxy-db-proxy-server-core + - proxy-db-proxy-runtime + - proxy-db-custom-logging + - proxy-db-logging-misc + - proxy-db-db-and-spend + - proxy-db-guardrails-hooks + - proxy-db-budgets + - proxy-db-endpoints-and-responses + base_ref: << pipeline.event.name == "pull_request" and pipeline.event.github.pull_request.base.ref or "" >> + pull_request_url: << pipeline.event.name == "pull_request" and pipeline.event.github.pull_request.url or "" >> + - unit: + name: unit-proxy-db-proxy-utils + flag: proxy-db-proxy-utils + shards: 1 + reruns: 2 + dist: worksteal + base_ref: << pipeline.event.name == "pull_request" and pipeline.event.github.pull_request.base.ref or "" >> + pull_request_url: << pipeline.event.name == "pull_request" and pipeline.event.github.pull_request.url or "" >> + - unit: + name: unit-proxy-db-key-generation + flag: proxy-db-key-generation + shards: 1 + workers: 0 + reruns: 2 base_ref: << pipeline.event.name == "pull_request" and pipeline.event.github.pull_request.base.ref or "" >> pull_request_url: << pipeline.event.name == "pull_request" and pipeline.event.github.pull_request.url or "" >> - documentation diff --git a/.github/scripts/assert_ci_coverage.py b/.github/scripts/assert_ci_coverage.py index d8246225a3b..01a01b1034b 100644 --- a/.github/scripts/assert_ci_coverage.py +++ b/.github/scripts/assert_ci_coverage.py @@ -34,7 +34,6 @@ GLOB_CHARS = frozenset("*?") # tests has to be named by some shard or it runs nowhere. A child listed here is # itself decomposed one level deeper and is checked through its own entry. SHARDED_ROOTS: tuple[str, ...] = ( - "tests/proxy_unit_tests", "tests/test_litellm", "tests/test_litellm/proxy", ) diff --git a/.github/workflows/test-unit-proxy-db.yml b/.github/workflows/test-unit-proxy-db.yml index 73015ac6e02..86b385d91a7 100644 --- a/.github/workflows/test-unit-proxy-db.yml +++ b/.github/workflows/test-unit-proxy-db.yml @@ -20,6 +20,12 @@ concurrency: # rather than alphabetical letter ranges. Adding a new test file means adding it # to whichever group it belongs to, not reshuffling slices. # +# `.circleci/tests.yml` runs each group's files on same-repo events under the +# `proxy-db-` Codecov flag; `.circleci/scripts/unit_selection.sh` holds +# the file lists. CircleCI does not build pull requests from forks, so `fork-flag` +# makes the shard run that list there. `test-path` keeps the files that still +# reach real providers and never left tests/proxy_unit_tests. +# # Design targets: # * Every shard runs in <= 7 minutes of wall-clock on the default runner. # Most of a shard's time is pytest plugin load + xdist worker imports + @@ -58,7 +64,7 @@ jobs: proxy-db: needs: assert-shard-coverage # Display only the semantic shard name in the checks UI instead of GHA's - # default "proxy-db (key-generation, tests/proxy_unit_tests/…, 0, loadscope, 20)" + # default "proxy-db (key-generation, tests/unit/proxy/…, 0, loadscope, 20)" # which includes every matrix field and gets truncated past the test-path. name: ${{ matrix.test-group }} permissions: @@ -71,132 +77,93 @@ jobs: include: # Must run serially — event-loop conflict with the logging worker. - test-group: key-generation - test-path: "tests/proxy_unit_tests/test_key_generate_prisma.py" + test-path: "" + fork-flag: proxy-db-key-generation workers: 0 dist: loadscope timeout: 20 # ---- auth: split into 2 shards ---- - test-group: auth-checks - test-path: >- - tests/proxy_unit_tests/test_auth_checks.py - tests/proxy_unit_tests/test_user_api_key_auth.py - tests/proxy_unit_tests/test_deprecated_key_grace_period.py + test-path: "" + fork-flag: proxy-db-auth-checks workers: 4 dist: loadscope timeout: 15 - test-group: jwt-and-keys - test-path: >- - tests/proxy_unit_tests/test_jwt.py - tests/proxy_unit_tests/test_jwt_key_mapping.py - tests/proxy_unit_tests/test_proxy_custom_auth.py - tests/proxy_unit_tests/test_key_generate_dynamodb.py + test-path: "" + fork-flag: proxy-db-jwt-and-keys workers: 4 dist: loadscope timeout: 15 # ---- test_proxy_utils.py, single shard, worksteal distribution ---- - test-group: proxy-utils - test-path: "tests/proxy_unit_tests/test_proxy_utils.py" + test-path: "" + fork-flag: proxy-db-proxy-utils workers: 4 dist: worksteal timeout: 15 # ---- proxy server: split into 2 shards ---- - test-group: proxy-server-core - test-path: >- - tests/proxy_unit_tests/test_proxy_server.py - tests/proxy_unit_tests/test_proxy_server_gemini_pass_through.py - tests/proxy_unit_tests/test_aproxy_startup.py + test-path: "tests/proxy_unit_tests/test_proxy_server_gemini_pass_through.py" + fork-flag: proxy-db-proxy-server-core workers: 4 dist: loadscope timeout: 15 - test-group: proxy-runtime - test-path: >- - tests/proxy_unit_tests/test_proxy_config_unit_test.py - tests/proxy_unit_tests/test_proxy_routes.py - tests/proxy_unit_tests/test_server_root_path.py - tests/proxy_unit_tests/test_proxy_token_counter.py - tests/proxy_unit_tests/test_request_size_limit_middleware.py - tests/proxy_unit_tests/test_multipart_bypass_repro.py + test-path: "" + fork-flag: proxy-db-proxy-runtime workers: 4 dist: loadscope timeout: 15 # ---- logging: split into 2 shards ---- - test-group: custom-logging - test-path: >- - tests/proxy_unit_tests/test_custom_callback_input.py - tests/proxy_unit_tests/test_custom_logger_s3_gcs.py - tests/proxy_unit_tests/test_proxy_custom_logger.py + test-path: "tests/proxy_unit_tests/test_proxy_custom_logger.py" + fork-flag: proxy-db-custom-logging workers: 4 dist: loadscope timeout: 15 - test-group: logging-misc - test-path: >- - tests/proxy_unit_tests/test_proxy_reject_logging.py - tests/proxy_unit_tests/test_audit_logs_proxy.py - tests/proxy_unit_tests/test_search_api_logging.py + test-path: "" + fork-flag: proxy-db-logging-misc workers: 4 dist: loadscope timeout: 15 - test-group: db-and-spend - test-path: >- - tests/proxy_unit_tests/test_prisma_client_backoff_retry.py - tests/proxy_unit_tests/test_db_schema_changes.py - tests/proxy_unit_tests/test_e2e_pod_lock_manager.py - tests/proxy_unit_tests/test_skills_db.py - tests/proxy_unit_tests/test_update_daily_tag_spend.py - tests/proxy_unit_tests/test_update_spend.py - tests/proxy_unit_tests/test_proxy_encrypt_decrypt.py + test-path: "" + fork-flag: proxy-db-db-and-spend workers: 4 dist: loadscope timeout: 15 # ---- guardrails + budget + hooks: split into 2 ---- - test-group: guardrails-hooks - test-path: >- - tests/proxy_unit_tests/test_proxy_setting_guardrails.py - tests/proxy_unit_tests/test_banned_keyword_list.py - tests/proxy_unit_tests/test_unit_test_proxy_hooks.py + test-path: "" + fork-flag: proxy-db-guardrails-hooks workers: 4 dist: loadscope timeout: 15 - test-group: budgets - test-path: >- - tests/proxy_unit_tests/test_default_end_user_budget_simple.py - tests/proxy_unit_tests/test_unit_test_max_model_budget_limiter.py - tests/proxy_unit_tests/test_zero_cost_model_budget_bypass.py + test-path: "" + fork-flag: proxy-db-budgets workers: 4 dist: loadscope timeout: 15 - test-group: endpoints-and-responses - test-path: >- - tests/proxy_unit_tests/test_blog_posts_endpoint.py - tests/proxy_unit_tests/test_models_fallback_endpoint.py - tests/proxy_unit_tests/test_google_endpoint_routing.py - tests/proxy_unit_tests/test_google_gemini_proxy_request.py - tests/proxy_unit_tests/test_gemini_agents_endpoints.py - tests/proxy_unit_tests/test_get_favicon.py - tests/proxy_unit_tests/test_get_image.py - tests/proxy_unit_tests/test_reducto_ocr_route.py - tests/proxy_unit_tests/test_ui_path_detection.py - tests/proxy_unit_tests/test_prompt_test_endpoint.py - tests/proxy_unit_tests/test_check_batch_cost.py - tests/proxy_unit_tests/test_check_responses_cost.py - tests/proxy_unit_tests/test_response_polling_handler.py - tests/proxy_unit_tests/test_response_polling_pre_call_checks.py - tests/proxy_unit_tests/test_realtime_cache.py - tests/proxy_unit_tests/test_proxy_exception_mapping.py - tests/proxy_unit_tests/test_custom_tokenizer_bug.py + test-path: "tests/proxy_unit_tests/test_proxy_exception_mapping.py" + fork-flag: proxy-db-endpoints-and-responses workers: 4 dist: loadscope timeout: 15 uses: ./.github/workflows/_test-unit-base.yml with: test-path: ${{ matrix.test-path }} + fork-flag: ${{ matrix.fork-flag }} workers: ${{ matrix.workers }} reruns: 2 timeout-minutes: ${{ matrix.timeout }} diff --git a/.github/workflows/test-unit.yml b/.github/workflows/test-unit.yml index 94de6040038..bf2e1602be8 100644 --- a/.github/workflows/test-unit.yml +++ b/.github/workflows/test-unit.yml @@ -31,10 +31,10 @@ concurrency: # number, so a partially-specified entry would fail the call rather than fall # back to the default. # -# tests/proxy_unit_tests keeps its own caller (test-unit-proxy-db.yml): it is -# already a matrix and carries a shard-coverage guard that reads that file by -# name. Folding it in here is a follow-up, together with generalising that guard -# into assert_ci_coverage.py. +# tests/unit/proxy keeps its own caller (test-unit-proxy-db.yml): it is already +# a matrix and carries a shard-coverage guard that reads that file by name. +# Folding it in here is a follow-up, together with generalising that guard into +# assert_ci_coverage.py. # # `fork-flag` names the `.circleci/tests.yml` job that now runs part of the # shard under the same Codecov flag. CircleCI does not build pull requests from diff --git a/Makefile b/Makefile index 6263b646c17..28daf589a23 100644 --- a/Makefile +++ b/Makefile @@ -51,8 +51,8 @@ help: @echo " make test-unit-core-utils - Run core utils tests (~32 files)" @echo " make test-unit-other - Run other tests (caching, responses, etc., ~69 files)" @echo " make test-unit-root - Run root-level tests (~34 files)" - @echo " make test-proxy-unit-a - Run proxy_unit_tests (a-o, ~20 files)" - @echo " make test-proxy-unit-b - Run proxy_unit_tests (p-z, ~28 files)" + @echo " make test-proxy-unit-a - Run tests/unit/proxy (a-o)" + @echo " make test-proxy-unit-b - Run tests/unit/proxy (p-z)" @echo " make test-integration - Run integration tests" @echo " make test-unit-helm - Run helm unit tests" @echo " make test-rust-extension - Build the Rust extension and run its public Python tests" @@ -337,12 +337,12 @@ test-unit-other: install-test-deps test-unit-root: install-test-deps $(UV_RUN) pytest tests/test_litellm/test_*.py --tb=short -vv -n 4 --durations=20 -# Proxy unit tests (tests/proxy_unit_tests split alphabetically) +# Proxy unit tests (tests/unit/proxy split alphabetically) test-proxy-unit-a: install-test-deps - $(UV_RUN) pytest tests/proxy_unit_tests/test_[a-o]*.py --tb=short -vv -n 2 --durations=20 + $(UV_RUN) pytest tests/unit/proxy --ignore-glob='tests/unit/proxy/test_[p-z]*.py' --tb=short -vv -n 2 --durations=20 test-proxy-unit-b: install-test-deps - $(UV_RUN) pytest tests/proxy_unit_tests/test_[p-z]*.py --tb=short -vv -n 2 --durations=20 + $(UV_RUN) pytest tests/unit/proxy/test_[p-z]*.py tests/unit/skills --tb=short -vv -n 2 --durations=20 test-integration: install-test-deps $(UV_RUN) pytest tests/ -k "not test_litellm" diff --git a/litellm/llms/litellm_proxy/skills/README.md b/litellm/llms/litellm_proxy/skills/README.md index a896aa1166e..ccbd394cddc 100644 --- a/litellm/llms/litellm_proxy/skills/README.md +++ b/litellm/llms/litellm_proxy/skills/README.md @@ -369,7 +369,7 @@ model LiteLLM_SkillsTable { Run the tests: ```bash -pytest tests/proxy_unit_tests/test_skills_db.py -v +pytest tests/unit/skills/test_skills_db.py -v ``` Tests cover: diff --git a/tests/code_coverage_tests/user_api_key_auth_code_coverage.py b/tests/code_coverage_tests/user_api_key_auth_code_coverage.py index a9c2f8ef15f..2f221a7ebe7 100644 --- a/tests/code_coverage_tests/user_api_key_auth_code_coverage.py +++ b/tests/code_coverage_tests/user_api_key_auth_code_coverage.py @@ -31,11 +31,11 @@ def get_function_names_from_file(file_path): def get_all_functions_called_in_tests(base_dir): """ Returns a set of function names that are called in test functions - inside 'local_testing' and 'proxy_unit_tests' directories, + inside 'local_testing' and 'unit/proxy' directories, specifically in files containing the word 'router'. """ called_functions = set() - test_dirs = ["local_testing", "proxy_unit_tests"] + test_dirs = ["local_testing", "unit/proxy"] for test_dir in test_dirs: dir_path = os.path.join(base_dir, test_dir) diff --git a/tests/test_litellm/proxy/image_endpoints/test_azure_routes.py b/tests/test_litellm/proxy/image_endpoints/test_azure_routes.py index 91fff717d25..46fe9a6f893 100644 --- a/tests/test_litellm/proxy/image_endpoints/test_azure_routes.py +++ b/tests/test_litellm/proxy/image_endpoints/test_azure_routes.py @@ -53,7 +53,8 @@ def client_no_auth(): config_fp = ( repo_root / "tests" - / "proxy_unit_tests" + / "unit" + / "proxy" / "test_configs" / "test_config_no_auth.yaml" ) diff --git a/tests/test_litellm/test_circleci_path_filter.py b/tests/test_litellm/test_circleci_path_filter.py index 84e2327057d..dcce7f57113 100644 --- a/tests/test_litellm/test_circleci_path_filter.py +++ b/tests/test_litellm/test_circleci_path_filter.py @@ -107,7 +107,7 @@ CI = [".github/workflows/test-litellm-ui-unit.yml"] ), ( "cost-map-only", - ["model_prices_and_context_window.json", "tests/proxy_unit_tests/test_y.py"], + ["model_prices_and_context_window.json", "tests/unit/proxy/test_y.py"], "run", ), ( diff --git a/tests/proxy_unit_tests/test_key_generate_dynamodb.py b/tests/unit/proxy/__init__.py similarity index 100% rename from tests/proxy_unit_tests/test_key_generate_dynamodb.py rename to tests/unit/proxy/__init__.py diff --git a/tests/unit/proxy/auth/__init__.py b/tests/unit/proxy/auth/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/proxy_unit_tests/test_auth_checks.py b/tests/unit/proxy/auth/test_auth_checks.py similarity index 100% rename from tests/proxy_unit_tests/test_auth_checks.py rename to tests/unit/proxy/auth/test_auth_checks.py diff --git a/tests/proxy_unit_tests/test_default_end_user_budget_simple.py b/tests/unit/proxy/auth/test_default_end_user_budget_simple.py similarity index 100% rename from tests/proxy_unit_tests/test_default_end_user_budget_simple.py rename to tests/unit/proxy/auth/test_default_end_user_budget_simple.py diff --git a/tests/proxy_unit_tests/test_jwt.py b/tests/unit/proxy/auth/test_jwt.py similarity index 100% rename from tests/proxy_unit_tests/test_jwt.py rename to tests/unit/proxy/auth/test_jwt.py diff --git a/tests/proxy_unit_tests/test_models_fallback_endpoint.py b/tests/unit/proxy/auth/test_models_fallback_endpoint.py similarity index 100% rename from tests/proxy_unit_tests/test_models_fallback_endpoint.py rename to tests/unit/proxy/auth/test_models_fallback_endpoint.py diff --git a/tests/proxy_unit_tests/test_multipart_bypass_repro.py b/tests/unit/proxy/auth/test_multipart_bypass_repro.py similarity index 100% rename from tests/proxy_unit_tests/test_multipart_bypass_repro.py rename to tests/unit/proxy/auth/test_multipart_bypass_repro.py diff --git a/tests/proxy_unit_tests/test_proxy_routes.py b/tests/unit/proxy/auth/test_proxy_routes.py similarity index 100% rename from tests/proxy_unit_tests/test_proxy_routes.py rename to tests/unit/proxy/auth/test_proxy_routes.py diff --git a/tests/proxy_unit_tests/test_user_api_key_auth.py b/tests/unit/proxy/auth/test_user_api_key_auth.py similarity index 100% rename from tests/proxy_unit_tests/test_user_api_key_auth.py rename to tests/unit/proxy/auth/test_user_api_key_auth.py diff --git a/tests/unit/proxy/common_utils/__init__.py b/tests/unit/proxy/common_utils/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/proxy_unit_tests/test_check_batch_cost.py b/tests/unit/proxy/common_utils/test_check_batch_cost.py similarity index 100% rename from tests/proxy_unit_tests/test_check_batch_cost.py rename to tests/unit/proxy/common_utils/test_check_batch_cost.py diff --git a/tests/proxy_unit_tests/test_check_responses_cost.py b/tests/unit/proxy/common_utils/test_check_responses_cost.py similarity index 100% rename from tests/proxy_unit_tests/test_check_responses_cost.py rename to tests/unit/proxy/common_utils/test_check_responses_cost.py diff --git a/tests/proxy_unit_tests/test_proxy_encrypt_decrypt.py b/tests/unit/proxy/common_utils/test_proxy_encrypt_decrypt.py similarity index 100% rename from tests/proxy_unit_tests/test_proxy_encrypt_decrypt.py rename to tests/unit/proxy/common_utils/test_proxy_encrypt_decrypt.py diff --git a/tests/proxy_unit_tests/test_realtime_cache.py b/tests/unit/proxy/common_utils/test_realtime_cache.py similarity index 100% rename from tests/proxy_unit_tests/test_realtime_cache.py rename to tests/unit/proxy/common_utils/test_realtime_cache.py diff --git a/tests/unit/proxy/conftest.py b/tests/unit/proxy/conftest.py new file mode 100644 index 00000000000..148751c33f2 --- /dev/null +++ b/tests/unit/proxy/conftest.py @@ -0,0 +1,150 @@ +# conftest.py + +import asyncio +import copy +import inspect +import warnings + +import pytest + + +import litellm +import litellm.proxy.proxy_server + + +# Top-level assignments of these types are the ones importlib.reload(litellm) +# would have effectively reset. We snapshot them at conftest import time and +# deep-copy the snapshot back before every test. +_SNAPSHOT_TYPES = (list, dict, set, tuple, str, int, float, bool, bytes) + + +def _snapshot_mutable_state(module): + """Capture a per-module snapshot of primitive and collection attributes.""" + snapshot = {} + for attr in list(vars(module)): + if attr.startswith("_"): + continue + try: + value = getattr(module, attr) + except Exception as exc: + warnings.warn( + f"conftest: could not read {module.__name__}.{attr} during snapshot: {exc}", + stacklevel=2, + ) + continue + if value is None or isinstance(value, _SNAPSHOT_TYPES): + try: + snapshot[attr] = copy.deepcopy(value) + except Exception as exc: + warnings.warn( + f"conftest: could not snapshot {module.__name__}.{attr}: {exc}", + stacklevel=2, + ) + return snapshot + + +def _restore_mutable_state(module, snapshot): + for attr, default in snapshot.items(): + try: + setattr(module, attr, copy.deepcopy(default)) + except Exception as exc: + warnings.warn( + f"conftest: could not restore {module.__name__}.{attr}: {exc}", + stacklevel=2, + ) + + +def _collect_flushable_caches(): + """Return (module, attr) pairs whose values expose flush_cache().""" + targets = [] + for module in (litellm, litellm.proxy.proxy_server): + for attr in list(vars(module)): + if attr.startswith("_"): + continue + try: + value = getattr(module, attr) + except Exception: + continue + # Only instances — a class reference has an unbound flush_cache + # that can't be called without a self argument. + if inspect.isclass(value) or inspect.ismodule(value): + continue + if callable(getattr(value, "flush_cache", None)): + targets.append((module, attr)) + return targets + + +def _flush_caches(targets): + for module, attr in targets: + try: + value = getattr(module, attr) + except Exception: + continue + flush = getattr(value, "flush_cache", None) + if callable(flush): + try: + flush() + except Exception as exc: + warnings.warn( + f"conftest: flush_cache failed on {module.__name__}.{attr}: {exc}", + stacklevel=2, + ) + + +# Snapshot once at conftest import — these are the "clean" module states. +_LITELLM_STATE = _snapshot_mutable_state(litellm) +_PROXY_SERVER_STATE = _snapshot_mutable_state(litellm.proxy.proxy_server) +_FLUSHABLE_CACHES = _collect_flushable_caches() + + +@pytest.fixture(scope="function", autouse=True) +def setup_and_teardown(): + """Reset mutable module state on litellm and proxy_server before each test. + + Replaces a previous importlib.reload(litellm) approach that cost ~17s + per test (re-executing the full litellm __init__ import chain). + + What IS reset: + - Top-level module attributes of type list / dict / set / tuple + / str / int / float / bool / bytes, and None-valued attributes. + These cover callback lists, general_settings, master_key, + premium_user, prisma_client, etc. — anything the old reload() reset + by re-executing the module body. + - Any module-level object instance that exposes flush_cache() (the + DualCache and LLMClientCache family), which handles cache state + that can't round-trip through deepcopy because of internal locks. + + What is NOT reset: + - Class instances without flush_cache() (e.g. ProxyLogging, + JWTHandler, FastAPI routers, loggers). If a test mutates such an + instance in-place (setattr on the instance, appending to one of + its internal lists, etc.), the mutation will leak into later tests. + Use pytest's monkeypatch.setattr() or a local fixture for those + cases — don't rely on this autouse fixture to undo them. + """ + _restore_mutable_state(litellm, _LITELLM_STATE) + _restore_mutable_state(litellm.proxy.proxy_server, _PROXY_SERVER_STATE) + _flush_caches(_FLUSHABLE_CACHES) + + loop = asyncio.get_event_loop_policy().new_event_loop() + asyncio.set_event_loop(loop) + try: + yield + finally: + loop.close() + asyncio.set_event_loop(None) + + +def pytest_collection_modifyitems(config, items): + # Separate tests in 'test_amazing_proxy_custom_logger.py' and other tests + custom_logger_tests = [ + item for item in items if "custom_logger" in item.parent.name + ] + other_tests = [item for item in items if "custom_logger" not in item.parent.name] + + # Sort tests based on their names + custom_logger_tests.sort(key=lambda x: x.name) + other_tests.sort(key=lambda x: x.name) + + # Reorder the items list + items[:] = custom_logger_tests + other_tests diff --git a/tests/unit/proxy/db/__init__.py b/tests/unit/proxy/db/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/proxy/db/db_transaction_queue/__init__.py b/tests/unit/proxy/db/db_transaction_queue/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/proxy_unit_tests/test_e2e_pod_lock_manager.py b/tests/unit/proxy/db/db_transaction_queue/test_e2e_pod_lock_manager.py similarity index 100% rename from tests/proxy_unit_tests/test_e2e_pod_lock_manager.py rename to tests/unit/proxy/db/db_transaction_queue/test_e2e_pod_lock_manager.py diff --git a/tests/proxy_unit_tests/test_update_daily_tag_spend.py b/tests/unit/proxy/db/test_update_daily_tag_spend.py similarity index 100% rename from tests/proxy_unit_tests/test_update_daily_tag_spend.py rename to tests/unit/proxy/db/test_update_daily_tag_spend.py diff --git a/tests/unit/proxy/example_config_yaml/__init__.py b/tests/unit/proxy/example_config_yaml/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/proxy_unit_tests/example_config_yaml/aliases_config.yaml b/tests/unit/proxy/example_config_yaml/aliases_config.yaml similarity index 100% rename from tests/proxy_unit_tests/example_config_yaml/aliases_config.yaml rename to tests/unit/proxy/example_config_yaml/aliases_config.yaml diff --git a/tests/proxy_unit_tests/example_config_yaml/azure_config.yaml b/tests/unit/proxy/example_config_yaml/azure_config.yaml similarity index 100% rename from tests/proxy_unit_tests/example_config_yaml/azure_config.yaml rename to tests/unit/proxy/example_config_yaml/azure_config.yaml diff --git a/tests/proxy_unit_tests/example_config_yaml/cache_no_params.yaml b/tests/unit/proxy/example_config_yaml/cache_no_params.yaml similarity index 100% rename from tests/proxy_unit_tests/example_config_yaml/cache_no_params.yaml rename to tests/unit/proxy/example_config_yaml/cache_no_params.yaml diff --git a/tests/proxy_unit_tests/example_config_yaml/cache_with_params.yaml b/tests/unit/proxy/example_config_yaml/cache_with_params.yaml similarity index 100% rename from tests/proxy_unit_tests/example_config_yaml/cache_with_params.yaml rename to tests/unit/proxy/example_config_yaml/cache_with_params.yaml diff --git a/tests/proxy_unit_tests/example_config_yaml/config_with_env_vars.yaml b/tests/unit/proxy/example_config_yaml/config_with_env_vars.yaml similarity index 100% rename from tests/proxy_unit_tests/example_config_yaml/config_with_env_vars.yaml rename to tests/unit/proxy/example_config_yaml/config_with_env_vars.yaml diff --git a/tests/proxy_unit_tests/example_config_yaml/config_with_include.yaml b/tests/unit/proxy/example_config_yaml/config_with_include.yaml similarity index 100% rename from tests/proxy_unit_tests/example_config_yaml/config_with_include.yaml rename to tests/unit/proxy/example_config_yaml/config_with_include.yaml diff --git a/tests/proxy_unit_tests/example_config_yaml/config_with_missing_include.yaml b/tests/unit/proxy/example_config_yaml/config_with_missing_include.yaml similarity index 100% rename from tests/proxy_unit_tests/example_config_yaml/config_with_missing_include.yaml rename to tests/unit/proxy/example_config_yaml/config_with_missing_include.yaml diff --git a/tests/proxy_unit_tests/example_config_yaml/config_with_multiple_includes.yaml b/tests/unit/proxy/example_config_yaml/config_with_multiple_includes.yaml similarity index 100% rename from tests/proxy_unit_tests/example_config_yaml/config_with_multiple_includes.yaml rename to tests/unit/proxy/example_config_yaml/config_with_multiple_includes.yaml diff --git a/tests/proxy_unit_tests/example_config_yaml/included_models.yaml b/tests/unit/proxy/example_config_yaml/included_models.yaml similarity index 100% rename from tests/proxy_unit_tests/example_config_yaml/included_models.yaml rename to tests/unit/proxy/example_config_yaml/included_models.yaml diff --git a/tests/proxy_unit_tests/example_config_yaml/langfuse_config.yaml b/tests/unit/proxy/example_config_yaml/langfuse_config.yaml similarity index 100% rename from tests/proxy_unit_tests/example_config_yaml/langfuse_config.yaml rename to tests/unit/proxy/example_config_yaml/langfuse_config.yaml diff --git a/tests/proxy_unit_tests/example_config_yaml/load_balancer.yaml b/tests/unit/proxy/example_config_yaml/load_balancer.yaml similarity index 100% rename from tests/proxy_unit_tests/example_config_yaml/load_balancer.yaml rename to tests/unit/proxy/example_config_yaml/load_balancer.yaml diff --git a/tests/proxy_unit_tests/example_config_yaml/models_file_1.yaml b/tests/unit/proxy/example_config_yaml/models_file_1.yaml similarity index 100% rename from tests/proxy_unit_tests/example_config_yaml/models_file_1.yaml rename to tests/unit/proxy/example_config_yaml/models_file_1.yaml diff --git a/tests/proxy_unit_tests/example_config_yaml/models_file_2.yaml b/tests/unit/proxy/example_config_yaml/models_file_2.yaml similarity index 100% rename from tests/proxy_unit_tests/example_config_yaml/models_file_2.yaml rename to tests/unit/proxy/example_config_yaml/models_file_2.yaml diff --git a/tests/proxy_unit_tests/example_config_yaml/opentelemetry_config.yaml b/tests/unit/proxy/example_config_yaml/opentelemetry_config.yaml similarity index 100% rename from tests/proxy_unit_tests/example_config_yaml/opentelemetry_config.yaml rename to tests/unit/proxy/example_config_yaml/opentelemetry_config.yaml diff --git a/tests/proxy_unit_tests/example_config_yaml/simple_config.yaml b/tests/unit/proxy/example_config_yaml/simple_config.yaml similarity index 100% rename from tests/proxy_unit_tests/example_config_yaml/simple_config.yaml rename to tests/unit/proxy/example_config_yaml/simple_config.yaml diff --git a/tests/unit/proxy/google_endpoints/__init__.py b/tests/unit/proxy/google_endpoints/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/proxy_unit_tests/test_gemini_agents_endpoints.py b/tests/unit/proxy/google_endpoints/test_gemini_agents_endpoints.py similarity index 100% rename from tests/proxy_unit_tests/test_gemini_agents_endpoints.py rename to tests/unit/proxy/google_endpoints/test_gemini_agents_endpoints.py diff --git a/tests/proxy_unit_tests/test_google_endpoint_routing.py b/tests/unit/proxy/google_endpoints/test_google_endpoint_routing.py similarity index 100% rename from tests/proxy_unit_tests/test_google_endpoint_routing.py rename to tests/unit/proxy/google_endpoints/test_google_endpoint_routing.py diff --git a/tests/proxy_unit_tests/test_google_gemini_proxy_request.py b/tests/unit/proxy/google_endpoints/test_google_gemini_proxy_request.py similarity index 100% rename from tests/proxy_unit_tests/test_google_gemini_proxy_request.py rename to tests/unit/proxy/google_endpoints/test_google_gemini_proxy_request.py diff --git a/tests/unit/proxy/hooks/__init__.py b/tests/unit/proxy/hooks/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/proxy_unit_tests/test_banned_keyword_list.py b/tests/unit/proxy/hooks/test_banned_keyword_list.py similarity index 100% rename from tests/proxy_unit_tests/test_banned_keyword_list.py rename to tests/unit/proxy/hooks/test_banned_keyword_list.py diff --git a/tests/proxy_unit_tests/test_unit_test_max_model_budget_limiter.py b/tests/unit/proxy/hooks/test_unit_test_max_model_budget_limiter.py similarity index 100% rename from tests/proxy_unit_tests/test_unit_test_max_model_budget_limiter.py rename to tests/unit/proxy/hooks/test_unit_test_max_model_budget_limiter.py diff --git a/tests/unit/proxy/management_endpoints/__init__.py b/tests/unit/proxy/management_endpoints/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/proxy_unit_tests/test_jwt_key_mapping.py b/tests/unit/proxy/management_endpoints/test_jwt_key_mapping.py similarity index 99% rename from tests/proxy_unit_tests/test_jwt_key_mapping.py rename to tests/unit/proxy/management_endpoints/test_jwt_key_mapping.py index e95ed42013b..50b7a5c03fd 100644 --- a/tests/proxy_unit_tests/test_jwt_key_mapping.py +++ b/tests/unit/proxy/management_endpoints/test_jwt_key_mapping.py @@ -5,7 +5,7 @@ from datetime import datetime, timezone from unittest.mock import AsyncMock, MagicMock, patch # Add project root to sys.path -sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../.."))) +sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../.."))) from litellm.proxy.auth.user_api_key_auth import ( _resolve_jwt_to_virtual_key, diff --git a/tests/proxy_unit_tests/test_key_generate_prisma.py b/tests/unit/proxy/management_endpoints/test_key_generate_prisma.py similarity index 100% rename from tests/proxy_unit_tests/test_key_generate_prisma.py rename to tests/unit/proxy/management_endpoints/test_key_generate_prisma.py diff --git a/tests/unit/proxy/management_helpers/__init__.py b/tests/unit/proxy/management_helpers/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/proxy_unit_tests/test_audit_logs_proxy.py b/tests/unit/proxy/management_helpers/test_audit_logs_proxy.py similarity index 100% rename from tests/proxy_unit_tests/test_audit_logs_proxy.py rename to tests/unit/proxy/management_helpers/test_audit_logs_proxy.py diff --git a/tests/unit/proxy/middleware/__init__.py b/tests/unit/proxy/middleware/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/proxy_unit_tests/test_request_size_limit_middleware.py b/tests/unit/proxy/middleware/test_request_size_limit_middleware.py similarity index 100% rename from tests/proxy_unit_tests/test_request_size_limit_middleware.py rename to tests/unit/proxy/middleware/test_request_size_limit_middleware.py diff --git a/tests/unit/proxy/public_endpoints/__init__.py b/tests/unit/proxy/public_endpoints/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/proxy_unit_tests/test_blog_posts_endpoint.py b/tests/unit/proxy/public_endpoints/test_blog_posts_endpoint.py similarity index 100% rename from tests/proxy_unit_tests/test_blog_posts_endpoint.py rename to tests/unit/proxy/public_endpoints/test_blog_posts_endpoint.py diff --git a/tests/unit/proxy/response_polling/__init__.py b/tests/unit/proxy/response_polling/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/proxy_unit_tests/test_response_polling_handler.py b/tests/unit/proxy/response_polling/test_response_polling_handler.py similarity index 100% rename from tests/proxy_unit_tests/test_response_polling_handler.py rename to tests/unit/proxy/response_polling/test_response_polling_handler.py diff --git a/tests/unit/proxy/spend_tracking/__init__.py b/tests/unit/proxy/spend_tracking/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/proxy_unit_tests/test_search_api_logging.py b/tests/unit/proxy/spend_tracking/test_search_api_logging.py similarity index 100% rename from tests/proxy_unit_tests/test_search_api_logging.py rename to tests/unit/proxy/spend_tracking/test_search_api_logging.py diff --git a/tests/proxy_unit_tests/test_aproxy_startup.py b/tests/unit/proxy/test_aproxy_startup.py similarity index 100% rename from tests/proxy_unit_tests/test_aproxy_startup.py rename to tests/unit/proxy/test_aproxy_startup.py diff --git a/tests/unit/proxy/test_configs/__init__.py b/tests/unit/proxy/test_configs/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/proxy_unit_tests/test_configs/custom_auth.py b/tests/unit/proxy/test_configs/custom_auth.py similarity index 100% rename from tests/proxy_unit_tests/test_configs/custom_auth.py rename to tests/unit/proxy/test_configs/custom_auth.py diff --git a/tests/proxy_unit_tests/test_configs/test_cloudflare_azure_with_cache_config.yaml b/tests/unit/proxy/test_configs/test_cloudflare_azure_with_cache_config.yaml similarity index 100% rename from tests/proxy_unit_tests/test_configs/test_cloudflare_azure_with_cache_config.yaml rename to tests/unit/proxy/test_configs/test_cloudflare_azure_with_cache_config.yaml diff --git a/tests/proxy_unit_tests/test_configs/test_config.yaml b/tests/unit/proxy/test_configs/test_config.yaml similarity index 100% rename from tests/proxy_unit_tests/test_configs/test_config.yaml rename to tests/unit/proxy/test_configs/test_config.yaml diff --git a/tests/proxy_unit_tests/test_configs/test_config_custom_auth.yaml b/tests/unit/proxy/test_configs/test_config_custom_auth.yaml similarity index 100% rename from tests/proxy_unit_tests/test_configs/test_config_custom_auth.yaml rename to tests/unit/proxy/test_configs/test_config_custom_auth.yaml diff --git a/tests/proxy_unit_tests/test_configs/test_config_no_auth.yaml b/tests/unit/proxy/test_configs/test_config_no_auth.yaml similarity index 100% rename from tests/proxy_unit_tests/test_configs/test_config_no_auth.yaml rename to tests/unit/proxy/test_configs/test_config_no_auth.yaml diff --git a/tests/proxy_unit_tests/test_configs/test_guardrails_config.yaml b/tests/unit/proxy/test_configs/test_guardrails_config.yaml similarity index 100% rename from tests/proxy_unit_tests/test_configs/test_guardrails_config.yaml rename to tests/unit/proxy/test_configs/test_guardrails_config.yaml diff --git a/tests/proxy_unit_tests/test_custom_callback_input.py b/tests/unit/proxy/test_custom_callback_input.py similarity index 100% rename from tests/proxy_unit_tests/test_custom_callback_input.py rename to tests/unit/proxy/test_custom_callback_input.py diff --git a/tests/proxy_unit_tests/test_custom_logger_s3_gcs.py b/tests/unit/proxy/test_custom_logger_s3_gcs.py similarity index 100% rename from tests/proxy_unit_tests/test_custom_logger_s3_gcs.py rename to tests/unit/proxy/test_custom_logger_s3_gcs.py diff --git a/tests/proxy_unit_tests/test_custom_tokenizer_bug.py b/tests/unit/proxy/test_custom_tokenizer_bug.py similarity index 100% rename from tests/proxy_unit_tests/test_custom_tokenizer_bug.py rename to tests/unit/proxy/test_custom_tokenizer_bug.py diff --git a/tests/proxy_unit_tests/test_db_schema_changes.py b/tests/unit/proxy/test_db_schema_changes.py similarity index 100% rename from tests/proxy_unit_tests/test_db_schema_changes.py rename to tests/unit/proxy/test_db_schema_changes.py diff --git a/tests/proxy_unit_tests/test_deprecated_key_grace_period.py b/tests/unit/proxy/test_deprecated_key_grace_period.py similarity index 100% rename from tests/proxy_unit_tests/test_deprecated_key_grace_period.py rename to tests/unit/proxy/test_deprecated_key_grace_period.py diff --git a/tests/proxy_unit_tests/test_get_favicon.py b/tests/unit/proxy/test_get_favicon.py similarity index 100% rename from tests/proxy_unit_tests/test_get_favicon.py rename to tests/unit/proxy/test_get_favicon.py diff --git a/tests/proxy_unit_tests/test_get_image.py b/tests/unit/proxy/test_get_image.py similarity index 100% rename from tests/proxy_unit_tests/test_get_image.py rename to tests/unit/proxy/test_get_image.py diff --git a/tests/proxy_unit_tests/test_prisma_client_backoff_retry.py b/tests/unit/proxy/test_prisma_client_backoff_retry.py similarity index 100% rename from tests/proxy_unit_tests/test_prisma_client_backoff_retry.py rename to tests/unit/proxy/test_prisma_client_backoff_retry.py diff --git a/tests/proxy_unit_tests/test_prompt_test_endpoint.py b/tests/unit/proxy/test_prompt_test_endpoint.py similarity index 100% rename from tests/proxy_unit_tests/test_prompt_test_endpoint.py rename to tests/unit/proxy/test_prompt_test_endpoint.py diff --git a/tests/proxy_unit_tests/test_proxy_config_unit_test.py b/tests/unit/proxy/test_proxy_config_unit_test.py similarity index 99% rename from tests/proxy_unit_tests/test_proxy_config_unit_test.py rename to tests/unit/proxy/test_proxy_config_unit_test.py index 5f236806685..2181c932586 100644 --- a/tests/proxy_unit_tests/test_proxy_config_unit_test.py +++ b/tests/unit/proxy/test_proxy_config_unit_test.py @@ -31,7 +31,7 @@ async def test_basic_reading_configs_from_files(): example_config_yaml_path = os.path.join(current_path, "example_config_yaml") # get all the files from example_config_yaml - files = os.listdir(example_config_yaml_path) + files = [f for f in os.listdir(example_config_yaml_path) if f.endswith((".yaml", ".yml"))] print(files) for file in files: diff --git a/tests/proxy_unit_tests/test_proxy_custom_auth.py b/tests/unit/proxy/test_proxy_custom_auth.py similarity index 100% rename from tests/proxy_unit_tests/test_proxy_custom_auth.py rename to tests/unit/proxy/test_proxy_custom_auth.py diff --git a/tests/proxy_unit_tests/test_proxy_reject_logging.py b/tests/unit/proxy/test_proxy_reject_logging.py similarity index 100% rename from tests/proxy_unit_tests/test_proxy_reject_logging.py rename to tests/unit/proxy/test_proxy_reject_logging.py diff --git a/tests/proxy_unit_tests/test_proxy_server.py b/tests/unit/proxy/test_proxy_server.py similarity index 99% rename from tests/proxy_unit_tests/test_proxy_server.py rename to tests/unit/proxy/test_proxy_server.py index 5be27b3ad72..eae80f311d8 100644 --- a/tests/proxy_unit_tests/test_proxy_server.py +++ b/tests/unit/proxy/test_proxy_server.py @@ -477,7 +477,7 @@ async def test_team_disable_guardrails(mock_acompletion, client_no_auth): assert e.code == str(403) -from test_custom_callback_input import CompletionCustomHandler +from tests.unit.proxy.test_custom_callback_input import CompletionCustomHandler @mock_patch_acompletion() @@ -1114,7 +1114,7 @@ from litellm.proxy._types import ( ) from litellm.proxy.management_endpoints.internal_user_endpoints import new_user from litellm.proxy.management_endpoints.team_endpoints import team_member_add -from test_key_generate_prisma import prisma_client +from tests.unit.proxy.management_endpoints.test_key_generate_prisma import prisma_client @pytest.fixture diff --git a/tests/proxy_unit_tests/test_proxy_setting_guardrails.py b/tests/unit/proxy/test_proxy_setting_guardrails.py similarity index 100% rename from tests/proxy_unit_tests/test_proxy_setting_guardrails.py rename to tests/unit/proxy/test_proxy_setting_guardrails.py diff --git a/tests/proxy_unit_tests/test_proxy_token_counter.py b/tests/unit/proxy/test_proxy_token_counter.py similarity index 100% rename from tests/proxy_unit_tests/test_proxy_token_counter.py rename to tests/unit/proxy/test_proxy_token_counter.py diff --git a/tests/proxy_unit_tests/test_proxy_utils.py b/tests/unit/proxy/test_proxy_utils.py similarity index 100% rename from tests/proxy_unit_tests/test_proxy_utils.py rename to tests/unit/proxy/test_proxy_utils.py diff --git a/tests/proxy_unit_tests/test_reducto_ocr_route.py b/tests/unit/proxy/test_reducto_ocr_route.py similarity index 100% rename from tests/proxy_unit_tests/test_reducto_ocr_route.py rename to tests/unit/proxy/test_reducto_ocr_route.py diff --git a/tests/proxy_unit_tests/test_response_polling_pre_call_checks.py b/tests/unit/proxy/test_response_polling_pre_call_checks.py similarity index 100% rename from tests/proxy_unit_tests/test_response_polling_pre_call_checks.py rename to tests/unit/proxy/test_response_polling_pre_call_checks.py diff --git a/tests/proxy_unit_tests/test_server_root_path.py b/tests/unit/proxy/test_server_root_path.py similarity index 100% rename from tests/proxy_unit_tests/test_server_root_path.py rename to tests/unit/proxy/test_server_root_path.py diff --git a/tests/proxy_unit_tests/test_ui_path_detection.py b/tests/unit/proxy/test_ui_path_detection.py similarity index 100% rename from tests/proxy_unit_tests/test_ui_path_detection.py rename to tests/unit/proxy/test_ui_path_detection.py diff --git a/tests/proxy_unit_tests/test_unit_test_proxy_hooks.py b/tests/unit/proxy/test_unit_test_proxy_hooks.py similarity index 100% rename from tests/proxy_unit_tests/test_unit_test_proxy_hooks.py rename to tests/unit/proxy/test_unit_test_proxy_hooks.py diff --git a/tests/proxy_unit_tests/test_update_spend.py b/tests/unit/proxy/test_update_spend.py similarity index 100% rename from tests/proxy_unit_tests/test_update_spend.py rename to tests/unit/proxy/test_update_spend.py diff --git a/tests/proxy_unit_tests/test_zero_cost_model_budget_bypass.py b/tests/unit/proxy/test_zero_cost_model_budget_bypass.py similarity index 100% rename from tests/proxy_unit_tests/test_zero_cost_model_budget_bypass.py rename to tests/unit/proxy/test_zero_cost_model_budget_bypass.py diff --git a/tests/proxy_unit_tests/vertex_key.json b/tests/unit/proxy/vertex_key.json similarity index 100% rename from tests/proxy_unit_tests/vertex_key.json rename to tests/unit/proxy/vertex_key.json diff --git a/tests/proxy_unit_tests/test_skills_db.py b/tests/unit/skills/test_skills_db.py similarity index 98% rename from tests/proxy_unit_tests/test_skills_db.py rename to tests/unit/skills/test_skills_db.py index 8eb07a5ad48..20ffed6fec1 100644 --- a/tests/proxy_unit_tests/test_skills_db.py +++ b/tests/unit/skills/test_skills_db.py @@ -42,7 +42,7 @@ def create_skill_zip(skill_name: str): The zip file is automatically cleaned up after use. """ - test_dir = Path(__file__).parent.parent / "llm_translation" / "test_skills_data" + test_dir = Path(__file__).parents[2] / "llm_translation" / "test_skills_data" skill_dir = test_dir / skill_name # Create a zip file containing the skill directory diff --git a/tests/unit/skills/test_skills_main.py b/tests/unit/skills/test_skills_main.py index e1c66c8d9ea..71d65d45a08 100644 --- a/tests/unit/skills/test_skills_main.py +++ b/tests/unit/skills/test_skills_main.py @@ -30,7 +30,7 @@ def test_create_skill_forwards_description_and_instructions_from_top_level_kwarg def test_create_skill_forwards_description_and_instructions_from_extra_body(monkeypatch) -> None: - """The SDK convention (see tests/proxy_unit_tests/test_skills_db.py) nests them under + """The SDK convention (see tests/unit/skills/test_skills_db.py) nests them under extra_body instead of passing them as top-level kwargs; both paths must reach the DB.""" handler = MagicMock() monkeypatch.setattr(skills_main, "_get_litellm_skills_handler", lambda: handler) From ba776469916bcdb35ab238dec92f33ad428d724d Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 24 Sep 2026 23:07:48 +0000 Subject: [PATCH 096/120] ci: move provider-independent MCP tests into tests/unit and run mcp-integration from litellm-tests (#42904) * ci: fix the litellm-tests unit job with sysmon coverage, an env allowlist and coverage upload on failure * test: replace key-dependent proxy, enterprise and mcp unit tests with synthetic values and integration and e2e coverage * test: drop key reads at the legacy proxy, enterprise and mcp paths and wire the gemini pass-through split * ci: move caching, proxy-extras, gateway and enterprise tests into tests/unit and run them from litellm-tests under their legacy flags * ci: move caching, proxy-extras, gateway and enterprise tests into tests/unit and run them from litellm-tests under their legacy flags * ci: move tests/proxy_unit_tests to tests/unit/proxy and run the proxy-db shards from litellm-tests * ci: move provider-independent MCP tests into tests/unit and run mcp-integration from litellm-tests * ci: fail the unit shard when circleci tests split errors * test: drop restating comments from the gemini pass-through split * build: point the local proxy unit targets at the nested tests/unit/proxy tree * ci: exit the unit shard cleanly when circleci tests split assigns it no files --------- Co-authored-by: yuneng --- .circleci/scripts/unit_selection.sh | 5 ++ .circleci/tests.yml | 21 +++++ .github/workflows/test-unit.yml | 1 + tests/unit/proxy/_experimental/__init__.py | 0 .../_experimental/mcp_server/__init__.py | 0 .../_experimental/mcp_server/conftest.py | 78 +++++++++++++++++++ .../test_mcp_auth_header_extraction.py | 0 .../mcp_server}/test_mcp_auth_priority.py | 0 .../mcp_server}/test_mcp_chat_completions.py | 0 .../mcp_server}/test_mcp_client_unit.py | 0 .../mcp_server}/test_mcp_logging.py | 0 .../mcp_server}/test_mcp_server.py | 0 .../mcp_server}/test_oauth2_mcp_config.yaml | 0 .../mcp_server}/test_openapi_spec_path_url.py | 0 .../mcp_server}/test_per_user_oauth_cache.py | 0 tests/unit/responses/__init__.py | 0 tests/unit/responses/mcp/__init__.py | 0 .../mcp}/test_aresponses_api_with_mcp.py | 0 18 files changed, 105 insertions(+) create mode 100644 tests/unit/proxy/_experimental/__init__.py create mode 100644 tests/unit/proxy/_experimental/mcp_server/__init__.py create mode 100644 tests/unit/proxy/_experimental/mcp_server/conftest.py rename tests/{mcp_tests => unit/proxy/_experimental/mcp_server}/test_mcp_auth_header_extraction.py (100%) rename tests/{mcp_tests => unit/proxy/_experimental/mcp_server}/test_mcp_auth_priority.py (100%) rename tests/{mcp_tests => unit/proxy/_experimental/mcp_server}/test_mcp_chat_completions.py (100%) rename tests/{mcp_tests => unit/proxy/_experimental/mcp_server}/test_mcp_client_unit.py (100%) rename tests/{mcp_tests => unit/proxy/_experimental/mcp_server}/test_mcp_logging.py (100%) rename tests/{mcp_tests => unit/proxy/_experimental/mcp_server}/test_mcp_server.py (100%) rename tests/{mcp_tests => unit/proxy/_experimental/mcp_server}/test_oauth2_mcp_config.yaml (100%) rename tests/{mcp_tests => unit/proxy/_experimental/mcp_server}/test_openapi_spec_path_url.py (100%) rename tests/{mcp_tests => unit/proxy/_experimental/mcp_server}/test_per_user_oauth_cache.py (100%) create mode 100644 tests/unit/responses/__init__.py create mode 100644 tests/unit/responses/mcp/__init__.py rename tests/{mcp_tests => unit/responses/mcp}/test_aresponses_api_with_mcp.py (100%) diff --git a/.circleci/scripts/unit_selection.sh b/.circleci/scripts/unit_selection.sh index 2c60c5b1334..f2ee7550df3 100755 --- a/.circleci/scripts/unit_selection.sh +++ b/.circleci/scripts/unit_selection.sh @@ -7,6 +7,7 @@ legacy_flags=( caching-local enterprise-package enterprise-routing + mcp-integration proxy-db-auth-checks proxy-db-budgets proxy-db-custom-logging @@ -46,6 +47,10 @@ legacy_paths() { echo tests/unit/enterprise/proxy/test_file_deletion_blocking.py echo tests/unit/enterprise/proxy/test_managed_files_access_check.py echo tests/unit/enterprise/proxy/test_managed_files_hook.py ;; + mcp-integration) + echo tests/unit/proxy/_experimental/mcp_server + echo tests/unit/responses/mcp + echo tests/mcp_tests/test_proxy_mcp_e2e.py ;; proxy-db-auth-checks) echo tests/unit/proxy/auth/test_auth_checks.py echo tests/unit/proxy/auth/test_user_api_key_auth.py diff --git a/.circleci/tests.yml b/.circleci/tests.yml index 08e735637b7..264d7695a94 100644 --- a/.circleci/tests.yml +++ b/.circleci/tests.yml @@ -183,6 +183,9 @@ jobs: pull_request_url: type: string default: "" + legacy_mcp_peer: + type: boolean + default: false reruns: type: integer default: 0 @@ -200,6 +203,15 @@ jobs: base_ref: << parameters.base_ref >> pull_request_url: << parameters.pull_request_url >> - setup_test_deps + - when: + condition: << parameters.legacy_mcp_peer >> + steps: + - run: + name: Install MCP SDK1 peer + command: | + uv venv --python 3.12 .venv-mcp-peer + uv pip install --python .venv-mcp-peer 'mcp==1.28.1' 'langchain-mcp-adapters==0.2.1' + echo "export MCP_TEST_PEER_PYTHON=$PWD/.venv-mcp-peer/bin/python" >> "$BASH_ENV" - run: name: "Run << parameters.flag >> shard" no_output_timeout: 20m @@ -215,6 +227,7 @@ jobs: rerun_args=(-p no:rerunfailures) if [ "<< parameters.reruns >>" -gt 0 ]; then rerun_args=(--reruns << parameters.reruns >> --reruns-delay 1 --rerun-except "from pytest-timeout"); fi test_env=(PATH="$PATH" HOME="$HOME" CI=true COVERAGE_CORE="$COVERAGE_CORE" LITELLM_LOCAL_MODEL_COST_MAP="$LITELLM_LOCAL_MODEL_COST_MAP") + if [ -n "${MCP_TEST_PEER_PYTHON:-}" ]; then test_env+=(MCP_TEST_PEER_PYTHON="$MCP_TEST_PEER_PYTHON"); fi set +e env -i "${test_env[@]}" \ uv run --no-sync pytest "${files[@]}" "${rerun_args[@]}" -p no:pytest-retry --timeout=90 "${xdist_args[@]}" --tb=short --durations=20 -o junit_family=xunit1 --junitxml=test-results/<< parameters.flag >>/junit.xml --cov=./litellm --cov=./enterprise/litellm_enterprise --cov-report=xml:coverage.xml --cov-config=pyproject.toml @@ -311,6 +324,14 @@ workflows: flag: [caching-local, proxy-extras, enterprise-routing] base_ref: << pipeline.event.name == "pull_request" and pipeline.event.github.pull_request.base.ref or "" >> pull_request_url: << pipeline.event.name == "pull_request" and pipeline.event.github.pull_request.url or "" >> + - unit: + name: unit-mcp-integration + flag: mcp-integration + shards: 1 + workers: 2 + legacy_mcp_peer: true + base_ref: << pipeline.event.name == "pull_request" and pipeline.event.github.pull_request.base.ref or "" >> + pull_request_url: << pipeline.event.name == "pull_request" and pipeline.event.github.pull_request.url or "" >> - unit: name: unit-<< matrix.flag >> shards: 1 diff --git a/.github/workflows/test-unit.yml b/.github/workflows/test-unit.yml index bf2e1602be8..126a6e26e6f 100644 --- a/.github/workflows/test-unit.yml +++ b/.github/workflows/test-unit.yml @@ -53,6 +53,7 @@ jobs: - shard: mcp-integration artifact-name: mcp-integration test-path: "tests/mcp_tests tests/test_litellm/experimental_mcp_client" + fork-flag: mcp-integration workers: 2 reruns: 0 timeout-minutes: 20 diff --git a/tests/unit/proxy/_experimental/__init__.py b/tests/unit/proxy/_experimental/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/proxy/_experimental/mcp_server/__init__.py b/tests/unit/proxy/_experimental/mcp_server/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/proxy/_experimental/mcp_server/conftest.py b/tests/unit/proxy/_experimental/mcp_server/conftest.py new file mode 100644 index 00000000000..d8b91e07467 --- /dev/null +++ b/tests/unit/proxy/_experimental/mcp_server/conftest.py @@ -0,0 +1,78 @@ +import asyncio +import importlib + +import pytest + +import litellm +from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER + + +@pytest.fixture(scope="session") +def event_loop(): + try: + loop = asyncio.get_running_loop() + except RuntimeError: + loop = asyncio.new_event_loop() + yield loop + loop.close() + + +@pytest.fixture(scope="function", autouse=True) +def setup_and_teardown(): + """ + This fixture reloads litellm before every function. To speed up testing by removing callbacks being chained. + """ + importlib.reload(litellm) + import asyncio + + loop = asyncio.get_event_loop_policy().new_event_loop() + asyncio.set_event_loop(loop) + yield + + # Teardown code (executes after the yield point) + # LoggingWorker carries still-queued coroutines onto the next test's loop, where they'd log into that test's callbacks + asyncio.run(GLOBAL_LOGGING_WORKER.clear_queue()) + loop.close() # Close the loop created earlier + asyncio.set_event_loop(None) # Remove the reference to the loop + + +@pytest.fixture(scope="function", autouse=True) +async def drain_logging_worker(): + """ + The logging queue is bound to the running loop, so anything left queued when a test's loop + goes away is carried onto the next test's loop and fires against its callbacks. + """ + from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER + + yield + + try: + await asyncio.wait_for(GLOBAL_LOGGING_WORKER.clear_queue(), timeout=10) + except asyncio.TimeoutError: + pass + + +def pytest_collection_modifyitems(config, items): + # Separate tests in 'test_amazing_proxy_custom_logger.py' and other tests + custom_logger_tests = [ + item for item in items if "custom_logger" in item.parent.name + ] + other_tests = [item for item in items if "custom_logger" not in item.parent.name] + + # Sort tests based on their names + custom_logger_tests.sort(key=lambda x: x.name) + other_tests.sort(key=lambda x: x.name) + + # Reorder the items list + items[:] = custom_logger_tests + other_tests + + +@pytest.fixture +def config_only_mcp_manager_factory(): + from litellm.proxy._experimental.mcp_server.mcp_server_manager import MCPServerManager + + class ConfigOnlyManager(MCPServerManager): + def initialize_tool_name_to_mcp_server_name_mapping(self): + return None + + return ConfigOnlyManager diff --git a/tests/mcp_tests/test_mcp_auth_header_extraction.py b/tests/unit/proxy/_experimental/mcp_server/test_mcp_auth_header_extraction.py similarity index 100% rename from tests/mcp_tests/test_mcp_auth_header_extraction.py rename to tests/unit/proxy/_experimental/mcp_server/test_mcp_auth_header_extraction.py diff --git a/tests/mcp_tests/test_mcp_auth_priority.py b/tests/unit/proxy/_experimental/mcp_server/test_mcp_auth_priority.py similarity index 100% rename from tests/mcp_tests/test_mcp_auth_priority.py rename to tests/unit/proxy/_experimental/mcp_server/test_mcp_auth_priority.py diff --git a/tests/mcp_tests/test_mcp_chat_completions.py b/tests/unit/proxy/_experimental/mcp_server/test_mcp_chat_completions.py similarity index 100% rename from tests/mcp_tests/test_mcp_chat_completions.py rename to tests/unit/proxy/_experimental/mcp_server/test_mcp_chat_completions.py diff --git a/tests/mcp_tests/test_mcp_client_unit.py b/tests/unit/proxy/_experimental/mcp_server/test_mcp_client_unit.py similarity index 100% rename from tests/mcp_tests/test_mcp_client_unit.py rename to tests/unit/proxy/_experimental/mcp_server/test_mcp_client_unit.py diff --git a/tests/mcp_tests/test_mcp_logging.py b/tests/unit/proxy/_experimental/mcp_server/test_mcp_logging.py similarity index 100% rename from tests/mcp_tests/test_mcp_logging.py rename to tests/unit/proxy/_experimental/mcp_server/test_mcp_logging.py diff --git a/tests/mcp_tests/test_mcp_server.py b/tests/unit/proxy/_experimental/mcp_server/test_mcp_server.py similarity index 100% rename from tests/mcp_tests/test_mcp_server.py rename to tests/unit/proxy/_experimental/mcp_server/test_mcp_server.py diff --git a/tests/mcp_tests/test_oauth2_mcp_config.yaml b/tests/unit/proxy/_experimental/mcp_server/test_oauth2_mcp_config.yaml similarity index 100% rename from tests/mcp_tests/test_oauth2_mcp_config.yaml rename to tests/unit/proxy/_experimental/mcp_server/test_oauth2_mcp_config.yaml diff --git a/tests/mcp_tests/test_openapi_spec_path_url.py b/tests/unit/proxy/_experimental/mcp_server/test_openapi_spec_path_url.py similarity index 100% rename from tests/mcp_tests/test_openapi_spec_path_url.py rename to tests/unit/proxy/_experimental/mcp_server/test_openapi_spec_path_url.py diff --git a/tests/mcp_tests/test_per_user_oauth_cache.py b/tests/unit/proxy/_experimental/mcp_server/test_per_user_oauth_cache.py similarity index 100% rename from tests/mcp_tests/test_per_user_oauth_cache.py rename to tests/unit/proxy/_experimental/mcp_server/test_per_user_oauth_cache.py diff --git a/tests/unit/responses/__init__.py b/tests/unit/responses/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/responses/mcp/__init__.py b/tests/unit/responses/mcp/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/mcp_tests/test_aresponses_api_with_mcp.py b/tests/unit/responses/mcp/test_aresponses_api_with_mcp.py similarity index 100% rename from tests/mcp_tests/test_aresponses_api_with_mcp.py rename to tests/unit/responses/mcp/test_aresponses_api_with_mcp.py From 3fa688223d97bb46345c8ed8032b583e03e29218 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 24 Sep 2026 16:15:58 -0700 Subject: [PATCH 097/120] fix(vertex_ai): translate /v1/responses batch rows through the Responses-to-Chat bridge (#43042) * fix(vertex_ai): translate /v1/responses batch rows through the Responses-to-Chat bridge Vertex batch uploads treated every non-embeddings JSONL row as a chat completions body, so a /v1/responses row lost its input and reached GCS as a blank text part. Route detection now recognizes /v1/responses rows and bridges them to chat through the same Responses-to-Chat bridge the real-time path uses. That bridge call moves out of the Bedrock files transformation into a shared helper both providers call, forwarding the record's fields as sent, like real time, instead of validating them against the SDK TypedDicts whose required keys clients omit. * chore(batches): type the Vertex responses test helper and drop the quoted input cast * fix(batches): translate developer messages to system on Vertex and Bedrock batch rows like real time --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- litellm/llms/base_llm/files/batch_records.py | 54 +++++++ litellm/llms/bedrock/files/transformation.py | 53 +----- .../llms/vertex_ai/files/transformation.py | 32 +++- .../test_bedrock_files_transformation.py | 43 +++++ .../test_vertex_ai_files_transformation.py | 153 ++++++++++++++++++ 5 files changed, 280 insertions(+), 55 deletions(-) create mode 100644 litellm/llms/base_llm/files/batch_records.py diff --git a/litellm/llms/base_llm/files/batch_records.py b/litellm/llms/base_llm/files/batch_records.py new file mode 100644 index 00000000000..6bb98456e69 --- /dev/null +++ b/litellm/llms/base_llm/files/batch_records.py @@ -0,0 +1,54 @@ +from collections.abc import Iterable, Mapping +from functools import cache +from types import MappingProxyType +from typing import Final, cast, get_type_hints + +from litellm.types.llms.openai import ResponseInputParam, ResponsesAPIOptionalRequestParams + + +def _frozen_mapping(items: Iterable[tuple[str, object]]) -> Mapping[str, object]: + return MappingProxyType(dict(items)) + + +@cache +def _responses_request_keys() -> frozenset[str]: + return frozenset(get_type_hints(ResponsesAPIOptionalRequestParams)) + + +def responses_batch_body_to_chat_body( + openai_request_body: Mapping[str, object], + custom_llm_provider: str | None = None, +) -> dict[str, object]: # mutable-ok: provider transforms take the bridged chat body as a plain dict + """ + Rewrite the body of an OpenAI `/v1/responses` batch record as a Chat Completions body. + + Batch providers translate chat bodies into their own request shape, so a Responses + record goes through the same Responses-to-Chat bridge the real-time path uses for + providers without a native Responses API: `input`, `instructions`, `max_output_tokens` + and the tool params translate identically in batch and real time. Like real time, the + record's fields are forwarded as sent instead of validated against the SDK TypedDicts, + whose required keys (a function tool's `strict`, an image part's `detail`) clients omit. + """ + from litellm.responses.litellm_completion_transformation.transformation import ( + LiteLLMCompletionResponsesConfig, + ) + + responses_input: Final = openai_request_body.get("input") + if responses_input is None: + raise ValueError( + "Batch record for /v1/responses is missing required `input` field: " + f"model={openai_request_body.get('model', '')}" + ) + model: Final = openai_request_body.get("model") + chat_input: Final = cast(str | ResponseInputParam, responses_input) # cast-ok: forwarded as sent + responses_request: Final = cast( # cast-ok: client-supplied fields forwarded verbatim, as real time does + ResponsesAPIOptionalRequestParams, + _frozen_mapping((key, value) for key, value in openai_request_body.items() if key in _responses_request_keys()), + ) + return LiteLLMCompletionResponsesConfig.transform_responses_api_request_to_chat_completion_request( # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # transformer declares a bare dict return + model=model if isinstance(model, str) else "", + input=chat_input, + responses_api_request=responses_request, + custom_llm_provider=custom_llm_provider, + metadata=openai_request_body.get("metadata"), + ) diff --git a/litellm/llms/bedrock/files/transformation.py b/litellm/llms/bedrock/files/transformation.py index 43faa7d79ea..a79f4de1e3d 100644 --- a/litellm/llms/bedrock/files/transformation.py +++ b/litellm/llms/bedrock/files/transformation.py @@ -8,7 +8,6 @@ from collections.abc import Iterable, Mapping, MutableMapping, Sequence from contextlib import suppress from dataclasses import dataclass from datetime import datetime -from functools import cache from itertools import chain from types import MappingProxyType from typing import Any, Final, Literal, TypeAlias, TypedDict @@ -17,7 +16,7 @@ from urllib.parse import quote, unquote, urlencode import httpx from httpx import Headers, Response from openai.types.file_deleted import FileDeleted -from pydantic import BaseModel, ConfigDict, Field, TypeAdapter +from pydantic import BaseModel, ConfigDict, Field from typing_extensions import ReadOnly from litellm._logging import verbose_logger @@ -41,7 +40,9 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( extract_file_data, text_completion_prompt_to_messages, ) +from litellm.llms.base_llm.base_utils import map_developer_role_to_system_role from litellm.llms.base_llm.chat.transformation import BaseLLMException +from litellm.llms.base_llm.files.batch_records import responses_batch_body_to_chat_body from litellm.llms.base_llm.files.transformation import ( BaseFilesConfig, LiteLLMLoggingObj, @@ -56,8 +57,6 @@ from litellm.types.llms.openai import ( OpenAICreateFileRequestOptionalParams, OpenAIFileObject, PathLike, - ResponseInputParam, - ResponsesAPIOptionalRequestParams, ) from litellm.types.utils import ExtractedFileData, LlmProviders, SpecialEnums from litellm.utils import get_llm_provider @@ -130,22 +129,6 @@ class _S3UploadResponse(TypedDict, total=False): ContentLength: ReadOnly[int] -# JSONL batch records are untyped json, so the `/v1/responses` fields are -# validated into their concrete Responses API types before being handed to the -# Responses-to-Chat bridge. Both adapters drop keys the Responses API doesn't -# define, which is what the bridge would ignore anyway. Built on first use -# rather than at import: `ResponseInputParam` is a deep union and only batch -# files carrying `/v1/responses` records need it. -@cache -def _responses_input_adapter() -> TypeAdapter[str | ResponseInputParam]: - return TypeAdapter(str | ResponseInputParam) - - -@cache -def _responses_request_adapter() -> TypeAdapter[ResponsesAPIOptionalRequestParams]: - return TypeAdapter(ResponsesAPIOptionalRequestParams) - - class _BedrockS3RequestParams(AwsAuthParams): """Typed view of the credential/region params the S3 GetObject path reads.""" @@ -859,33 +842,9 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): Delegates to the same Responses-to-Chat bridge the real-time path uses for providers without a native Responses API (which is every Bedrock model), so `input`, `instructions`, `max_output_tokens` and the tool - params translate identically in batch and real time. The bridge always - emits a `tools` key; an empty one is dropped rather than shipped as an - empty array inside `modelInput`. + params translate identically in batch and real time. """ - from litellm.responses.litellm_completion_transformation.transformation import ( - LiteLLMCompletionResponsesConfig, - ) - - responses_input: Final = openai_request_body.get("input") - if responses_input is None: - raise ValueError( - "Batch record for /v1/responses is missing required `input` field: " - f"model={openai_request_body.get('model', '')}" - ) - chat_body: Final[Mapping[str, object]] = ( - LiteLLMCompletionResponsesConfig.transform_responses_api_request_to_chat_completion_request( - model=openai_request_body.get("model", ""), - input=_responses_input_adapter().validate_python(responses_input), - responses_api_request=_responses_request_adapter().validate_python( - _frozen_mapping( - (key, value) for key, value in openai_request_body.items() if key not in ("model", "input") - ) - ), - metadata=openai_request_body.get("metadata"), - ) - ) - return _frozen_mapping((key, value) for key, value in chat_body.items() if key != "tools" or value) + return responses_batch_body_to_chat_body(openai_request_body) @staticmethod def _transform_batch_body_to_chat_body( @@ -922,7 +881,7 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): """ from litellm.types.utils import LlmProviders - messages: Final = openai_request_body.get("messages", []) + messages: Final = map_developer_role_to_system_role(openai_request_body.get("messages", [])) optional_params: Final = {k: v for k, v in openai_request_body.items() if k not in ["model", "messages"]} # --- Anthropic: use existing AmazonAnthropicClaudeConfig --- diff --git a/litellm/llms/vertex_ai/files/transformation.py b/litellm/llms/vertex_ai/files/transformation.py index 789b36ef3d0..dbb41b57348 100644 --- a/litellm/llms/vertex_ai/files/transformation.py +++ b/litellm/llms/vertex_ai/files/transformation.py @@ -35,7 +35,9 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( extract_file_data, extract_file_metadata, ) +from litellm.llms.base_llm.base_utils import map_developer_role_to_system_role from litellm.llms.base_llm.chat.transformation import BaseLLMException +from litellm.llms.base_llm.files.batch_records import responses_batch_body_to_chat_body from litellm.llms.base_llm.files.transformation import ( BaseFilesConfig, BaseFileUploadStream, @@ -529,21 +531,30 @@ def is_passthrough_batch_upload(create_file_data: Mapping[str, object], litellm_ return create_file_data.get("purpose") == "batch" and litellm_params.get("passthrough") is True -def _is_embeddings_batch_entry(openai_entry: Mapping[str, object]) -> bool: +def _batch_entry_route_path(openai_entry: Mapping[str, object]) -> str: """ - Whether an OpenAI batch JSONL line targets the embeddings endpoint. + The route an OpenAI batch JSONL line targets, without query string or trailing slash. OpenAI puts the target route on each line's `url` (e.g. `/v1/embeddings`); Vertex has no equivalent per-line field, so the route decides which Vertex request shape the line has to be translated into. """ - url = openai_entry.get("url") + url: Final = openai_entry.get("url") if not isinstance(url, str): - return False - path = url.split("?")[0].rstrip("/") + return "" + return url.split("?")[0].rstrip("/") + + +def _is_embeddings_batch_entry(openai_entry: Mapping[str, object]) -> bool: + path: Final = _batch_entry_route_path(openai_entry) return path == "embeddings" or path.endswith("/embeddings") +def _is_responses_batch_entry(openai_entry: Mapping[str, object]) -> bool: + path: Final = _batch_entry_route_path(openai_entry) + return path == "responses" or path.endswith("/responses") + + def _openai_embedding_input_elements( embedding_input: GeminiEmbeddingInput, ) -> tuple[str | list[str], ...]: @@ -665,10 +676,15 @@ def _openai_batch_jsonl_entry_to_vertex_rows( return _openai_batch_jsonl_entry_to_vertex_embeddings_rows(openai_entry) openai_request_body: Final = openai_entry.get("body") or {} + chat_request_body: Final = ( + responses_batch_body_to_chat_body(openai_request_body, custom_llm_provider="vertex_ai") + if _is_responses_batch_entry(openai_entry) + else openai_request_body + ) vertex_request_body: Final = _transform_request_body( - messages=openai_request_body.get("messages", []), - model=openai_request_body.get("model", ""), - optional_params=map_openai_to_vertex_params(openai_request_body), + messages=map_developer_role_to_system_role(chat_request_body.get("messages", [])), + model=chat_request_body.get("model", ""), + optional_params=map_openai_to_vertex_params(chat_request_body), custom_llm_provider="vertex_ai", litellm_params={}, cached_content=None, diff --git a/tests/unit/llms/bedrock/files/test_bedrock_files_transformation.py b/tests/unit/llms/bedrock/files/test_bedrock_files_transformation.py index d0921e68424..7a7159a1624 100644 --- a/tests/unit/llms/bedrock/files/test_bedrock_files_transformation.py +++ b/tests/unit/llms/bedrock/files/test_bedrock_files_transformation.py @@ -1672,6 +1672,49 @@ class TestBedrockBatchNonChatEndpointRecords: assert "input" not in model_input assert "max_output_tokens" not in model_input + def test_anthropic_responses_record_accepts_a_function_tool_without_strict(self): + """Clients omit the SDK's required `strict`; the record is forwarded like real time, not validated.""" + parameters = {"type": "object", "properties": {"city": {"type": "string"}}} + model_input = self._transform( + { + "custom_id": "4a", + "method": "POST", + "url": "/v1/responses", + "body": { + "model": self.ANTHROPIC_MODEL, + "input": "Weather in Paris?", + "tools": [{"type": "function", "name": "get_weather", "parameters": parameters}], + }, + } + ) + + assert model_input["messages"][0]["content"] == [{"type": "text", "text": "Weather in Paris?"}] + tool = model_input["tools"][0] + function = tool.get("function", tool) + assert (function["name"], function.get("parameters", function.get("input_schema"))) == ("get_weather", parameters) + + @pytest.mark.parametrize( + ("url", "body"), + [ + ( + "/v1/responses", + {"input": [{"role": "developer", "content": "be terse"}, {"role": "user", "content": "ping"}]}, + ), + ( + "/v1/chat/completions", + {"messages": [{"role": "developer", "content": "be terse"}, {"role": "user", "content": "ping"}]}, + ), + ], + ids=["responses", "chat"], + ) + def test_anthropic_developer_role_becomes_the_system_prompt_like_real_time(self, url, body): + model_input = self._transform( + {"custom_id": "4c", "method": "POST", "url": url, "body": {"model": self.ANTHROPIC_MODEL, **body}} + ) + + assert model_input["system"] == [{"type": "text", "text": "be terse"}] + assert [message["role"] for message in model_input["messages"]] == ["user"] + def test_responses_record_keeps_metadata(self): """`metadata` reaches the bridge, which reads it as its own kwarg.""" model_input = self._transform( diff --git a/tests/unit/llms/vertex_ai/files/test_vertex_ai_files_transformation.py b/tests/unit/llms/vertex_ai/files/test_vertex_ai_files_transformation.py index 48464e79876..7434eae72a4 100644 --- a/tests/unit/llms/vertex_ai/files/test_vertex_ai_files_transformation.py +++ b/tests/unit/llms/vertex_ai/files/test_vertex_ai_files_transformation.py @@ -5,6 +5,7 @@ Includes tests for Vertex AI batch output transformation to OpenAI format. import json import urllib.parse +from collections.abc import Mapping from types import MappingProxyType from urllib.parse import parse_qs, urlparse @@ -1447,6 +1448,158 @@ class TestVertexEmbeddingsBatchInputTranslation: assert "content" in embeddings_row["request"] +def _responses_entry( + body: Mapping[str, object] | None = None, + custom_id: str = "resp-1", + url: str = "/v1/responses", +) -> dict[str, object]: + return { + "custom_id": custom_id, + "method": "POST", + "url": url, + "body": body + if body is not None + else {"model": "gemini-2.5-flash", "input": "What was the top headline in world news yesterday?"}, + } + + +class TestVertexResponsesBatchInputTranslation: + """ + /v1/responses batch lines carry `input`, not `messages`, so they go through the + Responses-to-Chat bridge before the Gemini translation instead of uploading as an + empty text part. + """ + + def test_string_input_becomes_the_user_prompt(self): + (row,) = _wrap_entries([_responses_entry()]) + + assert row["request"]["contents"] == [ + {"role": "user", "parts": [{"text": "What was the top headline in world news yesterday?"}]} + ] + assert row["request"]["labels"]["litellm_custom_id"] == "resp-1" + + def test_instructions_and_input_items_map_like_real_time(self): + (row,) = _wrap_entries( + [ + _responses_entry( + body={ + "model": "gemini-2.5-flash", + "instructions": "be terse", + "input": [ + {"role": "user", "content": "what is 2+2?"}, + {"role": "assistant", "content": "4"}, + {"role": "user", "content": "and 3+3?"}, + ], + "max_output_tokens": 32, + "temperature": 0.2, + } + ) + ] + ) + + request = row["request"] + assert request["system_instruction"] == {"parts": [{"text": "be terse"}]} + assert [content["role"] for content in request["contents"]] == ["user", "model", "user"] + assert request["contents"][-1]["parts"] == [{"text": "and 3+3?"}] + assert request["generationConfig"]["max_output_tokens"] == 32 + assert request["generationConfig"]["temperature"] == 0.2 + + def test_web_search_tool_keeps_the_prompt(self): + (row,) = _wrap_entries( + [ + _responses_entry( + body={ + "model": "gemini-2.5-flash", + "input": "What was the top headline in world news yesterday?", + "tools": [{"type": "web_search"}], + } + ) + ] + ) + + assert row["request"]["contents"] == [ + {"role": "user", "parts": [{"text": "What was the top headline in world news yesterday?"}]} + ] + assert row["request"]["tools"] + + def test_sdk_optional_keys_are_not_required_like_real_time(self): + (row,) = _wrap_entries( + [ + _responses_entry( + body={ + "model": "gemini-2.5-flash", + "input": [ + { + "role": "user", + "content": [ + {"type": "input_text", "text": "Weather in the pictured city?"}, + {"type": "input_image", "image_url": "https://example.com/paris.png"}, + ], + } + ], + "tools": [ + { + "type": "function", + "name": "get_weather", + "parameters": {"type": "object", "properties": {"city": {"type": "string"}}}, + } + ], + } + ) + ] + ) + + request = row["request"] + assert request["contents"][0]["parts"] == [ + {"text": "Weather in the pictured city?"}, + {"file_data": {"mime_type": "image/png", "file_uri": "https://example.com/paris.png"}}, + ] + assert request["tools"][0]["function_declarations"][0]["name"] == "get_weather" + + @pytest.mark.parametrize( + "url", + ["/v1/responses", "/v1/responses/", "/v1/responses?beta=1", "responses", "https://api.openai.com/v1/responses"], + ) + def test_route_spellings_are_all_responses(self, url): + (row,) = _wrap_entries([_responses_entry(url=url)]) + + assert row["request"]["contents"][0]["parts"] == [ + {"text": "What was the top headline in world news yesterday?"} + ] + + def test_missing_input_fails_the_upload(self): + with pytest.raises(ValueError, match="missing required `input` field"): + _wrap_entries([_responses_entry(body={"model": "gemini-2.5-flash"})]) + + @pytest.mark.parametrize( + "entry", + [ + _responses_entry( + body={ + "model": "gemini-2.5-flash", + "input": [{"role": "developer", "content": "be terse"}, {"role": "user", "content": "ping"}], + } + ), + { + "custom_id": "chat-1", + "method": "POST", + "url": "/v1/chat/completions", + "body": { + "model": "gemini-2.5-flash", + "messages": [{"role": "developer", "content": "be terse"}, {"role": "user", "content": "ping"}], + }, + }, + ], + ids=["responses", "chat"], + ) + def test_developer_role_becomes_the_system_instruction_like_real_time(self, entry): + (row,) = _wrap_entries([entry]) + + request = row["request"] + assert request["system_instruction"] == {"parts": [{"text": "be terse"}]} + assert request["contents"] == [{"role": "user", "parts": [{"text": "ping"}]}] + + class TestVertexEmbeddingsBatchOutputTranslation: """Vertex Gemini Embedding batch output rows must come back as OpenAI batch rows.""" From 1dc3b62dbc160c5f97bd8420829dcdd629e1e0c7 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 24 Sep 2026 16:26:44 -0700 Subject: [PATCH 098/120] fix(cost-map): add vertex priority prices for gemini-3-pro-image-preview and batch price for gemini-embedding-001 (#43069) Co-authored-by: kerry Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/model_prices_and_context_window_backup.json | 13 +++++++++++++ model_prices_and_context_window.json | 13 +++++++++++++ 2 files changed, 26 insertions(+) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 9bec83f6b08..dc979cbd651 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -26312,10 +26312,14 @@ "gemini-3-pro-image-preview": { "input_cost_per_image": 0.0011, "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_priority": 3.6e-07, "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "cache_read_input_token_cost_above_200k_tokens_priority": 7.2e-07, "cache_read_input_token_cost_batches": 1e-07, "input_cost_per_token": 2e-06, + "input_cost_per_token_priority": 3.6e-06, "input_cost_per_token_above_200k_tokens": 4e-06, + "input_cost_per_token_above_200k_tokens_priority": 7.2e-06, "input_cost_per_token_batches": 1e-06, "litellm_provider": "vertex_ai-language-models", "max_input_tokens": 65536, @@ -26325,7 +26329,9 @@ "output_cost_per_image": 0.134, "output_cost_per_image_token": 0.00012, "output_cost_per_token": 1.2e-05, + "output_cost_per_token_priority": 2.16e-05, "output_cost_per_token_above_200k_tokens": 1.8e-05, + "output_cost_per_token_above_200k_tokens_priority": 3.24e-05, "output_cost_per_token_batches": 6e-06, "source": "https://ai.google.dev/gemini-api/docs/pricing", "supported_endpoints": [ @@ -27848,6 +27854,7 @@ "gemini-embedding-001": { "deprecation_date": "2028-05-20", "input_cost_per_token": 1.5e-07, + "input_cost_per_token_batches": 1.2e-07, "litellm_provider": "vertex_ai-embedding-models", "max_input_tokens": 2048, "max_tokens": 2048, @@ -49492,10 +49499,14 @@ "vertex_ai/gemini-3-pro-image-preview": { "input_cost_per_image": 0.0011, "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_priority": 3.6e-07, "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "cache_read_input_token_cost_above_200k_tokens_priority": 7.2e-07, "cache_read_input_token_cost_batches": 1e-07, "input_cost_per_token": 2e-06, + "input_cost_per_token_priority": 3.6e-06, "input_cost_per_token_above_200k_tokens": 4e-06, + "input_cost_per_token_above_200k_tokens_priority": 7.2e-06, "input_cost_per_token_batches": 1e-06, "litellm_provider": "vertex_ai-language-models", "max_input_tokens": 65536, @@ -49505,7 +49516,9 @@ "output_cost_per_image": 0.134, "output_cost_per_image_token": 0.00012, "output_cost_per_token": 1.2e-05, + "output_cost_per_token_priority": 2.16e-05, "output_cost_per_token_above_200k_tokens": 1.8e-05, + "output_cost_per_token_above_200k_tokens_priority": 3.24e-05, "output_cost_per_token_batches": 6e-06, "supports_reasoning": false, "source": "https://docs.cloud.google.com/vertex-ai/generative-ai/docs/models/gemini/3-pro-image" diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 9bec83f6b08..dc979cbd651 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -26312,10 +26312,14 @@ "gemini-3-pro-image-preview": { "input_cost_per_image": 0.0011, "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_priority": 3.6e-07, "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "cache_read_input_token_cost_above_200k_tokens_priority": 7.2e-07, "cache_read_input_token_cost_batches": 1e-07, "input_cost_per_token": 2e-06, + "input_cost_per_token_priority": 3.6e-06, "input_cost_per_token_above_200k_tokens": 4e-06, + "input_cost_per_token_above_200k_tokens_priority": 7.2e-06, "input_cost_per_token_batches": 1e-06, "litellm_provider": "vertex_ai-language-models", "max_input_tokens": 65536, @@ -26325,7 +26329,9 @@ "output_cost_per_image": 0.134, "output_cost_per_image_token": 0.00012, "output_cost_per_token": 1.2e-05, + "output_cost_per_token_priority": 2.16e-05, "output_cost_per_token_above_200k_tokens": 1.8e-05, + "output_cost_per_token_above_200k_tokens_priority": 3.24e-05, "output_cost_per_token_batches": 6e-06, "source": "https://ai.google.dev/gemini-api/docs/pricing", "supported_endpoints": [ @@ -27848,6 +27854,7 @@ "gemini-embedding-001": { "deprecation_date": "2028-05-20", "input_cost_per_token": 1.5e-07, + "input_cost_per_token_batches": 1.2e-07, "litellm_provider": "vertex_ai-embedding-models", "max_input_tokens": 2048, "max_tokens": 2048, @@ -49492,10 +49499,14 @@ "vertex_ai/gemini-3-pro-image-preview": { "input_cost_per_image": 0.0011, "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_priority": 3.6e-07, "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "cache_read_input_token_cost_above_200k_tokens_priority": 7.2e-07, "cache_read_input_token_cost_batches": 1e-07, "input_cost_per_token": 2e-06, + "input_cost_per_token_priority": 3.6e-06, "input_cost_per_token_above_200k_tokens": 4e-06, + "input_cost_per_token_above_200k_tokens_priority": 7.2e-06, "input_cost_per_token_batches": 1e-06, "litellm_provider": "vertex_ai-language-models", "max_input_tokens": 65536, @@ -49505,7 +49516,9 @@ "output_cost_per_image": 0.134, "output_cost_per_image_token": 0.00012, "output_cost_per_token": 1.2e-05, + "output_cost_per_token_priority": 2.16e-05, "output_cost_per_token_above_200k_tokens": 1.8e-05, + "output_cost_per_token_above_200k_tokens_priority": 3.24e-05, "output_cost_per_token_batches": 6e-06, "supports_reasoning": false, "source": "https://docs.cloud.google.com/vertex-ai/generative-ai/docs/models/gemini/3-pro-image" From 4584958574cd22b6f7f9b65aa6c92c58cdf940b6 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 24 Sep 2026 18:26:50 -0500 Subject: [PATCH 099/120] feat(agents): add optional per-agent kill switch webhook (#42841) --- .../migration.sql | 2 + .../litellm_proxy_extras/schema.prisma | 1 + litellm/constants.py | 2 + litellm/proxy/_lazy_openapi_snapshot.json | 293 ++++++++++++++++++ litellm/proxy/_types.py | 4 +- .../proxy/agent_endpoints/agent_registry.py | 50 ++- litellm/proxy/agent_endpoints/endpoints.py | 90 +++++- litellm/proxy/agent_endpoints/kill_switch.py | 239 ++++++++++++++ litellm/proxy/schema.prisma | 1 + litellm/types/agents.py | 76 ++++- litellm/types/llms/custom_http.py | 1 + schema.prisma | 1 + .../agent_endpoints/test_agent_registry.py | 162 +++++++++- .../proxy/agent_endpoints/test_endpoints.py | 223 ++++++++++++- .../proxy/agent_endpoints/test_kill_switch.py | 248 +++++++++++++++ .../proxy/auth/test_route_checks.py | 1 + .../agents/_components/AgentFormKit.tsx | 5 + .../AgentKillSwitchDangerZone.test.tsx | 124 ++++++++ .../_components/AgentKillSwitchDangerZone.tsx | 147 +++++++++ .../_components/KillSwitchFormFields.tsx | 223 +++++++++++++ .../agents/_components/agent_config.ts | 23 ++ .../agents/_components/agent_form_fields.tsx | 9 + .../agents/_components/agent_info.test.tsx | 28 ++ .../agents/_components/agent_info.tsx | 11 +- .../_components/agent_type_utils.test.ts | 26 ++ .../agents/_components/agent_type_utils.ts | 2 + .../dynamic_agent_form_fields.test.ts | 65 ++++ .../_components/dynamic_agent_form_fields.tsx | 21 +- .../_components/kill_switch_config.test.ts | 118 +++++++ .../agents/_components/kill_switch_config.ts | 124 ++++++++ .../src/components/agents/types.ts | 3 + .../src/components/networking.tsx | 11 + .../AuditLogDrawer/AuditLogDrawer.tsx | 1 + .../components/view_logs/AuditLogsTable.tsx | 2 + .../view_logs/AuditLogsTableColumns.tsx | 9 +- .../src/contexts/PluginModeContext.tsx | 10 +- ui/litellm-dashboard/src/lib/http/schema.d.ts | 149 +++++++++ 37 files changed, 2485 insertions(+), 20 deletions(-) create mode 100644 litellm-proxy-extras/litellm_proxy_extras/migrations/20260923000000_add_agent_kill_switch/migration.sql create mode 100644 litellm/proxy/agent_endpoints/kill_switch.py create mode 100644 tests/test_litellm/proxy/agent_endpoints/test_kill_switch.py create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/agents/_components/AgentKillSwitchDangerZone.test.tsx create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/agents/_components/AgentKillSwitchDangerZone.tsx create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/agents/_components/KillSwitchFormFields.tsx create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/agents/_components/dynamic_agent_form_fields.test.ts create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/agents/_components/kill_switch_config.test.ts create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/agents/_components/kill_switch_config.ts diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260923000000_add_agent_kill_switch/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260923000000_add_agent_kill_switch/migration.sql new file mode 100644 index 00000000000..dd21ed644eb --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260923000000_add_agent_kill_switch/migration.sql @@ -0,0 +1,2 @@ +-- AlterTable +ALTER TABLE "LiteLLM_AgentsTable" ADD COLUMN IF NOT EXISTS "kill_switch" JSONB; diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index 85996430bc5..69c63d9ecd6 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -72,6 +72,7 @@ model LiteLLM_AgentsTable { agent_card_params Json static_headers Json? @default("{}") extra_headers String[] @default([]) + kill_switch Json? agent_access_groups String[] @default([]) access_group_ids String[] @default([]) object_permission_id String? diff --git a/litellm/constants.py b/litellm/constants.py index 7b40f432446..807694c2f8a 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -557,6 +557,8 @@ SEMANTIC_CACHE_EMBEDDING_TIMEOUT_SECONDS: Final[float] = float( request_timeout: float = float(os.getenv("REQUEST_TIMEOUT", str(int(DEFAULT_REQUEST_TIMEOUT_SECONDS)))) request_timeout_explicitly_set: bool = "REQUEST_TIMEOUT" in os.environ DEFAULT_A2A_AGENT_TIMEOUT: Final[float] = float(os.getenv("DEFAULT_A2A_AGENT_TIMEOUT", 6000)) # 10 minutes +AGENT_KILL_SWITCH_TIMEOUT_SECONDS: Final = 10.0 +AGENT_KILL_SWITCH_RESPONSE_BODY_MAX_CHARS: Final = 2000 # Patterns that indicate a localhost/internal URL in A2A agent cards that should be # replaced with the original base_url. This is a common misconfiguration where # developers deploy agents with development URLs in their agent cards. diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index 0b43c3864ab..3d44315341b 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -2392,6 +2392,16 @@ ], "title": "Extra Headers" }, + "kill_switch": { + "anyOf": [ + { + "$ref": "#/components/schemas/AgentKillSwitchConfig" + }, + { + "type": "null" + } + ] + }, "litellm_params": { "additionalProperties": true, "title": "Litellm Params", @@ -2561,6 +2571,221 @@ "title": "AgentKeySummary", "type": "object" }, + "AgentKillSwitchApiKeyAuth": { + "additionalProperties": false, + "properties": { + "api_key": { + "title": "Api Key", + "type": "string" + }, + "header_name": { + "default": "x-api-key", + "title": "Header Name", + "type": "string" + }, + "type": { + "const": "api_key", + "title": "Type", + "type": "string" + } + }, + "required": [ + "type", + "api_key" + ], + "title": "AgentKillSwitchApiKeyAuth", + "type": "object" + }, + "AgentKillSwitchBasicAuth": { + "additionalProperties": false, + "properties": { + "password": { + "title": "Password", + "type": "string" + }, + "type": { + "const": "basic", + "title": "Type", + "type": "string" + }, + "username": { + "title": "Username", + "type": "string" + } + }, + "required": [ + "type", + "username", + "password" + ], + "title": "AgentKillSwitchBasicAuth", + "type": "object" + }, + "AgentKillSwitchBearerAuth": { + "additionalProperties": false, + "properties": { + "token": { + "title": "Token", + "type": "string" + }, + "type": { + "const": "bearer", + "title": "Type", + "type": "string" + } + }, + "required": [ + "type", + "token" + ], + "title": "AgentKillSwitchBearerAuth", + "type": "object" + }, + "AgentKillSwitchConfig": { + "additionalProperties": false, + "description": "Webhook an admin fires to shut an agent down out of band. LiteLLM only\nmakes the call; whatever the endpoint does with it is the agent's business.", + "properties": { + "auth": { + "anyOf": [ + { + "discriminator": { + "mapping": { + "api_key": "#/components/schemas/AgentKillSwitchApiKeyAuth", + "basic": "#/components/schemas/AgentKillSwitchBasicAuth", + "bearer": "#/components/schemas/AgentKillSwitchBearerAuth" + }, + "propertyName": "type" + }, + "oneOf": [ + { + "$ref": "#/components/schemas/AgentKillSwitchBearerAuth" + }, + { + "$ref": "#/components/schemas/AgentKillSwitchApiKeyAuth" + }, + { + "$ref": "#/components/schemas/AgentKillSwitchBasicAuth" + } + ] + }, + { + "type": "null" + } + ], + "title": "Auth" + }, + "body": { + "anyOf": [ + { + "additionalProperties": true, + "type": "object" + }, + { + "type": "null" + } + ], + "title": "Body" + }, + "headers": { + "additionalProperties": { + "type": "string" + }, + "title": "Headers", + "type": "object" + }, + "method": { + "default": "POST", + "enum": [ + "POST", + "PUT", + "PATCH", + "DELETE", + "GET" + ], + "title": "Method", + "type": "string" + }, + "query_params": { + "additionalProperties": { + "type": "string" + }, + "title": "Query Params", + "type": "object" + }, + "url": { + "title": "Url", + "type": "string" + } + }, + "required": [ + "url" + ], + "title": "AgentKillSwitchConfig", + "type": "object" + }, + "AgentKillSwitchResult": { + "properties": { + "agent_id": { + "title": "Agent Id", + "type": "string" + }, + "error": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Error" + }, + "method": { + "enum": [ + "POST", + "PUT", + "PATCH", + "DELETE", + "GET" + ], + "title": "Method", + "type": "string" + }, + "response_body": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Response Body" + }, + "status_code": { + "anyOf": [ + { + "type": "integer" + }, + { + "type": "null" + } + ], + "title": "Status Code" + }, + "url": { + "title": "Url", + "type": "string" + } + }, + "required": [ + "agent_id", + "url", + "method" + ], + "title": "AgentKillSwitchResult", + "type": "object" + }, "AgentMakePublicResponse": { "properties": { "message": { @@ -2775,6 +3000,16 @@ ], "title": "Keys" }, + "kill_switch": { + "anyOf": [ + { + "$ref": "#/components/schemas/AgentKillSwitchConfig" + }, + { + "type": "null" + } + ] + }, "litellm_params": { "anyOf": [ { @@ -3569,6 +3804,16 @@ ], "title": "Extra Headers" }, + "kill_switch": { + "anyOf": [ + { + "$ref": "#/components/schemas/AgentKillSwitchConfig" + }, + { + "type": "null" + } + ] + }, "litellm_params": { "additionalProperties": true, "title": "Litellm Params", @@ -4331,6 +4576,54 @@ ] } }, + "/v1/agents/{agent_id}/kill_switch": { + "post": { + "description": "Fire the agent's configured kill switch webhook. Proxy admin only.\n\nLiteLLM only makes the configured HTTP call and reports what came back; it\ndoes not change the agent's state in LiteLLM. Returns 200 when the webhook\nanswered 2xx, 502 with the same result body otherwise. Every attempt is\nwritten to the audit log as a `kill_switch_fired` row against the agent.\n\nExample Request:\n```bash\ncurl -X POST \"http://localhost:4000/v1/agents/123e4567-e89b-12d3-a456-426614174000/kill_switch\" \\\n -H \"Authorization: Bearer \"\n```", + "operationId": "trigger_agent_kill_switch_v1_agents__agent_id__kill_switch_post", + "parameters": [ + { + "in": "path", + "name": "agent_id", + "required": true, + "schema": { + "title": "Agent Id", + "type": "string" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/AgentKillSwitchResult" + } + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Trigger Agent Kill Switch", + "tags": [ + "agents" + ] + } + }, "/v1/agents/{agent_id}/make_public": { "post": { "description": "Make an agent publicly discoverable\n\nExample Request:\n```bash\ncurl -X POST \"http://localhost:4000/v1/agents/123e4567-e89b-12d3-a456-426614174000/make_public\" \\\n -H \"Authorization: Bearer \" \\\n -H \"Content-Type: application/json\"\n```\n\nExample Response:\n```json\n{\n \"agent_id\": \"123e4567-e89b-12d3-a456-426614174000\",\n \"agent_name\": \"my-custom-agent\",\n \"litellm_params\": {\n \"make_public\": true\n },\n \"agent_card_params\": {...},\n \"created_at\": \"2025-11-15T10:30:00Z\",\n \"updated_at\": \"2025-11-15T10:35:00Z\",\n \"created_by\": \"user123\",\n \"updated_by\": \"user123\"\n}\n```", diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index b6de36f8423..12b4d4b2412 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -238,6 +238,7 @@ class LitellmTableNames(str, enum.Enum): CONFIG_TABLE_NAME = "LiteLLM_Config" SSO_CONFIG_TABLE_NAME = "LiteLLM_SSOConfig" UI_SETTINGS_TABLE_NAME = "LiteLLM_UISettings" + AGENT_TABLE_NAME = "LiteLLM_AgentsTable" class Litellm_EntityType(enum.Enum): @@ -578,6 +579,7 @@ class LiteLLMRoutes(enum.Enum): "/v1/agents/{agent_id}", "/v1/agents/make_public", "/v1/agents/{agent_id}/make_public", + "/v1/agents/{agent_id}/kill_switch", ) # Backwards-compat union — virtual keys may be configured with @@ -3688,7 +3690,7 @@ from litellm.models.spend_logs import ( # noqa: E402 ) from litellm.models.tag import LiteLLM_TagTable as LiteLLM_TagTable # noqa: E402 -AUDIT_ACTIONS = Literal["created", "updated", "deleted", "blocked", "unblocked", "rotated"] +AUDIT_ACTIONS = Literal["created", "updated", "deleted", "blocked", "unblocked", "rotated", "kill_switch_fired"] class LiteLLM_AuditLogs(LiteLLMPydanticObjectBase): diff --git a/litellm/proxy/agent_endpoints/agent_registry.py b/litellm/proxy/agent_endpoints/agent_registry.py index 3d56c2b5326..3e775d7648e 100644 --- a/litellm/proxy/agent_endpoints/agent_registry.py +++ b/litellm/proxy/agent_endpoints/agent_registry.py @@ -13,13 +13,14 @@ import litellm from litellm.constants import REDACTED_BY_LITELM_STRING from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.litellm_core_utils.sensitive_data_masker import SensitiveDataMasker +from litellm.proxy.agent_endpoints.kill_switch import restore_kill_switch from litellm.proxy.management_helpers.object_permission_utils import ( handle_update_object_permission_common, ) from litellm.proxy.utils import PrismaClient from litellm.repositories.prisma_protocols import TableActions from litellm.repositories.table_repositories import AgentsRepository, ObjectPermissionRepository -from litellm.types.agents import AgentConfig, AgentResponse, PatchAgentRequest +from litellm.types.agents import AgentConfig, AgentKillSwitchConfig, AgentResponse, PatchAgentRequest if TYPE_CHECKING: from prisma import models as prisma_models @@ -31,6 +32,10 @@ class AgentObjectPermissionRecord(Protocol): def dict(self) -> dict[str, object]: ... +class AgentIdWhere(TypedDict): + agent_id: ReadOnly[str] + + class AgentRecordDump(TypedDict): agent_id: str agent_name: str @@ -38,6 +43,7 @@ class AgentRecordDump(TypedDict): agent_card_params: dict[str, object] static_headers: dict[str, str] | None extra_headers: list[str] | None + kill_switch: ReadOnly[AgentKillSwitchConfig | None] access_group_ids: ReadOnly[Sequence[str] | None] object_permission: dict[str, object] | None spend: float @@ -70,6 +76,9 @@ class AgentRecord(Protocol): @property def access_group_ids(self) -> Sequence[str] | None: ... + @property + def kill_switch(self) -> Mapping[str, object] | None: ... + @property def spend(self) -> float: ... @@ -211,6 +220,29 @@ def parse_agent_litellm_params(value: object) -> Mapping[str, object]: return _EMPTY_LITELLM_PARAMS +_KILL_SWITCH_ADAPTER: Final[TypeAdapter[AgentKillSwitchConfig | None]] = TypeAdapter(AgentKillSwitchConfig | None) + + +def parse_agent_kill_switch(value: object) -> AgentKillSwitchConfig | None: + if value is None: + return None + try: + if isinstance(value, str): + return _KILL_SWITCH_ADAPTER.validate_json(value) + return _KILL_SWITCH_ADAPTER.validate_python(value) + except ValidationError: + return None + + +def serialize_agent_kill_switch(incoming: object, existing: object) -> str: + """prisma-client-py drops ``None`` from update data, so a cleared kill switch is stored as the JSON literal + ``null`` (read back as ``None``), the same convention ``memory_endpoints`` uses for ``Json?`` columns.""" + restored: Final = restore_kill_switch( + _KILL_SWITCH_ADAPTER.validate_python(incoming), parse_agent_kill_switch(existing) + ) + return safe_dumps(restored.model_dump() if restored is not None else None) + + _MISSING_AGENT_PARAM: Final = object() _RESTORE_AGENT_PARAMS_MAX_DEPTH: Final = 10 @@ -293,6 +325,12 @@ def _patched_access_group_ids(agent: PatchAgentRequest) -> Mapping[str, object]: return MappingProxyType({"access_group_ids": tuple(dict.fromkeys(agent.get("access_group_ids") or ()))}) +def _patched_kill_switch(agent: PatchAgentRequest, existing: object) -> Mapping[str, object]: + if "kill_switch" not in agent: + return MappingProxyType({}) + return MappingProxyType({"kill_switch": serialize_agent_kill_switch(agent.get("kill_switch"), existing)}) + + def _restore_redacted_litellm_params( incoming: Mapping[str, object], existing: Mapping[str, object], @@ -531,6 +569,7 @@ class AgentRegistry: "agent_name": agent_name, "litellm_params": litellm_params, "agent_card_params": agent_card_params, + "kill_switch": serialize_agent_kill_switch(agent.get("kill_switch"), None), "created_by": created_by, "updated_by": created_by, "created_at": datetime.now(timezone.utc), @@ -613,7 +652,10 @@ class AgentRegistry: existing_agent: Final[Mapping[str, object]] = dict(existing_record) augment_agent: Final = {**existing_agent, **agent} - update_data: Final[dict[str, object]] = {**_patched_access_group_ids(agent)} + update_data: Final[dict[str, object]] = { + **_patched_access_group_ids(agent), + **_patched_kill_switch(agent, existing_agent.get("kill_switch")), + } if augment_agent.get("agent_name"): update_data["agent_name"] = augment_agent.get("agent_name") if "litellm_params" in agent: @@ -716,6 +758,9 @@ class AgentRegistry: ) extra_headers_val_u: Final = agent.get("extra_headers") or [] access_group_ids_val_u: Final = tuple(dict.fromkeys(agent.get("access_group_ids") or ())) + kill_switch_val_u: Final = serialize_agent_kill_switch( + agent.get("kill_switch"), existing_row.kill_switch if existing_row is not None else None + ) update_data: Final[dict[str, object]] = { "agent_name": agent_name, @@ -723,6 +768,7 @@ class AgentRegistry: "agent_card_params": agent_card_params, "static_headers": static_headers_val_u, "extra_headers": extra_headers_val_u, + "kill_switch": kill_switch_val_u, "access_group_ids": access_group_ids_val_u, "updated_by": updated_by, "updated_at": datetime.now(timezone.utc), diff --git a/litellm/proxy/agent_endpoints/endpoints.py b/litellm/proxy/agent_endpoints/endpoints.py index aa8979a73c6..28c82a715e0 100644 --- a/litellm/proxy/agent_endpoints/endpoints.py +++ b/litellm/proxy/agent_endpoints/endpoints.py @@ -33,6 +33,8 @@ from litellm.proxy.a2a.agent_card import ( normalize_protocol_version, ) from litellm.proxy.agent_endpoints.agent_registry import ( + AgentIdWhere, + parse_agent_kill_switch, parse_agent_litellm_params, redact_sensitive_agent_litellm_params, ) @@ -45,6 +47,15 @@ from litellm.proxy.agent_endpoints.agent_search import ( search_agents, ) from litellm.proxy.agent_endpoints.auth.agent_permission_handler import accessible_agents +from litellm.proxy.agent_endpoints.kill_switch import ( + KillSwitchAuditLogWriter, + KillSwitchHttpClient, + build_kill_switch_audit_log, + default_kill_switch_audit_log_writer, + default_kill_switch_http_client, + fire_kill_switch, + redact_kill_switch, +) from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.common_utils.rbac_utils import check_feature_access_for_user from litellm.proxy.management_endpoints.common_daily_activity import get_daily_activity @@ -53,6 +64,8 @@ from litellm.types.agents import ( AgentCard, AgentConfig, AgentKeySummary, + AgentKillSwitchConfig, + AgentKillSwitchResult, AgentMakePublicResponse, AgentResponse, MakeAgentsPublicRequest, @@ -160,9 +173,10 @@ def _redact_sensitive_agent_fields( ) -> list[AgentResponse]: """ Return copies of the given agents with credential-bearing litellm_params - values replaced by a fixed marker (never returned to ANY caller, - admin included) and, for non-admin callers, virtual-key and header - fields stripped entirely. The original objects are not modified. + values and kill-switch auth secrets replaced by a fixed marker (never + returned to ANY caller, admin included) and, for non-admin callers, + virtual-key, header and kill-switch fields stripped entirely. The original + objects are not modified. """ redacted: Final[list[AgentResponse]] = [] for agent in agents: @@ -171,8 +185,10 @@ def _redact_sensitive_agent_fields( copy.static_headers = None copy.extra_headers = None copy.keys = None + copy.kill_switch = None if copy.litellm_params: copy.litellm_params = _redact_agent_litellm_params_dict(copy.litellm_params) + copy.kill_switch = redact_kill_switch(copy.kill_switch) redacted.append(copy) return redacted @@ -872,6 +888,74 @@ async def delete_agent( raise HTTPException(status_code=500, detail=str(e)) +@router.post( + "/v1/agents/{agent_id}/kill_switch", + tags=["[beta] A2A Agents"], # mutable-ok: fastapi types tags as list[str | Enum] + dependencies=(Depends(user_api_key_auth),), + response_model=AgentKillSwitchResult, +) +async def trigger_agent_kill_switch( + agent_id: str, + user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)], + http_client: Annotated[KillSwitchHttpClient, Depends(default_kill_switch_http_client)], + audit_log_writer: Annotated[KillSwitchAuditLogWriter, Depends(default_kill_switch_audit_log_writer)], +): + """ + Fire the agent's configured kill switch webhook. Proxy admin only. + + LiteLLM only makes the configured HTTP call and reports what came back; it + does not change the agent's state in LiteLLM. Returns 200 when the webhook + answered 2xx, 502 with the same result body otherwise. Every attempt is + written to the audit log as a `kill_switch_fired` row against the agent. + + Example Request: + ```bash + curl -X POST "http://localhost:4000/v1/agents/123e4567-e89b-12d3-a456-426614174000/kill_switch" \\ + -H "Authorization: Bearer " + ``` + """ + from litellm.proxy.proxy_server import litellm_proxy_admin_name + + await check_feature_access_for_user(user_api_key_dict, "agents") + _check_agent_management_permission(user_api_key_dict) + + resolved: Final = await _resolve_agent_kill_switch(agent_id) + if resolved is None: + raise HTTPException(status_code=404, detail=f"Agent with ID {agent_id} not found") + resolved_agent_id, config = resolved + if config is None: + raise HTTPException(status_code=400, detail=f"Agent with ID {agent_id} has no kill_switch configured") + + result: Final = await fire_kill_switch(agent_id=resolved_agent_id, config=config, http_client=http_client) + await audit_log_writer( + build_kill_switch_audit_log( + result=result, + user_api_key_dict=user_api_key_dict, + litellm_proxy_admin_name=litellm_proxy_admin_name, + ) + ) + if not result.succeeded: + raise HTTPException(status_code=502, detail=result.model_dump()) + return result + + +async def _resolve_agent_kill_switch(agent_id: str) -> tuple[str, AgentKillSwitchConfig | None] | None: + """The DB row wins over this replica's in-memory registry so a trigger never fires a webhook another + replica has since changed; config.yaml agents have no row and fall back to the registry.""" + from litellm.proxy.proxy_server import prisma_client + + if prisma_client is not None: + where: Final[AgentIdWhere] = {"agent_id": agent_id} + row: Final = await agents_table(prisma_client).find_unique(where=where) + if row is not None: + return row.agent_id, parse_agent_kill_switch(row.kill_switch) + + agent: Final = AGENT_REGISTRY.get_agent_by_id(agent_id=agent_id) + if agent is None: + return None + return agent.agent_id, agent.kill_switch + + @router.post( "/v1/agents/{agent_id}/make_public", tags=["[beta] A2A Agents"], diff --git a/litellm/proxy/agent_endpoints/kill_switch.py b/litellm/proxy/agent_endpoints/kill_switch.py new file mode 100644 index 00000000000..8b3f64e74ee --- /dev/null +++ b/litellm/proxy/agent_endpoints/kill_switch.py @@ -0,0 +1,239 @@ +from base64 import b64encode +from collections.abc import AsyncIterator, Awaitable, Callable, Mapping +from dataclasses import dataclass +from datetime import datetime, timezone +from types import MappingProxyType +from typing import Final, Protocol, TypeAlias + +import httpx +from typing_extensions import assert_never + +from litellm._logging import verbose_proxy_logger +from litellm._uuid import uuid +from litellm.constants import ( + AGENT_KILL_SWITCH_RESPONSE_BODY_MAX_CHARS, + AGENT_KILL_SWITCH_TIMEOUT_SECONDS, + REDACTED_BY_LITELM_STRING, +) +from litellm.llms.custom_httpx.http_handler import ( + get_async_httpx_client, # pyright: ignore[reportUnknownVariableType] # its params arg is a bare dict in http_handler +) +from litellm.proxy._types import LiteLLM_AuditLogs, LitellmTableNames, UserAPIKeyAuth +from litellm.proxy.management_helpers.audit_logs import create_audit_log_for_update, get_audit_log_changed_by +from litellm.types.agents import ( + AgentKillSwitchApiKeyAuth, + AgentKillSwitchAuth, + AgentKillSwitchBasicAuth, + AgentKillSwitchBearerAuth, + AgentKillSwitchConfig, + AgentKillSwitchResult, +) +from litellm.types.llms.custom_http import httpxSpecialProvider + + +def _with_auth(config: AgentKillSwitchConfig, auth: AgentKillSwitchAuth) -> AgentKillSwitchConfig: + return AgentKillSwitchConfig( + url=config.url, + method=config.method, + headers=config.headers, + query_params=config.query_params, + body=config.body, + auth=auth, + ) + + +def redact_kill_switch(config: AgentKillSwitchConfig | None) -> AgentKillSwitchConfig | None: + if config is None or config.auth is None: + return config + return _with_auth(config, _redact_auth(config.auth)) + + +def _redact_auth(auth: AgentKillSwitchAuth) -> AgentKillSwitchAuth: + match auth: + case AgentKillSwitchBearerAuth(): + return AgentKillSwitchBearerAuth(type="bearer", token=REDACTED_BY_LITELM_STRING) + case AgentKillSwitchApiKeyAuth(): + return AgentKillSwitchApiKeyAuth( + type="api_key", header_name=auth.header_name, api_key=REDACTED_BY_LITELM_STRING + ) + case AgentKillSwitchBasicAuth(): + return AgentKillSwitchBasicAuth(type="basic", username=auth.username, password=REDACTED_BY_LITELM_STRING) + case _: + assert_never(auth) + + +def restore_kill_switch( + incoming: AgentKillSwitchConfig | None, + existing: AgentKillSwitchConfig | None, +) -> AgentKillSwitchConfig | None: + """Put the stored secret back behind an auth field echoed as the redaction + marker; a marker with no stored secret of the same auth type becomes "".""" + if incoming is None or incoming.auth is None: + return incoming + existing_auth: Final = existing.auth if existing is not None else None + return _with_auth(incoming, _restore_auth(incoming.auth, existing_auth)) + + +def _restore_secret(incoming_value: str, existing_value: str | None) -> str: + if incoming_value != REDACTED_BY_LITELM_STRING: + return incoming_value + return existing_value if existing_value is not None else "" + + +def _restore_auth(incoming: AgentKillSwitchAuth, existing: AgentKillSwitchAuth | None) -> AgentKillSwitchAuth: + match incoming: + case AgentKillSwitchBearerAuth(): + stored_token: Final = existing.token if isinstance(existing, AgentKillSwitchBearerAuth) else None + return AgentKillSwitchBearerAuth(type="bearer", token=_restore_secret(incoming.token, stored_token)) + case AgentKillSwitchApiKeyAuth(): + stored_key: Final = existing.api_key if isinstance(existing, AgentKillSwitchApiKeyAuth) else None + return AgentKillSwitchApiKeyAuth( + type="api_key", + header_name=incoming.header_name, + api_key=_restore_secret(incoming.api_key, stored_key), + ) + case AgentKillSwitchBasicAuth(): + stored_password: Final = existing.password if isinstance(existing, AgentKillSwitchBasicAuth) else None + return AgentKillSwitchBasicAuth( + type="basic", + username=incoming.username, + password=_restore_secret(incoming.password, stored_password), + ) + case _: + assert_never(incoming) + + +@dataclass(frozen=True, slots=True) +class KillSwitchRequest: + method: str + url: str + headers: Mapping[str, str] + json_body: Mapping[str, object] | None + + +def _auth_headers(auth: AgentKillSwitchAuth | None) -> Mapping[str, str]: + match auth: + case None: + return MappingProxyType({}) + case AgentKillSwitchBearerAuth(): + return MappingProxyType({"Authorization": f"Bearer {auth.token}"}) + case AgentKillSwitchApiKeyAuth(): + return MappingProxyType({auth.header_name: auth.api_key}) + case AgentKillSwitchBasicAuth(): + credentials: Final = b64encode(f"{auth.username}:{auth.password}".encode()).decode() + return MappingProxyType({"Authorization": f"Basic {credentials}"}) + case _: + assert_never(auth) + + +def build_kill_switch_request(config: AgentKillSwitchConfig) -> KillSwitchRequest: + url: Final = httpx.URL(config.url).copy_merge_params(config.query_params) + return KillSwitchRequest( + method=config.method, + url=str(url), + headers=MappingProxyType({**config.headers, **_auth_headers(config.auth)}), + json_body=config.body, + ) + + +class KillSwitchHttpClient(Protocol): + def build_request( + self, + method: str, + url: str, + *, + headers: Mapping[str, str], + json: Mapping[str, object] | None, + timeout: float, + ) -> httpx.Request: ... + + async def send(self, request: httpx.Request, *, stream: bool, follow_redirects: bool) -> httpx.Response: ... + + +def default_kill_switch_http_client() -> KillSwitchHttpClient: + return get_async_httpx_client(llm_provider=httpxSpecialProvider.AgentKillSwitch).client + + +KillSwitchAuditLogWriter: TypeAlias = Callable[[LiteLLM_AuditLogs], Awaitable[None]] # mutable-ok: Callable params + + +def default_kill_switch_audit_log_writer() -> KillSwitchAuditLogWriter: + return create_audit_log_for_update + + +def build_kill_switch_audit_log( + *, + result: AgentKillSwitchResult, + user_api_key_dict: UserAPIKeyAuth, + litellm_proxy_admin_name: str | None, +) -> LiteLLM_AuditLogs: + return LiteLLM_AuditLogs( + id=str(uuid.uuid4()), + updated_at=datetime.now(timezone.utc), + changed_by=get_audit_log_changed_by( + litellm_changed_by=None, + user_api_key_dict=user_api_key_dict, + litellm_proxy_admin_name=litellm_proxy_admin_name, + ), + changed_by_api_key=user_api_key_dict.api_key, + table_name=LitellmTableNames.AGENT_TABLE_NAME, + object_id=result.agent_id, + action="kill_switch_fired", + updated_values=result.model_dump_json(exclude_none=True), + ) + + +async def fire_kill_switch( + *, + agent_id: str, + config: AgentKillSwitchConfig, + http_client: KillSwitchHttpClient, + timeout: float = AGENT_KILL_SWITCH_TIMEOUT_SECONDS, +) -> AgentKillSwitchResult: + request: Final = build_kill_switch_request(config) + reported_url: Final = str(httpx.URL(request.url).copy_with(query=None)) + verbose_proxy_logger.info("Firing kill switch for agent %s: %s %s", agent_id, request.method, reported_url) + try: + response: Final = await http_client.send( + http_client.build_request( + request.method, + request.url, + headers=request.headers, + json=request.json_body, + timeout=timeout, + ), + stream=True, + follow_redirects=False, + ) + body: Final = await _read_text_prefix(response, AGENT_KILL_SWITCH_RESPONSE_BODY_MAX_CHARS) + except httpx.HTTPError as exc: + verbose_proxy_logger.warning("Kill switch for agent %s failed: %s", agent_id, type(exc).__name__) + return AgentKillSwitchResult( + agent_id=agent_id, + url=reported_url, + method=config.method, + error=type(exc).__name__, + ) + return AgentKillSwitchResult( + agent_id=agent_id, + url=reported_url, + method=config.method, + status_code=response.status_code, + response_body=body, + ) + + +async def _read_text_prefix(response: httpx.Response, max_chars: int) -> str: + try: + return await _take_text(response.aiter_text(), max_chars) + finally: + await response.aclose() + + +async def _take_text(chunks: AsyncIterator[str], max_chars: int) -> str: + taken = "" # rebind-ok: running prefix of a stream that is abandoned once the cap is hit + async for chunk in chunks: + taken += chunk # rebind-ok: see above + if len(taken) >= max_chars: + break + return taken[:max_chars] diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 85996430bc5..69c63d9ecd6 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -72,6 +72,7 @@ model LiteLLM_AgentsTable { agent_card_params Json static_headers Json? @default("{}") extra_headers String[] @default([]) + kill_switch Json? agent_access_groups String[] @default([]) access_group_ids String[] @default([]) object_permission_id String? diff --git a/litellm/types/agents.py b/litellm/types/agents.py index 7f8d8c6af66..f7aef09fa29 100644 --- a/litellm/types/agents.py +++ b/litellm/types/agents.py @@ -1,8 +1,9 @@ from collections.abc import Mapping, Sequence from datetime import datetime -from typing import TYPE_CHECKING, Any, Final, Literal +from typing import TYPE_CHECKING, Annotated, Any, Final, Literal, TypeAlias +from urllib.parse import urlsplit -from pydantic import BaseModel, ConfigDict, PrivateAttr, StrictInt +from pydantic import BaseModel, ConfigDict, Field, PrivateAttr, StrictInt, field_validator from typing_extensions import ReadOnly, Required, TypedDict from litellm.types.llms.base import LiteLLMPydanticObjectBase @@ -178,6 +179,74 @@ class AgentObjectPermission(TypedDict, total=False): agents: list[str] | None +class AgentKillSwitchBearerAuth(BaseModel): + model_config = ConfigDict(frozen=True, extra="forbid") + + type: Literal["bearer"] + token: str + + +class AgentKillSwitchApiKeyAuth(BaseModel): + model_config = ConfigDict(frozen=True, extra="forbid") + + type: Literal["api_key"] + header_name: str = "x-api-key" + api_key: str + + +class AgentKillSwitchBasicAuth(BaseModel): + model_config = ConfigDict(frozen=True, extra="forbid") + + type: Literal["basic"] + username: str + password: str + + +AgentKillSwitchAuth: TypeAlias = Annotated[ + AgentKillSwitchBearerAuth | AgentKillSwitchApiKeyAuth | AgentKillSwitchBasicAuth, + Field(discriminator="type"), +] + +AgentKillSwitchMethod: TypeAlias = Literal["POST", "PUT", "PATCH", "DELETE", "GET"] + + +class AgentKillSwitchConfig(BaseModel): + """Webhook an admin fires to shut an agent down out of band. LiteLLM only + makes the call; whatever the endpoint does with it is the agent's business.""" + + model_config = ConfigDict(frozen=True, extra="forbid") + + url: str + method: AgentKillSwitchMethod = "POST" + headers: Mapping[str, str] = Field(default_factory=dict) + query_params: Mapping[str, str] = Field(default_factory=dict) + body: Mapping[str, object] | None = None + auth: AgentKillSwitchAuth | None = None + + @field_validator("url") + @classmethod + def _require_absolute_http_url(cls, value: str) -> str: + parts: Final = urlsplit(value) + if parts.scheme not in ("http", "https") or not parts.netloc: + raise ValueError("kill_switch.url must be an absolute http(s) URL") + return value + + +class AgentKillSwitchResult(BaseModel): + model_config = ConfigDict(frozen=True) + + agent_id: str + url: str + method: AgentKillSwitchMethod + status_code: int | None = None + response_body: str | None = None + error: str | None = None + + @property + def succeeded(self) -> bool: + return self.status_code is not None and 200 <= self.status_code < 300 + + class AgentConfig(TypedDict, total=False): agent_name: Required[str] agent_card_params: Required[AgentCard] @@ -190,6 +259,7 @@ class AgentConfig(TypedDict, total=False): static_headers: dict[str, str] | None extra_headers: list[str] | None access_group_ids: ReadOnly[Sequence[str] | None] + kill_switch: ReadOnly[AgentKillSwitchConfig | None] class PatchAgentRequest(TypedDict, total=False): @@ -204,6 +274,7 @@ class PatchAgentRequest(TypedDict, total=False): static_headers: dict[str, str] | None extra_headers: list[str] | None access_group_ids: ReadOnly[Sequence[str] | None] + kill_switch: ReadOnly[AgentKillSwitchConfig | None] AGENT_CALLER_USER_ID_HEADER: Final = "x-litellm-user-id" @@ -243,6 +314,7 @@ class AgentResponse(BaseModel): static_headers: dict[str, str] | None = None extra_headers: list[str] | None = None access_group_ids: Sequence[str] | None = None + kill_switch: AgentKillSwitchConfig | None = None keys: list[AgentKeySummary] | None = None search_score: float | None = None created_at: datetime | None = None diff --git a/litellm/types/llms/custom_http.py b/litellm/types/llms/custom_http.py index 06982a16755..fa2d1373ea1 100644 --- a/litellm/types/llms/custom_http.py +++ b/litellm/types/llms/custom_http.py @@ -26,6 +26,7 @@ class httpxSpecialProvider(str, Enum): RAG = "rag" A2AProvider = "a2a_provider" AgentHealthCheck = "agent_health_check" + AgentKillSwitch = "agent_kill_switch" A2A = "a2a" PromptManagement = "prompt_management" UI = "ui" diff --git a/schema.prisma b/schema.prisma index 85996430bc5..69c63d9ecd6 100644 --- a/schema.prisma +++ b/schema.prisma @@ -72,6 +72,7 @@ model LiteLLM_AgentsTable { agent_card_params Json static_headers Json? @default("{}") extra_headers String[] @default([]) + kill_switch Json? agent_access_groups String[] @default([]) access_group_ids String[] @default([]) object_permission_id String? diff --git a/tests/test_litellm/proxy/agent_endpoints/test_agent_registry.py b/tests/test_litellm/proxy/agent_endpoints/test_agent_registry.py index b036e0dac4d..ef20e88c368 100644 --- a/tests/test_litellm/proxy/agent_endpoints/test_agent_registry.py +++ b/tests/test_litellm/proxy/agent_endpoints/test_agent_registry.py @@ -451,7 +451,7 @@ async def test_update_agent_in_db_raises_when_row_deleted_mid_update(): registry: Final = AgentRegistry() mock_prisma: Final = MagicMock() mock_prisma.db.litellm_agentstable.find_unique = AsyncMock( - return_value=SimpleNamespace(litellm_params={}, object_permission_id=None) + return_value=SimpleNamespace(litellm_params={}, object_permission_id=None, kill_switch=None) ) mock_prisma.db.litellm_agentstable.update = AsyncMock(return_value=None) @@ -736,6 +736,7 @@ async def test_update_agent_in_db_preserves_secret_when_echoed_back_redacted(): "model": "bedrock/agentcore/my-agent", }, object_permission_id=None, + kill_switch=None, ) ) updated_agent = MagicMock() @@ -784,6 +785,7 @@ async def test_update_agent_in_db_preserves_secret_when_key_omitted_entirely(): return_value=SimpleNamespace( litellm_params={"aws_secret_access_key": SENTINEL_AWS_SECRET_ACCESS_KEY}, object_permission_id=None, + kill_switch=None, ) ) updated_agent = MagicMock() @@ -830,6 +832,7 @@ async def test_update_agent_in_db_preserves_secret_nested_under_a_non_sensitive_ } }, object_permission_id=None, + kill_switch=None, ) ) updated_agent = MagicMock() @@ -878,6 +881,7 @@ async def test_update_agent_in_db_clears_secret_on_explicit_empty_value(): return_value=SimpleNamespace( litellm_params={"aws_secret_access_key": SENTINEL_AWS_SECRET_ACCESS_KEY}, object_permission_id=None, + kill_switch=None, ) ) updated_agent = MagicMock() @@ -1110,7 +1114,9 @@ async def test_update_agent_in_db_always_writes_access_group_ids(body_access_gro registry: Final = AgentRegistry() mock_prisma: Final = MagicMock() mock_prisma.db.litellm_agentstable.find_unique = AsyncMock( - return_value=SimpleNamespace(litellm_params={}, object_permission_id=None, access_group_ids=["ag-1"]) + return_value=SimpleNamespace( + litellm_params={}, object_permission_id=None, kill_switch=None, access_group_ids=["ag-1"] + ) ) mock_update = AsyncMock(return_value=_agent_row_mock(expected)) mock_prisma.db.litellm_agentstable.update = mock_update @@ -1126,3 +1132,155 @@ async def test_update_agent_in_db_always_writes_access_group_ids(body_access_gro ) assert tuple(mock_update.call_args.kwargs["data"]["access_group_ids"]) == tuple(expected) + + +_KILL_SWITCH: Final = { + "url": "https://ops.example.com/kill", + "method": "POST", + "headers": {"X-Env": "prod"}, + "query_params": {"reason": "manual"}, + "body": {"action": "stop"}, + "auth": {"type": "bearer", "token": "tok-real"}, +} + + +@pytest.mark.asyncio +async def test_add_agent_to_db_stores_kill_switch_json_and_a_json_null_when_unset(): + registry: Final = AgentRegistry() + mock_prisma: Final = MagicMock() + mock_create = AsyncMock(return_value=_agent_row_mock([])) + mock_prisma.db.litellm_agentstable.create = mock_create + + await registry.add_agent_to_db( + agent={ + "agent_name": "Test Agent", + "agent_card_params": _sample_agent_card_params(), + "kill_switch": _KILL_SWITCH, + }, + prisma_client=mock_prisma, + created_by="test-user", + ) + assert json.loads(mock_create.call_args.kwargs["data"]["kill_switch"]) == _KILL_SWITCH + + await registry.add_agent_to_db( + agent={"agent_name": "Plain Agent", "agent_card_params": _sample_agent_card_params()}, + prisma_client=mock_prisma, + created_by="test-user", + ) + assert mock_create.call_args.kwargs["data"]["kill_switch"] == json.dumps(None) + + +@pytest.mark.asyncio +async def test_add_agent_to_db_rejects_a_kill_switch_with_a_non_http_url(): + registry: Final = AgentRegistry() + mock_prisma: Final = MagicMock() + mock_prisma.db.litellm_agentstable.create = AsyncMock(return_value=_agent_row_mock([])) + + with pytest.raises(Exception, match="absolute http"): + await registry.add_agent_to_db( + agent={ + "agent_name": "Test Agent", + "agent_card_params": _sample_agent_card_params(), + "kill_switch": {**_KILL_SWITCH, "url": "ops.example.com/kill"}, + }, + prisma_client=mock_prisma, + created_by="test-user", + ) + mock_prisma.db.litellm_agentstable.create.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_patch_agent_in_db_keeps_kill_switch_when_omitted_and_clears_it_on_null(): + registry: Final = AgentRegistry() + mock_prisma: Final = MagicMock() + mock_prisma.db.litellm_agentstable.find_unique = AsyncMock( + return_value={ + "agent_id": "agent-123", + "agent_name": "Old", + "litellm_params": {}, + "object_permission_id": None, + "kill_switch": _KILL_SWITCH, + } + ) + mock_update = AsyncMock(return_value=_agent_row_mock([])) + mock_prisma.db.litellm_agentstable.update = mock_update + + await registry.patch_agent_in_db( + agent_id="agent-123", agent={"agent_name": "New"}, prisma_client=mock_prisma, updated_by="u" + ) + assert "kill_switch" not in mock_update.call_args.kwargs["data"] + + await registry.patch_agent_in_db( + agent_id="agent-123", agent={"kill_switch": None}, prisma_client=mock_prisma, updated_by="u" + ) + assert mock_update.call_args.kwargs["data"]["kill_switch"] == json.dumps(None), ( + "prisma-client-py silently drops None, so the clear must be written as the JSON literal null" + ) + + +@pytest.mark.asyncio +async def test_patch_agent_in_db_restores_the_stored_kill_switch_secret_behind_the_marker(): + registry: Final = AgentRegistry() + mock_prisma: Final = MagicMock() + mock_prisma.db.litellm_agentstable.find_unique = AsyncMock( + return_value={ + "agent_id": "agent-123", + "agent_name": "A", + "litellm_params": {}, + "object_permission_id": None, + "kill_switch": _KILL_SWITCH, + } + ) + mock_update = AsyncMock(return_value=_agent_row_mock([])) + mock_prisma.db.litellm_agentstable.update = mock_update + + await registry.patch_agent_in_db( + agent_id="agent-123", + agent={ + "kill_switch": { + **_KILL_SWITCH, + "url": "https://ops.example.com/v2/kill", + "auth": {"type": "bearer", "token": REDACTED_BY_LITELM_STRING}, + } + }, + prisma_client=mock_prisma, + updated_by="u", + ) + + assert json.loads(mock_update.call_args.kwargs["data"]["kill_switch"]) == { + **_KILL_SWITCH, + "url": "https://ops.example.com/v2/kill", + } + + +@pytest.mark.asyncio +async def test_update_agent_in_db_clears_kill_switch_when_omitted_and_restores_secret_when_echoed(): + registry: Final = AgentRegistry() + mock_prisma: Final = MagicMock() + mock_prisma.db.litellm_agentstable.find_unique = AsyncMock( + return_value=SimpleNamespace(litellm_params={}, object_permission_id=None, kill_switch=json.dumps(_KILL_SWITCH)) + ) + mock_update = AsyncMock(return_value=_agent_row_mock([])) + mock_prisma.db.litellm_agentstable.update = mock_update + base: Final = {"agent_name": "Test Agent", "agent_card_params": _sample_agent_card_params(), "litellm_params": {}} + + await registry.update_agent_in_db(agent_id="agent-123", agent=base, prisma_client=mock_prisma, updated_by="u") + assert mock_update.call_args.kwargs["data"]["kill_switch"] == json.dumps(None) + + echoed: Final = {**_KILL_SWITCH, "auth": {"type": "bearer", "token": REDACTED_BY_LITELM_STRING}} + await registry.update_agent_in_db( + agent_id="agent-123", agent={**base, "kill_switch": echoed}, prisma_client=mock_prisma, updated_by="u" + ) + assert json.loads(mock_update.call_args.kwargs["data"]["kill_switch"]) == _KILL_SWITCH + + +def test_load_agents_from_config_exposes_a_typed_kill_switch(): + registry: Final = AgentRegistry() + + registry.load_agents_from_config( + [{"agent_name": "cfg-agent", "agent_card_params": _sample_agent_card_params(), "kill_switch": _KILL_SWITCH}] + ) + + (agent,) = registry.get_agent_list() + assert agent.kill_switch is not None + assert agent.kill_switch.model_dump() == _KILL_SWITCH diff --git a/tests/test_litellm/proxy/agent_endpoints/test_endpoints.py b/tests/test_litellm/proxy/agent_endpoints/test_endpoints.py index 482294e7b92..526f24c5221 100644 --- a/tests/test_litellm/proxy/agent_endpoints/test_endpoints.py +++ b/tests/test_litellm/proxy/agent_endpoints/test_endpoints.py @@ -1,13 +1,15 @@ import json +from types import SimpleNamespace from typing import Final from unittest.mock import AsyncMock, MagicMock, patch +import httpx import pytest from fastapi import FastAPI from fastapi.testclient import TestClient from litellm.constants import REDACTED_BY_LITELM_STRING -from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth +from litellm.proxy._types import LiteLLM_AuditLogs, LitellmTableNames, LitellmUserRoles, UserAPIKeyAuth from litellm.proxy.agent_endpoints import endpoints as agent_endpoints from litellm.proxy.agent_endpoints.auth.agent_permission_handler import ( RestrictedAgentAccess, @@ -1136,3 +1138,222 @@ def test_make_agent_public_rejects_an_agent_published_only_in_the_db(monkeypatch assert duplicate.status_code == 400 assert "already in public agent groups" in duplicate.json()["detail"] + + +_KILL_SWITCH: Final = { + "url": "https://ops.example.com/kill", + "method": "POST", + "headers": {"X-Env": "prod"}, + "query_params": {"reason": "manual"}, + "body": {"action": "stop"}, + "auth": {"type": "bearer", "token": "tok-real"}, +} + + +def _agent_with_kill_switch() -> AgentResponse: + return AgentResponse( + agent_id="agent-123", + agent_name="Test Agent", + agent_card_params=_sample_agent_card_params(), + litellm_params={}, + kill_switch=_KILL_SWITCH, + ) + + +class _FakeKillSwitchClient: + def __init__(self, response: httpx.Response) -> None: + self.calls: list[tuple[str, str, dict[str, str], object, float]] = [] # mutable-ok: test double records calls + self._response: Final = response + + def build_request(self, method: str, url: str, *, headers, json, timeout: float) -> httpx.Request: + self.calls.append((method, url, dict(headers), json, timeout)) + return httpx.Request(method, url, headers=dict(headers), json=json) + + async def send(self, request: httpx.Request, *, stream: bool, follow_redirects: bool) -> httpx.Response: + return self._response + + +class _AuditLogRecorder: + def __init__(self) -> None: + self.rows: list[LiteLLM_AuditLogs] = [] # mutable-ok: test double records writes + + async def __call__(self, request_data: LiteLLM_AuditLogs) -> None: + self.rows.append(request_data) + + +def _kill_switch_app( + role: LitellmUserRoles, + http_client: _FakeKillSwitchClient, + audit_log: _AuditLogRecorder | None = None, +) -> TestClient: + test_client: Final = _make_app_with_role(role) + test_client.app.dependency_overrides[agent_endpoints.default_kill_switch_http_client] = lambda: http_client + test_client.app.dependency_overrides[agent_endpoints.default_kill_switch_audit_log_writer] = ( + lambda: audit_log or _AuditLogRecorder() + ) + return test_client + + +def test_kill_switch_trigger_fires_the_configured_webhook_and_returns_the_result(monkeypatch) -> None: + registry: Final = MagicMock() + registry.get_agent_by_id = MagicMock(return_value=_agent_with_kill_switch()) + monkeypatch.setattr(agent_endpoints, "AGENT_REGISTRY", registry) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None) + fake: Final = _FakeKillSwitchClient(httpx.Response(200, text="ok")) + + resp: Final = _kill_switch_app(LitellmUserRoles.PROXY_ADMIN, fake).post( + "/v1/agents/agent-123/kill_switch", headers={"Authorization": "Bearer k"} + ) + + assert resp.status_code == 200, resp.text + assert resp.json() == { + "agent_id": "agent-123", + "url": "https://ops.example.com/kill", + "method": "POST", + "status_code": 200, + "response_body": "ok", + "error": None, + } + (method, url, headers, body, _timeout) = fake.calls[0] + assert (method, url, body) == ("POST", "https://ops.example.com/kill?reason=manual", {"action": "stop"}) + assert headers == {"X-Env": "prod", "Authorization": "Bearer tok-real"} + + +def test_kill_switch_trigger_writes_an_audit_log_row_naming_the_admin_and_the_sanitized_result(monkeypatch) -> None: + registry: Final = MagicMock() + registry.get_agent_by_id = MagicMock(return_value=_agent_with_kill_switch()) + monkeypatch.setattr(agent_endpoints, "AGENT_REGISTRY", registry) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None) + fake: Final = _FakeKillSwitchClient(httpx.Response(202, text='{"stopped": true}')) + audit: Final = _AuditLogRecorder() + test_client: Final = _kill_switch_app(LitellmUserRoles.PROXY_ADMIN, fake, audit) + test_client.app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth( + user_id="test-user", user_role=LitellmUserRoles.PROXY_ADMIN, api_key="hashed-k" + ) + + resp: Final = test_client.post("/v1/agents/agent-123/kill_switch", headers={"Authorization": "Bearer k"}) + + assert resp.status_code == 200, resp.text + (row,) = audit.rows + assert (row.action, row.table_name, row.object_id) == ( + "kill_switch_fired", + LitellmTableNames.AGENT_TABLE_NAME, + "agent-123", + ) + assert (row.changed_by, row.changed_by_api_key) == ("test-user", "hashed-k") + assert row.before_value is None + assert json.loads(row.updated_values) == { + "agent_id": "agent-123", + "url": "https://ops.example.com/kill", + "method": "POST", + "status_code": 202, + "response_body": '{"stopped": true}', + } + assert "tok-real" not in row.model_dump_json() + + +def test_kill_switch_trigger_returns_502_and_still_audits_when_the_webhook_rejects(monkeypatch) -> None: + registry: Final = MagicMock() + registry.get_agent_by_id = MagicMock(return_value=_agent_with_kill_switch()) + monkeypatch.setattr(agent_endpoints, "AGENT_REGISTRY", registry) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None) + fake: Final = _FakeKillSwitchClient(httpx.Response(401, text="bad token")) + audit: Final = _AuditLogRecorder() + + resp: Final = _kill_switch_app(LitellmUserRoles.PROXY_ADMIN, fake, audit).post( + "/v1/agents/agent-123/kill_switch", headers={"Authorization": "Bearer k"} + ) + + assert resp.status_code == 502, resp.text + assert resp.json()["detail"]["status_code"] == 401 + assert resp.json()["detail"]["response_body"] == "bad token" + (row,) = audit.rows + assert row.action == "kill_switch_fired" + assert json.loads(row.updated_values)["status_code"] == 401 + + +@pytest.mark.parametrize("role", [LitellmUserRoles.INTERNAL_USER, LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY]) +def test_kill_switch_trigger_is_refused_before_any_webhook_call_for_non_admins(monkeypatch, role) -> None: + registry: Final = MagicMock() + registry.get_agent_by_id = MagicMock(return_value=_agent_with_kill_switch()) + monkeypatch.setattr(agent_endpoints, "AGENT_REGISTRY", registry) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None) + fake: Final = _FakeKillSwitchClient(httpx.Response(200)) + audit: Final = _AuditLogRecorder() + + resp: Final = _kill_switch_app(role, fake, audit).post( + "/v1/agents/agent-123/kill_switch", headers={"Authorization": "Bearer k"} + ) + + assert resp.status_code == 403, resp.text + assert fake.calls == [] + assert audit.rows == [] + + +def test_kill_switch_trigger_404s_unknown_agent_and_400s_an_agent_without_one(monkeypatch) -> None: + registry: Final = MagicMock() + registry.get_agent_by_id = MagicMock(side_effect=[None, _sample_agent_response()]) + monkeypatch.setattr(agent_endpoints, "AGENT_REGISTRY", registry) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None) + fake: Final = _FakeKillSwitchClient(httpx.Response(200)) + audit: Final = _AuditLogRecorder() + test_client: Final = _kill_switch_app(LitellmUserRoles.PROXY_ADMIN, fake, audit) + + missing: Final = test_client.post("/v1/agents/nope/kill_switch", headers={"Authorization": "Bearer k"}) + unconfigured: Final = test_client.post("/v1/agents/agent-123/kill_switch", headers={"Authorization": "Bearer k"}) + + assert missing.status_code == 404 + assert unconfigured.status_code == 400 + assert "no kill_switch configured" in unconfigured.json()["detail"] + assert fake.calls == [] + assert audit.rows == [] + + +def test_kill_switch_trigger_fires_the_db_row_config_over_a_stale_in_memory_copy(monkeypatch) -> None: + """Another replica may have updated the agent; the row is the source of truth for what gets fired.""" + registry: Final = MagicMock() + registry.get_agent_by_id = MagicMock(return_value=_agent_with_kill_switch()) + monkeypatch.setattr(agent_endpoints, "AGENT_REGISTRY", registry) + db_row: Final = SimpleNamespace( + agent_id="agent-123", + kill_switch={"url": "https://ops.example.com/kill-v2", "method": "DELETE", "auth": None}, + ) + prisma: Final = MagicMock() + prisma.db.litellm_agentstable.find_unique = AsyncMock(return_value=db_row) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", prisma) + fake: Final = _FakeKillSwitchClient(httpx.Response(204)) + + resp: Final = _kill_switch_app(LitellmUserRoles.PROXY_ADMIN, fake).post( + "/v1/agents/agent-123/kill_switch", headers={"Authorization": "Bearer k"} + ) + + assert resp.status_code == 200, resp.text + (method, url, headers, body, _timeout) = fake.calls[0] + assert (method, url, headers, body) == ("DELETE", "https://ops.example.com/kill-v2", {}, None) + assert prisma.db.litellm_agentstable.find_unique.await_args.kwargs == {"where": {"agent_id": "agent-123"}} + registry.get_agent_by_id.assert_not_called() + + +def test_get_agent_redacts_kill_switch_secret_for_admins_and_hides_it_from_others(monkeypatch) -> None: + registry: Final = MagicMock() + registry.get_agent_by_id = MagicMock(return_value=_agent_with_kill_switch()) + registry.ids_for_agent = MagicMock(return_value=("agent-123",)) + monkeypatch.setattr(agent_endpoints, "AGENT_REGISTRY", registry) + + def _get_as(role: LitellmUserRoles): + with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma: + mock_prisma.db.litellm_agentstable.find_unique = AsyncMock(return_value=None) + mock_prisma.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[]) + return _make_app_with_role(role).get("/v1/agents/agent-123", headers={"Authorization": "Bearer k"}) + + admin: Final = _get_as(LitellmUserRoles.PROXY_ADMIN) + assert admin.status_code == 200, admin.text + assert admin.json()["kill_switch"] == { + **_KILL_SWITCH, + "auth": {"type": "bearer", "token": REDACTED_BY_LITELM_STRING}, + } + + internal: Final = _get_as(LitellmUserRoles.INTERNAL_USER) + assert internal.status_code == 200, internal.text + assert internal.json()["kill_switch"] is None + assert "tok-real" not in internal.text diff --git a/tests/test_litellm/proxy/agent_endpoints/test_kill_switch.py b/tests/test_litellm/proxy/agent_endpoints/test_kill_switch.py new file mode 100644 index 00000000000..a6bb945713e --- /dev/null +++ b/tests/test_litellm/proxy/agent_endpoints/test_kill_switch.py @@ -0,0 +1,248 @@ +from base64 import b64encode +from collections.abc import Mapping +from dataclasses import dataclass +from typing import Final + +import httpx +import pytest +from pydantic import ValidationError + +from litellm.constants import REDACTED_BY_LITELM_STRING +from litellm.proxy.agent_endpoints.kill_switch import ( + build_kill_switch_request, + fire_kill_switch, + redact_kill_switch, + restore_kill_switch, +) +from litellm.types.agents import AgentKillSwitchConfig + + +@dataclass(frozen=True, slots=True) +class _SentRequest: + method: str + url: str + headers: Mapping[str, str] + json: Mapping[str, object] | None + timeout: float + + +class _RecordingClient: + def __init__(self, respond: httpx.Response | httpx.HTTPError) -> None: + self.sent: list[_SentRequest] = [] # mutable-ok: test double records calls + self.follow_redirects: list[bool] = [] # mutable-ok: test double records calls + self._respond: Final = respond + + def build_request( + self, + method: str, + url: str, + *, + headers: Mapping[str, str], + json: Mapping[str, object] | None, + timeout: float, + ) -> httpx.Request: + self.sent.append(_SentRequest(method, url, headers, json, timeout)) + return httpx.Request(method, url, headers=dict(headers), json=json) + + async def send(self, request: httpx.Request, *, stream: bool, follow_redirects: bool) -> httpx.Response: + self.follow_redirects.append(follow_redirects) + if isinstance(self._respond, httpx.HTTPError): + raise self._respond + return self._respond + + +class _CountingStream(httpx.AsyncByteStream): + def __init__(self, chunk: bytes, chunks: int) -> None: + self.pulled: int = 0 # rebind-ok: test double counts reads + self._chunk: Final = chunk + self._chunks: Final = chunks + + async def __aiter__(self): + for _ in range(self._chunks): + self.pulled += 1 # rebind-ok: test double counts reads + yield self._chunk + + +def _config(**overrides: object) -> AgentKillSwitchConfig: + return AgentKillSwitchConfig.model_validate({"url": "https://ops.example.com/agents/kill", **overrides}) + + +def test_request_carries_endpoint_method_query_params_headers_and_body() -> None: + request: Final = build_kill_switch_request( + _config( + url="https://ops.example.com/kill?env=prod", + method="PUT", + query_params={"agent": "billing-bot", "reason": "manual stop"}, + headers={"X-Trace": "abc"}, + body={"action": "stop", "hard": True}, + ) + ) + + assert request.method == "PUT" + assert str(httpx.URL(request.url)) == "https://ops.example.com/kill?env=prod&agent=billing-bot&reason=manual+stop" + assert dict(request.headers) == {"X-Trace": "abc"} + assert request.json_body == {"action": "stop", "hard": True} + + +def test_request_defaults_to_post_with_no_body_and_untouched_url() -> None: + request: Final = build_kill_switch_request(_config()) + + assert (request.method, request.url, dict(request.headers), request.json_body) == ( + "POST", + "https://ops.example.com/agents/kill", + {}, + None, + ) + + +@pytest.mark.parametrize( + ("auth", "expected_headers"), + [ + ({"type": "bearer", "token": "tok-123"}, {"Authorization": "Bearer tok-123"}), + ({"type": "api_key", "api_key": "k-456"}, {"x-api-key": "k-456"}), + ({"type": "api_key", "header_name": "X-Ops-Key", "api_key": "k-456"}, {"X-Ops-Key": "k-456"}), + ( + {"type": "basic", "username": "ops", "password": "pw:1"}, + {"Authorization": f"Basic {b64encode(b'ops:pw:1').decode()}"}, + ), + ], +) +def test_auth_becomes_the_matching_request_header(auth: Mapping[str, object], expected_headers: dict[str, str]) -> None: + request: Final = build_kill_switch_request(_config(auth=auth)) + + assert dict(request.headers) == expected_headers + + +def test_auth_header_wins_over_a_conflicting_custom_header() -> None: + request: Final = build_kill_switch_request( + _config(headers={"Authorization": "stale", "X-Env": "prod"}, auth={"type": "bearer", "token": "fresh"}) + ) + + assert dict(request.headers) == {"Authorization": "Bearer fresh", "X-Env": "prod"} + + +@pytest.mark.parametrize("url", ["ftp://ops.example.com/kill", "/relative/kill", "ops.example.com/kill", ""]) +def test_config_rejects_non_http_urls(url: str) -> None: + with pytest.raises(ValidationError, match="absolute http"): + _config(url=url) + + +def test_config_rejects_unknown_auth_type_and_unknown_fields() -> None: + with pytest.raises(ValidationError): + _config(auth={"type": "hmac", "secret": "x"}) + with pytest.raises(ValidationError): + _config(endpoint="https://typo.example.com") + + +@pytest.mark.parametrize( + ("auth", "secret_field"), + [ + ({"type": "bearer", "token": "tok-123"}, "token"), + ({"type": "api_key", "header_name": "X-K", "api_key": "k-456"}, "api_key"), + ({"type": "basic", "username": "ops", "password": "pw"}, "password"), + ], +) +def test_redact_replaces_only_the_secret_and_restore_puts_it_back(auth: dict[str, str], secret_field: str) -> None: + original: Final = _config(auth=auth) + + redacted: Final = redact_kill_switch(original) + assert redacted is not None and redacted.auth is not None + assert redacted.auth.model_dump() == {**auth, secret_field: REDACTED_BY_LITELM_STRING} + assert original.auth is not None and original.auth.model_dump() == auth, "redact must not mutate its input" + + restored: Final = restore_kill_switch(redacted, original) + assert restored == original + + +def test_restore_keeps_a_rotated_secret_and_never_stores_the_marker_itself() -> None: + rotated: Final = _config(auth={"type": "bearer", "token": "new-token"}) + stored: Final = _config(auth={"type": "bearer", "token": "old-token"}) + assert restore_kill_switch(rotated, stored) == rotated + assert restore_kill_switch(None, stored) is None + + marker_only: Final = _config(auth={"type": "bearer", "token": REDACTED_BY_LITELM_STRING}) + assert restore_kill_switch(marker_only, None) == _config(auth={"type": "bearer", "token": ""}) + + +def test_restore_does_not_borrow_a_secret_from_a_different_auth_type() -> None: + incoming: Final = _config(auth={"type": "bearer", "token": REDACTED_BY_LITELM_STRING}) + stored: Final = _config(auth={"type": "api_key", "api_key": "k-456"}) + + assert restore_kill_switch(incoming, stored) == _config(auth={"type": "bearer", "token": ""}) + + +def test_redact_passes_through_configs_without_auth() -> None: + assert redact_kill_switch(None) is None + plain: Final = _config(headers={"X-Env": "prod"}) + assert redact_kill_switch(plain) is plain + + +@pytest.mark.asyncio +async def test_fire_sends_exactly_the_built_request_and_reports_the_2xx_reply_without_the_query() -> None: + client: Final = _RecordingClient(httpx.Response(202, text="stopping")) + config: Final = _config( + method="DELETE", + query_params={"force": "1", "token": "qs-secret"}, + headers={"X-Env": "prod"}, + body={"agent": "billing-bot"}, + auth={"type": "bearer", "token": "tok-123"}, + ) + + result: Final = await fire_kill_switch(agent_id="agent-1", config=config, http_client=client, timeout=3.5) + + assert client.sent == [ + _SentRequest( + method="DELETE", + url="https://ops.example.com/agents/kill?force=1&token=qs-secret", + headers={"X-Env": "prod", "Authorization": "Bearer tok-123"}, + json={"agent": "billing-bot"}, + timeout=3.5, + ) + ] + assert client.follow_redirects == [False], "a redirecting webhook must not be followed to another host" + assert result.succeeded is True + assert result.model_dump() == { + "agent_id": "agent-1", + "url": "https://ops.example.com/agents/kill", + "method": "DELETE", + "status_code": 202, + "response_body": "stopping", + "error": None, + } + + +@pytest.mark.asyncio +async def test_fire_reports_a_non_2xx_reply_as_failure_with_the_body() -> None: + client: Final = _RecordingClient(httpx.Response(503, text="x" * 5000)) + + result: Final = await fire_kill_switch(agent_id="agent-1", config=_config(), http_client=client) + + assert result.succeeded is False + assert result.status_code == 503 + assert result.response_body == "x" * 2000 + assert result.error is None + + +@pytest.mark.asyncio +async def test_fire_stops_reading_the_body_at_the_cap_instead_of_buffering_the_whole_reply() -> None: + stream: Final = _CountingStream(b"y" * 500, chunks=100) + client: Final = _RecordingClient(httpx.Response(200, stream=stream)) + + result: Final = await fire_kill_switch(agent_id="agent-1", config=_config(), http_client=client) + + assert result.response_body == "y" * 2000 + assert stream.pulled == 4, f"read {stream.pulled} of 100 chunks for a 2000 char cap" + + +@pytest.mark.asyncio +async def test_fire_reports_a_transport_error_by_type_without_raising_or_echoing_the_url() -> None: + client: Final = _RecordingClient(httpx.ConnectError("boom https://ops.example.com/agents/kill?token=qs-secret")) + + result: Final = await fire_kill_switch( + agent_id="agent-1", config=_config(query_params={"token": "qs-secret"}), http_client=client + ) + + assert result.succeeded is False + assert (result.status_code, result.response_body) == (None, None) + assert result.error == "ConnectError" + assert "qs-secret" not in result.model_dump_json() diff --git a/tests/test_litellm/proxy/auth/test_route_checks.py b/tests/test_litellm/proxy/auth/test_route_checks.py index 7bb79a115dd..f76a02e8361 100644 --- a/tests/test_litellm/proxy/auth/test_route_checks.py +++ b/tests/test_litellm/proxy/auth/test_route_checks.py @@ -3776,6 +3776,7 @@ AGENT_MANAGEMENT_ROUTES = [ "/v1/agents/abc-123", "/v1/agents/make_public", "/v1/agents/abc-123/make_public", + "/v1/agents/abc-123/kill_switch", ] AGENT_INFERENCE_ROUTES = [ diff --git a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/AgentFormKit.tsx b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/AgentFormKit.tsx index 8e100d0c3ed..3d863036234 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/AgentFormKit.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/AgentFormKit.tsx @@ -27,6 +27,7 @@ import { Collapsible, CollapsibleContent, CollapsibleTrigger } from "@/component import { Input } from "@/components/ui/input"; import { Tooltip, TooltipContent, TooltipTrigger } from "@/components/ui/tooltip"; import { Field, FieldDescription, FieldError, FieldGroup, FieldLabel } from "@/components/ui/field"; +import type { KeyValueFormValue, KillSwitchConfig, KillSwitchFormValue } from "./kill_switch_config"; export interface AgentSkillFormValue { id?: string; @@ -54,6 +55,8 @@ export type AgentFormFieldValue = | string[] | AgentSkillFormValue[] | StaticHeaderFormValue[] + | KeyValueFormValue[] + | KillSwitchFormValue | McpServerSelection | Record | null @@ -82,6 +85,7 @@ export interface AgentFormValues { output_cost_per_token?: string | number; static_headers?: StaticHeaderFormValue[]; extra_headers?: string[]; + kill_switch?: KillSwitchFormValue; tpm_limit?: number | null; rpm_limit?: number | null; session_tpm_limit?: number | null; @@ -123,6 +127,7 @@ export interface AgentRequestPayload { litellm_params?: Record; object_permission?: Record; access_group_ids?: string[]; + kill_switch?: KillSwitchConfig | null; } interface AgentFormFieldProps { diff --git a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/AgentKillSwitchDangerZone.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/AgentKillSwitchDangerZone.test.tsx new file mode 100644 index 00000000000..aa4b83478b6 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/AgentKillSwitchDangerZone.test.tsx @@ -0,0 +1,124 @@ +import React from "react"; +import { fireEvent, render, screen, waitFor, within } from "@testing-library/react"; +import { describe, it, expect, vi, beforeEach } from "vitest"; +import AgentKillSwitchDangerZone from "./AgentKillSwitchDangerZone"; +import * as networking from "@/components/networking"; +import { toast } from "@/lib/toast"; + +vi.mock("@/components/networking", () => ({ + triggerAgentKillSwitchCall: vi.fn(), +})); + +vi.mock("@/lib/toast", () => ({ + toast: { success: vi.fn(), error: vi.fn() }, +})); + +const killSwitch = { url: "https://ops.example.com/kill", method: "DELETE" as const }; + +const renderZone = (props: Partial> = {}) => + render( + , + ); + +const openDialog = () => { + fireEvent.click(screen.getByRole("button", { name: "Fire Kill Switch" })); + return screen.getByRole("dialog"); +}; + +const dialogFireButton = () => within(screen.getByRole("dialog")).getByRole("button", { name: "Fire Kill Switch" }); + +describe("AgentKillSwitchDangerZone", () => { + beforeEach(() => { + vi.mocked(networking.triggerAgentKillSwitchCall).mockReset(); + vi.mocked(toast.success).mockReset(); + vi.mocked(toast.error).mockReset(); + }); + + it("renders nothing for non-admins", () => { + const { container } = renderZone({ isAdmin: false }); + + expect(container).toBeEmptyDOMElement(); + }); + + it("shows the webhook target and an outage warning inside a Danger Zone region", () => { + renderZone(); + + const region = screen.getByRole("region", { name: "Danger Zone" }); + expect(region).toHaveTextContent("DELETE https://ops.example.com/kill"); + expect(region).toHaveTextContent("can cause an outage"); + expect(screen.getByRole("button", { name: "Fire Kill Switch" })).toBeEnabled(); + }); + + it("shows an unconfigured notice without a fire button when no kill switch is set", () => { + renderZone({ killSwitch: null }); + + expect(screen.getByRole("region", { name: "Danger Zone" })).toHaveTextContent("Not configured"); + expect(screen.queryByRole("button", { name: "Fire Kill Switch" })).not.toBeInTheDocument(); + }); + + it("keeps the confirm button disabled until the exact agent name is typed", () => { + renderZone(); + openDialog(); + + expect(dialogFireButton()).toBeDisabled(); + + fireEvent.change(screen.getByLabelText("Confirm agent name"), { target: { value: "support-agen" } }); + expect(dialogFireButton()).toBeDisabled(); + + fireEvent.change(screen.getByLabelText("Confirm agent name"), { target: { value: "support-agent" } }); + expect(dialogFireButton()).toBeEnabled(); + expect(networking.triggerAgentKillSwitchCall).not.toHaveBeenCalled(); + }); + + it("fires the webhook after typed confirmation, closes the dialog and shows the sanitized result", async () => { + const firedResult = { + agent_id: "agent-1", + url: killSwitch.url, + method: "DELETE" as const, + status_code: 202, + response_body: '{"stopped": true}', + }; + vi.mocked(networking.triggerAgentKillSwitchCall).mockResolvedValue(firedResult); + renderZone(); + openDialog(); + + fireEvent.change(screen.getByLabelText("Confirm agent name"), { target: { value: "support-agent" } }); + fireEvent.click(dialogFireButton()); + + expect(await screen.findByRole("status")).toHaveTextContent('Last result: HTTP 202 {"stopped": true}'); + await waitFor(() => expect(screen.queryByRole("dialog")).not.toBeInTheDocument()); + expect(networking.triggerAgentKillSwitchCall).toHaveBeenCalledWith("sk-test", "agent-1"); + expect(toast.success).toHaveBeenCalledWith("Kill switch fired (HTTP 202)"); + }); + + it("does not call the webhook when the dialog is cancelled", () => { + renderZone(); + openDialog(); + + fireEvent.change(screen.getByLabelText("Confirm agent name"), { target: { value: "support-agent" } }); + fireEvent.click(screen.getByRole("button", { name: "Cancel" })); + + expect(networking.triggerAgentKillSwitchCall).not.toHaveBeenCalled(); + expect(screen.queryByRole("status")).not.toBeInTheDocument(); + }); + + it("surfaces a failed webhook as an error toast and keeps the dialog open", async () => { + vi.mocked(networking.triggerAgentKillSwitchCall).mockRejectedValue(new Error("Kill switch webhook returned 500")); + renderZone(); + openDialog(); + + fireEvent.change(screen.getByLabelText("Confirm agent name"), { target: { value: "support-agent" } }); + fireEvent.click(dialogFireButton()); + + await waitFor(() => expect(toast.error).toHaveBeenCalledWith("Kill switch webhook returned 500")); + expect(screen.getByRole("dialog")).toBeInTheDocument(); + expect(screen.queryByRole("status")).not.toBeInTheDocument(); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/AgentKillSwitchDangerZone.tsx b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/AgentKillSwitchDangerZone.tsx new file mode 100644 index 00000000000..8d68a8f921b --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/AgentKillSwitchDangerZone.tsx @@ -0,0 +1,147 @@ +import { CircleAlert } from "lucide-react"; +import React, { useState } from "react"; +import { Alert, AlertDescription, AlertTitle } from "@/components/shared/Alert"; +import { Button } from "@/components/ui/button"; +import { Dialog, DialogContent, DialogFooter, DialogHeader, DialogTitle } from "@/components/ui/dialog"; +import { InputGroup, InputGroupAddon, InputGroupInput } from "@/components/ui/input-group"; +import { UiLoadingSpinner } from "@/components/ui/ui-loading-spinner"; +import { toast } from "@/lib/toast"; +import { AgentKillSwitchResult, triggerAgentKillSwitchCall } from "@/components/networking"; +import { KillSwitchConfig } from "./kill_switch_config"; + +interface AgentKillSwitchDangerZoneProps { + agentId: string; + agentName: string; + killSwitch: KillSwitchConfig | null | undefined; + accessToken: string | null; + isAdmin: boolean; +} + +const AgentKillSwitchDangerZone: React.FC = ({ + agentId, + agentName, + killSwitch, + accessToken, + isAdmin, +}) => { + const [isConfirmOpen, setIsConfirmOpen] = useState(false); + const [confirmationInput, setConfirmationInput] = useState(""); + const [isFiring, setIsFiring] = useState(false); + const [lastResult, setLastResult] = useState(null); + + if (!isAdmin) return null; + + const openConfirm = () => { + setConfirmationInput(""); + setIsConfirmOpen(true); + }; + + const fire = async () => { + if (!accessToken) return; + setIsFiring(true); + setLastResult(null); + try { + const result = await triggerAgentKillSwitchCall(accessToken, agentId); + setLastResult(result); + setIsConfirmOpen(false); + toast.success(`Kill switch fired (HTTP ${result.status_code})`); + } catch (error) { + toast.error(error instanceof Error ? error.message : "Failed to fire kill switch"); + } finally { + setIsFiring(false); + } + }; + + return ( +
+

+ Danger Zone +

+
+
+
+

Kill switch

+ {killSwitch ? ( + <> +

+ Calls the configured webhook to stop this agent's upstream runtime. This can cause an outage for + everyone using the agent and cannot be undone from LiteLLM +

+

+ {killSwitch.method ?? "POST"} {killSwitch.url} +

+ + ) : ( +

+ Not configured. Add a kill switch webhook under Settings to enable this action +

+ )} +
+ {killSwitch && ( + + )} +
+ {lastResult && ( +

+ Last result: HTTP {lastResult.status_code} + {lastResult.response_body ? ` ${lastResult.response_body}` : ""} +

+ )} +
+ + !open && !isFiring && setIsConfirmOpen(false)}> + + + Fire kill switch for {agentName}? + +
+ + + This can cause an outage + + LiteLLM will call {killSwitch?.method ?? "POST"} {killSwitch?.url} immediately. Whatever that webhook + does to the agent is outside LiteLLM's control and cannot be reverted here + + +
+

+ Type {agentName} to confirm: +

+ + + + + setConfirmationInput(e.target.value)} + placeholder={agentName} + autoFocus + /> + +
+
+ + + + + + +
+ ); +}; + +export default AgentKillSwitchDangerZone; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/KillSwitchFormFields.tsx b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/KillSwitchFormFields.tsx new file mode 100644 index 00000000000..578d80da6ed --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/KillSwitchFormFields.tsx @@ -0,0 +1,223 @@ +import React from "react"; +import { useFieldArray, useFormContext, useWatch } from "react-hook-form"; +import { Plus, Trash2 } from "lucide-react"; +import { Button } from "@/components/ui/button"; +import { Input } from "@/components/ui/input"; +import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select"; +import { Textarea } from "@/components/ui/textarea"; +import { Field, FieldTitle } from "@/components/ui/field"; +import { PasswordInput } from "@/components/shared/PasswordInput"; +import { AgentFormField, AgentFormValues, labelWithHint } from "./AgentFormKit"; +import { KILL_SWITCH_AUTH_TYPES, KILL_SWITCH_METHODS, validateKillSwitchBody } from "./kill_switch_config"; + +const KeyValueFieldArray = ({ + name, + addLabel, + keyPlaceholder, + valuePlaceholder, +}: { + name: "kill_switch.headers" | "kill_switch.query_params"; + addLabel: string; + keyPlaceholder: string; + valuePlaceholder: string; +}) => { + const { control } = useFormContext(); + const { fields, append, remove } = useFieldArray({ control, name }); + + return ( +
+ {fields.map((item, index) => ( +
+ + {({ value, onChange, ref, ...control }) => ( + + )} + + + {({ value, onChange, ref, ...control }) => ( + + )} + + +
+ ))} + +
+ ); +}; + +const TextField = ({ + name, + label, + placeholder, + required, + secret, +}: { + name: `kill_switch.${string}`; + label: React.ReactNode; + placeholder?: string; + required?: string; + secret?: boolean; +}) => ( + + {({ value, onChange, ref, ...control }) => + secret ? ( + + ) : ( + + ) + } + +); + +const KillSwitchAuthFields = () => { + const { control } = useFormContext(); + const authType = useWatch({ control, name: "kill_switch.auth_type" }); + + switch (authType) { + case "bearer": + return ; + case "api_key": + return ( + <> + + + + ); + case "basic": + return ( + <> + + + + ); + default: + return null; + } +}; + +const KillSwitchFormFields = () => ( + <> + + + + {({ value, onChange, ref: _ref, ...control }) => ( + + )} + + + + Headers + + + + + Query Parameters + + + + validateKillSwitchBody(typeof value === "string" ? value : "") }} + > + {({ value, onChange, ref, ...control }) => ( +